From d92632cddea97335d342f4eff34fbf7d462b9186 Mon Sep 17 00:00:00 2001 From: cuipengfei Date: Wed, 5 Aug 2026 00:16:22 +0800 Subject: [PATCH] fix(advisor): preserve split update safety --- .../coding-agent/src/advisor/delta-split.ts | 16 +- packages/coding-agent/src/advisor/runtime.ts | 23 +-- .../coding-agent/test/advisor/advisor.test.ts | 143 +++++++++++++++++- .../advisor/delta-split-obfuscation.test.ts | 20 +++ .../test/advisor/delta-split.test.ts | 16 ++ 5 files changed, 204 insertions(+), 14 deletions(-) diff --git a/packages/coding-agent/src/advisor/delta-split.ts b/packages/coding-agent/src/advisor/delta-split.ts index 28cfbe602..4147ea1a6 100644 --- a/packages/coding-agent/src/advisor/delta-split.ts +++ b/packages/coding-agent/src/advisor/delta-split.ts @@ -73,13 +73,23 @@ export function renderAdvisorDeltaChunks( // marker append below is a plain field access (no double-cast). const chunks: { role: "user"; content: TextContent[]; timestamp: number }[] = []; for (let i = 0; i < delta.length; i++) { - let text = renderChunk([delta[i]]); + const text = renderChunk([delta[i]]); if (!text.trim()) continue; - if (opts.obfuscator) text = opts.obfuscator.obfuscate(text, opts.advisorRegexSecretValues); - if (i === 0) text = `${heading}\n\n${text}`; chunks.push({ role: "user", content: [{ type: "text", text }], timestamp: Date.now() }); } if (chunks.length === 0) return null; + if (opts.obfuscator) { + const fullText = chunks.map(chunk => chunk.content[0].text).join("\n"); + const individuallyObfuscated = chunks.map(chunk => + opts.obfuscator!.obfuscate(chunk.content[0].text, opts.advisorRegexSecretValues), + ); + if (opts.obfuscator.obfuscate(fullText, opts.advisorRegexSecretValues) !== individuallyObfuscated.join("\n")) { + return null; + } + for (let i = 0; i < chunks.length; i++) chunks[i].content[0].text = individuallyObfuscated[i]; + } + chunks[0].content[0].text = `${heading}\n\n${chunks[0].content[0].text}`; + if (chunks.length === 0) return null; if (opts.wip) { const last = chunks[chunks.length - 1]; last.content[0].text += `\n\n---\n\n[in progress — more steps follow]`; diff --git a/packages/coding-agent/src/advisor/runtime.ts b/packages/coding-agent/src/advisor/runtime.ts index 9684c357b..501810337 100644 --- a/packages/coding-agent/src/advisor/runtime.ts +++ b/packages/coding-agent/src/advisor/runtime.ts @@ -882,16 +882,19 @@ export class AdvisorRuntime { * append-only context); falls back to truncating `state.messages` for tests * that hand-roll a minimal facade. */ + #restoreSeenContextInFlight(): void { + if (!this.#seenContextInFlight) return; + this.#seenContext.clear(); + for (const [key, value] of this.#seenContextInFlight) this.#seenContext.set(key, value); + this.#seenContextInFlight = undefined; + } + #rollbackFailedTurn(snapshot: number): void { // Restore the primary-context dedup map to its pre-batch state: the // failed turn never reached the advisor, so first-time context collapsed // to "(unchanged…)" by this batch's #prepareBatch must expand again on // the retry/requeue pass. - if (this.#seenContextInFlight) { - this.#seenContext.clear(); - for (const [key, value] of this.#seenContextInFlight) this.#seenContext.set(key, value); - this.#seenContextInFlight = undefined; - } + this.#restoreSeenContextInFlight(); const messages = this.agent.state.messages; if (messages.length <= snapshot) return; try { @@ -1102,11 +1105,10 @@ export class AdvisorRuntime { const epoch = this.#epoch; for (const delta of popped) { if (delta.renderRevision === this.#renderRevision) continue; - // Batch text is finalized by #collectAndMaintainBatch -> #prepareBatch - // (single dedup+render pass). Refreshing here would run - // #dedupContextMessage a SECOND time and double-fold re-injected - // primary-context custom messages ("(unchanged…)" on first - // delivery). Mark the revision so the delta is not re-refreshed. + // Context maintenance estimates this preview before #prepareBatch makes + // its final deduped render. Rebuild stale text against the new context + // so the maintenance budget cannot undercount an expanded re-injection. + delta.text = this.#formatRawDelta(delta.rawMessages, delta.wip, false) ?? delta.text; delta.renderRevision = this.#renderRevision; } const recoveringOverflow = popped.some(delta => delta.overflowRecovery === true); @@ -1120,6 +1122,7 @@ export class AdvisorRuntime { // Epoch was invalidated during batch collection; restart the loop. if (result === null) continue; if (this.#sessionTransitionPaused) { + this.#restoreSeenContextInFlight(); this.#pending.unshift(...popped); continue; } diff --git a/packages/coding-agent/test/advisor/advisor.test.ts b/packages/coding-agent/test/advisor/advisor.test.ts index c03d57a38..d80b4e1f5 100644 --- a/packages/coding-agent/test/advisor/advisor.test.ts +++ b/packages/coding-agent/test/advisor/advisor.test.ts @@ -1438,6 +1438,29 @@ describe("advisor", () => { expect(promptText(promptInputs[0])).not.toContain(secret); }); + it("falls back to one redacted update when a regex secret spans source messages", async () => { + const obfuscator = new SecretObfuscator([{ type: "regex", content: "BEGIN[\\s\\S]*END" }]); + const promptInputs: Array = []; + const agent = makeAgent(promptInputs); + const messages: AgentMessage[] = [ + { role: "user", content: "BEGIN", timestamp: 1 } as AgentMessage, + { role: "user", content: "sensitive END", timestamp: 2 } as AgentMessage, + ]; + const runtime = new AdvisorRuntime(agent, { + snapshotMessages: () => messages, + enqueueAdvice: () => {}, + obfuscator, + }); + + runtime.onTurnEnd(messages); + await Promise.resolve(); + + expect(promptInputs).toHaveLength(1); + expect(typeof promptInputs[0]).toBe("string"); + expect(promptText(promptInputs[0]!)).not.toContain("BEGIN"); + expect(promptText(promptInputs[0]!)).not.toContain("sensitive END"); + }); + it("redacts expanded primary context before XML escaping", async () => { const secret = "ADVISOR&SECRET123"; const obfuscator = new SecretObfuscator([{ type: "plain", content: secret }]); @@ -2581,6 +2604,78 @@ describe("advisor", () => { expect(resetCount).toBe(1); }); + it("re-renders a queued primary context before its maintenance budget after overflow resets advisor context", async () => { + const overflowMessage = "context_length_exceeded: Your input exceeds the context window of this model."; + const firstOverflowPromptStarted = Promise.withResolvers(); + const releaseOverflowPrompt = Promise.withResolvers(); + const fourthMaintenance = Promise.withResolvers(); + const maintenanceTokens: number[] = []; + const state: { messages: AgentMessage[]; error?: string } = { messages: [] }; + let promptCalls = 0; + const agent: AdvisorAgent = { + prompt: async () => { + promptCalls++; + if (promptCalls === 2) { + firstOverflowPromptStarted.resolve(); + await releaseOverflowPrompt.promise; + state.error = overflowMessage; + } else { + state.error = undefined; + } + }, + abort: () => {}, + reset: () => { + state.messages.length = 0; + state.error = undefined; + }, + state, + }; + const planRule = "keep-expanded ".repeat(300); + const messages: AgentMessage[] = [ + { role: "user", content: "seed", timestamp: 1 } as AgentMessage, + { + role: "custom", + customType: "plan-mode-context", + content: planRule, + display: false, + timestamp: 2, + } as AgentMessage, + ]; + const runtime = new AdvisorRuntime( + agent, + { + snapshotMessages: () => messages, + enqueueAdvice: () => {}, + maintainContext: async incomingTokens => { + maintenanceTokens.push(incomingTokens); + if (maintenanceTokens.length === 4) fourthMaintenance.resolve(); + return false; + }, + }, + 0, + ); + runtime.onTurnEnd(messages); + await settleUntil(() => promptCalls === 1 && runtime.backlog === 0); + + messages.push({ role: "user", content: "overflow", timestamp: 3 } as AgentMessage); + runtime.onTurnEnd(messages); + await firstOverflowPromptStarted.promise; + messages.push({ role: "user", content: "after overflow", timestamp: 4 } as AgentMessage); + messages.push({ + role: "custom", + customType: "plan-mode-context", + content: planRule, + display: false, + timestamp: 5, + } as AgentMessage); + runtime.onTurnEnd(messages); + releaseOverflowPrompt.resolve(); + await fourthMaintenance.promise; + + expect(maintenanceTokens).toHaveLength(4); + expect(maintenanceTokens[3]).toBeGreaterThan(500); + }); + it("does not double-fold first-time primary context on overflow-recovery retry", async () => { // Regression: the recovery render previews the retry batch; if it advances // #seenContext, the retry's #prepareBatch re-dedup would collapse first-time @@ -2612,7 +2707,6 @@ describe("advisor", () => { }; const runtime = new AdvisorRuntime(agent, host, 0); runtime.seedTo(messages.length); - const rule = "Plan mode is active. You MUST perform READ-ONLY work only:\n- You NEVER create, edit, or delete files — except the single plan file named below."; messages.push({ role: "user", content: "overflowing-current-update", timestamp: 2 } as AgentMessage); @@ -4368,6 +4462,53 @@ describe("advisor", () => { expect(promptText(promptInputs[1])).toContain("keep me"); }); + it("re-expands first-time primary context after a session transition pauses before dispatch", async () => { + const promptInputs: Array = []; + const planRule = "Plan mode is active. Keep this first delivery expanded."; + const maintenancePaused = Promise.withResolvers(); + const prompted = Promise.withResolvers(); + let maintenanceCalls = 0; + let runtime: AdvisorRuntime; + const agent: AdvisorAgent = { + prompt: async input => { + promptInputs.push(input); + prompted.resolve(); + }, + abort: () => {}, + reset: () => {}, + state: { messages: [] }, + }; + const messages: AgentMessage[] = [ + { role: "user", content: "start planning", timestamp: 1 } as AgentMessage, + { + role: "custom", + customType: "plan-mode-context", + content: planRule, + display: false, + timestamp: 2, + } as AgentMessage, + ]; + runtime = new AdvisorRuntime(agent, { + snapshotMessages: () => messages, + enqueueAdvice: () => {}, + maintainContext: async () => { + if (++maintenanceCalls === 1) { + void runtime.pauseForSessionTransition(); + maintenancePaused.resolve(); + } + return false; + }, + }); + + runtime.onTurnEnd(messages); + await maintenancePaused.promise; + runtime.resumeAfterSessionTransition(); + await prompted.promise; + + expect(promptText(promptInputs[0]!)).toContain(planRule); + expect(promptText(promptInputs[0]!)).not.toContain("unchanged — still in effect"); + }); + it.each(["success", "error"] as const)( "releases blocked %s hooks so reset can run replacement work", async hookKind => { diff --git a/packages/coding-agent/test/advisor/delta-split-obfuscation.test.ts b/packages/coding-agent/test/advisor/delta-split-obfuscation.test.ts index f8881f095..9af009a2c 100644 --- a/packages/coding-agent/test/advisor/delta-split-obfuscation.test.ts +++ b/packages/coding-agent/test/advisor/delta-split-obfuscation.test.ts @@ -55,4 +55,24 @@ describe("renderAdvisorDeltaChunks obfuscation", () => { expect(text).not.toContain("SECRETVALUE123"); expect(text).toContain("[REDACTED]"); }); + + it("falls back when full-delta redaction spans source chunks", () => { + const crossChunkObfuscator: AdvisorObfuscator = { + obfuscate: text => text.replace(/first[\s\S]*second/g, "[REDACTED]"), + }; + const chunks = renderAdvisorDeltaChunks( + [ + { role: "user", content: "first", timestamp: 1 } as AgentMessage, + { role: "user", content: "second", timestamp: 2 } as AgentMessage, + ], + { + wip: false, + includeThinking: true, + obfuscator: crossChunkObfuscator, + advisorRegexSecretValues: new Set(), + }, + ); + + expect(chunks).toBeNull(); + }); }); diff --git a/packages/coding-agent/test/advisor/delta-split.test.ts b/packages/coding-agent/test/advisor/delta-split.test.ts index b808c1316..816644623 100644 --- a/packages/coding-agent/test/advisor/delta-split.test.ts +++ b/packages/coding-agent/test/advisor/delta-split.test.ts @@ -128,4 +128,20 @@ describe("renderAdvisorDeltaChunks (delta-split)", () => { }); expect(chunks!.length).toBeGreaterThan(1); }); + + it("puts the heading on the first emitted chunk when earlier messages render empty", () => { + const empty = { + role: "custom", + customType: "advisor", + content: "internal advice", + display: false, + timestamp: 1, + } as AgentMessage; + const chunks = renderAdvisorDeltaChunks([empty, user("visible", 2)], { + wip: false, + includeThinking: true, + advisorRegexSecretValues: new Set(), + }); + expect(chunksToText(chunks)).toStartWith("### Session update\n\n"); + }); });