Skip to content

Commit 81924cf

Browse files
joelnishanthjoelnishanth
authored
fix(ollama): preserve length-limited responses (#89160)
Preserve Ollama length termination and surface incomplete token-limited embedded-agent turns without discarding durable tool output. Fixes #89051. Co-authored-by: joelnishanth <[email protected]>
1 parent 6134c00 commit 81924cf

14 files changed

Lines changed: 498 additions & 18 deletions

‎extensions/ollama/src/stream.test.ts‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ function makeOllamaResponse(params: {
1515
content?: string;
1616
thinking?: string;
1717
reasoning?: string;
18+
done_reason?: string;
1819
tool_calls?: Array<{ function: { name: string; arguments: Record<string, unknown> } }>;
1920
}) {
2021
return {
@@ -28,6 +29,7 @@ function makeOllamaResponse(params: {
2829
...(params.tool_calls ? { tool_calls: params.tool_calls } : {}),
2930
},
3031
done: true,
32+
...(params.done_reason ? { done_reason: params.done_reason } : {}),
3133
prompt_eval_count: 100,
3234
eval_count: 50,
3335
};
@@ -86,6 +88,24 @@ describe("buildAssistantMessage", () => {
8688
expect(msg.content).toHaveLength(1);
8789
expect(msg.content[0]).toEqual({ type: "text", text: "Just text" });
8890
});
91+
92+
it("preserves output-budget length stops", () => {
93+
const response = makeOllamaResponse({
94+
content: "Partial answer",
95+
done_reason: "length",
96+
});
97+
const msg = buildAssistantMessage(response, MODEL_INFO);
98+
expect(msg.stopReason).toBe("length");
99+
});
100+
101+
it("keeps tool use authoritative over a length stop", () => {
102+
const response = makeOllamaResponse({
103+
done_reason: "length",
104+
tool_calls: [{ function: { name: "read", arguments: { path: "README.md" } } }],
105+
});
106+
const msg = buildAssistantMessage(response, MODEL_INFO);
107+
expect(msg.stopReason).toBe("toolUse");
108+
});
89109
});
90110

91111
describe("createOllamaStreamFn thinking events", () => {
@@ -235,6 +255,33 @@ describe("createOllamaStreamFn thinking events", () => {
235255
expect(textStart?.contentIndex).toBe(0);
236256
});
237257

258+
it("emits length for a token-limited native stream", async () => {
259+
const events = await streamOllamaEvents([
260+
{
261+
model: "qwen3.5",
262+
created_at: "2026-01-01T00:00:00Z",
263+
message: { role: "assistant", content: "Partial answer" },
264+
done: false,
265+
},
266+
{
267+
model: "qwen3.5",
268+
created_at: "2026-01-01T00:00:01Z",
269+
message: { role: "assistant", content: "" },
270+
done: true,
271+
done_reason: "length",
272+
prompt_eval_count: 10,
273+
eval_count: 5,
274+
},
275+
]);
276+
277+
const done = events.find((event) => event.type === "done") as {
278+
reason?: string;
279+
message?: { stopReason?: string };
280+
};
281+
expect(done.reason).toBe("length");
282+
expect(done.message?.stopReason).toBe("length");
283+
});
284+
238285
it("uses generic stream timeout for Ollama request timeout", async () => {
239286
await streamOllamaEvents([makeOllamaResponse({ content: "ok" })], { timeoutMs: 2500 });
240287

‎extensions/ollama/src/stream.ts‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -653,6 +653,13 @@ function estimateTokensFromChars(chars: number): number {
653653
return Math.max(1, Math.round(chars / CHARS_PER_TOKEN_ESTIMATE));
654654
}
655655

656+
function resolveOllamaStopReason(response: OllamaChatResponse) {
657+
if (response.message.tool_calls?.length) {
658+
return "toolUse" as const;
659+
}
660+
return response.done_reason === "length" ? ("length" as const) : ("stop" as const);
661+
}
662+
656663
function estimateOllamaPromptTokens(params: {
657664
messages: OllamaChatMessage[];
658665
tools: OllamaTool[];
@@ -1061,7 +1068,7 @@ export function buildAssistantMessage(
10611068
return buildStreamAssistantMessage({
10621069
model: modelInfo,
10631070
content,
1064-
stopReason: toolCalls && toolCalls.length > 0 ? "toolUse" : "stop",
1071+
stopReason: resolveOllamaStopReason(response),
10651072
usage: buildUsageWithNoCost({
10661073
input: resolveUsageCount(response.prompt_eval_count, usageFallback?.input),
10671074
output: resolveUsageCount(response.eval_count, usageFallback?.output),
@@ -1442,7 +1449,7 @@ function createRawOllamaStreamFn(
14421449

14431450
stream.push({
14441451
type: "done",
1445-
reason: assistantMessage.stopReason === "toolUse" ? "toolUse" : "stop",
1452+
reason: resolveOllamaStopReason(finalResponse),
14461453
message: assistantMessage,
14471454
});
14481455
} finally {

‎src/agents/embedded-agent-runner/run.incomplete-turn.test.ts‎

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1400,6 +1400,26 @@ describe("runEmbeddedAgent incomplete-turn safety", () => {
14001400
lastAssistant: { stopReason: "end_turn" },
14011401
}),
14021402
).toBe(false);
1403+
expect(
1404+
isIncompleteTerminalAssistantTurn({
1405+
hasAssistantVisibleText: true,
1406+
lastAssistant: { stopReason: "length" },
1407+
}),
1408+
).toBe(true);
1409+
expect(
1410+
isIncompleteTerminalAssistantTurn({
1411+
hasAssistantVisibleText: true,
1412+
hasTerminalOutput: true,
1413+
lastAssistant: { stopReason: "length" },
1414+
}),
1415+
).toBe(false);
1416+
expect(
1417+
isIncompleteTerminalAssistantTurn({
1418+
hasAssistantVisibleText: true,
1419+
hasTerminalOutput: true,
1420+
lastAssistant: { stopReason: "toolUse" },
1421+
}),
1422+
).toBe(true);
14031423
});
14041424

14051425
it("surfaces no-visible-answer recovery for app-server interrupted tool-only output", () => {
@@ -2567,6 +2587,109 @@ describe("runEmbeddedAgent incomplete-turn safety", () => {
25672587
expect(incompleteTurnText).toContain("couldn't generate a response");
25682588
});
25692589

2590+
it("surfaces incomplete-turn text for token-limited partial answers", () => {
2591+
const incompleteTurnText = resolveIncompleteTurnPayloadText({
2592+
payloadCount: 1,
2593+
aborted: false,
2594+
timedOut: false,
2595+
attempt: makeAttemptResult({
2596+
assistantTexts: ["Partial answer"],
2597+
lastAssistant: {
2598+
role: "assistant",
2599+
stopReason: "length",
2600+
provider: "ollama",
2601+
model: "qwen3.5",
2602+
content: [{ type: "text", text: "Partial answer" }],
2603+
} as unknown as EmbeddedRunAttemptResult["lastAssistant"],
2604+
}),
2605+
});
2606+
2607+
expect(incompleteTurnText).toContain("couldn't generate a response");
2608+
});
2609+
2610+
it("keeps complete visible stop turns successful", () => {
2611+
const incompleteTurnText = resolveIncompleteTurnPayloadText({
2612+
payloadCount: 1,
2613+
aborted: false,
2614+
timedOut: false,
2615+
attempt: makeAttemptResult({
2616+
assistantTexts: ["Complete answer"],
2617+
lastAssistant: {
2618+
role: "assistant",
2619+
stopReason: "stop",
2620+
provider: "ollama",
2621+
model: "qwen3.5",
2622+
content: [{ type: "text", text: "Complete answer" }],
2623+
} as unknown as EmbeddedRunAttemptResult["lastAssistant"],
2624+
}),
2625+
});
2626+
2627+
expect(incompleteTurnText).toBeNull();
2628+
});
2629+
2630+
it("preserves terminal tool media on token-limited turns", () => {
2631+
const incompleteTurnText = resolveIncompleteTurnPayloadText({
2632+
payloadCount: 1,
2633+
aborted: false,
2634+
timedOut: false,
2635+
attempt: makeAttemptResult({
2636+
assistantTexts: ["Partial answer"],
2637+
toolMediaUrls: ["file:///tmp/render.png"],
2638+
lastAssistant: {
2639+
role: "assistant",
2640+
stopReason: "length",
2641+
provider: "ollama",
2642+
model: "qwen3.5",
2643+
content: [{ type: "text", text: "Partial answer" }],
2644+
} as unknown as EmbeddedRunAttemptResult["lastAssistant"],
2645+
}),
2646+
});
2647+
2648+
expect(incompleteTurnText).toBeNull();
2649+
});
2650+
2651+
it("preserves tool media already delivered through block replies", () => {
2652+
const incompleteTurnText = resolveIncompleteTurnPayloadText({
2653+
payloadCount: 1,
2654+
aborted: false,
2655+
timedOut: false,
2656+
attempt: makeAttemptResult({
2657+
assistantTexts: ["Partial answer"],
2658+
hasToolMediaBlockReply: true,
2659+
lastAssistant: {
2660+
role: "assistant",
2661+
stopReason: "length",
2662+
provider: "ollama",
2663+
model: "qwen3.5",
2664+
content: [{ type: "text", text: "Partial answer" }],
2665+
} as unknown as EmbeddedRunAttemptResult["lastAssistant"],
2666+
}),
2667+
});
2668+
2669+
expect(incompleteTurnText).toBeNull();
2670+
});
2671+
2672+
it("preserves successful cron progress on token-limited turns", () => {
2673+
const incompleteTurnText = resolveIncompleteTurnPayloadText({
2674+
payloadCount: 1,
2675+
aborted: false,
2676+
timedOut: false,
2677+
attempt: makeAttemptResult({
2678+
assistantTexts: ["Partial answer"],
2679+
successfulCronAdds: 1,
2680+
lastAssistant: {
2681+
role: "assistant",
2682+
stopReason: "length",
2683+
provider: "ollama",
2684+
model: "qwen3.5",
2685+
content: [{ type: "text", text: "Partial answer" }],
2686+
} as unknown as EmbeddedRunAttemptResult["lastAssistant"],
2687+
}),
2688+
});
2689+
2690+
expect(incompleteTurnText).toBeNull();
2691+
});
2692+
25702693
it.each([
25712694
[
25722695
"heartbeat responses",

‎src/agents/embedded-agent-runner/run/attempt-trajectory-status.test.ts‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,47 @@ describe("attempt trajectory status", () => {
4343
).toEqual({ status: "success" });
4444
});
4545

46+
it("marks length-limited visible text as non-deliverable without terminal output", () => {
47+
expect(
48+
resolveAttemptTrajectoryTerminal(
49+
baseParams({
50+
assistantTexts: ["Partial answer."],
51+
lastAssistantStopReason: "length",
52+
}),
53+
),
54+
).toEqual({
55+
status: "error",
56+
terminalError: NON_DELIVERABLE_TERMINAL_TURN_REASON,
57+
});
58+
});
59+
60+
it("does not treat streamed partial payloads as completed length-limited output", () => {
61+
expect(
62+
resolveAttemptTrajectoryTerminal(
63+
baseParams({
64+
assistantTexts: ["Partial answer."],
65+
synthesizedPayloadCount: 1,
66+
lastAssistantStopReason: "length",
67+
}),
68+
),
69+
).toEqual({
70+
status: "error",
71+
terminalError: NON_DELIVERABLE_TERMINAL_TURN_REASON,
72+
});
73+
});
74+
75+
it("keeps length-limited turns successful when terminal output was delivered", () => {
76+
expect(
77+
resolveAttemptTrajectoryTerminal(
78+
baseParams({
79+
assistantTexts: [],
80+
lastAssistantStopReason: "length",
81+
hasTerminalOutput: true,
82+
}),
83+
),
84+
).toEqual({ status: "success" });
85+
});
86+
4687
it("keeps committed messaging tool delivery as success even without assistant text", () => {
4788
expect(
4889
resolveAttemptTrajectoryTerminal(
@@ -52,6 +93,15 @@ describe("attempt trajectory status", () => {
5293
}),
5394
),
5495
).toEqual({ status: "success" });
96+
expect(
97+
resolveAttemptTrajectoryTerminal(
98+
baseParams({
99+
didSendViaMessagingTool: true,
100+
messagingToolSentTargets: [{ channel: "telegram" }],
101+
lastAssistantStopReason: "length",
102+
}),
103+
),
104+
).toEqual({ status: "success" });
55105
});
56106

57107
it("keeps accepted session spawns as terminal progress", () => {

‎src/agents/embedded-agent-runner/run/attempt-trajectory-status.ts‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ export type ResolveAttemptTrajectoryTerminalParams = {
4646
silentExpected?: boolean;
4747
emptyAssistantReplyIsSilent?: boolean;
4848
lastAssistantStopReason?: string;
49+
hasTerminalOutput?: boolean;
4950
};
5051

5152
/**
@@ -118,7 +119,6 @@ export function resolveAttemptTrajectoryTerminal(
118119
params.didSendDeterministicApprovalPrompt ||
119120
hasCommittedMessagingDeliveryEvidence(params) ||
120121
hasAcceptedSessionSpawn(params.acceptedSessionSpawns) ||
121-
params.synthesizedPayloadCount > 0 ||
122122
params.heartbeatToolResponse !== undefined ||
123123
(params.clientToolCalls?.length ?? 0) > 0 ||
124124
params.yieldDetected === true ||
@@ -131,9 +131,21 @@ export function resolveAttemptTrajectoryTerminal(
131131
terminalError: NON_DELIVERABLE_TERMINAL_TURN_REASON,
132132
};
133133
}
134+
if (
135+
params.lastAssistantStopReason === "length" &&
136+
!params.hasTerminalOutput &&
137+
!hasExplicitTerminalDelivery
138+
) {
139+
return {
140+
status: "error",
141+
terminalError: NON_DELIVERABLE_TERMINAL_TURN_REASON,
142+
};
143+
}
134144

135145
const hasDeliverableOrProgress =
136146
hasExplicitTerminalDelivery ||
147+
params.hasTerminalOutput ||
148+
params.synthesizedPayloadCount > 0 ||
137149
hasNonEmptyAssistantText(params.assistantTexts) ||
138150
params.successfulCronAdds > 0;
139151

‎src/agents/embedded-agent-runner/run/attempt.spawn-workspace.test-support.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,7 @@ export function createSubscriptionMock(): SubscriptionMock {
118118
getMessagingToolSourceReplyPayloads: () => [] as MessagingToolSourceReplyPayload[],
119119
getHeartbeatToolResponse: () => undefined,
120120
getPendingToolMediaReply: () => null,
121+
hasToolMediaBlockReply: () => false,
121122
getVisibleBlockReplyCount: () => 0,
122123
getSuccessfulCronAdds: () => 0,
123124
getReplayState: () => ({

0 commit comments

Comments
 (0)