diff --git a/docs/users/integration-jetbrains.md b/docs/users/integration-jetbrains.md index baced8149fa..f1efc7f55bd 100644 --- a/docs/users/integration-jetbrains.md +++ b/docs/users/integration-jetbrains.md @@ -8,6 +8,7 @@ - **Agent Client Protocol**: Full support for ACP enabling advanced IDE interactions - **Symbol management**: #-mention files to add them to the conversation context - **Conversation history**: Access to past conversations within the IDE +- **Context usage**: See the current context-window occupancy while Qwen Code works ### Requirements diff --git a/integration-tests/cli/acp-integration.test.ts b/integration-tests/cli/acp-integration.test.ts index 4f0b21f2959..731b2946296 100644 --- a/integration-tests/cli/acp-integration.test.ts +++ b/integration-tests/cli/acp-integration.test.ts @@ -11,6 +11,7 @@ import { createInterface } from 'node:readline'; import { setTimeout as delay } from 'node:timers/promises'; import { describe, expect, it } from 'vitest'; import { TestRig } from '../test-helper.js'; +import { startFakeOpenAIServer } from '../fake-openai-server.js'; const REQUEST_TIMEOUT_MS = 60_000; const INITIAL_PROMPT = 'Create a quick note (smoke test).'; @@ -25,6 +26,7 @@ type PendingRequest = { }; type UsageMetadata = { + inputTokens?: number | null; promptTokens?: number | null; completionTokens?: number | null; thoughtsTokens?: number | null; @@ -47,6 +49,8 @@ type SessionUpdateNotification = { }; modeId?: string; currentModeId?: string; + used?: number; + size?: number; _meta?: { usage?: UsageMetadata; }; @@ -86,7 +90,11 @@ type PermissionHandler = ( */ function setupAcpTest( rig: TestRig, - options?: { permissionHandler?: PermissionHandler; useNewFlag?: boolean }, + options?: { + permissionHandler?: PermissionHandler; + useNewFlag?: boolean; + env?: NodeJS.ProcessEnv; + }, ) { const pending = new Map(); let nextRequestId = 1; @@ -125,7 +133,7 @@ function setupAcpTest( { cwd: rig.testDir!, stdio: ['pipe', 'pipe', 'pipe'], - env: { ...process.env, QWEN_HOME: qwenHome }, + env: { ...process.env, ...options?.env, QWEN_HOME: qwenHome }, }, ); @@ -924,11 +932,34 @@ function setupAcpTest( } }); - it('receives usage metadata in agent_message_chunk updates', async () => { + it('receives private usage metadata and standard ACP usage updates', async () => { + const fakeServer = await startFakeOpenAIServer(() => ({ + content: 'hello', + usage: { + prompt_tokens: 321, + completion_tokens: 1, + total_tokens: 322, + }, + })); const rig = new TestRig(); - rig.setup('acp usage metadata'); + rig.setup('acp usage metadata', { + settings: { + model: { + generationConfig: { contextWindowSize: 128_000 }, + }, + }, + }); - const { sendRequest, cleanup, stderr, sessionUpdates } = setupAcpTest(rig); + const { sendRequest, cleanup, stderr, sessionUpdates } = setupAcpTest(rig, { + env: { + OPENAI_API_KEY: 'fake-key', + OPENAI_BASE_URL: fakeServer.baseUrl, + OPENAI_MODEL: 'fake-model', + QWEN_MODEL: 'fake-model', + NO_PROXY: '127.0.0.1,localhost', + no_proxy: '127.0.0.1,localhost', + }, + }); try { await sendRequest('initialize', { @@ -961,14 +992,29 @@ function setupAcpTest( const usage = updatesWithUsage[0].update?._meta?.usage; expect(usage).toBeDefined(); expect( - typeof usage?.promptTokens === 'number' || + typeof usage?.inputTokens === 'number' || + typeof usage?.promptTokens === 'number' || typeof usage?.totalTokens === 'number', ).toBe(true); + + const standardUsageUpdates = sessionUpdates.filter( + (u) => u.update?.sessionUpdate === 'usage_update', + ); + expect(standardUsageUpdates.length).toBeGreaterThan(0); + + const standardUsage = standardUsageUpdates.at(-1)?.update; + expect(standardUsage).toMatchObject({ used: 321, size: 128_000 }); + + const privateInputTokens = usage?.inputTokens ?? usage?.promptTokens; + if (typeof privateInputTokens === 'number') { + expect(standardUsage?.used).toBe(privateInputTokens); + } } catch (e) { if (stderr.length) console.error('Agent stderr:', stderr.join('')); throw e; } finally { await cleanup(); + await fakeServer.close(); } }); }); diff --git a/packages/acp-bridge/src/bridge.test.ts b/packages/acp-bridge/src/bridge.test.ts index 59b358d6248..5c896f6498e 100644 --- a/packages/acp-bridge/src/bridge.test.ts +++ b/packages/acp-bridge/src/bridge.test.ts @@ -13404,7 +13404,7 @@ describe('createAcpSessionBridge', () => { ); }); - it('publishes session_update events to subscribers when the agent sends them', async () => { + it('publishes ACP usage updates to subscribers unchanged', async () => { let capturedConn: AgentSideConnection | undefined; const factory: ChannelFactory = async () => { // Build a channel pair where we capture the agent-side connection @@ -13434,8 +13434,9 @@ describe('createAcpSessionBridge', () => { void capturedConn!.sessionUpdate({ sessionId: session.sessionId, update: { - sessionUpdate: 'agent_message_chunk', - content: { type: 'text', text: 'hi' }, + sessionUpdate: 'usage_update', + used: 42, + size: 128_000, }, }); @@ -13446,6 +13447,14 @@ describe('createAcpSessionBridge', () => { } expect(collected[0]?.type).toBe('session_update'); expect(collected[0]?.id).toBe(1); + expect(collected[0]?.data).toMatchObject({ + sessionId: session.sessionId, + update: { + sessionUpdate: 'usage_update', + used: 42, + size: 128_000, + }, + }); abort.abort(); await bridge.shutdown(); diff --git a/packages/cli/src/acp-integration/session/emitters/MessageEmitter.test.ts b/packages/cli/src/acp-integration/session/emitters/MessageEmitter.test.ts index f1a2016b0d1..f8afa7c2102 100644 --- a/packages/cli/src/acp-integration/session/emitters/MessageEmitter.test.ts +++ b/packages/cli/src/acp-integration/session/emitters/MessageEmitter.test.ts @@ -22,7 +22,11 @@ describe('MessageEmitter', () => { sendUpdateSpy = vi.fn().mockResolvedValue(undefined); mockContext = { sessionId: 'test-session-id', - config: {} as Config, + config: { + getContentGeneratorConfig: vi.fn().mockReturnValue({ + contextWindowSize: 128_000, + }), + } as unknown as Config, sendUpdate: sendUpdateSpy, }; emitter = new MessageEmitter(mockContext); @@ -257,6 +261,85 @@ describe('MessageEmitter', () => { }); describe('emitUsageMetadata', () => { + it('emits standard usage updates with the latest context occupancy, not a cumulative sum', async () => { + await emitter.emitUsageMetadata( + { promptTokenCount: 100, totalTokenCount: 175 }, + '', + 20, + ); + await emitter.emitUsageMetadata( + { promptTokenCount: 120, totalTokenCount: 210 }, + '', + 30, + ); + + const usageUpdates = sendUpdateSpy.mock.calls + .map(([update]) => update) + .filter((update) => update.sessionUpdate === 'usage_update'); + expect(usageUpdates).toEqual([ + { sessionUpdate: 'usage_update', used: 100, size: 128_000 }, + { sessionUpdate: 'usage_update', used: 120, size: 128_000 }, + ]); + + // Keep emitting the existing Qwen extension for current consumers. + expect(sendUpdateSpy).toHaveBeenCalledWith( + expect.objectContaining({ + sessionUpdate: 'agent_message_chunk', + _meta: expect.objectContaining({ + usage: expect.objectContaining({ inputTokens: 120 }), + }), + }), + ); + }); + + it('does not replace main-session context usage with replay or subagent usage', async () => { + await emitter.emitUsageMetadata({ promptTokenCount: 90 }); + await emitter.emitUsageMetadata({ promptTokenCount: 40 }, '', 10, { + parentToolCallId: 'agent-parent-1', + subagentType: 'general-purpose', + }); + + expect( + sendUpdateSpy.mock.calls + .map(([update]) => update) + .filter((update) => update.sessionUpdate === 'usage_update'), + ).toEqual([]); + expect(sendUpdateSpy).toHaveBeenCalledTimes(2); + }); + + it('falls back to total tokens when a live provider omits prompt tokens', async () => { + await emitter.emitUsageMetadata({ totalTokenCount: 75 }, '', 10); + + expect(sendUpdateSpy).toHaveBeenLastCalledWith({ + sessionUpdate: 'usage_update', + used: 75, + size: 128_000, + }); + }); + + it('keeps private usage metadata when the context window is unresolved', async () => { + const ctx: SessionContext = { + ...mockContext, + config: { + getContentGeneratorConfig: () => undefined, + } as unknown as Config, + }; + + await new MessageEmitter(ctx).emitUsageMetadata( + { promptTokenCount: 75 }, + '', + 10, + ); + + expect(sendUpdateSpy).toHaveBeenCalledTimes(1); + expect(sendUpdateSpy).toHaveBeenCalledWith( + expect.objectContaining({ + sessionUpdate: 'agent_message_chunk', + _meta: expect.objectContaining({ usage: expect.any(Object) }), + }), + ); + }); + it('should emit agent_message_chunk with _meta.usage containing token counts', async () => { const usageMetadata = { promptTokenCount: 100, @@ -317,7 +400,7 @@ describe('MessageEmitter', () => { // Live round (durationMs present) → the counts are drained and stamped. await emitter.emitUsageMetadata({ totalTokenCount: 1 }, '', 500); - expect(sendUpdateSpy).toHaveBeenLastCalledWith( + expect(sendUpdateSpy).toHaveBeenCalledWith( expect.objectContaining({ _meta: expect.objectContaining({ apiErrors: 2, apiRetries: 1 }), }), @@ -326,7 +409,14 @@ describe('MessageEmitter', () => { // A second live round with nothing pending carries neither key (the first // emit drained the tracker to zero). await emitter.emitUsageMetadata({ totalTokenCount: 1 }, '', 500); - const secondMeta = sendUpdateSpy.mock.lastCall?.[0]._meta; + const privateUsageUpdates = sendUpdateSpy.mock.calls + .map(([update]) => update) + .filter( + (update) => + update.sessionUpdate === 'agent_message_chunk' && + update._meta?.usage, + ); + const secondMeta = privateUsageUpdates.at(-1)?._meta; expect(secondMeta).not.toHaveProperty('apiErrors'); expect(secondMeta).not.toHaveProperty('apiRetries'); }); @@ -352,7 +442,7 @@ describe('MessageEmitter', () => { }; const ctx: SessionContext = { sessionId: 'test-session-id', - config: {} as Config, + config: mockContext.config, sendUpdate: sendUpdateSpy, cumulativeUsage, }; diff --git a/packages/cli/src/acp-integration/session/emitters/MessageEmitter.ts b/packages/cli/src/acp-integration/session/emitters/MessageEmitter.ts index 3de8933862d..f8f625a4955 100644 --- a/packages/cli/src/acp-integration/session/emitters/MessageEmitter.ts +++ b/packages/cli/src/acp-integration/session/emitters/MessageEmitter.ts @@ -5,7 +5,7 @@ */ import type { GenerateContentResponseUsageMetadata } from '@google/genai'; -import type { SubagentMeta } from '../types.js'; +import { hasFullSessionContext, type SubagentMeta } from '../types.js'; import { createTranscriptMessageUpdate, createTranscriptUsageUpdate, @@ -234,6 +234,39 @@ export class MessageEmitter extends BaseEmitter { }, }), ); + + // ACP clients such as JetBrains render context occupancy from the + // standard usage_update frame rather than Qwen's private `_meta.usage`. + // Emit it only for a live main-session model round: replay frames do not + // have a duration, and subagent usage describes a separate context window + // that must not replace the parent session's indicator. + if ( + !Number.isFinite(durationMs) || + subagentMeta || + !hasFullSessionContext(this.ctx) + ) { + return; + } + + const used = + usageMetadata.promptTokenCount ?? usageMetadata.totalTokenCount; + const size = this.ctx.config.getContentGeneratorConfig()?.contextWindowSize; + if ( + typeof used !== 'number' || + !Number.isSafeInteger(used) || + used < 0 || + typeof size !== 'number' || + !Number.isSafeInteger(size) || + size <= 0 + ) { + return; + } + + await this.sendUpdate({ + sessionUpdate: 'usage_update', + used, + size, + }); } /**