diff --git a/integration-tests/acp-telemetry.test.ts b/integration-tests/acp-telemetry.test.ts new file mode 100644 index 0000000000..970239de9e --- /dev/null +++ b/integration-tests/acp-telemetry.test.ts @@ -0,0 +1,116 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, beforeEach, afterEach } from 'vitest'; +import { TestRig } from './test-helper.js'; +import { spawn, ChildProcess } from 'node:child_process'; +import { join } from 'node:path'; +import { readFileSync, existsSync } from 'node:fs'; +import { Writable, Readable } from 'node:stream'; +import { env } from 'node:process'; +import * as acp from '@agentclientprotocol/sdk'; + +// Skip in sandbox mode - test spawns CLI directly which behaves differently in containers +const sandboxEnv = env['GEMINI_SANDBOX']; +const itMaybe = sandboxEnv && sandboxEnv !== 'false' ? it.skip : it; + +// Reuse existing fake responses that return a simple "Hello" response +const SIMPLE_RESPONSE_PATH = 'hooks-system.session-startup.responses'; + +class SessionUpdateCollector implements acp.Client { + updates: acp.SessionNotification[] = []; + + sessionUpdate = async (params: acp.SessionNotification) => { + this.updates.push(params); + }; + + requestPermission = async (): Promise => { + throw new Error('unexpected'); + }; +} + +describe('ACP telemetry', () => { + let rig: TestRig; + let child: ChildProcess | undefined; + + beforeEach(() => { + rig = new TestRig(); + }); + + afterEach(async () => { + child?.kill(); + child = undefined; + await rig.cleanup(); + }); + + itMaybe('should flush telemetry when connection closes', async () => { + rig.setup('acp-telemetry-flush', { + fakeResponsesPath: join(import.meta.dirname, SIMPLE_RESPONSE_PATH), + }); + + const telemetryPath = join(rig.homeDir!, 'telemetry.log'); + const bundlePath = join(import.meta.dirname, '..', 'bundle/gemini.js'); + + child = spawn( + 'node', + [ + bundlePath, + '--experimental-acp', + '--fake-responses', + join(rig.testDir!, 'fake-responses.json'), + ], + { + cwd: rig.testDir!, + stdio: ['pipe', 'pipe', 'inherit'], + env: { + ...process.env, + GEMINI_API_KEY: 'fake-key', + GEMINI_CLI_HOME: rig.homeDir!, + GEMINI_TELEMETRY_ENABLED: 'true', + GEMINI_TELEMETRY_TARGET: 'local', + GEMINI_TELEMETRY_OUTFILE: telemetryPath, + // GEMINI_DEV_TRACING not set: fake responses aren't instrumented for spans + }, + }, + ); + + const input = Writable.toWeb(child.stdin!); + const output = Readable.toWeb(child.stdout!) as ReadableStream; + const testClient = new SessionUpdateCollector(); + const stream = acp.ndJsonStream(input, output); + const connection = new acp.ClientSideConnection(() => testClient, stream); + + await connection.initialize({ + protocolVersion: acp.PROTOCOL_VERSION, + clientCapabilities: { fs: { readTextFile: false, writeTextFile: false } }, + }); + + const { sessionId } = await connection.newSession({ + cwd: rig.testDir!, + mcpServers: [], + }); + + await connection.prompt({ + sessionId, + prompt: [{ type: 'text', text: 'Say hello' }], + }); + + expect(JSON.stringify(testClient.updates)).toContain('Hello'); + + // Close stdin to trigger telemetry flush via runExitCleanup() + child.stdin!.end(); + await new Promise((resolve) => { + child!.on('close', () => resolve()); + }); + child = undefined; + + // gen_ai.output.messages is the last OTEL log emitted (after prompt response) + expect(existsSync(telemetryPath)).toBe(true); + expect(readFileSync(telemetryPath, 'utf-8')).toContain( + 'gen_ai.output.messages', + ); + }); +}); diff --git a/package-lock.json b/package-lock.json index 6fb8dad5a0..56b985c5e7 100644 --- a/package-lock.json +++ b/package-lock.json @@ -19,6 +19,7 @@ "gemini": "bundle/gemini.js" }, "devDependencies": { + "@agentclientprotocol/sdk": "^0.12.0", "@octokit/rest": "^22.0.0", "@types/marked": "^5.0.2", "@types/mime-types": "^3.0.1", @@ -110,9 +111,9 @@ } }, "node_modules/@agentclientprotocol/sdk": { - "version": "0.11.0", - "resolved": "https://registry.npmjs.org/@agentclientprotocol/sdk/-/sdk-0.11.0.tgz", - "integrity": "sha512-hngnMwQ13DCC7oEr0BUnrx+vTDFf/ToCLhF0YcCMWRs+v4X60rKQyAENsx0PdbQF21jC1VjMFkh2+vwNBLh6fQ==", + "version": "0.12.0", + "resolved": "https://registry.npmjs.org/@agentclientprotocol/sdk/-/sdk-0.12.0.tgz", + "integrity": "sha512-V8uH/KK1t7utqyJmTA7y7DzKu6+jKFIXM+ZVouz8E55j8Ej2RV42rEvPKn3/PpBJlliI5crcGk1qQhZ7VwaepA==", "license": "Apache-2.0", "peerDependencies": { "zod": "^3.25.0 || ^4.0.0" @@ -18695,7 +18696,7 @@ "version": "0.26.0-nightly.20260114.bb6c57414", "license": "Apache-2.0", "dependencies": { - "@agentclientprotocol/sdk": "^0.11.0", + "@agentclientprotocol/sdk": "^0.12.0", "@google/gemini-cli-core": "file:../core", "@google/genai": "1.30.0", "@iarna/toml": "^2.2.5", diff --git a/package.json b/package.json index ff97e64715..e4602937ba 100644 --- a/package.json +++ b/package.json @@ -79,6 +79,7 @@ "LICENSE" ], "devDependencies": { + "@agentclientprotocol/sdk": "^0.12.0", "@octokit/rest": "^22.0.0", "@types/marked": "^5.0.2", "@types/mime-types": "^3.0.1", diff --git a/packages/cli/package.json b/packages/cli/package.json index ea1737738e..2f8e5ec8c2 100644 --- a/packages/cli/package.json +++ b/packages/cli/package.json @@ -29,7 +29,7 @@ "sandboxImageUri": "us-docker.pkg.dev/gemini-code-dev/gemini-cli/sandbox:0.26.0-nightly.20260114.bb6c57414" }, "dependencies": { - "@agentclientprotocol/sdk": "^0.11.0", + "@agentclientprotocol/sdk": "^0.12.0", "@google/gemini-cli-core": "file:../core", "@google/genai": "1.30.0", "@iarna/toml": "^2.2.5", diff --git a/packages/cli/src/zed-integration/zedIntegration.ts b/packages/cli/src/zed-integration/zedIntegration.ts index d4381efc0e..0d7d3262ac 100644 --- a/packages/cli/src/zed-integration/zedIntegration.ts +++ b/packages/cli/src/zed-integration/zedIntegration.ts @@ -44,6 +44,7 @@ import { z } from 'zod'; import { randomUUID } from 'node:crypto'; import type { CliArgs } from '../config/config.js'; import { loadCliConfig } from '../config/config.js'; +import { runExitCleanup } from '../utils/cleanup.js'; export async function runZedIntegration( config: Config, @@ -55,10 +56,15 @@ export async function runZedIntegration( const stdin = Readable.toWeb(process.stdin) as ReadableStream; const stream = acp.ndJsonStream(stdout, stdin); - new acp.AgentSideConnection( + const connection = new acp.AgentSideConnection( (connection) => new GeminiAgent(config, settings, argv, connection), stream, ); + + // SIGTERM/SIGINT handlers (in sdk.ts) don't fire when stdin closes. + // We must explicitly await the connection close to flush telemetry. + // Use finally() to ensure cleanup runs even on stream errors. + await connection.closed.finally(runExitCleanup); } export class GeminiAgent { diff --git a/packages/core/src/core/contentGenerator.ts b/packages/core/src/core/contentGenerator.ts index 12e07790cc..740bede47c 100644 --- a/packages/core/src/core/contentGenerator.ts +++ b/packages/core/src/core/contentGenerator.ts @@ -114,10 +114,10 @@ export async function createContentGenerator( ): Promise { const generator = await (async () => { if (gcConfig.fakeResponses) { - return new LoggingContentGenerator( - await FakeContentGenerator.fromFile(gcConfig.fakeResponses), - gcConfig, + const fakeGenerator = await FakeContentGenerator.fromFile( + gcConfig.fakeResponses, ); + return new LoggingContentGenerator(fakeGenerator, gcConfig); } const version = await getVersion(); const model = resolveModel(