vrr/qwen-code
2
0
Fork 0
mirror of https://github.com/QwenLM/qwen-code.git synced 2026-09-02 12:15:06 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 14

fix(core): persist partial assistant turn when stream errors mid tool_use

Browse source 
Weak-network failures during an Anthropic-compatible stream (DeepSeek,
api.anthropic.com, etc.) can drop the SSE between a tool_use
`content_block_stop` and the terminal `message_stop`. The functionCall
chunk is already yielded at content_block_stop, so:

  - Turn.run records a ToolCallRequest event.
  - useGeminiStream's for-await exits and schedules the tool.
  - handleCompletedTools eventually fires submitQuery(..., ToolResult)
    and pushes a user[functionResponse] into history.
  - But processStreamResponse's history.push for the model turn never
    ran (the for-await threw first), so the matching tool_use is gone.

The next request body has `user → user[tool_result]` with no tool_use
in between, and the server rejects with HTTP 400:
"tool_use_id ... must have a corresponding tool_use block in the
previous message". Ctrl+Y can't recover because
stripOrphanedUserEntriesFromHistory only strips trailing user
entries — the lost tool_use is unrecoverable, and the session is
wedged.

Wrap the for-await loop in processStreamResponse with try/catch.
When the stream throws AND any functionCall chunk was already
yielded (hasToolCall=true), persist the partial assistant turn to
history before re-throwing. The eventual tool_result submission
then has a matching tool_use and the session can continue.

Plain-text partial turns (no functionCall yielded) are intentionally
NOT persisted: the Retry path pops the trailing user prompt and
re-issues it, so a stale partial-text model turn between them would
either bias the retry or surface as duplicate output.
This commit is contained in:
wenshao 2026-05-15 21:13:55 +08:00
parent 435f711e33
commit fe35e37783
2 changed files with 201 additions and 37 deletions
Download patch file Download diff file Expand all files Collapse all files

110
packages/core/src/core/geminiChat.test.ts
View file

@ -699,6 +699,116 @@ describe('GeminiChat', async () => {
).resolves.not.toThrow();
});
it('persists partial assistant turn when stream throws after a tool_use chunk', async () => {
// Weak-network scenario: Anthropic-compatible providers emit the
// `functionCall` part on `content_block_stop`; the SSE may then drop
// before `message_stop`. The yielded chunk is enough for `Turn.run`
// to queue a `ToolCallRequest`, the tool scheduler will eventually
// submit a `functionResponse` user turn — without a matching
// tool_use in history, the next request body shows
// `user → user[tool_result]` and DeepSeek/Anthropic rejects with
// "tool_use_id ... must have a corresponding tool_use block in the
// previous message". `processStreamResponse` must persist the
// partial model turn before re-throwing so the pairing is intact.
mockRetryWithBackoff.mockImplementation(async (apiCall) => apiCall());
const networkError = new Error('SSE connection reset by peer');
const streamThatThrowsAfterToolCall = (async function* () {
yield {
candidates: [
{
content: {
role: 'model',
parts: [
{
functionCall: {
id: 'call_00_CeJrKJB0PSmXUZTCWHET7332',
name: 'read_file',
args: { path: '/tmp/x.txt' },
},
},
],
},
},
],
} as unknown as GenerateContentResponse;
throw networkError;
})();
vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(
streamThatThrowsAfterToolCall,
);
const stream = await chat.sendMessageStream(
'test-model',
{ message: 'open /tmp/x.txt please' },
'prompt-weak-network-tool',
);
await expect(
(async () => {
for await (const _ of stream) {
/* drain */
}
})(),
).rejects.toBe(networkError);
const history = chat.getHistory();
expect(history.length).toBe(2);
expect(history[0]!.role).toBe('user');
const modelTurn = history[1]!;
expect(modelTurn.role).toBe('model');
expect(modelTurn.parts).toBeDefined();
const functionCallPart = modelTurn.parts!.find((p) => p.functionCall);
expect(functionCallPart?.functionCall?.id).toBe(
'call_00_CeJrKJB0PSmXUZTCWHET7332',
);
expect(functionCallPart?.functionCall?.name).toBe('read_file');
});
it('does NOT persist partial assistant turn when stream throws before any tool_use chunk', async () => {
// Plain-text partial responses are deliberately dropped on stream
// error: the Retry path pops the trailing user prompt and re-issues
// it, so a stale partial-text model turn between them would bias
// the retry or surface as duplicate output. Only tool_use turns
// need the partial-history bridge to preserve the tool_use →
// tool_result invariant — text alone has no such invariant.
mockRetryWithBackoff.mockImplementation(async (apiCall) => apiCall());
const networkError = new Error('connection reset');
const streamThatThrowsAfterText = (async function* () {
yield {
candidates: [
{
content: {
role: 'model',
parts: [{ text: 'partial reply that will be lost' }],
},
},
],
} as unknown as GenerateContentResponse;
throw networkError;
})();
vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(
streamThatThrowsAfterText,
);
const stream = await chat.sendMessageStream(
'test-model',
{ message: 'hello' },
'prompt-weak-network-text',
);
await expect(
(async () => {
for await (const _ of stream) {
/* drain */
}
})(),
).rejects.toBe(networkError);
const history = chat.getHistory();
// Only the user turn is in history — the partial-text model turn is
// intentionally not persisted.
expect(history.length).toBe(1);
expect(history[0]!.role).toBe('user');
});
it('should throw InvalidStreamError when no tool call and no finish reason', async () => {
vi.useFakeTimers();
try {

128
packages/core/src/core/geminiChat.ts
View file

@ -1257,48 +1257,62 @@ export class GeminiChat {
let hasToolCall = false;
let hasFinishReason = false;
// Captured if the upstream stream throws mid-iteration (typical on weak
// networks: SSE drops between `content_block_stop` of a tool_use and the
// terminal `message_stop`). We still build / record / push a partial
// assistant turn below before re-throwing — see the dedicated branch in
// the post-loop block for why this is needed to keep tool_use/tool_result
// pairing intact across the failure.
let streamError: unknown = null;
for await (const chunk of streamResponse) {
// Use ||= to avoid later usage-only chunks (no candidates) overwriting
// a finishReason that was already seen in an earlier chunk.
hasFinishReason ||=
chunk?.candidates?.some((candidate) => candidate.finishReason) ?? false;
try {
for await (const chunk of streamResponse) {
// Use ||= to avoid later usage-only chunks (no candidates) overwriting
// a finishReason that was already seen in an earlier chunk.
hasFinishReason ||=
chunk?.candidates?.some((candidate) => candidate.finishReason) ??
false;
if (isValidResponse(chunk)) {
const content = chunk.candidates?.[0]?.content;
if (content?.parts) {
if (content.parts.some((part) => part.functionCall)) {
hasToolCall = true;
if (isValidResponse(chunk)) {
const content = chunk.candidates?.[0]?.content;
if (content?.parts) {
if (content.parts.some((part) => part.functionCall)) {
hasToolCall = true;
}
// Collect all parts for recording
allModelParts.push(...content.parts);
}
// Collect all parts for recording
allModelParts.push(...content.parts);
}
// Collect token usage for consolidated recording
if (chunk.usageMetadata) {
usageMetadata = chunk.usageMetadata;
// Context usage tracks prompt size; output isn't in history yet.
const lastPromptTokenCount =
usageMetadata.promptTokenCount || usageMetadata.totalTokenCount;
if (lastPromptTokenCount) {
// Always update the per-chat counter so this chat (including
// subagents) can make its own compaction decisions.
this.lastPromptTokenCount = lastPromptTokenCount;
// Mirror to the global telemetry only when wired — subagents
// pass `telemetryService=undefined` to keep their context usage
// out of the main session's UI counters.
this.telemetryService?.setLastPromptTokenCount(
lastPromptTokenCount,
);
}
if (usageMetadata.cachedContentTokenCount && this.telemetryService) {
this.telemetryService.setLastCachedContentTokenCount(
usageMetadata.cachedContentTokenCount,
);
}
}
yield chunk; // Yield every chunk to the UI immediately.
}
// Collect token usage for consolidated recording
if (chunk.usageMetadata) {
usageMetadata = chunk.usageMetadata;
// Context usage tracks prompt size; output isn't in history yet.
const lastPromptTokenCount =
usageMetadata.promptTokenCount || usageMetadata.totalTokenCount;
if (lastPromptTokenCount) {
// Always update the per-chat counter so this chat (including
// subagents) can make its own compaction decisions.
this.lastPromptTokenCount = lastPromptTokenCount;
// Mirror to the global telemetry only when wired — subagents
// pass `telemetryService=undefined` to keep their context usage
// out of the main session's UI counters.
this.telemetryService?.setLastPromptTokenCount(lastPromptTokenCount);
}
if (usageMetadata.cachedContentTokenCount && this.telemetryService) {
this.telemetryService.setLastCachedContentTokenCount(
usageMetadata.cachedContentTokenCount,
);
}
}
yield chunk; // Yield every chunk to the UI immediately.
} catch (e) {
streamError = e;
}
let thoughtContentPart: Part | undefined;
@ -1369,6 +1383,46 @@ export class GeminiChat {
});
}
// Mid-stream failure recovery: if the upstream stream threw (typical on
// weak networks — SSE cut between a tool_use `content_block_stop` and
// the terminal `message_stop`) AND any `functionCall` chunk was already
// yielded to consumers, we must persist the partial assistant turn here.
//
// The content generator (Anthropic / OpenAI) emits a `functionCall` part
// only at the end of a tool_use block. Once yielded, `Turn.run` registers
// a `ToolCallRequest` event, the React tool scheduler queues the call,
// and `handleCompletedTools` will fire `submitQuery(..., ToolResult)` —
// pushing a user message with `functionResponse` into history — even
// though the parent stream errored. Without preserving the matching
// tool_use on the model side, the next request body would have
// `user → user[tool_result]` with no tool_use in between, and the
// Anthropic-compatible API (DeepSeek, Anthropic, etc.) rejects with
// "tool_use_id ... must have a corresponding tool_use block in the
// previous message"
// — an unrecoverable state because Ctrl+Y's `stripOrphanedUserEntries`
// only strips trailing user entries; the lost tool_use can't be
// resurrected.
//
// Plain-text partial turns (no functionCall yielded) are deliberately
// NOT persisted — the Retry path pops the trailing user prompt and
// re-issues it; a stale partial-text model turn between them would
// either bias the retry or surface as a duplicate.
if (streamError !== null) {
if (
hasToolCall &&
(thoughtContentPart || consolidatedHistoryParts.length > 0)
) {
this.history.push({
role: 'model',
parts: [
...(thoughtContentPart ? [thoughtContentPart] : []),
...consolidatedHistoryParts,
],
});
}
throw streamError;
}
// Stream validation logic: A stream is considered successful if:
// 1. There's a tool call (tool calls can end without explicit finish reasons), OR
// 2. There's a finish reason AND we have non-empty response text or thought text