mirror of
https://github.com/microsoft/vscode.git
synced 2026-05-21 23:59:34 +01:00
feat: support Harbor/ATIF trajectory capture and export (#2893)
* Initial plan * Add trajectory format types and core implementation Co-authored-by: zhichli <57812115+zhichli@users.noreply.github.com> * Add unit tests for trajectory logger and fix implementation Co-authored-by: zhichli <57812115+zhichli@users.noreply.github.com> * Add comprehensive documentation for trajectory format Co-authored-by: zhichli <57812115+zhichli@users.noreply.github.com> * Add implementation status and next steps documentation Co-authored-by: zhichli <57812115+zhichli@users.noreply.github.com> * Add trajectory implementation quick reference summary * Add trajectory export commands and enhance trajectory logging functionality * Refactor trajectory metrics calculation and update schema version to ATIF v1.5 * trajectory scaffolding * Enhance trajectory tracking by adding subAgentName and agentName to tool metadata in SearchSubagentTool and TrajectoryLoggerAdapter * add export trajectory cmd for tree nodes * use sessionId for main trajectory * Update command categories from 'Copilot' to 'Chat' in package.json * Update trajectory schema version to ATIF-v1.5 and enhance error handling in export commands * Remove obsolete trajectory documentation and integration files * Refactor trajectory export commands and update README for clarity * Update src/platform/trajectory/common/trajectoryLogger.ts Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update src/platform/trajectory/node/trajectoryLogger.ts Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update src/platform/trajectory/node/trajectoryLoggerAdapter.ts Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update src/extension/trajectory/vscode-node/trajectoryExportCommands.ts Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Refactor adapter to avoid double trajectory start and add comprehensive tests - Remove redundant second startTrajectory call in adapter sync loop - Add 12 comprehensive tests for TrajectoryLoggerAdapter covering: - Basic trajectory creation from request logs - User message deduplication - Tool call correlation (single, parallel, orphan) - Subagent trajectory linking - Metrics tracking - Session ID management - Non-conversation request handling - Update TestRequestLogger to expose toolMetadata property - Add async wait in tests for proper event processing Co-authored-by: zhichli <57812115+zhichli@users.noreply.github.com> * rm readme * enhance trajectory export functionality with session ID mapping and folder selection * add comment * add arch * add bounding for logs and rm tests for now --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: zhichli <57812115+zhichli@users.noreply.github.com> Co-authored-by: Zhichao Li <zhichli@microsoft.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
This commit is contained in:
Copilot
GitHub
@@ -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": [
|
||||
|
||||
@@ -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),
|
||||
];
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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<ISearchSubagentToolCallingLoopOptions> {
|
||||
@@ -54,7 +56,7 @@ export class SearchSubagentToolCallingLoop extends ToolCallingLoop<ISearchSubage
|
||||
context.tools = {
|
||||
...context.tools,
|
||||
toolReferences: [],
|
||||
subAgentInvocationId: randomUUID(),
|
||||
subAgentInvocationId: this.options.subAgentInvocationId ?? randomUUID(),
|
||||
subAgentName: 'search'
|
||||
};
|
||||
}
|
||||
|
||||
@@ -55,6 +55,11 @@ class SearchSubagentTool implements ICopilotTool<ISearchSubagentParams> {
|
||||
|
||||
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<ISearchSubagentParams> {
|
||||
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<ISearchSubagentParams> {
|
||||
|
||||
// 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<ISearchSubagentParams> {
|
||||
// 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 = '';
|
||||
|
||||
@@ -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<string> # entry IDs already synced │ │
|
||||
│ │ ├── processedToolCalls: Set<string> # tool call IDs processed │ │
|
||||
│ │ ├── lastUserMessageBySession: Map<sessionId, hash> │ │
|
||||
│ │ ├── requestToStepContext: Map<requestId, StepInfo> │ │
|
||||
│ │ └── runSubagentToolCallToSessionId: Map<toolCallId, sessionId> │ │
|
||||
│ │ │ │
|
||||
│ │ TOKEN MAPPING (WeakMap - GC-friendly) │ │
|
||||
│ │ ├── sessionMap: WeakMap<CapturingToken, sessionId> │ │
|
||||
│ │ └── tokenToSessionId: WeakMap<CapturingToken, sessionId> │ │
|
||||
│ └─────────────────────────────────────────────────────────────────────────┘ │
|
||||
│ │ │
|
||||
│ Converts LoggedInfo → TrajectoryStep │
|
||||
└─────────────────────────────────────┬───────────────────────────────────────────┘
|
||||
│
|
||||
▼
|
||||
┌─────────────────────────────────────────────────────────────────────────────────┐
|
||||
│ TrajectoryLogger │
|
||||
│ ┌─────────────────────────────────────────────────────────────────────────┐ │
|
||||
│ │ TRAJECTORY STORAGE (⚠️ UNBOUNDED) │ │
|
||||
│ │ ├── trajectories: Map<sessionId, TrajectoryBuilder> │ │
|
||||
│ │ │ └── steps: ITrajectoryStep[] │ │
|
||||
│ │ └── subagentTrajectories: Map<sessionId, IAgentTrajectory> │ │
|
||||
│ │ │ │
|
||||
│ │ 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<boolean> {
|
||||
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<string, TrajectoryBuilder>();
|
||||
private subagentTrajectories = new Map<string, IAgentTrajectory>();
|
||||
```
|
||||
|
||||
### 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<CapturingToken, string>();
|
||||
private tokenToSessionId = new WeakMap<CapturingToken, string>();
|
||||
|
||||
// ⚠️ UNBOUNDED - grows indefinitely
|
||||
private processedEntries = new Set<string>(); // entry.id strings
|
||||
private processedToolCalls = new Set<string>(); // tool call ID strings
|
||||
private lastUserMessageBySession = new Map<string, string>();
|
||||
private requestToStepContext = new Map<string, {...}>();
|
||||
private runSubagentToolCallToSessionId = new Map<string, string>();
|
||||
```
|
||||
|
||||
## 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<sessionId, IAgentTrajectory>
|
||||
└──────────────────────────┘
|
||||
```
|
||||
|
||||
## 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<void> {
|
||||
// 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<sessionId, IAgentTrajectory> │
|
||||
└───────────────────────────────────┘
|
||||
│
|
||||
▼
|
||||
┌───────────────────────────────────┐
|
||||
│ 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.
|
||||
@@ -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<string, IAgentTrajectory>): Map<string, string> {
|
||||
const mapping = new Map<string, string>();
|
||||
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, string>): 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<string, IAgentTrajectory>,
|
||||
saveDir: vscode.Uri,
|
||||
pathMapping: Map<string, string>
|
||||
): Promise<void> {
|
||||
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<vscode.Uri | undefined> {
|
||||
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<void> {
|
||||
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<void> {
|
||||
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<string, IAgentTrajectory>
|
||||
): Map<string, IAgentTrajectory> {
|
||||
const result = new Map<string, IAgentTrajectory>();
|
||||
const visited = new Set<string>();
|
||||
|
||||
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, '_');
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
) { }
|
||||
}
|
||||
|
||||
@@ -124,6 +124,7 @@ export interface ILoggedToolCall {
|
||||
token: CapturingToken | undefined;
|
||||
time: number;
|
||||
thinking?: ThinkingData;
|
||||
toolMetadata?: unknown;
|
||||
toJSON(): Promise<object>;
|
||||
}
|
||||
|
||||
|
||||
@@ -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<object> {
|
||||
return {
|
||||
|
||||
@@ -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>('ITrajectoryLogger');
|
||||
|
||||
export interface ITrajectoryLogger {
|
||||
readonly _serviceBrand: undefined;
|
||||
|
||||
/**
|
||||
* Event fired when the trajectory is updated
|
||||
*/
|
||||
readonly onDidUpdateTrajectory: Event<void>;
|
||||
|
||||
/**
|
||||
* 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<string, IAgentTrajectory>;
|
||||
|
||||
/**
|
||||
* 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;
|
||||
}
|
||||
@@ -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<string, unknown>;
|
||||
}
|
||||
|
||||
/**
|
||||
* 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<string, unknown>;
|
||||
}
|
||||
|
||||
/**
|
||||
* Tool definition in OpenAI function calling format
|
||||
*/
|
||||
export interface IToolDefinition {
|
||||
readonly type: 'function';
|
||||
readonly function: {
|
||||
readonly name: string;
|
||||
readonly description: string;
|
||||
readonly parameters?: Record<string, unknown>;
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* 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<string, unknown>;
|
||||
}
|
||||
|
||||
/**
|
||||
* 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<string, unknown>;
|
||||
}
|
||||
|
||||
/**
|
||||
* 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<string, unknown>;
|
||||
}
|
||||
|
||||
/**
|
||||
* 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<string, unknown>;
|
||||
}
|
||||
|
||||
/**
|
||||
* 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<string, unknown>;
|
||||
}
|
||||
|
||||
/**
|
||||
* 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';
|
||||
@@ -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<string, TrajectoryBuilder>();
|
||||
private currentSessionId: string | undefined;
|
||||
private subagentTrajectories = new Map<string, IAgentTrajectory>();
|
||||
|
||||
private readonly _onDidUpdateTrajectory = this._register(new Emitter<void>());
|
||||
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<string, IAgentTrajectory> {
|
||||
const trajectories = new Map<string, IAgentTrajectory>();
|
||||
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<ITrajectoryStep> = {
|
||||
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<ITrajectoryStep>,
|
||||
private readonly onComplete: (step: Partial<ITrajectoryStep>) => 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);
|
||||
}
|
||||
}
|
||||
@@ -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<string>;
|
||||
|
||||
/**
|
||||
* 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<CapturingToken, string>();
|
||||
// Also maintain a map for lookup by token reference without preventing GC of tokens
|
||||
private tokenToSessionId = new WeakMap<CapturingToken, string>();
|
||||
private processedEntries = new Set<string>();
|
||||
private processedToolCalls = new Set<string>(); // Track processed tool calls by their ID
|
||||
private lastUserMessageBySession = new Map<string, string>();
|
||||
// Track pending step contexts by both session and request ID to handle parallel tool calls
|
||||
private pendingStepContexts = new Map<string, IAgentStepContext>();
|
||||
private requestToStepContext = new Map<string, { sessionId: string; context: IAgentStepContext; toolCallCount: number; processedToolCalls: number }>();
|
||||
private runSubagentToolCallToSessionId = new Map<string, string>();
|
||||
|
||||
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<void> {
|
||||
// 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<void> {
|
||||
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<void> {
|
||||
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<void> {
|
||||
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<void> {
|
||||
// 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<string, unknown> {
|
||||
if (typeof args === 'string') {
|
||||
try {
|
||||
return JSON.parse(args) as Record<string, unknown>;
|
||||
} catch {
|
||||
return { raw: args };
|
||||
}
|
||||
}
|
||||
if (typeof args === 'object' && args !== null) {
|
||||
return args as Record<string, unknown>;
|
||||
}
|
||||
return {};
|
||||
}
|
||||
|
||||
private async extractToolResultContent(entry: ILoggedToolCall): Promise<string> {
|
||||
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<string> {
|
||||
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 `<decode error: ${part.data.length} bytes>`;
|
||||
}
|
||||
}
|
||||
|
||||
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<Record<string, unknown>> | 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<string, unknown> } | 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();
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user