diff --git a/lib/compress/protected-content.ts b/lib/compress/protected-content.ts index 45b27355..206e051a 100644 --- a/lib/compress/protected-content.ts +++ b/lib/compress/protected-content.ts @@ -7,6 +7,7 @@ import { } from "../protected-patterns" import { buildSubagentResultText, + extractTaskResultBody, getSubAgentId, mergeSubagentResult, } from "../subagents/subagent-results" @@ -167,25 +168,31 @@ export async function appendProtectedTools( ) } } else { - const subAgentSessionId = getSubAgentId(part) - if (subAgentSessionId) { - let subAgentResultText = "" - try { - const subAgentMessages = await fetchSessionMessages( - client, - subAgentSessionId, - ) - subAgentResultText = buildSubagentResultText(subAgentMessages) - } catch { - subAgentResultText = "" - } - - if (subAgentResultText) { - state.subAgentResultCache.set(part.callID, subAgentResultText) - output = mergeSubagentResult( - part.state.output, - subAgentResultText, - ) + const extractedResult = extractTaskResultBody(part.state.output) + if (extractedResult) { + state.subAgentResultCache.set(part.callID, extractedResult) + output = mergeSubagentResult(part.state.output, extractedResult) + } else { + const subAgentSessionId = getSubAgentId(part) + if (subAgentSessionId) { + let subAgentResultText = "" + try { + const subAgentMessages = await fetchSessionMessages( + client, + subAgentSessionId, + ) + subAgentResultText = buildSubagentResultText(subAgentMessages) + } catch { + subAgentResultText = "" + } + + if (subAgentResultText) { + state.subAgentResultCache.set(part.callID, subAgentResultText) + output = mergeSubagentResult( + part.state.output, + subAgentResultText, + ) + } } } } diff --git a/lib/messages/inject/subagent-results.ts b/lib/messages/inject/subagent-results.ts index 8ca3d1d5..9e25d374 100644 --- a/lib/messages/inject/subagent-results.ts +++ b/lib/messages/inject/subagent-results.ts @@ -3,6 +3,7 @@ import type { SessionState, WithParts } from "../../state" import { filterMessages } from "../shape" import { buildSubagentResultText, + extractTaskResultBody, getSubAgentId, mergeSubagentResult, } from "../../subagents/subagent-results" @@ -51,6 +52,12 @@ export const injectExtendedSubAgentResults = async ( continue } + const extractedResult = extractTaskResultBody(part.state.output) + if (extractedResult) { + state.subAgentResultCache.set(part.callID, extractedResult) + continue + } + const subAgentSessionId = getSubAgentId(part) if (!subAgentSessionId) { continue diff --git a/lib/subagents/subagent-results.ts b/lib/subagents/subagent-results.ts index bea4af49..eb4a924d 100644 --- a/lib/subagents/subagent-results.ts +++ b/lib/subagents/subagent-results.ts @@ -2,6 +2,20 @@ import type { WithParts } from "../state" const SUB_AGENT_RESULT_BLOCK_REGEX = /(\s*)([\s\S]*?)(\s*<\/task_result>)/i +export function extractTaskResultBody(output: string): string | null { + if (typeof output !== "string") { + return null + } + + const match = output.match(SUB_AGENT_RESULT_BLOCK_REGEX) + if (!match) { + return null + } + + const body = match[2]?.trim() + return body && body.length > 0 ? body : null +} + export function getSubAgentId(part: any): string | null { const sessionId = part?.state?.metadata?.sessionId if (typeof sessionId !== "string") { diff --git a/tests/subagent-results.test.ts b/tests/subagent-results.test.ts new file mode 100644 index 00000000..fb515dbf --- /dev/null +++ b/tests/subagent-results.test.ts @@ -0,0 +1,126 @@ +import assert from "node:assert/strict" +import test from "node:test" +import { injectExtendedSubAgentResults } from "../lib/messages/inject/subagent-results" +import { Logger } from "../lib/logger" +import { createSessionState, resetSessionState, type WithParts } from "../lib/state" +import { + extractTaskResultBody, +} from "../lib/subagents/subagent-results" + +function buildTaskPart(callID: string, sessionId: string, marker: string) { + return { + type: "tool", + tool: "task", + callID, + state: { + status: "completed", + metadata: { sessionId }, + output: `\n${marker}\n`, + }, + } +} + +function buildTaskMessage(callID: string, sessionId: string, marker: string): WithParts { + return { + info: { + id: `msg-${callID}`, + role: "assistant", + sessionID: "parent-session", + agent: "assistant", + time: { created: 1 }, + } as WithParts["info"], + parts: [buildTaskPart(callID, sessionId, marker)], + } +} + +test("extractTaskResultBody returns task_result body content", () => { + assert.equal( + extractTaskResultBody("\nROUND-1-MARKER\n"), + "ROUND-1-MARKER", + ) + assert.equal(extractTaskResultBody("no task result here"), null) +}) + +test("injectExtendedSubAgentResults keeps each resumed task round distinct", async () => { + const subAgentSessionId = "subagent-session-1" + const messages = [ + buildTaskMessage("call-1", subAgentSessionId, "ROUND-1-MARKER answer A"), + buildTaskMessage("call-2", subAgentSessionId, "ROUND-2-MARKER answer B"), + buildTaskMessage("call-3", subAgentSessionId, "ROUND-3-MARKER answer C"), + ] + + const client = { + session: { + messages: async () => ({ + data: [ + { + info: { role: "assistant" }, + parts: [{ type: "text", text: "ROUND-3-MARKER answer C" }], + }, + ], + }), + }, + } + + const state = createSessionState() + const logger = new Logger(false) + + await injectExtendedSubAgentResults(client, state, logger, messages, true) + + const roundOne = messages[0].parts[0].state.output as string + const roundTwo = messages[1].parts[0].state.output as string + const roundThree = messages[2].parts[0].state.output as string + + assert.match(roundOne, /ROUND-1-MARKER answer A/) + assert.match(roundTwo, /ROUND-2-MARKER answer B/) + assert.match(roundThree, /ROUND-3-MARKER answer C/) + assert.doesNotMatch(roundOne, /ROUND-3-MARKER/) + assert.doesNotMatch(roundTwo, /ROUND-3-MARKER/) + + resetSessionState(state) + assert.equal(state.subAgentResultCache.size, 0) + + await injectExtendedSubAgentResults(client, state, logger, messages, true) + + assert.match(messages[0].parts[0].state.output as string, /ROUND-1-MARKER answer A/) + assert.match(messages[1].parts[0].state.output as string, /ROUND-2-MARKER answer B/) + assert.match(messages[2].parts[0].state.output as string, /ROUND-3-MARKER answer C/) +}) + +test("injectExtendedSubAgentResults still expands empty task_result from subagent session", async () => { + const subAgentSessionId = "subagent-session-2" + const message = buildTaskMessage("call-4", subAgentSessionId, "") + message.parts[0].state.output = "\n" + + const client = { + session: { + messages: async () => ({ + data: [ + { + info: { + id: "sub-msg-1", + role: "assistant", + sessionID: subAgentSessionId, + time: { created: 1 }, + }, + parts: [{ type: "text", text: "Expanded subagent reply" }], + }, + ], + }), + }, + } + + const state = createSessionState() + await injectExtendedSubAgentResults( + client, + state, + new Logger(false), + [message], + true, + ) + + assert.equal( + extractTaskResultBody(message.parts[0].state.output as string), + "Expanded subagent reply", + ) +})