Add Context Window Widget support for Claude Agent (#3816)

* Initial plan

* Add usage handler mechanism for Claude context window widget

- Add UsageHandler type and methods to IClaudeSessionStateService
- Set usage handler in claudeChatSessionContentProvider request handler
- Wire up usage reporting in ClaudeStreamingPassThroughEndpoint
- Extract usage from AnthropicMessagesProcessor completion and call handler

Co-authored-by: TylerLeonhardt <2644648+TylerLeonhardt@users.noreply.github.com>

* Add unit tests for usage handler in ClaudeSessionStateService

- Add tests for getUsageHandlerForSession
- Add tests for setUsageHandlerForSession
- Verify handler preservation across state changes
- Verify handler can be called correctly

Co-authored-by: TylerLeonhardt <2644648+TylerLeonhardt@users.noreply.github.com>

* don't do the separate cache section

* a test

* include compaction

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: TylerLeonhardt <2644648+TylerLeonhardt@users.noreply.github.com>
Co-authored-by: Tyler Leonhardt <tyleonha@microsoft.com>

extensions/copilot/src/extension/agents/claude/node/claudeCodeAgent.ts

@@ -639,6 +639,8 @@ export class ClaudeCodeSession extends Disposable {
 					this.handleAssistantMessage(message, this._currentRequest.stream, unprocessedToolCalls);
 				} else if (message.type === 'user') {
 					this.handleUserMessage(message, this._currentRequest.stream, unprocessedToolCalls, this._currentRequest.toolInvocationToken, this._currentRequest.token);
+				} else if (message.type === 'system' && message.subtype === 'compact_boundary') {
+					this._currentRequest.stream.markdown('*Conversation compacted*');
 				} else if (message.type === 'result') {
 					this.handleResultMessage(message, this._currentRequest.stream);
 					// Clear the capturing token so subsequent requests get their own

extensions/copilot/src/extension/agents/claude/node/claudeLanguageModelServer.ts

@@ -209,7 +209,8 @@ export class ClaudeLanguageModelServer extends Disposable {
 				{
 					modelMaxPromptTokens: DEFAULT_MAX_TOKENS - DEFAULT_MAX_OUTPUT_TOKENS,
 					maxOutputTokens: DEFAULT_MAX_OUTPUT_TOKENS
-				}
+				},
+				sessionId
 			);
 
 			let messagesForLogging: Raw.ChatMessage[] = [];
@@ -477,8 +478,10 @@ class ClaudeStreamingPassThroughEndpoint implements IChatEndpoint {
 		private readonly requestHeaders: http.IncomingHttpHeaders,
 		private readonly userAgentPrefix: string,
 		private readonly contextWindowOverride: { modelMaxPromptTokens?: number; maxOutputTokens?: number },
+		private readonly sessionId: string | undefined,
 		@IChatMLFetcher private readonly chatMLFetcher: IChatMLFetcher,
-		@IInstantiationService private readonly instantiationService: IInstantiationService
+		@IInstantiationService private readonly instantiationService: IInstantiationService,
+		@IClaudeSessionStateService private readonly sessionStateService: IClaudeSessionStateService
 	) { }
 
 	public get urlOrRequestMetadata(): string | RequestMetadata {
@@ -655,6 +658,18 @@ class ClaudeStreamingPassThroughEndpoint implements IChatEndpoint {
 					const completion = processor.push({ ...parsed, type }, finishCallback);
 					if (completion) {
 						feed.emitOne(completion);
+
+						// Report usage to the usage handler if available
+						if (completion.usage && this.sessionId) {
+							const usageHandler = this.sessionStateService.getUsageHandlerForSession(this.sessionId);
+							if (usageHandler) {
+								usageHandler({
+									// Could we bucketize these token counts somehow for the details?
+									promptTokens: completion.usage.prompt_tokens,
+									completionTokens: completion.usage.completion_tokens
+								});
+							}
+						}
 					}
 				} catch (e) {
 					feed.reject(e);

extensions/copilot/src/extension/agents/claude/node/claudeSessionStateService.ts

@@ -4,6 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { PermissionMode } from '@anthropic-ai/claude-agent-sdk';
+import type * as vscode from 'vscode';
 import { CapturingToken } from '../../../../platform/requestLogger/common/capturingToken';
 import { createServiceIdentifier } from '../../../../util/common/services';
 import { arrayEquals } from '../../../../util/vs/base/common/equals';
@@ -11,11 +12,17 @@ import { Emitter, Event } from '../../../../util/vs/base/common/event';
 import { Disposable } from '../../../../util/vs/base/common/lifecycle';
 import type { ClaudeFolderInfo } from '../common/claudeFolderInfo';
 
+/**
+ * Usage handler function type for reporting token usage to stream.
+ */
+export type UsageHandler = (usage: vscode.ChatResultUsage) => void;
+
 export interface SessionState {
 	modelId: string | undefined;
 	permissionMode: PermissionMode;
 	capturingToken: CapturingToken | undefined;
 	folderInfo: ClaudeFolderInfo | undefined;
+	usageHandler: UsageHandler | undefined;
 }
 
 /**
@@ -75,6 +82,16 @@ export interface IClaudeSessionStateService {
 	 * Sets the folder info for a session.
 	 */
 	setFolderInfoForSession(sessionId: string, folderInfo: ClaudeFolderInfo): void;
+
+	/**
+	 * Gets the usage handler for a session.
+	 */
+	getUsageHandlerForSession(sessionId: string): UsageHandler | undefined;
+
+	/**
+	 * Sets the usage handler for a session.
+	 */
+	setUsageHandlerForSession(sessionId: string, handler: UsageHandler | undefined): void;
 }
 
 export const IClaudeSessionStateService = createServiceIdentifier<IClaudeSessionStateService>('IClaudeSessionStateService');
@@ -108,6 +125,7 @@ export class ClaudeSessionStateService extends Disposable implements IClaudeSess
 			permissionMode: existing?.permissionMode ?? 'acceptEdits',
 			capturingToken: existing?.capturingToken,
 			folderInfo: existing?.folderInfo,
+			usageHandler: existing?.usageHandler,
 		});
 		this._onDidChangeSessionState.fire({ sessionId, modelId });
 	}
@@ -126,6 +144,7 @@ export class ClaudeSessionStateService extends Disposable implements IClaudeSess
 			permissionMode: mode,
 			capturingToken: existing?.capturingToken,
 			folderInfo: existing?.folderInfo,
+			usageHandler: existing?.usageHandler,
 		});
 		this._onDidChangeSessionState.fire({ sessionId, permissionMode: mode });
 	}
@@ -141,6 +160,7 @@ export class ClaudeSessionStateService extends Disposable implements IClaudeSess
 			permissionMode: existing?.permissionMode ?? 'acceptEdits',
 			capturingToken: token,
 			folderInfo: existing?.folderInfo,
+			usageHandler: existing?.usageHandler,
 		});
 	}
 
@@ -158,10 +178,26 @@ export class ClaudeSessionStateService extends Disposable implements IClaudeSess
 			permissionMode: existing?.permissionMode ?? 'acceptEdits',
 			capturingToken: existing?.capturingToken,
 			folderInfo,
+			usageHandler: existing?.usageHandler,
 		});
 		this._onDidChangeSessionState.fire({ sessionId, folderInfo });
 	}
 
+	getUsageHandlerForSession(sessionId: string): UsageHandler | undefined {
+		return this._sessionState.get(sessionId)?.usageHandler;
+	}
+
+	setUsageHandlerForSession(sessionId: string, handler: UsageHandler | undefined): void {
+		const existing = this._sessionState.get(sessionId);
+		this._sessionState.set(sessionId, {
+			modelId: existing?.modelId,
+			permissionMode: existing?.permissionMode ?? 'acceptEdits',
+			capturingToken: existing?.capturingToken,
+			folderInfo: existing?.folderInfo,
+			usageHandler: handler,
+		});
+	}
+
 	override dispose(): void {
 		this._sessionState.clear();
 		super.dispose();

extensions/copilot/src/extension/agents/claude/node/test/claudeSessionStateService.spec.ts

@@ -207,4 +207,95 @@ describe('ClaudeSessionStateService', () => {
 			newService.dispose();
 		});
 	});
+
+	describe('getUsageHandlerForSession', () => {
+		it('should return undefined when no usage handler is set', () => {
+			const handler = service.getUsageHandlerForSession('session-1');
+			assert.strictEqual(handler, undefined);
+		});
+
+		it('should return the set usage handler', () => {
+			const mockHandler = sinon.stub();
+			service.setUsageHandlerForSession('session-1', mockHandler);
+			const handler = service.getUsageHandlerForSession('session-1');
+			assert.strictEqual(handler, mockHandler);
+		});
+
+		it('should return different handlers for different sessions', () => {
+			const handler1 = sinon.stub();
+			const handler2 = sinon.stub();
+			service.setUsageHandlerForSession('session-1', handler1);
+			service.setUsageHandlerForSession('session-2', handler2);
+
+			const retrieved1 = service.getUsageHandlerForSession('session-1');
+			const retrieved2 = service.getUsageHandlerForSession('session-2');
+
+			assert.strictEqual(retrieved1, handler1);
+			assert.strictEqual(retrieved2, handler2);
+		});
+	});
+
+	describe('setUsageHandlerForSession', () => {
+		it('should allow setting a usage handler', () => {
+			const mockHandler = sinon.stub();
+			service.setUsageHandlerForSession('session-1', mockHandler);
+
+			const handler = service.getUsageHandlerForSession('session-1');
+			assert.strictEqual(handler, mockHandler);
+		});
+
+		it('should allow clearing a usage handler', () => {
+			const mockHandler = sinon.stub();
+			service.setUsageHandlerForSession('session-1', mockHandler);
+			service.setUsageHandlerForSession('session-1', undefined);
+
+			const handler = service.getUsageHandlerForSession('session-1');
+			assert.strictEqual(handler, undefined);
+		});
+
+		it('should preserve other state when setting usage handler', () => {
+			service.setModelIdForSession('session-1', 'claude-opus-4-20250514');
+			service.setPermissionModeForSession('session-1', 'bypassPermissions');
+
+			const mockHandler = sinon.stub();
+			service.setUsageHandlerForSession('session-1', mockHandler);
+
+			const modelId = service.getModelIdForSession('session-1');
+			assert.strictEqual(modelId, 'claude-opus-4-20250514');
+			const permissionMode = service.getPermissionModeForSession('session-1');
+			assert.strictEqual(permissionMode, 'bypassPermissions');
+		});
+
+		it('should allow usage handler to be called after setting', () => {
+			const mockHandler = sinon.stub();
+			service.setUsageHandlerForSession('session-1', mockHandler);
+
+			const handler = service.getUsageHandlerForSession('session-1');
+			handler?.({ promptTokens: 100, completionTokens: 50 });
+
+			assert.strictEqual(mockHandler.callCount, 1);
+			assert.deepStrictEqual(mockHandler.firstCall.args[0], { promptTokens: 100, completionTokens: 50 });
+		});
+
+		it('should not fire onDidChangeSessionState event', () => {
+			const events: SessionStateChangeEvent[] = [];
+			service.onDidChangeSessionState(e => events.push(e));
+
+			const mockHandler = sinon.stub();
+			service.setUsageHandlerForSession('session-1', mockHandler);
+
+			assert.strictEqual(events.length, 0);
+		});
+
+		it('should initialize defaults when session has no prior state', () => {
+			const mockHandler = sinon.stub();
+			service.setUsageHandlerForSession('new-session', mockHandler);
+
+			assert.strictEqual(service.getModelIdForSession('new-session'), undefined);
+			assert.strictEqual(service.getPermissionModeForSession('new-session'), 'acceptEdits');
+			assert.strictEqual(service.getCapturingTokenForSession('new-session'), undefined);
+			assert.strictEqual(service.getFolderInfoForSession('new-session'), undefined);
+			assert.strictEqual(service.getUsageHandlerForSession('new-session'), mockHandler);
+		});
+	});
 });

extensions/copilot/src/extension/chatSessions/vscode-node/claudeChatSessionContentProvider.ts

@@ -320,11 +320,19 @@ export class ClaudeChatSessionContentProvider extends Disposable implements vsco
 			this.sessionStateService.setPermissionModeForSession(effectiveSessionId, permissionMode);
 			this.sessionStateService.setFolderInfoForSession(effectiveSessionId, folderInfo);
 
+			// Set usage handler to report token usage for context window widget
+			this.sessionStateService.setUsageHandlerForSession(effectiveSessionId, (usage) => {
+				stream.usage(usage);
+			});
+
 			const prompt = request.prompt;
 			this._controller.updateItemStatus(effectiveSessionId, vscode.ChatSessionStatus.InProgress, prompt);
 			const result = await this.claudeAgentManager.handleRequest(effectiveSessionId, request, context, stream, token, isNewSession, yieldRequested);
 			this._controller.updateItemStatus(effectiveSessionId, vscode.ChatSessionStatus.Completed, prompt);
 
+			// Clear usage handler after request completes
+			this.sessionStateService.setUsageHandlerForSession(effectiveSessionId, undefined);
+
 			return result.errorDetails ? { errorDetails: result.errorDetails } : {};
 		};
 	}