Skip to content

Commit 0bde143

Browse files
benbrandtbenbrandt
committed
feat: Add message IDs to text session chunks
1 parent 5506fba commit 0bde143

12 files changed

Lines changed: 137 additions & 134 deletions

‎src/CodexAcpServer.ts‎

Lines changed: 39 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,11 @@ import {
6666
import packageJson from "../package.json";
6767
import {isJetBrains2026_1Client} from "./JBUtils";
6868
import {resolveTerminalOutputMode, type TerminalOutputMode} from "./TerminalOutputMode";
69+
import {
70+
createAgentTextMessageChunk,
71+
createAgentTextThoughtChunk,
72+
createUserMessageChunk,
73+
} from "./ContentChunks";
6974

7075
export interface ThreadGoalSnapshot {
7176
objective: string;
@@ -966,6 +971,7 @@ export class CodexAcpServer {
966971
case "agentMessage":
967972
return [{
968973
sessionUpdate: "agent_message_chunk",
974+
messageId: item.id,
969975
content: { type: "text", text: item.text },
970976
}];
971977
case "reasoning":
@@ -1005,24 +1011,20 @@ export class CodexAcpServer {
10051011

10061012
private createUserMessageUpdates(item: ThreadItem & { type: "userMessage" }): UpdateSessionEvent[] {
10071013
const updates: UpdateSessionEvent[] = [];
1014+
const messageId = item.id;
10081015
for (const input of item.content) {
10091016
const blocks = this.userInputToContentBlocks(input);
10101017
for (const block of blocks) {
1011-
updates.push({
1012-
sessionUpdate: "user_message_chunk",
1013-
content: block,
1014-
});
1018+
updates.push(createUserMessageChunk(block, messageId));
10151019
}
10161020
}
10171021
return updates;
10181022
}
10191023

10201024
private createReasoningUpdates(item: ThreadItem & { type: "reasoning" }): UpdateSessionEvent[] {
10211025
const parts = item.summary.length > 0 ? item.summary : item.content;
1022-
return parts.map((text) => ({
1023-
sessionUpdate: "agent_thought_chunk",
1024-
content: { type: "text", text: text },
1025-
}));
1026+
const messageId = item.id;
1027+
return parts.map((text) => createAgentTextThoughtChunk(text, messageId));
10261028
}
10271029

10281030
private createWebSearchUpdate(
@@ -1589,13 +1591,7 @@ export class CodexAcpServer {
15891591
}
15901592
await this.connection.notify(acp.methods.client.session.update, {
15911593
sessionId,
1592-
update: {
1593-
sessionUpdate: "agent_message_chunk",
1594-
content: {
1595-
type: "text",
1596-
text: "*Conversation interrupted*"
1597-
}
1598-
}
1594+
update: createAgentTextMessageChunk("*Conversation interrupted*"),
15991595
});
16001596
}
16011597

@@ -1674,26 +1670,33 @@ function mergeHistoryUpdates(
16741670
merged.push(update);
16751671
};
16761672

1677-
const flushFallbackThrough = (targetKey: string): boolean => {
1673+
const flushFallbackBeforeMatchingDuplicate = (targetUpdate: UpdateSessionEvent): void => {
1674+
const targetKey = historyUpdateKey(targetUpdate);
1675+
const targetContentKey = historyUpdateContentKey(targetUpdate);
1676+
if (!targetKey && !targetContentKey) {
1677+
return;
1678+
}
1679+
16781680
const matchIndex = responseItemFallbackUpdates.findIndex((update, index) => (
1679-
index >= fallbackIndex && historyUpdateKey(update) === targetKey
1681+
index >= fallbackIndex
1682+
&& (
1683+
(targetKey !== null && historyUpdateKey(update) === targetKey)
1684+
|| (targetContentKey !== null && historyUpdateContentKey(update) === targetContentKey)
1685+
)
16801686
));
16811687
if (matchIndex === -1) {
1682-
return false;
1688+
return;
16831689
}
16841690

1685-
while (fallbackIndex <= matchIndex) {
1691+
while (fallbackIndex < matchIndex) {
16861692
pushUpdate(responseItemFallbackUpdates[fallbackIndex]!);
16871693
fallbackIndex += 1;
16881694
}
1689-
return true;
1695+
fallbackIndex += 1;
16901696
};
16911697

16921698
for (const update of threadUpdates) {
1693-
const key = historyUpdateKey(update);
1694-
if (key && flushFallbackThrough(key)) {
1695-
continue;
1696-
}
1699+
flushFallbackBeforeMatchingDuplicate(update);
16971700
pushUpdate(update);
16981701
}
16991702

@@ -1710,7 +1713,7 @@ function historyUpdateKey(update: UpdateSessionEvent): string | null {
17101713
case "user_message_chunk":
17111714
case "agent_message_chunk":
17121715
case "agent_thought_chunk":
1713-
return `${update.sessionUpdate}:${JSON.stringify(update.content)}`;
1716+
return `${update.sessionUpdate}:${update.messageId ?? ""}:${JSON.stringify(update.content)}`;
17141717
case "tool_call":
17151718
return `tool_call:${update.toolCallId}:start`;
17161719
case "tool_call_update":
@@ -1720,6 +1723,17 @@ function historyUpdateKey(update: UpdateSessionEvent): string | null {
17201723
}
17211724
}
17221725

1726+
function historyUpdateContentKey(update: UpdateSessionEvent): string | null {
1727+
switch (update.sessionUpdate) {
1728+
case "user_message_chunk":
1729+
case "agent_message_chunk":
1730+
case "agent_thought_chunk":
1731+
return `${update.sessionUpdate}:${JSON.stringify(update.content)}`;
1732+
default:
1733+
return historyUpdateKey(update);
1734+
}
1735+
}
1736+
17231737
function getRequestedMcpServerNames(mcpServers: Array<acp.McpServer>): Array<string> {
17241738
return Array.from(new Set(mcpServers.map(server => sanitizeMcpServerName(server.name))));
17251739
}

‎src/CodexCommands.ts‎

Lines changed: 8 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import type {SessionState} from "./CodexAcpServer";
77
import type {RateLimitsMap} from "./RateLimitsMap";
88
import type {TokenCount} from "./TokenCount";
99
import {logger} from "./Logger";
10+
import {createAgentTextMessageChunk} from "./ContentChunks";
1011

1112
type ParsedSlashCommand = {
1213
name: string;
@@ -210,20 +211,14 @@ export class CodexCommands {
210211
case "status": {
211212
const session = new ACPSessionConnection(this.connection, sessionId);
212213
const message = this.buildStatusMessage(sessionState);
213-
await session.update({
214-
sessionUpdate: "agent_message_chunk",
215-
content: { type: "text", text: message }
216-
});
214+
await session.update(createAgentTextMessageChunk(message));
217215
return { handled: true };
218216
}
219217
case "logout": {
220218
await this.runWithProcessCheck(() => this.codexAcpClient.logout());
221219
await this.onLogout();
222220
const session = new ACPSessionConnection(this.connection, sessionId);
223-
await session.update({
224-
sessionUpdate: "agent_message_chunk",
225-
content: { type: "text", text: "Logged out from Codex account." }
226-
});
221+
await session.update(createAgentTextMessageChunk("Logged out from Codex account."));
227222
return { handled: true };
228223
}
229224
case "skills": {
@@ -237,10 +232,7 @@ export class CodexCommands {
237232
? ["Available skills:", ...lines].join("\n")
238233
: "No skills configured.";
239234
const session = new ACPSessionConnection(this.connection, sessionId);
240-
await session.update({
241-
sessionUpdate: "agent_message_chunk",
242-
content: { type: "text", text }
243-
});
235+
await session.update(createAgentTextMessageChunk(text));
244236
return { handled: true };
245237
}
246238
case "mcp": {
@@ -258,10 +250,7 @@ export class CodexCommands {
258250
? ["Configured MCP servers:", ...lines].join("\n")
259251
: "No MCP servers configured.";
260252
const session = new ACPSessionConnection(this.connection, sessionId);
261-
await session.update({
262-
sessionUpdate: "agent_message_chunk",
263-
content: { type: "text", text }
264-
});
253+
await session.update(createAgentTextMessageChunk(text));
265254
return { handled: true };
266255
}
267256
default:
@@ -316,13 +305,7 @@ export class CodexCommands {
316305

317306
if (argument.length > 4000) {
318307
const session = new ACPSessionConnection(this.connection, sessionId);
319-
await session.update({
320-
sessionUpdate: "agent_message_chunk",
321-
content: {
322-
type: "text",
323-
text: 'Command "/goal" requires goal text of at most 4000 characters.'
324-
}
325-
});
308+
await session.update(createAgentTextMessageChunk('Command "/goal" requires goal text of at most 4000 characters.'));
326309
return { handled: true };
327310
}
328311

@@ -371,13 +354,7 @@ export class CodexCommands {
371354

372355
private async sendCommandUsageMessage(name: string, inputHint: string, sessionId: string): Promise<void> {
373356
const session = new ACPSessionConnection(this.connection, sessionId);
374-
await session.update({
375-
sessionUpdate: "agent_message_chunk",
376-
content: {
377-
type: "text",
378-
text: `Command "/${name}" requires ${inputHint}.`
379-
}
380-
});
357+
await session.update(createAgentTextMessageChunk(`Command "/${name}" requires ${inputHint}.`));
381358
}
382359

383360
private async sendUnknownCommandMessage(name: string, sessionId: string): Promise<void> {
@@ -390,10 +367,7 @@ export class CodexCommands {
390367
text.push(...lines);
391368
}
392369
const session = new ACPSessionConnection(this.connection, sessionId);
393-
await session.update({
394-
sessionUpdate: "agent_message_chunk",
395-
content: { type: "text", text: text.join("\n") }
396-
});
370+
await session.update(createAgentTextMessageChunk(text.join("\n")));
397371
}
398372

399373
private buildStatusMessage(sessionState: SessionState): string {

‎src/CodexEventHandler.ts‎

Lines changed: 18 additions & 74 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,10 @@ import {
5555
} from "./CodexToolCallMapper";
5656
import { stripShellPrefix } from "./CommandUtils";
5757
import {createTerminalOutputMeta, type TerminalOutputMode} from "./TerminalOutputMode";
58+
import {
59+
createAgentTextMessageChunk,
60+
createAgentTextThoughtChunk,
61+
} from "./ContentChunks";
5862

5963
export { stripShellPrefix };
6064

@@ -226,44 +230,20 @@ export class CodexEventHandler {
226230
}
227231

228232
private async createTextEvent(event: AgentMessageDeltaNotification): Promise<UpdateSessionEvent> {
229-
return {
230-
sessionUpdate: "agent_message_chunk",
231-
content: {
232-
type: "text",
233-
text: event.delta
234-
}
235-
}
233+
return createAgentTextMessageChunk(event.delta, event.itemId);
236234
}
237235

238236
private async createConfigWarningEvent(event: ConfigWarningNotification): Promise<UpdateSessionEvent> {
239237
const detailsText = event.details ? `\n\n${event.details}` : "";
240-
return {
241-
sessionUpdate: "agent_message_chunk",
242-
content: {
243-
type: "text",
244-
text: `Config warning: ${event.summary}${detailsText}\n\n`
245-
}
246-
}
238+
return createAgentTextMessageChunk(`Config warning: ${event.summary}${detailsText}\n\n`);
247239
}
248240

249241
private createWarningEvent(event: WarningNotification): UpdateSessionEvent {
250-
return {
251-
sessionUpdate: "agent_message_chunk",
252-
content: {
253-
type: "text",
254-
text: `Warning: ${event.message}\n\n`
255-
}
256-
};
242+
return createAgentTextMessageChunk(`Warning: ${event.message}\n\n`);
257243
}
258244

259245
private createModelReroutedEvent(event: ModelReroutedNotification): UpdateSessionEvent {
260-
return {
261-
sessionUpdate: "agent_thought_chunk",
262-
content: {
263-
type: "text",
264-
text: `Model rerouted from ${event.fromModel} to ${event.toModel} (${event.reason}).\n\n`
265-
}
266-
};
246+
return createAgentTextThoughtChunk(`Model rerouted from ${event.fromModel} to ${event.toModel} (${event.reason}).\n\n`);
267247
}
268248

269249
private createThreadGoalUpdatedEvent(event: ThreadGoalUpdatedNotification): UpdateSessionEvent | null {
@@ -278,13 +258,7 @@ export class CodexEventHandler {
278258
const text = objective.includes("\n")
279259
? `Goal updated (${status}):\n${objective}`
280260
: `Goal updated (${status}): ${objective}`;
281-
return {
282-
sessionUpdate: "agent_message_chunk",
283-
content: {
284-
type: "text",
285-
text: `\n\n${text}\n\n`,
286-
},
287-
};
261+
return createAgentTextMessageChunk(`\n\n${text}\n\n`);
288262
}
289263

290264
private formatThreadGoalStatus(status: ThreadGoalUpdatedNotification["goal"]["status"]): string {
@@ -310,13 +284,7 @@ export class CodexEventHandler {
310284
}
311285
this.sessionState.currentGoal = null;
312286

313-
return {
314-
sessionUpdate: "agent_message_chunk",
315-
content: {
316-
type: "text",
317-
text: "\n\nGoal cleared.\n\n",
318-
},
319-
};
287+
return createAgentTextMessageChunk("\n\nGoal cleared.\n\n");
320288
}
321289

322290
private createThreadGoalSnapshot(event: ThreadGoalUpdatedNotification): ThreadGoalSnapshot {
@@ -342,22 +310,16 @@ export class CodexEventHandler {
342310
event: ReasoningSummaryTextDeltaNotification | ReasoningTextDeltaNotification
343311
): UpdateSessionEvent {
344312
this.seenReasoningDeltaItemIds.add(event.itemId);
345-
return this.createAgentThoughtEvent(event.delta);
313+
return this.createAgentThoughtEvent(event.delta, event.itemId);
346314
}
347315

348316
private createReasoningSectionBreakEvent(event: ReasoningSummaryPartAddedNotification): UpdateSessionEvent {
349317
this.seenReasoningDeltaItemIds.add(event.itemId);
350-
return this.createAgentThoughtEvent("\n\n");
318+
return this.createAgentThoughtEvent("\n\n", event.itemId);
351319
}
352320

353-
private createAgentThoughtEvent(text: string): UpdateSessionEvent {
354-
return {
355-
sessionUpdate: "agent_thought_chunk",
356-
content: {
357-
type: "text",
358-
text,
359-
}
360-
};
321+
private createAgentThoughtEvent(text: string, messageId: string): UpdateSessionEvent {
322+
return createAgentTextThoughtChunk(text, messageId);
361323
}
362324

363325
private async createItemEvent(event: ItemStartedNotification): Promise<UpdateSessionEvent | null> {
@@ -462,31 +424,19 @@ export class CodexEventHandler {
462424
if (text.length === 0) {
463425
return null;
464426
}
465-
return this.createAgentThoughtEvent(text);
427+
return this.createAgentThoughtEvent(text, item.id);
466428
}
467429

468430
private createExitedReviewModeEvent(item: ThreadItem & { type: "exitedReviewMode" }): UpdateSessionEvent | null {
469431
const text = item.review.trim();
470432
if (text.length === 0) {
471433
return null;
472434
}
473-
return {
474-
sessionUpdate: "agent_message_chunk",
475-
content: {
476-
type: "text",
477-
text,
478-
}
479-
};
435+
return createAgentTextMessageChunk(text);
480436
}
481437

482438
private createContextCompactedEvent(): UpdateSessionEvent {
483-
return {
484-
sessionUpdate: "agent_message_chunk",
485-
content: {
486-
type: "text",
487-
text: "*Context compacted to fit the model's context window.*\n\n"
488-
}
489-
};
439+
return createAgentTextMessageChunk("*Context compacted to fit the model's context window.*\n\n");
490440
}
491441

492442
private createCommandOutputDeltaEvent(event: CommandExecutionOutputDeltaNotification): UpdateSessionEvent {
@@ -629,13 +579,7 @@ export class CodexEventHandler {
629579
? RequestError.internalError(this.createTurnErrorData(params.error))
630580
: RequestError.authRequired(this.createTurnErrorData(params.error), params.error.message);
631581
}
632-
return {
633-
sessionUpdate: "agent_message_chunk",
634-
content: {
635-
type: "text",
636-
text: `${params.error.message}\n\n`
637-
}
638-
}
582+
return createAgentTextMessageChunk(`${params.error.message}\n\n`);
639583
}
640584

641585
private isAuthenticationRequiredError(error: CodexErrorInfo | null): boolean {

0 commit comments

Comments
 (0)