fix(cli): prevent ACP stdout pollution from SessionEnd hooks (#26125)

This commit is contained in:
Coco Sheng
2026-04-28 13:29:49 -04:00
committed by GitHub
parent 7a3f7c383e
commit 4b8d5e7624
2 changed files with 132 additions and 2 deletions
+7 -1
View File
@@ -385,7 +385,6 @@ export async function main() {
},
});
consolePatcher.patch();
registerCleanup(consolePatcher.cleanup);
dns.setDefaultResultOrder(
validateDnsResolutionOrder(settings.merged.advanced.dnsResolutionOrder),
@@ -411,6 +410,7 @@ export async function main() {
const partialConfig = await loadCliConfig(settings.merged, sessionId, argv, {
projectHooks: settings.workspace.settings.hooks,
});
adminControlsListner.setConfig(partialConfig);
// Refresh auth to fetch remote admin settings from CCPA and before entering
@@ -568,6 +568,12 @@ export async function main() {
await config.getHookSystem()?.fireSessionEndEvent(SessionEndReason.Exit);
});
// Register ConsolePatcher cleanup last to ensure logs from shutdown hooks
// are correctly redirected to stderr (especially for non-interactive JSON output).
if (!config.getAcpMode()) {
registerCleanup(consolePatcher.cleanup);
}
// Launch cleanup expired sessions as a background task
cleanupExpiredSessions(config, settings.merged).catch((e) => {
debugLogger.error('Failed to cleanup expired sessions:', e);
+125 -1
View File
@@ -68,6 +68,13 @@ vi.mock('./config/settings.js', async (importOriginal) => {
};
});
vi.mock('./ui/utils/ConsolePatcher.js', () => ({
ConsolePatcher: vi.fn().mockImplementation(() => ({
patch: vi.fn(),
cleanup: vi.fn(),
})),
}));
vi.mock('./config/config.js', () => ({
loadCliConfig: vi.fn().mockResolvedValue({
getSandbox: vi.fn(() => false),
@@ -150,6 +157,10 @@ vi.mock('./utils/cleanup.js', async (importOriginal) => {
};
});
vi.mock('./acp/acpClient.js', () => ({
runAcpClient: vi.fn().mockResolvedValue(undefined),
}));
vi.mock('./zed-integration/zedIntegration.js', () => ({
runZedIntegration: vi.fn().mockResolvedValue(undefined),
}));
@@ -296,6 +307,120 @@ describe('gemini.tsx main function cleanup', () => {
);
});
it('should not register ConsolePatcher cleanup in ACP mode', async () => {
const { registerCleanup } = await import('./utils/cleanup.js');
const { ConsolePatcher } = await import('./ui/utils/ConsolePatcher.js');
const { loadCliConfig, parseArguments } = await import(
'./config/config.js'
);
const { loadSettings } = await import('./config/settings.js');
vi.mocked(parseArguments).mockResolvedValue({
acp: true,
startupMessages: [],
// eslint-disable-next-line @typescript-eslint/no-explicit-any
} as any);
vi.mocked(loadSettings).mockReturnValue({
merged: {
tools: { allowed: [], exclude: [] },
advanced: { dnsResolutionOrder: 'ipv4first' },
security: { auth: { selectedType: 'google' } },
ui: { theme: 'default' },
},
workspace: { settings: {} },
errors: [],
subscribe: vi.fn(),
getSnapshot: vi.fn(),
setValue: vi.fn(),
// eslint-disable-next-line @typescript-eslint/no-explicit-any
} as any);
vi.mocked(loadCliConfig).mockResolvedValue(
buildMockConfig({
getAcpMode: () => true,
}),
);
let capturedCleanup: () => void;
vi.mocked(ConsolePatcher).mockImplementation(() => {
const instance = {
patch: vi.fn(),
cleanup: vi.fn(),
};
capturedCleanup = instance.cleanup;
// eslint-disable-next-line @typescript-eslint/no-explicit-any
return instance as any;
});
await main();
const registeredFunctions = vi
.mocked(registerCleanup)
.mock.calls.map((call) => call[0]);
expect(registeredFunctions).not.toContain(capturedCleanup!);
});
it('should register ConsolePatcher cleanup in non-ACP mode', async () => {
const { registerCleanup } = await import('./utils/cleanup.js');
const { ConsolePatcher } = await import('./ui/utils/ConsolePatcher.js');
const { loadCliConfig, parseArguments } = await import(
'./config/config.js'
);
const { loadSettings } = await import('./config/settings.js');
vi.mocked(parseArguments).mockResolvedValue({
acp: false,
query: 'test',
startupMessages: [],
// eslint-disable-next-line @typescript-eslint/no-explicit-any
} as any);
vi.mocked(loadSettings).mockReturnValue({
merged: {
tools: { allowed: [], exclude: [] },
advanced: { dnsResolutionOrder: 'ipv4first' },
security: { auth: { selectedType: 'google' } },
ui: { theme: 'default' },
},
workspace: { settings: {} },
errors: [],
subscribe: vi.fn(),
getSnapshot: vi.fn(),
setValue: vi.fn(),
// eslint-disable-next-line @typescript-eslint/no-explicit-any
} as any);
vi.mocked(loadCliConfig).mockResolvedValue(
buildMockConfig({
getAcpMode: () => false,
getQuestion: () => 'test',
}),
);
let capturedCleanup: () => void;
vi.mocked(ConsolePatcher).mockImplementation(() => {
const instance = {
patch: vi.fn(),
cleanup: vi.fn(),
};
capturedCleanup = instance.cleanup;
// eslint-disable-next-line @typescript-eslint/no-explicit-any
return instance as any;
});
try {
await main();
} catch {
// Ignore errors from incomplete mocks in full main() execution
}
const registeredFunctions = vi
.mocked(registerCleanup)
.mock.calls.map((call) => call[0]);
expect(registeredFunctions).toContain(capturedCleanup!);
});
function buildMockConfig(overrides: Partial<Config> = {}): Config {
return {
isInteractive: vi.fn(() => false),
@@ -319,7 +444,6 @@ describe('gemini.tsx main function cleanup', () => {
getListExtensions: vi.fn(() => false),
getListSessions: vi.fn(() => false),
getDeleteSession: vi.fn(() => undefined),
getToolRegistry: vi.fn(),
getExtensions: vi.fn(() => []),
getModel: vi.fn(() => 'gemini-pro'),
getEmbeddingModel: vi.fn(() => 'embedding-001'),