diff --git a/lib/messages/inject/subagent-results.ts b/lib/messages/inject/subagent-results.ts index 967121f8..41d7f9d5 100644 --- a/lib/messages/inject/subagent-results.ts +++ b/lib/messages/inject/subagent-results.ts @@ -16,16 +16,20 @@ async function fetchSubAgentMessages(client: any, sessionId: string): Promise => { - if (!allowSubAgents) { - return - } +interface ExpandablePart { + part: any + callID: string + subAgentSessionId: string +} + +// Resuming a subagent with task_id keeps one subagent session alive across +// rounds, so several task results can point at the same session id. Only the +// newest of them describes the subagent's current last reply; for the earlier +// ones that reply belongs to a later round. Overwriting them rewrote history, +// and because the pinning cache lives in session state it was cold on every +// parent turn that followed a session switch (#595). +function collectExpandableParts(messages: WithParts[], state: SessionState): ExpandablePart[] { + const expandable: ExpandablePart[] = [] for (const message of messages) { const parts = Array.isArray(message.parts) ? message.parts : [] @@ -41,44 +45,69 @@ export const injectExtendedSubAgentResults = async ( continue } - const cachedResult = state.subAgentResultCache.get(part.callID) - if (cachedResult !== undefined) { - if (cachedResult) { - part.state.output = stripHallucinationsFromString( - mergeSubagentResult(part.state.output, cachedResult), - state.idFormat, - ) - } - continue - } - const subAgentSessionId = getSubAgentId(part) if (!subAgentSessionId) { continue } - let subAgentMessages: WithParts[] = [] - try { - subAgentMessages = await fetchSubAgentMessages(client, subAgentSessionId) - } catch (error) { - logger.warn("Failed to fetch subagent session for output expansion", { - subAgentSessionId, - callID: part.callID, - error: error instanceof Error ? error.message : String(error), - }) - continue - } + expandable.push({ part, callID: part.callID, subAgentSessionId }) + } + } - const subAgentResultText = buildSubagentResultText(subAgentMessages) - if (!subAgentResultText) { - continue + const newestIndexBySession = new Map() + expandable.forEach((entry, index) => { + newestIndexBySession.set(entry.subAgentSessionId, index) + }) + + return expandable.filter( + (entry, index) => newestIndexBySession.get(entry.subAgentSessionId) === index, + ) +} + +export const injectExtendedSubAgentResults = async ( + client: any, + state: SessionState, + logger: Logger, + messages: WithParts[], + allowSubAgents: boolean, +): Promise => { + if (!allowSubAgents) { + return + } + + for (const { part, callID, subAgentSessionId } of collectExpandableParts(messages, state)) { + const cachedResult = state.subAgentResultCache.get(callID) + if (cachedResult !== undefined) { + if (cachedResult) { + part.state.output = stripHallucinationsFromString( + mergeSubagentResult(part.state.output, cachedResult), + state.idFormat, + ) } + continue + } + + let subAgentMessages: WithParts[] = [] + try { + subAgentMessages = await fetchSubAgentMessages(client, subAgentSessionId) + } catch (error) { + logger.warn("Failed to fetch subagent session for output expansion", { + subAgentSessionId, + callID, + error: error instanceof Error ? error.message : String(error), + }) + continue + } - state.subAgentResultCache.set(part.callID, subAgentResultText) - part.state.output = stripHallucinationsFromString( - mergeSubagentResult(part.state.output, subAgentResultText), - state.idFormat, - ) + const subAgentResultText = buildSubagentResultText(subAgentMessages) + if (!subAgentResultText) { + continue } + + state.subAgentResultCache.set(callID, subAgentResultText) + part.state.output = stripHallucinationsFromString( + mergeSubagentResult(part.state.output, subAgentResultText), + state.idFormat, + ) } } diff --git a/tests/subagent-rounds.test.ts b/tests/subagent-rounds.test.ts new file mode 100644 index 00000000..3ed4fb63 --- /dev/null +++ b/tests/subagent-rounds.test.ts @@ -0,0 +1,212 @@ +import assert from "node:assert/strict" +import test from "node:test" +import { mkdtempSync, rmSync } from "node:fs" +import { join } from "node:path" +import { tmpdir } from "node:os" +import type { WithParts } from "../lib/state" + +const root = mkdtempSync(join(tmpdir(), "dcp-subagent-rounds-")) +process.env.XDG_DATA_HOME = join(root, "data") +process.env.XDG_CONFIG_HOME = join(root, "config") +process.env.OPENCODE_CONFIG_DIR = join(root, "config", "opencode") +test.after(() => rmSync(root, { recursive: true, force: true })) + +const { Logger } = await import("../lib/logger") +const { createSessionState, resetSessionState } = await import("../lib/state") +const { injectExtendedSubAgentResults } = await import("../lib/messages/inject/subagent-results") + +const logger = new Logger(false) + +function taskMessage( + sessionID: string, + id: string, + created: number, + subAgentSessionId: string, + recorded: string, +): WithParts { + return { + info: { + id, + role: "assistant", + sessionID, + agent: "build", + model: { providerID: "lab", modelID: "test" }, + time: { created }, + }, + parts: [ + { + id: `prt_${id}`, + messageID: id, + sessionID, + type: "tool", + tool: "task", + callID: `call_${id}`, + state: { + status: "completed", + input: { description: `round ${created}` }, + output: `${recorded}`, + metadata: { sessionId: subAgentSessionId }, + }, + }, + ], + } as unknown as WithParts +} + +function subAgentSession(subAgentSessionId: string, replies: string[]): WithParts[] { + return replies.map((reply, index) => ({ + info: { + id: `sub_msg_${index + 1}`, + role: "assistant", + sessionID: subAgentSessionId, + agent: "subagent", + model: { providerID: "lab", modelID: "test" }, + time: { created: index + 1 }, + }, + parts: [ + { + id: `sub_prt_${index + 1}`, + messageID: `sub_msg_${index + 1}`, + sessionID: subAgentSessionId, + type: "text", + text: reply, + }, + ], + })) as WithParts[] +} + +function clientFor(subAgentSessionId: string, replies: string[]) { + return { + session: { + messages: async ({ path }: { path: { id: string } }) => ({ + data: + path.id === subAgentSessionId + ? subAgentSessionId && subAgentSession(subAgentSessionId, replies) + : [], + }), + get: async () => ({ data: { parentID: null } }), + }, + } +} + +function recordedBodies(messages: WithParts[]): string[] { + return messages.map((message) => { + const part = (message.parts as any[])[0] + const match = part.state.output.match(/([\s\S]*?)<\/task_result>/) + return match ? match[1] : "" + }) +} + +test("resuming a subagent does not rewrite the results of earlier rounds", async () => { + const parent = "ses_parent" + const subAgent = "ses_subagent" + + // One subagent session, three rounds, all reachable through the same id. + const messages = [ + taskMessage(parent, "msg_1", 1, subAgent, "ROUND-1 answer A"), + taskMessage(parent, "msg_2", 2, subAgent, "ROUND-2 answer B"), + taskMessage(parent, "msg_3", 3, subAgent, "ROUND-3 answer C"), + ] + const state = createSessionState() + state.sessionId = parent + const client = clientFor(subAgent, ["ROUND-1 answer A", "ROUND-2 answer B", "ROUND-3 answer C"]) + + await injectExtendedSubAgentResults(client as any, state, logger, messages, true) + + assert.deepEqual(recordedBodies(messages), [ + "ROUND-1 answer A", + "ROUND-2 answer B", + "ROUND-3 answer C", + ]) +}) + +test("the result stays pinned after the session cache is dropped", async () => { + const parent = "ses_parent_cold" + const subAgent = "ses_subagent_cold" + const replies = ["A", "B", "C"] + + const build = () => [ + taskMessage(parent, "msg_1", 1, subAgent, "A"), + taskMessage(parent, "msg_2", 2, subAgent, "B"), + taskMessage(parent, "msg_3", 3, subAgent, "C"), + ] + + const client = clientFor(subAgent, replies) + + // First pass fills the cache. + const warm = build() + const warmState = createSessionState() + warmState.sessionId = parent + await injectExtendedSubAgentResults(client as any, warmState, logger, warm, true) + const warmBodies = recordedBodies(warm) + + // A session switch wipes subAgentResultCache, so the second pass runs cold. + // This is the case the report describes: a resumed subagent makes its own + // model calls through the same hook between the parent's turns. + const cold = build() + const coldState = createSessionState() + coldState.sessionId = parent + assert.equal(coldState.subAgentResultCache.size, 0) + await injectExtendedSubAgentResults(client as any, coldState, logger, cold, true) + + assert.deepEqual(recordedBodies(cold), warmBodies) + assert.deepEqual(warmBodies, ["A", "B", "C"]) +}) + +test("a fresh subagent session still gets its result expanded", async () => { + const parent = "ses_parent_fresh" + const subAgent = "ses_subagent_fresh" + const messages = [taskMessage(parent, "msg_1", 1, subAgent, "truncated by the host")] + const state = createSessionState() + state.sessionId = parent + const client = clientFor(subAgent, ["the full subagent reply"]) + + await injectExtendedSubAgentResults(client as any, state, logger, messages, true) + + assert.deepEqual(recordedBodies(messages), ["the full subagent reply"]) +}) + +test("separate subagent sessions do not interfere with each other", async () => { + const parent = "ses_parent_multi" + const messages = [ + taskMessage(parent, "msg_1", 1, "ses_sub_a", "A recorded"), + taskMessage(parent, "msg_2", 2, "ses_sub_b", "B recorded"), + taskMessage(parent, "msg_3", 3, "ses_sub_a", "A recorded again"), + ] + const state = createSessionState() + state.sessionId = parent + const client = { + session: { + messages: async ({ path }: { path: { id: string } }) => ({ + data: + path.id === "ses_sub_a" + ? subAgentSession("ses_sub_a", ["A1", "A2"]) + : subAgentSession("ses_sub_b", ["B1"]), + }), + get: async () => ({ data: { parentID: null } }), + }, + } + + await injectExtendedSubAgentResults(client as any, state, logger, messages, true) + + const bodies = recordedBodies(messages) + assert.equal(bodies[0], "A recorded", "the first call into sub A was rewritten by round 2") + assert.equal(bodies[1], "B1") + assert.equal(bodies[2], "A2") +}) + +test("allowSubAgents off leaves every result untouched", async () => { + const parent = "ses_parent_off" + const messages = [taskMessage(parent, "msg_1", 1, "ses_sub_x", "original")] + const state = createSessionState() + state.sessionId = parent + + await injectExtendedSubAgentResults( + clientFor("ses_sub_x", ["hijacked"]) as any, + state, + logger, + messages, + false, + ) + + assert.deepEqual(recordedBodies(messages), ["original"]) +})