fix: preserve per-round subagent task results on resume

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,
+                                        )
+                                    }
                                 }
                             }
                         }

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

lib/subagents/subagent-results.ts

@@ -2,6 +2,20 @@ import type { WithParts } from "../state"
 
 const SUB_AGENT_RESULT_BLOCK_REGEX = /(<task_result>\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") {

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: `<task_result>\n${marker}\n</task_result>`,
+        },
+    }
+}
+
+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("<task_result>\nROUND-1-MARKER\n</task_result>"),
+        "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 = "<task_result>\n</task_result>"
+
+    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",
+    )
+})