Skip to content

Commit 6ed5a00

Browse files
LHMQ878LHMQ878
committed
fix: preserve per-round subagent task results on resume
 Seed subagent result cache from each task_result body before falling back to the shared subagent session, which previously overwrote every prior task result with the latest reply after session switches.
1 parent 85b6f5c commit 6ed5a00

4 files changed

Lines changed: 173 additions & 19 deletions

File tree

‎lib/compress/protected-content.ts‎

Lines changed: 26 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import {
77
} from "../protected-patterns"
88
import {
99
buildSubagentResultText,
10+
extractTaskResultBody,
1011
getSubAgentId,
1112
mergeSubagentResult,
1213
} from "../subagents/subagent-results"
@@ -167,25 +168,31 @@ export async function appendProtectedTools(
167168
)
168169
}
169170
} else {
170-
const subAgentSessionId = getSubAgentId(part)
171-
if (subAgentSessionId) {
172-
let subAgentResultText = ""
173-
try {
174-
const subAgentMessages = await fetchSessionMessages(
175-
client,
176-
subAgentSessionId,
177-
)
178-
subAgentResultText = buildSubagentResultText(subAgentMessages)
179-
} catch {
180-
subAgentResultText = ""
181-
}
182-
183-
if (subAgentResultText) {
184-
state.subAgentResultCache.set(part.callID, subAgentResultText)
185-
output = mergeSubagentResult(
186-
part.state.output,
187-
subAgentResultText,
188-
)
171+
const extractedResult = extractTaskResultBody(part.state.output)
172+
if (extractedResult) {
173+
state.subAgentResultCache.set(part.callID, extractedResult)
174+
output = mergeSubagentResult(part.state.output, extractedResult)
175+
} else {
176+
const subAgentSessionId = getSubAgentId(part)
177+
if (subAgentSessionId) {
178+
let subAgentResultText = ""
179+
try {
180+
const subAgentMessages = await fetchSessionMessages(
181+
client,
182+
subAgentSessionId,
183+
)
184+
subAgentResultText = buildSubagentResultText(subAgentMessages)
185+
} catch {
186+
subAgentResultText = ""
187+
}
188+
189+
if (subAgentResultText) {
190+
state.subAgentResultCache.set(part.callID, subAgentResultText)
191+
output = mergeSubagentResult(
192+
part.state.output,
193+
subAgentResultText,
194+
)
195+
}
189196
}
190197
}
191198
}

‎lib/messages/inject/subagent-results.ts‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ import type { SessionState, WithParts } from "../../state"
33
import { filterMessages } from "../shape"
44
import {
55
buildSubagentResultText,
6+
extractTaskResultBody,
67
getSubAgentId,
78
mergeSubagentResult,
89
} from "../../subagents/subagent-results"
@@ -51,6 +52,12 @@ export const injectExtendedSubAgentResults = async (
5152
continue
5253
}
5354

55+
const extractedResult = extractTaskResultBody(part.state.output)
56+
if (extractedResult) {
57+
state.subAgentResultCache.set(part.callID, extractedResult)
58+
continue
59+
}
60+
5461
const subAgentSessionId = getSubAgentId(part)
5562
if (!subAgentSessionId) {
5663
continue

‎lib/subagents/subagent-results.ts‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,20 @@ import type { WithParts } from "../state"
22

33
const SUB_AGENT_RESULT_BLOCK_REGEX = /(<task_result>\s*)([\s\S]*?)(\s*<\/task_result>)/i
44

5+
export function extractTaskResultBody(output: string): string | null {
6+
if (typeof output !== "string") {
7+
return null
8+
}
9+
10+
const match = output.match(SUB_AGENT_RESULT_BLOCK_REGEX)
11+
if (!match) {
12+
return null
13+
}
14+
15+
const body = match[2]?.trim()
16+
return body && body.length > 0 ? body : null
17+
}
18+
519
export function getSubAgentId(part: any): string | null {
620
const sessionId = part?.state?.metadata?.sessionId
721
if (typeof sessionId !== "string") {

‎tests/subagent-results.test.ts‎

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
1+
import assert from "node:assert/strict"
2+
import test from "node:test"
3+
import { injectExtendedSubAgentResults } from "../lib/messages/inject/subagent-results"
4+
import { Logger } from "../lib/logger"
5+
import { createSessionState, resetSessionState, type WithParts } from "../lib/state"
6+
import {
7+
extractTaskResultBody,
8+
} from "../lib/subagents/subagent-results"
9+
10+
function buildTaskPart(callID: string, sessionId: string, marker: string) {
11+
return {
12+
type: "tool",
13+
tool: "task",
14+
callID,
15+
state: {
16+
status: "completed",
17+
metadata: { sessionId },
18+
output: `<task_result>\n${marker}\n</task_result>`,
19+
},
20+
}
21+
}
22+
23+
function buildTaskMessage(callID: string, sessionId: string, marker: string): WithParts {
24+
return {
25+
info: {
26+
id: `msg-${callID}`,
27+
role: "assistant",
28+
sessionID: "parent-session",
29+
agent: "assistant",
30+
time: { created: 1 },
31+
} as WithParts["info"],
32+
parts: [buildTaskPart(callID, sessionId, marker)],
33+
}
34+
}
35+
36+
test("extractTaskResultBody returns task_result body content", () => {
37+
assert.equal(
38+
extractTaskResultBody("<task_result>\nROUND-1-MARKER\n</task_result>"),
39+
"ROUND-1-MARKER",
40+
)
41+
assert.equal(extractTaskResultBody("no task result here"), null)
42+
})
43+
44+
test("injectExtendedSubAgentResults keeps each resumed task round distinct", async () => {
45+
const subAgentSessionId = "subagent-session-1"
46+
const messages = [
47+
buildTaskMessage("call-1", subAgentSessionId, "ROUND-1-MARKER answer A"),
48+
buildTaskMessage("call-2", subAgentSessionId, "ROUND-2-MARKER answer B"),
49+
buildTaskMessage("call-3", subAgentSessionId, "ROUND-3-MARKER answer C"),
50+
]
51+
52+
const client = {
53+
session: {
54+
messages: async () => ({
55+
data: [
56+
{
57+
info: { role: "assistant" },
58+
parts: [{ type: "text", text: "ROUND-3-MARKER answer C" }],
59+
},
60+
],
61+
}),
62+
},
63+
}
64+
65+
const state = createSessionState()
66+
const logger = new Logger(false)
67+
68+
await injectExtendedSubAgentResults(client, state, logger, messages, true)
69+
70+
const roundOne = messages[0].parts[0].state.output as string
71+
const roundTwo = messages[1].parts[0].state.output as string
72+
const roundThree = messages[2].parts[0].state.output as string
73+
74+
assert.match(roundOne, /ROUND-1-MARKER answer A/)
75+
assert.match(roundTwo, /ROUND-2-MARKER answer B/)
76+
assert.match(roundThree, /ROUND-3-MARKER answer C/)
77+
assert.doesNotMatch(roundOne, /ROUND-3-MARKER/)
78+
assert.doesNotMatch(roundTwo, /ROUND-3-MARKER/)
79+
80+
resetSessionState(state)
81+
assert.equal(state.subAgentResultCache.size, 0)
82+
83+
await injectExtendedSubAgentResults(client, state, logger, messages, true)
84+
85+
assert.match(messages[0].parts[0].state.output as string, /ROUND-1-MARKER answer A/)
86+
assert.match(messages[1].parts[0].state.output as string, /ROUND-2-MARKER answer B/)
87+
assert.match(messages[2].parts[0].state.output as string, /ROUND-3-MARKER answer C/)
88+
})
89+
90+
test("injectExtendedSubAgentResults still expands empty task_result from subagent session", async () => {
91+
const subAgentSessionId = "subagent-session-2"
92+
const message = buildTaskMessage("call-4", subAgentSessionId, "")
93+
message.parts[0].state.output = "<task_result>\n</task_result>"
94+
95+
const client = {
96+
session: {
97+
messages: async () => ({
98+
data: [
99+
{
100+
info: {
101+
id: "sub-msg-1",
102+
role: "assistant",
103+
sessionID: subAgentSessionId,
104+
time: { created: 1 },
105+
},
106+
parts: [{ type: "text", text: "Expanded subagent reply" }],
107+
},
108+
],
109+
}),
110+
},
111+
}
112+
113+
const state = createSessionState()
114+
await injectExtendedSubAgentResults(
115+
client,
116+
state,
117+
new Logger(false),
118+
[message],
119+
true,
120+
)
121+
122+
assert.equal(
123+
extractTaskResultBody(message.parts[0].state.output as string),
124+
"Expanded subagent reply",
125+
)
126+
})

0 commit comments

Comments
 (0)