Skip to content

Commit 2a11323

Browse files
DragonnZhangDragonnZhang
committed
fix(vscode): address ACP notification review feedback
1 parent c157d58 commit 2a11323

10 files changed

Lines changed: 313 additions & 56 deletions

File tree

‎packages/cli/src/acp-integration/session/Session.test.ts‎

Lines changed: 197 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import { SettingScope } from '../../config/settings.js';
1717
import type {
1818
AgentSideConnection,
1919
PromptRequest,
20+
SessionNotification,
2021
} from '@agentclientprotocol/sdk';
2122
import type { LoadedSettings } from '../../config/settings.js';
2223
import * as nonInteractiveCliCommands from '../../nonInteractiveCliCommands.js';
@@ -983,6 +984,202 @@ describe('Session', () => {
983984
});
984985
});
985986

987+
it('aborts an in-flight background notification before accepting a user prompt', async () => {
988+
const noopCompression = {
989+
originalTokenCount: 0,
990+
newTokenCount: 0,
991+
compressionStatus: core.CompressionStatus.NOOP,
992+
};
993+
let notificationSignal: AbortSignal | undefined;
994+
mockGeminiClient.tryCompressChat = vi
995+
.fn()
996+
.mockResolvedValueOnce(noopCompression)
997+
.mockImplementationOnce(
998+
async (_promptId: string, _force: boolean, signal: AbortSignal) => {
999+
notificationSignal = signal;
1000+
await new Promise<void>((resolve) => {
1001+
signal.addEventListener('abort', () => resolve(), {
1002+
once: true,
1003+
});
1004+
});
1005+
return noopCompression;
1006+
},
1007+
)
1008+
.mockResolvedValue(noopCompression);
1009+
mockChat.sendMessageStream = vi
1010+
.fn()
1011+
.mockResolvedValue(createEmptyStream());
1012+
1013+
await session.prompt({
1014+
sessionId: 'test-session-id',
1015+
prompt: [{ type: 'text', text: 'start background work' }],
1016+
});
1017+
1018+
const callback = mockBackgroundTaskRegistry.setNotificationCallback.mock
1019+
.calls[0][0] as (
1020+
displayText: string,
1021+
modelText: string,
1022+
meta: { agentId: string; status: string; toolUseId?: string },
1023+
) => void;
1024+
1025+
callback('done', '<task-notification />', {
1026+
agentId: 'agent-1',
1027+
status: 'completed',
1028+
});
1029+
1030+
await vi.waitFor(() => {
1031+
expect(notificationSignal).toBeDefined();
1032+
});
1033+
1034+
await expect(
1035+
session.prompt({
1036+
sessionId: 'test-session-id',
1037+
prompt: [{ type: 'text', text: 'interrupt notification' }],
1038+
}),
1039+
).resolves.toEqual({ stopReason: 'end_turn' });
1040+
1041+
expect(notificationSignal?.aborted).toBe(true);
1042+
});
1043+
1044+
it('drops oldest background notifications when the queue reaches its cap', () => {
1045+
(
1046+
session as unknown as {
1047+
pendingPrompt: AbortController | null;
1048+
}
1049+
).pendingPrompt = new AbortController();
1050+
1051+
const callback = mockBackgroundTaskRegistry.setNotificationCallback.mock
1052+
.calls[0][0] as (
1053+
displayText: string,
1054+
modelText: string,
1055+
meta: { agentId: string; status: string; toolUseId?: string },
1056+
) => void;
1057+
1058+
for (let index = 0; index < 25; index++) {
1059+
callback(
1060+
`done ${index}`,
1061+
`<task-notification>${index}</task-notification>`,
1062+
{
1063+
agentId: `agent-${index}`,
1064+
status: 'completed',
1065+
},
1066+
);
1067+
}
1068+
1069+
const queued = (
1070+
session as unknown as {
1071+
notificationQueue: Array<{ taskId: string }>;
1072+
}
1073+
).notificationQueue;
1074+
expect(queued).toHaveLength(20);
1075+
expect(queued[0]?.taskId).toBe('agent-5');
1076+
expect(queued.at(-1)?.taskId).toBe('agent-24');
1077+
});
1078+
1079+
it('emits end_turn even when notification error display fails', async () => {
1080+
mockChat.sendMessageStream = vi
1081+
.fn()
1082+
.mockResolvedValueOnce(createEmptyStream())
1083+
.mockRejectedValueOnce(new Error('notification blew up'));
1084+
mockClient.sessionUpdate = vi.fn().mockImplementation(async (params) => {
1085+
const text = (
1086+
(params as SessionNotification).update as {
1087+
content?: { text?: string };
1088+
}
1089+
)?.content?.text;
1090+
if (text?.includes('[notification error]')) {
1091+
throw new Error('display failed');
1092+
}
1093+
});
1094+
1095+
await session.prompt({
1096+
sessionId: 'test-session-id',
1097+
prompt: [{ type: 'text', text: 'start background work' }],
1098+
});
1099+
1100+
const callback = mockBackgroundTaskRegistry.setNotificationCallback.mock
1101+
.calls[0][0] as (
1102+
displayText: string,
1103+
modelText: string,
1104+
meta: { agentId: string; status: string; toolUseId?: string },
1105+
) => void;
1106+
1107+
callback('done', '<task-notification />', {
1108+
agentId: 'agent-1',
1109+
status: 'completed',
1110+
});
1111+
1112+
await vi.waitFor(() => {
1113+
expect(mockClient.sessionUpdate).toHaveBeenCalledWith({
1114+
sessionId: 'test-session-id',
1115+
update: expect.objectContaining({
1116+
content: expect.objectContaining({
1117+
text: expect.stringContaining('[notification error]'),
1118+
}),
1119+
}),
1120+
});
1121+
expect(mockClient.extNotification).toHaveBeenCalledWith(
1122+
'_qwencode/end_turn',
1123+
{
1124+
sessionId: 'test-session-id',
1125+
reason: 'end_turn',
1126+
source: 'background_notification',
1127+
},
1128+
);
1129+
});
1130+
});
1131+
1132+
it('flushes notification rewrite metadata even without usage metadata', async () => {
1133+
const flushTurn = vi.fn().mockResolvedValue(undefined);
1134+
const waitForPendingRewrites = vi.fn().mockResolvedValue(undefined);
1135+
const interceptUpdate = vi.fn().mockResolvedValue(undefined);
1136+
session.messageRewriter = {
1137+
interceptUpdate,
1138+
flushTurn,
1139+
waitForPendingRewrites,
1140+
} as unknown as Session['messageRewriter'];
1141+
mockChat.sendMessageStream = vi
1142+
.fn()
1143+
.mockResolvedValueOnce(createEmptyStream())
1144+
.mockResolvedValueOnce(
1145+
createStreamWithChunks([
1146+
{
1147+
type: core.StreamEventType.CHUNK,
1148+
value: {
1149+
candidates: [
1150+
{
1151+
content: {
1152+
parts: [{ text: 'notification response' }],
1153+
},
1154+
},
1155+
],
1156+
},
1157+
},
1158+
]),
1159+
);
1160+
1161+
await session.prompt({
1162+
sessionId: 'test-session-id',
1163+
prompt: [{ type: 'text', text: 'start background work' }],
1164+
});
1165+
1166+
const callback = mockBackgroundTaskRegistry.setNotificationCallback.mock
1167+
.calls[0][0] as (
1168+
displayText: string,
1169+
modelText: string,
1170+
meta: { agentId: string; status: string; toolUseId?: string },
1171+
) => void;
1172+
1173+
callback('done', '<task-notification />', {
1174+
agentId: 'agent-1',
1175+
status: 'completed',
1176+
});
1177+
1178+
await vi.waitFor(() => {
1179+
expect(flushTurn).toHaveBeenCalled();
1180+
});
1181+
});
1182+
9861183
it('does not enqueue running monitor notifications for model follow-up', async () => {
9871184
mockChat.sendMessageStream = vi
9881185
.fn()

‎packages/cli/src/acp-integration/session/Session.ts‎

Lines changed: 31 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,8 @@ interface BackgroundNotificationQueueItem {
138138
toolUseId?: string;
139139
}
140140

141+
const MAX_NOTIFICATION_QUEUE = 20;
142+
141143
export function computeInitialTurnFromHistory(
142144
records: ChatRecord[],
143145
sessionId: string,
@@ -533,6 +535,7 @@ export class Session implements SessionContext {
533535
this.notificationAbortController = null;
534536
}
535537
this.notificationQueue = [];
538+
this.notificationProcessing = false;
536539

537540
// Stop scheduler and emit exit summary
538541
const scheduler = this.config.isCronEnabled()
@@ -579,9 +582,15 @@ export class Session implements SessionContext {
579582
}
580583
}
581584

582-
// A background notification turn mutates the same chat history as a
583-
// user prompt. Wait for an already-running drain before accepting the
584-
// new prompt so ACP never sends two turns against the same chat at once.
585+
// A background notification turn mutates the same chat history as a user
586+
// prompt. Abort it before awaiting the drain so user input is not blocked
587+
// behind notification tool calls.
588+
if (this.notificationAbortController) {
589+
this.notificationAbortController.abort();
590+
this.notificationAbortController = null;
591+
this.notificationQueue = [];
592+
this.notificationProcessing = false;
593+
}
585594
if (this.notificationCompletion) {
586595
try {
587596
await this.notificationCompletion;
@@ -1628,6 +1637,9 @@ export class Session implements SessionContext {
16281637
}
16291638

16301639
#enqueueBackgroundNotification(item: BackgroundNotificationQueueItem): void {
1640+
while (this.notificationQueue.length >= MAX_NOTIFICATION_QUEUE) {
1641+
this.notificationQueue.shift();
1642+
}
16311643
this.notificationQueue.push(item);
16321644
void this.#drainNotificationQueue();
16331645
}
@@ -1782,11 +1794,12 @@ export class Session implements SessionContext {
17821794
);
17831795
}
17841796

1797+
if (this.messageRewriter) {
1798+
await this.messageRewriter.flushTurn(ac.signal);
1799+
}
1800+
17851801
if (usageMetadata) {
17861802
this.#recordPromptTokenCount(usageMetadata);
1787-
if (this.messageRewriter) {
1788-
this.messageRewriter.flushTurn(ac.signal);
1789-
}
17901803
const durationMs = Date.now() - streamStartTime;
17911804
await this.messageEmitter.emitUsageMetadata(
17921805
usageMetadata,
@@ -1817,10 +1830,18 @@ export class Session implements SessionContext {
18171830
}
18181831
debugLogger.error('Error processing background notification:', error);
18191832
const msg = error instanceof Error ? error.message : String(error);
1820-
await this.messageEmitter.emitAgentMessage(
1821-
`[notification error] ${msg}`,
1822-
);
1823-
await this.#emitBackgroundNotificationEndTurn('end_turn');
1833+
try {
1834+
await this.messageEmitter.emitAgentMessage(
1835+
`[notification error] ${msg}`,
1836+
);
1837+
} catch (emitError) {
1838+
debugLogger.error(
1839+
'Failed to emit background notification error:',
1840+
emitError,
1841+
);
1842+
} finally {
1843+
await this.#emitBackgroundNotificationEndTurn('end_turn');
1844+
}
18241845
} finally {
18251846
if (this.notificationAbortController === ac) {
18261847
this.notificationAbortController = null;

‎packages/cli/src/acp-integration/session/rewrite/MessageRewriteMiddleware.test.ts‎

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -203,7 +203,7 @@ describe('MessageRewriteMiddleware', () => {
203203
expect(meta['turnIndex']).toBe(1);
204204
});
205205

206-
it('preserves background discrete metadata on rewritten messages', async () => {
206+
it('omits background discrete metadata on rewritten messages', async () => {
207207
const { middleware, mockSendUpdate } = createMiddleware('message');
208208

209209
await middleware.interceptUpdate({
@@ -218,6 +218,7 @@ describe('MessageRewriteMiddleware', () => {
218218
kind: 'monitor',
219219
toolUseId: 'tool-1',
220220
},
221+
customTraceId: 'trace-1',
221222
},
222223
} as unknown as SessionUpdate);
223224

@@ -234,14 +235,7 @@ describe('MessageRewriteMiddleware', () => {
234235
);
235236
expect(rewriteCall).toBeDefined();
236237
expect((rewriteCall![0] as Record<string, unknown>)['_meta']).toEqual({
237-
source: 'background_notification_response',
238-
qwenDiscreteMessage: true,
239-
backgroundTask: {
240-
taskId: 'monitor-1',
241-
status: 'completed',
242-
kind: 'monitor',
243-
toolUseId: 'tool-1',
244-
},
238+
customTraceId: 'trace-1',
245239
rewritten: true,
246240
turnIndex: 1,
247241
});

‎packages/cli/src/acp-integration/session/rewrite/MessageRewriteMiddleware.ts‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,11 @@ const debugLogger = createDebugLogger('MESSAGE_REWRITE');
2828
* 4. Rewritten text is emitted as agent_message_chunk with _meta.rewritten=true
2929
*/
3030
const DEFAULT_REWRITE_TIMEOUT_MS = 30_000;
31+
const REWRITE_META_EXCLUDED_KEYS = new Set([
32+
'backgroundTask',
33+
'qwenDiscreteMessage',
34+
'source',
35+
]);
3136

3237
export class MessageRewriteMiddleware {
3338
private readonly turnBuffer: TurnBuffer;
@@ -167,9 +172,17 @@ export class MessageRewriteMiddleware {
167172
return;
168173
}
169174

175+
const safeMeta = Object.fromEntries(
176+
Object.entries(meta as Record<string, unknown>).filter(
177+
([key]) => !REWRITE_META_EXCLUDED_KEYS.has(key),
178+
),
179+
);
180+
181+
if (Object.keys(safeMeta).length === 0) return;
182+
170183
this.turnMeta = {
171184
...this.turnMeta,
172-
...(meta as Record<string, unknown>),
185+
...safeMeta,
173186
};
174187
}
175188

‎packages/core/src/services/backgroundShellRegistry.test.ts‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,22 @@ describe('BackgroundShellRegistry', () => {
8080
);
8181
});
8282

83+
it('strips display control characters from notification XML summaries', () => {
84+
const reg = new BackgroundShellRegistry();
85+
const callback = vi.fn();
86+
reg.setNotificationCallback(callback);
87+
reg.register(
88+
makeEntry({ shellId: 'a', command: 'npm \x1b[31mtest\x00' }),
89+
);
90+
91+
reg.complete('a', 0, 2000);
92+
93+
const modelText = callback.mock.calls[0][1] as string;
94+
expect(modelText).toContain('<summary>Shell "npm [31mtest" completed.');
95+
expect(modelText).not.toContain('\x1b');
96+
expect(modelText).not.toContain('\x00');
97+
});
98+
8399
it('is a no-op when entry is not running', () => {
84100
const reg = new BackgroundShellRegistry();
85101
reg.register(makeEntry({ shellId: 'a' }));

‎packages/core/src/services/backgroundShellRegistry.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -304,7 +304,7 @@ export class BackgroundShellRegistry {
304304
`<task-id>${escapeXml(entry.shellId)}</task-id>`,
305305
'<kind>shell</kind>',
306306
`<status>${escapeXml(entry.status)}</status>`,
307-
`<summary>Shell "${escapeXml(entry.command)}" ${statusText}.</summary>`,
307+
`<summary>Shell "${escapeXml(this.stripDisplayControlChars(entry.command))}" ${statusText}.</summary>`,
308308
`<output-file>${escapeXml(entry.outputFile)}</output-file>`,
309309
];
310310
if (entry.exitCode !== undefined) {

0 commit comments

Comments
 (0)