Skip to content

Commit 8d9ce2a

Browse files
nikita-ashihminnikita-ashihmin
committed
Expose Codex agent message phases in ACP chunks
1 parent c516b9a commit 8d9ce2a

7 files changed

Lines changed: 201 additions & 15 deletions

‎src/CodexAcpServer.ts‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,7 @@ import packageJson from "../package.json";
6969
import {isJetBrains2026_1Client} from "./JBUtils";
7070
import {resolveTerminalOutputMode, type TerminalOutputMode} from "./TerminalOutputMode";
7171
import {
72+
createCodexMessagePhaseMeta,
7273
createAgentTextMessageChunk,
7374
createAgentTextThoughtChunk,
7475
createUserMessageChunk,
@@ -983,12 +984,15 @@ export class CodexAcpServer {
983984
case "subAgentActivity":
984985
case "sleep":
985986
return [];
986-
case "agentMessage":
987+
case "agentMessage": {
988+
const meta = createCodexMessagePhaseMeta(item.phase);
987989
return [{
988990
sessionUpdate: "agent_message_chunk",
989991
messageId: item.id,
990992
content: { type: "text", text: item.text },
993+
...(meta ? { _meta: meta } : {}),
991994
}];
995+
}
992996
case "reasoning":
993997
return this.createReasoningUpdates(item);
994998
case "fileChange":

‎src/CodexEventHandler.ts‎

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@ import {
5656
import { stripShellPrefix } from "./CommandUtils";
5757
import {createTerminalOutputMeta, type TerminalOutputMode} from "./TerminalOutputMode";
5858
import {
59+
createCodexMessagePhaseMeta,
5960
createAgentTextMessageChunk,
6061
createAgentTextThoughtChunk,
6162
} from "./ContentChunks";
@@ -74,6 +75,7 @@ export class CodexEventHandler {
7475
private readonly seenReasoningDeltaItemIds = new Set<string>();
7576
private readonly terminalCommandIds = new Set<string>();
7677
private readonly terminalCommandOutputIds = new Set<string>();
78+
private readonly agentMessagePhases = new Map<string, string | null>();
7779

7880
constructor(connection: AcpClientConnection, sessionState: SessionState) {
7981
this.connection = connection;
@@ -230,7 +232,8 @@ export class CodexEventHandler {
230232
}
231233

232234
private async createTextEvent(event: AgentMessageDeltaNotification): Promise<UpdateSessionEvent> {
233-
return createAgentTextMessageChunk(event.delta, event.itemId);
235+
const phase = this.agentMessagePhases.get(event.itemId) ?? null;
236+
return createAgentTextMessageChunk(event.delta, event.itemId, createCodexMessagePhaseMeta(phase));
234237
}
235238

236239
private async createConfigWarningEvent(event: ConfigWarningNotification): Promise<UpdateSessionEvent> {
@@ -331,11 +334,13 @@ export class CodexEventHandler {
331334
return createImageGenerationStartUpdate(event.item);
332335
case "collabAgentToolCall":
333336
return createCollabAgentToolCallUpdate(event.item);
337+
case "agentMessage":
338+
this.rememberAgentMessagePhase(event.item);
339+
return null;
334340
case "subAgentActivity":
335341
case "sleep":
336342
case "userMessage":
337343
case "hookPrompt":
338-
case "agentMessage":
339344
case "reasoning":
340345
case "enteredReviewMode":
341346
case "exitedReviewMode":
@@ -383,6 +388,9 @@ export class CodexEventHandler {
383388
return createWebSearchCompleteUpdate(event.item);
384389
case "collabAgentToolCall":
385390
return createCollabAgentToolCallCompleteUpdate(event.item);
391+
case "agentMessage":
392+
this.rememberAgentMessagePhase(event.item);
393+
return null;
386394
case "exitedReviewMode":
387395
return this.createExitedReviewModeEvent(event.item);
388396
case "contextCompaction":
@@ -392,14 +400,17 @@ export class CodexEventHandler {
392400
case "sleep":
393401
case "userMessage":
394402
case "hookPrompt":
395-
case "agentMessage":
396403
case "enteredReviewMode":
397404
case "plan":
398405
return null;
399406

400407
}
401408
}
402409

410+
private rememberAgentMessagePhase(item: ThreadItem & { type: "agentMessage" }): void {
411+
this.agentMessagePhases.set(item.id, item.phase);
412+
}
413+
403414
private createCompletedReasoningEvent(item: ThreadItem & { type: "reasoning" }): UpdateSessionEvent | null {
404415
const parts = item.summary.length > 0 ? item.summary : item.content;
405416
const text = parts.filter(part => part.length > 0).join("\n\n");

‎src/ContentChunks.ts‎

Lines changed: 22 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,52 +1,67 @@
11
import type {ContentBlock} from "@agentclientprotocol/sdk";
22
import type {UpdateSessionEvent} from "./ACPSessionConnection";
33

4-
export function createUserMessageChunk(content: ContentBlock, messageId?: string): UpdateSessionEvent {
4+
type AcpMeta = Record<string, unknown>;
5+
6+
export function createCodexMessagePhaseMeta(phase: string | null | undefined): AcpMeta | undefined {
7+
if (!phase) {
8+
return undefined;
9+
}
10+
return { codex: { phase } };
11+
}
12+
13+
export function createUserMessageChunk(content: ContentBlock, messageId?: string, meta?: AcpMeta): UpdateSessionEvent {
514
if (messageId) {
615
return {
716
sessionUpdate: "user_message_chunk",
817
messageId,
918
content,
19+
...(meta ? { _meta: meta } : {}),
1020
};
1121
}
1222
return {
1323
sessionUpdate: "user_message_chunk",
1424
content,
25+
...(meta ? { _meta: meta } : {}),
1526
};
1627
}
1728

18-
export function createAgentMessageChunk(content: ContentBlock, messageId?: string): UpdateSessionEvent {
29+
export function createAgentMessageChunk(content: ContentBlock, messageId?: string, meta?: AcpMeta): UpdateSessionEvent {
1930
if (messageId) {
2031
return {
2132
sessionUpdate: "agent_message_chunk",
2233
messageId,
2334
content,
35+
...(meta ? { _meta: meta } : {}),
2436
};
2537
}
2638
return {
2739
sessionUpdate: "agent_message_chunk",
2840
content,
41+
...(meta ? { _meta: meta } : {}),
2942
};
3043
}
3144

32-
export function createAgentThoughtChunk(content: ContentBlock, messageId?: string): UpdateSessionEvent {
45+
export function createAgentThoughtChunk(content: ContentBlock, messageId?: string, meta?: AcpMeta): UpdateSessionEvent {
3346
if (messageId) {
3447
return {
3548
sessionUpdate: "agent_thought_chunk",
3649
messageId,
3750
content,
51+
...(meta ? { _meta: meta } : {}),
3852
};
3953
}
4054
return {
4155
sessionUpdate: "agent_thought_chunk",
4256
content,
57+
...(meta ? { _meta: meta } : {}),
4358
};
4459
}
4560

46-
export function createAgentTextMessageChunk(text: string, messageId?: string): UpdateSessionEvent {
47-
return createAgentMessageChunk({type: "text", text}, messageId);
61+
export function createAgentTextMessageChunk(text: string, messageId?: string, meta?: AcpMeta): UpdateSessionEvent {
62+
return createAgentMessageChunk({type: "text", text}, messageId, meta);
4863
}
4964

50-
export function createAgentTextThoughtChunk(text: string, messageId?: string): UpdateSessionEvent {
51-
return createAgentThoughtChunk({type: "text", text}, messageId);
65+
export function createAgentTextThoughtChunk(text: string, messageId?: string, meta?: AcpMeta): UpdateSessionEvent {
66+
return createAgentThoughtChunk({type: "text", text}, messageId, meta);
5267
}

‎src/ResponseItemHistoryFallback.ts‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import { stripShellPrefix } from "./CommandUtils";
66
import type { CommandAction, Thread, ThreadItem } from "./app-server/v2";
77
import { createCommandActionEvent } from "./CodexToolCallMapper";
88
import { createTerminalOutputMeta, type TerminalOutputMode } from "./TerminalOutputMode";
9+
import { createAgentMessageChunk, createCodexMessagePhaseMeta } from "./ContentChunks";
910

1011
type JsonRecord = Record<string, unknown>;
1112
type AcpToolCallEvent = Extract<UpdateSessionEvent, { sessionUpdate: "tool_call" }>;
@@ -234,10 +235,10 @@ function createMessageUpdates(item: JsonRecord): UpdateSessionEvent[] {
234235
return [];
235236
}
236237

237-
return contentBlocksFromResponseContent(item["content"]).map((content) => ({
238-
sessionUpdate: "agent_message_chunk",
239-
content,
240-
}));
238+
const phase = stringValue(item["phase"]);
239+
return contentBlocksFromResponseContent(item["content"]).map((content) => (
240+
createAgentMessageChunk(content, undefined, createCodexMessagePhaseMeta(phase))
241+
));
241242
}
242243

243244
function createEventMsgUpdates(record: JsonRecord): UpdateSessionEvent[] | null {

‎src/__tests__/CodexACPAgent/agent-message-events.test.ts‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
import { beforeEach, describe, expect, it, vi } from "vitest";
2+
import type { ServerNotification } from "../../app-server";
3+
import type { SessionState } from "../../CodexAcpServer";
4+
import { AgentMode } from "../../AgentMode";
5+
import {
6+
createCodexMockTestFixture,
7+
createTestSessionState,
8+
setupPromptAndSendNotifications,
9+
type CodexMockTestFixture
10+
} from "../acp-test-utils";
11+
12+
describe("CodexEventHandler - agent message events", () => {
13+
let mockFixture: CodexMockTestFixture;
14+
const sessionId = "test-session-id";
15+
16+
beforeEach(() => {
17+
mockFixture = createCodexMockTestFixture();
18+
vi.clearAllMocks();
19+
});
20+
21+
const sessionState: SessionState = createTestSessionState({
22+
sessionId,
23+
currentModelId: "model-id[effort]",
24+
agentMode: AgentMode.DEFAULT_AGENT_MODE
25+
});
26+
27+
it("includes Codex message phase metadata on streamed agent message chunks", async () => {
28+
const notifications: ServerNotification[] = [
29+
{
30+
method: "item/started",
31+
params: {
32+
threadId: sessionId,
33+
turnId: "turn-1",
34+
startedAtMs: 0,
35+
item: {
36+
type: "agentMessage",
37+
id: "commentary-message",
38+
text: "",
39+
phase: "commentary",
40+
memoryCitation: null,
41+
},
42+
},
43+
},
44+
{
45+
method: "item/agentMessage/delta",
46+
params: {
47+
threadId: sessionId,
48+
turnId: "turn-1",
49+
itemId: "commentary-message",
50+
delta: "Checking the relevant event mapping.",
51+
},
52+
},
53+
{
54+
method: "item/started",
55+
params: {
56+
threadId: sessionId,
57+
turnId: "turn-1",
58+
startedAtMs: 10,
59+
item: {
60+
type: "agentMessage",
61+
id: "final-message",
62+
text: "",
63+
phase: "final_answer",
64+
memoryCitation: null,
65+
},
66+
},
67+
},
68+
{
69+
method: "item/agentMessage/delta",
70+
params: {
71+
threadId: sessionId,
72+
turnId: "turn-1",
73+
itemId: "final-message",
74+
delta: "Yes, here is the answer.",
75+
},
76+
},
77+
];
78+
79+
await setupPromptAndSendNotifications(mockFixture, sessionId, sessionState, notifications);
80+
81+
await expect(mockFixture.getAcpConnectionDump([])).toMatchFileSnapshot(
82+
"data/agent-message-phases.json"
83+
);
84+
});
85+
});

‎src/__tests__/CodexACPAgent/data/agent-message-phases.json‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
{
2+
"method": "sessionUpdate",
3+
"args": [
4+
{
5+
"sessionId": "test-session-id",
6+
"update": {
7+
"sessionUpdate": "agent_message_chunk",
8+
"messageId": "commentary-message",
9+
"content": {
10+
"type": "text",
11+
"text": "Checking the relevant event mapping."
12+
},
13+
"_meta": {
14+
"codex": {
15+
"phase": "commentary"
16+
}
17+
}
18+
}
19+
}
20+
]
21+
}
22+
{
23+
"method": "sessionUpdate",
24+
"args": [
25+
{
26+
"sessionId": "test-session-id",
27+
"update": {
28+
"sessionUpdate": "agent_message_chunk",
29+
"messageId": "final-message",
30+
"content": {
31+
"type": "text",
32+
"text": "Yes, here is the answer."
33+
},
34+
"_meta": {
35+
"codex": {
36+
"phase": "final_answer"
37+
}
38+
}
39+
}
40+
}
41+
]
42+
}

‎src/__tests__/CodexACPAgent/response-item-history-fallback.test.ts‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,26 @@ describe("ResponseItemHistoryFallback", () => {
5555
expect(thoughtTexts(updates)).toEqual(["Need to inspect the directory."]);
5656
});
5757

58+
it("preserves assistant message phase metadata from response items", () => {
59+
const updates = parseResponseItemHistoryFallback(jsonl([
60+
{
61+
type: "response_item",
62+
payload: {
63+
type: "message",
64+
role: "assistant",
65+
content: [{ type: "output_text", text: "Final answer text." }],
66+
phase: "final_answer",
67+
},
68+
},
69+
functionCall("call-missing", "ls"),
70+
functionCallOutput("call-missing", "Chunk ID: missing\nProcess exited with code 0\nOutput:\nREADME.md\n"),
71+
]), "terminal_output");
72+
73+
expect(agentMessageMetas(updates)).toEqual([
74+
{ codex: { phase: "final_answer" } },
75+
]);
76+
});
77+
5878
it("marks exec command outputs without exit footers failed when they report command errors", () => {
5979
const updates = parseResponseItemHistoryFallback(jsonl([
6080
functionCall("call-read-failed", "cat missing.txt"),
@@ -133,3 +153,11 @@ function thoughtTexts(updates: UpdateSessionEvent[] | null): string[] {
133153
))
134154
.flatMap((update) => update.content.type === "text" ? [update.content.text] : []);
135155
}
156+
157+
function agentMessageMetas(updates: UpdateSessionEvent[] | null): unknown[] {
158+
return (updates ?? [])
159+
.filter((update): update is Extract<UpdateSessionEvent, { sessionUpdate: "agent_message_chunk" }> => (
160+
update.sessionUpdate === "agent_message_chunk"
161+
))
162+
.map((update) => update._meta);
163+
}

0 commit comments

Comments
 (0)