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

@@ -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();