fix: restore codex streaming responsiveness

This commit is contained in:
Mohamed Boudra
2026-05-01 16:18:46 +07:00
parent 1c4c22b060
commit af38c6e135
7 changed files with 546 additions and 40 deletions
@@ -103,7 +103,7 @@ describe("AgentStreamCoalescer", () => {
expect(coalescer.handle("agent-1", assistant("lo"))).toBe(true);
expect(flushes).toEqual([]);
await vi.advanceTimersByTimeAsync(199);
await vi.advanceTimersByTimeAsync(59);
expect(flushes).toEqual([]);
await vi.advanceTimersByTimeAsync(1);
@@ -119,7 +119,7 @@ describe("AgentStreamCoalescer", () => {
test("uses constructor windowMs instead of a hard-coded value", async () => {
const { coalescer, flushes } = createHarness(10);
expect(AGENT_STREAM_COALESCE_DEFAULT_WINDOW_MS).toBe(200);
expect(AGENT_STREAM_COALESCE_DEFAULT_WINDOW_MS).toBe(60);
expect(coalescer.handle("agent-1", assistant("fast"))).toBe(true);
await vi.advanceTimersByTimeAsync(9);
@@ -203,7 +203,7 @@ describe("AgentStreamCoalescer", () => {
coalescer.handle("agent-1", reasoning("r1"));
coalescer.handle("agent-1", assistant("a2"));
await vi.advanceTimersByTimeAsync(200);
await vi.advanceTimersByTimeAsync(60);
expect(flushes).toEqual([
{
agentId: "agent-1",
@@ -231,7 +231,7 @@ describe("AgentStreamCoalescer", () => {
coalescer.handle("agent-1", assistant("a2"));
coalescer.handle("agent-1", reasoning("r2"));
await vi.advanceTimersByTimeAsync(200);
await vi.advanceTimersByTimeAsync(60);
expect(flushes.map((flush) => flush.item)).toEqual([
{ type: "assistant_message", text: "a1" },
{ type: "reasoning", text: "r1" },
@@ -246,7 +246,7 @@ describe("AgentStreamCoalescer", () => {
coalescer.handle("agent-1", assistant("c", { provider: "codex" }));
coalescer.handle("agent-1", assistant("o", { provider: "opencode" }));
await vi.advanceTimersByTimeAsync(200);
await vi.advanceTimersByTimeAsync(60);
expect(flushes).toEqual([
{
agentId: "agent-1",
@@ -267,7 +267,7 @@ describe("AgentStreamCoalescer", () => {
coalescer.handle("agent-1", assistant("one", { turnId: "turn-1" }));
coalescer.handle("agent-1", assistant("two", { turnId: "turn-2" }));
await vi.advanceTimersByTimeAsync(200);
await vi.advanceTimersByTimeAsync(60);
expect(flushes).toEqual([
{
agentId: "agent-1",
@@ -290,7 +290,7 @@ describe("AgentStreamCoalescer", () => {
expect(coalescer.handle("agent-1", assistant(""))).toBe(true);
expect(coalescer.handle("agent-2", reasoning(""))).toBe(true);
await vi.advanceTimersByTimeAsync(200);
await vi.advanceTimersByTimeAsync(60);
expect(flushes).toEqual([]);
});
@@ -302,7 +302,7 @@ describe("AgentStreamCoalescer", () => {
coalescer.handle("agent-1", assistant("\t"));
coalescer.handle("agent-1", assistant(" mixed \n\t whitespace "));
await vi.advanceTimersByTimeAsync(200);
await vi.advanceTimersByTimeAsync(60);
expect(flushes).toEqual([
{
agentId: "agent-1",
@@ -332,7 +332,7 @@ describe("AgentStreamCoalescer", () => {
coalescer.handle("agent-1", assistant(chunk));
}
await vi.advanceTimersByTimeAsync(200);
await vi.advanceTimersByTimeAsync(60);
expect(flushes).toEqual([
{
agentId: "agent-1",
@@ -350,7 +350,7 @@ describe("AgentStreamCoalescer", () => {
coalescer.handle("agent-1", assistant("b"));
coalescer.handle("agent-2", assistant("y"));
await vi.advanceTimersByTimeAsync(200);
await vi.advanceTimersByTimeAsync(60);
expect(flushes).toEqual([
{
agentId: "agent-1",
@@ -431,7 +431,7 @@ describe("AgentStreamCoalescer", () => {
},
]);
await vi.advanceTimersByTimeAsync(200);
await vi.advanceTimersByTimeAsync(60);
expect(flushes).toEqual([
{
agentId: "agent-1",
@@ -452,7 +452,7 @@ describe("AgentStreamCoalescer", () => {
coalescer.handle("agent-1", assistant("manual"));
coalescer.flushFor("agent-1");
await vi.advanceTimersByTimeAsync(200);
await vi.advanceTimersByTimeAsync(60);
expect(flushes).toEqual([
{
agentId: "agent-1",
@@ -468,7 +468,7 @@ describe("AgentStreamCoalescer", () => {
coalescer.handle("agent-1", assistant("durable"));
coalescer.flushAndDiscard("agent-1");
await vi.advanceTimersByTimeAsync(200);
await vi.advanceTimersByTimeAsync(60);
expect(flushes).toEqual([
{
agentId: "agent-1",
@@ -532,7 +532,7 @@ describe("AgentStreamCoalescer", () => {
coalescer.handle("agent-1", timeline(firstItem));
coalescer.handle("agent-1", assistant("llo"));
await vi.advanceTimersByTimeAsync(200);
await vi.advanceTimersByTimeAsync(60);
expect(flushes).toEqual([
{
agentId: "agent-1",
@@ -552,7 +552,7 @@ describe("AgentStreamCoalescer", () => {
expect(coalescer.handle("agent-1", toolCall({ output: "first" }))).toBe(true);
expect(coalescer.handle("agent-1", toolCall({ output: "second" }))).toBe(true);
await vi.advanceTimersByTimeAsync(200);
await vi.advanceTimersByTimeAsync(60);
expect(flushes).toEqual([
{
@@ -583,7 +583,7 @@ describe("AgentStreamCoalescer", () => {
coalescer.handle("agent-1", toolCall({ callId: "tool-1", output: "one-b" }));
coalescer.handle("agent-1", toolCall({ callId: "tool-2", output: "two-b" }));
await vi.advanceTimersByTimeAsync(200);
await vi.advanceTimersByTimeAsync(60);
expect(flushes.map((flush) => flush.item)).toEqual([
{
@@ -651,7 +651,7 @@ describe("AgentStreamCoalescer", () => {
coalescer.handle("agent-1", toolCall({ output: "latest" }));
coalescer.handle("agent-1", assistant("b"));
await vi.advanceTimersByTimeAsync(200);
await vi.advanceTimersByTimeAsync(60);
expect(flushes.map((flush) => flush.item)).toEqual([
{ type: "assistant_message", text: "a" },
@@ -1,6 +1,6 @@
import type { AgentProvider, AgentStreamEvent, AgentTimelineItem } from "./agent-sdk-types.js";
export const AGENT_STREAM_COALESCE_DEFAULT_WINDOW_MS = 200;
export const AGENT_STREAM_COALESCE_DEFAULT_WINDOW_MS = 60;
type CoalescableTextKind = "assistant_message" | "reasoning";
type CoalescableTimelineKind = CoalescableTextKind | "tool_call";
@@ -903,6 +903,166 @@ describe("Codex app-server provider", () => {
});
});
test("streams Codex assistant message deltas and does not replay completed text", () => {
const session = createSession();
const events: AgentStreamEvent[] = [];
session.subscribe((event) => events.push(event));
asInternals(session).handleNotification("item/agentMessage/delta", {
itemId: "assistant-item-1",
delta: "Hel",
});
asInternals(session).handleNotification("item/agentMessage/delta", {
itemId: "assistant-item-1",
delta: "lo",
});
asInternals(session).handleNotification("item/completed", {
item: {
id: "assistant-item-1",
type: "agentMessage",
text: "Hello",
},
});
expect(events).toEqual([
{
type: "timeline",
provider: "codex",
turnId: "test-turn",
item: { type: "assistant_message", text: "Hel" },
},
{
type: "timeline",
provider: "codex",
turnId: "test-turn",
item: { type: "assistant_message", text: "lo" },
},
]);
});
test("emits only the missing assistant suffix when completed text extends streamed deltas", () => {
const session = createSession();
const events: AgentStreamEvent[] = [];
session.subscribe((event) => events.push(event));
asInternals(session).handleNotification("item/agentMessage/delta", {
itemId: "assistant-item-2",
delta: "Hel",
});
asInternals(session).handleNotification("item/agentMessage/delta", {
itemId: "assistant-item-2",
delta: "lo",
});
asInternals(session).handleNotification("item/completed", {
item: {
id: "assistant-item-2",
type: "agentMessage",
text: "Hello!",
},
});
expect(events).toEqual([
{
type: "timeline",
provider: "codex",
turnId: "test-turn",
item: { type: "assistant_message", text: "Hel" },
},
{
type: "timeline",
provider: "codex",
turnId: "test-turn",
item: { type: "assistant_message", text: "lo" },
},
{
type: "timeline",
provider: "codex",
turnId: "test-turn",
item: { type: "assistant_message", text: "!" },
},
]);
});
test("streams Codex reasoning deltas and does not replay completed reasoning", () => {
const session = createSession();
const events: AgentStreamEvent[] = [];
session.subscribe((event) => events.push(event));
asInternals(session).handleNotification("item/reasoning/summaryTextDelta", {
itemId: "reasoning-item-1",
delta: "Think",
});
asInternals(session).handleNotification("item/reasoning/summaryTextDelta", {
itemId: "reasoning-item-1",
delta: "ing",
});
asInternals(session).handleNotification("item/completed", {
item: {
id: "reasoning-item-1",
type: "reasoning",
summary: ["Thinking"],
},
});
expect(events).toEqual([
{
type: "timeline",
provider: "codex",
turnId: "test-turn",
item: { type: "reasoning", text: "Think" },
},
{
type: "timeline",
provider: "codex",
turnId: "test-turn",
item: { type: "reasoning", text: "ing" },
},
]);
});
test("emits only the missing reasoning suffix when completed reasoning extends streamed deltas", () => {
const session = createSession();
const events: AgentStreamEvent[] = [];
session.subscribe((event) => events.push(event));
asInternals(session).handleNotification("item/reasoning/summaryTextDelta", {
itemId: "reasoning-item-2",
delta: "Think",
});
asInternals(session).handleNotification("item/reasoning/summaryTextDelta", {
itemId: "reasoning-item-2",
delta: "ing",
});
asInternals(session).handleNotification("item/completed", {
item: {
id: "reasoning-item-2",
type: "reasoning",
summary: ["Thinking!"],
},
});
expect(events).toEqual([
{
type: "timeline",
provider: "codex",
turnId: "test-turn",
item: { type: "reasoning", text: "Think" },
},
{
type: "timeline",
provider: "codex",
turnId: "test-turn",
item: { type: "reasoning", text: "ing" },
},
{
type: "timeline",
provider: "codex",
turnId: "test-turn",
item: { type: "reasoning", text: "!" },
},
]);
});
test("approving a synthetic Codex plan permission disables plan mode, preserves fast mode, and returns follow-up prompt", async () => {
const session = createSession({
featureValues: { plan_mode: true, fast_mode: true },
@@ -3541,11 +3541,6 @@ class CodexAppServerAgentSession implements AgentSession {
}
private emitEvent(event: AgentStreamEvent): void {
if (event.type === "timeline") {
if (event.item.type === "assistant_message") {
this.pendingAgentMessages.clear();
}
}
this.notifySubscribers(event);
}
@@ -3638,12 +3633,22 @@ class CodexAppServerAgentSession implements AgentSession {
if (parsed.kind === "agent_message_delta") {
const prev = this.pendingAgentMessages.get(parsed.itemId) ?? "";
this.pendingAgentMessages.set(parsed.itemId, prev + parsed.delta);
this.emitEvent({
type: "timeline",
provider: CODEX_PROVIDER,
item: { type: "assistant_message", text: parsed.delta },
});
return;
}
if (parsed.kind === "reasoning_delta") {
const prev = this.pendingReasoning.get(parsed.itemId) ?? [];
prev.push(parsed.delta);
this.pendingReasoning.set(parsed.itemId, prev);
this.emitEvent({
type: "timeline",
provider: CODEX_PROVIDER,
item: { type: "reasoning", text: parsed.delta },
});
return;
}
if (parsed.kind === "exec_command_output_delta") {
@@ -3705,6 +3710,8 @@ class CodexAppServerAgentSession implements AgentSession {
this.emittedItemCompletedIds.clear();
this.emittedExecCommandStartedCallIds.clear();
this.emittedExecCommandCompletedCallIds.clear();
this.pendingAgentMessages.clear();
this.pendingReasoning.clear();
this.pendingCommandOutputDeltas.clear();
this.pendingFileChangeOutputDeltas.clear();
this.warnedIncompleteEditToolCallIds.clear();
@@ -3885,6 +3892,13 @@ class CodexAppServerAgentSession implements AgentSession {
if (itemId && this.emittedItemCompletedIds.has(itemId)) {
return;
}
if (this.consumeStreamedTextCompletion(timelineItem, itemId)) {
if (itemId) {
this.emittedItemCompletedIds.add(itemId);
this.emittedItemStartedIds.delete(itemId);
}
return;
}
this.applyBufferedDeltaTextToTimelineItem(timelineItem, itemId);
if (timelineItem.type === "tool_call") {
if (timelineItem.detail.type === "plan") {
@@ -3906,6 +3920,47 @@ class CodexAppServerAgentSession implements AgentSession {
}
}
private consumeStreamedTextCompletion(
timelineItem: AgentTimelineItem,
itemId: string | null | undefined,
): boolean {
if (!itemId) {
return false;
}
if (timelineItem.type === "assistant_message" && this.pendingAgentMessages.has(itemId)) {
const streamedText = this.pendingAgentMessages.get(itemId) ?? "";
this.pendingAgentMessages.delete(itemId);
this.emitMissingFinalTextSuffix(timelineItem, streamedText);
return true;
}
if (timelineItem.type === "reasoning" && this.pendingReasoning.has(itemId)) {
const streamedText = this.pendingReasoning.get(itemId)?.join("") ?? "";
this.pendingReasoning.delete(itemId);
this.emitMissingFinalTextSuffix(timelineItem, streamedText);
return true;
}
return false;
}
private emitMissingFinalTextSuffix(
timelineItem: Extract<AgentTimelineItem, { type: "assistant_message" | "reasoning" }>,
streamedText: string,
): void {
if (!timelineItem.text.startsWith(streamedText)) {
this.emitEvent({ type: "timeline", provider: CODEX_PROVIDER, item: timelineItem });
return;
}
const suffix = timelineItem.text.slice(streamedText.length);
if (!suffix) {
return;
}
this.emitEvent({
type: "timeline",
provider: CODEX_PROVIDER,
item: { type: timelineItem.type, text: suffix },
});
}
private applyBufferedDeltaTextToTimelineItem(
timelineItem: AgentTimelineItem,
itemId: string | null | undefined,