diff --git a/extensions/copilot/package.json b/extensions/copilot/package.json index f4b32bb21ff..9a5c3a55e01 100644 --- a/extensions/copilot/package.json +++ b/extensions/copilot/package.json @@ -2602,6 +2602,16 @@ "title": "Export All Prompt Logs as JSON...", "icon": "$(export)" }, + { + "command": "github.copilot.chat.debug.exportTrajectories", + "title": "Export Agent Trajectories", + "category": "Chat" + }, + { + "command": "github.copilot.chat.debug.exportSingleTrajectory", + "title": "Export Trajectory...", + "category": "Chat" + }, { "command": "github.copilot.nes.captureExpected.start", "title": "Record Expected Edit (NES)", @@ -4427,6 +4437,10 @@ } ], "commandPalette": [ + { + "command": "github.copilot.chat.debug.exportSingleTrajectory", + "when": "false" + }, { "command": "github.copilot.chat.triggerPermissiveSignIn", "when": "false" @@ -4722,6 +4736,11 @@ "command": "github.copilot.chat.debug.exportPromptLogsAsJson", "when": "view == copilot-chat && viewItem == chatprompt", "group": "export@3" + }, + { + "command": "github.copilot.chat.debug.exportSingleTrajectory", + "when": "view == copilot-chat && viewItem == chatprompt", + "group": "export@4" } ], "searchPanel/aiResults/commands": [ diff --git a/extensions/copilot/src/extension/extension/vscode-node/contributions.ts b/extensions/copilot/src/extension/extension/vscode-node/contributions.ts index e4cff67c943..27cdae932f8 100644 --- a/extensions/copilot/src/extension/extension/vscode-node/contributions.ts +++ b/extensions/copilot/src/extension/extension/vscode-node/contributions.ts @@ -45,6 +45,7 @@ import { SettingsSchemaFeature } from '../../settingsSchema/vscode-node/settings import { SurveyCommandContribution } from '../../survey/vscode-node/surveyCommands'; import { SetupTestsContribution } from '../../testing/vscode/setupTestContributions'; import { ToolsContribution } from '../../tools/vscode-node/tools'; +import { TrajectoryExportCommands } from '../../trajectory/vscode-node/trajectoryExportCommands'; import { InlineCompletionContribution } from '../../typescriptContext/vscode-node/languageContextService'; import { NesRenameContribution } from '../../typescriptContext/vscode-node/nesRenameService'; import * as workspaceChunkSearchContribution from '../../workspaceChunkSearch/node/workspaceChunkSearch.contribution'; @@ -120,4 +121,5 @@ export const vscodeNodeChatContributions: IExtensionContributionFactory[] = [ asContributionFactory(LanguageModelProxyContrib), asContributionFactory(PromptFileContribution), newWorkspaceContribution, + asContributionFactory(TrajectoryExportCommands), ]; diff --git a/extensions/copilot/src/extension/extension/vscode-node/services.ts b/extensions/copilot/src/extension/extension/vscode-node/services.ts index 4a6ab455fa8..9e8403fcbdd 100644 --- a/extensions/copilot/src/extension/extension/vscode-node/services.ts +++ b/extensions/copilot/src/extension/extension/vscode-node/services.ts @@ -68,6 +68,8 @@ import { IWorkspaceMutationManager } from '../../../platform/testing/common/work import { ISetupTestsDetector, SetupTestsDetector } from '../../../platform/testing/node/setupTestDetector'; import { ITestDepsResolver, TestDepsResolver } from '../../../platform/testing/node/testDepsResolver'; import { ITokenizerProvider, TokenizerProvider } from '../../../platform/tokenizer/node/tokenizer'; +import { ITrajectoryLogger } from '../../../platform/trajectory/common/trajectoryLogger'; +import { TrajectoryLogger } from '../../../platform/trajectory/node/trajectoryLogger'; import { GithubAvailableEmbeddingTypesService, IGithubAvailableEmbeddingTypesService } from '../../../platform/workspaceChunkSearch/common/githubAvailableEmbeddingTypes'; import { IRerankerService, RerankerService } from '../../../platform/workspaceChunkSearch/common/rerankerService'; import { IWorkspaceChunkSearchService, WorkspaceChunkSearchService } from '../../../platform/workspaceChunkSearch/node/workspaceChunkSearchService'; @@ -221,6 +223,7 @@ export function registerServices(builder: IInstantiationServiceBuilder, extensio builder.define(IUndesiredModelsManager, new SyncDescriptor(UndesiredModels.Manager)); builder.define(ICopilotInlineCompletionItemProviderService, new SyncDescriptor(CopilotInlineCompletionItemProviderService)); builder.define(IGitHubOrgChatResourcesService, new SyncDescriptor(GitHubOrgChatResourcesService)); + builder.define(ITrajectoryLogger, new SyncDescriptor(TrajectoryLogger)); } function setupMSFTExperimentationService(builder: IInstantiationServiceBuilder, extensionContext: ExtensionContext) { diff --git a/extensions/copilot/src/extension/prompt/node/defaultIntentRequestHandler.ts b/extensions/copilot/src/extension/prompt/node/defaultIntentRequestHandler.ts index 4b42958deca..65713487af6 100644 --- a/extensions/copilot/src/extension/prompt/node/defaultIntentRequestHandler.ts +++ b/extensions/copilot/src/extension/prompt/node/defaultIntentRequestHandler.ts @@ -128,7 +128,19 @@ export class DefaultIntentRequestHandler { return confirmationResult; } - const resultDetails = await this._requestLogger.captureInvocation(new CapturingToken(this.request.prompt, 'comment', false), () => this.runWithToolCalling(intentInvocation)); + // For subagent requests, use the subAgentInvocationId as the session ID. + // This enables explicit linking between the parent's runSubagent tool call and the subagent trajectory. + // For main requests, use the VS Code chat sessionId directly as the trajectory session ID. + const capturingToken = new CapturingToken( + this.request.prompt, + 'comment', + false, + false, + this.request.subAgentInvocationId, + this.request.subAgentName, + this.request.sessionId, + ); + const resultDetails = await this._requestLogger.captureInvocation(capturingToken, () => this.runWithToolCalling(intentInvocation)); let chatResult = resultDetails.chatResult || {}; this._surveyService.signalUsage(`${this.location === ChatLocation.Editor ? 'inline' : 'panel'}.${this.intent.id}`, this.documentContext?.document.languageId); diff --git a/extensions/copilot/src/extension/prompt/node/searchSubagentToolCallingLoop.ts b/extensions/copilot/src/extension/prompt/node/searchSubagentToolCallingLoop.ts index 0cbcb922981..dd09a322eda 100644 --- a/extensions/copilot/src/extension/prompt/node/searchSubagentToolCallingLoop.ts +++ b/extensions/copilot/src/extension/prompt/node/searchSubagentToolCallingLoop.ts @@ -27,6 +27,8 @@ export interface ISearchSubagentToolCallingLoopOptions extends IToolCallingLoopO request: ChatRequest; location: ChatLocation; promptText: string; + /** Optional pre-generated subagent invocation ID. If not provided, a new UUID will be generated. */ + subAgentInvocationId?: string; } export class SearchSubagentToolCallingLoop extends ToolCallingLoop { @@ -54,7 +56,7 @@ export class SearchSubagentToolCallingLoop extends ToolCallingLoop { const request = this._inputContext!.request!; const parentSessionId = this._inputContext?.conversation?.sessionId ?? generateUuid(); + // Generate a stable session ID for this subagent invocation that will be used: + // 1. As subAgentInvocationId in the subagent's tool context + // 2. As subAgentInvocationId in toolMetadata for parent trajectory linking + // 3. As the session_id in the subagent's own trajectory + const subAgentInvocationId = generateUuid(); const toolCallLimit = this.configurationService.getExperimentBasedConfig(ConfigKey.Advanced.SearchSubagentToolCallLimit, this.experimentationService); @@ -64,6 +69,7 @@ class SearchSubagentTool implements ICopilotTool { request: request, location: request.location, promptText: options.input.query, + subAgentInvocationId: subAgentInvocationId, }); const stream = this._inputContext?.stream && ChatResponseStreamImpl.filter( @@ -73,10 +79,14 @@ class SearchSubagentTool implements ICopilotTool { // Create a new capturing token to group this search subagent and all its nested tool calls // Similar to how DefaultIntentRequestHandler does it + // Pass the subAgentInvocationId so the trajectory uses this ID for explicit linking const searchSubagentToken = new CapturingToken( `Search: ${options.input.query.substring(0, 50)}${options.input.query.length > 50 ? '...' : ''}`, 'search', - false + false, + false, + subAgentInvocationId, + 'search' // subAgentName for trajectory tracking ); // Wrap the loop execution in captureInvocation with the new token @@ -87,8 +97,10 @@ class SearchSubagentTool implements ICopilotTool { // All nested tool calls are already logged by ToolCallingLoop.logToolResult() const toolMetadata = { query: options.input.query, - description: options.input.description - // details: options.input.details + description: options.input.description, + // The subAgentInvocationId links this tool call to the subagent's trajectory + subAgentInvocationId: subAgentInvocationId, + agentName: 'search' }; let subagentResponse = ''; diff --git a/extensions/copilot/src/extension/trajectory/ARCHITECTURE.md b/extensions/copilot/src/extension/trajectory/ARCHITECTURE.md new file mode 100644 index 00000000000..8e5d23f656e --- /dev/null +++ b/extensions/copilot/src/extension/trajectory/ARCHITECTURE.md @@ -0,0 +1,431 @@ +# Trajectory Logging Architecture + +This document explains the architecture of trajectory logging in VS Code Copilot Chat, +including how it relates to the existing request logging system. + +## Overview + +The trajectory logging system captures agent execution traces in the **ATIF (Agent Trajectory +Interchange Format)** for analysis, debugging, and potential benchmarking. It builds on top +of the existing `RequestLogger` infrastructure. + +``` +┌─────────────────────────────────────────────────────────────────────────────────┐ +│ Runtime Data Flow │ +└─────────────────────────────────────────────────────────────────────────────────┘ + + Chat Request Tool Execution LLM Response + │ │ │ + ▼ ▼ ▼ +┌─────────────────────────────────────────────────────────────────────────────────┐ +│ RequestLogger │ +│ ┌─────────────────────────────────────────────────────────────────────────┐ │ +│ │ _entries: LoggedInfo[] (bounded by RequestLoggerMaxEntries) │ │ +│ │ ├── LoggedRequestInfo { id, entry: LoggedRequest, token } │ │ +│ │ ├── LoggedToolCall { id, name, args, response, token } │ │ +│ │ └── LoggedElementInfo { id, name, tokens, trace, token } │ │ +│ └─────────────────────────────────────────────────────────────────────────┘ │ +│ │ │ +│ onDidChangeRequests (Event) │ +└─────────────────────────────────────┬───────────────────────────────────────────┘ + │ + ▼ +┌─────────────────────────────────────────────────────────────────────────────────┐ +│ TrajectoryLoggerAdapter │ +│ ┌─────────────────────────────────────────────────────────────────────────┐ │ +│ │ TRACKING STATE (⚠️ UNBOUNDED - Memory Leak Risk) │ │ +│ │ ├── processedEntries: Set # entry IDs already synced │ │ +│ │ ├── processedToolCalls: Set # tool call IDs processed │ │ +│ │ ├── lastUserMessageBySession: Map │ │ +│ │ ├── requestToStepContext: Map │ │ +│ │ └── runSubagentToolCallToSessionId: Map │ │ +│ │ │ │ +│ │ TOKEN MAPPING (WeakMap - GC-friendly) │ │ +│ │ ├── sessionMap: WeakMap │ │ +│ │ └── tokenToSessionId: WeakMap │ │ +│ └─────────────────────────────────────────────────────────────────────────┘ │ +│ │ │ +│ Converts LoggedInfo → TrajectoryStep │ +└─────────────────────────────────────┬───────────────────────────────────────────┘ + │ + ▼ +┌─────────────────────────────────────────────────────────────────────────────────┐ +│ TrajectoryLogger │ +│ ┌─────────────────────────────────────────────────────────────────────────┐ │ +│ │ TRAJECTORY STORAGE (⚠️ UNBOUNDED) │ │ +│ │ ├── trajectories: Map │ │ +│ │ │ └── steps: ITrajectoryStep[] │ │ +│ │ └── subagentTrajectories: Map │ │ +│ │ │ │ +│ │ currentSessionId: string | undefined │ │ +│ └─────────────────────────────────────────────────────────────────────────┘ │ +└─────────────────────────────────────────────────────────────────────────────────┘ +``` + +## Component Responsibilities + +### 1. RequestLogger (`src/extension/prompt/vscode-node/requestLoggerImpl.ts`) + +**Purpose**: Captures all LLM requests, tool calls, and prompt traces for debugging. + +**Storage Pattern**: +- **Bounded array** with configurable max size (`RequestLoggerMaxEntries`) +- Old entries are shifted out when limit is reached +- Uses `AsyncLocalStorage` to associate requests with `CapturingToken` + +```typescript +// Bounded storage - old entries evicted +private readonly _entries: LoggedInfo[] = []; + +private async _addEntry(entry: LoggedInfo): Promise { + this._entries.push(entry); + const maxEntries = this._configService.getConfig(ConfigKey.Advanced.RequestLoggerMaxEntries); + if (this._entries.length > maxEntries) { + this._entries.shift(); // ✅ Bounded - evicts oldest + } + // ... +} +``` + +### 2. TrajectoryLogger (`src/platform/trajectory/node/trajectoryLogger.ts`) + +**Purpose**: Builds structured trajectory objects in ATIF format. + +**Storage Pattern**: +- **Unbounded Maps** - trajectories persist until explicit `clearTrajectory()` call +- Each session gets its own `TrajectoryBuilder` +- Steps accumulate within builders + +```typescript +// ⚠️ Unbounded storage - never auto-cleared +private readonly trajectories = new Map(); +private subagentTrajectories = new Map(); +``` + +### 3. TrajectoryLoggerAdapter (`src/platform/trajectory/node/trajectoryLoggerAdapter.ts`) + +**Purpose**: Bridge between RequestLogger events and TrajectoryLogger. + +**Storage Pattern**: +- **WeakMaps** for token→sessionId mapping (GC-friendly ✅) +- **Sets/Maps** for deduplication tracking (⚠️ Unbounded) + +```typescript +// ✅ GC-friendly - tokens can be collected +private sessionMap = new WeakMap(); +private tokenToSessionId = new WeakMap(); + +// ⚠️ UNBOUNDED - grows indefinitely +private processedEntries = new Set(); // entry.id strings +private processedToolCalls = new Set(); // tool call ID strings +private lastUserMessageBySession = new Map(); +private requestToStepContext = new Map(); +private runSubagentToolCallToSessionId = new Map(); +``` + +## Data Flow Sequence + +``` +┌────────────────────────────────────────────────────────────────────────────────┐ +│ Sequence: Agent Request Flow │ +└────────────────────────────────────────────────────────────────────────────────┘ + +User Input ────┐ + │ + ▼ +┌──────────────────────────┐ +│ captureInvocation() │ CapturingToken created with: +│ with CapturingToken │ - chatSessionId (if main chat) +└──────────────────────────┘ - subAgentInvocationId (if subagent) + │ + ▼ +┌──────────────────────────┐ +│ LLM Request Sent │ +│ logChatRequest() called │ +└──────────────────────────┘ + │ + ▼ +┌──────────────────────────┐ +│ RequestLogger._addEntry │ LoggedRequestInfo stored +│ fires onDidChangeReqs │ (bounded, oldest evicted) +└──────────────────────────┘ + │ + ▼ +┌──────────────────────────┐ +│ Adapter.syncTrajectories│ Iterates all RequestLogger entries +│ │ Skips if entry.id in processedEntries +└──────────────────────────┘ + │ + ▼ +┌──────────────────────────┐ +│ TrajectoryLogger │ +│ .startTrajectory() │ Creates/updates TrajectoryBuilder +│ .beginAgentStep() │ Adds step to builder +└──────────────────────────┘ + │ + ▼ +┌──────────────────────────┐ If LLM response includes tool calls: +│ Tool Execution │ - logToolCall() → RequestLogger +│ │ - Adapter processes as observation +└──────────────────────────┘ + │ + ▼ +┌──────────────────────────┐ +│ Trajectory Complete │ getAllTrajectories() returns +│ (on export) │ Map +└──────────────────────────┘ +``` + +## Session ID Resolution + +The adapter determines session IDs with the following priority: + +``` +┌─────────────────────────────────────────────────────────────────┐ +│ Session ID Priority │ +├─────────────────────────────────────────────────────────────────┤ +│ 1. token.subAgentInvocationId │ Explicit subagent linking │ +│ 2. token.chatSessionId │ VS Code chat session ID │ +│ 3. generateSessionId(label) │ Fallback: hash + timestamp │ +└─────────────────────────────────────────────────────────────────┘ +``` + +This enables: +- **Main trajectories**: Use VS Code's chat session ID for 1:1 mapping +- **Subagent trajectories**: Use pre-assigned invocation ID for parent↔child linking +- **Parent references child**: Tool call observation includes `subagent_trajectory_ref` + +## Memory Lifecycle Comparison + +``` +┌────────────────────────────────────────────────────────────────────────────────┐ +│ Memory Growth Over Extension Lifetime │ +├────────────────────────────────────────────────────────────────────────────────┤ +│ │ +│ RequestLogger._entries │ +│ ┌────────────────────────────────────────────┐ │ +│ │████████████████████████████████████████████│ ← Bounded (shifts oldest) │ +│ └────────────────────────────────────────────┘ │ +│ Max = RequestLoggerMaxEntries │ +│ │ +│ ───────────────────────────────────────────────────────────────────────────── │ +│ │ +│ TrajectoryLoggerAdapter.processedEntries │ +│ ┌────────────────────────────────────────────────────────────────────────────┐│ +│ │████████████████████████████████████████████████████████████████████████████││ +│ └────────────────────────────────────────────────────────────────────────────┘│ +│ ⚠️ UNBOUNDED - keeps all entry IDs ever seen │ +│ │ +│ ───────────────────────────────────────────────────────────────────────────── │ +│ │ +│ TrajectoryLogger.trajectories │ +│ ┌────────────────────────────────────────────────────────────────────────────┐│ +│ │████████████████████████████████████████████████████████████████████████████││ +│ └────────────────────────────────────────────────────────────────────────────┘│ +│ ⚠️ UNBOUNDED - accumulates all sessions until clearTrajectory() │ +│ │ +└────────────────────────────────────────────────────────────────────────────────┘ + Time → +``` + +## The Memory Leak Problem + +### Root Cause + +When RequestLogger evicts old entries (bounded), the adapter's tracking sets still retain: +- Entry IDs in `processedEntries` +- Tool call IDs in `processedToolCalls` +- Session data in `lastUserMessageBySession`, `requestToStepContext`, etc. + +This creates **orphaned references** that can never be cleaned up. + +### Problematic Scenario + +``` +Time T1: Entry "abc123" logged → adapter tracks in processedEntries +Time T2: RequestLogger evicts "abc123" (hit max entries) +Time T3: processedEntries still contains "abc123" forever ❌ + +Result: Set grows unboundedly with orphaned string IDs +``` + +## Proposed Fix: Session-Scoped Cleanup + +Add lifecycle management to clear adapter state when trajectories are cleared: + +```typescript +// In TrajectoryLoggerAdapter +public clearSessionState(sessionId?: string): void { + if (sessionId) { + // Clear session-specific data + this.lastUserMessageBySession.delete(sessionId); + this.pendingStepContexts.delete(sessionId); + // Clear requestToStepContext entries for this session + for (const [key, info] of this.requestToStepContext) { + if (info.sessionId === sessionId) { + this.requestToStepContext.delete(key); + } + } + } else { + // Clear all state + this.processedEntries.clear(); + this.processedToolCalls.clear(); + this.lastUserMessageBySession.clear(); + this.pendingStepContexts.clear(); + this.requestToStepContext.clear(); + this.runSubagentToolCallToSessionId.clear(); + } +} +``` + +Additionally, add a bounded safety valve (similar to RequestLogger): + +```typescript +const MAX_PROCESSED_ENTRIES = 10000; // Reasonable upper bound + +private async syncTrajectories(): Promise { + // Safety valve: prevent unbounded growth + if (this.processedEntries.size > MAX_PROCESSED_ENTRIES) { + // Clear oldest entries (convert to array, slice, rebuild set) + const entries = [...this.processedEntries]; + this.processedEntries = new Set(entries.slice(-MAX_PROCESSED_ENTRIES / 2)); + } + // ... rest of sync logic +} +``` + +## File Structure + +``` +src/ +├── platform/ +│ ├── requestLogger/ +│ │ ├── common/ +│ │ │ └── capturingToken.ts # Token for grouping requests +│ │ └── node/ +│ │ └── requestLogger.ts # Abstract + interfaces +│ │ +│ └── trajectory/ +│ ├── common/ +│ │ ├── trajectoryLogger.ts # Interface definition +│ │ └── trajectoryTypes.ts # ATIF schema types +│ ├── node/ +│ │ ├── trajectoryLogger.ts # Concrete implementation +│ │ └── trajectoryLoggerAdapter.ts # Bridge to RequestLogger +│ └── test/ +│ └── node/ +│ └── trajectoryLoggerAdapter.spec.ts +│ +└── extension/ + ├── prompt/ + │ └── vscode-node/ + │ └── requestLoggerImpl.ts # Concrete RequestLogger + │ + └── trajectory/ + └── vscode-node/ + └── trajectoryExportCommands.ts # Export commands +``` + +## Export Flow + +``` +┌─────────────────────────────────────────────────────────────────────────────────┐ +│ Export Trajectory Command │ +└─────────────────────────────────────────────────────────────────────────────────┘ + + User clicks "Export Trajectory" on tree item + │ + ▼ + ┌───────────────────────────────────┐ + │ Get CapturingToken from treeItem │ + └───────────────────────────────────┘ + │ + ▼ + ┌───────────────────────────────────┐ + │ adapter.getSessionIdForToken() │ WeakMap lookup + └───────────────────────────────────┘ + │ + ▼ + ┌───────────────────────────────────┐ + │ trajectoryLogger.getAllTrajectories() │ + │ → Map │ + └───────────────────────────────────┘ + │ + ▼ + ┌───────────────────────────────────┐ + │ collectTrajectoryWithSubagents() │ Recursively collect + │ (follows subagent_trajectory_ref)│ linked subagent trajectories + └───────────────────────────────────┘ + │ + ▼ + ┌───────────────────────────────────┐ + │ Write .trajectory.json files │ + │ to user-selected folder │ + └───────────────────────────────────┘ +``` + +## Output Format (ATIF v1.5) + +```json +{ + "schema_version": "ATIF-v1.5", + "session_id": "chat-session-abc123", + "agent": { + "name": "GitHub Copilot Chat", + "version": "1.0.0", + "tool_definitions": [...] + }, + "steps": [ + { + "step_id": 1, + "timestamp": "2026-01-28T10:00:00.000Z", + "source": "user", + "message": "Create a hello world function" + }, + { + "step_id": 2, + "timestamp": "2026-01-28T10:00:01.000Z", + "source": "agent", + "model_name": "gpt-4o", + "message": "I'll create a hello world function.", + "tool_calls": [ + { + "tool_call_id": "call_123", + "function_name": "create_file", + "arguments": { "filePath": "/hello.ts", "content": "..." } + } + ], + "observation": { + "results": [ + { "source_call_id": "call_123", "content": "File created successfully" } + ] + }, + "metrics": { + "prompt_tokens": 1500, + "completion_tokens": 200, + "duration_ms": 2340 + } + } + ], + "final_metrics": { + "total_prompt_tokens": 1500, + "total_completion_tokens": 200, + "total_steps": 2, + "total_tool_calls": 1 + } +} +``` + +## Summary + +| Component | Storage Pattern | Bounded? | Cleanup Mechanism | +|-----------|----------------|----------|-------------------| +| RequestLogger._entries | Array | ✅ Yes | Auto-shift oldest | +| TrajectoryLogger.trajectories | Map | ❌ No | Manual clearTrajectory() | +| Adapter.processedEntries | Set | ❌ No | **None (memory leak)** | +| Adapter.processedToolCalls | Set | ❌ No | **None (memory leak)** | +| Adapter.sessionMap | WeakMap | ✅ Yes | GC when token collected | + +**Key Insight**: The adapter serves as a "translation layer" that watches RequestLogger events +and populates TrajectoryLogger. However, its deduplication tracking (Sets/Maps) grows unboundedly, +creating a memory leak in long-running sessions. diff --git a/extensions/copilot/src/extension/trajectory/vscode-node/trajectoryExportCommands.ts b/extensions/copilot/src/extension/trajectory/vscode-node/trajectoryExportCommands.ts new file mode 100644 index 00000000000..4f3f171e344 --- /dev/null +++ b/extensions/copilot/src/extension/trajectory/vscode-node/trajectoryExportCommands.ts @@ -0,0 +1,278 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import * as os from 'os'; +import * as vscode from 'vscode'; +import { IConfigurationService } from '../../../platform/configuration/common/configurationService'; +import { CapturingToken } from '../../../platform/requestLogger/common/capturingToken'; +import { IRequestLogger } from '../../../platform/requestLogger/node/requestLogger'; +import { ITrajectoryLogger } from '../../../platform/trajectory/common/trajectoryLogger'; +import { TRAJECTORY_FILE_EXTENSION, type IAgentTrajectory, type IObservationResult, type ITrajectoryStep } from '../../../platform/trajectory/common/trajectoryTypes'; +import { TrajectoryLoggerAdapter } from '../../../platform/trajectory/node/trajectoryLoggerAdapter'; +import { Disposable } from '../../../util/vs/base/common/lifecycle'; +import { IExtensionContribution } from '../../common/contributions'; +import { renderToolResultToStringNoBudget } from '../../prompt/vscode-node/requestLoggerToolResult'; + +const exportTrajectoriesCommand = 'github.copilot.chat.debug.exportTrajectories'; +const exportSingleTrajectoryCommand = 'github.copilot.chat.debug.exportSingleTrajectory'; + +/** + * Command contribution for exporting agent trajectories + */ +export class TrajectoryExportCommands extends Disposable implements IExtensionContribution { + readonly id = 'trajectoryExportCommands'; + private readonly adapter: TrajectoryLoggerAdapter; + + constructor( + @ITrajectoryLogger private readonly trajectoryLogger: ITrajectoryLogger, + @IRequestLogger requestLogger: IRequestLogger, + @IConfigurationService configService: IConfigurationService, + ) { + super(); + // Initialize adapter to bridge RequestLogger to TrajectoryLogger + // The adapter subscribes to RequestLogger events and populates TrajectoryLogger + this.adapter = this._register(new TrajectoryLoggerAdapter(requestLogger, trajectoryLogger, configService, renderToolResultToStringNoBudget)); + this.registerCommands(); + } + + private registerCommands(): void { + this._register(vscode.commands.registerCommand(exportTrajectoriesCommand, async (savePath?: string) => { + await this.exportTrajectories(savePath); + })); + + this._register(vscode.commands.registerCommand(exportSingleTrajectoryCommand, async (treeItem?: { token?: CapturingToken }) => { + await this.exportSingleTrajectory(treeItem); + })); + } + + /** + * Build a mapping from sessionId to trajectory_path by scanning subagent refs + */ + private buildTrajectoryPathMapping(trajectories: Map): Map { + const mapping = new Map(); + for (const trajectory of trajectories.values()) { + const steps: ITrajectoryStep[] = Array.isArray(trajectory?.steps) ? trajectory.steps : []; + for (const step of steps) { + const results: IObservationResult[] = Array.isArray(step.observation?.results) ? step.observation.results : []; + for (const r of results) { + for (const ref of r.subagent_trajectory_ref ?? []) { + if (ref.session_id && ref.trajectory_path && !mapping.has(ref.session_id)) { + mapping.set(ref.session_id, ref.trajectory_path); + } + } + } + } + } + return mapping; + } + + /** + * Get the filename for a trajectory, using referenced path if available + */ + private getTrajectoryFilename(sessionId: string, pathMapping: Map): string { + const referencedPath = pathMapping.get(sessionId); + const rawFilename = referencedPath + ? this.sanitizeFilename(referencedPath) + : this.sanitizeFilename(sessionId); + return rawFilename.endsWith(TRAJECTORY_FILE_EXTENSION) + ? rawFilename + : `${rawFilename}${TRAJECTORY_FILE_EXTENSION}`; + } + + /** + * Write multiple trajectories to a folder + */ + private async writeTrajectoriesToFolder( + trajectories: Map, + saveDir: vscode.Uri, + pathMapping: Map + ): Promise { + for (const [sessionId, trajectory] of trajectories) { + const filename = this.getTrajectoryFilename(sessionId, pathMapping); + const fileUri = vscode.Uri.joinPath(saveDir, filename); + const content = JSON.stringify(trajectory, null, 2); + await vscode.workspace.fs.writeFile(fileUri, Buffer.from(content, 'utf8')); + } + } + + /** + * Prompt user for folder selection + */ + private async promptForFolder(title: string): Promise { + const dialogResult = await vscode.window.showOpenDialog({ + canSelectFiles: false, + canSelectFolders: true, + canSelectMany: false, + title, + defaultUri: vscode.Uri.file(os.homedir()) + }); + return dialogResult?.[0]; + } + + private async exportTrajectories(savePath?: string): Promise { + const trajectories = this.trajectoryLogger.getAllTrajectories(); + + if (trajectories.size === 0) { + vscode.window.showInformationMessage('No trajectories found to export.'); + return; + } + + const saveDir = savePath + ? vscode.Uri.file(savePath) + : await this.promptForFolder('Select Folder to Export Trajectories'); + + if (!saveDir) { + return; // User cancelled + } + + try { + const pathMapping = this.buildTrajectoryPathMapping(trajectories); + await this.writeTrajectoriesToFolder(trajectories, saveDir, pathMapping); + + const revealAction = 'Reveal in Explorer'; + const result = await vscode.window.showInformationMessage( + `Successfully exported ${trajectories.size} trajectories to ${saveDir.fsPath}`, + revealAction + ); + + if (result === revealAction) { + await vscode.commands.executeCommand('revealFileInOS', saveDir); + } + } catch (error) { + vscode.window.showErrorMessage(`Failed to export trajectories: ${error instanceof Error ? error.message : String(error)}`); + } + } + + /** + * Export a single trajectory and its referenced subagent trajectories + * @param treeItem The tree item containing the capturing token + */ + private async exportSingleTrajectory(treeItem?: { token?: CapturingToken }): Promise { + if (!treeItem?.token) { + vscode.window.showWarningMessage('No trajectory available for this item.'); + return; + } + + const sessionId = this.adapter.getSessionIdForToken(treeItem.token); + if (!sessionId) { + vscode.window.showWarningMessage('No trajectory found for this request. Try running the request first.'); + return; + } + + const allTrajectories = this.trajectoryLogger.getAllTrajectories(); + const mainTrajectory = allTrajectories.get(sessionId); + + if (!mainTrajectory) { + vscode.window.showWarningMessage('Trajectory data not found.'); + return; + } + + // Collect the main trajectory and all referenced subagent trajectories + const trajectoriesToExport = this.collectTrajectoryWithSubagents(mainTrajectory, allTrajectories); + + if (trajectoriesToExport.size === 0) { + vscode.window.showWarningMessage('No trajectory data to export.'); + return; + } + + const pathMapping = this.buildTrajectoryPathMapping(trajectoriesToExport); + const isSingleFile = trajectoriesToExport.size === 1; + + let saveDir: vscode.Uri; + let singleFileUri: vscode.Uri | undefined; + + if (isSingleFile) { + // Use showSaveDialog with predetermined filename for single file export + const suggestedFilename = this.getTrajectoryFilename(sessionId, pathMapping); + const saveResult = await vscode.window.showSaveDialog({ + title: 'Export Trajectory', + defaultUri: vscode.Uri.joinPath(vscode.Uri.file(os.homedir()), suggestedFilename), + filters: { 'Trajectory Files': [TRAJECTORY_FILE_EXTENSION.slice(1)] } + }); + + if (!saveResult) { + return; // User cancelled + } + + singleFileUri = saveResult; + saveDir = vscode.Uri.joinPath(saveResult, '..'); + } else { + // Use folder selection for multiple files + const folderUri = await this.promptForFolder('Select Folder to Export Trajectories'); + if (!folderUri) { + return; // User cancelled + } + saveDir = folderUri; + } + + try { + if (isSingleFile && singleFileUri) { + // Export single file using the user-specified path + const [, trajectory] = [...trajectoriesToExport][0]; + const content = JSON.stringify(trajectory, null, 2); + await vscode.workspace.fs.writeFile(singleFileUri, Buffer.from(content, 'utf8')); + } else { + await this.writeTrajectoriesToFolder(trajectoriesToExport, saveDir, pathMapping); + } + + const subagentCount = trajectoriesToExport.size - 1; + const subagentMsg = subagentCount > 0 ? ` (including ${subagentCount} subagent ${subagentCount === 1 ? 'trajectory' : 'trajectories'})` : ''; + const exportPath = isSingleFile && singleFileUri ? singleFileUri.fsPath : saveDir.fsPath; + + const revealAction = 'Reveal in Explorer'; + const result = await vscode.window.showInformationMessage( + `Successfully exported trajectory${subagentMsg} to ${exportPath}`, + revealAction + ); + + if (result === revealAction) { + await vscode.commands.executeCommand('revealFileInOS', singleFileUri ?? saveDir); + } + } catch (error) { + vscode.window.showErrorMessage(`Failed to export trajectory: ${error instanceof Error ? error.message : String(error)}`); + } + } + + /** + * Recursively collect a trajectory and all its referenced subagent trajectories + */ + private collectTrajectoryWithSubagents( + mainTrajectory: IAgentTrajectory, + allTrajectories: Map + ): Map { + const result = new Map(); + const visited = new Set(); + + const collect = (trajectory: IAgentTrajectory) => { + if (visited.has(trajectory.session_id)) { + return; + } + visited.add(trajectory.session_id); + result.set(trajectory.session_id, trajectory); + + // Find subagent references in this trajectory's steps + const steps: ITrajectoryStep[] = Array.isArray(trajectory?.steps) ? trajectory.steps : []; + for (const step of steps) { + const results: IObservationResult[] = Array.isArray(step.observation?.results) ? step.observation.results : []; + for (const r of results) { + for (const ref of r.subagent_trajectory_ref ?? []) { + const subagentTrajectory = allTrajectories.get(ref.session_id); + if (subagentTrajectory) { + collect(subagentTrajectory); + } + } + } + } + }; + + collect(mainTrajectory); + return result; + } + + private sanitizeFilename(name: string): string { + // Remove invalid filename characters + return name.replace(/[<>:"/\\|?*]/g, '_').replace(/\s+/g, '_'); + } +} diff --git a/extensions/copilot/src/platform/requestLogger/common/capturingToken.ts b/extensions/copilot/src/platform/requestLogger/common/capturingToken.ts index 0bd31309039..5742331a24e 100644 --- a/extensions/copilot/src/platform/requestLogger/common/capturingToken.ts +++ b/extensions/copilot/src/platform/requestLogger/common/capturingToken.ts @@ -26,5 +26,22 @@ export class CapturingToken { * excluded from the children list and its content is shown when clicking the parent. */ public readonly promoteMainEntry: boolean = false, + /** + * Optional pre-assigned subAgentInvocationId as session id for trajectory tracking. + * When set, the trajectory will use this ID instead of generating a new one, + * enabling explicit linking between parent tool calls and subagent trajectories. + */ + public readonly subAgentInvocationId?: string, + /** + * Optional name of the subagent being invoked. + * Used alongside subAgentInvocationId to identify the subagent in trajectory tracking. + */ + public readonly subAgentName?: string, + /** + * Optional VS Code chat session ID for trajectory tracking. + * When set, this ID is used directly as the trajectory session ID for the main chat, + * providing a 1:1 mapping between chat sessions and trajectories. + */ + public readonly chatSessionId?: string, ) { } } diff --git a/extensions/copilot/src/platform/requestLogger/node/requestLogger.ts b/extensions/copilot/src/platform/requestLogger/node/requestLogger.ts index 261140217de..1d0ff538b22 100644 --- a/extensions/copilot/src/platform/requestLogger/node/requestLogger.ts +++ b/extensions/copilot/src/platform/requestLogger/node/requestLogger.ts @@ -124,6 +124,7 @@ export interface ILoggedToolCall { token: CapturingToken | undefined; time: number; thinking?: ThinkingData; + toolMetadata?: unknown; toJSON(): Promise; } diff --git a/extensions/copilot/src/platform/requestLogger/test/node/testRequestLogger.ts b/extensions/copilot/src/platform/requestLogger/test/node/testRequestLogger.ts index 846841a883a..e7e272d0053 100644 --- a/extensions/copilot/src/platform/requestLogger/test/node/testRequestLogger.ts +++ b/extensions/copilot/src/platform/requestLogger/test/node/testRequestLogger.ts @@ -171,6 +171,7 @@ class TestLoggedRequestInfo implements ILoggedRequestInfo { class TestLoggedToolCall { public readonly kind = LoggedInfoKind.ToolCall; + public readonly toolMetadata: unknown; constructor( public readonly id: string, @@ -180,7 +181,10 @@ class TestLoggedToolCall { public readonly token: CapturingToken | undefined, public readonly time: number, public readonly thinking?: ThinkingData, - ) { } + ) { + // Extract toolMetadata from response if it exists + this.toolMetadata = 'toolMetadata' in response ? (response as { toolMetadata?: unknown }).toolMetadata : undefined; + } async toJSON(): Promise { return { diff --git a/extensions/copilot/src/platform/trajectory/common/trajectoryLogger.ts b/extensions/copilot/src/platform/trajectory/common/trajectoryLogger.ts new file mode 100644 index 00000000000..d9446d81f4d --- /dev/null +++ b/extensions/copilot/src/platform/trajectory/common/trajectoryLogger.ts @@ -0,0 +1,140 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { createServiceIdentifier } from '../../../util/common/services'; +import { Event } from '../../../util/vs/base/common/event'; +import type { + IAgentInfo, + IAgentTrajectory, + IObservationResult, + IStepMetrics, + ISubagentTrajectoryRef, + IToolCall, + ITrajectoryStep +} from './trajectoryTypes'; + +export type { + IAgentInfo, + IAgentTrajectory, + IObservationResult, + IStepMetrics, + ISubagentTrajectoryRef, + IToolCall, + ITrajectoryStep +}; + +/** + * Service for building and managing agent trajectories. + * This service tracks agent execution steps, tool calls, and observations + * to construct a complete trajectory that can be exported and analyzed. + */ +export const ITrajectoryLogger = createServiceIdentifier('ITrajectoryLogger'); + +export interface ITrajectoryLogger { + readonly _serviceBrand: undefined; + + /** + * Event fired when the trajectory is updated + */ + readonly onDidUpdateTrajectory: Event; + + /** + * Start a new trajectory session + * @param sessionId Unique identifier for this session + * @param agentInfo Agent configuration information + */ + startTrajectory(sessionId: string, agentInfo: IAgentInfo): void; + + /** + * Add a system message step (e.g., initial system prompt) + * @param message The system message content + * @param timestamp Optional ISO 8601 timestamp + */ + addSystemStep(message: string, timestamp?: string): void; + + /** + * Add a user message step + * @param message The user's message content + * @param timestamp Optional ISO 8601 timestamp + */ + addUserStep(message: string, timestamp?: string): void; + + /** + * Begin an agent step (LLM inference) + * @param message The agent's response message + * @param modelName The model used for this step + * @param reasoningContent Optional internal reasoning content + * @param timestamp Optional ISO 8601 timestamp + * @returns A step context for adding tool calls and observations + */ + beginAgentStep( + message: string, + modelName?: string, + reasoningContent?: string, + timestamp?: string + ): IAgentStepContext; + + /** + * Get the complete trajectory for the current session + * @returns The complete trajectory or undefined if no session is active + */ + getTrajectory(): IAgentTrajectory | undefined; + + /** + * Get all trajectories (main and subagent) for export + * @returns Map of session IDs to trajectories + */ + getAllTrajectories(): Map; + + /** + * Clear the current trajectory session + */ + clearTrajectory(): void; + + /** + * Check if a trajectory is currently being tracked + */ + hasActiveTrajectory(): boolean; + + /** + * Get the current session ID + */ + getCurrentSessionId(): string | undefined; +} + +/** + * Context for building an agent step with tool calls and observations + */ +export interface IAgentStepContext { + /** + * Add tool calls to this step + * @param toolCalls Array of tool calls made by the agent + */ + addToolCalls(toolCalls: IToolCall[]): void; + + /** + * Add observations (tool results) to this step + * @param results Array of observation results + */ + addObservation(results: IObservationResult[]): void; + + /** + * Add a subagent trajectory reference + * @param toolCallId The tool call ID that spawned the subagent + * @param subagentRef Reference to the subagent's trajectory + */ + addSubagentReference(toolCallId: string, subagentRef: ISubagentTrajectoryRef): void; + + /** + * Set metrics for this step + * @param metrics Step metrics including token counts and costs + */ + setMetrics(metrics: IStepMetrics): void; + + /** + * Finalize this step and add it to the trajectory + */ + complete(): void; +} diff --git a/extensions/copilot/src/platform/trajectory/common/trajectoryTypes.ts b/extensions/copilot/src/platform/trajectory/common/trajectoryTypes.ts new file mode 100644 index 00000000000..26efea47892 --- /dev/null +++ b/extensions/copilot/src/platform/trajectory/common/trajectoryTypes.ts @@ -0,0 +1,194 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +/** + * Agent Trajectory Format + * + * This format is inspired by the Harbor ATIF (Agent Trajectory Interchange Format) + * specification, adapted for VS Code Copilot Chat's multi-agent, parallel tool calling, + * and MCP integration capabilities. + * + * Key design goals: + * - Capture hierarchical agent/subagent relationships + * - Support parallel tool calling with proper correlation + * - Include MCP tool calls with server context + * - Enable trajectory replay and visualization + * - Maintain compatibility with existing logging infrastructure + */ + +/** + * Root trajectory object representing a complete agent session + */ +export interface IAgentTrajectory { + /** Schema version for compatibility tracking */ + readonly schema_version: string; + /** Unique identifier for this trajectory session */ + readonly session_id: string; + /** Agent configuration and metadata */ + readonly agent: IAgentInfo; + /** Sequential steps in the agent's execution */ + readonly steps: ITrajectoryStep[]; + /** Optional aggregate metrics for the entire trajectory */ + readonly final_metrics?: IFinalMetrics; + /** Optional notes or metadata about this trajectory */ + readonly notes?: string; + /** Optional reference to continuation trajectory if split */ + readonly continued_trajectory_ref?: string; + /** Custom metadata not covered by core schema */ + readonly extra?: Record; +} + +/** + * Agent information and configuration + */ +export interface IAgentInfo { + /** Name of the agent system (e.g., "copilot-agent", "search-subagent") */ + readonly name: string; + /** Version of the agent system */ + readonly version: string; + /** Default model used for this trajectory */ + readonly model_name?: string; + /** Tool definitions available to the agent (OpenAI function calling format) */ + readonly tool_definitions?: IToolDefinition[]; + /** Custom agent configuration */ + readonly extra?: Record; +} + +/** + * Tool definition in OpenAI function calling format + */ +export interface IToolDefinition { + readonly type: 'function'; + readonly function: { + readonly name: string; + readonly description: string; + readonly parameters?: Record; + }; +} + +/** + * A single step in the trajectory, representing one interaction turn + */ +export interface ITrajectoryStep { + /** Ordinal step number starting from 1 */ + readonly step_id: number; + /** ISO 8601 timestamp for this step */ + readonly timestamp?: string; + /** Source of this step: system, user, or agent */ + readonly source: 'system' | 'user' | 'agent'; + /** Model used for this step (only for agent steps) */ + readonly model_name?: string; + /** The message content for this step */ + readonly message: string; + /** Agent's internal reasoning (only for agent steps) */ + readonly reasoning_content?: string; + /** Tool calls made in this step (only for agent steps) */ + readonly tool_calls?: IToolCall[]; + /** Results from tool calls or system events */ + readonly observation?: IObservation; + /** Metrics for this step (only for agent steps) */ + readonly metrics?: IStepMetrics; + /** Custom step metadata */ + readonly extra?: Record; +} + +/** + * A tool call made by the agent + */ +export interface IToolCall { + /** Unique identifier for this tool call */ + readonly tool_call_id: string; + /** Name of the function/tool being called */ + readonly function_name: string; + /** Arguments passed to the tool (structured object) */ + readonly arguments: Record; +} + +/** + * Observation containing results from tool calls or system events + */ +export interface IObservation { + /** Array of results from tool executions */ + readonly results: IObservationResult[]; +} + +/** + * Result from a single tool call or system event + */ +export interface IObservationResult { + /** The tool_call_id this result corresponds to (null for non-tool events) */ + readonly source_call_id?: string; + /** Textual content of the result */ + readonly content?: string; + /** Reference to subagent trajectory if this was a subagent invocation */ + readonly subagent_trajectory_ref?: ISubagentTrajectoryRef[]; +} + +/** + * Reference to a subagent trajectory + */ +export interface ISubagentTrajectoryRef { + /** Session ID of the subagent */ + readonly session_id: string; + /** Optional path to the subagent's trajectory file */ + readonly trajectory_path?: string; + /** Summary or metadata about the subagent execution */ + readonly extra?: Record; +} + +/** + * Metrics for a single step + */ +export interface IStepMetrics { + /** Total input tokens (including cached) */ + readonly prompt_tokens?: number; + /** Tokens generated by the model */ + readonly completion_tokens?: number; + /** Cached input tokens (subset of prompt_tokens) */ + readonly cached_tokens?: number; + /** Cost in USD for this step */ + readonly cost_usd?: number; + /** Time to first token in milliseconds */ + readonly time_to_first_token_ms?: number; + /** Total duration of the step in milliseconds */ + readonly duration_ms?: number; + /** Provider-specific metrics */ + readonly extra?: Record; +} + +/** + * Aggregate metrics for the entire trajectory + */ +export interface IFinalMetrics { + /** Total prompt tokens across all steps */ + readonly total_prompt_tokens?: number; + /** Total completion tokens across all steps */ + readonly total_completion_tokens?: number; + /** Total cached tokens across all steps */ + readonly total_cached_tokens?: number; + /** Total cost in USD */ + readonly total_cost_usd?: number; + /** Total number of steps */ + readonly total_steps?: number; + /** Total number of tool calls across all steps */ + readonly total_tool_calls?: number; + /** Custom aggregate metrics */ + readonly extra?: Record; +} + +/** + * Current trajectory format version + */ +export const TRAJECTORY_SCHEMA_VERSION = 'ATIF-v1.5'; + +/** + * File extension for trajectory files + */ +export const TRAJECTORY_FILE_EXTENSION = '.trajectory.json'; + +/** + * File extension for a trajectory bundle file (contains multiple trajectories). + */ +export const TRAJECTORY_BUNDLE_FILE_EXTENSION = '.trajectory.bundle.json'; diff --git a/extensions/copilot/src/platform/trajectory/node/trajectoryLogger.ts b/extensions/copilot/src/platform/trajectory/node/trajectoryLogger.ts new file mode 100644 index 00000000000..0af4d6d7391 --- /dev/null +++ b/extensions/copilot/src/platform/trajectory/node/trajectoryLogger.ts @@ -0,0 +1,339 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { Emitter } from '../../../util/vs/base/common/event'; +import { Disposable } from '../../../util/vs/base/common/lifecycle'; +import type { + IAgentInfo, + IAgentStepContext, + IAgentTrajectory, + IObservationResult, + IStepMetrics, + ISubagentTrajectoryRef, + IToolCall, + ITrajectoryLogger, + ITrajectoryStep +} from '../common/trajectoryLogger'; +import { TRAJECTORY_SCHEMA_VERSION } from '../common/trajectoryTypes'; + +/** + * Concrete implementation of the trajectory logger + */ +export class TrajectoryLogger extends Disposable implements ITrajectoryLogger { + declare readonly _serviceBrand: undefined; + + private readonly trajectories = new Map(); + private currentSessionId: string | undefined; + private subagentTrajectories = new Map(); + + private readonly _onDidUpdateTrajectory = this._register(new Emitter()); + public readonly onDidUpdateTrajectory = this._onDidUpdateTrajectory.event; + + public startTrajectory(sessionId: string, agentInfo: IAgentInfo): void { + let builder = this.trajectories.get(sessionId); + if (!builder) { + builder = new TrajectoryBuilder(sessionId, agentInfo); + this.trajectories.set(sessionId, builder); + } else { + builder.updateAgentInfo(agentInfo); + } + this.currentSessionId = sessionId; + this._onDidUpdateTrajectory.fire(); + } + + private getCurrentTrajectoryBuilder(): TrajectoryBuilder | undefined { + if (!this.currentSessionId) { + return undefined; + } + return this.trajectories.get(this.currentSessionId); + } + + public addSystemStep(message: string, timestamp?: string): void { + const current = this.getCurrentTrajectoryBuilder(); + if (!current) { + return; + } + current.addSystemStep(message, timestamp); + this._onDidUpdateTrajectory.fire(); + } + + public addUserStep(message: string, timestamp?: string): void { + const current = this.getCurrentTrajectoryBuilder(); + if (!current) { + return; + } + current.addUserStep(message, timestamp); + this._onDidUpdateTrajectory.fire(); + } + + public beginAgentStep( + message: string, + modelName?: string, + reasoningContent?: string, + timestamp?: string + ): IAgentStepContext { + const current = this.getCurrentTrajectoryBuilder(); + if (!current) { + throw new Error('No active trajectory. Call startTrajectory first.'); + } + const context = current.beginAgentStep(message, modelName, reasoningContent, timestamp); + return { + addToolCalls: (toolCalls) => context.addToolCalls(toolCalls), + addObservation: (results) => context.addObservation(results), + addSubagentReference: (toolCallId, subagentRef) => context.addSubagentReference(toolCallId, subagentRef), + setMetrics: (metrics) => context.setMetrics(metrics), + complete: () => { + context.complete(); + this._onDidUpdateTrajectory.fire(); + } + }; + } + + public getTrajectory(): IAgentTrajectory | undefined { + return this.getCurrentTrajectoryBuilder()?.build(); + } + + public getAllTrajectories(): Map { + const trajectories = new Map(); + for (const builder of this.trajectories.values()) { + const trajectory = builder.build(); + trajectories.set(trajectory.session_id, trajectory); + } + for (const [sessionId, trajectory] of this.subagentTrajectories) { + trajectories.set(sessionId, trajectory); + } + return trajectories; + } + + public clearTrajectory(): void { + this.trajectories.clear(); + this.currentSessionId = undefined; + this.subagentTrajectories.clear(); + this._onDidUpdateTrajectory.fire(); + } + + public hasActiveTrajectory(): boolean { + return this.currentSessionId !== undefined; + } + + public getCurrentSessionId(): string | undefined { + return this.currentSessionId; + } + + /** + * Register a subagent trajectory + * @internal Used by subagent implementations + */ + public registerSubagentTrajectory(trajectory: IAgentTrajectory): void { + this.subagentTrajectories.set(trajectory.session_id, trajectory); + this._onDidUpdateTrajectory.fire(); + } +} + +/** + * Builder for constructing a trajectory incrementally + */ +class TrajectoryBuilder { + private steps: ITrajectoryStep[] = []; + private stepCounter = 0; + + constructor( + private readonly sessionId: string, + private agentInfo: IAgentInfo + ) { } + + public updateAgentInfo(agentInfo: IAgentInfo): void { + this.agentInfo = { + ...this.agentInfo, + ...agentInfo, + tool_definitions: agentInfo.tool_definitions ?? this.agentInfo.tool_definitions + }; + } + + public getSessionId(): string { + return this.sessionId; + } + + public addSystemStep(message: string, timestamp?: string): void { + this.steps.push({ + step_id: ++this.stepCounter, + timestamp: timestamp || new Date().toISOString(), + source: 'system', + message + }); + } + + public addUserStep(message: string, timestamp?: string): void { + this.steps.push({ + step_id: ++this.stepCounter, + timestamp: timestamp || new Date().toISOString(), + source: 'user', + message + }); + } + + public beginAgentStep( + message: string, + modelName?: string, + reasoningContent?: string, + timestamp?: string + ): IAgentStepContext { + const stepId = ++this.stepCounter; + const stepTimestamp = timestamp || new Date().toISOString(); + + const step: Partial = { + step_id: stepId, + timestamp: stepTimestamp, + source: 'agent', + message, + model_name: modelName, + reasoning_content: reasoningContent + }; + + return new AgentStepContext(step, (completedStep) => { + this.steps.push(completedStep as ITrajectoryStep); + }); + } + + public build(): IAgentTrajectory { + // Infer a default model name for the trajectory if not provided at start. + // ATIF allows a root-level agent.model_name which step-level model_name can override. + let inferredModelName: string | undefined; + if (!this.agentInfo.model_name) { + for (const step of this.steps) { + if (step.source === 'agent' && step.model_name) { + inferredModelName = step.model_name; + break; + } + } + } + + // Calculate final metrics (ATIF v1.5): only include fields that actually + // appeared in per-step metrics. The final_metrics object itself is optional. + let hasAnyStepMetrics = false; + let sawPromptTokens = false; + let sawCompletionTokens = false; + let sawCachedTokens = false; + let sawCostUsd = false; + + let totalPromptTokens = 0; + let totalCompletionTokens = 0; + let totalCachedTokens = 0; + let totalCostUsd = 0; + let totalToolCalls = 0; + + for (const step of this.steps) { + const metrics = step.metrics; + if (!metrics) { + continue; + } + hasAnyStepMetrics = true; + if (metrics.prompt_tokens !== undefined) { + sawPromptTokens = true; + totalPromptTokens += metrics.prompt_tokens; + } + if (metrics.completion_tokens !== undefined) { + sawCompletionTokens = true; + totalCompletionTokens += metrics.completion_tokens; + } + if (metrics.cached_tokens !== undefined) { + sawCachedTokens = true; + totalCachedTokens += metrics.cached_tokens; + } + if (metrics.cost_usd !== undefined) { + sawCostUsd = true; + totalCostUsd += metrics.cost_usd; + } + } + + // Count total tool calls across all steps + for (const step of this.steps) { + if (step.tool_calls) { + totalToolCalls += step.tool_calls.length; + } + } + + const finalMetrics = hasAnyStepMetrics || totalToolCalls > 0 ? { + ...(sawPromptTokens ? { total_prompt_tokens: totalPromptTokens } : {}), + ...(sawCompletionTokens ? { total_completion_tokens: totalCompletionTokens } : {}), + ...(sawCachedTokens ? { total_cached_tokens: totalCachedTokens } : {}), + ...(sawCostUsd ? { total_cost_usd: totalCostUsd } : {}), + total_steps: this.steps.length, + ...(totalToolCalls > 0 ? { total_tool_calls: totalToolCalls } : {}) + } : undefined; + + const agent = inferredModelName ? { ...this.agentInfo, model_name: inferredModelName } : this.agentInfo; + + return { + schema_version: TRAJECTORY_SCHEMA_VERSION, + session_id: this.sessionId, + agent, + steps: [...this.steps], + final_metrics: finalMetrics + }; + } +} + +/** + * Context for building an agent step + */ +class AgentStepContext implements IAgentStepContext { + private toolCalls: IToolCall[] = []; + private observationResults: IObservationResult[] = []; + private metrics: IStepMetrics | undefined; + + constructor( + private readonly step: Partial, + private readonly onComplete: (step: Partial) => void + ) { } + + public addToolCalls(toolCalls: IToolCall[]): void { + this.toolCalls.push(...toolCalls); + } + + public addObservation(results: IObservationResult[]): void { + this.observationResults.push(...results); + } + + public addSubagentReference(toolCallId: string, subagentRef: ISubagentTrajectoryRef): void { + // Find or create observation result for this tool call + let result = this.observationResults.find(r => r.source_call_id === toolCallId); + if (!result) { + result = { source_call_id: toolCallId }; + this.observationResults.push(result); + } + + // Add subagent reference + const mutableResult = result as { subagent_trajectory_ref?: ISubagentTrajectoryRef[] }; + if (!mutableResult.subagent_trajectory_ref) { + mutableResult.subagent_trajectory_ref = []; + } + mutableResult.subagent_trajectory_ref.push(subagentRef); + } + + public setMetrics(metrics: IStepMetrics): void { + this.metrics = metrics; + } + + public complete(): void { + // Finalize the step (cast to mutable for assignment) + const mutableStep = this.step as { + tool_calls?: IToolCall[]; + observation?: { results: IObservationResult[] }; + metrics?: IStepMetrics; + }; + if (this.toolCalls.length > 0) { + mutableStep.tool_calls = this.toolCalls; + } + if (this.observationResults.length > 0) { + mutableStep.observation = { results: this.observationResults }; + } + if (this.metrics) { + mutableStep.metrics = this.metrics; + } + + this.onComplete(this.step); + } +} diff --git a/extensions/copilot/src/platform/trajectory/node/trajectoryLoggerAdapter.ts b/extensions/copilot/src/platform/trajectory/node/trajectoryLoggerAdapter.ts new file mode 100644 index 00000000000..feaa13abaee --- /dev/null +++ b/extensions/copilot/src/platform/trajectory/node/trajectoryLoggerAdapter.ts @@ -0,0 +1,563 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { Raw } from '@vscode/prompt-tsx'; +import { Disposable } from '../../../util/vs/base/common/lifecycle'; +import { LanguageModelDataPart, LanguageModelPromptTsxPart, LanguageModelTextPart } from '../../../vscodeTypes'; +import { ConfigKey, IConfigurationService } from '../../configuration/common/configurationService'; +import { CapturingToken } from '../../requestLogger/common/capturingToken'; +import { ILoggedToolCall, IRequestLogger, LoggedInfo, LoggedInfoKind, LoggedRequest, LoggedRequestKind } from '../../requestLogger/node/requestLogger'; +import { IAgentStepContext, type IObservationResult, type IStepMetrics, type IToolCall, ITrajectoryLogger } from '../common/trajectoryLogger'; +import { IToolDefinition, TRAJECTORY_FILE_EXTENSION } from '../common/trajectoryTypes'; + +const AGENT_NAME = 'GitHub Copilot Chat'; + +/** + * Function type for rendering PromptTsx parts to strings. + * Injected from extension layer to avoid layering violations. + */ +export type PromptTsxRenderer = (part: LanguageModelPromptTsxPart) => Promise; + +/** + * Adapter that converts request logger entries to trajectory format. + * This is a bridge between the existing logging system and the new trajectory format. + */ +export class TrajectoryLoggerAdapter extends Disposable { + private sessionMap = new WeakMap(); + // Also maintain a map for lookup by token reference without preventing GC of tokens + private tokenToSessionId = new WeakMap(); + private processedEntries = new Set(); + private processedToolCalls = new Set(); // Track processed tool calls by their ID + private lastUserMessageBySession = new Map(); + // Track pending step contexts by both session and request ID to handle parallel tool calls + private pendingStepContexts = new Map(); + private requestToStepContext = new Map(); + private runSubagentToolCallToSessionId = new Map(); + + constructor( + private readonly requestLogger: IRequestLogger, + private readonly trajectoryLogger: ITrajectoryLogger, + private readonly configService: IConfigurationService, + private readonly promptTsxRenderer?: PromptTsxRenderer + ) { + super(); + // Subscribe to request logger updates + this._register(this.requestLogger.onDidChangeRequests(() => { + this.syncTrajectories(); + })); + } + + /** + * Synchronize trajectories with request logger state + */ + private async syncTrajectories(): Promise { + // Safety valve: prevent unbounded growth of tracking sets + // This handles the case where RequestLogger evicts entries but our sets retain orphaned IDs + // Use the same max entries setting as RequestLogger for consistency + const maxEntries = this.configService.getConfig(ConfigKey.Advanced.RequestLoggerMaxEntries); + if (this.processedEntries.size > maxEntries) { + const entries = [...this.processedEntries]; + this.processedEntries = new Set(entries.slice(-maxEntries / 2)); + } + if (this.processedToolCalls.size > maxEntries) { + const toolCalls = [...this.processedToolCalls]; + this.processedToolCalls = new Set(toolCalls.slice(-maxEntries / 2)); + } + + const requests = this.requestLogger.getRequests(); + + for (const entry of requests) { + // Skip already processed entries + if (this.processedEntries.has(entry.id)) { + continue; + } + + // Only process entries with capturing tokens (grouped requests) + if (!entry.token) { + continue; + } + + // Get or create session for this token + let sessionId = this.sessionMap.get(entry.token); + // Use subAgentName as agent name for subagent trajectories + const agentName = entry.token.subAgentName ?? AGENT_NAME; + + if (!sessionId) { + // Use the following priority for session ID: + // 1. subAgentInvocationId for explicit subagent linking + // 2. chatSessionId for main chat sessions (provides 1:1 mapping with VS Code chat) + // 3. Generate a new one as fallback + sessionId = entry.token.subAgentInvocationId ?? entry.token.chatSessionId ?? this.generateSessionId(entry.token.label); + this.sessionMap.set(entry.token, sessionId); + this.tokenToSessionId.set(entry.token, sessionId); + } + + // Start or switch to the trajectory for this session. This ensures the correct + // trajectory is active before processing the entry, and updates agent info if needed. + this.trajectoryLogger.startTrajectory(sessionId, { + name: agentName, + version: '1.0.0', + tool_definitions: this.extractToolDefinitionsFromEntry(entry) + }); + + await this.processEntry(entry, sessionId); + this.processedEntries.add(entry.id); + } + } + + private async processEntry(entry: LoggedInfo, sessionId: string): Promise { + switch (entry.kind) { + case LoggedInfoKind.Request: + await this.processRequestInfo(entry, sessionId); + break; + case LoggedInfoKind.ToolCall: + await this.processToolCall(entry, sessionId); + break; + case LoggedInfoKind.Element: + // Elements are debug info, not relevant for trajectories + break; + } + } + + private async processRequestInfo(entry: LoggedInfo, sessionId: string): Promise { + if (entry.kind !== LoggedInfoKind.Request) { + return; + } + + const loggedRequest = entry.entry; + + // Skip non-conversation requests + if (loggedRequest.isConversationRequest === false) { + return; + } + + // Handle different request types + if (loggedRequest.type === LoggedRequestKind.ChatMLSuccess) { + await this.processSuccessfulRequest(loggedRequest, sessionId); + } + } + + private async processSuccessfulRequest(request: LoggedRequest & { type: LoggedRequestKind.ChatMLSuccess }, sessionId: string): Promise { + this.maybeAddUserStepFromRequest(request, sessionId); + + const message = Array.isArray(request.result.value) + ? request.result.value.join('\n') + : String(request.result.value); + + const modelName = request.chatEndpoint.model; + + // Extract reasoning content from deltas if available + let reasoningContent: string | undefined; + if (request.deltas) { + const thinkingDeltas = request.deltas.filter(d => d.thinking); + if (thinkingDeltas.length > 0) { + reasoningContent = thinkingDeltas.map(d => d.thinking?.text || '').join('\n'); + } + } + + const stepContext = this.trajectoryLogger.beginAgentStep( + message, + modelName, + reasoningContent, + request.startTime.toISOString() + ); + + // Add metrics + if (request.usage) { + const metrics: IStepMetrics = { + prompt_tokens: request.usage.prompt_tokens, + completion_tokens: request.usage.completion_tokens, + cached_tokens: request.usage.prompt_tokens_details?.cached_tokens, + time_to_first_token_ms: request.timeToFirstToken, + duration_ms: request.endTime.getTime() - request.startTime.getTime() + }; + stepContext.setMetrics(metrics); + } + + // Count expected tool calls from this request + let toolCallCount = 0; + if (request.deltas) { + for (const delta of request.deltas) { + if (delta.copilotToolCalls) { + toolCallCount += delta.copilotToolCalls.length; + } + } + } + + // If no tool calls expected, complete immediately + if (toolCallCount === 0) { + stepContext.complete(); + } else { + // Store context with request ID for tool calls to attach to + // Use a unique request ID based on startTime + const requestId = `${sessionId}-${request.startTime.getTime()}`; + this.requestToStepContext.set(requestId, { + sessionId, + context: stepContext, + toolCallCount, + processedToolCalls: 0 + }); + // Also store in pendingStepContexts for backwards compatibility + this.pendingStepContexts.set(sessionId, stepContext); + } + } + + private maybeAddUserStepFromRequest(request: LoggedRequest & { startTime: Date }, sessionId: string): void { + const messages = this.getChatMessagesFromRequest(request); + if (!Array.isArray(messages) || messages.length === 0) { + return; + } + + const lastUser = this.getLastUserMessageText(messages); + if (!lastUser) { + return; + } + + const lastKey = this.lastUserMessageBySession.get(sessionId); + const key = this.simpleHash(lastUser) + ':' + lastUser.length; + if (lastKey === key) { + return; + } + + this.lastUserMessageBySession.set(sessionId, key); + const timestamp = request.startTime.toISOString(); + this.trajectoryLogger.addUserStep(lastUser, timestamp); + } + + private getChatMessagesFromRequest(request: LoggedRequest): Raw.ChatMessage[] | undefined { + const messages = (request as unknown as { chatParams?: { messages?: unknown } }).chatParams?.messages; + if (!Array.isArray(messages)) { + return undefined; + } + return messages as Raw.ChatMessage[]; + } + + private getLastUserMessageText(messages: Raw.ChatMessage[]): string | undefined { + for (let i = messages.length - 1; i >= 0; i--) { + const m = messages[i]; + if (m.role !== Raw.ChatRole.User) { + continue; + } + + const content = (m as unknown as { content?: unknown }).content; + if (typeof content === 'string') { + return content.trim() || undefined; + } + + if (!Array.isArray(content)) { + return undefined; + } + + const text = content + .map(part => { + const partType = (part as { type?: unknown }).type; + if (partType === Raw.ChatCompletionContentPartKind.Text) { + const t = (part as { text?: unknown }).text; + return typeof t === 'string' ? t : undefined; + } + return undefined; + }) + .filter((t): t is string => typeof t === 'string' && t.length > 0) + .join('\n') + .trim(); + + return text || undefined; + } + return undefined; + } + + private async processToolCall(entry: ILoggedToolCall, sessionId: string): Promise { + // Skip already processed tool calls (prevents duplicates from multiple event fires) + if (this.processedToolCalls.has(entry.id)) { + return; + } + this.processedToolCalls.add(entry.id); + + // Find the step context for this tool call + // Try to find by iterating request contexts for this session + let stepInfo: { sessionId: string; context: IAgentStepContext; toolCallCount: number; processedToolCalls: number } | undefined; + let requestKey: string | undefined; + + for (const [key, info] of this.requestToStepContext) { + if (info.sessionId === sessionId) { + stepInfo = info; + requestKey = key; + break; + } + } + + let stepContext: IAgentStepContext; + let shouldComplete = true; + + if (stepInfo) { + stepContext = stepInfo.context; + stepInfo.processedToolCalls++; + // Only complete after all tool calls are processed + shouldComplete = stepInfo.processedToolCalls >= stepInfo.toolCallCount; + } else { + // No pending context - create a new step for this orphan tool call + stepContext = this.trajectoryLogger.beginAgentStep('', undefined, + entry.thinking?.text ? (Array.isArray(entry.thinking.text) ? entry.thinking.text.join('\n') : entry.thinking.text) : undefined); + } + + // Parse tool call + const toolCall: IToolCall = { + tool_call_id: entry.id, + function_name: entry.name, + arguments: this.parseArguments(entry.args) + }; + + stepContext.addToolCalls([toolCall]); + + // Extract observation result + const content = await this.extractToolResultContent(entry); + const observationResult: IObservationResult = { + source_call_id: entry.id, + content + }; + + // Add observation first so subagent references merge into the same result. + stepContext.addObservation([observationResult]); + + // Check if this is a subagent tool call (runSubagent or search_subagent) + if (entry.name === 'runSubagent' || entry.name === 'search_subagent') { + const resolvedSubagentSessionId = this.resolveSubagentSessionIdForSubagentTool(entry); + if (resolvedSubagentSessionId) { + const subagentDescription = this.extractSubagentDescription(entry); + const trajectoryPath = `${this.sanitizeSubagentDescriptionForFilename(subagentDescription)}-${resolvedSubagentSessionId}${TRAJECTORY_FILE_EXTENSION}`; + stepContext.addSubagentReference(entry.id, { + session_id: resolvedSubagentSessionId, + trajectory_path: trajectoryPath + }); + } + } + + // Only complete when all tool calls from this request are processed + if (shouldComplete) { + stepContext.complete(); + if (requestKey) { + this.requestToStepContext.delete(requestKey); + } + this.pendingStepContexts.delete(sessionId); + } + } + + private resolveSubagentSessionIdForSubagentTool(entry: ILoggedToolCall): string | undefined { + const cached = this.runSubagentToolCallToSessionId.get(entry.id); + if (cached) { + return cached; + } + + // Use subAgentInvocationId from toolMetadata (set at capture time) + // This is the principled approach: the subagent session ID is captured when + // the subagent is launched and attached to the tool result metadata + const metadata = entry.toolMetadata as { subAgentInvocationId?: string } | undefined; + if (metadata?.subAgentInvocationId) { + this.runSubagentToolCallToSessionId.set(entry.id, metadata.subAgentInvocationId); + return metadata.subAgentInvocationId; + } + + return undefined; + } + + private extractSubagentDescription(entry: ILoggedToolCall): string { + const metadata = entry.toolMetadata as { agentName?: unknown; description?: unknown } | undefined; + + // Prefer agentName for consistent naming + if (metadata && typeof metadata.agentName === 'string' && metadata.agentName.trim().length > 0) { + return metadata.agentName; + } + + if (metadata && typeof metadata.description === 'string' && metadata.description.trim().length > 0) { + return metadata.description; + } + + const args = this.parseArguments(entry.args); + const description = args.description; + if (typeof description === 'string' && description.trim().length > 0) { + return description; + } + + return entry.name; + } + + private sanitizeSubagentDescriptionForFilename(description: string): string { + // Keep filenames stable and readable across platforms. + const sanitized = description + .toLowerCase() + .replace(/[^a-z0-9]+/g, '-') + .replace(/^-+|-+$/g, ''); + + return (sanitized || 'subagent').substring(0, 50); + } + + private parseArguments(args: unknown): Record { + if (typeof args === 'string') { + try { + return JSON.parse(args) as Record; + } catch { + return { raw: args }; + } + } + if (typeof args === 'object' && args !== null) { + return args as Record; + } + return {}; + } + + private async extractToolResultContent(entry: ILoggedToolCall): Promise { + const parts: string[] = []; + + for (const content of entry.response.content) { + if (content instanceof LanguageModelTextPart) { + parts.push(content.value); + } else if (content instanceof LanguageModelDataPart) { + parts.push(this.renderDataPartToString(content)); + } else if (content instanceof LanguageModelPromptTsxPart) { + parts.push(await this.renderPromptTsxPartToStringNoBudget(content)); + } + } + + return parts.join('\n'); + } + + private async renderPromptTsxPartToStringNoBudget(part: LanguageModelPromptTsxPart): Promise { + if (this.promptTsxRenderer) { + try { + return await this.promptTsxRenderer(part); + } catch { + // Fall through to fallback + } + } + // Fallback: serialize the value as JSON + try { + return JSON.stringify(part.value, null, 2); + } catch { + return String(part.value); + } + } + + private renderDataPartToString(part: LanguageModelDataPart): string { + const mimeType = typeof part.mimeType === 'string' ? part.mimeType : ''; + const isImage = mimeType.startsWith('image/'); + + if (isImage) { + const base64 = Buffer.from(part.data).toString('base64'); + return `data:${mimeType};base64,${base64}`; + } + + try { + return new TextDecoder().decode(part.data); + } catch { + return ``; + } + } + + private generateSessionId(label: string): string { + // Create a short hash from the label for uniqueness + const hash = this.simpleHash(label); + // Truncate and sanitize the label to create a readable prefix + const sanitized = label.toLowerCase() + .replace(/[^a-z0-9]+/g, '-') + .replace(/^-+|-+$/g, '') + .substring(0, 30); // Limit to 30 chars for readability + return `${sanitized}-${hash}-${Date.now()}`; + } + + private simpleHash(str: string): string { + let hash = 0; + for (let i = 0; i < str.length; i++) { + const char = str.charCodeAt(i); + hash = ((hash << 5) - hash) + char; + hash = hash & hash; // Convert to 32bit integer + } + return Math.abs(hash).toString(36); + } + + private extractToolDefinitionsFromEntry(entry: LoggedInfo): IToolDefinition[] | undefined { + if (entry.kind !== LoggedInfoKind.Request) { + return undefined; + } + + const request = entry.entry; + if (!('chatParams' in request)) { + return undefined; + } + + const tools = request.chatParams.body?.tools as Array> | undefined; + if (!Array.isArray(tools) || tools.length === 0) { + return undefined; + } + + const definitions: IToolDefinition[] = []; + for (const tool of tools) { + const type = tool.type; + if (type !== 'function') { + continue; + } + + const fn = tool.function as { name?: unknown; description?: unknown; parameters?: Record } | undefined; + if (!fn || typeof fn.name !== 'string') { + continue; + } + + definitions.push({ + type: 'function', + function: { + name: fn.name, + description: typeof fn.description === 'string' ? fn.description : '', + parameters: fn.parameters + } + }); + } + + return definitions.length > 0 ? definitions : undefined; + } + + /** + * Get the session ID associated with a capturing token + * @param token The capturing token to look up + * @returns The session ID or undefined if not found + */ + public getSessionIdForToken(token: CapturingToken): string | undefined { + return this.tokenToSessionId.get(token); + } + + /** + * Clear adapter state for a specific session or all sessions. + * This should be called when trajectories are cleared to prevent memory leaks + * from orphaned tracking data. + * @param sessionId Optional session ID to clear. If omitted, clears all state. + */ + public clearSessionState(sessionId?: string): void { + if (sessionId) { + // Clear session-specific data + this.lastUserMessageBySession.delete(sessionId); + this.pendingStepContexts.delete(sessionId); + + // Clear requestToStepContext entries for this session + for (const [key, info] of this.requestToStepContext) { + if (info.sessionId === sessionId) { + this.requestToStepContext.delete(key); + } + } + + // Clear runSubagentToolCallToSessionId entries pointing to this session + for (const [toolCallId, mappedSessionId] of this.runSubagentToolCallToSessionId) { + if (mappedSessionId === sessionId) { + this.runSubagentToolCallToSessionId.delete(toolCallId); + } + } + } else { + // Clear all state + this.processedEntries.clear(); + this.processedToolCalls.clear(); + this.lastUserMessageBySession.clear(); + this.pendingStepContexts.clear(); + this.requestToStepContext.clear(); + this.runSubagentToolCallToSessionId.clear(); + } + } +}