fix(coding-agent): harden chat stream completion

Flush final text before closing each AI SDK text block, surface event-processing failures to chat callers, and clear the remaining Companion OS check blockers.

fixes #273

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Harivansh Rathi 2026-03-09 12:48:21 -07:00
parent 6b2a639fb6
commit 3c0f74c1dc
5 changed files with 124 additions and 24 deletions

View file

@ -26,7 +26,6 @@ import type {
GatewaySessionState,
GatewaySessionSnapshot,
HistoryMessage,
HistoryPart,
ModelInfo,
} from "./types.js";
import {
@ -54,9 +53,9 @@ export type {
GatewaySessionState,
GatewaySessionSnapshot,
HistoryMessage,
HistoryPart,
ModelInfo,
} from "./types.js";
export type { HistoryPart } from "./types.js";
let activeGatewayRuntime: GatewayRuntime | null = null;

View file

@ -129,7 +129,8 @@ export function createVercelStreamListener(
};
const emitTextDelta = (contentIndex: number, delta: string): void => {
if (delta.length === 0) {
const state = getTextState(contentIndex);
if (delta.length === 0 || state.ended) {
return;
}
emitTextStart(contentIndex);
@ -138,7 +139,7 @@ export function createVercelStreamListener(
id: `text_${contentIndex}`,
delta,
});
getTextState(contentIndex).streamedText += delta;
state.streamedText += delta;
};
const emitTextEnd = (contentIndex: number): void => {
@ -154,20 +155,36 @@ export function createVercelStreamListener(
state.ended = true;
};
const flushTextPart = (
contentIndex: number,
fullText: string,
close: boolean,
): void => {
const state = getTextState(contentIndex);
if (state.ended) {
return;
}
if (!fullText.startsWith(state.streamedText)) {
if (close && state.started) {
emitTextEnd(contentIndex);
}
return;
}
const suffix = fullText.slice(state.streamedText.length);
if (suffix.length > 0) {
emitTextDelta(contentIndex, suffix);
}
if (close) {
emitTextEnd(contentIndex);
}
};
const flushAssistantMessageText = (
event: Extract<AgentSessionEvent, { type: "message_end" }>,
): void => {
for (const part of getAssistantTextParts(event)) {
const state = getTextState(part.contentIndex);
if (!part.text.startsWith(state.streamedText)) {
continue;
}
const suffix = part.text.slice(state.streamedText.length);
if (suffix.length > 0) {
emitTextDelta(part.contentIndex, suffix);
}
emitTextEnd(part.contentIndex);
flushTextPart(part.contentIndex, part.text, true);
}
};
@ -206,7 +223,7 @@ export function createVercelStreamListener(
emitTextDelta(inner.contentIndex, inner.delta);
return;
case "text_end":
emitTextEnd(inner.contentIndex);
flushTextPart(inner.contentIndex, inner.content, true);
return;
case "toolcall_start": {
const content = inner.partial.content[inner.contentIndex];