Preserve Codex assistant message ids (#989)

This commit is contained in:
Mohamed Boudra
2026-05-13 21:18:31 +08:00
committed by GitHub
parent 49fa72d70d
commit e2ecae0e72
10 changed files with 283 additions and 15 deletions

View File

@@ -5,12 +5,13 @@ import { applyStreamEvent } from "@/types/stream";
const baseTimestamp = new Date(0);
const assistantChunk = (text: string): AgentStreamEventPayload => ({
const assistantChunk = (text: string, messageId?: string): AgentStreamEventPayload => ({
type: "timeline",
provider: "codex",
item: {
type: "assistant_message",
text,
...(messageId ? { messageId } : {}),
},
});
@@ -131,6 +132,36 @@ describe("applyStreamEvent", () => {
expect(result.tail[0].kind).toBe("assistant_message");
});
it("does not continue a tail assistant message when the incoming message id differs", () => {
const result = applyStreamEvent({
tail: [
{
kind: "assistant_message",
id: "msg-first",
messageId: "msg-first",
text: "First answer.",
timestamp: baseTimestamp,
},
],
head: [],
event: assistantChunk("Second answer.", "msg-second"),
timestamp: baseTimestamp,
});
expect(result.tail).toHaveLength(1);
expect(result.head).toHaveLength(1);
expect(result.tail[0].kind).toBe("assistant_message");
expect(result.head[0].kind).toBe("assistant_message");
if (
result.tail[0].kind === "assistant_message" &&
result.head[0].kind === "assistant_message"
) {
expect(result.tail[0].text).toBe("First answer.");
expect(result.head[0].text).toBe("Second answer.");
expect(result.head[0].messageId).toBe("msg-second");
}
});
it("flushes reasoning when assistant message starts", () => {
let result = applyStreamEvent({
tail: [],

View File

@@ -18,11 +18,12 @@ type CanonicalToolStatus = "running" | "completed" | "failed" | "canceled";
function assistantTimeline(
text: string,
provider: AgentProvider = "claude",
messageId?: string,
): AgentStreamEventPayload {
return {
type: "timeline",
provider,
item: { type: "assistant_message", text },
item: { type: "assistant_message", text, ...(messageId ? { messageId } : {}) },
};
}
@@ -312,6 +313,65 @@ describe("stream reducer canonical tool calls", () => {
assert.strictEqual(assistantMessage?.text, "Hello world");
});
it("keeps adjacent assistant timeline items separate when message ids differ", () => {
const state = hydrateStreamState([
{
event: assistantTimeline("First answer.", "codex", "msg-first"),
timestamp: new Date("2025-01-01T10:01:00Z"),
},
{
event: assistantTimeline("Second answer.", "codex", "msg-second"),
timestamp: new Date("2025-01-01T10:01:01Z"),
},
]);
assert.deepStrictEqual(
state.map((item) => (item.kind === "assistant_message" ? item.text : item.kind)),
["First answer.", "Second answer."],
);
assert.deepStrictEqual(
state.map((item) => (item.kind === "assistant_message" ? item.messageId : null)),
["msg-first", "msg-second"],
);
});
it("merges adjacent assistant deltas when message ids match", () => {
const state = hydrateStreamState([
{
event: assistantTimeline("Hel", "codex", "msg-same"),
timestamp: new Date("2025-01-01T10:02:00Z"),
},
{
event: assistantTimeline("lo", "codex", "msg-same"),
timestamp: new Date("2025-01-01T10:02:01Z"),
},
]);
assert.strictEqual(state.length, 1);
assert.strictEqual(state[0]?.kind, "assistant_message");
if (state[0]?.kind === "assistant_message") {
assert.strictEqual(state[0].text, "Hello");
assert.strictEqual(state[0].id, "msg-same");
assert.strictEqual(state[0].messageId, "msg-same");
}
});
it("preserves old assistant merge behavior when message ids are absent", () => {
const state = hydrateStreamState([
{
event: assistantTimeline("Hel", "codex"),
timestamp: new Date("2025-01-01T10:03:00Z"),
},
{
event: assistantTimeline("lo", "codex"),
timestamp: new Date("2025-01-01T10:03:01Z"),
},
]);
assert.strictEqual(state.length, 1);
assert.strictEqual(state[0]?.kind === "assistant_message" ? state[0].text : null, "Hello");
});
it("merges running and completed events by callId", () => {
const callId = "tool-merge-1";
const updates = [

View File

@@ -66,6 +66,7 @@ export interface UserMessageItem {
export interface AssistantMessageItem {
kind: "assistant_message";
id: string;
messageId?: string;
text: string;
timestamp: Date;
blockGroupId?: string;
@@ -227,6 +228,7 @@ function appendAssistantMessage(
text: string,
timestamp: Date,
source: StreamUpdateSource,
messageId?: string,
): StreamItem[] {
const { chunk, hasContent } = normalizeChunk(text);
if (!chunk) {
@@ -234,7 +236,11 @@ function appendAssistantMessage(
}
const last = state[state.length - 1];
if (last && last.kind === "assistant_message") {
const shouldAppendToLast =
last &&
last.kind === "assistant_message" &&
(messageId === undefined || last.messageId === messageId);
if (shouldAppendToLast) {
const updated: AssistantMessageItem = {
...last,
text: `${last.text}${chunk}`,
@@ -249,7 +255,8 @@ function appendAssistantMessage(
if (
source === "live" &&
last?.kind === "user_message" &&
secondLast?.kind === "assistant_message"
secondLast?.kind === "assistant_message" &&
(messageId === undefined || secondLast.messageId === messageId)
) {
const updated: AssistantMessageItem = {
...secondLast,
@@ -264,9 +271,11 @@ function appendAssistantMessage(
}
const idSeed = chunk.trim() || chunk;
const entryId = messageId ?? createUniqueTimelineId(state, "assistant", idSeed, timestamp);
const item: AssistantMessageItem = {
kind: "assistant_message",
id: createUniqueTimelineId(state, "assistant", idSeed, timestamp),
id: entryId,
...(messageId ? { messageId } : {}),
text: chunk,
timestamp,
};
@@ -647,7 +656,9 @@ function reduceTimelineEvent(
case "user_message":
return finalizeActiveThoughts(appendUserMessage(state, item.text, timestamp, item.messageId));
case "assistant_message":
return finalizeActiveThoughts(appendAssistantMessage(state, item.text, timestamp, source));
return finalizeActiveThoughts(
appendAssistantMessage(state, item.text, timestamp, source, item.messageId),
);
case "reasoning":
return appendThought(state, item.text, timestamp);
case "tool_call":
@@ -993,7 +1004,14 @@ export function applyStreamEvent(params: {
if (incomingKind === "assistant_message" && nextHead.length === 0) {
const tailAssistant = nextTail.at(-1);
if (tailAssistant?.kind === "assistant_message") {
const incomingMessageId =
event.type === "timeline" && event.item.type === "assistant_message"
? event.item.messageId
: undefined;
const shouldContinueTailAssistant =
tailAssistant?.kind === "assistant_message" &&
(incomingMessageId === undefined || tailAssistant.messageId === incomingMessageId);
if (shouldContinueTailAssistant) {
nextTail = nextTail.slice(0, -1);
nextHead = [tailAssistant];
changedTail = true;

View File

@@ -304,7 +304,7 @@ export interface CompactionTimelineItem {
export type AgentTimelineItem =
| { type: "user_message"; text: string; messageId?: string }
| { type: "assistant_message"; text: string }
| { type: "assistant_message"; text: string; messageId?: string }
| { type: "reasoning"; text: string }
| ToolCallTimelineItem
| { type: "todo"; items: { text: string; completed: boolean }[] }

View File

@@ -1315,6 +1315,66 @@ describe("Codex app-server provider", () => {
item: {
type: "assistant_message",
text: "History loaded.",
messageId: "message-history",
},
},
]);
});
test("preserves Codex app-server assistant item ids in persisted history", async () => {
const session = createSession();
session.client = {
request: vi.fn(async (method: string) => {
if (method !== "thread/read") {
return {};
}
return {
thread: {
turns: [
{
items: [
{
type: "agentMessage",
id: "before-tool-message",
text: "I checked the workspace.",
},
{
type: "agentMessage",
id: "after-tool-message",
text: "The tests are green.",
},
],
},
],
},
};
}),
};
await asInternals(session).loadPersistedHistory();
const history: AgentStreamEvent[] = [];
for await (const event of session.streamHistory()) {
history.push(event);
}
expect(history).toEqual([
{
type: "timeline",
provider: "codex",
item: {
type: "assistant_message",
text: "I checked the workspace.",
messageId: "before-tool-message",
},
},
{
type: "timeline",
provider: "codex",
item: {
type: "assistant_message",
text: "The tests are green.",
messageId: "after-tool-message",
},
},
]);
@@ -1768,13 +1828,13 @@ describe("Codex app-server provider", () => {
type: "timeline",
provider: "codex",
turnId: "test-turn",
item: { type: "assistant_message", text: "Hel" },
item: { type: "assistant_message", text: "Hel", messageId: "assistant-item-1" },
},
{
type: "timeline",
provider: "codex",
turnId: "test-turn",
item: { type: "assistant_message", text: "lo" },
item: { type: "assistant_message", text: "lo", messageId: "assistant-item-1" },
},
]);
});
@@ -1805,19 +1865,19 @@ describe("Codex app-server provider", () => {
type: "timeline",
provider: "codex",
turnId: "test-turn",
item: { type: "assistant_message", text: "Hel" },
item: { type: "assistant_message", text: "Hel", messageId: "assistant-item-2" },
},
{
type: "timeline",
provider: "codex",
turnId: "test-turn",
item: { type: "assistant_message", text: "lo" },
item: { type: "assistant_message", text: "lo", messageId: "assistant-item-2" },
},
{
type: "timeline",
provider: "codex",
turnId: "test-turn",
item: { type: "assistant_message", text: "!" },
item: { type: "assistant_message", text: "!", messageId: "assistant-item-2" },
},
]);
});
@@ -1852,6 +1912,7 @@ describe("Codex app-server provider", () => {
turnId: "test-turn",
item: {
type: "assistant_message",
messageId: "assistant-item-3",
text: "Im in the waiting phase now. The next read is intentionally delayed so we get meaningful CI state instead of churn.",
},
},
@@ -1861,6 +1922,7 @@ describe("Codex app-server provider", () => {
turnId: "test-turn",
item: {
type: "assistant_message",
messageId: "assistant-item-4",
text: "\n\n---\n\nCI is still cooking. Im staying on the current run rather than jumping around, because the first red job will tell us exactly whether anything else needs work.",
},
},

View File

@@ -1534,11 +1534,14 @@ function threadItemToTimeline(
switch (normalizedType) {
case "userMessage":
return mapCodexThreadUserMessageItem(normalizedItem, includeUserMessage);
case "agentMessage":
case "agentMessage": {
const messageId = nonEmptyString(normalizedItem.id);
return {
type: "assistant_message",
text: typeof normalizedItem.text === "string" ? normalizedItem.text : "",
...(messageId ? { messageId } : {}),
};
}
case "plan":
return mapCodexThreadPlanItem(normalizedItem);
case "reasoning":
@@ -4073,6 +4076,7 @@ class CodexAppServerAgentSession implements AgentSession {
if (subAgentCallId) {
this.upsertSubAgentChildItem(subAgentCallId, parsed.itemId, {
type: "assistant_message",
messageId: parsed.itemId,
text,
});
this.emitSubAgentActivityUpdate(subAgentCallId, "running");
@@ -4084,6 +4088,7 @@ class CodexAppServerAgentSession implements AgentSession {
provider: CODEX_PROVIDER,
item: {
type: "assistant_message",
messageId: parsed.itemId,
text:
isFirstDeltaForItem && this.pendingAssistantMessageBoundary
? `${ASSISTANT_MESSAGE_BOUNDARY_MARKDOWN}${parsed.delta}`
@@ -4443,7 +4448,14 @@ class CodexAppServerAgentSession implements AgentSession {
this.emitEvent({
type: "timeline",
provider: CODEX_PROVIDER,
item: { type: timelineItem.type, text: suffix },
item:
timelineItem.type === "assistant_message"
? {
type: timelineItem.type,
text: suffix,
...(timelineItem.messageId ? { messageId: timelineItem.messageId } : {}),
}
: { type: timelineItem.type, text: suffix },
});
}

View File

@@ -39,6 +39,59 @@ describe("projectTimelineRows", () => {
expect(projected[0]?.collapsed).toContain("assistant_merge");
});
test("merges adjacent assistant chunks with the same message id in projected mode", () => {
const rows: AgentTimelineRow[] = [
{
seq: 1,
timestamp: "2026-02-13T00:00:00.000Z",
item: { type: "assistant_message", text: "Hel", messageId: "msg-1" },
},
{
seq: 2,
timestamp: "2026-02-13T00:00:00.100Z",
item: { type: "assistant_message", text: "lo", messageId: "msg-1" },
},
];
const projected = projectTimelineRows({ rows, mode: "projected" });
expect(projected).toHaveLength(1);
expect(projected[0]?.item).toEqual({
type: "assistant_message",
text: "Hello",
messageId: "msg-1",
});
});
test("keeps adjacent assistant chunks with different message ids separate in projected mode", () => {
const rows: AgentTimelineRow[] = [
{
seq: 1,
timestamp: "2026-02-13T00:00:00.000Z",
item: { type: "assistant_message", text: "First answer.", messageId: "msg-1" },
},
{
seq: 2,
timestamp: "2026-02-13T00:00:00.100Z",
item: { type: "assistant_message", text: "Second answer.", messageId: "msg-2" },
},
];
const projected = projectTimelineRows({ rows, mode: "projected" });
expect(projected).toHaveLength(2);
expect(projected[0]?.item).toEqual({
type: "assistant_message",
text: "First answer.",
messageId: "msg-1",
});
expect(projected[1]?.item).toEqual({
type: "assistant_message",
text: "Second answer.",
messageId: "msg-2",
});
});
test("merges adjacent reasoning chunks in projected mode", () => {
const rows: AgentTimelineRow[] = [
{

View File

@@ -210,6 +210,13 @@ function mergeAssistantChunks(entries: readonly WorkingEntry[]): WorkingEntry[]
{ type: "assistant_message" }
>;
const entryAssistant = entry.item as Extract<AgentTimelineItem, { type: "assistant_message" }>;
if (
entryAssistant.messageId !== undefined &&
previousAssistant.messageId !== entryAssistant.messageId
) {
output.push(entry);
continue;
}
const collapsedKinds = new Set<TimelineProjectionKind>([
...previous.collapsed,
@@ -222,6 +229,7 @@ function mergeAssistantChunks(entries: readonly WorkingEntry[]): WorkingEntry[]
item: {
type: "assistant_message",
text: `${previousAssistant.text}${entryAssistant.text}`,
...(previousAssistant.messageId ? { messageId: previousAssistant.messageId } : {}),
},
seqEnd: entry.seqEnd,
sourceSeqRanges: mergeSeqRanges(previous.sourceSeqRanges, entry.sourceSeqRanges),

View File

@@ -505,6 +505,7 @@ export const AgentTimelineItemPayloadSchema: z.ZodType<AgentTimelineItem, z.ZodT
z.object({
type: z.literal("assistant_message"),
text: z.string(),
messageId: z.string().optional(),
}),
z.object({
type: z.literal("reasoning"),

View File

@@ -194,6 +194,29 @@ async function emitTimelineResponse(
}
describe("wire compatibility", () => {
test("assistant timeline message ids are optional on the wire", () => {
expect(
AgentTimelineItemPayloadSchema.parse({
type: "assistant_message",
text: "old daemon shape",
}),
).toEqual({
type: "assistant_message",
text: "old daemon shape",
});
expect(
AgentTimelineItemPayloadSchema.parse({
type: "assistant_message",
text: "new daemon shape",
messageId: "msg-1",
}),
).toEqual({
type: "assistant_message",
text: "new daemon shape",
messageId: "msg-1",
});
});
test("downgrades reasoning_merge for clients that do not declare the capability", async () => {
const response = await emitTimelineResponse();