From da5fb11b3bc01bd23c9b8bf27a1e2e1059c56080 Mon Sep 17 00:00:00 2001 From: Connor Peet Date: Thu, 26 Mar 2026 15:46:46 -0700 Subject: [PATCH 1/8] agentHost: add session-specific metadata Adds a SQLite DB for session-specific metadata. Stores edits in there. It can _almost_ restore edits, but I still need to make undoStops be similarly persisted. That is a project for later this evening. --- .../common/agentHostFileSystemProvider.ts | 8 +- .../platform/agentHost/common/agentHostUri.ts | 3 +- .../agentHost/common/sessionDataService.ts | 86 ++++ .../platform/agentHost/node/agentService.ts | 8 +- .../agentHost/node/agentSideEffects.ts | 29 +- .../agentHost/node/copilot/copilotAgent.ts | 123 ++--- .../agentHost/node/copilot/fileEditTracker.ts | 138 ++++-- .../node/copilot/mapSessionEvents.ts | 216 +++++++++ .../agentHost/node/sessionDataService.ts | 33 +- .../agentHost/node/sessionDatabase.ts | 315 +++++++++++++ .../agentHost/node/sessionStateManager.ts | 12 +- .../agentHost/test/node/agentService.test.ts | 1 + .../test/node/agentSideEffects.test.ts | 1 + .../test/node/fileEditTracker.test.ts | 142 ++++-- .../test/node/mapSessionEvents.test.ts | 238 ++++++++++ .../test/node/sessionDataService.test.ts | 74 +++ .../test/node/sessionDatabase.test.ts | 421 ++++++++++++++++++ .../test/node/sessionStateManager.test.ts | 15 +- .../browser/chatEditing/chatEditingSession.ts | 2 +- 19 files changed, 1688 insertions(+), 177 deletions(-) create mode 100644 src/vs/platform/agentHost/node/copilot/mapSessionEvents.ts create mode 100644 src/vs/platform/agentHost/node/sessionDatabase.ts create mode 100644 src/vs/platform/agentHost/test/node/mapSessionEvents.test.ts create mode 100644 src/vs/platform/agentHost/test/node/sessionDatabase.test.ts diff --git a/src/vs/platform/agentHost/common/agentHostFileSystemProvider.ts b/src/vs/platform/agentHost/common/agentHostFileSystemProvider.ts index 54eba157fb3..2842c3e2081 100644 --- a/src/vs/platform/agentHost/common/agentHostFileSystemProvider.ts +++ b/src/vs/platform/agentHost/common/agentHostFileSystemProvider.ts @@ -80,8 +80,12 @@ export class AgentHostFileSystemProvider extends Disposable implements IFileSyst if (path === '/' || path === '') { return { type: FileType.Directory, mtime: 0, ctime: 0, size: 0, permissions: FilePermission.Readonly }; } - const decodedPath = fromAgentHostUri(resource).path; - if (decodedPath === '/' || decodedPath === '') { + const decoded = fromAgentHostUri(resource); + if (decoded.scheme === 'session-db') { + return { type: FileType.File, mtime: 0, ctime: 0, size: 0, permissions: FilePermission.Readonly }; + } + + if (decoded.path === '/' || decoded.path === '') { return { type: FileType.Directory, mtime: 0, ctime: 0, size: 0, permissions: FilePermission.Readonly }; } diff --git a/src/vs/platform/agentHost/common/agentHostUri.ts b/src/vs/platform/agentHost/common/agentHostUri.ts index 773bce4e676..157277a9213 100644 --- a/src/vs/platform/agentHost/common/agentHostUri.ts +++ b/src/vs/platform/agentHost/common/agentHostUri.ts @@ -4,6 +4,7 @@ *--------------------------------------------------------------------------------------------*/ import { encodeBase64, VSBuffer } from '../../../base/common/buffer.js'; +import { Schemas } from '../../../base/common/network.js'; import { URI } from '../../../base/common/uri.js'; import type { ResourceLabelFormatter } from '../../label/common/label.js'; @@ -34,7 +35,7 @@ export const AGENT_HOST_SCHEME = 'vscode-agent-host'; * the URI authority (from {@link agentHostAuthority}). */ export function toAgentHostUri(originalUri: URI, connectionAuthority: string): URI { - if (connectionAuthority === 'local') { + if (connectionAuthority === 'local' && originalUri.scheme === Schemas.file) { return originalUri; } diff --git a/src/vs/platform/agentHost/common/sessionDataService.ts b/src/vs/platform/agentHost/common/sessionDataService.ts index dcc52fde556..0f6f2c45763 100644 --- a/src/vs/platform/agentHost/common/sessionDataService.ts +++ b/src/vs/platform/agentHost/common/sessionDataService.ts @@ -3,11 +3,86 @@ * Licensed under the MIT License. See License.txt in the project root for license information. *--------------------------------------------------------------------------------------------*/ +import { IDisposable, IReference } from '../../../base/common/lifecycle.js'; import { URI } from '../../../base/common/uri.js'; import { createDecorator } from '../../instantiation/common/instantiation.js'; export const ISessionDataService = createDecorator('sessionDataService'); +// ---- File-edit types ---------------------------------------------------- + +/** + * Lightweight metadata for a file edit. Returned by {@link ISessionDatabase.getFileEdits} + * without the (potentially large) file content blobs. + */ +export interface IFileEditRecord { + /** The turn that owns this file edit. */ + turnId: string; + /** The tool call that produced this edit. */ + toolCallId: string; + /** Absolute file path that was edited. */ + filePath: string; + /** Number of lines added (informational, for diff metadata). */ + addedLines: number | undefined; + /** Number of lines removed (informational, for diff metadata). */ + removedLines: number | undefined; +} + +/** + * The before/after content blobs for a single file edit. + * Retrieved on demand via {@link ISessionDatabase.readFileEditContent}. + */ +export interface IFileEditContent { + /** File content before the edit (may be empty for newly created files). */ + beforeContent: Uint8Array; + /** File content after the edit. */ + afterContent: Uint8Array; +} + +// ---- Session database --------------------------------------------------- + +/** + * A disposable handle to a per-session SQLite database backed by + * `@vscode/sqlite3`. + * + * Callers obtain an instance via {@link ISessionDataService.openDatabase} and + * **must** dispose it when finished to close the underlying database connection. + */ +export interface ISessionDatabase extends IDisposable { + /** + * Create a turn record. Must be called before storing file edits that + * reference this turn. + */ + createTurn(turnId: string): Promise; + + /** + * Delete a turn and all of its associated file edits (cascade). + */ + deleteTurn(turnId: string): Promise; + + /** + * Store a file-edit snapshot (metadata + content) for a tool invocation + * within a turn. + * + * If a record for the same `toolCallId` and `filePath` already exists + * it is replaced. + */ + storeFileEdit(edit: IFileEditRecord & IFileEditContent): Promise; + + /** + * Retrieve file-edit metadata for the given tool call IDs. + * Content blobs are **not** included — use {@link readFileEditContent} + * to fetch them on demand. Results are returned in insertion order. + */ + getFileEdits(toolCallIds: string[]): Promise; + + /** + * Read the before/after content blobs for a single file edit. + * Returns `undefined` if no edit exists for the given key. + */ + readFileEditContent(toolCallId: string, filePath: string): Promise; +} + /** * Provides persistent, per-session data directories on disk. * @@ -34,6 +109,17 @@ export interface ISessionDataService { */ getSessionDataDirById(sessionId: string): URI; + /** + * Opens (or creates) a per-session SQLite database. The database file is + * stored at `{sessionDataDir}/session.db`. Migrations are applied + * automatically on first use. + * + * Returns a ref-counted reference. Multiple callers for the same session + * share the same underlying connection. The connection is closed when + * the last reference is disposed. + */ + openDatabase(session: URI): IReference; + /** * Recursively deletes the data directory for a session, if it exists. */ diff --git a/src/vs/platform/agentHost/node/agentService.ts b/src/vs/platform/agentHost/node/agentService.ts index e4d16a198b2..50a79dfa965 100644 --- a/src/vs/platform/agentHost/node/agentService.ts +++ b/src/vs/platform/agentHost/node/agentService.ts @@ -171,7 +171,7 @@ export class AgentService extends Disposable implements IAgentService { await provider.disposeSession(session); this._sessionToProvider.delete(session.toString()); } - this._stateManager.removeSession(session.toString()); + this._stateManager.deleteSession(session.toString()); this._sessionDataService.deleteSessionData(session); } @@ -179,7 +179,11 @@ export class AgentService extends Disposable implements IAgentService { async subscribe(resource: URI): Promise { this._logService.trace(`[AgentService] subscribe: ${resource.toString()}`); - const snapshot = this._stateManager.getSnapshot(resource.toString()); + let snapshot = this._stateManager.getSnapshot(resource.toString()); + if (!snapshot) { + await this.restoreSession(resource); + snapshot = this._stateManager.getSnapshot(resource.toString()); + } if (!snapshot) { throw new Error(`Cannot subscribe to unknown resource: ${resource.toString()}`); } diff --git a/src/vs/platform/agentHost/node/agentSideEffects.ts b/src/vs/platform/agentHost/node/agentSideEffects.ts index b7592c59f47..64132425be3 100644 --- a/src/vs/platform/agentHost/node/agentSideEffects.ts +++ b/src/vs/platform/agentHost/node/agentSideEffects.ts @@ -29,6 +29,7 @@ import { type URI as ProtocolURI, } from '../common/state/sessionState.js'; import { AgentEventMapper } from './agentEventMapper.js'; +import { ISessionDbUriFields, parseSessionDbUri } from './copilot/fileEditTracker.js'; import type { IProtocolSideEffectHandler } from './protocolServerHandler.js'; import { SessionStateManager } from './sessionStateManager.js'; @@ -339,7 +340,7 @@ export class AgentSideEffects extends Disposable implements IProtocolSideEffectH handleDisposeSession(session: ProtocolURI): void { const agent = this._options.getAgent(session); agent?.disposeSession(URI.parse(session)).catch(() => { }); - this._stateManager.removeSession(session); + this._stateManager.deleteSession(session); this._options.sessionDataService.deleteSessionData(URI.parse(session)); } @@ -567,6 +568,13 @@ export class AgentSideEffects extends Disposable implements IProtocolSideEffectH } async handleFetchContent(uri: ProtocolURI): Promise { + // Handle session-db: URIs that reference file-edit content stored + // in a per-session SQLite database. + const dbFields = parseSessionDbUri(uri); + if (dbFields) { + return this._fetchSessionDbContent(dbFields); + } + try { const content = await this._fileService.readFile(URI.parse(uri)); return { @@ -579,6 +587,25 @@ export class AgentSideEffects extends Disposable implements IProtocolSideEffectH } } + private async _fetchSessionDbContent(fields: ISessionDbUriFields): Promise { + const sessionUri = URI.parse(fields.sessionUri); + const ref = this._options.sessionDataService.openDatabase(sessionUri); + try { + const content = await ref.object.readFileEditContent(fields.toolCallId, fields.filePath); + if (!content) { + throw new ProtocolError(AhpErrorCodes.NotFound, `File edit not found: toolCallId=${fields.toolCallId}, filePath=${fields.filePath}`); + } + const bytes = fields.part === 'before' ? content.beforeContent : content.afterContent; + return { + data: new TextDecoder().decode(bytes), + encoding: ContentEncoding.Utf8, + contentType: 'text/plain', + }; + } finally { + ref.dispose(); + } + } + override dispose(): void { this._toolCallAgents.clear(); super.dispose(); diff --git a/src/vs/platform/agentHost/node/copilot/copilotAgent.ts b/src/vs/platform/agentHost/node/copilot/copilotAgent.ts index 31a4fad42e9..443082caeff 100644 --- a/src/vs/platform/agentHost/node/copilot/copilotAgent.ts +++ b/src/vs/platform/agentHost/node/copilot/copilotAgent.ts @@ -3,11 +3,11 @@ * Licensed under the MIT License. See License.txt in the project root for license information. *--------------------------------------------------------------------------------------------*/ -import { CopilotClient, CopilotSession, type SessionEvent, type SessionEventPayload } from '@github/copilot-sdk'; +import { CopilotClient, CopilotSession } from '@github/copilot-sdk'; import { rgPath } from '@vscode/ripgrep'; import { DeferredPromise } from '../../../../base/common/async.js'; import { Emitter } from '../../../../base/common/event.js'; -import { Disposable, DisposableMap } from '../../../../base/common/lifecycle.js'; +import { Disposable, DisposableMap, IReference } from '../../../../base/common/lifecycle.js'; import { FileAccess } from '../../../../base/common/network.js'; import type { IAuthorizationProtectedResourceMetadata } from '../../../../base/common/oauth.js'; import { delimiter, dirname } from '../../../../base/common/path.js'; @@ -16,11 +16,12 @@ import { IFileService } from '../../../files/common/files.js'; import { ILogService } from '../../../log/common/log.js'; import { localize } from '../../../../nls.js'; import { AgentSession, IAgent, IAgentAttachment, IAgentCreateSessionConfig, IAgentDescriptor, IAgentMessageEvent, IAgentModelInfo, IAgentProgressEvent, IAgentSessionMetadata, IAgentToolCompleteEvent, IAgentToolStartEvent } from '../../common/agentService.js'; -import { ISessionDataService } from '../../common/sessionDataService.js'; +import { ISessionDatabase, ISessionDataService } from '../../common/sessionDataService.js'; import { ToolResultContentType, type IPendingMessage, type IToolResultContent, type PolicyState } from '../../common/state/sessionState.js'; import { CopilotSessionWrapper } from './copilotSessionWrapper.js'; import { getEditFilePath, getInvocationMessage, getPastTenseMessage, getShellLanguage, getToolDisplayName, getToolInputString, getToolKind, isEditTool, isHiddenTool } from './copilotToolDisplay.js'; import { FileEditTracker } from './fileEditTracker.js'; +import { mapSessionEvents } from './mapSessionEvents.js'; function tryStringify(value: unknown): string | undefined { try { @@ -51,6 +52,8 @@ export class CopilotAgent extends Disposable implements IAgent { private readonly _sessionWorkingDirs = new Map(); /** File edit trackers per session, keyed by raw session ID. */ private readonly _editTrackers = new Map(); + /** Session database references, keyed by raw session ID. */ + private readonly _sessionDatabases = this._register(new DisposableMap>()); constructor( @ILogService private readonly _logService: ILogService, @@ -270,7 +273,13 @@ export class CopilotAgent extends Disposable implements IAgent { } const events = await entry.session.getMessages(); - return this._mapSessionEvents(session, events); + let db: ISessionDatabase | undefined; + try { + db = this._getSessionDatabase(sessionId); + } catch { + // Database may not exist yet — that's fine + } + return mapSessionEvents(session, db, events); } async disposeSession(session: URI): Promise { @@ -278,6 +287,7 @@ export class CopilotAgent extends Disposable implements IAgent { this._sessions.deleteAndDispose(sessionId); this._clearToolCallsForSession(sessionId); this._sessionWorkingDirs.delete(sessionId); + this._sessionDatabases.deleteAndDispose(sessionId); this._denyPendingPermissionsForSession(sessionId); } @@ -306,6 +316,7 @@ export class CopilotAgent extends Disposable implements IAgent { this._activeToolCalls.clear(); this._sessionWorkingDirs.clear(); this._denyPendingPermissions(); + this._sessionDatabases.clearAndDisposeAll(); await this._client?.stop(); this._client = undefined; } @@ -439,10 +450,22 @@ export class CopilotAgent extends Disposable implements IAgent { } } + private _getSessionDatabase(rawSessionId: string): ISessionDatabase { + let ref = this._sessionDatabases.get(rawSessionId); + if (!ref) { + const session = AgentSession.uri(this.id, rawSessionId); + ref = this._sessionDataService.openDatabase(session); + this._sessionDatabases.set(rawSessionId, ref); + } + return ref.object; + } + private _getOrCreateEditTracker(rawSessionId: string): FileEditTracker { let tracker = this._editTrackers.get(rawSessionId); if (!tracker) { - tracker = new FileEditTracker(rawSessionId, this._sessionDataService, this._fileService, this._logService); + const session = AgentSession.uri(this.id, rawSessionId); + const db = this._getSessionDatabase(rawSessionId); + tracker = new FileEditTracker(session.toString(), db, this._fileService, this._logService); this._editTrackers.set(rawSessionId, tracker); } return tracker; @@ -547,6 +570,11 @@ export class CopilotAgent extends Disposable implements IAgent { }); }); + let turnId: string = ''; + wrapper.onTurnStart(e => { + turnId = e.data.turnId; + }); + wrapper.onToolComplete(e => { const trackingKey = `${rawId}:${e.data.toolCallId}`; const tracked = this._activeToolCalls.get(trackingKey); @@ -567,7 +595,7 @@ export class CopilotAgent extends Disposable implements IAgent { const tracker = this._editTrackers.get(rawId); const filePath = isEditTool(tracked.toolName) ? getEditFilePath(tracked.parameters) : undefined; if (tracker && filePath) { - const fileEdit = tracker.takeCompletedEdit(filePath); + const fileEdit = tracker.takeCompletedEdit(turnId, e.data.toolCallId, filePath); if (fileEdit) { content.push(fileEdit); } @@ -761,89 +789,6 @@ export class CopilotAgent extends Disposable implements IAgent { return this._trackSession(raw, sessionId); } - private _mapSessionEvents(session: URI, events: readonly SessionEvent[]): (IAgentMessageEvent | IAgentToolStartEvent | IAgentToolCompleteEvent)[] { - const result: (IAgentMessageEvent | IAgentToolStartEvent | IAgentToolCompleteEvent)[] = []; - const toolInfoByCallId = new Map | undefined }>(); - - for (const e of events) { - if (e.type === 'assistant.message' || e.type === 'user.message') { - const d = (e as SessionEventPayload<'assistant.message'>).data; - result.push({ - session, - type: 'message', - role: e.type === 'user.message' ? 'user' : 'assistant', - messageId: d?.messageId ?? '', - content: d?.content ?? '', - toolRequests: d?.toolRequests?.map((tr: { toolCallId: string; name: string; arguments?: unknown; type?: 'function' | 'custom' }) => ({ - toolCallId: tr.toolCallId, - name: tr.name, - arguments: tr.arguments !== undefined ? tryStringify(tr.arguments) : undefined, - type: tr.type, - })), - reasoningOpaque: d?.reasoningOpaque, - reasoningText: d?.reasoningText, - encryptedContent: d?.encryptedContent, - parentToolCallId: d?.parentToolCallId, - }); - } else if (e.type === 'tool.execution_start') { - const d = (e as SessionEventPayload<'tool.execution_start'>).data; - if (isHiddenTool(d.toolName)) { - continue; - } - const toolArgs = d.arguments !== undefined ? tryStringify(d.arguments) : undefined; - let parameters: Record | undefined; - if (toolArgs) { - try { parameters = JSON.parse(toolArgs) as Record; } catch { /* ignore */ } - } - toolInfoByCallId.set(d.toolCallId, { toolName: d.toolName, parameters }); - const displayName = getToolDisplayName(d.toolName); - const toolKind = getToolKind(d.toolName); - result.push({ - session, - type: 'tool_start', - toolCallId: d.toolCallId, - toolName: d.toolName, - displayName, - invocationMessage: getInvocationMessage(d.toolName, displayName, parameters), - toolInput: getToolInputString(d.toolName, parameters, toolArgs), - toolKind, - language: toolKind === 'terminal' ? getShellLanguage(d.toolName) : undefined, - toolArguments: toolArgs, - mcpServerName: d.mcpServerName, - mcpToolName: d.mcpToolName, - parentToolCallId: d.parentToolCallId, - }); - } else if (e.type === 'tool.execution_complete') { - const d = (e as SessionEventPayload<'tool.execution_complete'>).data; - const info = toolInfoByCallId.get(d.toolCallId); - if (!info) { - continue; - } - toolInfoByCallId.delete(d.toolCallId); - const displayName = getToolDisplayName(info.toolName); - const toolOutput = d.error?.message ?? d.result?.content; - const content: IToolResultContent[] = []; - if (toolOutput !== undefined) { - content.push({ type: ToolResultContentType.Text, text: toolOutput }); - } - result.push({ - session, - type: 'tool_complete', - toolCallId: d.toolCallId, - result: { - success: d.success, - pastTenseMessage: getPastTenseMessage(info.toolName, displayName, info.parameters, d.success), - content: content.length > 0 ? content : undefined, - error: d.error, - }, - isUserRequested: d.isUserRequested, - toolTelemetry: d.toolTelemetry !== undefined ? tryStringify(d.toolTelemetry) : undefined, - }); - } - } - return result; - } - override dispose(): void { this._denyPendingPermissions(); this._client?.stop().catch(() => { /* best-effort */ }); diff --git a/src/vs/platform/agentHost/node/copilot/fileEditTracker.ts b/src/vs/platform/agentHost/node/copilot/fileEditTracker.ts index 8ee6b5a0893..d6aa590242f 100644 --- a/src/vs/platform/agentHost/node/copilot/fileEditTracker.ts +++ b/src/vs/platform/agentHost/node/copilot/fileEditTracker.ts @@ -3,16 +3,60 @@ * Licensed under the MIT License. See License.txt in the project root for license information. *--------------------------------------------------------------------------------------------*/ -import { VSBuffer } from '../../../../base/common/buffer.js'; +import { decodeHex, encodeHex, VSBuffer } from '../../../../base/common/buffer.js'; import { URI } from '../../../../base/common/uri.js'; import { IFileService } from '../../../files/common/files.js'; import { ILogService } from '../../../log/common/log.js'; -import { ISessionDataService } from '../../common/sessionDataService.js'; +import { ISessionDatabase } from '../../common/sessionDataService.js'; import { ToolResultContentType, type IToolResultFileEditContent } from '../../common/state/sessionState.js'; +const SESSION_DB_SCHEME = 'session-db'; + +/** + * Builds a `session-db:` URI that references a file-edit content blob + * stored in the session database. Parsed by {@link parseSessionDbUri}. + */ +export function buildSessionDbUri(sessionUri: string, toolCallId: string, filePath: string, part: 'before' | 'after'): string { + return URI.from({ + scheme: SESSION_DB_SCHEME, + authority: encodeHex(VSBuffer.fromString(sessionUri)).toString(), + path: `/${encodeURIComponent(toolCallId)}/${encodeHex(VSBuffer.fromString(filePath))}/${part}`, + }).toString(); +} + +/** Parsed fields from a `session-db:` content URI. */ +export interface ISessionDbUriFields { + sessionUri: string; + toolCallId: string; + filePath: string; + part: 'before' | 'after'; +} + +/** + * Parses a `session-db:` URI produced by {@link buildSessionDbUri}. + * Returns `undefined` if the URI is not a valid `session-db:` URI. + */ +export function parseSessionDbUri(raw: string): ISessionDbUriFields | undefined { + const parsed = URI.parse(raw); + if (parsed.scheme !== SESSION_DB_SCHEME) { + return undefined; + } + const [, toolCallId, filePath, part] = parsed.path.split('/'); + if (!toolCallId || !filePath || (part !== 'before' && part !== 'after')) { + return undefined; + } + return { + sessionUri: decodeHex(parsed.authority).toString(), + toolCallId: decodeURIComponent(toolCallId), + filePath: decodeHex(filePath).toString(), + part + }; +} + /** * Tracks file edits made by tools in a session by snapshotting file content - * before and after each edit tool invocation. + * before and after each edit tool invocation, persisting snapshots into the + * session database. */ export class FileEditTracker { @@ -20,45 +64,41 @@ export class FileEditTracker { * Pending edits keyed by file path. The `onPreToolUse` hook stores * entries here; `completeEdit` pops them when the tool finishes. */ - private readonly _pendingEdits = new Map }>(); + private readonly _pendingEdits = new Map }>(); /** * Completed edits keyed by file path. The `onPostToolUse` hook stores - * entries here; `takeCompletedEdit` retrieves them synchronously from - * the `onToolComplete` handler. + * entries here; `takeCompletedEdit` retrieves them from the + * `onToolComplete` handler and persists to the database. */ - private readonly _completedEdits = new Map(); + private readonly _completedEdits = new Map(); constructor( - private readonly _sessionId: string, - private readonly _sessionDataService: ISessionDataService, + private readonly _sessionUri: string, + private readonly _db: ISessionDatabase, private readonly _fileService: IFileService, private readonly _logService: ILogService, ) { } /** * Call from the `onPreToolUse` hook before an edit tool runs. - * Snapshots the file's current content as the "before" state. + * Reads the file's current content into memory as the "before" state. * The hook blocks the SDK until this returns, ensuring the snapshot * captures pre-edit content. * * @param filePath - Absolute path of the file being edited. */ async trackEditStart(filePath: string): Promise { - const editKey = generateEditKey(); - const sessionDataDir = this._sessionDataService.getSessionDataDirById(this._sessionId); - const beforeUri = URI.joinPath(sessionDataDir, 'file-edits', editKey, 'before'); - - const snapshotDone = this._snapshotFile(filePath, beforeUri); - this._pendingEdits.set(filePath, { editKey, beforeUri, snapshotDone }); - await snapshotDone; + const snapshotDone = this._readFile(filePath); + const entry = { beforeContent: VSBuffer.fromString(''), snapshotDone: snapshotDone.then(buf => { entry.beforeContent = buf; }) }; + this._pendingEdits.set(filePath, entry); + await entry.snapshotDone; } /** * Call from the `onPostToolUse` hook after an edit tool finishes. - * Stores the result for later synchronous retrieval via {@link takeCompletedEdit}. - * The `beforeURI` points to the stored snapshot; the `afterURI` is - * the real file path (the tool already modified it on disk). + * Reads the file content again as the "after" state and stores the + * result for later retrieval via {@link takeCompletedEdit}. * * @param filePath - Absolute path of the file that was edited. */ @@ -70,44 +110,56 @@ export class FileEditTracker { this._pendingEdits.delete(filePath); await pending.snapshotDone; - // Snapshot the after-content into session data so it remains - // stable even if the file is modified again later. - const sessionDataDir = this._sessionDataService.getSessionDataDirById(this._sessionId); - const afterUri = URI.joinPath(sessionDataDir, 'file-edits', pending.editKey, 'after'); - await this._snapshotFile(filePath, afterUri); + const afterContent = await this._readFile(filePath); this._completedEdits.set(filePath, { - type: ToolResultContentType.FileEdit, - beforeURI: pending.beforeUri.toString(), - afterURI: afterUri.toString(), + beforeContent: pending.beforeContent, + afterContent, }); } /** - * Synchronously retrieves and removes a completed edit for the given - * file path. Call from the `onToolComplete` handler to include the - * edit in the tool result without async work. + * Retrieves and removes a completed edit for the given file path, + * persists it to the session database, and returns the result as an + * {@link IToolResultFileEditContent} for inclusion in the tool result. + * + * @param toolCallId - The tool call that produced this edit. + * @param filePath - Absolute path of the edited file. */ - takeCompletedEdit(filePath: string): IToolResultFileEditContent | undefined { + takeCompletedEdit(turnId: string, toolCallId: string, filePath: string): IToolResultFileEditContent | undefined { const edit = this._completedEdits.get(filePath); - if (edit) { - this._completedEdits.delete(filePath); + if (!edit) { + return undefined; } - return edit; + this._completedEdits.delete(filePath); + + const beforeBytes = edit.beforeContent.buffer; + const afterBytes = edit.afterContent.buffer; + + this._db.storeFileEdit({ + turnId, + toolCallId, + filePath, + beforeContent: beforeBytes, + afterContent: afterBytes, + addedLines: undefined, + removedLines: undefined, + }).catch(err => this._logService.warn(`[FileEditTracker] Failed to persist file edit to database: ${filePath}`, err)); + + return { + type: ToolResultContentType.FileEdit, + beforeURI: buildSessionDbUri(this._sessionUri, toolCallId, filePath, 'before'), + afterURI: buildSessionDbUri(this._sessionUri, toolCallId, filePath, 'after'), + }; } - private async _snapshotFile(filePath: string, targetUri: URI): Promise { + private async _readFile(filePath: string): Promise { try { const content = await this._fileService.readFile(URI.file(filePath)); - await this._fileService.writeFile(targetUri, content.value); + return content.value; } catch (err) { this._logService.trace(`[FileEditTracker] Could not read file for snapshot: ${filePath}`, err); - await this._fileService.writeFile(targetUri, VSBuffer.fromString('')).catch(() => { }); + return VSBuffer.fromString(''); } } } - -let _editKeyCounter = 0; -function generateEditKey(): string { - return `${Date.now()}-${_editKeyCounter++}`; -} diff --git a/src/vs/platform/agentHost/node/copilot/mapSessionEvents.ts b/src/vs/platform/agentHost/node/copilot/mapSessionEvents.ts new file mode 100644 index 00000000000..56c9a09e79a --- /dev/null +++ b/src/vs/platform/agentHost/node/copilot/mapSessionEvents.ts @@ -0,0 +1,216 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { URI } from '../../../../base/common/uri.js'; +import { IAgentMessageEvent, IAgentToolCompleteEvent, IAgentToolStartEvent } from '../../common/agentService.js'; +import { IFileEditRecord, ISessionDatabase } from '../../common/sessionDataService.js'; +import { ToolResultContentType, type IToolResultContent } from '../../common/state/sessionState.js'; +import { getInvocationMessage, getPastTenseMessage, getShellLanguage, getToolDisplayName, getToolInputString, getToolKind, isEditTool, isHiddenTool } from './copilotToolDisplay.js'; +import { buildSessionDbUri } from './fileEditTracker.js'; + +function tryStringify(value: unknown): string | undefined { + try { + return JSON.stringify(value); + } catch { + return undefined; + } +} + +// ---- Minimal event shapes matching the SDK's SessionEvent union --------- +// Defined here so tests can construct events without importing the SDK. + +export interface ISessionEventToolStart { + type: 'tool.execution_start'; + data: { + toolCallId: string; + toolName: string; + arguments?: unknown; + mcpServerName?: string; + mcpToolName?: string; + parentToolCallId?: string; + }; +} + +export interface ISessionEventToolComplete { + type: 'tool.execution_complete'; + data: { + toolCallId: string; + success: boolean; + result?: { content?: string }; + error?: { message: string; code?: string }; + isUserRequested?: boolean; + toolTelemetry?: unknown; + parentToolCallId?: string; + }; +} + +export interface ISessionEventMessage { + type: 'assistant.message' | 'user.message'; + data?: { + messageId?: string; + content?: string; + toolRequests?: readonly { toolCallId: string; name: string; arguments?: unknown; type?: 'function' | 'custom' }[]; + reasoningOpaque?: string; + reasoningText?: string; + encryptedContent?: string; + parentToolCallId?: string; + }; +} + +/** Minimal event shape for session history mapping. */ +export type ISessionEvent = ISessionEventToolStart | ISessionEventToolComplete | ISessionEventMessage | { type: string; data?: unknown }; + +/** + * Maps raw SDK session events into agent protocol events, restoring + * stored file-edit metadata from the session database when available. + * + * Extracted as a standalone function so it can be tested without the + * full CopilotAgent or SDK dependencies. + */ +export async function mapSessionEvents( + session: URI, + db: ISessionDatabase | undefined, + events: readonly ISessionEvent[], +): Promise<(IAgentMessageEvent | IAgentToolStartEvent | IAgentToolCompleteEvent)[]> { + const result: (IAgentMessageEvent | IAgentToolStartEvent | IAgentToolCompleteEvent)[] = []; + const toolInfoByCallId = new Map | undefined }>(); + + // Collect all tool call IDs for edit tools so we can batch-query the database + const editToolCallIds: string[] = []; + + // First pass: collect tool info and identify edit tool calls + for (const e of events) { + if (e.type === 'tool.execution_start') { + const d = (e as ISessionEventToolStart).data; + if (isHiddenTool(d.toolName)) { + continue; + } + const toolArgs = d.arguments !== undefined ? tryStringify(d.arguments) : undefined; + let parameters: Record | undefined; + if (toolArgs) { + try { parameters = JSON.parse(toolArgs) as Record; } catch { /* ignore */ } + } + toolInfoByCallId.set(d.toolCallId, { toolName: d.toolName, parameters }); + if (isEditTool(d.toolName)) { + editToolCallIds.push(d.toolCallId); + } + } + } + + // Query the database for stored file edits (metadata only) + let storedEdits: Map | undefined; + if (db && editToolCallIds.length > 0) { + try { + const records = await db.getFileEdits(editToolCallIds); + if (records.length > 0) { + storedEdits = new Map(); + for (const r of records) { + let list = storedEdits.get(r.toolCallId); + if (!list) { + list = []; + storedEdits.set(r.toolCallId, list); + } + list.push(r); + } + } + } catch { + // Database may not exist yet for new sessions — that's fine + } + } + + const sessionUriStr = session.toString(); + + // Second pass: build result events + for (const e of events) { + if (e.type === 'assistant.message' || e.type === 'user.message') { + const d = (e as ISessionEventMessage).data; + result.push({ + session, + type: 'message', + role: e.type === 'user.message' ? 'user' : 'assistant', + messageId: d?.messageId ?? '', + content: d?.content ?? '', + toolRequests: d?.toolRequests?.map((tr) => ({ + toolCallId: tr.toolCallId, + name: tr.name, + arguments: tr.arguments !== undefined ? tryStringify(tr.arguments) : undefined, + type: tr.type, + })), + reasoningOpaque: d?.reasoningOpaque, + reasoningText: d?.reasoningText, + encryptedContent: d?.encryptedContent, + parentToolCallId: d?.parentToolCallId, + }); + } else if (e.type === 'tool.execution_start') { + const d = (e as ISessionEventToolStart).data; + if (isHiddenTool(d.toolName)) { + continue; + } + const info = toolInfoByCallId.get(d.toolCallId); + const displayName = getToolDisplayName(d.toolName); + const toolKind = getToolKind(d.toolName); + const toolArgs = d.arguments !== undefined ? tryStringify(d.arguments) : undefined; + result.push({ + session, + type: 'tool_start', + toolCallId: d.toolCallId, + toolName: d.toolName, + displayName, + invocationMessage: getInvocationMessage(d.toolName, displayName, info?.parameters), + toolInput: getToolInputString(d.toolName, info?.parameters, toolArgs), + toolKind, + language: toolKind === 'terminal' ? getShellLanguage(d.toolName) : undefined, + toolArguments: toolArgs, + mcpServerName: d.mcpServerName, + mcpToolName: d.mcpToolName, + parentToolCallId: d.parentToolCallId, + }); + } else if (e.type === 'tool.execution_complete') { + const d = (e as ISessionEventToolComplete).data; + const info = toolInfoByCallId.get(d.toolCallId); + if (!info) { + continue; + } + toolInfoByCallId.delete(d.toolCallId); + const displayName = getToolDisplayName(info.toolName); + const toolOutput = d.error?.message ?? d.result?.content; + const content: IToolResultContent[] = []; + if (toolOutput !== undefined) { + content.push({ type: ToolResultContentType.Text, text: toolOutput }); + } + + // Restore file edit content references from the database + const edits = storedEdits?.get(d.toolCallId); + if (edits) { + for (const edit of edits) { + content.push({ + type: ToolResultContentType.FileEdit, + beforeURI: buildSessionDbUri(sessionUriStr, edit.toolCallId, edit.filePath, 'before'), + afterURI: buildSessionDbUri(sessionUriStr, edit.toolCallId, edit.filePath, 'after'), + diff: (edit.addedLines !== undefined || edit.removedLines !== undefined) + ? { added: edit.addedLines, removed: edit.removedLines } + : undefined, + }); + } + } + + result.push({ + session, + type: 'tool_complete', + toolCallId: d.toolCallId, + result: { + success: d.success, + pastTenseMessage: getPastTenseMessage(info.toolName, displayName, info.parameters, d.success), + content: content.length > 0 ? content : undefined, + error: d.error, + }, + isUserRequested: d.isUserRequested, + toolTelemetry: d.toolTelemetry !== undefined ? tryStringify(d.toolTelemetry) : undefined, + parentToolCallId: d.parentToolCallId, + }); + } + } + return result; +} diff --git a/src/vs/platform/agentHost/node/sessionDataService.ts b/src/vs/platform/agentHost/node/sessionDataService.ts index 5bbcb89dc07..007be7ee4f9 100644 --- a/src/vs/platform/agentHost/node/sessionDataService.ts +++ b/src/vs/platform/agentHost/node/sessionDataService.ts @@ -3,11 +3,32 @@ * Licensed under the MIT License. See License.txt in the project root for license information. *--------------------------------------------------------------------------------------------*/ +import { IReference, ReferenceCollection } from '../../../base/common/lifecycle.js'; import { URI } from '../../../base/common/uri.js'; import { IFileService } from '../../files/common/files.js'; import { ILogService } from '../../log/common/log.js'; import { AgentSession } from '../common/agentService.js'; -import { ISessionDataService } from '../common/sessionDataService.js'; +import { ISessionDatabase, ISessionDataService } from '../common/sessionDataService.js'; +import { SessionDatabase } from './sessionDatabase.js'; + +class SessionDatabaseCollection extends ReferenceCollection { + constructor( + private readonly _getDbPath: (key: string) => string, + private readonly _logService: ILogService, + ) { + super(); + } + + protected createReferencedObject(key: string): ISessionDatabase { + const dbPath = this._getDbPath(key); + this._logService.trace(`[SessionDataService] Opening database: ${dbPath}`); + return new SessionDatabase(dbPath); + } + + protected destroyReferencedObject(_key: string, object: ISessionDatabase): void { + object.dispose(); + } +} /** * Implementation of {@link ISessionDataService} that stores per-session data @@ -17,6 +38,7 @@ export class SessionDataService implements ISessionDataService { declare readonly _serviceBrand: undefined; private readonly _basePath: URI; + private readonly _databases: SessionDatabaseCollection; constructor( userDataPath: URI, @@ -24,6 +46,10 @@ export class SessionDataService implements ISessionDataService { @ILogService private readonly _logService: ILogService, ) { this._basePath = URI.joinPath(userDataPath, 'agentSessionData'); + this._databases = new SessionDatabaseCollection( + key => URI.joinPath(this._basePath, key, 'session.db').fsPath, + this._logService, + ); } getSessionDataDir(session: URI): URI { @@ -35,6 +61,11 @@ export class SessionDataService implements ISessionDataService { return URI.joinPath(this._basePath, sanitized); } + openDatabase(session: URI): IReference { + const sanitized = AgentSession.id(session).replace(/[^a-zA-Z0-9_.-]/g, '-'); + return this._databases.acquire(sanitized); + } + async deleteSessionData(session: URI): Promise { const dir = this.getSessionDataDir(session); try { diff --git a/src/vs/platform/agentHost/node/sessionDatabase.ts b/src/vs/platform/agentHost/node/sessionDatabase.ts new file mode 100644 index 00000000000..2696d6a9de1 --- /dev/null +++ b/src/vs/platform/agentHost/node/sessionDatabase.ts @@ -0,0 +1,315 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import * as fs from 'fs'; +import { SequencerByKey } from '../../../base/common/async.js'; +import type { Database, RunResult } from '@vscode/sqlite3'; +import type { IFileEditContent, IFileEditRecord, ISessionDatabase } from '../common/sessionDataService.js'; +import { dirname } from '../../../base/common/path.js'; + +/** + * A single numbered migration. Migrations are applied in order of + * {@link version} and tracked via `PRAGMA user_version`. + */ +export interface ISessionDatabaseMigration { + /** Monotonically-increasing version number (1-based). */ + readonly version: number; + /** SQL to execute for this migration. */ + readonly sql: string; +} + +/** + * The set of migrations that define the current session database schema. + * New migrations should be **appended** to this array with the next version + * number. Never reorder or mutate existing entries. + */ +export const sessionDatabaseMigrations: readonly ISessionDatabaseMigration[] = [ + { + version: 1, + sql: [ + `CREATE TABLE IF NOT EXISTS turns ( + id TEXT PRIMARY KEY NOT NULL + )`, + `CREATE TABLE IF NOT EXISTS file_edits ( + turn_id TEXT NOT NULL REFERENCES turns(id) ON DELETE CASCADE, + tool_call_id TEXT NOT NULL, + file_path TEXT NOT NULL, + before_content BLOB NOT NULL, + after_content BLOB NOT NULL, + added_lines INTEGER, + removed_lines INTEGER, + PRIMARY KEY (tool_call_id, file_path) + )`, + ].join(';\n'), + }, +]; + +// ---- Promise wrappers around callback-based @vscode/sqlite3 API ----------- + +function dbExec(db: Database, sql: string): Promise { + return new Promise((resolve, reject) => { + db.exec(sql, err => err ? reject(err) : resolve()); + }); +} + +function dbRun(db: Database, sql: string, params: unknown[]): Promise<{ changes: number; lastID: number }> { + return new Promise((resolve, reject) => { + db.run(sql, params, function (this: RunResult, err: Error | null) { + if (err) { + return reject(err); + } + resolve({ changes: this.changes, lastID: this.lastID }); + }); + }); +} + +function dbGet(db: Database, sql: string, params: unknown[]): Promise | undefined> { + return new Promise((resolve, reject) => { + db.get(sql, params, (err: Error | null, row: Record | undefined) => { + if (err) { + return reject(err); + } + resolve(row); + }); + }); +} + +function dbAll(db: Database, sql: string, params: unknown[]): Promise[]> { + return new Promise((resolve, reject) => { + db.all(sql, params, (err: Error | null, rows: Record[]) => { + if (err) { + return reject(err); + } + resolve(rows); + }); + }); +} + +function dbClose(db: Database): Promise { + return new Promise((resolve, reject) => { + db.close(err => err ? reject(err) : resolve()); + }); +} + +function dbOpen(path: string): Promise { + return new Promise((resolve, reject) => { + import('@vscode/sqlite3').then(sqlite3 => { + const db = new sqlite3.default.Database(path, (err: Error | null) => { + if (err) { + return reject(err); + } + resolve(db); + }); + }, reject); + }); +} + +/** + * Applies any pending {@link ISessionDatabaseMigration migrations} to a + * database. Migrations whose version is greater than the current + * `PRAGMA user_version` are run inside a serialized transaction. After all + * migrations complete the pragma is updated to the highest applied version. + */ +async function runMigrations(db: Database, migrations: readonly ISessionDatabaseMigration[]): Promise { + // Enable foreign key enforcement — must be set outside a transaction + // and every time a connection is opened. + await dbExec(db, 'PRAGMA foreign_keys = ON'); + + const row = await dbGet(db, 'PRAGMA user_version', []); + const currentVersion = (row?.user_version as number | undefined) ?? 0; + + const pending = migrations + .filter(m => m.version > currentVersion) + .sort((a, b) => a.version - b.version); + + if (pending.length === 0) { + return; + } + + await dbExec(db, 'BEGIN TRANSACTION'); + try { + for (const migration of pending) { + await dbExec(db, migration.sql); + // PRAGMA cannot be parameterized; the version is a trusted literal. + await dbExec(db, `PRAGMA user_version = ${migration.version}`); + } + await dbExec(db, 'COMMIT'); + } catch (err) { + await dbExec(db, 'ROLLBACK'); + throw err; + } +} + +/** + * A wrapper around a `@vscode/sqlite3` {@link Database} instance with + * lazy initialisation. + * + * The underlying connection is opened on the first async method call + * (not at construction time), allowing the object to be created + * synchronously and shared via a {@link ReferenceCollection}. + * + * Calling {@link dispose} closes the connection. + */ +export class SessionDatabase implements ISessionDatabase { + + private _dbPromise: Promise | undefined; + private _disposed = false; + private readonly _fileEditSequencer = new SequencerByKey(); + + constructor( + private readonly _path: string, + private readonly _migrations: readonly ISessionDatabaseMigration[] = sessionDatabaseMigrations, + ) { } + + /** + * Opens (or creates) a SQLite database at {@link path} and applies + * any pending migrations. Only used in tests where synchronous + * construction + immediate readiness is desired. + */ + static async open(path: string, migrations: readonly ISessionDatabaseMigration[] = sessionDatabaseMigrations): Promise { + const inst = new SessionDatabase(path, migrations); + await inst._ensureDb(); + return inst; + } + + private _ensureDb(): Promise { + if (this._disposed) { + return Promise.reject(new Error('SessionDatabase has been disposed')); + } + if (!this._dbPromise) { + this._dbPromise = (async () => { + // Ensure the parent directory exists before SQLite tries to + // create the database file. + await fs.promises.mkdir(dirname(this._path), { recursive: true }); + const db = await dbOpen(this._path); + try { + await runMigrations(db, this._migrations); + } catch (err) { + await dbClose(db); + this._dbPromise = undefined; + throw err; + } + // If dispose() was called while we were opening, close immediately. + if (this._disposed) { + await dbClose(db); + throw new Error('SessionDatabase has been disposed'); + } + return db; + })(); + } + return this._dbPromise; + } + + /** + * Returns the names of all user-created tables in the database. + * Useful for testing migration behavior. + */ + async getAllTables(): Promise { + const db = await this._ensureDb(); + const rows = await dbAll(db, `SELECT name FROM sqlite_master WHERE type='table' ORDER BY name`, []); + return rows.map(r => r.name as string); + } + + // ---- Turns ---------------------------------------------------------- + + async createTurn(turnId: string): Promise { + const db = await this._ensureDb(); + await dbRun(db, 'INSERT OR IGNORE INTO turns (id) VALUES (?)', [turnId]); + } + + async deleteTurn(turnId: string): Promise { + const db = await this._ensureDb(); + await dbRun(db, 'DELETE FROM turns WHERE id = ?', [turnId]); + } + + // ---- File edits ----------------------------------------------------- + + async storeFileEdit(edit: IFileEditRecord & IFileEditContent): Promise { + return this._fileEditSequencer.queue(edit.filePath, async () => { + const db = await this._ensureDb(); + // Ensure the turn exists — the onTurnStart event that calls + // createTurn() is fire-and-forget and may not have completed yet. + await dbRun(db, 'INSERT OR IGNORE INTO turns (id) VALUES (?)', [edit.turnId]); + await dbRun( + db, + `INSERT OR REPLACE INTO file_edits + (turn_id, tool_call_id, file_path, before_content, after_content, added_lines, removed_lines) + VALUES (?, ?, ?, ?, ?, ?, ?)`, + [ + edit.turnId, + edit.toolCallId, + edit.filePath, + Buffer.from(edit.beforeContent), + Buffer.from(edit.afterContent), + edit.addedLines ?? null, + edit.removedLines ?? null, + ], + ); + }); + } + + async getFileEdits(toolCallIds: string[]): Promise { + if (toolCallIds.length === 0) { + return []; + } + const db = await this._ensureDb(); + const placeholders = toolCallIds.map(() => '?').join(','); + const rows = await dbAll( + db, + `SELECT turn_id, tool_call_id, file_path, added_lines, removed_lines + FROM file_edits + WHERE tool_call_id IN (${placeholders}) + ORDER BY rowid`, + toolCallIds, + ); + return rows.map(row => ({ + turnId: row.turn_id as string, + toolCallId: row.tool_call_id as string, + filePath: row.file_path as string, + addedLines: row.added_lines as number | undefined ?? undefined, + removedLines: row.removed_lines as number | undefined ?? undefined, + })); + } + + async readFileEditContent(toolCallId: string, filePath: string): Promise { + return this._fileEditSequencer.queue(filePath, async () => { + const db = await this._ensureDb(); + const row = await dbGet( + db, + `SELECT before_content, after_content + FROM file_edits + WHERE tool_call_id = ? AND file_path = ?`, + [toolCallId, filePath], + ); + if (!row) { + return undefined; + } + return { + beforeContent: toUint8Array(row.before_content), + afterContent: toUint8Array(row.after_content), + }; + }); + } + + dispose(): void { + if (!this._disposed) { + this._disposed = true; + this._dbPromise?.then(db => db.close()).catch(() => { }); + } + } +} + +function toUint8Array(value: unknown): Uint8Array { + if (value instanceof Buffer) { + return new Uint8Array(value.buffer, value.byteOffset, value.byteLength); + } + if (value instanceof Uint8Array) { + return value; + } + if (typeof value === 'string') { + return new TextEncoder().encode(value); + } + return new Uint8Array(0); +} diff --git a/src/vs/platform/agentHost/node/sessionStateManager.ts b/src/vs/platform/agentHost/node/sessionStateManager.ts index 531c18dd6bc..21e2091cca0 100644 --- a/src/vs/platform/agentHost/node/sessionStateManager.ts +++ b/src/vs/platform/agentHost/node/sessionStateManager.ts @@ -143,7 +143,9 @@ export class SessionStateManager extends Disposable { } /** - * Removes a session from state and emits a sessionRemoved notification. + * Removes a session from in-memory state without emitting a notification. + * Use {@link deleteSession} when the session is being permanently deleted + * and clients need to be notified. */ removeSession(session: URI): void { const state = this._sessionStates.get(session); @@ -158,7 +160,15 @@ export class SessionStateManager extends Disposable { this._sessionStates.delete(session); this._logService.trace(`[SessionStateManager] Removed session: ${session}`); + } + /** + * Permanently deletes a session from state and emits a + * {@link NotificationType.SessionRemoved} notification so that clients + * know the session is no longer accessible. + */ + deleteSession(session: URI): void { + this.removeSession(session); this._onDidEmitNotification.fire({ type: NotificationType.SessionRemoved, session, diff --git a/src/vs/platform/agentHost/test/node/agentService.test.ts b/src/vs/platform/agentHost/test/node/agentService.test.ts index 724044d8c3a..bbe7695e5e9 100644 --- a/src/vs/platform/agentHost/test/node/agentService.test.ts +++ b/src/vs/platform/agentHost/test/node/agentService.test.ts @@ -26,6 +26,7 @@ suite('AgentService (node dispatcher)', () => { _serviceBrand: undefined, getSessionDataDir: () => URI.parse('inmemory:/session-data'), getSessionDataDirById: () => URI.parse('inmemory:/session-data'), + openDatabase: () => { throw new Error('not implemented'); }, deleteSessionData: async () => { }, cleanupOrphanedData: async () => { }, }; diff --git a/src/vs/platform/agentHost/test/node/agentSideEffects.test.ts b/src/vs/platform/agentHost/test/node/agentSideEffects.test.ts index 3d36928963b..332100af045 100644 --- a/src/vs/platform/agentHost/test/node/agentSideEffects.test.ts +++ b/src/vs/platform/agentHost/test/node/agentSideEffects.test.ts @@ -74,6 +74,7 @@ suite('AgentSideEffects', () => { _serviceBrand: undefined, getSessionDataDir: () => URI.from({ scheme: Schemas.inMemory, path: '/session-data' }), getSessionDataDirById: () => URI.from({ scheme: Schemas.inMemory, path: '/session-data' }), + openDatabase: () => { throw new Error('not implemented'); }, deleteSessionData: async () => { }, cleanupOrphanedData: async () => { }, } satisfies ISessionDataService, diff --git a/src/vs/platform/agentHost/test/node/fileEditTracker.test.ts b/src/vs/platform/agentHost/test/node/fileEditTracker.test.ts index 8679dab91e4..54955b60bf6 100644 --- a/src/vs/platform/agentHost/test/node/fileEditTracker.test.ts +++ b/src/vs/platform/agentHost/test/node/fileEditTracker.test.ts @@ -4,89 +4,161 @@ *--------------------------------------------------------------------------------------------*/ import assert from 'assert'; +import { tmpdir } from 'os'; +import { randomUUID } from 'crypto'; +import { mkdirSync, rmSync } from 'fs'; import { VSBuffer } from '../../../../base/common/buffer.js'; import { DisposableStore } from '../../../../base/common/lifecycle.js'; -import { Schemas } from '../../../../base/common/network.js'; import { URI } from '../../../../base/common/uri.js'; import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../base/test/common/utils.js'; import { FileService } from '../../../files/common/fileService.js'; import { InMemoryFileSystemProvider } from '../../../files/common/inMemoryFilesystemProvider.js'; import { NullLogService } from '../../../log/common/log.js'; -import { ISessionDataService } from '../../common/sessionDataService.js'; import { ToolResultContentType } from '../../common/state/sessionState.js'; -import { SessionDataService } from '../../node/sessionDataService.js'; -import { FileEditTracker } from '../../node/copilot/fileEditTracker.js'; +import { SessionDatabase } from '../../node/sessionDatabase.js'; +import { FileEditTracker, buildSessionDbUri, parseSessionDbUri } from '../../node/copilot/fileEditTracker.js'; +import { join } from '../../../../base/common/path.js'; suite('FileEditTracker', () => { const disposables = new DisposableStore(); let fileService: FileService; - let sessionDataService: ISessionDataService; + let db: SessionDatabase; let tracker: FileEditTracker; + let testDir: string; - const basePath = URI.from({ scheme: Schemas.inMemory, path: '/userData' }); + setup(async () => { + testDir = join(tmpdir(), `vscode-edit-tracker-test-${randomUUID()}`); + mkdirSync(testDir, { recursive: true }); - setup(() => { fileService = disposables.add(new FileService(new NullLogService())); - disposables.add(fileService.registerProvider(Schemas.inMemory, disposables.add(new InMemoryFileSystemProvider()))); - sessionDataService = new SessionDataService(basePath, fileService, new NullLogService()); - tracker = new FileEditTracker('test-session', sessionDataService, fileService, new NullLogService()); + const sourceFs = disposables.add(new InMemoryFileSystemProvider()); + disposables.add(fileService.registerProvider('file', sourceFs)); + + db = disposables.add(await SessionDatabase.open(join(testDir, 'session.db'))); + await db.createTurn('turn-1'); + + tracker = new FileEditTracker('copilot:/test-session', db, fileService, new NullLogService()); }); - teardown(() => disposables.clear()); + teardown(() => { + disposables.clear(); + rmSync(testDir, { recursive: true, force: true }); + }); ensureNoDisposablesAreLeakedInTestSuite(); test('tracks edit start and complete for existing file', async () => { - const sourceFs = disposables.add(new InMemoryFileSystemProvider()); - disposables.add(fileService.registerProvider(Schemas.file, sourceFs)); await fileService.writeFile(URI.file('/workspace/test.txt'), VSBuffer.fromString('original content\nline 2')); await tracker.trackEditStart('/workspace/test.txt'); await fileService.writeFile(URI.file('/workspace/test.txt'), VSBuffer.fromString('modified content\nline 2\nline 3')); await tracker.completeEdit('/workspace/test.txt'); - const fileEdit = tracker.takeCompletedEdit('/workspace/test.txt'); + const fileEdit = await tracker.takeCompletedEdit('turn-1', 'tc-1', '/workspace/test.txt'); assert.ok(fileEdit); assert.strictEqual(fileEdit.type, ToolResultContentType.FileEdit); - // Both URIs point to snapshots in the session data directory - const sessionDir = sessionDataService.getSessionDataDirById('test-session'); - assert.ok(fileEdit.beforeURI.startsWith(sessionDir.toString())); - assert.ok(fileEdit.afterURI.startsWith(sessionDir.toString())); + + // URIs are parseable session-db: URIs + const beforeFields = parseSessionDbUri(fileEdit.beforeURI); + assert.ok(beforeFields); + assert.strictEqual(beforeFields.sessionUri, 'copilot:/test-session'); + assert.strictEqual(beforeFields.toolCallId, 'tc-1'); + assert.strictEqual(beforeFields.filePath, '/workspace/test.txt'); + assert.strictEqual(beforeFields.part, 'before'); + + const afterFields = parseSessionDbUri(fileEdit.afterURI); + assert.ok(afterFields); + assert.strictEqual(afterFields.part, 'after'); + + // Content is persisted in the database (wait for fire-and-forget write) + await new Promise(r => setTimeout(r, 50)); + + const content = await db.readFileEditContent('tc-1', '/workspace/test.txt'); + assert.ok(content); + assert.strictEqual(new TextDecoder().decode(content.beforeContent), 'original content\nline 2'); + assert.strictEqual(new TextDecoder().decode(content.afterContent), 'modified content\nline 2\nline 3'); }); test('tracks edit for newly created file (no before content)', async () => { - const sourceFs = disposables.add(new InMemoryFileSystemProvider()); - disposables.add(fileService.registerProvider(Schemas.file, sourceFs)); - await tracker.trackEditStart('/workspace/new-file.txt'); await fileService.writeFile(URI.file('/workspace/new-file.txt'), VSBuffer.fromString('new file\ncontent')); await tracker.completeEdit('/workspace/new-file.txt'); - const fileEdit = tracker.takeCompletedEdit('/workspace/new-file.txt'); + const fileEdit = await tracker.takeCompletedEdit('turn-1', 'tc-2', '/workspace/new-file.txt'); assert.ok(fileEdit); - const sessionDir = sessionDataService.getSessionDataDirById('test-session'); - assert.ok(fileEdit.afterURI.startsWith(sessionDir.toString())); + + // Wait for the fire-and-forget DB write to complete + await new Promise(r => setTimeout(r, 50)); + + const content = await db.readFileEditContent('tc-2', '/workspace/new-file.txt'); + assert.ok(content); + assert.strictEqual(new TextDecoder().decode(content.beforeContent), ''); + assert.strictEqual(new TextDecoder().decode(content.afterContent), 'new file\ncontent'); }); - test('takeCompletedEdit returns undefined for unknown file path', () => { - const result = tracker.takeCompletedEdit('/nonexistent'); + test('takeCompletedEdit returns undefined for unknown file path', async () => { + const result = await tracker.takeCompletedEdit('turn-1', 'tc-x', '/nonexistent'); assert.strictEqual(result, undefined); }); - test('before and after snapshot content can be read back', async () => { - const sourceFs = disposables.add(new InMemoryFileSystemProvider()); - disposables.add(fileService.registerProvider(Schemas.file, sourceFs)); + test('before and after content can be read from database', async () => { await fileService.writeFile(URI.file('/workspace/file.ts'), VSBuffer.fromString('original')); await tracker.trackEditStart('/workspace/file.ts'); await fileService.writeFile(URI.file('/workspace/file.ts'), VSBuffer.fromString('modified')); await tracker.completeEdit('/workspace/file.ts'); - const fileEdit = tracker.takeCompletedEdit('/workspace/file.ts'); - assert.ok(fileEdit); - const beforeContent = await fileService.readFile(URI.parse(fileEdit.beforeURI)); - assert.strictEqual(beforeContent.value.toString(), 'original'); - const afterContent = await fileService.readFile(URI.parse(fileEdit.afterURI)); - assert.strictEqual(afterContent.value.toString(), 'modified'); + await tracker.takeCompletedEdit('turn-1', 'tc-3', '/workspace/file.ts'); + + // Wait for the fire-and-forget DB write to complete + await new Promise(r => setTimeout(r, 50)); + + const content = await db.readFileEditContent('tc-3', '/workspace/file.ts'); + assert.ok(content); + assert.strictEqual(new TextDecoder().decode(content.beforeContent), 'original'); + assert.strictEqual(new TextDecoder().decode(content.afterContent), 'modified'); + }); +}); + +suite('buildSessionDbUri / parseSessionDbUri', () => { + + ensureNoDisposablesAreLeakedInTestSuite(); + + test('round-trips a simple URI', () => { + const uri = buildSessionDbUri('copilot:/abc-123', 'tc-1', '/workspace/file.ts', 'before'); + const parsed = parseSessionDbUri(uri); + assert.ok(parsed); + assert.deepStrictEqual(parsed, { + sessionUri: 'copilot:/abc-123', + toolCallId: 'tc-1', + filePath: '/workspace/file.ts', + part: 'before', + }); + }); + + test('round-trips with special characters in filePath', () => { + const uri = buildSessionDbUri('copilot:/s1', 'tc-2', '/work space/file (1).ts', 'after'); + const parsed = parseSessionDbUri(uri); + assert.ok(parsed); + assert.strictEqual(parsed.filePath, '/work space/file (1).ts'); + assert.strictEqual(parsed.part, 'after'); + }); + + test('round-trips with special characters in toolCallId', () => { + const uri = buildSessionDbUri('copilot:/s1', 'call_abc=123&x', '/file.ts', 'before'); + const parsed = parseSessionDbUri(uri); + assert.ok(parsed); + assert.strictEqual(parsed.toolCallId, 'call_abc=123&x'); + }); + + test('parseSessionDbUri returns undefined for non-session-db URIs', () => { + assert.strictEqual(parseSessionDbUri('file:///foo/bar'), undefined); + assert.strictEqual(parseSessionDbUri('https://example.com'), undefined); + }); + + test('parseSessionDbUri returns undefined for malformed session-db URIs', () => { + assert.strictEqual(parseSessionDbUri('session-db:copilot:/s1'), undefined); + assert.strictEqual(parseSessionDbUri('session-db:copilot:/s1?toolCallId=tc-1'), undefined); + assert.strictEqual(parseSessionDbUri('session-db:copilot:/s1?toolCallId=tc-1&filePath=/f&part=middle'), undefined); }); }); diff --git a/src/vs/platform/agentHost/test/node/mapSessionEvents.test.ts b/src/vs/platform/agentHost/test/node/mapSessionEvents.test.ts new file mode 100644 index 00000000000..d61905c5458 --- /dev/null +++ b/src/vs/platform/agentHost/test/node/mapSessionEvents.test.ts @@ -0,0 +1,238 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import assert from 'assert'; +import { tmpdir } from 'os'; +import { randomUUID } from 'crypto'; +import { mkdirSync, rmSync } from 'fs'; +import { DisposableStore } from '../../../../base/common/lifecycle.js'; +import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../base/test/common/utils.js'; +import { AgentSession } from '../../common/agentService.js'; +import { ToolResultContentType } from '../../common/state/sessionState.js'; +import { SessionDatabase } from '../../node/sessionDatabase.js'; +import { parseSessionDbUri } from '../../node/copilot/fileEditTracker.js'; +import { mapSessionEvents, type ISessionEvent } from '../../node/copilot/mapSessionEvents.js'; +import { join } from '../../../../base/common/path.js'; + +suite('mapSessionEvents', () => { + + const disposables = new DisposableStore(); + let testDir: string; + const session = AgentSession.uri('copilot', 'test-session'); + + setup(() => { + testDir = join(tmpdir(), `vscode-map-events-test-${randomUUID()}`); + mkdirSync(testDir, { recursive: true }); + }); + + teardown(() => { + disposables.clear(); + rmSync(testDir, { recursive: true, force: true }); + }); + ensureNoDisposablesAreLeakedInTestSuite(); + + function dbPath(): string { + return join(testDir, 'session.db'); + } + + // ---- Basic event mapping -------------------------------------------- + + test('maps user and assistant messages', async () => { + const events: ISessionEvent[] = [ + { type: 'user.message', data: { messageId: 'msg-1', content: 'hello' } }, + { type: 'assistant.message', data: { messageId: 'msg-2', content: 'world' } }, + ]; + + const result = await mapSessionEvents(session, undefined, events); + assert.strictEqual(result.length, 2); + assert.deepStrictEqual(result[0], { + session, + type: 'message', + role: 'user', + messageId: 'msg-1', + content: 'hello', + toolRequests: undefined, + reasoningOpaque: undefined, + reasoningText: undefined, + encryptedContent: undefined, + parentToolCallId: undefined, + }); + assert.strictEqual(result[1].type, 'message'); + assert.strictEqual((result[1] as { role: string }).role, 'assistant'); + }); + + test('maps tool start and complete events', async () => { + const events: ISessionEvent[] = [ + { + type: 'tool.execution_start', + data: { toolCallId: 'tc-1', toolName: 'shell', arguments: { command: 'echo hi' } }, + }, + { + type: 'tool.execution_complete', + data: { toolCallId: 'tc-1', success: true, result: { content: 'hi\n' } }, + }, + ]; + + const result = await mapSessionEvents(session, undefined, events); + assert.strictEqual(result.length, 2); + assert.strictEqual(result[0].type, 'tool_start'); + assert.strictEqual(result[1].type, 'tool_complete'); + + const complete = result[1] as { result: { content?: readonly { type: string; text?: string }[] } }; + assert.ok(complete.result.content); + assert.strictEqual(complete.result.content[0].type, ToolResultContentType.Text); + }); + + test('skips tool_complete without matching tool_start', async () => { + const events: ISessionEvent[] = [ + { type: 'tool.execution_complete', data: { toolCallId: 'orphan', success: true } }, + ]; + + const result = await mapSessionEvents(session, undefined, events); + assert.strictEqual(result.length, 0); + }); + + test('ignores unknown event types', async () => { + const events: ISessionEvent[] = [ + { type: 'some.unknown.event', data: {} }, + { type: 'user.message', data: { messageId: 'msg-1', content: 'test' } }, + ]; + + const result = await mapSessionEvents(session, undefined, events); + assert.strictEqual(result.length, 1); + }); + + // ---- File edit restoration ------------------------------------------ + + suite('file edit restoration', () => { + + test('restores file edits from database for edit tools', async () => { + const db = disposables.add(await SessionDatabase.open(dbPath())); + await db.createTurn('turn-1'); + await db.storeFileEdit({ + turnId: 'turn-1', + toolCallId: 'tc-edit', + filePath: '/workspace/file.ts', + beforeContent: new TextEncoder().encode('before'), + afterContent: new TextEncoder().encode('after'), + addedLines: 3, + removedLines: 1, + }); + + const events: ISessionEvent[] = [ + { + type: 'tool.execution_start', + data: { toolCallId: 'tc-edit', toolName: 'edit', arguments: { filePath: '/workspace/file.ts' } }, + }, + { + type: 'tool.execution_complete', + data: { toolCallId: 'tc-edit', success: true, result: { content: 'Edited file.ts' } }, + }, + ]; + + const result = await mapSessionEvents(session, db, events); + const complete = result[1]; + assert.strictEqual(complete.type, 'tool_complete'); + + const content = (complete as { result: { content?: readonly Record[] } }).result.content; + assert.ok(content); + // Should have text content + file edit + assert.strictEqual(content.length, 2); + assert.strictEqual(content[0].type, ToolResultContentType.Text); + assert.strictEqual(content[1].type, ToolResultContentType.FileEdit); + + // File edit URIs should be parseable + const fileEdit = content[1] as { beforeURI: string; afterURI: string; diff?: { added?: number; removed?: number } }; + const beforeFields = parseSessionDbUri(fileEdit.beforeURI); + assert.ok(beforeFields); + assert.strictEqual(beforeFields.toolCallId, 'tc-edit'); + assert.strictEqual(beforeFields.filePath, '/workspace/file.ts'); + assert.strictEqual(beforeFields.part, 'before'); + assert.deepStrictEqual(fileEdit.diff, { added: 3, removed: 1 }); + }); + + test('handles multiple file edits for one tool call', async () => { + const db = disposables.add(await SessionDatabase.open(dbPath())); + await db.createTurn('turn-1'); + await db.storeFileEdit({ + turnId: 'turn-1', + toolCallId: 'tc-multi', + filePath: '/workspace/a.ts', + beforeContent: new Uint8Array(0), + afterContent: new TextEncoder().encode('a'), + addedLines: undefined, + removedLines: undefined, + }); + await db.storeFileEdit({ + turnId: 'turn-1', + toolCallId: 'tc-multi', + filePath: '/workspace/b.ts', + beforeContent: new Uint8Array(0), + afterContent: new TextEncoder().encode('b'), + addedLines: undefined, + removedLines: undefined, + }); + + const events: ISessionEvent[] = [ + { + type: 'tool.execution_start', + data: { toolCallId: 'tc-multi', toolName: 'write' }, + }, + { + type: 'tool.execution_complete', + data: { toolCallId: 'tc-multi', success: true }, + }, + ]; + + const result = await mapSessionEvents(session, db, events); + const content = (result[1] as { result: { content?: readonly Record[] } }).result.content; + assert.ok(content); + // Two file edits (no text since result had no content) + const fileEdits = content.filter(c => c.type === ToolResultContentType.FileEdit); + assert.strictEqual(fileEdits.length, 2); + }); + + test('works without database (no file edits restored)', async () => { + const events: ISessionEvent[] = [ + { + type: 'tool.execution_start', + data: { toolCallId: 'tc-1', toolName: 'edit', arguments: { filePath: '/workspace/file.ts' } }, + }, + { + type: 'tool.execution_complete', + data: { toolCallId: 'tc-1', success: true, result: { content: 'done' } }, + }, + ]; + + const result = await mapSessionEvents(session, undefined, events); + const content = (result[1] as { result: { content?: readonly Record[] } }).result.content; + assert.ok(content); + // Only text content, no file edits + assert.strictEqual(content.length, 1); + assert.strictEqual(content[0].type, ToolResultContentType.Text); + }); + + test('non-edit tools do not get file edits even if db has data', async () => { + const db = disposables.add(await SessionDatabase.open(dbPath())); + + const events: ISessionEvent[] = [ + { + type: 'tool.execution_start', + data: { toolCallId: 'tc-1', toolName: 'shell', arguments: { command: 'ls' } }, + }, + { + type: 'tool.execution_complete', + data: { toolCallId: 'tc-1', success: true, result: { content: 'files' } }, + }, + ]; + + const result = await mapSessionEvents(session, db, events); + const content = (result[1] as { result: { content?: readonly Record[] } }).result.content; + assert.ok(content); + assert.strictEqual(content.length, 1); + assert.strictEqual(content[0].type, ToolResultContentType.Text); + }); + }); +}); diff --git a/src/vs/platform/agentHost/test/node/sessionDataService.test.ts b/src/vs/platform/agentHost/test/node/sessionDataService.test.ts index 5adefa25f08..9675c4dac0b 100644 --- a/src/vs/platform/agentHost/test/node/sessionDataService.test.ts +++ b/src/vs/platform/agentHost/test/node/sessionDataService.test.ts @@ -4,16 +4,21 @@ *--------------------------------------------------------------------------------------------*/ import assert from 'assert'; +import { tmpdir } from 'os'; +import { randomUUID } from 'crypto'; +import { mkdirSync, rmSync } from 'fs'; import { VSBuffer } from '../../../../base/common/buffer.js'; import { DisposableStore } from '../../../../base/common/lifecycle.js'; import { Schemas } from '../../../../base/common/network.js'; import { URI } from '../../../../base/common/uri.js'; import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../base/test/common/utils.js'; import { FileService } from '../../../files/common/fileService.js'; +import { DiskFileSystemProvider } from '../../../files/node/diskFileSystemProvider.js'; import { InMemoryFileSystemProvider } from '../../../files/common/inMemoryFilesystemProvider.js'; import { NullLogService } from '../../../log/common/log.js'; import { AgentSession } from '../../common/agentService.js'; import { SessionDataService } from '../../node/sessionDataService.js'; +import { join } from '../../../../base/common/path.js'; suite('SessionDataService', () => { @@ -80,3 +85,72 @@ suite('SessionDataService', () => { await service.cleanupOrphanedData(new Set()); }); }); + +suite('SessionDataService — openDatabase ref-counting', () => { + + const disposables = new DisposableStore(); + let service: SessionDataService; + let testDir: string; + + setup(() => { + testDir = join(tmpdir(), `vscode-session-data-test-${randomUUID()}`); + mkdirSync(testDir, { recursive: true }); + + const fileService = disposables.add(new FileService(new NullLogService())); + disposables.add(fileService.registerProvider(Schemas.file, disposables.add(new DiskFileSystemProvider(new NullLogService())))); + service = new SessionDataService(URI.file(testDir), fileService, new NullLogService()); + }); + + teardown(() => { + disposables.clear(); + rmSync(testDir, { recursive: true, force: true }); + }); + ensureNoDisposablesAreLeakedInTestSuite(); + + test('returns a functional database reference', async () => { + const session = AgentSession.uri('copilot', 'ref-test'); + const ref = service.openDatabase(session); + disposables.add(ref); + + await ref.object.createTurn('turn-1'); + const edits = await ref.object.getFileEdits([]); + assert.deepStrictEqual(edits, []); + }); + + test('multiple references share the same database', async () => { + const session = AgentSession.uri('copilot', 'shared-test'); + const ref1 = service.openDatabase(session); + const ref2 = service.openDatabase(session); + + assert.strictEqual(ref1.object, ref2.object); + + ref1.dispose(); + ref2.dispose(); + }); + + test('database remains usable until last reference is disposed', async () => { + const session = AgentSession.uri('copilot', 'refcount-test'); + const ref1 = service.openDatabase(session); + const ref2 = service.openDatabase(session); + + ref1.dispose(); + + // ref2 still works + await ref2.object.createTurn('turn-1'); + + ref2.dispose(); + }); + + test('new reference after all disposed gets a fresh database', async () => { + const session = AgentSession.uri('copilot', 'reopen-test'); + const ref1 = service.openDatabase(session); + const db1 = ref1.object; + ref1.dispose(); + + const ref2 = service.openDatabase(session); + disposables.add(ref2); + // New reference — may or may not be the same object, but must be functional + await ref2.object.createTurn('turn-1'); + assert.notStrictEqual(ref2.object, db1); + }); +}); diff --git a/src/vs/platform/agentHost/test/node/sessionDatabase.test.ts b/src/vs/platform/agentHost/test/node/sessionDatabase.test.ts new file mode 100644 index 00000000000..fe7bfe96b6b --- /dev/null +++ b/src/vs/platform/agentHost/test/node/sessionDatabase.test.ts @@ -0,0 +1,421 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import assert from 'assert'; +import { tmpdir } from 'os'; +import { randomUUID } from 'crypto'; +import { mkdirSync, rmSync } from 'fs'; +import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../base/test/common/utils.js'; +import { DisposableStore } from '../../../../base/common/lifecycle.js'; +import { SessionDatabase, type ISessionDatabaseMigration } from '../../node/sessionDatabase.js'; +import { join } from '../../../../base/common/path.js'; + +suite('SessionDatabase', () => { + + const disposables = new DisposableStore(); + let testDir: string; + + setup(() => { + testDir = join(tmpdir(), `vscode-session-db-test-${randomUUID()}`); + mkdirSync(testDir, { recursive: true }); + }); + + teardown(() => { + disposables.clear(); + rmSync(testDir, { recursive: true, force: true }); + }); + ensureNoDisposablesAreLeakedInTestSuite(); + + function dbPath(name = 'session.db'): string { + return join(testDir, name); + } + + // ---- Migration system ----------------------------------------------- + + suite('migrations', () => { + + test('applies all migrations on a fresh database', async () => { + const migrations: ISessionDatabaseMigration[] = [ + { version: 1, sql: 'CREATE TABLE t1 (id INTEGER PRIMARY KEY)' }, + { version: 2, sql: 'CREATE TABLE t2 (id INTEGER PRIMARY KEY)' }, + ]; + + const db = disposables.add(await SessionDatabase.open(dbPath(), migrations)); + + const tables = (await db.getAllTables()).sort(); + assert.deepStrictEqual(tables, ['t1', 't2']); + }); + + test('reopening with same migrations is a no-op', async () => { + const migrations: ISessionDatabaseMigration[] = [ + { version: 1, sql: 'CREATE TABLE t1 (id INTEGER PRIMARY KEY)' }, + ]; + + const db1 = await SessionDatabase.open(dbPath(), migrations); + db1.dispose(); + + // Reopen — should not throw (table already exists, migration skipped) + const db2 = disposables.add(await SessionDatabase.open(dbPath(), migrations)); + assert.deepStrictEqual(await db2.getAllTables(), ['t1']); + }); + + test('only applies new migrations on reopen', async () => { + const v1: ISessionDatabaseMigration[] = [ + { version: 1, sql: 'CREATE TABLE t1 (id INTEGER PRIMARY KEY)' }, + ]; + const db1 = await SessionDatabase.open(dbPath(), v1); + db1.dispose(); + + const v2: ISessionDatabaseMigration[] = [ + ...v1, + { version: 2, sql: 'CREATE TABLE t2 (id INTEGER PRIMARY KEY)' }, + ]; + const db2 = disposables.add(await SessionDatabase.open(dbPath(), v2)); + + const tables = (await db2.getAllTables()).sort(); + assert.deepStrictEqual(tables, ['t1', 't2']); + }); + + test('rolls back on migration failure', async () => { + const migrations: ISessionDatabaseMigration[] = [ + { version: 1, sql: 'CREATE TABLE t1 (id INTEGER PRIMARY KEY)' }, + { version: 2, sql: 'THIS IS INVALID SQL' }, + ]; + + await assert.rejects(() => SessionDatabase.open(dbPath(), migrations)); + + // Reopen with only v1 — t1 should not exist because the whole + // transaction was rolled back + const db = disposables.add(await SessionDatabase.open(dbPath(), [ + { version: 1, sql: 'CREATE TABLE t1 (id INTEGER PRIMARY KEY)' }, + ])); + assert.deepStrictEqual(await db.getAllTables(), ['t1']); + }); + }); + + // ---- File edits ----------------------------------------------------- + + suite('file edits', () => { + + test('store and retrieve a file edit', async () => { + const db = disposables.add(await SessionDatabase.open(dbPath())); + + await db.createTurn('turn-1'); + await db.storeFileEdit({ + turnId: 'turn-1', + toolCallId: 'tc-1', + filePath: '/workspace/file.ts', + beforeContent: new TextEncoder().encode('before'), + afterContent: new TextEncoder().encode('after'), + addedLines: 5, + removedLines: 2, + }); + + const edits = await db.getFileEdits(['tc-1']); + assert.deepStrictEqual(edits, [{ + turnId: 'turn-1', + toolCallId: 'tc-1', + filePath: '/workspace/file.ts', + addedLines: 5, + removedLines: 2, + }]); + }); + + test('retrieve multiple edits for a single tool call', async () => { + const db = disposables.add(await SessionDatabase.open(dbPath())); + + await db.createTurn('turn-1'); + await db.storeFileEdit({ + turnId: 'turn-1', + toolCallId: 'tc-1', + filePath: '/workspace/a.ts', + beforeContent: new TextEncoder().encode('a-before'), + afterContent: new TextEncoder().encode('a-after'), + addedLines: undefined, + removedLines: undefined, + }); + await db.storeFileEdit({ + turnId: 'turn-1', + toolCallId: 'tc-1', + filePath: '/workspace/b.ts', + beforeContent: new TextEncoder().encode('b-before'), + afterContent: new TextEncoder().encode('b-after'), + addedLines: 1, + removedLines: 0, + }); + + const edits = await db.getFileEdits(['tc-1']); + assert.strictEqual(edits.length, 2); + assert.strictEqual(edits[0].filePath, '/workspace/a.ts'); + assert.strictEqual(edits[1].filePath, '/workspace/b.ts'); + }); + + test('retrieve edits across multiple tool calls', async () => { + const db = disposables.add(await SessionDatabase.open(dbPath())); + + await db.createTurn('turn-1'); + await db.storeFileEdit({ + turnId: 'turn-1', + toolCallId: 'tc-1', + filePath: '/workspace/a.ts', + beforeContent: new Uint8Array(0), + afterContent: new TextEncoder().encode('hello'), + addedLines: undefined, + removedLines: undefined, + }); + await db.storeFileEdit({ + turnId: 'turn-1', + toolCallId: 'tc-2', + filePath: '/workspace/b.ts', + beforeContent: new Uint8Array(0), + afterContent: new TextEncoder().encode('world'), + addedLines: undefined, + removedLines: undefined, + }); + + const edits = await db.getFileEdits(['tc-1', 'tc-2']); + assert.strictEqual(edits.length, 2); + + // Only tc-2 + const edits2 = await db.getFileEdits(['tc-2']); + assert.strictEqual(edits2.length, 1); + assert.strictEqual(edits2[0].toolCallId, 'tc-2'); + }); + + test('returns empty array for unknown tool call IDs', async () => { + const db = disposables.add(await SessionDatabase.open(dbPath())); + const edits = await db.getFileEdits(['nonexistent']); + assert.deepStrictEqual(edits, []); + }); + + test('returns empty array when given empty array', async () => { + const db = disposables.add(await SessionDatabase.open(dbPath())); + const edits = await db.getFileEdits([]); + assert.deepStrictEqual(edits, []); + }); + + test('replace on conflict (same toolCallId + filePath)', async () => { + const db = disposables.add(await SessionDatabase.open(dbPath())); + + await db.createTurn('turn-1'); + await db.storeFileEdit({ + turnId: 'turn-1', + toolCallId: 'tc-1', + filePath: '/workspace/file.ts', + beforeContent: new TextEncoder().encode('v1'), + afterContent: new TextEncoder().encode('v1-after'), + addedLines: 1, + removedLines: 0, + }); + await db.storeFileEdit({ + turnId: 'turn-1', + toolCallId: 'tc-1', + filePath: '/workspace/file.ts', + beforeContent: new TextEncoder().encode('v2'), + afterContent: new TextEncoder().encode('v2-after'), + addedLines: 3, + removedLines: 1, + }); + + const edits = await db.getFileEdits(['tc-1']); + assert.strictEqual(edits.length, 1); + assert.strictEqual(edits[0].addedLines, 3); + + const content = await db.readFileEditContent('tc-1', '/workspace/file.ts'); + assert.ok(content); + assert.deepStrictEqual(new TextDecoder().decode(content.beforeContent), 'v2'); + }); + + test('readFileEditContent returns content on demand', async () => { + const db = disposables.add(await SessionDatabase.open(dbPath())); + + await db.createTurn('turn-1'); + await db.storeFileEdit({ + turnId: 'turn-1', + toolCallId: 'tc-1', + filePath: '/workspace/file.ts', + beforeContent: new TextEncoder().encode('before'), + afterContent: new TextEncoder().encode('after'), + addedLines: undefined, + removedLines: undefined, + }); + + const content = await db.readFileEditContent('tc-1', '/workspace/file.ts'); + assert.ok(content); + assert.deepStrictEqual(content.beforeContent, new TextEncoder().encode('before')); + assert.deepStrictEqual(content.afterContent, new TextEncoder().encode('after')); + }); + + test('readFileEditContent returns undefined for missing edit', async () => { + const db = disposables.add(await SessionDatabase.open(dbPath())); + const content = await db.readFileEditContent('tc-missing', '/no/such/file'); + assert.strictEqual(content, undefined); + }); + + test('persists binary content correctly', async () => { + const db = disposables.add(await SessionDatabase.open(dbPath())); + const binary = new Uint8Array([0, 1, 2, 255, 128, 64]); + + await db.createTurn('turn-1'); + await db.storeFileEdit({ + turnId: 'turn-1', + toolCallId: 'tc-bin', + filePath: '/workspace/image.png', + beforeContent: new Uint8Array(0), + afterContent: binary, + addedLines: undefined, + removedLines: undefined, + }); + + const content = await db.readFileEditContent('tc-bin', '/workspace/image.png'); + assert.ok(content); + assert.deepStrictEqual(content.afterContent, binary); + }); + + test('auto-creates turn if it does not exist', async () => { + const db = disposables.add(await SessionDatabase.open(dbPath())); + + // storeFileEdit should succeed even without a prior createTurn call + await db.storeFileEdit({ + turnId: 'auto-turn', + toolCallId: 'tc-1', + filePath: '/x', + beforeContent: new Uint8Array(0), + afterContent: new Uint8Array(0), + addedLines: undefined, + removedLines: undefined, + }); + + const edits = await db.getFileEdits(['tc-1']); + assert.strictEqual(edits.length, 1); + assert.strictEqual(edits[0].turnId, 'auto-turn'); + }); + }); + + // ---- Turns ---------------------------------------------------------- + + suite('turns', () => { + + test('createTurn is idempotent', async () => { + const db = disposables.add(await SessionDatabase.open(dbPath())); + await db.createTurn('turn-1'); + await db.createTurn('turn-1'); // should not throw + }); + + test('deleteTurn cascades to file edits', async () => { + const db = disposables.add(await SessionDatabase.open(dbPath())); + + await db.createTurn('turn-1'); + await db.storeFileEdit({ + turnId: 'turn-1', + toolCallId: 'tc-1', + filePath: '/workspace/a.ts', + beforeContent: new TextEncoder().encode('before'), + afterContent: new TextEncoder().encode('after'), + addedLines: undefined, + removedLines: undefined, + }); + + // Edits exist + assert.strictEqual((await db.getFileEdits(['tc-1'])).length, 1); + + // Delete the turn — edits should be gone + await db.deleteTurn('turn-1'); + assert.deepStrictEqual(await db.getFileEdits(['tc-1']), []); + }); + + test('deleteTurn only removes its own edits', async () => { + const db = disposables.add(await SessionDatabase.open(dbPath())); + + await db.createTurn('turn-1'); + await db.createTurn('turn-2'); + await db.storeFileEdit({ + turnId: 'turn-1', + toolCallId: 'tc-1', + filePath: '/workspace/a.ts', + beforeContent: new Uint8Array(0), + afterContent: new TextEncoder().encode('a'), + addedLines: undefined, + removedLines: undefined, + }); + await db.storeFileEdit({ + turnId: 'turn-2', + toolCallId: 'tc-2', + filePath: '/workspace/b.ts', + beforeContent: new Uint8Array(0), + afterContent: new TextEncoder().encode('b'), + addedLines: undefined, + removedLines: undefined, + }); + + await db.deleteTurn('turn-1'); + + assert.deepStrictEqual(await db.getFileEdits(['tc-1']), []); + assert.strictEqual((await db.getFileEdits(['tc-2'])).length, 1); + }); + + test('deleteTurn is a no-op for unknown turn', async () => { + const db = disposables.add(await SessionDatabase.open(dbPath())); + await db.deleteTurn('nonexistent'); // should not throw + }); + }); + + // ---- Dispose -------------------------------------------------------- + + suite('dispose', () => { + + test('methods throw after dispose', async () => { + const db = await SessionDatabase.open(dbPath()); + db.dispose(); + + await assert.rejects( + () => db.createTurn('turn-1'), + /disposed/, + ); + }); + + test('double dispose is safe', async () => { + const db = await SessionDatabase.open(dbPath()); + db.dispose(); + db.dispose(); // should not throw + }); + }); + + // ---- Lazy open ------------------------------------------------------ + + suite('lazy open', () => { + + test('constructor does not open the database', () => { + // Should not throw even if path does not exist yet + const db = new SessionDatabase(join(testDir, 'lazy', 'session.db')); + disposables.add(db); + // No error — the database is not opened until first use + }); + + test('first async call opens and migrates the database', async () => { + const db = disposables.add(new SessionDatabase(dbPath())); + // Database file may not exist yet — first call triggers open + await db.createTurn('turn-1'); + const edits = await db.getFileEdits(['nonexistent']); + assert.deepStrictEqual(edits, []); + }); + + test('multiple concurrent calls share the same open promise', async () => { + const db = disposables.add(new SessionDatabase(dbPath())); + // Fire multiple calls concurrently — all should succeed + await Promise.all([ + db.createTurn('turn-1'), + db.createTurn('turn-2'), + db.getFileEdits([]), + ]); + }); + + test('dispose during open rejects subsequent calls', async () => { + const db = new SessionDatabase(dbPath()); + db.dispose(); + await assert.rejects(() => db.createTurn('turn-1'), /disposed/); + }); + }); +}); diff --git a/src/vs/platform/agentHost/test/node/sessionStateManager.test.ts b/src/vs/platform/agentHost/test/node/sessionStateManager.test.ts index 852e8cf58fc..e1e7a08830f 100644 --- a/src/vs/platform/agentHost/test/node/sessionStateManager.test.ts +++ b/src/vs/platform/agentHost/test/node/sessionStateManager.test.ts @@ -121,7 +121,7 @@ suite('SessionStateManager', () => { assert.deepStrictEqual(envelopes[0].origin, origin); }); - test('removeSession clears state and emits notification', () => { + test('removeSession clears state without notification', () => { manager.createSession(makeSessionSummary()); const notifications: INotification[] = []; @@ -129,6 +129,19 @@ suite('SessionStateManager', () => { manager.removeSession(sessionUri); + assert.strictEqual(manager.getSessionState(sessionUri), undefined); + assert.strictEqual(manager.getSnapshot(sessionUri), undefined); + assert.strictEqual(notifications.length, 0); + }); + + test('deleteSession clears state and emits notification', () => { + manager.createSession(makeSessionSummary()); + + const notifications: INotification[] = []; + disposables.add(manager.onDidEmitNotification(n => notifications.push(n))); + + manager.deleteSession(sessionUri); + assert.strictEqual(manager.getSessionState(sessionUri), undefined); assert.strictEqual(manager.getSnapshot(sessionUri), undefined); assert.strictEqual(notifications.length, 1); diff --git a/src/vs/workbench/contrib/chat/browser/chatEditing/chatEditingSession.ts b/src/vs/workbench/contrib/chat/browser/chatEditing/chatEditingSession.ts index 37a0f2466aa..f585c7f0996 100644 --- a/src/vs/workbench/contrib/chat/browser/chatEditing/chatEditingSession.ts +++ b/src/vs/workbench/contrib/chat/browser/chatEditing/chatEditingSession.ts @@ -786,7 +786,7 @@ export class ChatEditingSession extends Disposable implements IChatEditingSessio try { const data = await this._fileService.readFile(contentSource); afterSnapshot = data.value.toString(); - } catch { + } catch (_e) { afterSnapshot = ''; } } else { From b0caf28f7e445f4c5c8778d9b228c2e70c80101c Mon Sep 17 00:00:00 2001 From: Connor Peet Date: Thu, 26 Mar 2026 18:20:51 -0700 Subject: [PATCH 2/8] comments --- .../platform/agentHost/node/agentService.ts | 1 - .../agentHost/node/agentSideEffects.ts | 1 - .../agentHost/node/copilot/fileEditTracker.ts | 16 ++++--- .../node/copilot/mapSessionEvents.ts | 2 +- .../agentHost/stateToProgressAdapter.ts | 46 +++++++++++++++++-- 5 files changed, 52 insertions(+), 14 deletions(-) diff --git a/src/vs/platform/agentHost/node/agentService.ts b/src/vs/platform/agentHost/node/agentService.ts index 50a79dfa965..b90fa91c0c0 100644 --- a/src/vs/platform/agentHost/node/agentService.ts +++ b/src/vs/platform/agentHost/node/agentService.ts @@ -172,7 +172,6 @@ export class AgentService extends Disposable implements IAgentService { this._sessionToProvider.delete(session.toString()); } this._stateManager.deleteSession(session.toString()); - this._sessionDataService.deleteSessionData(session); } // ---- Protocol methods --------------------------------------------------- diff --git a/src/vs/platform/agentHost/node/agentSideEffects.ts b/src/vs/platform/agentHost/node/agentSideEffects.ts index 64132425be3..fa90488d644 100644 --- a/src/vs/platform/agentHost/node/agentSideEffects.ts +++ b/src/vs/platform/agentHost/node/agentSideEffects.ts @@ -341,7 +341,6 @@ export class AgentSideEffects extends Disposable implements IProtocolSideEffectH const agent = this._options.getAgent(session); agent?.disposeSession(URI.parse(session)).catch(() => { }); this._stateManager.deleteSession(session); - this._options.sessionDataService.deleteSessionData(URI.parse(session)); } async handleListSessions(): Promise { diff --git a/src/vs/platform/agentHost/node/copilot/fileEditTracker.ts b/src/vs/platform/agentHost/node/copilot/fileEditTracker.ts index d6aa590242f..7eee79b096e 100644 --- a/src/vs/platform/agentHost/node/copilot/fileEditTracker.ts +++ b/src/vs/platform/agentHost/node/copilot/fileEditTracker.ts @@ -45,12 +45,16 @@ export function parseSessionDbUri(raw: string): ISessionDbUriFields | undefined if (!toolCallId || !filePath || (part !== 'before' && part !== 'after')) { return undefined; } - return { - sessionUri: decodeHex(parsed.authority).toString(), - toolCallId: decodeURIComponent(toolCallId), - filePath: decodeHex(filePath).toString(), - part - }; + try { + return { + sessionUri: decodeHex(parsed.authority).toString(), + toolCallId: decodeURIComponent(toolCallId), + filePath: decodeHex(filePath).toString(), + part + }; + } catch { + return undefined; + } } /** diff --git a/src/vs/platform/agentHost/node/copilot/mapSessionEvents.ts b/src/vs/platform/agentHost/node/copilot/mapSessionEvents.ts index 56c9a09e79a..6f697f02cae 100644 --- a/src/vs/platform/agentHost/node/copilot/mapSessionEvents.ts +++ b/src/vs/platform/agentHost/node/copilot/mapSessionEvents.ts @@ -115,7 +115,7 @@ export async function mapSessionEvents( list.push(r); } } - } catch { + } catch (_e) { // Database may not exist yet for new sessions — that's fine } } diff --git a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/stateToProgressAdapter.ts b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/stateToProgressAdapter.ts index dceeac89e09..5a3e30493f8 100644 --- a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/stateToProgressAdapter.ts +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/stateToProgressAdapter.ts @@ -4,7 +4,6 @@ *--------------------------------------------------------------------------------------------*/ import { IMarkdownString, MarkdownString } from '../../../../../../base/common/htmlContent.js'; -import { generateUuid } from '../../../../../../base/common/uuid.js'; import { URI } from '../../../../../../base/common/uri.js'; import { ToolCallStatus, TurnState, ResponsePartKind, getToolFileEdits, getToolOutputText, type ICompletedToolCall, type IToolCallState, type ITurn } from '../../../../../../platform/agentHost/common/state/sessionState.js'; import { getToolKind, getToolLanguage } from '../../../../../../platform/agentHost/common/state/sessionReducers.js'; @@ -46,12 +45,20 @@ export function turnsToHistory(turns: readonly ITurn[], participantId: string): switch (rp.kind) { case ResponsePartKind.Markdown: if (rp.content) { - parts.push({ kind: 'markdownContent', content: new MarkdownString(rp.content) }); + parts.push({ kind: 'markdownContent', content: new MarkdownString(rp.content, { supportHtml: true }) }); } break; - case ResponsePartKind.ToolCall: - parts.push(completedToolCallToSerialized(rp.toolCall as ICompletedToolCall)); + case ResponsePartKind.ToolCall: { + const tc = rp.toolCall as ICompletedToolCall; + const fileEditParts = completedToolCallToEditParts(tc); + const serialized = completedToolCallToSerialized(tc); + if (fileEditParts.length > 0) { + serialized.presentation = ToolInvocationPresentation.Hidden; + } + parts.push(serialized); + parts.push(...fileEditParts); break; + } case ResponsePartKind.Reasoning: if (rp.content) { parts.push({ kind: 'thinking', value: rp.content }); @@ -122,6 +129,35 @@ function completedToolCallToSerialized(tc: ICompletedToolCall): IChatToolInvocat }; } +/** + * Builds edit-structure progress parts for a completed tool call that + * produced file edits. Returns an empty array if the tool call has no edits. + * These parts replay the undo stops and code-block UI when restoring history. + */ +function completedToolCallToEditParts(tc: ICompletedToolCall): IChatProgress[] { + if (tc.status !== ToolCallStatus.Completed) { + return []; + } + const fileEdits = getToolFileEdits(tc); + if (fileEdits.length === 0) { + return []; + } + const filePath = getFilePathFromToolInput(tc); + if (!filePath) { + return []; + } + const fileUri = URI.file(filePath); + const parts: IChatProgress[] = []; + for (const _edit of fileEdits) { + parts.push({ kind: 'markdownContent', content: new MarkdownString('\n````\n') }); + parts.push({ kind: 'codeblockUri', uri: fileUri, isEdit: true, undoStopId: tc.toolCallId }); + parts.push({ kind: 'textEdit', uri: fileUri, edits: [], done: false, isExternalEdit: true }); + parts.push({ kind: 'textEdit', uri: fileUri, edits: [], done: true, isExternalEdit: true }); + parts.push({ kind: 'markdownContent', content: new MarkdownString('\n````\n') }); + } + return parts; +} + /** * Creates a live {@link ChatToolInvocation} from the protocol's tool-call * state. Used during active turns to represent running tool calls in the UI. @@ -278,7 +314,7 @@ function fileEditsToExternalEdits(tc: IToolCallState): IToolCallFileEdit[] { resource: URI.file(filePath), beforeContentUri: URI.parse(edit.beforeURI), afterContentUri: URI.parse(edit.afterURI), - undoStopId: generateUuid(), + undoStopId: tc.toolCallId, }); } } From 133627a2231191784900b2d4c5bacadf2bca72c1 Mon Sep 17 00:00:00 2001 From: Connor Peet Date: Fri, 27 Mar 2026 08:59:30 -0700 Subject: [PATCH 3/8] fix tests --- .../agentHost/node/sessionDatabase.ts | 18 +++--- .../test/node/sessionDatabase.test.ts | 57 ++++++++++--------- 2 files changed, 41 insertions(+), 34 deletions(-) diff --git a/src/vs/platform/agentHost/node/sessionDatabase.ts b/src/vs/platform/agentHost/node/sessionDatabase.ts index 2696d6a9de1..e8f0a9acecb 100644 --- a/src/vs/platform/agentHost/node/sessionDatabase.ts +++ b/src/vs/platform/agentHost/node/sessionDatabase.ts @@ -155,7 +155,7 @@ async function runMigrations(db: Database, migrations: readonly ISessionDatabase export class SessionDatabase implements ISessionDatabase { private _dbPromise: Promise | undefined; - private _disposed = false; + private _closed = false; private readonly _fileEditSequencer = new SequencerByKey(); constructor( @@ -175,7 +175,7 @@ export class SessionDatabase implements ISessionDatabase { } private _ensureDb(): Promise { - if (this._disposed) { + if (this._closed) { return Promise.reject(new Error('SessionDatabase has been disposed')); } if (!this._dbPromise) { @@ -192,7 +192,7 @@ export class SessionDatabase implements ISessionDatabase { throw err; } // If dispose() was called while we were opening, close immediately. - if (this._disposed) { + if (this._closed) { await dbClose(db); throw new Error('SessionDatabase has been disposed'); } @@ -293,12 +293,16 @@ export class SessionDatabase implements ISessionDatabase { }); } - dispose(): void { - if (!this._disposed) { - this._disposed = true; - this._dbPromise?.then(db => db.close()).catch(() => { }); + async close() { + if (!this._closed) { + this._closed = true; + await this._dbPromise?.then(db => db.close()).catch(() => { }); } } + + dispose(): void { + this.close(); + } } function toUint8Array(value: unknown): Uint8Array { diff --git a/src/vs/platform/agentHost/test/node/sessionDatabase.test.ts b/src/vs/platform/agentHost/test/node/sessionDatabase.test.ts index fe7bfe96b6b..e540ace59f9 100644 --- a/src/vs/platform/agentHost/test/node/sessionDatabase.test.ts +++ b/src/vs/platform/agentHost/test/node/sessionDatabase.test.ts @@ -16,14 +16,17 @@ suite('SessionDatabase', () => { const disposables = new DisposableStore(); let testDir: string; + let db: SessionDatabase | undefined; + let db2: SessionDatabase | undefined; setup(() => { testDir = join(tmpdir(), `vscode-session-db-test-${randomUUID()}`); mkdirSync(testDir, { recursive: true }); }); - teardown(() => { + teardown(async () => { disposables.clear(); + await Promise.all([db?.close(), db2?.close()]); rmSync(testDir, { recursive: true, force: true }); }); ensureNoDisposablesAreLeakedInTestSuite(); @@ -42,7 +45,7 @@ suite('SessionDatabase', () => { { version: 2, sql: 'CREATE TABLE t2 (id INTEGER PRIMARY KEY)' }, ]; - const db = disposables.add(await SessionDatabase.open(dbPath(), migrations)); + db = disposables.add(await SessionDatabase.open(dbPath(), migrations)); const tables = (await db.getAllTables()).sort(); assert.deepStrictEqual(tables, ['t1', 't2']); @@ -54,10 +57,10 @@ suite('SessionDatabase', () => { ]; const db1 = await SessionDatabase.open(dbPath(), migrations); - db1.dispose(); + await db1.close(); // Reopen — should not throw (table already exists, migration skipped) - const db2 = disposables.add(await SessionDatabase.open(dbPath(), migrations)); + db2 = disposables.add(await SessionDatabase.open(dbPath(), migrations)); assert.deepStrictEqual(await db2.getAllTables(), ['t1']); }); @@ -66,13 +69,13 @@ suite('SessionDatabase', () => { { version: 1, sql: 'CREATE TABLE t1 (id INTEGER PRIMARY KEY)' }, ]; const db1 = await SessionDatabase.open(dbPath(), v1); - db1.dispose(); + await db1.close(); const v2: ISessionDatabaseMigration[] = [ ...v1, { version: 2, sql: 'CREATE TABLE t2 (id INTEGER PRIMARY KEY)' }, ]; - const db2 = disposables.add(await SessionDatabase.open(dbPath(), v2)); + db2 = disposables.add(await SessionDatabase.open(dbPath(), v2)); const tables = (await db2.getAllTables()).sort(); assert.deepStrictEqual(tables, ['t1', 't2']); @@ -88,7 +91,7 @@ suite('SessionDatabase', () => { // Reopen with only v1 — t1 should not exist because the whole // transaction was rolled back - const db = disposables.add(await SessionDatabase.open(dbPath(), [ + db = disposables.add(await SessionDatabase.open(dbPath(), [ { version: 1, sql: 'CREATE TABLE t1 (id INTEGER PRIMARY KEY)' }, ])); assert.deepStrictEqual(await db.getAllTables(), ['t1']); @@ -100,7 +103,7 @@ suite('SessionDatabase', () => { suite('file edits', () => { test('store and retrieve a file edit', async () => { - const db = disposables.add(await SessionDatabase.open(dbPath())); + db = disposables.add(await SessionDatabase.open(dbPath())); await db.createTurn('turn-1'); await db.storeFileEdit({ @@ -124,7 +127,7 @@ suite('SessionDatabase', () => { }); test('retrieve multiple edits for a single tool call', async () => { - const db = disposables.add(await SessionDatabase.open(dbPath())); + db = disposables.add(await SessionDatabase.open(dbPath())); await db.createTurn('turn-1'); await db.storeFileEdit({ @@ -153,7 +156,7 @@ suite('SessionDatabase', () => { }); test('retrieve edits across multiple tool calls', async () => { - const db = disposables.add(await SessionDatabase.open(dbPath())); + db = disposables.add(await SessionDatabase.open(dbPath())); await db.createTurn('turn-1'); await db.storeFileEdit({ @@ -185,19 +188,19 @@ suite('SessionDatabase', () => { }); test('returns empty array for unknown tool call IDs', async () => { - const db = disposables.add(await SessionDatabase.open(dbPath())); + db = disposables.add(await SessionDatabase.open(dbPath())); const edits = await db.getFileEdits(['nonexistent']); assert.deepStrictEqual(edits, []); }); test('returns empty array when given empty array', async () => { - const db = disposables.add(await SessionDatabase.open(dbPath())); + db = disposables.add(await SessionDatabase.open(dbPath())); const edits = await db.getFileEdits([]); assert.deepStrictEqual(edits, []); }); test('replace on conflict (same toolCallId + filePath)', async () => { - const db = disposables.add(await SessionDatabase.open(dbPath())); + db = disposables.add(await SessionDatabase.open(dbPath())); await db.createTurn('turn-1'); await db.storeFileEdit({ @@ -229,7 +232,7 @@ suite('SessionDatabase', () => { }); test('readFileEditContent returns content on demand', async () => { - const db = disposables.add(await SessionDatabase.open(dbPath())); + db = disposables.add(await SessionDatabase.open(dbPath())); await db.createTurn('turn-1'); await db.storeFileEdit({ @@ -249,13 +252,13 @@ suite('SessionDatabase', () => { }); test('readFileEditContent returns undefined for missing edit', async () => { - const db = disposables.add(await SessionDatabase.open(dbPath())); + db = disposables.add(await SessionDatabase.open(dbPath())); const content = await db.readFileEditContent('tc-missing', '/no/such/file'); assert.strictEqual(content, undefined); }); test('persists binary content correctly', async () => { - const db = disposables.add(await SessionDatabase.open(dbPath())); + db = disposables.add(await SessionDatabase.open(dbPath())); const binary = new Uint8Array([0, 1, 2, 255, 128, 64]); await db.createTurn('turn-1'); @@ -275,7 +278,7 @@ suite('SessionDatabase', () => { }); test('auto-creates turn if it does not exist', async () => { - const db = disposables.add(await SessionDatabase.open(dbPath())); + db = disposables.add(await SessionDatabase.open(dbPath())); // storeFileEdit should succeed even without a prior createTurn call await db.storeFileEdit({ @@ -299,13 +302,13 @@ suite('SessionDatabase', () => { suite('turns', () => { test('createTurn is idempotent', async () => { - const db = disposables.add(await SessionDatabase.open(dbPath())); + db = disposables.add(await SessionDatabase.open(dbPath())); await db.createTurn('turn-1'); await db.createTurn('turn-1'); // should not throw }); test('deleteTurn cascades to file edits', async () => { - const db = disposables.add(await SessionDatabase.open(dbPath())); + db = disposables.add(await SessionDatabase.open(dbPath())); await db.createTurn('turn-1'); await db.storeFileEdit({ @@ -327,7 +330,7 @@ suite('SessionDatabase', () => { }); test('deleteTurn only removes its own edits', async () => { - const db = disposables.add(await SessionDatabase.open(dbPath())); + db = disposables.add(await SessionDatabase.open(dbPath())); await db.createTurn('turn-1'); await db.createTurn('turn-2'); @@ -357,7 +360,7 @@ suite('SessionDatabase', () => { }); test('deleteTurn is a no-op for unknown turn', async () => { - const db = disposables.add(await SessionDatabase.open(dbPath())); + db = disposables.add(await SessionDatabase.open(dbPath())); await db.deleteTurn('nonexistent'); // should not throw }); }); @@ -367,7 +370,7 @@ suite('SessionDatabase', () => { suite('dispose', () => { test('methods throw after dispose', async () => { - const db = await SessionDatabase.open(dbPath()); + db = await SessionDatabase.open(dbPath()); db.dispose(); await assert.rejects( @@ -377,7 +380,7 @@ suite('SessionDatabase', () => { }); test('double dispose is safe', async () => { - const db = await SessionDatabase.open(dbPath()); + db = await SessionDatabase.open(dbPath()); db.dispose(); db.dispose(); // should not throw }); @@ -389,13 +392,13 @@ suite('SessionDatabase', () => { test('constructor does not open the database', () => { // Should not throw even if path does not exist yet - const db = new SessionDatabase(join(testDir, 'lazy', 'session.db')); + db = new SessionDatabase(join(testDir, 'lazy', 'session.db')); disposables.add(db); // No error — the database is not opened until first use }); test('first async call opens and migrates the database', async () => { - const db = disposables.add(new SessionDatabase(dbPath())); + db = disposables.add(new SessionDatabase(dbPath())); // Database file may not exist yet — first call triggers open await db.createTurn('turn-1'); const edits = await db.getFileEdits(['nonexistent']); @@ -403,7 +406,7 @@ suite('SessionDatabase', () => { }); test('multiple concurrent calls share the same open promise', async () => { - const db = disposables.add(new SessionDatabase(dbPath())); + db = disposables.add(new SessionDatabase(dbPath())); // Fire multiple calls concurrently — all should succeed await Promise.all([ db.createTurn('turn-1'), @@ -413,7 +416,7 @@ suite('SessionDatabase', () => { }); test('dispose during open rejects subsequent calls', async () => { - const db = new SessionDatabase(dbPath()); + db = new SessionDatabase(dbPath()); db.dispose(); await assert.rejects(() => db.createTurn('turn-1'), /disposed/); }); From cdd9c01e978f40d25cb964039724a1b4e17f5f4b Mon Sep 17 00:00:00 2001 From: Connor Peet Date: Fri, 27 Mar 2026 09:26:30 -0700 Subject: [PATCH 4/8] fix --- .../platform/agentHost/common/sessionDataService.ts | 6 ++++++ .../agentHost/test/node/sessionDataService.test.ts | 6 ++++++ .../agentHost/test/node/sessionDatabase.test.ts | 12 ++++++------ 3 files changed, 18 insertions(+), 6 deletions(-) diff --git a/src/vs/platform/agentHost/common/sessionDataService.ts b/src/vs/platform/agentHost/common/sessionDataService.ts index 0f6f2c45763..61bc8ba757b 100644 --- a/src/vs/platform/agentHost/common/sessionDataService.ts +++ b/src/vs/platform/agentHost/common/sessionDataService.ts @@ -81,6 +81,12 @@ export interface ISessionDatabase extends IDisposable { * Returns `undefined` if no edit exists for the given key. */ readFileEditContent(toolCallId: string, filePath: string): Promise; + + /** + * Close the database connection. After calling this method, the object is + * considered disposed and all other methods will reject with an error. + */ + close(): Promise; } /** diff --git a/src/vs/platform/agentHost/test/node/sessionDataService.test.ts b/src/vs/platform/agentHost/test/node/sessionDataService.test.ts index 9675c4dac0b..4faf0573777 100644 --- a/src/vs/platform/agentHost/test/node/sessionDataService.test.ts +++ b/src/vs/platform/agentHost/test/node/sessionDataService.test.ts @@ -115,6 +115,7 @@ suite('SessionDataService — openDatabase ref-counting', () => { await ref.object.createTurn('turn-1'); const edits = await ref.object.getFileEdits([]); assert.deepStrictEqual(edits, []); + await ref.object.close(); }); test('multiple references share the same database', async () => { @@ -126,6 +127,7 @@ suite('SessionDataService — openDatabase ref-counting', () => { ref1.dispose(); ref2.dispose(); + await ref1.object.close(); }); test('database remains usable until last reference is disposed', async () => { @@ -139,6 +141,8 @@ suite('SessionDataService — openDatabase ref-counting', () => { await ref2.object.createTurn('turn-1'); ref2.dispose(); + + await ref1.object.close(); }); test('new reference after all disposed gets a fresh database', async () => { @@ -152,5 +156,7 @@ suite('SessionDataService — openDatabase ref-counting', () => { // New reference — may or may not be the same object, but must be functional await ref2.object.createTurn('turn-1'); assert.notStrictEqual(ref2.object, db1); + + await ref2.object.close(); }); }); diff --git a/src/vs/platform/agentHost/test/node/sessionDatabase.test.ts b/src/vs/platform/agentHost/test/node/sessionDatabase.test.ts index e540ace59f9..8c4734b31c4 100644 --- a/src/vs/platform/agentHost/test/node/sessionDatabase.test.ts +++ b/src/vs/platform/agentHost/test/node/sessionDatabase.test.ts @@ -371,18 +371,18 @@ suite('SessionDatabase', () => { test('methods throw after dispose', async () => { db = await SessionDatabase.open(dbPath()); - db.dispose(); + db.close(); await assert.rejects( - () => db.createTurn('turn-1'), + () => db!.createTurn('turn-1'), /disposed/, ); }); test('double dispose is safe', async () => { db = await SessionDatabase.open(dbPath()); - db.dispose(); - db.dispose(); // should not throw + await db.close(); + await db.close(); // should not throw }); }); @@ -417,8 +417,8 @@ suite('SessionDatabase', () => { test('dispose during open rejects subsequent calls', async () => { db = new SessionDatabase(dbPath()); - db.dispose(); - await assert.rejects(() => db.createTurn('turn-1'), /disposed/); + await db.close(); + await assert.rejects(() => db!.createTurn('turn-1'), /disposed/); }); }); }); From 676164bf1567d4b28db782ac8aaf16ad4e533f36 Mon Sep 17 00:00:00 2001 From: Connor Peet Date: Fri, 27 Mar 2026 09:36:05 -0700 Subject: [PATCH 5/8] fix --- .../agentHost/test/node/mapSessionEvents.test.ts | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/src/vs/platform/agentHost/test/node/mapSessionEvents.test.ts b/src/vs/platform/agentHost/test/node/mapSessionEvents.test.ts index d61905c5458..9ebeb2fec2c 100644 --- a/src/vs/platform/agentHost/test/node/mapSessionEvents.test.ts +++ b/src/vs/platform/agentHost/test/node/mapSessionEvents.test.ts @@ -20,6 +20,7 @@ suite('mapSessionEvents', () => { const disposables = new DisposableStore(); let testDir: string; + let db: SessionDatabase | undefined; const session = AgentSession.uri('copilot', 'test-session'); setup(() => { @@ -27,8 +28,9 @@ suite('mapSessionEvents', () => { mkdirSync(testDir, { recursive: true }); }); - teardown(() => { + teardown(async () => { disposables.clear(); + await db?.close(); rmSync(testDir, { recursive: true, force: true }); }); ensureNoDisposablesAreLeakedInTestSuite(); @@ -109,7 +111,7 @@ suite('mapSessionEvents', () => { suite('file edit restoration', () => { test('restores file edits from database for edit tools', async () => { - const db = disposables.add(await SessionDatabase.open(dbPath())); + db = disposables.add(await SessionDatabase.open(dbPath())); await db.createTurn('turn-1'); await db.storeFileEdit({ turnId: 'turn-1', @@ -154,7 +156,7 @@ suite('mapSessionEvents', () => { }); test('handles multiple file edits for one tool call', async () => { - const db = disposables.add(await SessionDatabase.open(dbPath())); + db = disposables.add(await SessionDatabase.open(dbPath())); await db.createTurn('turn-1'); await db.storeFileEdit({ turnId: 'turn-1', @@ -215,7 +217,7 @@ suite('mapSessionEvents', () => { }); test('non-edit tools do not get file edits even if db has data', async () => { - const db = disposables.add(await SessionDatabase.open(dbPath())); + db = disposables.add(await SessionDatabase.open(dbPath())); const events: ISessionEvent[] = [ { From 3746fd5a92c48562fa401065d681212c40841232 Mon Sep 17 00:00:00 2001 From: Connor Peet Date: Fri, 27 Mar 2026 09:45:04 -0700 Subject: [PATCH 6/8] fix --- src/vs/platform/agentHost/test/node/fileEditTracker.test.ts | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/vs/platform/agentHost/test/node/fileEditTracker.test.ts b/src/vs/platform/agentHost/test/node/fileEditTracker.test.ts index 54955b60bf6..bb62eadc929 100644 --- a/src/vs/platform/agentHost/test/node/fileEditTracker.test.ts +++ b/src/vs/platform/agentHost/test/node/fileEditTracker.test.ts @@ -41,8 +41,9 @@ suite('FileEditTracker', () => { tracker = new FileEditTracker('copilot:/test-session', db, fileService, new NullLogService()); }); - teardown(() => { + teardown(async () => { disposables.clear(); + await db.close(); rmSync(testDir, { recursive: true, force: true }); }); ensureNoDisposablesAreLeakedInTestSuite(); From 39473a200312636c3286b12d1ace5f9856fb809e Mon Sep 17 00:00:00 2001 From: Connor Peet Date: Fri, 27 Mar 2026 10:01:42 -0700 Subject: [PATCH 7/8] fix race in teardown --- src/vs/platform/agentHost/node/sessionDatabase.ts | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/src/vs/platform/agentHost/node/sessionDatabase.ts b/src/vs/platform/agentHost/node/sessionDatabase.ts index e8f0a9acecb..296e1b3dfbb 100644 --- a/src/vs/platform/agentHost/node/sessionDatabase.ts +++ b/src/vs/platform/agentHost/node/sessionDatabase.ts @@ -155,7 +155,7 @@ async function runMigrations(db: Database, migrations: readonly ISessionDatabase export class SessionDatabase implements ISessionDatabase { private _dbPromise: Promise | undefined; - private _closed = false; + private _closed: Promise | undefined; private readonly _fileEditSequencer = new SequencerByKey(); constructor( @@ -294,10 +294,7 @@ export class SessionDatabase implements ISessionDatabase { } async close() { - if (!this._closed) { - this._closed = true; - await this._dbPromise?.then(db => db.close()).catch(() => { }); - } + await (this._closed ??= this._dbPromise?.then(db => db.close()).catch(() => { })); } dispose(): void { From e55a46787c42ab6a3bb6e2d8ee4eb1505b23c576 Mon Sep 17 00:00:00 2001 From: Connor Peet Date: Fri, 27 Mar 2026 10:13:50 -0700 Subject: [PATCH 8/8] last fix --- src/vs/platform/agentHost/node/sessionDatabase.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/vs/platform/agentHost/node/sessionDatabase.ts b/src/vs/platform/agentHost/node/sessionDatabase.ts index 296e1b3dfbb..7f960f89d2a 100644 --- a/src/vs/platform/agentHost/node/sessionDatabase.ts +++ b/src/vs/platform/agentHost/node/sessionDatabase.ts @@ -155,7 +155,7 @@ async function runMigrations(db: Database, migrations: readonly ISessionDatabase export class SessionDatabase implements ISessionDatabase { private _dbPromise: Promise | undefined; - private _closed: Promise | undefined; + private _closed: Promise | true | undefined; private readonly _fileEditSequencer = new SequencerByKey(); constructor( @@ -294,7 +294,7 @@ export class SessionDatabase implements ISessionDatabase { } async close() { - await (this._closed ??= this._dbPromise?.then(db => db.close()).catch(() => { })); + await (this._closed ??= this._dbPromise?.then(db => db.close()).catch(() => { }) || true); } dispose(): void {