Fix BYOK request logging: preserve context across IPC and remove duplicates#3474
Fix BYOK request logging: preserve context across IPC and remove duplicates#3474
Conversation
- Add correlation ID mechanism to restore AsyncLocalStorage context after IPC - BYOK providers (AnthropicBYOK, GeminiNativeBYOK, CopilotLanguageModelWrapper) now log requests as children of their prompt items with correct token usage - Remove duplicate request logging from ExtensionContributedChatEndpoint - Simplify requestLogTree.ts export logic (remove lastChatRequestItem fallback)
There was a problem hiding this comment.
Pull request overview
This PR fixes BYOK request logging by implementing a correlation ID mechanism to preserve AsyncLocalStorage context across IPC boundaries, eliminating duplicate log entries and ensuring requests appear correctly as children of their prompt items.
Changes:
- Added correlation ID mechanism to preserve CapturingToken context across VS Code IPC boundaries
- Removed duplicate logging from ExtensionContributedChatEndpoint (now handled by BYOK providers)
- Updated BYOK providers to restore context using correlation IDs
- Simplified request log export logic by removing workarounds for orphaned requests
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| src/platform/requestLogger/node/requestLogger.ts | Added correlation map and helper functions to store/retrieve CapturingToken across IPC |
| src/platform/endpoint/vscode-node/extChatEndpoint.ts | Store correlation ID before IPC call, removed duplicate logging and unused dependencies |
| src/extension/byok/vscode-node/anthropicProvider.ts | Retrieve correlation ID and wrap request execution in restored context |
| src/extension/byok/vscode-node/geminiNativeProvider.ts | Retrieve correlation ID and wrap request execution in restored context |
| src/extension/conversation/vscode-node/languageModelAccess.ts | CopilotLanguageModelWrapper retrieves correlation ID for OpenAI-compatible providers |
| src/extension/log/vscode-node/requestLogTree.ts | Removed workaround logic for re-associating orphaned requests |
| // BYOK providers (Anthropic, Gemini, CopilotLanguageModelWrapper) handle their own | ||
| // logging with correct token usage. Logging here would create duplicates with | ||
| // incorrect (0) token counts since we don't have access to actual usage stats. | ||
| storeCapturingTokenForCorrelation(ourRequestId); |
There was a problem hiding this comment.
Potential memory leak: If the request fails before reaching a BYOK provider (e.g., languageModel.sendRequest throws immediately, or if the model is not a BYOK provider), the correlation ID stored here will never be retrieved, leaving entries in capturingTokenCorrelationMap that are never cleaned up.
The correlation ID should be retrieved in a finally block or there should be a timeout-based cleanup mechanism for the correlation map. Alternatively, always attempt to retrieve (and discard) the correlation ID in the catch/finally block of this method to ensure cleanup even when errors occur.
There was a problem hiding this comment.
@copilot open a new pull request to apply changes based on this feedback
Problem
BYOK model requests (Anthropic, Gemini, CustomOAI/Azure via CopilotLanguageModelWrapper) were not being logged correctly in the request log tree:
vscode.lm.sendRequest(), theAsyncLocalStoragecontext is lost. This caused BYOK requests to appear as top-level items instead of children of their prompt items.Duplicate Logging:
ExtensionContributedChatEndpointwas logging requests with 0 token usage, while BYOK providers logged the same requests with correct token counts, creating duplicates.Incorrect Export: The log export logic had workarounds to re-associate orphaned requests, which was fragile.
Solution
Correlation ID Mechanism
capturingTokenCorrelationMapinrequestLogger.tsto storeCapturingTokenwith a correlation ID before IPCmodelOptions._capturingTokenCorrelationIdrunWithCapturingToken()Changes
requestLogger.tsstoreCapturingTokenForCorrelation()- stores token before IPCretrieveCapturingTokenByCorrelation()- retrieves token after IPCrunWithCapturingToken()- executes code within restored contextextChatEndpoint.tssendRequest()IPC call_requestLoggerand_endpointProviderdependenciesanthropicProvider.ts/geminiNativeProvider.ts/languageModelAccess.tsmodelOptionsCapturingTokencontextrequestLogTree.tslastChatRequestItemfallback workaroundResult
Verified Examples