From 2dbeafcab82b75205fbd2d189102f5bfba1cb304 Mon Sep 17 00:00:00 2001 From: terion Date: Fri, 3 Apr 2026 21:47:35 +0300 Subject: [PATCH 01/46] lsp integration draft --- src/common/constants/experiments.ts | 10 + src/common/types/tools.ts | 4 + src/common/utils/tools/toolAvailability.ts | 2 + src/common/utils/tools/toolDefinitions.ts | 113 ++++ src/common/utils/tools/tools.test.ts | 46 ++ src/common/utils/tools/tools.ts | 12 + src/constants/lsp.ts | 7 + src/node/services/aiService.ts | 15 +- src/node/services/coreServices.ts | 6 + src/node/services/lsp/lspClient.ts | 440 +++++++++++++++ src/node/services/lsp/lspManager.test.ts | 101 ++++ src/node/services/lsp/lspManager.ts | 517 ++++++++++++++++++ src/node/services/lsp/lspPathMapper.test.ts | 38 ++ src/node/services/lsp/lspPathMapper.ts | 117 ++++ src/node/services/lsp/lspServerRegistry.ts | 72 +++ .../services/lsp/lspStdioTransport.test.ts | 84 +++ src/node/services/lsp/lspStdioTransport.ts | 186 +++++++ src/node/services/lsp/types.ts | 134 +++++ src/node/services/serviceContainer.ts | 3 + src/node/services/systemMessage.ts | 5 +- src/node/services/tools/lsp_query.test.ts | 83 +++ src/node/services/tools/lsp_query.ts | 95 ++++ src/node/services/workspaceService.ts | 24 + 23 files changed, 2112 insertions(+), 2 deletions(-) create mode 100644 src/constants/lsp.ts create mode 100644 src/node/services/lsp/lspClient.ts create mode 100644 src/node/services/lsp/lspManager.test.ts create mode 100644 src/node/services/lsp/lspManager.ts create mode 100644 src/node/services/lsp/lspPathMapper.test.ts create mode 100644 src/node/services/lsp/lspPathMapper.ts create mode 100644 src/node/services/lsp/lspServerRegistry.ts create mode 100644 src/node/services/lsp/lspStdioTransport.test.ts create mode 100644 src/node/services/lsp/lspStdioTransport.ts create mode 100644 src/node/services/lsp/types.ts create mode 100644 src/node/services/tools/lsp_query.test.ts create mode 100644 src/node/services/tools/lsp_query.ts diff --git a/src/common/constants/experiments.ts b/src/common/constants/experiments.ts index f5e8ae5078..030fcb1552 100644 --- a/src/common/constants/experiments.ts +++ b/src/common/constants/experiments.ts @@ -10,6 +10,7 @@ export const EXPERIMENT_IDS = { PROGRAMMATIC_TOOL_CALLING_EXCLUSIVE: "programmatic-tool-calling-exclusive", CONFIGURABLE_BIND_URL: "configurable-bind-url", SYSTEM_1: "system-1", + LSP_QUERY: "lsp-query", EXEC_SUBAGENT_HARD_RESTART: "exec-subagent-hard-restart", MUX_GOVERNOR: "mux-governor", MULTI_PROJECT_WORKSPACES: "multi-project-workspaces", @@ -81,6 +82,15 @@ export const EXPERIMENTS: Record = { userOverridable: true, showInSettings: true, }, + [EXPERIMENT_IDS.LSP_QUERY]: { + id: EXPERIMENT_IDS.LSP_QUERY, + name: "LSP Query Tool", + description: + "Enable the built-in lsp_query tool for definitions, references, hover, and symbol lookup", + enabledByDefault: false, + userOverridable: true, + showInSettings: true, + }, [EXPERIMENT_IDS.EXEC_SUBAGENT_HARD_RESTART]: { id: EXPERIMENT_IDS.EXEC_SUBAGENT_HARD_RESTART, name: "Exec sub-agent hard restart", diff --git a/src/common/types/tools.ts b/src/common/types/tools.ts index f715f0df4c..c0e144ec74 100644 --- a/src/common/types/tools.ts +++ b/src/common/types/tools.ts @@ -22,6 +22,7 @@ import type { MuxAgentsReadToolResultSchema, MuxAgentsWriteToolResultSchema, FileReadToolResultSchema, + LspQueryToolResultSchema, AttachFileToolResultSchema, TaskToolResultSchema, TaskAwaitToolResultSchema, @@ -127,6 +128,9 @@ export interface ToolOutputUiOnlyFields { // FileReadToolResult derived from Zod schema (single source of truth) export type FileReadToolResult = z.infer; +export type LspQueryToolArgs = z.infer; +export type LspQueryToolResult = z.infer; + // AttachFileToolResult derived from Zod schema (single source of truth) export type AttachFileToolResult = z.infer; diff --git a/src/common/utils/tools/toolAvailability.ts b/src/common/utils/tools/toolAvailability.ts index 5088280031..2326f94a15 100644 --- a/src/common/utils/tools/toolAvailability.ts +++ b/src/common/utils/tools/toolAvailability.ts @@ -1,6 +1,7 @@ export interface ToolAvailabilityContext { workspaceId: string; parentWorkspaceId?: string | null; + enableLspQuery?: boolean; } /** @@ -10,6 +11,7 @@ export interface ToolAvailabilityContext { export function getToolAvailabilityOptions(context: ToolAvailabilityContext) { return { enableAgentReport: Boolean(context.parentWorkspaceId), + enableLspQuery: context.enableLspQuery === true, // skills_catalog_* tools are always available; agent tool policy controls access. } as const; } diff --git a/src/common/utils/tools/toolDefinitions.ts b/src/common/utils/tools/toolDefinitions.ts index 2ebf37e986..29626ddd2f 100644 --- a/src/common/utils/tools/toolDefinitions.ts +++ b/src/common/utils/tools/toolDefinitions.ts @@ -938,6 +938,51 @@ export const TOOL_DEFINITIONS = { }) ), }, + lsp_query: { + description: + "Query the built-in language server for code intelligence. " + + "Use this for hover, definitions, references, implementations, and symbol lookup. " + + "Provide line and column using 1-based positions for hover/definition/reference/implementation. " + + "For workspace_symbols, provide a representative file path to select the correct language server plus a non-empty query.", + schema: z.preprocess( + normalizeFilePath, + z + .object({ + operation: z.enum([ + "hover", + "definition", + "references", + "implementation", + "document_symbols", + "workspace_symbols", + ]), + path: FILE_TOOL_PATH.describe( + "Path to the file being queried (absolute or relative to the current workspace)" + ), + line: z + .number() + .int() + .positive() + .nullish() + .describe("1-based line number for position-based queries"), + column: z + .number() + .int() + .positive() + .nullish() + .describe("1-based column number for position-based queries"), + query: z + .string() + .nullish() + .describe("Required for workspace_symbols; ignored by other operations"), + includeDeclaration: z + .boolean() + .nullish() + .describe("For references only: whether declarations should be included"), + }) + .strict() + ), + }, attach_file: { description: "Attach a supported file from the filesystem so later model steps receive it as a real attachment instead of a huge base64 JSON blob. " + @@ -1840,6 +1885,69 @@ export const FileReadToolResultSchema = z.union([ }), ]); +const LspQueryResultRangeSchema = z + .object({ + start: z.object({ + line: z.number().int().positive(), + character: z.number().int().positive(), + }), + end: z.object({ + line: z.number().int().positive(), + character: z.number().int().positive(), + }), + }) + .strict(); + +const LspQueryLocationSchema = z + .object({ + path: z.string(), + uri: z.string(), + range: LspQueryResultRangeSchema, + preview: z.string().optional(), + }) + .strict(); + +const LspQuerySymbolSchema = z + .object({ + name: z.string(), + kind: z.number().int(), + detail: z.string().optional(), + containerName: z.string().optional(), + path: z.string(), + range: LspQueryResultRangeSchema, + preview: z.string().optional(), + }) + .strict(); + +export const LspQueryToolResultSchema = z.union([ + z + .object({ + success: z.literal(true), + operation: z.enum([ + "hover", + "definition", + "references", + "implementation", + "document_symbols", + "workspace_symbols", + ]), + serverId: z.string(), + rootUri: z.string(), + hover: z.string().optional(), + locations: z.array(LspQueryLocationSchema).optional(), + symbols: z.array(LspQuerySymbolSchema).optional(), + warning: z.string().optional(), + }) + .strict(), + z + .object({ + success: z.literal(false), + error: z.string(), + warning: z.string().optional(), + }) + .strict(), +]); + const AttachFileToolTextPartSchema = z .object({ type: z.literal("text"), @@ -1963,6 +2071,7 @@ export type BridgeableToolName = | "bash_background_list" | "bash_background_terminate" | "file_read" + | "lsp_query" | "attach_file" | "agent_skill_read" | "agent_skill_read_file" @@ -1990,6 +2099,7 @@ export const RESULT_SCHEMAS: Record = { bash_background_list: BashBackgroundListResultSchema, bash_background_terminate: BashBackgroundTerminateResultSchema, file_read: FileReadToolResultSchema, + lsp_query: LspQueryToolResultSchema, attach_file: AttachFileToolResultSchema, agent_skill_read: AgentSkillReadToolResultSchema, agent_skill_read_file: AgentSkillReadFileToolResultSchema, @@ -2033,6 +2143,7 @@ export function getAvailableTools( options?: { enableAgentReport?: boolean; enableAnalyticsQuery?: boolean; + enableLspQuery?: boolean; /** @deprecated Mux global tools are always included. */ enableMuxGlobalAgentsTools?: boolean; } @@ -2040,6 +2151,7 @@ export function getAvailableTools( const [provider] = modelString.split(":"); const enableAgentReport = options?.enableAgentReport ?? true; const enableAnalyticsQuery = options?.enableAnalyticsQuery ?? true; + const enableLspQuery = options?.enableLspQuery ?? false; // Base tools available for all models // Note: Tool availability is controlled by agent tool policy (allowlist), not mode checks here. @@ -2066,6 +2178,7 @@ export function getAvailableTools( "agent_skill_read", "agent_skill_read_file", "file_edit_replace_string", + ...(enableLspQuery ? ["lsp_query"] : []), // "file_edit_replace_lines", // DISABLED: causes models to break repo state "file_edit_insert", "ask_user_question", diff --git a/src/common/utils/tools/tools.test.ts b/src/common/utils/tools/tools.test.ts index d664793f7b..51bb756440 100644 --- a/src/common/utils/tools/tools.test.ts +++ b/src/common/utils/tools/tools.test.ts @@ -3,6 +3,7 @@ import { describe, expect, mock, test } from "bun:test"; import type { InitStateManager } from "@/node/services/initStateManager"; import type { DesktopSessionManager } from "@/node/services/desktop/DesktopSessionManager"; import { LocalRuntime } from "@/node/runtime/LocalRuntime"; +import { LspManager } from "@/node/services/lsp/lspManager"; import { getToolsForModel } from "./tools"; const DESKTOP_TOOL_NAMES = [ @@ -156,4 +157,49 @@ describe("getToolsForModel", () => { expect(Object.keys(tools).filter((toolName) => toolName.startsWith("desktop_"))).toEqual([]); }); + + test("only includes lsp_query when the experiment is enabled and a manager is available", async () => { + const runtime = new LocalRuntime(process.cwd()); + const initStateManager = createInitStateManager(); + const lspManager = new LspManager({ registry: [] }); + lspManager.query = mock(() => + Promise.resolve({ + operation: "hover" as const, + serverId: "typescript", + rootUri: "file:///tmp/workspace", + hover: "", + }) + ); + + try { + const toolsWithoutLsp = await getToolsForModel( + "noop:model", + { + cwd: process.cwd(), + runtime, + runtimeTempDir: "/tmp", + lspQueryEnabled: false, + }, + "ws-1", + initStateManager + ); + expect(toolsWithoutLsp.lsp_query).toBeUndefined(); + + const toolsWithLsp = await getToolsForModel( + "noop:model", + { + cwd: process.cwd(), + runtime, + runtimeTempDir: "/tmp", + lspManager, + lspQueryEnabled: true, + }, + "ws-1", + initStateManager + ); + expect(toolsWithLsp.lsp_query).toBeDefined(); + } finally { + await lspManager.dispose(); + } + }); }); diff --git a/src/common/utils/tools/tools.ts b/src/common/utils/tools/tools.ts index c454337656..2eb536545a 100644 --- a/src/common/utils/tools/tools.ts +++ b/src/common/utils/tools/tools.ts @@ -15,6 +15,7 @@ import { createTodoWriteTool, createTodoReadTool } from "@/node/services/tools/t import { createNotifyTool } from "@/node/services/tools/notify"; import { createAnalyticsQueryTool } from "@/node/services/tools/analyticsQuery"; import { createDesktopTools } from "@/node/services/tools/desktopTools"; +import { createLspQueryTool } from "@/node/services/tools/lsp_query"; import type { MuxToolScope } from "@/common/types/toolScope"; import { createTaskTool } from "@/node/services/tools/task"; import { createTaskApplyGitPatchTool } from "@/node/services/tools/task_apply_git_patch"; @@ -54,6 +55,7 @@ import type { FileState } from "@/node/services/agentSession"; import type { AgentDefinitionDescriptor } from "@/common/types/agentDefinition"; import type { AgentSkillDescriptor } from "@/common/types/agentSkill"; import type { ProjectRef } from "@/common/types/workspace"; +import type { LspManager } from "@/node/services/lsp/lspManager"; /** * Configuration for tools that need runtime context @@ -120,6 +122,10 @@ export interface ToolConfiguration { }; /** Desktop session manager for desktop automation tools */ desktopSessionManager?: DesktopSessionManager; + /** Shared workspace-scoped LSP manager for built-in query tooling */ + lspManager?: LspManager; + /** Whether the experiment-gated lsp_query tool should be exposed for this request */ + lspQueryEnabled?: boolean; } /** @@ -346,6 +352,11 @@ export async function getToolsForModel( agent_skill_read_file: wrap(createAgentSkillReadFileTool(config)), file_edit_replace_string: wrap(createFileEditReplaceStringTool(config)), file_edit_insert: wrap(createFileEditInsertTool(config)), + ...(config.lspManager && config.lspQueryEnabled + ? { + lsp_query: wrap(createLspQueryTool(config)), + } + : {}), // DISABLED: file_edit_replace_lines - causes models (particularly GPT-5-Codex) // to leave repository in broken state due to issues with concurrent file modifications // and line number miscalculations. Use file_edit_replace_string instead. @@ -492,6 +503,7 @@ export async function getToolsForModel( getAvailableTools(modelString, { enableAgentReport: config.enableAgentReport, enableAnalyticsQuery: Boolean(config.analyticsService), + enableLspQuery: config.lspManager != null && config.lspQueryEnabled === true, // Mux global tools are always created; tool policy (agent frontmatter) // controls which agents can actually use them. enableMuxGlobalAgentsTools: true, diff --git a/src/constants/lsp.ts b/src/constants/lsp.ts new file mode 100644 index 0000000000..0d64cd72e9 --- /dev/null +++ b/src/constants/lsp.ts @@ -0,0 +1,7 @@ +export const LSP_IDLE_TIMEOUT_MS = 10 * 60 * 1000; +export const LSP_IDLE_CHECK_INTERVAL_MS = 60 * 1000; +export const LSP_REQUEST_TIMEOUT_MS = 10 * 1000; +export const LSP_START_TIMEOUT_MS = 5 * 1000; +export const LSP_MAX_LOCATIONS = 25; +export const LSP_MAX_SYMBOLS = 100; +export const LSP_PREVIEW_CONTEXT_LINES = 1; diff --git a/src/node/services/aiService.ts b/src/node/services/aiService.ts index cc1873cfd1..56addc2e94 100644 --- a/src/node/services/aiService.ts +++ b/src/node/services/aiService.ts @@ -99,6 +99,7 @@ import { import { applyToolPolicyAndExperiments, captureMcpToolTelemetry } from "./toolAssembly"; import { getErrorMessage } from "@/common/utils/errors"; import { isProjectTrusted } from "@/node/utils/projectTrust"; +import type { LspManager } from "@/node/services/lsp/lspManager"; const STREAM_STARTUP_DIAGNOSTIC_THRESHOLD_MS = 1_000; @@ -211,6 +212,7 @@ export class AIService extends EventEmitter { private readonly config: Config; private readonly workspaceMcpOverridesService: WorkspaceMcpOverridesService; private mcpServerManager?: MCPServerManager; + private lspManager?: LspManager; private readonly policyService?: PolicyService; private readonly telemetryService?: TelemetryService; private readonly opResolver?: ExternalSecretResolver; @@ -312,6 +314,10 @@ export class AIService extends EventEmitter { this.streamManager.setMCPServerManager(manager); } + setLspManager(manager: LspManager): void { + this.lspManager = manager; + } + setTaskService(taskService: TaskService): void { this.taskService = taskService; } @@ -1106,6 +1112,8 @@ export class AIService extends EventEmitter { }; const desktopSessionManager = this.desktopSessionManager; + const lspQueryEnabled = + this.experimentsService?.isExperimentEnabled(EXPERIMENT_IDS.LSP_QUERY) ?? false; let desktopCapabilityPromise: ReturnType | undefined; const loadDesktopCapability = desktopSessionManager == null @@ -1204,7 +1212,10 @@ export class AIService extends EventEmitter { runtime, workspacePath, modelString, - agentSystemPrompt + agentSystemPrompt, + { + enableLspQuery: lspQueryEnabled, + } ); recordStartupPhaseTiming("readToolInstructionsMs", readToolInstructionsStartedAt); @@ -1273,6 +1284,8 @@ export class AIService extends EventEmitter { taskService: this.taskService, analyticsService: this.analyticsService, desktopSessionManager: this.desktopSessionManager, + lspManager: this.lspManager, + lspQueryEnabled, // PTC experiments for inheritance to subagents experiments, // Dynamic context for tool descriptions (moved from system prompt for better model attention) diff --git a/src/node/services/coreServices.ts b/src/node/services/coreServices.ts index e5b4a42e48..49366a4aab 100644 --- a/src/node/services/coreServices.ts +++ b/src/node/services/coreServices.ts @@ -13,6 +13,7 @@ import { BackgroundProcessManager } from "@/node/services/backgroundProcessManag import { SessionUsageService } from "@/node/services/sessionUsageService"; import { MCPConfigService } from "@/node/services/mcpConfigService"; import { MCPServerManager, type MCPServerManagerOptions } from "@/node/services/mcpServerManager"; +import { LspManager } from "@/node/services/lsp/lspManager"; import { ExtensionMetadataService } from "@/node/services/ExtensionMetadataService"; import { WorkspaceService } from "@/node/services/workspaceService"; import { TaskService } from "@/node/services/taskService"; @@ -49,6 +50,7 @@ export interface CoreServices { aiService: AIService; mcpConfigService: MCPConfigService; mcpServerManager: MCPServerManager; + lspManager: LspManager; extensionMetadata: ExtensionMetadataService; workspaceService: WorkspaceService; taskService: TaskService; @@ -88,6 +90,8 @@ export function createCoreServices(opts: CoreServicesOptions): CoreServices { opts.policyService ); aiService.setMCPServerManager(mcpServerManager); + const lspManager = new LspManager(); + aiService.setLspManager(lspManager); const extensionMetadata = new ExtensionMetadataService(extensionMetadataPath); @@ -106,6 +110,7 @@ export function createCoreServices(opts: CoreServicesOptions): CoreServices { opts.opResolver ); workspaceService.setMCPServerManager(mcpServerManager); + workspaceService.setLspManager(lspManager); const taskService = new TaskService( config, @@ -127,6 +132,7 @@ export function createCoreServices(opts: CoreServicesOptions): CoreServices { aiService, mcpConfigService, mcpServerManager, + lspManager, extensionMetadata, workspaceService, taskService, diff --git a/src/node/services/lsp/lspClient.ts b/src/node/services/lsp/lspClient.ts new file mode 100644 index 0000000000..a4b99848b9 --- /dev/null +++ b/src/node/services/lsp/lspClient.ts @@ -0,0 +1,440 @@ +import { readFileString } from "@/node/utils/runtime/helpers"; +import { shellQuote } from "@/common/utils/shell"; +import { + LSP_REQUEST_TIMEOUT_MS, + LSP_START_TIMEOUT_MS, +} from "@/constants/lsp"; +import { log } from "@/node/services/log"; +import { LspStdioTransport, type LspJsonRpcMessage } from "./lspStdioTransport"; +import type { + CreateLspClientOptions, + LspClientFileHandle, + LspClientInstance, + LspClientQueryRequest, + LspClientQueryResult, + LspDocumentSymbol, + LspHover, + LspLocation, + LspLocationLink, + LspMarkedString, + LspMarkupContent, + LspSymbolInformation, +} from "./types"; + +interface PendingRequest { + resolve: (value: unknown) => void; + reject: (error: Error) => void; + timeoutId: ReturnType; +} + +interface OpenDocumentState { + version: number; + text: string; +} + +interface InitializeResult { + capabilities?: Record; +} + +export class LspClient implements LspClientInstance { + private readonly transport: LspStdioTransport; + private readonly pendingRequests = new Map(); + private readonly openDocuments = new Map(); + private nextRequestId = 1; + private initialized = false; + isClosed = false; + + private constructor(private readonly options: CreateLspClientOptions, transport: LspStdioTransport) { + this.transport = transport; + this.transport.onmessage = (message) => this.handleMessage(message); + this.transport.onclose = () => this.handleClose(); + this.transport.onerror = (error) => this.handleTransportError(error); + } + + static async create(options: CreateLspClientOptions): Promise { + const command = [shellQuote(options.descriptor.command), ...options.descriptor.args.map((arg) => shellQuote(arg))].join(" "); + const execStream = await options.runtime.exec(command, { + cwd: options.rootPath, + // LSP servers are long-lived by design; timeout would kill healthy clients mid-session. + }); + + const transport = new LspStdioTransport(execStream); + const client = new LspClient(options, transport); + await client.start(); + return client; + } + + async ensureFile(file: LspClientFileHandle): Promise { + this.ensureStarted(); + + const text = await readFileString(this.options.runtime, file.readablePath); + const existingState = this.openDocuments.get(file.uri); + + if (!existingState) { + this.openDocuments.set(file.uri, { + version: 1, + text, + }); + await this.notify("textDocument/didOpen", { + textDocument: { + uri: file.uri, + languageId: file.languageId, + version: 1, + text, + }, + }); + return; + } + + if (existingState.text === text) { + return; + } + + const nextVersion = existingState.version + 1; + this.openDocuments.set(file.uri, { + version: nextVersion, + text, + }); + await this.notify("textDocument/didChange", { + textDocument: { + uri: file.uri, + version: nextVersion, + }, + contentChanges: [{ text }], + }); + } + + async query(request: LspClientQueryRequest): Promise { + this.ensureStarted(); + + switch (request.operation) { + case "hover": { + const result = (await this.request( + "textDocument/hover", + this.createTextDocumentPositionParams(request), + LSP_REQUEST_TIMEOUT_MS + )) as LspHover | null; + return { + operation: request.operation, + hover: result ? normalizeHoverContents(result.contents) : "", + }; + } + case "definition": { + const result = (await this.request( + "textDocument/definition", + this.createTextDocumentPositionParams(request), + LSP_REQUEST_TIMEOUT_MS + )) as LspLocation | LspLocationLink | Array | null; + return { + operation: request.operation, + locations: normalizeLocations(result), + }; + } + case "references": { + const result = (await this.request( + "textDocument/references", + { + ...this.createTextDocumentPositionParams(request), + context: { + includeDeclaration: request.includeDeclaration === true, + }, + }, + LSP_REQUEST_TIMEOUT_MS + )) as LspLocation[] | null; + return { + operation: request.operation, + locations: normalizeLocations(result), + }; + } + case "implementation": { + const result = (await this.request( + "textDocument/implementation", + this.createTextDocumentPositionParams(request), + LSP_REQUEST_TIMEOUT_MS + )) as LspLocation | LspLocationLink | Array | null; + return { + operation: request.operation, + locations: normalizeLocations(result), + }; + } + case "document_symbols": { + const result = (await this.request( + "textDocument/documentSymbol", + { + textDocument: { + uri: request.file.uri, + }, + }, + LSP_REQUEST_TIMEOUT_MS + )) as Array | null; + return { + operation: request.operation, + symbols: result ?? [], + }; + } + case "workspace_symbols": { + const result = (await this.request( + "workspace/symbol", + { + query: request.query ?? "", + }, + LSP_REQUEST_TIMEOUT_MS + )) as Array | null; + return { + operation: request.operation, + symbols: result ?? [], + }; + } + } + } + + async close(): Promise { + if (this.isClosed) { + return; + } + + try { + if (this.initialized) { + await this.request("shutdown", undefined, 3000); + await this.notify("exit"); + } + } catch (error) { + log.debug("Failed to shut down LSP client cleanly", { + serverId: this.options.descriptor.id, + error, + }); + } finally { + this.isClosed = true; + await this.transport.close(); + } + } + + private async start(): Promise { + this.transport.start(); + const response = (await this.request( + "initialize", + { + processId: process.pid, + rootUri: this.options.rootUri, + rootPath: this.options.rootPath, + workspaceFolders: [ + { + uri: this.options.rootUri, + name: this.options.descriptor.id, + }, + ], + capabilities: { + workspace: { + workspaceFolders: true, + }, + textDocument: { + hover: { + contentFormat: ["markdown", "plaintext"], + }, + definition: { + linkSupport: true, + }, + implementation: { + linkSupport: true, + }, + references: {}, + documentSymbol: { + hierarchicalDocumentSymbolSupport: true, + }, + }, + }, + }, + LSP_START_TIMEOUT_MS + )) as InitializeResult; + + this.initialized = true; + await this.notify("initialized", {}); + + log.debug("Started LSP client", { + serverId: this.options.descriptor.id, + rootUri: this.options.rootUri, + capabilities: response.capabilities ? Object.keys(response.capabilities) : [], + }); + } + + private ensureStarted(): void { + if (!this.initialized) { + throw new Error(`LSP client for ${this.options.descriptor.id} is not initialized`); + } + if (this.isClosed || this.transport.isClosed()) { + throw new Error( + `LSP client for ${this.options.descriptor.id} is closed${this.buildStderrSuffix()}` + ); + } + } + + private createTextDocumentPositionParams(request: LspClientQueryRequest) { + if (request.line == null || request.character == null) { + throw new Error(`${request.operation} requires a line and column`); + } + + return { + textDocument: { + uri: request.file.uri, + }, + position: { + line: request.line, + character: request.character, + }, + }; + } + + private async request(method: string, params: unknown, timeoutMs: number): Promise { + if (this.isClosed || this.transport.isClosed()) { + throw new Error( + `Cannot send ${method} to closed LSP client (${this.options.descriptor.id})${this.buildStderrSuffix()}` + ); + } + + const requestId = this.nextRequestId++; + + return await new Promise((resolve, reject) => { + const timeoutId = setTimeout(() => { + this.pendingRequests.delete(requestId); + reject( + new Error( + `LSP request timed out: ${this.options.descriptor.id} ${method}${this.buildStderrSuffix()}` + ) + ); + }, timeoutMs); + + this.pendingRequests.set(requestId, { + resolve, + reject, + timeoutId, + }); + + void this.transport.send({ + jsonrpc: "2.0", + id: requestId, + method, + ...(params !== undefined ? { params } : {}), + }).catch((error) => { + const pending = this.pendingRequests.get(requestId); + if (!pending) { + return; + } + + clearTimeout(pending.timeoutId); + this.pendingRequests.delete(requestId); + reject(error instanceof Error ? error : new Error(String(error))); + }); + }); + } + + private async notify(method: string, params?: unknown): Promise { + if (this.isClosed || this.transport.isClosed()) { + return; + } + + await this.transport.send({ + jsonrpc: "2.0", + method, + ...(params !== undefined ? { params } : {}), + }); + } + + private handleMessage(message: LspJsonRpcMessage): void { + if (typeof message.id !== "number") { + return; + } + + const pending = this.pendingRequests.get(message.id); + if (!pending) { + return; + } + + clearTimeout(pending.timeoutId); + this.pendingRequests.delete(message.id); + + if (message.error) { + pending.reject( + new Error( + `LSP ${this.options.descriptor.id} request failed: ${message.error.message}${this.buildStderrSuffix()}` + ) + ); + return; + } + + pending.resolve(message.result); + } + + private handleClose(): void { + this.isClosed = true; + const error = new Error( + `LSP client for ${this.options.descriptor.id} closed unexpectedly${this.buildStderrSuffix()}` + ); + this.rejectAllPending(error); + } + + private handleTransportError(error: Error): void { + this.isClosed = true; + this.rejectAllPending( + new Error( + `LSP transport error for ${this.options.descriptor.id}: ${error.message}${this.buildStderrSuffix()}` + ) + ); + } + + private rejectAllPending(error: Error): void { + for (const [requestId, pending] of this.pendingRequests) { + clearTimeout(pending.timeoutId); + pending.reject(error); + this.pendingRequests.delete(requestId); + } + } + + private buildStderrSuffix(): string { + const stderrTail = this.transport.getStderrTail(); + return stderrTail.length > 0 ? `; stderr: ${stderrTail}` : ""; + } +} + +function normalizeLocations( + value: LspLocation | LspLocationLink | Array | null | undefined +): LspLocation[] { + if (!value) { + return []; + } + + const values = Array.isArray(value) ? value : [value]; + return values.map((entry) => { + if ("targetUri" in entry) { + return { + uri: entry.targetUri, + range: entry.targetRange, + }; + } + + return entry; + }); +} + +function normalizeHoverContents( + contents: LspHover["contents"] +): string { + if (typeof contents === "string") { + return contents; + } + + if (Array.isArray(contents)) { + return contents.map((entry) => normalizeHoverContents(entry)).filter((entry) => entry.length > 0).join("\n\n"); + } + + if (isMarkupContent(contents)) { + return contents.value; + } + + return contents.value; +} + +function isMarkupContent( + value: LspMarkupContent | LspMarkedString +): value is LspMarkupContent { + return "kind" in value; +} diff --git a/src/node/services/lsp/lspManager.test.ts b/src/node/services/lsp/lspManager.test.ts new file mode 100644 index 0000000000..2a4f13a56f --- /dev/null +++ b/src/node/services/lsp/lspManager.test.ts @@ -0,0 +1,101 @@ +import * as fs from "node:fs/promises"; +import * as os from "node:os"; +import * as path from "node:path"; +import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test"; +import { LocalRuntime } from "@/node/runtime/LocalRuntime"; +import type { CreateLspClientOptions, LspClientInstance, LspServerDescriptor } from "./types"; +import { LspManager } from "./lspManager"; + +describe("LspManager", () => { + let workspacePath: string; + + beforeEach(async () => { + workspacePath = await fs.mkdtemp(path.join(os.tmpdir(), "mux-lsp-manager-")); + await fs.mkdir(path.join(workspacePath, ".git")); + await fs.mkdir(path.join(workspacePath, "src"), { recursive: true }); + await fs.writeFile(path.join(workspacePath, "src", "example.ts"), "export const value = 1;\n"); + }); + + afterEach(async () => { + await fs.rm(workspacePath, { recursive: true, force: true }); + }); + + test("reuses clients for the same workspace root and forwards normalized positions", async () => { + const ensureFile = mock(() => Promise.resolve(undefined)); + let lastQueryRequest: Parameters[0] | undefined; + const query = mock((request: Parameters[0]) => { + lastQueryRequest = request; + return Promise.resolve({ + operation: "hover" as const, + hover: "const value: 1", + }); + }); + const close = mock(() => Promise.resolve(undefined)); + const client: LspClientInstance = { + isClosed: false, + ensureFile, + query, + close, + }; + let clientFactoryOptions: CreateLspClientOptions | undefined; + const clientFactory = mock((options: CreateLspClientOptions): Promise => { + clientFactoryOptions = options; + return Promise.resolve(client); + }); + const registry: readonly LspServerDescriptor[] = [ + { + id: "typescript", + extensions: [".ts"], + command: "fake-lsp", + args: ["--stdio"], + rootMarkers: ["package.json", ".git"], + languageIdForPath: () => "typescript", + }, + ]; + + const manager = new LspManager({ + registry, + clientFactory, + }); + const runtime = new LocalRuntime(workspacePath); + + const firstResult = await manager.query({ + workspaceId: "ws-1", + runtime, + workspacePath, + filePath: "src/example.ts", + operation: "hover", + line: 2, + column: 3, + }); + const secondResult = await manager.query({ + workspaceId: "ws-1", + runtime, + workspacePath, + filePath: "src/example.ts", + operation: "hover", + line: 2, + column: 3, + }); + + expect(firstResult.hover).toBe("const value: 1"); + expect(secondResult.hover).toBe("const value: 1"); + expect(clientFactory).toHaveBeenCalledTimes(1); + expect(ensureFile).toHaveBeenCalledTimes(2); + expect(clientFactoryOptions).toBeDefined(); + if (!clientFactoryOptions) { + throw new Error("Expected the LSP client factory to receive a call"); + } + expect(clientFactoryOptions.rootPath).toBe(workspacePath); + + expect(lastQueryRequest).toBeDefined(); + if (!lastQueryRequest) { + throw new Error("Expected the LSP client to receive a query"); + } + expect(lastQueryRequest.line).toBe(1); + expect(lastQueryRequest.character).toBe(2); + + await manager.dispose(); + expect(close).toHaveBeenCalledTimes(1); + }); +}); diff --git a/src/node/services/lsp/lspManager.ts b/src/node/services/lsp/lspManager.ts new file mode 100644 index 0000000000..8224e920be --- /dev/null +++ b/src/node/services/lsp/lspManager.ts @@ -0,0 +1,517 @@ +import * as path from "node:path"; +import { readFileString } from "@/node/utils/runtime/helpers"; +import { + LSP_IDLE_CHECK_INTERVAL_MS, + LSP_IDLE_TIMEOUT_MS, + LSP_MAX_LOCATIONS, + LSP_MAX_SYMBOLS, + LSP_PREVIEW_CONTEXT_LINES, +} from "@/constants/lsp"; +import type { Runtime } from "@/node/runtime/Runtime"; +import { log } from "@/node/services/log"; +import { LspClient } from "./lspClient"; +import { LspPathMapper } from "./lspPathMapper"; +import { BUILTIN_LSP_SERVERS, findLspServerForFile } from "./lspServerRegistry"; +import type { + CreateLspClientOptions, + LspClientFileHandle, + LspClientInstance, + LspClientQueryResult, + LspDocumentSymbol, + LspLocation, + LspManagerQueryResult, + LspQueryOperation, + LspRange, + LspServerDescriptor, + LspSymbolInformation, +} from "./types"; + +interface WorkspaceClients { + clients: Map; + lastActivity: number; +} + +type LspClientFactory = (options: CreateLspClientOptions) => Promise; + +export interface LspManagerOptions { + registry?: readonly LspServerDescriptor[]; + clientFactory?: LspClientFactory; + idleTimeoutMs?: number; + idleCheckIntervalMs?: number; +} + +export interface LspManagerQueryOptions { + workspaceId: string; + runtime: Runtime; + workspacePath: string; + filePath: string; + operation: LspQueryOperation; + line?: number; + column?: number; + query?: string; + includeDeclaration?: boolean; +} + +export class LspManager { + private readonly workspaceClients = new Map(); + private readonly workspaceLeases = new Map(); + private readonly registry: readonly LspServerDescriptor[]; + private readonly clientFactory: LspClientFactory; + private readonly idleTimeoutMs: number; + private readonly idleCheckIntervalMs: number; + private readonly idleInterval: ReturnType; + + constructor(options?: LspManagerOptions) { + this.registry = options?.registry ?? BUILTIN_LSP_SERVERS; + this.clientFactory = options?.clientFactory ?? ((clientOptions) => LspClient.create(clientOptions)); + this.idleTimeoutMs = options?.idleTimeoutMs ?? LSP_IDLE_TIMEOUT_MS; + this.idleCheckIntervalMs = options?.idleCheckIntervalMs ?? LSP_IDLE_CHECK_INTERVAL_MS; + this.idleInterval = setInterval(() => { + void this.cleanupIdleWorkspaces(); + }, this.idleCheckIntervalMs); + this.idleInterval.unref?.(); + } + + acquireLease(workspaceId: string): void { + const currentLeaseCount = this.workspaceLeases.get(workspaceId) ?? 0; + this.workspaceLeases.set(workspaceId, currentLeaseCount + 1); + this.markActivity(workspaceId); + } + + releaseLease(workspaceId: string): void { + const currentLeaseCount = this.workspaceLeases.get(workspaceId) ?? 0; + if (currentLeaseCount <= 1) { + this.workspaceLeases.delete(workspaceId); + return; + } + + this.workspaceLeases.set(workspaceId, currentLeaseCount - 1); + } + + async query(options: LspManagerQueryOptions): Promise { + this.acquireLease(options.workspaceId); + + try { + const pathMapper = new LspPathMapper({ + runtime: options.runtime, + workspacePath: options.workspacePath, + }); + const runtimeFilePath = pathMapper.toRuntimePath(options.filePath); + if (!pathMapper.isWithinWorkspace(runtimeFilePath)) { + throw new Error(`LSP paths must stay inside the workspace (got ${options.filePath})`); + } + + const descriptor = findLspServerForFile(runtimeFilePath, this.registry); + if (!descriptor) { + const extension = path.extname(runtimeFilePath) || "(no extension)"; + throw new Error(`No built-in LSP server is configured for ${extension} files`); + } + + const rootPath = await this.resolveRootPath( + options.runtime, + runtimeFilePath, + pathMapper.getWorkspaceRuntimePath(), + descriptor.rootMarkers + ); + const rootUri = pathMapper.toUri(rootPath); + const fileHandle: LspClientFileHandle = { + runtimePath: runtimeFilePath, + readablePath: pathMapper.toReadablePath(runtimeFilePath), + uri: pathMapper.toUri(runtimeFilePath), + languageId: descriptor.languageIdForPath(runtimeFilePath), + }; + + const client = await this.getOrCreateClient(options.workspaceId, descriptor, options.runtime, rootPath, rootUri); + await client.ensureFile(fileHandle); + + const rawResult = await client.query({ + operation: options.operation, + file: fileHandle, + line: options.line != null ? Math.max(0, options.line - 1) : undefined, + character: options.column != null ? Math.max(0, options.column - 1) : undefined, + query: options.query, + includeDeclaration: options.includeDeclaration, + }); + + this.markActivity(options.workspaceId); + return await this.normalizeQueryResult(pathMapper, options.runtime, fileHandle, descriptor.id, rootUri, rawResult); + } finally { + this.releaseLease(options.workspaceId); + } + } + + async disposeWorkspace(workspaceId: string): Promise { + const entry = this.workspaceClients.get(workspaceId); + if (!entry) { + return; + } + + this.workspaceClients.delete(workspaceId); + await Promise.all( + [...entry.clients.values()].map(async (client) => { + try { + await client.close(); + } catch (error) { + log.debug("Failed to close LSP client during workspace disposal", { + workspaceId, + error, + }); + } + }) + ); + } + + async dispose(): Promise { + clearInterval(this.idleInterval); + const workspaceIds = [...this.workspaceClients.keys()]; + await Promise.all(workspaceIds.map(async (workspaceId) => this.disposeWorkspace(workspaceId))); + } + + private async getOrCreateClient( + workspaceId: string, + descriptor: LspServerDescriptor, + runtime: Runtime, + rootPath: string, + rootUri: string + ): Promise { + const workspaceEntry = this.workspaceClients.get(workspaceId) ?? { + clients: new Map(), + lastActivity: Date.now(), + }; + this.workspaceClients.set(workspaceId, workspaceEntry); + + const clientKey = `${descriptor.id}:${rootUri}`; + const existingClient = workspaceEntry.clients.get(clientKey); + if (existingClient && !existingClient.isClosed) { + workspaceEntry.lastActivity = Date.now(); + return existingClient; + } + + if (existingClient?.isClosed) { + workspaceEntry.clients.delete(clientKey); + } + + const client = await this.clientFactory({ + descriptor, + runtime, + rootPath, + rootUri, + }); + workspaceEntry.clients.set(clientKey, client); + workspaceEntry.lastActivity = Date.now(); + return client; + } + + private markActivity(workspaceId: string): void { + const entry = this.workspaceClients.get(workspaceId); + if (!entry) { + return; + } + + entry.lastActivity = Date.now(); + } + + private async cleanupIdleWorkspaces(): Promise { + const now = Date.now(); + for (const [workspaceId, entry] of this.workspaceClients) { + if ((this.workspaceLeases.get(workspaceId) ?? 0) > 0) { + continue; + } + + if (now - entry.lastActivity < this.idleTimeoutMs) { + continue; + } + + log.info("Stopping idle LSP clients", { + workspaceId, + idleMinutes: Math.round((now - entry.lastActivity) / 60_000), + }); + await this.disposeWorkspace(workspaceId); + } + } + + private async resolveRootPath( + runtime: Runtime, + filePath: string, + workspaceRuntimePath: string, + rootMarkers: readonly string[] + ): Promise { + const pathModule = selectPathModule(filePath); + let currentPath = pathModule.dirname(filePath); + + while (true) { + for (const marker of rootMarkers) { + const markerPath = runtime.normalizePath(marker, currentPath); + if (await pathExists(runtime, markerPath)) { + return currentPath; + } + } + + if (currentPath === workspaceRuntimePath) { + return workspaceRuntimePath; + } + + const parentPath = pathModule.dirname(currentPath); + if (parentPath === currentPath) { + return workspaceRuntimePath; + } + currentPath = parentPath; + } + } + + private async normalizeQueryResult( + pathMapper: LspPathMapper, + runtime: Runtime, + fileHandle: LspClientFileHandle, + serverId: string, + rootUri: string, + rawResult: LspClientQueryResult + ): Promise { + switch (rawResult.operation) { + case "hover": + return { + operation: rawResult.operation, + serverId, + rootUri, + hover: rawResult.hover ?? "", + }; + case "definition": + case "references": + case "implementation": { + const warning = + rawResult.locations && rawResult.locations.length > LSP_MAX_LOCATIONS + ? `Results truncated to the first ${LSP_MAX_LOCATIONS} locations` + : undefined; + const locations = await Promise.all( + (rawResult.locations ?? []).slice(0, LSP_MAX_LOCATIONS).map(async (location) => + this.buildLocationResult(pathMapper, runtime, location) + ) + ); + return { + operation: rawResult.operation, + serverId, + rootUri, + locations, + ...(warning ? { warning } : {}), + }; + } + case "document_symbols": { + const flattenedSymbols = flattenDocumentSymbols(rawResult.symbols ?? [], fileHandle.uri); + const warning = + flattenedSymbols.length > LSP_MAX_SYMBOLS + ? `Results truncated to the first ${LSP_MAX_SYMBOLS} symbols` + : undefined; + const symbols = await Promise.all( + flattenedSymbols.slice(0, LSP_MAX_SYMBOLS).map(async (symbol) => + this.buildSymbolResult(pathMapper, runtime, symbol.uri, symbol.range, symbol.name, symbol.kind, { + detail: symbol.detail, + containerName: symbol.containerName, + }) + ) + ); + return { + operation: rawResult.operation, + serverId, + rootUri, + symbols, + ...(warning ? { warning } : {}), + }; + } + case "workspace_symbols": { + const workspaceSymbols = flattenWorkspaceSymbols(rawResult.symbols ?? []); + const warning = + workspaceSymbols.length > LSP_MAX_SYMBOLS + ? `Results truncated to the first ${LSP_MAX_SYMBOLS} symbols` + : undefined; + const symbols = await Promise.all( + workspaceSymbols.slice(0, LSP_MAX_SYMBOLS).map(async (symbol) => + this.buildSymbolResult(pathMapper, runtime, symbol.uri, symbol.range, symbol.name, symbol.kind, { + detail: symbol.detail, + containerName: symbol.containerName, + }) + ) + ); + return { + operation: rawResult.operation, + serverId, + rootUri, + symbols, + ...(warning ? { warning } : {}), + }; + } + } + } + + private async buildLocationResult(pathMapper: LspPathMapper, runtime: Runtime, location: LspLocation) { + const runtimePath = pathMapper.fromUri(location.uri); + const outputPath = pathMapper.toOutputPath(runtimePath); + return { + path: outputPath, + uri: location.uri, + range: toOneBasedRange(location.range), + preview: await buildPreview(runtime, pathMapper.toReadablePath(runtimePath), location.range), + }; + } + + private async buildSymbolResult( + pathMapper: LspPathMapper, + runtime: Runtime, + uri: string, + range: LspRange, + name: string, + kind: number, + extra?: { detail?: string; containerName?: string } + ) { + const runtimePath = pathMapper.fromUri(uri); + const outputPath = pathMapper.toOutputPath(runtimePath); + return { + name, + kind, + path: outputPath, + range: toOneBasedRange(range), + preview: await buildPreview(runtime, pathMapper.toReadablePath(runtimePath), range), + ...(extra?.detail ? { detail: extra.detail } : {}), + ...(extra?.containerName ? { containerName: extra.containerName } : {}), + }; + } +} + +type PathModule = typeof path.posix; + +function selectPathModule(filePath: string): PathModule { + if (/^[A-Za-z]:[\\/]/.test(filePath) || filePath.includes("\\")) { + return path.win32; + } + return path.posix; +} + +async function pathExists(runtime: Runtime, candidatePath: string): Promise { + try { + await runtime.stat(candidatePath); + return true; + } catch { + return false; + } +} + +function toOneBasedRange(range: LspRange): LspRange { + return { + start: { + line: range.start.line + 1, + character: range.start.character + 1, + }, + end: { + line: range.end.line + 1, + character: range.end.character + 1, + }, + }; +} + +async function buildPreview( + runtime: Runtime, + filePath: string, + range: LspRange +): Promise { + try { + const contents = await readFileString(runtime, filePath); + const lines = contents.split("\n"); + const startLine = Math.max(1, range.start.line + 1 - LSP_PREVIEW_CONTEXT_LINES); + const endLine = Math.min( + lines.length, + Math.max(range.start.line + 1, range.end.line + 1) + LSP_PREVIEW_CONTEXT_LINES + ); + + return lines + .slice(startLine - 1, endLine) + .map((line, index) => `${startLine + index}\t${line}`) + .join("\n"); + } catch { + return undefined; + } +} + +function flattenDocumentSymbols( + symbols: Array, + fallbackUri: string, + containerName?: string +): Array<{ + name: string; + kind: number; + detail?: string; + containerName?: string; + range: LspRange; + uri: string; +}> { + const flattened: Array<{ + name: string; + kind: number; + detail?: string; + containerName?: string; + range: LspRange; + uri: string; + }> = []; + + for (const symbol of symbols) { + if (!isDocumentSymbol(symbol)) { + continue; + } + + const documentSymbol = symbol; + flattened.push({ + name: documentSymbol.name, + kind: documentSymbol.kind, + detail: documentSymbol.detail, + containerName, + range: documentSymbol.selectionRange, + uri: documentSymbol.uri ?? fallbackUri, + }); + + if (documentSymbol.children?.length) { + flattened.push( + ...flattenDocumentSymbols(documentSymbol.children, fallbackUri, documentSymbol.name) + ); + } + } + + return flattened; +} + +function flattenWorkspaceSymbols( + symbols: Array +): Array<{ + name: string; + kind: number; + detail?: string; + containerName?: string; + range: LspRange; + uri: string; +}> { + return symbols + .map((symbol) => { + if (!isWorkspaceSymbolInformation(symbol) || !symbol.location || !("range" in symbol.location)) { + return null; + } + + return { + name: symbol.name, + kind: symbol.kind, + detail: "detail" in symbol ? symbol.detail : undefined, + containerName: symbol.containerName, + range: symbol.location.range, + uri: symbol.location.uri, + }; + }) + .filter((symbol): symbol is NonNullable => symbol !== null); +} + +function isDocumentSymbol( + symbol: LspDocumentSymbol | LspSymbolInformation +): symbol is LspDocumentSymbol { + return "selectionRange" in symbol; +} + +function isWorkspaceSymbolInformation( + symbol: LspDocumentSymbol | LspSymbolInformation +): symbol is LspSymbolInformation { + return "location" in symbol; +} diff --git a/src/node/services/lsp/lspPathMapper.test.ts b/src/node/services/lsp/lspPathMapper.test.ts new file mode 100644 index 0000000000..43288af562 --- /dev/null +++ b/src/node/services/lsp/lspPathMapper.test.ts @@ -0,0 +1,38 @@ +import * as path from "node:path"; +import { describe, expect, test } from "bun:test"; +import { LocalRuntime } from "@/node/runtime/LocalRuntime"; +import { DevcontainerRuntime } from "@/node/runtime/DevcontainerRuntime"; +import { LspPathMapper } from "./lspPathMapper"; + +describe("LspPathMapper", () => { + test("keeps local runtime paths unchanged", () => { + const workspacePath = path.join(process.cwd(), "src"); + const mapper = new LspPathMapper({ + runtime: new LocalRuntime(workspacePath), + workspacePath, + }); + + const runtimePath = mapper.toRuntimePath("browser/App.tsx"); + expect(runtimePath).toBe(path.resolve(workspacePath, "browser/App.tsx")); + expect(mapper.toOutputPath(runtimePath)).toBe(runtimePath); + expect(mapper.fromUri(mapper.toUri(runtimePath))).toBe(runtimePath); + }); + + test("maps devcontainer host paths to remote workspace paths", () => { + const workspacePath = path.join("/Users", "mux", "repo"); + const runtime = Object.create(DevcontainerRuntime.prototype) as DevcontainerRuntime; + runtime.getRemoteWorkspaceFolder = () => "/workspaces/repo"; + + const mapper = new LspPathMapper({ + runtime, + workspacePath, + }); + + const hostPath = path.join(workspacePath, "src", "main.ts"); + const runtimePath = mapper.toRuntimePath(hostPath); + + expect(runtimePath).toBe("/workspaces/repo/src/main.ts"); + expect(mapper.toOutputPath(runtimePath)).toBe(path.join(workspacePath, "src", "main.ts")); + expect(mapper.isWithinWorkspace(runtimePath)).toBe(true); + }); +}); diff --git a/src/node/services/lsp/lspPathMapper.ts b/src/node/services/lsp/lspPathMapper.ts new file mode 100644 index 0000000000..3da7b3a7fe --- /dev/null +++ b/src/node/services/lsp/lspPathMapper.ts @@ -0,0 +1,117 @@ +import * as path from "node:path"; +import { pathToFileURL } from "node:url"; +import { DevcontainerRuntime } from "@/node/runtime/DevcontainerRuntime"; +import type { Runtime } from "@/node/runtime/Runtime"; + +interface LspPathMapperOptions { + runtime: Runtime; + workspacePath: string; +} + +type PathModule = typeof path.posix; + +function selectPathModule(filePath: string): PathModule { + if (/^[A-Za-z]:[\\/]/.test(filePath) || filePath.includes("\\")) { + return path.win32; + } + return path.posix; +} + +function isPathInside(rootPath: string, targetPath: string): boolean { + const pathModule = selectPathModule(targetPath); + const relativePath = pathModule.relative(rootPath, targetPath); + return relativePath === "" || (!relativePath.startsWith("..") && !pathModule.isAbsolute(relativePath)); +} + +function fileUriToPath(uri: string): string { + const parsed = new URL(uri); + const decodedPath = decodeURIComponent(parsed.pathname); + if (/^\/[A-Za-z]:\//.test(decodedPath)) { + return decodedPath.slice(1); + } + return decodedPath; +} + +export class LspPathMapper { + constructor(private readonly options: LspPathMapperOptions) {} + + getWorkspaceRuntimePath(): string { + if (this.options.runtime instanceof DevcontainerRuntime) { + const remoteWorkspaceFolder = this.options.runtime.getRemoteWorkspaceFolder(); + if (!remoteWorkspaceFolder) { + throw new Error("Devcontainer runtime is missing its remote workspace folder"); + } + return remoteWorkspaceFolder; + } + + return this.options.workspacePath; + } + + resolveToolPath(filePath: string): string { + return this.options.runtime.normalizePath(filePath, this.options.workspacePath); + } + + toRuntimePath(filePath: string): string { + const toolPath = this.resolveToolPath(filePath); + + if (!(this.options.runtime instanceof DevcontainerRuntime)) { + return toolPath; + } + + const remoteWorkspaceFolder = this.options.runtime.getRemoteWorkspaceFolder(); + if (!remoteWorkspaceFolder) { + throw new Error("Devcontainer runtime is missing its remote workspace folder"); + } + + if (toolPath === remoteWorkspaceFolder || toolPath.startsWith(`${remoteWorkspaceFolder}/`)) { + return toolPath; + } + + if (!isPathInside(this.options.workspacePath, toolPath)) { + throw new Error( + `LSP paths must stay inside the workspace for devcontainer runtimes (got ${toolPath})` + ); + } + + const relativePath = path.relative(this.options.workspacePath, toolPath); + return relativePath.length === 0 + ? remoteWorkspaceFolder + : path.posix.join(remoteWorkspaceFolder, relativePath.split(path.sep).join("/")); + } + + toOutputPath(runtimePath: string): string { + if (!(this.options.runtime instanceof DevcontainerRuntime)) { + return runtimePath; + } + + const remoteWorkspaceFolder = this.options.runtime.getRemoteWorkspaceFolder(); + if (!remoteWorkspaceFolder) { + return runtimePath; + } + + if (runtimePath === remoteWorkspaceFolder || runtimePath.startsWith(`${remoteWorkspaceFolder}/`)) { + const relativePath = runtimePath.slice(remoteWorkspaceFolder.length).replace(/^\/+/, ""); + return relativePath.length === 0 + ? this.options.workspacePath + : path.join(this.options.workspacePath, relativePath); + } + + return runtimePath; + } + + toReadablePath(runtimePath: string): string { + return this.toOutputPath(runtimePath); + } + + toUri(runtimePath: string): string { + return pathToFileURL(runtimePath).href; + } + + fromUri(uri: string): string { + return fileUriToPath(uri); + } + + isWithinWorkspace(runtimePath: string): boolean { + return isPathInside(this.getWorkspaceRuntimePath(), runtimePath); + } +} diff --git a/src/node/services/lsp/lspServerRegistry.ts b/src/node/services/lsp/lspServerRegistry.ts new file mode 100644 index 0000000000..9341bbb180 --- /dev/null +++ b/src/node/services/lsp/lspServerRegistry.ts @@ -0,0 +1,72 @@ +import * as path from "node:path"; +import type { LspServerDescriptor } from "./types"; + +function ext(filePath: string): string { + return path.extname(filePath).toLowerCase(); +} + +const TYPESCRIPT_EXTENSIONS = [".ts", ".tsx", ".js", ".jsx", ".mjs", ".cjs", ".mts", ".cts"]; + +const typescriptServer: LspServerDescriptor = { + id: "typescript", + extensions: TYPESCRIPT_EXTENSIONS, + command: "typescript-language-server", + args: ["--stdio"], + rootMarkers: ["tsconfig.json", "jsconfig.json", "package.json", ".git"], + languageIdForPath(filePath: string): string { + const fileExt = ext(filePath); + if (fileExt === ".tsx") return "typescriptreact"; + if (fileExt === ".jsx") return "javascriptreact"; + if (fileExt === ".js" || fileExt === ".mjs" || fileExt === ".cjs") return "javascript"; + return "typescript"; + }, +}; + +const pythonServer: LspServerDescriptor = { + id: "python", + extensions: [".py"], + command: "pyright-langserver", + args: ["--stdio"], + rootMarkers: ["pyproject.toml", "setup.py", "setup.cfg", "requirements.txt", ".git"], + languageIdForPath(): string { + return "python"; + }, +}; + +const goServer: LspServerDescriptor = { + id: "go", + extensions: [".go"], + command: "gopls", + args: [], + rootMarkers: ["go.work", "go.mod", ".git"], + languageIdForPath(): string { + return "go"; + }, +}; + +const rustServer: LspServerDescriptor = { + id: "rust", + extensions: [".rs"], + command: "rust-analyzer", + args: [], + rootMarkers: ["Cargo.toml", "rust-project.json", ".git"], + languageIdForPath(): string { + return "rust"; + }, +}; + +export const BUILTIN_LSP_SERVERS: readonly LspServerDescriptor[] = [ + typescriptServer, + pythonServer, + goServer, + rustServer, +]; + +export function findLspServerForFile( + filePath: string, + registry: readonly LspServerDescriptor[] = BUILTIN_LSP_SERVERS +): LspServerDescriptor | null { + const fileExt = ext(filePath); + return registry.find((descriptor) => descriptor.extensions.includes(fileExt)) ?? null; +} + diff --git a/src/node/services/lsp/lspStdioTransport.test.ts b/src/node/services/lsp/lspStdioTransport.test.ts new file mode 100644 index 0000000000..b7e735eac1 --- /dev/null +++ b/src/node/services/lsp/lspStdioTransport.test.ts @@ -0,0 +1,84 @@ +import { describe, expect, test } from "bun:test"; +import type { ExecStream } from "@/node/runtime/Runtime"; +import { LspStdioTransport } from "./lspStdioTransport"; + +const encoder = new TextEncoder(); +const decoder = new TextDecoder(); + +function createReadable(chunks: string[]): ReadableStream { + return new ReadableStream({ + start(controller) { + for (const chunk of chunks) { + controller.enqueue(encoder.encode(chunk)); + } + controller.close(); + }, + }); +} + +function createExecStream(options: { + stdoutChunks: string[]; + stderrChunks?: string[]; + writes?: string[]; + exitCode?: Promise; +}): ExecStream { + return { + stdout: createReadable(options.stdoutChunks), + stderr: createReadable(options.stderrChunks ?? []), + stdin: new WritableStream({ + write(chunk) { + options.writes?.push(decoder.decode(chunk)); + }, + }), + exitCode: options.exitCode ?? Promise.resolve(0), + duration: Promise.resolve(0), + }; +} + +describe("LspStdioTransport", () => { + test("parses Content-Length framed messages", async () => { + const writes: string[] = []; + const message = JSON.stringify({ + jsonrpc: "2.0", + id: 1, + result: { ok: true }, + }); + const framed = `Content-Length: ${message.length}\r\n\r\n${message}`; + const transport = new LspStdioTransport( + createExecStream({ + stdoutChunks: [framed], + stderrChunks: ["warning from server"], + writes, + exitCode: new Promise(() => undefined), + }) + ); + + const messages: unknown[] = []; + transport.onmessage = (parsed) => { + messages.push(parsed); + }; + + transport.start(); + await new Promise((resolve) => setTimeout(resolve, 10)); + + expect(messages).toEqual([ + { + jsonrpc: "2.0", + id: 1, + result: { ok: true }, + }, + ]); + expect(transport.getStderrTail()).toContain("warning from server"); + + await transport.send({ + jsonrpc: "2.0", + id: 2, + method: "initialize", + }); + + expect(writes.join("")).toContain("Content-Length:"); + expect(writes.join("")).toContain("\"method\":\"initialize\""); + + await transport.close(); + }); +}); diff --git a/src/node/services/lsp/lspStdioTransport.ts b/src/node/services/lsp/lspStdioTransport.ts new file mode 100644 index 0000000000..f57f911311 --- /dev/null +++ b/src/node/services/lsp/lspStdioTransport.ts @@ -0,0 +1,186 @@ +import { TextDecoder, TextEncoder } from "node:util"; +import type { ExecStream } from "@/node/runtime/Runtime"; +import { log } from "@/node/services/log"; + +export interface LspJsonRpcMessage { + jsonrpc: "2.0"; + id?: number | string; + method?: string; + params?: unknown; + result?: unknown; + error?: { code: number; message: string; data?: unknown }; +} + +export class LspStdioTransport { + private readonly decoder = new TextDecoder(); + private readonly encoder = new TextEncoder(); + private readonly stdoutReader: ReadableStreamDefaultReader; + private readonly stderrReader: ReadableStreamDefaultReader; + private readonly stdinWriter: WritableStreamDefaultWriter; + private readonly exitPromise: Promise; + private buffer = Buffer.alloc(0); + private running = false; + private closed = false; + private stderrTail = ""; + + onclose?: () => void; + onerror?: (error: Error) => void; + onmessage?: (message: LspJsonRpcMessage) => void; + + constructor(execStream: ExecStream) { + this.stdoutReader = execStream.stdout.getReader(); + this.stderrReader = execStream.stderr.getReader(); + this.stdinWriter = execStream.stdin.getWriter(); + this.exitPromise = execStream.exitCode; + void this.exitPromise.then(() => { + this.closed = true; + this.onclose?.(); + }); + } + + start(): void { + if (this.running) { + return; + } + + this.running = true; + void this.readLoop(); + void this.drainStderr(); + } + + async send(message: LspJsonRpcMessage): Promise { + const body = this.encoder.encode(JSON.stringify(message)); + const header = this.encoder.encode(`Content-Length: ${body.byteLength}\r\n\r\n`); + await this.stdinWriter.write(header); + await this.stdinWriter.write(body); + } + + getStderrTail(): string { + return this.stderrTail.trim(); + } + + isClosed(): boolean { + return this.closed; + } + + async close(): Promise { + this.closed = true; + + try { + await this.stdinWriter.close(); + } catch (error) { + log.debug("Failed to close LSP stdin writer", { error }); + } + + try { + await this.stdoutReader.cancel(); + } catch (error) { + log.debug("Failed to cancel LSP stdout reader", { error }); + } + + try { + await this.stderrReader.cancel(); + } catch (error) { + log.debug("Failed to cancel LSP stderr reader", { error }); + } + } + + private async drainStderr(): Promise { + try { + while (true) { + const { value, done } = await this.stderrReader.read(); + if (done) { + break; + } + + if (!value) { + continue; + } + + this.stderrTail += this.decoder.decode(value, { stream: true }); + if (this.stderrTail.length > 4096) { + this.stderrTail = this.stderrTail.slice(-4096); + } + } + } catch (error) { + this.handleError(error); + } + } + + private async readLoop(): Promise { + try { + while (true) { + const { value, done } = await this.stdoutReader.read(); + if (done) { + break; + } + + if (!value) { + continue; + } + + this.buffer = Buffer.concat([this.buffer, Buffer.from(value)]); + this.processBuffer(); + } + } catch (error) { + this.handleError(error); + } finally { + this.closed = true; + this.onclose?.(); + } + } + + private processBuffer(): void { + while (true) { + const headerEnd = this.buffer.indexOf("\r\n\r\n"); + if (headerEnd === -1) { + return; + } + + const headerText = this.buffer.subarray(0, headerEnd).toString("utf8"); + const contentLength = this.parseContentLength(headerText); + if (contentLength == null) { + this.handleError(new Error(`Invalid LSP header block: ${headerText}`)); + return; + } + + const bodyStart = headerEnd + 4; + const bodyEnd = bodyStart + contentLength; + if (this.buffer.length < bodyEnd) { + return; + } + + const body = this.buffer.subarray(bodyStart, bodyEnd); + this.buffer = this.buffer.subarray(bodyEnd); + + try { + const parsed = JSON.parse(body.toString("utf8")) as LspJsonRpcMessage; + this.onmessage?.(parsed); + } catch (error) { + this.handleError(error); + } + } + } + + private parseContentLength(headerText: string): number | null { + const headerLines = headerText.split("\r\n"); + const contentLengthLine = headerLines.find((line) => line.toLowerCase().startsWith("content-length:")); + if (!contentLengthLine) { + return null; + } + + const rawLength = contentLengthLine.split(":")[1]?.trim(); + const parsedLength = rawLength ? Number.parseInt(rawLength, 10) : Number.NaN; + return Number.isFinite(parsedLength) ? parsedLength : null; + } + + private handleError(error: unknown): void { + const typedError = error instanceof Error ? error : new Error(String(error)); + if (this.onerror) { + this.onerror(typedError); + return; + } + log.error("LSP stdio transport error", { error: typedError }); + } +} + diff --git a/src/node/services/lsp/types.ts b/src/node/services/lsp/types.ts new file mode 100644 index 0000000000..e68862c1ad --- /dev/null +++ b/src/node/services/lsp/types.ts @@ -0,0 +1,134 @@ +import type { Runtime } from "@/node/runtime/Runtime"; + +export type LspQueryOperation = + | "hover" + | "definition" + | "references" + | "implementation" + | "document_symbols" + | "workspace_symbols"; + +export interface LspPosition { + line: number; + character: number; +} + +export interface LspRange { + start: LspPosition; + end: LspPosition; +} + +export interface LspLocation { + uri: string; + range: LspRange; +} + +export interface LspLocationLink { + targetUri: string; + targetRange: LspRange; +} + +export interface LspMarkupContent { + kind: "plaintext" | "markdown"; + value: string; +} + +export interface LspMarkedString { + language: string; + value: string; +} + +export interface LspHover { + contents: string | LspMarkupContent | LspMarkedString | Array; +} + +export interface LspDocumentSymbol { + name: string; + kind: number; + detail?: string; + range: LspRange; + selectionRange: LspRange; + uri?: string; + children?: LspDocumentSymbol[]; +} + +export interface LspSymbolInformation { + name: string; + kind: number; + detail?: string; + location?: LspLocation | { uri: string }; + containerName?: string; +} + +export interface LspSymbolResult { + name: string; + kind: number; + detail?: string; + containerName?: string; + path: string; + range: LspRange; + preview?: string; +} + +export interface LspLocationResult { + path: string; + uri: string; + range: LspRange; + preview?: string; +} + +export interface LspManagerQueryResult { + operation: LspQueryOperation; + serverId: string; + rootUri: string; + hover?: string; + locations?: LspLocationResult[]; + symbols?: LspSymbolResult[]; + warning?: string; +} + +export interface LspServerDescriptor { + id: string; + extensions: readonly string[]; + command: string; + args: readonly string[]; + rootMarkers: readonly string[]; + languageIdForPath(filePath: string): string; +} + +export interface LspClientFileHandle { + runtimePath: string; + readablePath: string; + uri: string; + languageId: string; +} + +export interface LspClientQueryRequest { + operation: LspQueryOperation; + file: LspClientFileHandle; + line?: number; + character?: number; + query?: string; + includeDeclaration?: boolean; +} + +export interface LspClientQueryResult { + operation: LspQueryOperation; + hover?: string; + locations?: LspLocation[]; + symbols?: Array; +} + +export interface LspClientInstance { + readonly isClosed: boolean; + ensureFile(file: LspClientFileHandle): Promise; + query(request: LspClientQueryRequest): Promise; + close(): Promise; +} + +export interface CreateLspClientOptions { + descriptor: LspServerDescriptor; + runtime: Runtime; + rootPath: string; + rootUri: string; +} diff --git a/src/node/services/serviceContainer.ts b/src/node/services/serviceContainer.ts index 5aa65bb83e..5457066a40 100644 --- a/src/node/services/serviceContainer.ts +++ b/src/node/services/serviceContainer.ts @@ -110,6 +110,7 @@ export class ServiceContainer { public readonly providerService: CoreServices["providerService"]; public readonly mcpConfigService: CoreServices["mcpConfigService"]; public readonly mcpServerManager: CoreServices["mcpServerManager"]; + public readonly lspManager: CoreServices["lspManager"]; public readonly sessionUsageService: CoreServices["sessionUsageService"]; private readonly extensionMetadata: CoreServices["extensionMetadata"]; private readonly backgroundProcessManager: CoreServices["backgroundProcessManager"]; @@ -233,6 +234,7 @@ export class ServiceContainer { this.providerService = core.providerService; this.mcpConfigService = core.mcpConfigService; this.mcpServerManager = core.mcpServerManager; + this.lspManager = core.lspManager; this.sessionUsageService = core.sessionUsageService; this.extensionMetadata = core.extensionMetadata; this.backgroundProcessManager = core.backgroundProcessManager; @@ -687,6 +689,7 @@ export class ServiceContainer { await this.analyticsService.dispose(); this.policyService.dispose(); this.mcpServerManager.dispose(); + await this.lspManager.dispose(); await this.mcpOauthService.dispose(); await this.muxGatewayOauthService.dispose(); await this.muxGovernorOauthService.dispose(); diff --git a/src/node/services/systemMessage.ts b/src/node/services/systemMessage.ts index 15f4ce179e..5ad5445c6a 100644 --- a/src/node/services/systemMessage.ts +++ b/src/node/services/systemMessage.ts @@ -256,6 +256,7 @@ export function extractToolInstructions( options?: { enableAgentReport?: boolean; enableMuxGlobalAgentsTools?: boolean; + enableLspQuery?: boolean; agentInstructions?: string; } ): Record { @@ -293,7 +294,8 @@ export async function readToolInstructions( runtime: Runtime, workspacePath: string, modelString: string, - agentInstructions?: string + agentInstructions?: string, + options?: { enableLspQuery?: boolean } ): Promise> { const [globalInstructions, contextInstructions] = await readInstructionSources( metadata, @@ -305,6 +307,7 @@ export async function readToolInstructions( ...getToolAvailabilityOptions({ workspaceId: metadata.id, parentWorkspaceId: metadata.parentWorkspaceId, + enableLspQuery: options?.enableLspQuery, }), agentInstructions, }); diff --git a/src/node/services/tools/lsp_query.test.ts b/src/node/services/tools/lsp_query.test.ts new file mode 100644 index 0000000000..89323bac58 --- /dev/null +++ b/src/node/services/tools/lsp_query.test.ts @@ -0,0 +1,83 @@ +import { describe, expect, mock, test } from "bun:test"; +import type { ToolExecutionOptions } from "ai"; +import { LspManager } from "@/node/services/lsp/lspManager"; +import { createTestToolConfig } from "./testHelpers"; +import { createLspQueryTool } from "./lsp_query"; +import type { LspQueryToolResult } from "@/common/types/tools"; + +const mockToolCallOptions: ToolExecutionOptions = { + toolCallId: "tool-call-id", + messages: [], +}; + +describe("lsp_query tool", () => { + test("returns formatted LSP data from the manager", async () => { + const lspManager = new LspManager({ registry: [] }); + const query = mock(() => + Promise.resolve({ + operation: "hover" as const, + serverId: "typescript", + rootUri: "file:///tmp/workspace", + hover: "const value: 1", + }) + ); + lspManager.query = query; + const config = createTestToolConfig(process.cwd()); + config.lspManager = lspManager; + const configuredTool = createLspQueryTool(config); + + try { + const result = (await configuredTool.execute!( + { + operation: "hover", + path: "src/browser/App.tsx", + line: 1, + column: 1, + }, + mockToolCallOptions + )) as LspQueryToolResult; + + expect(result.success).toBe(true); + if (result.success) { + expect(result.serverId).toBe("typescript"); + expect(result.hover).toBe("const value: 1"); + } + expect(query).toHaveBeenCalledTimes(1); + } finally { + await lspManager.dispose(); + } + }); + + test("validates required position arguments for hover-like operations", async () => { + const config = createTestToolConfig(process.cwd()); + const lspManager = new LspManager({ registry: [] }); + const query = mock(() => + Promise.resolve({ + operation: "hover" as const, + serverId: "typescript", + rootUri: "file:///tmp/workspace", + hover: "", + }) + ); + lspManager.query = query; + config.lspManager = lspManager; + const tool = createLspQueryTool(config); + + try { + const result = (await tool.execute!( + { + operation: "definition", + path: "src/browser/App.tsx", + }, + mockToolCallOptions + )) as LspQueryToolResult; + + expect(result).toEqual({ + success: false, + error: "definition requires both line and column", + }); + } finally { + await lspManager.dispose(); + } + }); +}); diff --git a/src/node/services/tools/lsp_query.ts b/src/node/services/tools/lsp_query.ts new file mode 100644 index 0000000000..363f2a24d4 --- /dev/null +++ b/src/node/services/tools/lsp_query.ts @@ -0,0 +1,95 @@ +import { tool } from "ai"; +import type { ToolFactory, ToolConfiguration } from "@/common/utils/tools/tools"; +import { TOOL_DEFINITIONS } from "@/common/utils/tools/toolDefinitions"; +import type { LspQueryToolResult } from "@/common/types/tools"; +import { resolvePathWithinCwd, validatePathInCwd } from "./fileCommon"; + +function validateArguments(args: { + operation: string; + line?: number | null; + column?: number | null; + query?: string | null; +}): string | null { + switch (args.operation) { + case "hover": + case "definition": + case "references": + case "implementation": + if (args.line == null || args.column == null) { + return `${args.operation} requires both line and column`; + } + return null; + case "workspace_symbols": + if (!args.query?.trim()) { + return "workspace_symbols requires a non-empty query"; + } + return null; + case "document_symbols": + return null; + default: + return `Unsupported LSP operation: ${args.operation}`; + } +} + +export const createLspQueryTool: ToolFactory = (config: ToolConfiguration) => + tool({ + description: TOOL_DEFINITIONS.lsp_query.description, + inputSchema: TOOL_DEFINITIONS.lsp_query.schema, + execute: async (args): Promise => { + if (!config.lspManager || !config.workspaceId) { + return { + success: false, + error: "LSP query tool is unavailable in this workspace", + }; + } + + const argumentError = validateArguments(args); + if (argumentError) { + return { + success: false, + error: argumentError, + }; + } + + const pathError = validatePathInCwd(args.path, config.cwd, config.runtime); + if (pathError) { + return { + success: false, + error: pathError.error, + }; + } + + const { resolvedPath, warning: pathWarning } = resolvePathWithinCwd( + args.path, + config.cwd, + config.runtime + ); + + try { + const result = await config.lspManager.query({ + workspaceId: config.workspaceId, + runtime: config.runtime, + workspacePath: config.cwd, + filePath: resolvedPath, + operation: args.operation, + line: args.line ?? undefined, + column: args.column ?? undefined, + query: args.query ?? undefined, + includeDeclaration: args.includeDeclaration ?? undefined, + }); + + return { + success: true, + ...result, + ...(pathWarning ? { warning: pathWarning } : {}), + }; + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + return { + success: false, + error: message, + ...(pathWarning ? { warning: pathWarning } : {}), + }; + } + }, + }); diff --git a/src/node/services/workspaceService.ts b/src/node/services/workspaceService.ts index 2f691308d3..67f436b97b 100644 --- a/src/node/services/workspaceService.ts +++ b/src/node/services/workspaceService.ts @@ -141,6 +141,7 @@ import type { BackgroundProcessManager } from "@/node/services/backgroundProcess import type { WorkspaceLifecycleHooks } from "@/node/services/workspaceLifecycleHooks"; import type { TaskService } from "@/node/services/taskService"; import type { WorktreeArchiveSnapshotService } from "@/node/services/worktreeArchiveSnapshotService"; +import type { LspManager } from "@/node/services/lsp/lspManager"; import { DisposableTempDir } from "@/node/services/tempDir"; import { createBashTool } from "@/node/services/tools/bash"; @@ -1241,6 +1242,7 @@ export class WorkspaceService extends EventEmitter { private readonly telemetryService?: TelemetryService; private readonly experimentsService?: ExperimentsService; private mcpServerManager?: MCPServerManager; + private lspManager?: LspManager; // Optional services for workspace cleanup during archive/remove lifecycle operations. private terminalService?: TerminalService; private desktopSessionManager?: DesktopSessionManager; @@ -1257,6 +1259,10 @@ export class WorkspaceService extends EventEmitter { this.mcpServerManager = manager; } + setLspManager(manager: LspManager): void { + this.lspManager = manager; + } + /** * Set the terminal service for cleanup on workspace removal. */ @@ -1281,6 +1287,19 @@ export class WorkspaceService extends EventEmitter { } } + private async disposeLspWorkspaceBestEffort( + workspaceId: string, + reason: "archive" | "remove" | "rename" + ): Promise { + try { + await this.lspManager?.disposeWorkspace(workspaceId); + } catch (error) { + log.debug( + `Failed to dispose LSP state during ${reason} for workspace ${workspaceId}: ${getErrorMessage(error)}` + ); + } + } + setWorkspaceLifecycleHooks(hooks: WorkspaceLifecycleHooks): void { this.workspaceLifecycleHooks = hooks; } @@ -1428,6 +1447,7 @@ export class WorkspaceService extends EventEmitter { // Archiving hides workspace UI; do not leave terminal PTYs or desktop sessions running headless. this.terminalService?.closeWorkspaceSessions(workspaceId); await this.closeDesktopSessionBestEffort(workspaceId, "archive"); + await this.disposeLspWorkspaceBestEffort(workspaceId, "archive"); } /** @@ -2813,6 +2833,8 @@ export class WorkspaceService extends EventEmitter { } } + await this.disposeLspWorkspaceBestEffort(workspaceId, "remove"); + let parentWorkspaceId: string | null = null; let childTaskModelString: string | undefined; let childTaskThinkingLevel: ThinkingLevel | undefined; @@ -3398,6 +3420,8 @@ export class WorkspaceService extends EventEmitter { return Ok({ newWorkspaceId: workspaceId }); } + await this.disposeLspWorkspaceBestEffort(workspaceId, "rename"); + const allWorkspaces = await this.config.getAllWorkspaceMetadata(); const collision = allWorkspaces.find( (ws) => (ws.name === newName || ws.id === newName) && ws.id !== workspaceId From 0fb38f144397f788dc2834eb3370701ba4b6ba02 Mon Sep 17 00:00:00 2001 From: terion Date: Mon, 13 Apr 2026 11:59:04 +0300 Subject: [PATCH 02/46] address review comments --- src/node/services/lsp/lspManager.test.ts | 83 +++++++++++++++++ src/node/services/lsp/lspManager.ts | 56 ++++++++++-- .../services/lsp/lspStdioTransport.test.ts | 88 +++++++++++++++++-- src/node/services/lsp/lspStdioTransport.ts | 15 +++- 4 files changed, 228 insertions(+), 14 deletions(-) diff --git a/src/node/services/lsp/lspManager.test.ts b/src/node/services/lsp/lspManager.test.ts index 2a4f13a56f..c6008e65ef 100644 --- a/src/node/services/lsp/lspManager.test.ts +++ b/src/node/services/lsp/lspManager.test.ts @@ -6,6 +6,16 @@ import { LocalRuntime } from "@/node/runtime/LocalRuntime"; import type { CreateLspClientOptions, LspClientInstance, LspServerDescriptor } from "./types"; import { LspManager } from "./lspManager"; +function createDeferred() { + let resolve!: (value: T | PromiseLike) => void; + let reject!: (reason?: unknown) => void; + const promise = new Promise((promiseResolve, promiseReject) => { + resolve = promiseResolve; + reject = promiseReject; + }); + return { promise, resolve, reject }; +} + describe("LspManager", () => { let workspacePath: string; @@ -98,4 +108,77 @@ describe("LspManager", () => { await manager.dispose(); expect(close).toHaveBeenCalledTimes(1); }); + + test("deduplicates concurrent client creation for the same workspace root", async () => { + const ensureFile = mock(() => Promise.resolve(undefined)); + const query = mock(() => + Promise.resolve({ + operation: "hover" as const, + hover: "const value: 1", + }) + ); + const close = mock(() => Promise.resolve(undefined)); + const client: LspClientInstance = { + isClosed: false, + ensureFile, + query, + close, + }; + const clientReady = createDeferred(); + const clientFactoryStarted = createDeferred(); + const clientFactory = mock((_options: CreateLspClientOptions): Promise => { + clientFactoryStarted.resolve(); + return clientReady.promise; + }); + const registry: readonly LspServerDescriptor[] = [ + { + id: "typescript", + extensions: [".ts"], + command: "fake-lsp", + args: ["--stdio"], + rootMarkers: ["package.json", ".git"], + languageIdForPath: () => "typescript", + }, + ]; + + const manager = new LspManager({ + registry, + clientFactory, + }); + const runtime = new LocalRuntime(workspacePath); + + const firstQuery = manager.query({ + workspaceId: "ws-1", + runtime, + workspacePath, + filePath: "src/example.ts", + operation: "hover", + line: 1, + column: 1, + }); + await clientFactoryStarted.promise; + + const secondQuery = manager.query({ + workspaceId: "ws-1", + runtime, + workspacePath, + filePath: "src/example.ts", + operation: "hover", + line: 1, + column: 1, + }); + + await new Promise((resolve) => setTimeout(resolve, 10)); + expect(clientFactory).toHaveBeenCalledTimes(1); + + clientReady.resolve(client); + const [firstResult, secondResult] = await Promise.all([firstQuery, secondQuery]); + + expect(firstResult.hover).toBe("const value: 1"); + expect(secondResult.hover).toBe("const value: 1"); + expect(ensureFile).toHaveBeenCalledTimes(2); + + await manager.dispose(); + expect(close).toHaveBeenCalledTimes(1); + }); }); diff --git a/src/node/services/lsp/lspManager.ts b/src/node/services/lsp/lspManager.ts index 8224e920be..2554ef429e 100644 --- a/src/node/services/lsp/lspManager.ts +++ b/src/node/services/lsp/lspManager.ts @@ -28,6 +28,7 @@ import type { interface WorkspaceClients { clients: Map; + pendingClients: Map>; lastActivity: number; } @@ -147,8 +148,16 @@ export class LspManager { } this.workspaceClients.delete(workspaceId); + const pendingClients = await Promise.allSettled(entry.pendingClients.values()); + const clientsToClose = new Set(entry.clients.values()); + for (const result of pendingClients) { + if (result.status === "fulfilled") { + clientsToClose.add(result.value); + } + } + await Promise.all( - [...entry.clients.values()].map(async (client) => { + [...clientsToClose].map(async (client) => { try { await client.close(); } catch (error) { @@ -176,6 +185,7 @@ export class LspManager { ): Promise { const workspaceEntry = this.workspaceClients.get(workspaceId) ?? { clients: new Map(), + pendingClients: new Map>(), lastActivity: Date.now(), }; this.workspaceClients.set(workspaceId, workspaceEntry); @@ -191,15 +201,49 @@ export class LspManager { workspaceEntry.clients.delete(clientKey); } - const client = await this.clientFactory({ + const pendingClient = workspaceEntry.pendingClients.get(clientKey); + if (pendingClient) { + return await pendingClient; + } + + // Deduplicate concurrent queries for the same workspace/root so we never spawn + // orphaned LSP processes that escape manager tracking. + const clientPromise = this.clientFactory({ descriptor, runtime, rootPath, rootUri, - }); - workspaceEntry.clients.set(clientKey, client); - workspaceEntry.lastActivity = Date.now(); - return client; + }) + .then(async (client) => { + const currentWorkspaceEntry = this.workspaceClients.get(workspaceId); + if (currentWorkspaceEntry !== workspaceEntry) { + try { + await client.close(); + } catch (error) { + log.debug("Failed to close LSP client created after workspace disposal", { + workspaceId, + clientKey, + error, + }); + } + throw new Error( + `LSP workspace ${workspaceId} was disposed while starting ${descriptor.id}` + ); + } + + workspaceEntry.clients.set(clientKey, client); + workspaceEntry.lastActivity = Date.now(); + return client; + }) + .finally(() => { + const currentWorkspaceEntry = this.workspaceClients.get(workspaceId); + if (currentWorkspaceEntry === workspaceEntry) { + workspaceEntry.pendingClients.delete(clientKey); + } + }); + + workspaceEntry.pendingClients.set(clientKey, clientPromise); + return await clientPromise; } private markActivity(workspaceId: string): void { diff --git a/src/node/services/lsp/lspStdioTransport.test.ts b/src/node/services/lsp/lspStdioTransport.test.ts index b7e735eac1..483d4afd8e 100644 --- a/src/node/services/lsp/lspStdioTransport.test.ts +++ b/src/node/services/lsp/lspStdioTransport.test.ts @@ -21,20 +21,33 @@ function createExecStream(options: { stderrChunks?: string[]; writes?: string[]; exitCode?: Promise; + stdin?: WritableStream; }): ExecStream { return { stdout: createReadable(options.stdoutChunks), stderr: createReadable(options.stderrChunks ?? []), - stdin: new WritableStream({ - write(chunk) { - options.writes?.push(decoder.decode(chunk)); - }, - }), + stdin: + options.stdin ?? + new WritableStream({ + write(chunk) { + options.writes?.push(decoder.decode(chunk)); + }, + }), exitCode: options.exitCode ?? Promise.resolve(0), duration: Promise.resolve(0), }; } +function createDeferred() { + let resolve!: (value: T | PromiseLike) => void; + let reject!: (reason?: unknown) => void; + const promise = new Promise((promiseResolve, promiseReject) => { + resolve = promiseResolve; + reject = promiseReject; + }); + return { promise, resolve, reject }; +} + describe("LspStdioTransport", () => { test("parses Content-Length framed messages", async () => { const writes: string[] = []; @@ -81,4 +94,69 @@ describe("LspStdioTransport", () => { await transport.close(); }); + + test("serializes concurrent writes so framed messages do not interleave", async () => { + const writes: string[] = []; + let activeWrites = 0; + let maxActiveWrites = 0; + let writeCount = 0; + const releaseFirstWrite = createDeferred(); + const stdin = new WritableStream({ + async write(chunk) { + writeCount += 1; + activeWrites += 1; + maxActiveWrites = Math.max(maxActiveWrites, activeWrites); + writes.push(decoder.decode(chunk)); + + if (writeCount === 1) { + await releaseFirstWrite.promise; + } + + activeWrites -= 1; + }, + }); + const transport = new LspStdioTransport( + createExecStream({ + stdoutChunks: [], + stdin, + exitCode: new Promise(() => undefined), + }) + ); + + const firstSend = transport.send({ + jsonrpc: "2.0", + id: 1, + method: "initialize", + }); + const secondSend = transport.send({ + jsonrpc: "2.0", + id: 2, + method: "initialized", + }); + + await new Promise((resolve) => setTimeout(resolve, 10)); + expect(writeCount).toBe(1); + + releaseFirstWrite.resolve(); + await Promise.all([firstSend, secondSend]); + + const expectedFirstBody = JSON.stringify({ + jsonrpc: "2.0", + id: 1, + method: "initialize", + }); + const expectedSecondBody = JSON.stringify({ + jsonrpc: "2.0", + id: 2, + method: "initialized", + }); + + expect(maxActiveWrites).toBe(1); + expect(writes).toEqual([ + `Content-Length: ${expectedFirstBody.length}\r\n\r\n${expectedFirstBody}`, + `Content-Length: ${expectedSecondBody.length}\r\n\r\n${expectedSecondBody}`, + ]); + + await transport.close(); + }); }); diff --git a/src/node/services/lsp/lspStdioTransport.ts b/src/node/services/lsp/lspStdioTransport.ts index f57f911311..a0f3bf1158 100644 --- a/src/node/services/lsp/lspStdioTransport.ts +++ b/src/node/services/lsp/lspStdioTransport.ts @@ -22,6 +22,7 @@ export class LspStdioTransport { private running = false; private closed = false; private stderrTail = ""; + private sendQueue = Promise.resolve(); onclose?: () => void; onerror?: (error: Error) => void; @@ -51,8 +52,17 @@ export class LspStdioTransport { async send(message: LspJsonRpcMessage): Promise { const body = this.encoder.encode(JSON.stringify(message)); const header = this.encoder.encode(`Content-Length: ${body.byteLength}\r\n\r\n`); - await this.stdinWriter.write(header); - await this.stdinWriter.write(body); + const frame = new Uint8Array(header.byteLength + body.byteLength); + frame.set(header); + frame.set(body, header.byteLength); + + // Serialize framed writes so concurrent requests cannot interleave bytes and + // corrupt Content-Length boundaries on the shared stdio stream. + const writePromise = this.sendQueue.then(async () => { + await this.stdinWriter.write(frame); + }); + this.sendQueue = writePromise.catch(() => undefined); + await writePromise; } getStderrTail(): string { @@ -183,4 +193,3 @@ export class LspStdioTransport { log.error("LSP stdio transport error", { error: typedError }); } } - From 115e5f7f08cd1e2b37eba86596172b0cc46aa713 Mon Sep 17 00:00:00 2001 From: terion Date: Mon, 13 Apr 2026 13:16:20 +0300 Subject: [PATCH 03/46] Add post-edit LSP diagnostics for mutating tools --- src/common/utils/tools/tools.ts | 2 + src/constants/lsp.ts | 1 + src/node/services/aiService.ts | 98 ++++ src/node/services/lsp/lspClient.test.ts | 115 ++++ src/node/services/lsp/lspClient.ts | 186 +++++-- src/node/services/lsp/lspManager.test.ts | 166 +++++- src/node/services/lsp/lspManager.ts | 512 ++++++++++++++++-- src/node/services/lsp/types.ts | 27 +- .../services/tools/file_edit_insert.test.ts | 27 + src/node/services/tools/file_edit_insert.ts | 19 +- .../tools/file_edit_operation.test.ts | 54 ++ .../services/tools/file_edit_operation.ts | 24 +- .../tools/task_apply_git_patch.test.ts | 109 ++++ .../services/tools/task_apply_git_patch.ts | 80 ++- src/node/services/tools/testHelpers.ts | 2 + 15 files changed, 1289 insertions(+), 133 deletions(-) create mode 100644 src/node/services/lsp/lspClient.test.ts diff --git a/src/common/utils/tools/tools.ts b/src/common/utils/tools/tools.ts index 2eb536545a..3bf4d7e660 100644 --- a/src/common/utils/tools/tools.ts +++ b/src/common/utils/tools/tools.ts @@ -98,6 +98,8 @@ export interface ToolConfiguration { muxScope?: MuxToolScope; /** Callback to record file state for external edit detection (plan files) */ recordFileState?: (filePath: string, state: FileState) => Promise; + /** Optional callback for file-mutation follow-up work (for example, post-edit diagnostics). */ + onFilesMutated?: (params: { filePaths: string[] }) => Promise; /** Callback to notify that provider/config was written (triggers hot-reload). */ onConfigChanged?: () => void; /** Task orchestration for sub-agent tasks */ diff --git a/src/constants/lsp.ts b/src/constants/lsp.ts index 0d64cd72e9..dab5134cb4 100644 --- a/src/constants/lsp.ts +++ b/src/constants/lsp.ts @@ -2,6 +2,7 @@ export const LSP_IDLE_TIMEOUT_MS = 10 * 60 * 1000; export const LSP_IDLE_CHECK_INTERVAL_MS = 60 * 1000; export const LSP_REQUEST_TIMEOUT_MS = 10 * 1000; export const LSP_START_TIMEOUT_MS = 5 * 1000; +export const LSP_POST_MUTATION_DIAGNOSTICS_TIMEOUT_MS = 1_500; export const LSP_MAX_LOCATIONS = 25; export const LSP_MAX_SYMBOLS = 100; export const LSP_PREVIEW_CONTEXT_LINES = 1; diff --git a/src/node/services/aiService.ts b/src/node/services/aiService.ts index 56addc2e94..e142ba599b 100644 --- a/src/node/services/aiService.ts +++ b/src/node/services/aiService.ts @@ -1,3 +1,4 @@ +import * as path from "node:path"; import * as fs from "fs/promises"; import { EventEmitter } from "events"; @@ -100,6 +101,7 @@ import { applyToolPolicyAndExperiments, captureMcpToolTelemetry } from "./toolAs import { getErrorMessage } from "@/common/utils/errors"; import { isProjectTrusted } from "@/node/utils/projectTrust"; import type { LspManager } from "@/node/services/lsp/lspManager"; +import type { LspDiagnostic, LspFileDiagnostics } from "@/node/services/lsp/types"; const STREAM_STARTUP_DIAGNOSTIC_THRESHOLD_MS = 1_000; @@ -206,6 +208,80 @@ function derivePromptCacheScope(metadata: WorkspaceMetadata): string { return `${metadata.projectName}-${uniqueSuffix([metadata.projectPath])}`; } +const MAX_POST_EDIT_DIAGNOSTIC_LINES = 12; + +function formatPostMutationDiagnostics( + diagnostics: LspFileDiagnostics[], + workspacePath: string +): string | undefined { + if (diagnostics.length === 0) { + return undefined; + } + + const lines: string[] = []; + let omittedCount = 0; + + for (const fileDiagnostics of diagnostics) { + for (const diagnostic of fileDiagnostics.diagnostics) { + if (lines.length >= MAX_POST_EDIT_DIAGNOSTIC_LINES) { + omittedCount += 1; + continue; + } + + lines.push(formatPostMutationDiagnosticLine(fileDiagnostics, diagnostic, workspacePath)); + } + } + + if (lines.length === 0) { + return undefined; + } + + return [ + "Post-edit LSP diagnostics:", + ...lines.map((line) => `- ${line}`), + ...(omittedCount > 0 ? [`- ...and ${omittedCount} more diagnostics`] : []), + ].join("\n"); +} + +function formatPostMutationDiagnosticLine( + fileDiagnostics: LspFileDiagnostics, + diagnostic: LspDiagnostic, + workspacePath: string +): string { + const relativePath = toWorkspaceRelativePath(fileDiagnostics.path, workspacePath); + const line = diagnostic.range.start.line + 1; + const column = diagnostic.range.start.character + 1; + const severity = formatDiagnosticSeverity(diagnostic.severity); + const source = diagnostic.source ? ` ${diagnostic.source}` : ""; + const code = diagnostic.code != null ? ` ${String(diagnostic.code)}` : ""; + const message = diagnostic.message.replace(/\s+/g, " ").trim(); + return `${relativePath}:${line}:${column} ${severity}${source}${code}: ${message}`; +} + +function formatDiagnosticSeverity(severity: LspDiagnostic["severity"]): string { + switch (severity) { + case 1: + return "error"; + case 2: + return "warning"; + case 3: + return "info"; + case 4: + return "hint"; + default: + return "diagnostic"; + } +} + +function toWorkspaceRelativePath(filePath: string, workspacePath: string): string { + const relativePath = path.relative(workspacePath, filePath); + if (relativePath.length === 0) { + return "."; + } + + return relativePath.startsWith("..") ? filePath : relativePath; +} + export class AIService extends EventEmitter { private readonly streamManager: StreamManager; private readonly historyService: HistoryService; @@ -1280,6 +1356,28 @@ export class AIService extends EventEmitter { enableAgentReport: Boolean(metadata.parentWorkspaceId), // External edit detection callback recordFileState, + onFilesMutated: async ({ filePaths }) => { + if (!this.lspManager) { + return undefined; + } + + try { + const diagnostics = await this.lspManager.collectPostMutationDiagnostics({ + workspaceId, + runtime, + workspacePath, + filePaths, + }); + return formatPostMutationDiagnostics(diagnostics, workspacePath); + } catch (error) { + log.debug("Failed to collect post-mutation LSP diagnostics", { + workspaceId, + filePaths, + error, + }); + return undefined; + } + }, onConfigChanged: () => this.providerService.notifyConfigChanged(), taskService: this.taskService, analyticsService: this.analyticsService, diff --git a/src/node/services/lsp/lspClient.test.ts b/src/node/services/lsp/lspClient.test.ts new file mode 100644 index 0000000000..971c5d43ac --- /dev/null +++ b/src/node/services/lsp/lspClient.test.ts @@ -0,0 +1,115 @@ +import { describe, expect, it, mock } from "bun:test"; +import { LspClient } from "./lspClient"; +import type { CreateLspClientOptions, LspPublishDiagnosticsParams, LspServerDescriptor } from "./types"; + +function createDescriptor(): LspServerDescriptor { + return { + id: "typescript", + extensions: [".ts"], + command: "fake-lsp", + args: ["--stdio"], + rootMarkers: ["package.json", ".git"], + languageIdForPath: () => "typescript", + }; +} + +function createTransport() { + return { + onmessage: undefined as ((message: { jsonrpc: "2.0"; id?: number | string; method?: string; params?: unknown }) => void) | undefined, + onclose: undefined as (() => void) | undefined, + onerror: undefined as ((error: Error) => void) | undefined, + start: mock(() => undefined), + send: mock(() => Promise.resolve()), + close: mock(() => Promise.resolve()), + isClosed: () => false, + getStderrTail: () => "", + }; +} + +function createClient(options?: Partial) { + const transport = createTransport(); + const client = new (LspClient as unknown as { + new (clientOptions: CreateLspClientOptions, transport: ReturnType): LspClient; + })( + { + descriptor: createDescriptor(), + runtime: {} as never, + rootPath: "/tmp/workspace", + rootUri: "file:///tmp/workspace", + ...options, + }, + transport + ); + + return { + client, + transport, + }; +} + +describe("LspClient publishDiagnostics handling", () => { + it("forwards valid publishDiagnostics notifications", () => { + const onPublishDiagnostics = mock((_params: LspPublishDiagnosticsParams) => undefined); + const { transport } = createClient({ onPublishDiagnostics }); + + transport.onmessage?.({ + jsonrpc: "2.0", + method: "textDocument/publishDiagnostics", + params: { + uri: "file:///tmp/workspace/src/example.ts", + version: 2, + diagnostics: [ + { + range: { + start: { line: 0, character: 4 }, + end: { line: 0, character: 9 }, + }, + severity: 1, + code: "TS2322", + source: "tsserver", + message: "Type 'string' is not assignable to type 'number'.", + }, + ], + }, + }); + + expect(onPublishDiagnostics).toHaveBeenCalledTimes(1); + expect(onPublishDiagnostics.mock.calls[0]?.[0]).toEqual({ + uri: "file:///tmp/workspace/src/example.ts", + version: 2, + diagnostics: [ + { + range: { + start: { line: 0, character: 4 }, + end: { line: 0, character: 9 }, + }, + severity: 1, + code: "TS2322", + source: "tsserver", + message: "Type 'string' is not assignable to type 'number'.", + }, + ], + }); + }); + + it("ignores malformed publishDiagnostics notifications", () => { + const onPublishDiagnostics = mock((_params: LspPublishDiagnosticsParams) => undefined); + const { transport } = createClient({ onPublishDiagnostics }); + + transport.onmessage?.({ + jsonrpc: "2.0", + method: "textDocument/publishDiagnostics", + params: { + uri: "file:///tmp/workspace/src/example.ts", + diagnostics: [{ message: "missing range" }], + }, + }); + + expect(onPublishDiagnostics).toHaveBeenCalledTimes(1); + expect(onPublishDiagnostics.mock.calls[0]?.[0]).toEqual({ + uri: "file:///tmp/workspace/src/example.ts", + version: undefined, + diagnostics: [], + }); + }); +}); diff --git a/src/node/services/lsp/lspClient.ts b/src/node/services/lsp/lspClient.ts index a4b99848b9..cbcdff62c7 100644 --- a/src/node/services/lsp/lspClient.ts +++ b/src/node/services/lsp/lspClient.ts @@ -1,9 +1,6 @@ import { readFileString } from "@/node/utils/runtime/helpers"; import { shellQuote } from "@/common/utils/shell"; -import { - LSP_REQUEST_TIMEOUT_MS, - LSP_START_TIMEOUT_MS, -} from "@/constants/lsp"; +import { LSP_REQUEST_TIMEOUT_MS, LSP_START_TIMEOUT_MS } from "@/constants/lsp"; import { log } from "@/node/services/log"; import { LspStdioTransport, type LspJsonRpcMessage } from "./lspStdioTransport"; import type { @@ -12,12 +9,14 @@ import type { LspClientInstance, LspClientQueryRequest, LspClientQueryResult, + LspDiagnostic, LspDocumentSymbol, LspHover, LspLocation, LspLocationLink, LspMarkedString, LspMarkupContent, + LspPublishDiagnosticsParams, LspSymbolInformation, } from "./types"; @@ -44,7 +43,10 @@ export class LspClient implements LspClientInstance { private initialized = false; isClosed = false; - private constructor(private readonly options: CreateLspClientOptions, transport: LspStdioTransport) { + private constructor( + private readonly options: CreateLspClientOptions, + transport: LspStdioTransport + ) { this.transport = transport; this.transport.onmessage = (message) => this.handleMessage(message); this.transport.onclose = () => this.handleClose(); @@ -52,7 +54,10 @@ export class LspClient implements LspClientInstance { } static async create(options: CreateLspClientOptions): Promise { - const command = [shellQuote(options.descriptor.command), ...options.descriptor.args.map((arg) => shellQuote(arg))].join(" "); + const command = [ + shellQuote(options.descriptor.command), + ...options.descriptor.args.map((arg) => shellQuote(arg)), + ].join(" "); const execStream = await options.runtime.exec(command, { cwd: options.rootPath, // LSP servers are long-lived by design; timeout would kill healthy clients mid-session. @@ -64,7 +69,7 @@ export class LspClient implements LspClientInstance { return client; } - async ensureFile(file: LspClientFileHandle): Promise { + async ensureFile(file: LspClientFileHandle): Promise { this.ensureStarted(); const text = await readFileString(this.options.runtime, file.readablePath); @@ -83,11 +88,11 @@ export class LspClient implements LspClientInstance { text, }, }); - return; + return 1; } if (existingState.text === text) { - return; + return existingState.version; } const nextVersion = existingState.version + 1; @@ -102,6 +107,7 @@ export class LspClient implements LspClientInstance { }, contentChanges: [{ text }], }); + return nextVersion; } async query(request: LspClientQueryRequest): Promise { @@ -309,21 +315,23 @@ export class LspClient implements LspClientInstance { timeoutId, }); - void this.transport.send({ - jsonrpc: "2.0", - id: requestId, - method, - ...(params !== undefined ? { params } : {}), - }).catch((error) => { - const pending = this.pendingRequests.get(requestId); - if (!pending) { - return; - } - - clearTimeout(pending.timeoutId); - this.pendingRequests.delete(requestId); - reject(error instanceof Error ? error : new Error(String(error))); - }); + void this.transport + .send({ + jsonrpc: "2.0", + id: requestId, + method, + ...(params !== undefined ? { params } : {}), + }) + .catch((error) => { + const pending = this.pendingRequests.get(requestId); + if (!pending) { + return; + } + + clearTimeout(pending.timeoutId); + this.pendingRequests.delete(requestId); + reject(error instanceof Error ? error : new Error(String(error))); + }); }); } @@ -340,28 +348,34 @@ export class LspClient implements LspClientInstance { } private handleMessage(message: LspJsonRpcMessage): void { - if (typeof message.id !== "number") { - return; - } + if (typeof message.id === "number") { + const pending = this.pendingRequests.get(message.id); + if (!pending) { + return; + } - const pending = this.pendingRequests.get(message.id); - if (!pending) { - return; - } + clearTimeout(pending.timeoutId); + this.pendingRequests.delete(message.id); - clearTimeout(pending.timeoutId); - this.pendingRequests.delete(message.id); + if (message.error) { + pending.reject( + new Error( + `LSP ${this.options.descriptor.id} request failed: ${message.error.message}${this.buildStderrSuffix()}` + ) + ); + return; + } - if (message.error) { - pending.reject( - new Error( - `LSP ${this.options.descriptor.id} request failed: ${message.error.message}${this.buildStderrSuffix()}` - ) - ); + pending.resolve(message.result); return; } - pending.resolve(message.result); + if (message.method === "textDocument/publishDiagnostics") { + const params = parsePublishDiagnosticsParams(message.params); + if (params) { + this.options.onPublishDiagnostics?.(params); + } + } } private handleClose(): void { @@ -395,6 +409,85 @@ export class LspClient implements LspClientInstance { } } +function parsePublishDiagnosticsParams(params: unknown): LspPublishDiagnosticsParams | null { + if (typeof params !== "object" || params == null || Array.isArray(params)) { + return null; + } + + const record = params as Record; + if (typeof record.uri !== "string" || !Array.isArray(record.diagnostics)) { + return null; + } + + const diagnostics = record.diagnostics + .map((diagnostic) => parseDiagnostic(diagnostic)) + .filter((diagnostic): diagnostic is LspDiagnostic => diagnostic !== null); + + return { + uri: record.uri, + version: typeof record.version === "number" ? record.version : undefined, + diagnostics, + }; +} + +function parseDiagnostic(value: unknown): LspDiagnostic | null { + if (typeof value !== "object" || value == null || Array.isArray(value)) { + return null; + } + + const record = value as Record; + const range = parseRange(record.range); + if (!range || typeof record.message !== "string") { + return null; + } + + const severity = + typeof record.severity === "number" && record.severity >= 1 && record.severity <= 4 + ? (record.severity as 1 | 2 | 3 | 4) + : undefined; + const code = + typeof record.code === "string" || typeof record.code === "number" ? record.code : undefined; + + return { + range, + message: record.message, + severity, + code, + source: typeof record.source === "string" ? record.source : undefined, + }; +} + +function parseRange(value: unknown): LspLocation["range"] | null { + if (typeof value !== "object" || value == null || Array.isArray(value)) { + return null; + } + + const record = value as Record; + const start = parsePosition(record.start); + const end = parsePosition(record.end); + if (!start || !end) { + return null; + } + + return { start, end }; +} + +function parsePosition(value: unknown): LspLocation["range"]["start"] | null { + if (typeof value !== "object" || value == null || Array.isArray(value)) { + return null; + } + + const record = value as Record; + if (typeof record.line !== "number" || typeof record.character !== "number") { + return null; + } + + return { + line: record.line, + character: record.character, + }; +} + function normalizeLocations( value: LspLocation | LspLocationLink | Array | null | undefined ): LspLocation[] { @@ -415,15 +508,16 @@ function normalizeLocations( }); } -function normalizeHoverContents( - contents: LspHover["contents"] -): string { +function normalizeHoverContents(contents: LspHover["contents"]): string { if (typeof contents === "string") { return contents; } if (Array.isArray(contents)) { - return contents.map((entry) => normalizeHoverContents(entry)).filter((entry) => entry.length > 0).join("\n\n"); + return contents + .map((entry) => normalizeHoverContents(entry)) + .filter((entry) => entry.length > 0) + .join("\n\n"); } if (isMarkupContent(contents)) { @@ -433,8 +527,6 @@ function normalizeHoverContents( return contents.value; } -function isMarkupContent( - value: LspMarkupContent | LspMarkedString -): value is LspMarkupContent { +function isMarkupContent(value: LspMarkupContent | LspMarkedString): value is LspMarkupContent { return "kind" in value; } diff --git a/src/node/services/lsp/lspManager.test.ts b/src/node/services/lsp/lspManager.test.ts index c6008e65ef..2eefd6da66 100644 --- a/src/node/services/lsp/lspManager.test.ts +++ b/src/node/services/lsp/lspManager.test.ts @@ -3,7 +3,12 @@ import * as os from "node:os"; import * as path from "node:path"; import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test"; import { LocalRuntime } from "@/node/runtime/LocalRuntime"; -import type { CreateLspClientOptions, LspClientInstance, LspServerDescriptor } from "./types"; +import type { + CreateLspClientOptions, + LspClientInstance, + LspDiagnostic, + LspServerDescriptor, +} from "./types"; import { LspManager } from "./lspManager"; function createDeferred() { @@ -16,6 +21,31 @@ function createDeferred() { return { promise, resolve, reject }; } +function createRegistry(): readonly LspServerDescriptor[] { + return [ + { + id: "typescript", + extensions: [".ts"], + command: "fake-lsp", + args: ["--stdio"], + rootMarkers: ["package.json", ".git"], + languageIdForPath: () => "typescript", + }, + ]; +} + +function createDiagnostic(message: string): LspDiagnostic { + return { + range: { + start: { line: 0, character: 0 }, + end: { line: 0, character: 5 }, + }, + severity: 1, + source: "tsserver", + message, + }; +} + describe("LspManager", () => { let workspacePath: string; @@ -23,7 +53,14 @@ describe("LspManager", () => { workspacePath = await fs.mkdtemp(path.join(os.tmpdir(), "mux-lsp-manager-")); await fs.mkdir(path.join(workspacePath, ".git")); await fs.mkdir(path.join(workspacePath, "src"), { recursive: true }); + await fs.mkdir(path.join(workspacePath, "packages", "pkg", "src"), { recursive: true }); + await fs.writeFile(path.join(workspacePath, "package.json"), "{}\n"); await fs.writeFile(path.join(workspacePath, "src", "example.ts"), "export const value = 1;\n"); + await fs.writeFile(path.join(workspacePath, "packages", "pkg", "package.json"), "{}\n"); + await fs.writeFile( + path.join(workspacePath, "packages", "pkg", "src", "nested.ts"), + "export const nested = 1;\n" + ); }); afterEach(async () => { @@ -31,7 +68,7 @@ describe("LspManager", () => { }); test("reuses clients for the same workspace root and forwards normalized positions", async () => { - const ensureFile = mock(() => Promise.resolve(undefined)); + const ensureFile = mock(async () => 1); let lastQueryRequest: Parameters[0] | undefined; const query = mock((request: Parameters[0]) => { lastQueryRequest = request; @@ -52,19 +89,9 @@ describe("LspManager", () => { clientFactoryOptions = options; return Promise.resolve(client); }); - const registry: readonly LspServerDescriptor[] = [ - { - id: "typescript", - extensions: [".ts"], - command: "fake-lsp", - args: ["--stdio"], - rootMarkers: ["package.json", ".git"], - languageIdForPath: () => "typescript", - }, - ]; const manager = new LspManager({ - registry, + registry: createRegistry(), clientFactory, }); const runtime = new LocalRuntime(workspacePath); @@ -110,7 +137,7 @@ describe("LspManager", () => { }); test("deduplicates concurrent client creation for the same workspace root", async () => { - const ensureFile = mock(() => Promise.resolve(undefined)); + const ensureFile = mock(async () => 1); const query = mock(() => Promise.resolve({ operation: "hover" as const, @@ -130,19 +157,9 @@ describe("LspManager", () => { clientFactoryStarted.resolve(); return clientReady.promise; }); - const registry: readonly LspServerDescriptor[] = [ - { - id: "typescript", - extensions: [".ts"], - command: "fake-lsp", - args: ["--stdio"], - rootMarkers: ["package.json", ".git"], - languageIdForPath: () => "typescript", - }, - ]; const manager = new LspManager({ - registry, + registry: createRegistry(), clientFactory, }); const runtime = new LocalRuntime(workspacePath); @@ -181,4 +198,103 @@ describe("LspManager", () => { await manager.dispose(); expect(close).toHaveBeenCalledTimes(1); }); + + test("collects post-mutation diagnostics, clears empty publishes, and clears workspace cache on dispose", async () => { + const ensureFile = mock((file: Parameters[0]) => { + const version = ensureFile.mock.calls.length; + const diagnostics = version === 1 ? [createDiagnostic("first pass")] : []; + clientOptions?.onPublishDiagnostics?.({ + uri: file.uri, + version, + diagnostics, + }); + return Promise.resolve(version); + }); + const close = mock(() => Promise.resolve(undefined)); + let clientOptions: CreateLspClientOptions | undefined; + const clientFactory = mock((options: CreateLspClientOptions): Promise => { + clientOptions = options; + return Promise.resolve({ + isClosed: false, + ensureFile, + query: mock(() => Promise.resolve({ operation: "hover" as const, hover: "unused" })), + close, + }); + }); + + const manager = new LspManager({ + registry: createRegistry(), + clientFactory, + }); + const runtime = new LocalRuntime(workspacePath); + + const first = await manager.collectPostMutationDiagnostics({ + workspaceId: "ws-1", + runtime, + workspacePath, + filePaths: ["src/example.ts"], + timeoutMs: 20, + }); + expect(first).toHaveLength(1); + expect(first[0]?.diagnostics[0]?.message).toBe("first pass"); + + const second = await manager.collectPostMutationDiagnostics({ + workspaceId: "ws-1", + runtime, + workspacePath, + filePaths: ["src/example.ts"], + timeoutMs: 20, + }); + expect(second).toEqual([]); + + const workspaceDiagnostics = (manager as unknown as { + workspaceDiagnostics: Map>>; + }).workspaceDiagnostics; + expect(workspaceDiagnostics.get("ws-1")?.values().next().value?.size ?? 0).toBe(0); + + await manager.disposeWorkspace("ws-1"); + expect(workspaceDiagnostics.has("ws-1")).toBe(false); + expect(close).toHaveBeenCalledTimes(1); + }); + + test("keeps diagnostics isolated by root uri within the same workspace", async () => { + const clientFactory = mock((options: CreateLspClientOptions): Promise => { + const ensureFile = mock((file: Parameters[0]) => { + options.onPublishDiagnostics?.({ + uri: file.uri, + version: 1, + diagnostics: [createDiagnostic(`diagnostic for ${options.rootPath}`)], + }); + return Promise.resolve(1); + }); + + return Promise.resolve({ + isClosed: false, + ensureFile, + query: mock(() => Promise.resolve({ operation: "hover" as const, hover: "unused" })), + close: mock(() => Promise.resolve(undefined)), + }); + }); + + const manager = new LspManager({ + registry: createRegistry(), + clientFactory, + }); + const runtime = new LocalRuntime(workspacePath); + + const diagnostics = await manager.collectPostMutationDiagnostics({ + workspaceId: "ws-1", + runtime, + workspacePath, + filePaths: ["src/example.ts", "packages/pkg/src/nested.ts"], + timeoutMs: 20, + }); + + expect(diagnostics).toHaveLength(2); + expect(new Set(diagnostics.map((entry) => entry.rootUri)).size).toBe(2); + expect(diagnostics.map((entry) => entry.path)).toEqual([ + path.join(workspacePath, "packages", "pkg", "src", "nested.ts"), + path.join(workspacePath, "src", "example.ts"), + ]); + }); }); diff --git a/src/node/services/lsp/lspManager.ts b/src/node/services/lsp/lspManager.ts index 2554ef429e..51fcfcd789 100644 --- a/src/node/services/lsp/lspManager.ts +++ b/src/node/services/lsp/lspManager.ts @@ -5,6 +5,7 @@ import { LSP_IDLE_TIMEOUT_MS, LSP_MAX_LOCATIONS, LSP_MAX_SYMBOLS, + LSP_POST_MUTATION_DIAGNOSTICS_TIMEOUT_MS, LSP_PREVIEW_CONTEXT_LINES, } from "@/constants/lsp"; import type { Runtime } from "@/node/runtime/Runtime"; @@ -18,8 +19,10 @@ import type { LspClientInstance, LspClientQueryResult, LspDocumentSymbol, + LspFileDiagnostics, LspLocation, LspManagerQueryResult, + LspPublishDiagnosticsParams, LspQueryOperation, LspRange, LspServerDescriptor, @@ -32,6 +35,21 @@ interface WorkspaceClients { lastActivity: number; } +interface ResolvedLspClientContext { + client: LspClientInstance; + clientKey: string; + descriptor: LspServerDescriptor; + fileHandle: LspClientFileHandle; + pathMapper: LspPathMapper; + rootUri: string; +} + +interface LspDiagnosticPublishReceipt { + version?: number; + receivedAtMs: number; +} + +type LspDiagnosticWaiter = (publish: LspDiagnosticPublishReceipt) => void; type LspClientFactory = (options: CreateLspClientOptions) => Promise; export interface LspManagerOptions { @@ -53,9 +71,29 @@ export interface LspManagerQueryOptions { includeDeclaration?: boolean; } +export interface LspManagerCollectDiagnosticsOptions { + workspaceId: string; + runtime: Runtime; + workspacePath: string; + filePaths: string[]; + timeoutMs?: number; +} + export class LspManager { private readonly workspaceClients = new Map(); private readonly workspaceLeases = new Map(); + private readonly workspaceDiagnostics = new Map< + string, + Map> + >(); + private readonly workspaceDiagnosticPublishes = new Map< + string, + Map> + >(); + private readonly diagnosticWaiters = new Map< + string, + Map>> + >(); private readonly registry: readonly LspServerDescriptor[]; private readonly clientFactory: LspClientFactory; private readonly idleTimeoutMs: number; @@ -64,7 +102,8 @@ export class LspManager { constructor(options?: LspManagerOptions) { this.registry = options?.registry ?? BUILTIN_LSP_SERVERS; - this.clientFactory = options?.clientFactory ?? ((clientOptions) => LspClient.create(clientOptions)); + this.clientFactory = + options?.clientFactory ?? ((clientOptions) => LspClient.create(clientOptions)); this.idleTimeoutMs = options?.idleTimeoutMs ?? LSP_IDLE_TIMEOUT_MS; this.idleCheckIntervalMs = options?.idleCheckIntervalMs ?? LSP_IDLE_CHECK_INTERVAL_MS; this.idleInterval = setInterval(() => { @@ -93,41 +132,12 @@ export class LspManager { this.acquireLease(options.workspaceId); try { - const pathMapper = new LspPathMapper({ - runtime: options.runtime, - workspacePath: options.workspacePath, - }); - const runtimeFilePath = pathMapper.toRuntimePath(options.filePath); - if (!pathMapper.isWithinWorkspace(runtimeFilePath)) { - throw new Error(`LSP paths must stay inside the workspace (got ${options.filePath})`); - } - - const descriptor = findLspServerForFile(runtimeFilePath, this.registry); - if (!descriptor) { - const extension = path.extname(runtimeFilePath) || "(no extension)"; - throw new Error(`No built-in LSP server is configured for ${extension} files`); - } + const context = await this.resolveClientContext(options); + await context.client.ensureFile(context.fileHandle); - const rootPath = await this.resolveRootPath( - options.runtime, - runtimeFilePath, - pathMapper.getWorkspaceRuntimePath(), - descriptor.rootMarkers - ); - const rootUri = pathMapper.toUri(rootPath); - const fileHandle: LspClientFileHandle = { - runtimePath: runtimeFilePath, - readablePath: pathMapper.toReadablePath(runtimeFilePath), - uri: pathMapper.toUri(runtimeFilePath), - languageId: descriptor.languageIdForPath(runtimeFilePath), - }; - - const client = await this.getOrCreateClient(options.workspaceId, descriptor, options.runtime, rootPath, rootUri); - await client.ensureFile(fileHandle); - - const rawResult = await client.query({ + const rawResult = await context.client.query({ operation: options.operation, - file: fileHandle, + file: context.fileHandle, line: options.line != null ? Math.max(0, options.line - 1) : undefined, character: options.column != null ? Math.max(0, options.column - 1) : undefined, query: options.query, @@ -135,7 +145,79 @@ export class LspManager { }); this.markActivity(options.workspaceId); - return await this.normalizeQueryResult(pathMapper, options.runtime, fileHandle, descriptor.id, rootUri, rawResult); + return await this.normalizeQueryResult( + context.pathMapper, + options.runtime, + context.fileHandle, + context.descriptor.id, + context.rootUri, + rawResult + ); + } finally { + this.releaseLease(options.workspaceId); + } + } + + async collectPostMutationDiagnostics( + options: LspManagerCollectDiagnosticsOptions + ): Promise { + this.acquireLease(options.workspaceId); + + try { + const diagnostics: LspFileDiagnostics[] = []; + const seenFilePaths = new Set(); + + for (const filePath of options.filePaths) { + if (seenFilePaths.has(filePath)) { + continue; + } + seenFilePaths.add(filePath); + + try { + const context = await this.resolveClientContext({ + workspaceId: options.workspaceId, + runtime: options.runtime, + workspacePath: options.workspacePath, + filePath, + }); + const previousPublish = this.getLatestDiagnosticPublish( + options.workspaceId, + context.clientKey, + context.fileHandle.uri + ); + const expectedVersion = await context.client.ensureFile(context.fileHandle); + const freshPublish = await this.waitForFreshDiagnostics({ + workspaceId: options.workspaceId, + clientKey: context.clientKey, + uri: context.fileHandle.uri, + previousReceivedAtMs: previousPublish?.receivedAtMs, + expectedVersion, + timeoutMs: options.timeoutMs ?? LSP_POST_MUTATION_DIAGNOSTICS_TIMEOUT_MS, + }); + + if (!freshPublish) { + continue; + } + + const snapshot = this.getCachedDiagnostics( + options.workspaceId, + context.clientKey, + context.fileHandle.uri + ); + if (snapshot && snapshot.receivedAtMs === freshPublish.receivedAtMs) { + diagnostics.push(snapshot); + } + } catch (error) { + log.debug("Skipping post-mutation LSP diagnostics for file", { + workspaceId: options.workspaceId, + filePath, + error, + }); + } + } + + this.markActivity(options.workspaceId); + return diagnostics.sort((left, right) => left.path.localeCompare(right.path)); } finally { this.releaseLease(options.workspaceId); } @@ -144,10 +226,12 @@ export class LspManager { async disposeWorkspace(workspaceId: string): Promise { const entry = this.workspaceClients.get(workspaceId); if (!entry) { + this.clearWorkspaceDiagnostics(workspaceId); return; } this.workspaceClients.delete(workspaceId); + this.clearWorkspaceDiagnostics(workspaceId); const pendingClients = await Promise.allSettled(entry.pendingClients.values()); const clientsToClose = new Set(entry.clients.values()); for (const result of pendingClients) { @@ -180,9 +264,10 @@ export class LspManager { workspaceId: string, descriptor: LspServerDescriptor, runtime: Runtime, + pathMapper: LspPathMapper, rootPath: string, rootUri: string - ): Promise { + ): Promise<{ client: LspClientInstance; clientKey: string }> { const workspaceEntry = this.workspaceClients.get(workspaceId) ?? { clients: new Map(), pendingClients: new Map>(), @@ -190,11 +275,11 @@ export class LspManager { }; this.workspaceClients.set(workspaceId, workspaceEntry); - const clientKey = `${descriptor.id}:${rootUri}`; + const clientKey = this.getClientKey(descriptor.id, rootUri); const existingClient = workspaceEntry.clients.get(clientKey); if (existingClient && !existingClient.isClosed) { workspaceEntry.lastActivity = Date.now(); - return existingClient; + return { client: existingClient, clientKey }; } if (existingClient?.isClosed) { @@ -203,7 +288,10 @@ export class LspManager { const pendingClient = workspaceEntry.pendingClients.get(clientKey); if (pendingClient) { - return await pendingClient; + return { + client: await pendingClient, + clientKey, + }; } // Deduplicate concurrent queries for the same workspace/root so we never spawn @@ -213,6 +301,15 @@ export class LspManager { runtime, rootPath, rootUri, + onPublishDiagnostics: (params) => + this.handlePublishDiagnostics({ + workspaceId, + clientKey, + serverId: descriptor.id, + rootUri, + pathMapper, + params, + }), }) .then(async (client) => { const currentWorkspaceEntry = this.workspaceClients.get(workspaceId); @@ -243,7 +340,64 @@ export class LspManager { }); workspaceEntry.pendingClients.set(clientKey, clientPromise); - return await clientPromise; + return { + client: await clientPromise, + clientKey, + }; + } + + private async resolveClientContext(options: { + workspaceId: string; + runtime: Runtime; + workspacePath: string; + filePath: string; + }): Promise { + const pathMapper = new LspPathMapper({ + runtime: options.runtime, + workspacePath: options.workspacePath, + }); + const runtimeFilePath = pathMapper.toRuntimePath(options.filePath); + if (!pathMapper.isWithinWorkspace(runtimeFilePath)) { + throw new Error(`LSP paths must stay inside the workspace (got ${options.filePath})`); + } + + const descriptor = findLspServerForFile(runtimeFilePath, this.registry); + if (!descriptor) { + const extension = path.extname(runtimeFilePath) || "(no extension)"; + throw new Error(`No built-in LSP server is configured for ${extension} files`); + } + + const rootPath = await this.resolveRootPath( + options.runtime, + runtimeFilePath, + pathMapper.getWorkspaceRuntimePath(), + descriptor.rootMarkers + ); + const rootUri = pathMapper.toUri(rootPath); + const fileHandle: LspClientFileHandle = { + runtimePath: runtimeFilePath, + readablePath: pathMapper.toReadablePath(runtimeFilePath), + uri: pathMapper.toUri(runtimeFilePath), + languageId: descriptor.languageIdForPath(runtimeFilePath), + }; + + const clientResult = await this.getOrCreateClient( + options.workspaceId, + descriptor, + options.runtime, + pathMapper, + rootPath, + rootUri + ); + + return { + client: clientResult.client, + clientKey: clientResult.clientKey, + descriptor, + fileHandle, + pathMapper, + rootUri, + }; } private markActivity(workspaceId: string): void { @@ -303,6 +457,214 @@ export class LspManager { } } + private handlePublishDiagnostics(context: { + workspaceId: string; + clientKey: string; + serverId: string; + rootUri: string; + pathMapper: LspPathMapper; + params: LspPublishDiagnosticsParams; + }): void { + const receivedAtMs = Date.now(); + const publishReceipt: LspDiagnosticPublishReceipt = { + version: context.params.version, + receivedAtMs, + }; + const publishesForWorkspace = this.getOrCreateDiagnosticPublishes( + context.workspaceId, + context.clientKey + ); + publishesForWorkspace.set(context.params.uri, publishReceipt); + + const diagnosticsForWorkspace = this.getOrCreateWorkspaceDiagnostics( + context.workspaceId, + context.clientKey + ); + if (context.params.diagnostics.length === 0) { + diagnosticsForWorkspace.delete(context.params.uri); + } else { + const runtimePath = context.pathMapper.fromUri(context.params.uri); + diagnosticsForWorkspace.set(context.params.uri, { + uri: context.params.uri, + path: context.pathMapper.toOutputPath(runtimePath), + serverId: context.serverId, + rootUri: context.rootUri, + version: context.params.version, + diagnostics: context.params.diagnostics, + receivedAtMs, + }); + } + + this.notifyDiagnosticWaiters( + context.workspaceId, + context.clientKey, + context.params.uri, + publishReceipt + ); + } + + private getOrCreateWorkspaceDiagnostics( + workspaceId: string, + clientKey: string + ): Map { + const workspaceDiagnostics = this.workspaceDiagnostics.get(workspaceId) ?? new Map(); + this.workspaceDiagnostics.set(workspaceId, workspaceDiagnostics); + const diagnosticsForClient = + workspaceDiagnostics.get(clientKey) ?? new Map(); + workspaceDiagnostics.set(clientKey, diagnosticsForClient); + return diagnosticsForClient; + } + + private getOrCreateDiagnosticPublishes( + workspaceId: string, + clientKey: string + ): Map { + const workspacePublishes = this.workspaceDiagnosticPublishes.get(workspaceId) ?? new Map(); + this.workspaceDiagnosticPublishes.set(workspaceId, workspacePublishes); + const publishesForClient = + workspacePublishes.get(clientKey) ?? new Map(); + workspacePublishes.set(clientKey, publishesForClient); + return publishesForClient; + } + + private getCachedDiagnostics( + workspaceId: string, + clientKey: string, + uri: string + ): LspFileDiagnostics | undefined { + return this.workspaceDiagnostics.get(workspaceId)?.get(clientKey)?.get(uri); + } + + private getLatestDiagnosticPublish( + workspaceId: string, + clientKey: string, + uri: string + ): LspDiagnosticPublishReceipt | undefined { + return this.workspaceDiagnosticPublishes.get(workspaceId)?.get(clientKey)?.get(uri); + } + + private async waitForFreshDiagnostics(params: { + workspaceId: string; + clientKey: string; + uri: string; + previousReceivedAtMs?: number; + expectedVersion: number; + timeoutMs: number; + }): Promise { + const existingPublish = this.getLatestDiagnosticPublish( + params.workspaceId, + params.clientKey, + params.uri + ); + if ( + isFreshDiagnosticPublish(existingPublish, params.previousReceivedAtMs, params.expectedVersion) + ) { + return existingPublish; + } + + return await new Promise((resolve) => { + let settled = false; + const workspaceWaiters = this.getOrCreateDiagnosticWaiters( + params.workspaceId, + params.clientKey, + params.uri + ); + + const finish = (publish?: LspDiagnosticPublishReceipt) => { + if (settled) { + return; + } + settled = true; + clearTimeout(timeoutId); + workspaceWaiters.delete(onPublish); + if (workspaceWaiters.size === 0) { + this.deleteDiagnosticWaiters(params.workspaceId, params.clientKey, params.uri); + } + resolve(publish); + }; + + const onPublish: LspDiagnosticWaiter = (publish) => { + if ( + !isFreshDiagnosticPublish(publish, params.previousReceivedAtMs, params.expectedVersion) + ) { + return; + } + finish(publish); + }; + + const timeoutId = setTimeout(() => finish(undefined), params.timeoutMs); + workspaceWaiters.add(onPublish); + + const latestPublish = this.getLatestDiagnosticPublish( + params.workspaceId, + params.clientKey, + params.uri + ); + if ( + isFreshDiagnosticPublish(latestPublish, params.previousReceivedAtMs, params.expectedVersion) + ) { + finish(latestPublish); + } + }); + } + + private getOrCreateDiagnosticWaiters( + workspaceId: string, + clientKey: string, + uri: string + ): Set { + const workspaceWaiters = this.diagnosticWaiters.get(workspaceId) ?? new Map(); + this.diagnosticWaiters.set(workspaceId, workspaceWaiters); + const clientWaiters = + workspaceWaiters.get(clientKey) ?? new Map>(); + workspaceWaiters.set(clientKey, clientWaiters); + const uriWaiters = clientWaiters.get(uri) ?? new Set(); + clientWaiters.set(uri, uriWaiters); + return uriWaiters; + } + + private deleteDiagnosticWaiters(workspaceId: string, clientKey: string, uri: string): void { + const workspaceWaiters = this.diagnosticWaiters.get(workspaceId); + const clientWaiters = workspaceWaiters?.get(clientKey); + if (!clientWaiters) { + return; + } + + clientWaiters.delete(uri); + if (clientWaiters.size === 0) { + workspaceWaiters?.delete(clientKey); + } + if (workspaceWaiters && workspaceWaiters.size === 0) { + this.diagnosticWaiters.delete(workspaceId); + } + } + + private notifyDiagnosticWaiters( + workspaceId: string, + clientKey: string, + uri: string, + publish: LspDiagnosticPublishReceipt + ): void { + const waiters = this.diagnosticWaiters.get(workspaceId)?.get(clientKey)?.get(uri); + if (!waiters || waiters.size === 0) { + return; + } + + for (const waiter of [...waiters]) { + waiter(publish); + } + } + + private clearWorkspaceDiagnostics(workspaceId: string): void { + this.workspaceDiagnostics.delete(workspaceId); + this.workspaceDiagnosticPublishes.delete(workspaceId); + this.diagnosticWaiters.delete(workspaceId); + } + + private getClientKey(serverId: string, rootUri: string): string { + return `${serverId}:${rootUri}`; + } + private async normalizeQueryResult( pathMapper: LspPathMapper, runtime: Runtime, @@ -327,9 +689,9 @@ export class LspManager { ? `Results truncated to the first ${LSP_MAX_LOCATIONS} locations` : undefined; const locations = await Promise.all( - (rawResult.locations ?? []).slice(0, LSP_MAX_LOCATIONS).map(async (location) => - this.buildLocationResult(pathMapper, runtime, location) - ) + (rawResult.locations ?? []) + .slice(0, LSP_MAX_LOCATIONS) + .map(async (location) => this.buildLocationResult(pathMapper, runtime, location)) ); return { operation: rawResult.operation, @@ -347,10 +709,18 @@ export class LspManager { : undefined; const symbols = await Promise.all( flattenedSymbols.slice(0, LSP_MAX_SYMBOLS).map(async (symbol) => - this.buildSymbolResult(pathMapper, runtime, symbol.uri, symbol.range, symbol.name, symbol.kind, { - detail: symbol.detail, - containerName: symbol.containerName, - }) + this.buildSymbolResult( + pathMapper, + runtime, + symbol.uri, + symbol.range, + symbol.name, + symbol.kind, + { + detail: symbol.detail, + containerName: symbol.containerName, + } + ) ) ); return { @@ -369,10 +739,18 @@ export class LspManager { : undefined; const symbols = await Promise.all( workspaceSymbols.slice(0, LSP_MAX_SYMBOLS).map(async (symbol) => - this.buildSymbolResult(pathMapper, runtime, symbol.uri, symbol.range, symbol.name, symbol.kind, { - detail: symbol.detail, - containerName: symbol.containerName, - }) + this.buildSymbolResult( + pathMapper, + runtime, + symbol.uri, + symbol.range, + symbol.name, + symbol.kind, + { + detail: symbol.detail, + containerName: symbol.containerName, + } + ) ) ); return { @@ -386,7 +764,11 @@ export class LspManager { } } - private async buildLocationResult(pathMapper: LspPathMapper, runtime: Runtime, location: LspLocation) { + private async buildLocationResult( + pathMapper: LspPathMapper, + runtime: Runtime, + location: LspLocation + ) { const runtimePath = pathMapper.fromUri(location.uri); const outputPath = pathMapper.toOutputPath(runtimePath); return { @@ -429,6 +811,22 @@ function selectPathModule(filePath: string): PathModule { return path.posix; } +function isFreshDiagnosticPublish( + publish: LspDiagnosticPublishReceipt | undefined, + previousReceivedAtMs: number | undefined, + expectedVersion: number +): publish is LspDiagnosticPublishReceipt { + if (!publish) { + return false; + } + + if (publish.version != null) { + return publish.version >= expectedVersion; + } + + return publish.receivedAtMs > (previousReceivedAtMs ?? 0); +} + async function pathExists(runtime: Runtime, candidatePath: string): Promise { try { await runtime.stat(candidatePath); @@ -520,9 +918,7 @@ function flattenDocumentSymbols( return flattened; } -function flattenWorkspaceSymbols( - symbols: Array -): Array<{ +function flattenWorkspaceSymbols(symbols: Array): Array<{ name: string; kind: number; detail?: string; @@ -532,7 +928,11 @@ function flattenWorkspaceSymbols( }> { return symbols .map((symbol) => { - if (!isWorkspaceSymbolInformation(symbol) || !symbol.location || !("range" in symbol.location)) { + if ( + !isWorkspaceSymbolInformation(symbol) || + !symbol.location || + !("range" in symbol.location) + ) { return null; } diff --git a/src/node/services/lsp/types.ts b/src/node/services/lsp/types.ts index e68862c1ad..bdc51b6430 100644 --- a/src/node/services/lsp/types.ts +++ b/src/node/services/lsp/types.ts @@ -42,6 +42,30 @@ export interface LspHover { contents: string | LspMarkupContent | LspMarkedString | Array; } +export interface LspDiagnostic { + range: LspRange; + severity?: 1 | 2 | 3 | 4; + code?: string | number; + source?: string; + message: string; +} + +export interface LspPublishDiagnosticsParams { + uri: string; + version?: number; + diagnostics: LspDiagnostic[]; +} + +export interface LspFileDiagnostics { + uri: string; + path: string; + serverId: string; + rootUri: string; + version?: number; + diagnostics: LspDiagnostic[]; + receivedAtMs: number; +} + export interface LspDocumentSymbol { name: string; kind: number; @@ -121,7 +145,7 @@ export interface LspClientQueryResult { export interface LspClientInstance { readonly isClosed: boolean; - ensureFile(file: LspClientFileHandle): Promise; + ensureFile(file: LspClientFileHandle): Promise; query(request: LspClientQueryRequest): Promise; close(): Promise; } @@ -131,4 +155,5 @@ export interface CreateLspClientOptions { runtime: Runtime; rootPath: string; rootUri: string; + onPublishDiagnostics?: (params: LspPublishDiagnosticsParams) => void; } diff --git a/src/node/services/tools/file_edit_insert.test.ts b/src/node/services/tools/file_edit_insert.test.ts index 32adbd8a33..65b467e05f 100644 --- a/src/node/services/tools/file_edit_insert.test.ts +++ b/src/node/services/tools/file_edit_insert.test.ts @@ -142,6 +142,33 @@ describe("file_edit_insert tool", () => { expect(await fs.readFile(newFile, "utf-8")).toBe("Hello world!\n"); }); + it("appends post-mutation warnings when creating a new file", async () => { + const newFile = path.join(testDir, "diagnostics.txt"); + const mutationCalls: Array<{ filePaths: string[] }> = []; + const tool = createFileEditInsertTool({ + ...getTestDeps(), + cwd: testDir, + runtime: createRuntime({ type: "local", srcBaseDir: testDir }), + runtimeTempDir: testDir, + onFilesMutated: async (params) => { + mutationCalls.push(params); + return "Post-edit LSP diagnostics:\n- diagnostics.txt:1:1 warning TS1000: created"; + }, + }); + const args: FileEditInsertToolArgs = { + path: path.relative(testDir, newFile), + content: "Hello world!\n", + }; + + const result = (await tool.execute!(args, mockToolCallOptions)) as FileEditInsertToolResult; + + expect(result.success).toBe(true); + expect(mutationCalls).toEqual([{ filePaths: [newFile] }]); + if (result.success) { + expect(result.warning).toContain("Post-edit LSP diagnostics:"); + } + }); + it("fails when no guards are provided", async () => { const tool = createTestTool(testDir); const args: FileEditInsertToolArgs = { diff --git a/src/node/services/tools/file_edit_insert.ts b/src/node/services/tools/file_edit_insert.ts index 31d69f1a09..e118bbfa3e 100644 --- a/src/node/services/tools/file_edit_insert.ts +++ b/src/node/services/tools/file_edit_insert.ts @@ -8,12 +8,13 @@ import { import type { ToolConfiguration, ToolFactory } from "@/common/utils/tools/tools"; import { TOOL_DEFINITIONS } from "@/common/utils/tools/toolDefinitions"; import { generateDiff, resolvePathWithinCwd, validatePlanModeAccess } from "./fileCommon"; -import { executeFileEditOperation } from "./file_edit_operation"; +import { executeFileEditOperation, mergeFileMutationWarnings } from "./file_edit_operation"; import { convertNewlines, detectFileEol } from "./eol"; import { fileExists } from "@/node/utils/runtime/fileExists"; import { writeFileString } from "@/node/utils/runtime/helpers"; import { RuntimeError } from "@/node/runtime/Runtime"; import { getErrorMessage } from "@/common/utils/errors"; +import { log } from "@/node/services/log"; const READ_AND_RETRY_NOTE = `${EDIT_FAILED_NOTE_PREFIX} ${NOTE_READ_FILE_RETRY}`; @@ -96,6 +97,18 @@ export const createFileEditInsertTool: ToolFactory = (config: ToolConfiguration) } } + let postMutationWarning: string | undefined; + if (config.onFilesMutated) { + try { + postMutationWarning = await config.onFilesMutated({ filePaths: [resolvedPath] }); + } catch (error) { + log.debug("Failed to collect post-mutation file warnings", { + resolvedPath, + error, + }); + } + } + const diff = generateDiff(resolvedPath, "", content); return { success: true, @@ -105,7 +118,9 @@ export const createFileEditInsertTool: ToolFactory = (config: ToolConfiguration) diff, }, }, - ...(pathWarning && { warning: pathWarning }), + ...(mergeFileMutationWarnings(pathWarning, postMutationWarning) + ? { warning: mergeFileMutationWarnings(pathWarning, postMutationWarning) } + : {}), }; } diff --git a/src/node/services/tools/file_edit_operation.test.ts b/src/node/services/tools/file_edit_operation.test.ts index 519c1ac7e0..90e7c490ee 100644 --- a/src/node/services/tools/file_edit_operation.test.ts +++ b/src/node/services/tools/file_edit_operation.test.ts @@ -73,6 +73,60 @@ describe("executeFileEditOperation", () => { }); }); +describe("executeFileEditOperation post-mutation warnings", () => { + test("should append post-mutation warnings after a successful edit", async () => { + using tempDir = new TestTempDir("post-mutation-warning"); + + const testFile = path.join(tempDir.path, "main.ts"); + await fs.writeFile(testFile, "const value = 1;\n"); + const onFilesMutated = jest + .fn<(params: { filePaths: string[] }) => Promise>() + .mockResolvedValue("Post-edit LSP diagnostics:\n- main.ts:1:1 error TS1000: broken"); + + const result = await executeFileEditOperation({ + config: { + cwd: tempDir.path, + runtime: new LocalRuntime(tempDir.path), + runtimeTempDir: tempDir.path, + onFilesMutated, + }, + filePath: testFile, + operation: () => ({ success: true, newContent: "const value = 2;\n", metadata: {} }), + }); + + expect(result.success).toBe(true); + expect(onFilesMutated).toHaveBeenCalledTimes(1); + expect(onFilesMutated).toHaveBeenCalledWith({ filePaths: [testFile] }); + if (result.success) { + expect(result.warning).toContain("Post-edit LSP diagnostics:"); + } + }); + + test("should not call onFilesMutated when the edit fails before writing", async () => { + using tempDir = new TestTempDir("post-mutation-warning-failure"); + + const testFile = path.join(tempDir.path, "main.ts"); + await fs.writeFile(testFile, "const value = 1;\n"); + const onFilesMutated = jest + .fn<(params: { filePaths: string[] }) => Promise>() + .mockResolvedValue("unused"); + + const result = await executeFileEditOperation({ + config: { + cwd: tempDir.path, + runtime: new LocalRuntime(tempDir.path), + runtimeTempDir: tempDir.path, + onFilesMutated, + }, + filePath: testFile, + operation: () => ({ success: false, error: "no-op" }), + }); + + expect(result.success).toBe(false); + expect(onFilesMutated).not.toHaveBeenCalled(); + }); +}); + describe("executeFileEditOperation outside-cwd access", () => { test("should allow traversal outside cwd", async () => { using tempDir = new TestTempDir("outside-cwd-traversal"); diff --git a/src/node/services/tools/file_edit_operation.ts b/src/node/services/tools/file_edit_operation.ts index a99c76e33e..14fc2e265f 100644 --- a/src/node/services/tools/file_edit_operation.ts +++ b/src/node/services/tools/file_edit_operation.ts @@ -14,6 +14,7 @@ import { RuntimeError } from "@/node/runtime/Runtime"; import type { Runtime } from "@/node/runtime/Runtime"; import { readFileString, writeFileString } from "@/node/utils/runtime/helpers"; import { getErrorMessage } from "@/common/utils/errors"; +import { log } from "@/node/services/log"; import { MutexMap } from "@/node/utils/concurrency/mutexMap"; type FileEditOperationResult = @@ -100,6 +101,13 @@ async function waitForFileEditOrAbort( }); } +export function mergeFileMutationWarnings( + ...warnings: Array +): string | undefined { + const mergedWarnings = warnings.filter((warning): warning is string => Boolean(warning?.trim())); + return mergedWarnings.length > 0 ? mergedWarnings.join("\n") : undefined; +} + /** * Shared execution pipeline for file edit tools. * Handles validation, file IO, diff generation, and common error handling. @@ -233,6 +241,18 @@ export async function executeFileEditOperation({ } } + let postMutationWarning: string | undefined; + if (config.onFilesMutated) { + try { + postMutationWarning = await config.onFilesMutated({ filePaths: [resolvedPath] }); + } catch (error) { + log.debug("Failed to collect post-mutation file warnings", { + resolvedPath, + error, + }); + } + } + const diff = generateDiff(resolvedPath, originalContent, operationResult.newContent); return { @@ -244,7 +264,9 @@ export async function executeFileEditOperation({ }, }, ...operationResult.metadata, - ...(pathWarning && { warning: pathWarning }), + ...(mergeFileMutationWarnings(pathWarning, postMutationWarning) + ? { warning: mergeFileMutationWarnings(pathWarning, postMutationWarning) } + : {}), }; }); diff --git a/src/node/services/tools/task_apply_git_patch.test.ts b/src/node/services/tools/task_apply_git_patch.test.ts index 50857fb10f..f517d9baaf 100644 --- a/src/node/services/tools/task_apply_git_patch.test.ts +++ b/src/node/services/tools/task_apply_git_patch.test.ts @@ -149,6 +149,60 @@ async function writeWorkspaceConfig(params: { ); } +async function setupSingleProjectPatchFixture(rootDir: string, name: string) { + const childRepo = path.join(rootDir, `${name}-child`); + const targetRepo = path.join(rootDir, `${name}-target`); + for (const repo of [childRepo, targetRepo]) { + await fsPromises.mkdir(repo, { recursive: true }); + initGitRepo(repo); + } + + await commitFile(childRepo, "README.md", "hello", "base"); + await commitFile(targetRepo, "README.md", "hello", "base"); + const baseSha = execSync("git rev-parse HEAD", { cwd: childRepo, encoding: "utf-8" }).trim(); + + await commitFile(childRepo, "README.md", "hello\nupdated", "child change"); + const headSha = execSync("git rev-parse HEAD", { cwd: childRepo, encoding: "utf-8" }).trim(); + + const muxRoot = path.join(rootDir, `${name}-mux`); + const workspaceId = `${name}-workspace`; + const sessionDir = path.join(muxRoot, "sessions", workspaceId); + await fsPromises.mkdir(sessionDir, { recursive: true }); + await writeWorkspaceConfig({ + muxRoot, + workspaceId, + workspaceName: name, + primaryProjectPath: targetRepo, + projects: [{ projectPath: targetRepo, projectName: name }], + }); + + const childTaskId = `${name}-task`; + await writePatchArtifact({ + sessionDir, + workspaceId, + childTaskId, + projectArtifacts: [ + await buildReadyProjectArtifact({ + sessionDir, + childTaskId, + storageKey: name, + projectPath: targetRepo, + projectName: name, + childRepo, + baseSha, + headSha, + }), + ], + }); + + return { + childTaskId, + sessionDir, + targetRepo, + workspaceId, + }; +} + describe("task_apply_git_patch tool", () => { let rootDir: string; @@ -711,4 +765,59 @@ describe("task_apply_git_patch tool", () => { expect(artifact?.projectArtifacts[0]?.appliedAtMs).toBe(appliedAtMs); expect(await readSubagentGitPatchArtifact(currentSessionDir, childTaskId)).toBeNull(); }, 20_000); + + it("appends post-apply diagnostics notes for real applies", async () => { + const fixture = await setupSingleProjectPatchFixture(rootDir, "diagnostics"); + const mutationCalls: Array<{ filePaths: string[] }> = []; + const tool = createTaskApplyGitPatchTool({ + ...getTestDeps(), + workspaceId: fixture.workspaceId, + cwd: fixture.targetRepo, + runtime: createRuntime({ type: "local", srcBaseDir: "/tmp" }), + runtimeTempDir: "/tmp", + workspaceSessionDir: fixture.sessionDir, + onFilesMutated: async (params) => { + mutationCalls.push(params); + return "Post-edit LSP diagnostics:\n- README.md:2:1 error TS1000: patch issue"; + }, + }); + + const result = (await tool.execute!({ task_id: fixture.childTaskId }, mockToolCallOptions)) as { + success: boolean; + note?: string; + projectResults: Array<{ note?: string }>; + }; + + expect(result.success).toBe(true); + expect(mutationCalls).toEqual([{ filePaths: [path.join(fixture.targetRepo, "README.md")] }]); + expect(result.note).toContain("Post-edit LSP diagnostics:"); + expect(result.projectResults[0]?.note).toContain("Post-edit LSP diagnostics:"); + }, 20_000); + + it("does not request post-apply diagnostics for dry runs", async () => { + const fixture = await setupSingleProjectPatchFixture(rootDir, "dry-run"); + const onFilesMutated = async (_params: { filePaths: string[] }) => { + throw new Error("should not be called"); + }; + const tool = createTaskApplyGitPatchTool({ + ...getTestDeps(), + workspaceId: fixture.workspaceId, + cwd: fixture.targetRepo, + runtime: createRuntime({ type: "local", srcBaseDir: "/tmp" }), + runtimeTempDir: "/tmp", + workspaceSessionDir: fixture.sessionDir, + onFilesMutated, + }); + + const result = (await tool.execute!( + { task_id: fixture.childTaskId, dry_run: true }, + mockToolCallOptions + )) as { + success: boolean; + note?: string; + }; + + expect(result.success).toBe(true); + expect(result.note).not.toContain("Post-edit LSP diagnostics:"); + }, 20_000); }); diff --git a/src/node/services/tools/task_apply_git_patch.ts b/src/node/services/tools/task_apply_git_patch.ts index 1d5bd760fc..a9794d36cf 100644 --- a/src/node/services/tools/task_apply_git_patch.ts +++ b/src/node/services/tools/task_apply_git_patch.ts @@ -171,6 +171,61 @@ async function getAppliedCommits(params: { return []; } +function selectPathModule(filePath: string): typeof path.posix { + if (/^[A-Za-z]:[\\/]/.test(filePath) || filePath.includes("\\")) { + return path.win32; + } + return path.posix; +} + +function joinRepoRelativePath(repoCwd: string, relativePath: string): string { + const pathModule = selectPathModule(repoCwd); + const normalizedSegments = relativePath.split("/").filter((segment) => segment.length > 0); + return pathModule.join(repoCwd, ...normalizedSegments); +} + +async function getChangedFilesForAppliedRange(params: { + runtime: ToolConfiguration["runtime"]; + cwd: string; + beforeHeadSha: string | undefined; +}): Promise { + const diffCommand = params.beforeHeadSha + ? `git diff --name-only ${params.beforeHeadSha}..HEAD --` + : "git diff-tree --no-commit-id --name-only -r HEAD"; + + try { + const result = await execBuffered(params.runtime, diffCommand, { + cwd: params.cwd, + timeout: 30, + }); + if (result.exitCode !== 0) { + log.debug("task_apply_git_patch: git diff --name-only failed", { + cwd: params.cwd, + exitCode: result.exitCode, + stderr: result.stderr.trim(), + stdout: result.stdout.trim(), + }); + return []; + } + + return Array.from( + new Set( + result.stdout + .split("\n") + .map((line) => line.replace(/\r$/, "").trim()) + .filter((line) => line.length > 0) + .map((line) => joinRepoRelativePath(params.cwd, line)) + ) + ); + } catch (error) { + log.debug("task_apply_git_patch: git diff --name-only threw", { + cwd: params.cwd, + error, + }); + return []; + } +} + const MAX_PARENT_WORKSPACE_DEPTH = 32; function inferMuxRootFromWorkspaceSessionDir(workspaceSessionDir: string): string | undefined { @@ -569,6 +624,7 @@ async function applyProjectPatch(params: { threeWay: boolean; force: boolean; isReplay: boolean; + onFilesMutated?: (params: { filePaths: string[] }) => Promise; abortSignal?: AbortSignal; }): Promise<{ success: boolean; projectResult: TaskApplyGitPatchProjectResult }> { const patchResolution = await resolvePatchPath({ @@ -865,6 +921,27 @@ async function applyProjectPatch(params: { includeSha: true, }); + let postMutationNote: string | undefined; + if (params.onFilesMutated) { + const changedFiles = await getChangedFilesForAppliedRange({ + runtime: params.runtime, + cwd: params.repoCwd, + beforeHeadSha, + }); + if (changedFiles.length > 0) { + try { + postMutationNote = await params.onFilesMutated({ filePaths: changedFiles }); + } catch (error) { + log.debug("task_apply_git_patch: failed to collect post-apply warnings", { + taskId: params.taskId, + workspaceId: params.workspaceId, + cwd: params.repoCwd, + error, + }); + } + } + } + if (!params.isReplay) { await markSubagentGitPatchArtifactApplied({ workspaceId: params.artifactWorkspaceId, @@ -883,7 +960,7 @@ async function applyProjectPatch(params: { status: "applied", appliedCommits, headCommitSha, - note: patchResolution.note, + note: mergeNotes(patchResolution.note, postMutationNote), }, }; } @@ -1071,6 +1148,7 @@ export const createTaskApplyGitPatchTool: ToolFactory = (config: ToolConfigurati threeWay, force, isReplay, + onFilesMutated: config.onFilesMutated, abortSignal, }); projectResults.push(applyResult.projectResult); diff --git a/src/node/services/tools/testHelpers.ts b/src/node/services/tools/testHelpers.ts index 643321ddc9..adbefcbda9 100644 --- a/src/node/services/tools/testHelpers.ts +++ b/src/node/services/tools/testHelpers.ts @@ -59,6 +59,7 @@ export function createTestToolConfig( sessionsDir?: string; runtime?: Runtime; muxScope?: MuxToolScope; + onFilesMutated?: (params: { filePaths: string[] }) => Promise; } ): ToolConfiguration { return { @@ -67,6 +68,7 @@ export function createTestToolConfig( runtime: options?.runtime ?? new LocalRuntime(tempDir), runtimeTempDir: tempDir, workspaceId: options?.workspaceId ?? "test-workspace", + onFilesMutated: options?.onFilesMutated, muxScope: options?.muxScope ?? { type: "global", muxHome: tempDir, From c11119b09a59aa6ad1456655094131aec23427ea Mon Sep 17 00:00:00 2001 From: terion Date: Mon, 13 Apr 2026 13:41:12 +0300 Subject: [PATCH 04/46] Fix post-edit diagnostics freshness and patch file tracking --- src/node/services/lsp/lspClient.test.ts | 46 +++++- src/node/services/lsp/lspClient.ts | 1 + src/node/services/lsp/lspManager.test.ts | 142 +++++++++++++++++- src/node/services/lsp/lspManager.ts | 134 +++++++++++------ src/node/services/lsp/types.ts | 7 +- .../tools/task_apply_git_patch.test.ts | 103 +++++++++++-- .../services/tools/task_apply_git_patch.ts | 73 ++++----- 7 files changed, 404 insertions(+), 102 deletions(-) diff --git a/src/node/services/lsp/lspClient.test.ts b/src/node/services/lsp/lspClient.test.ts index 971c5d43ac..c1f360a33d 100644 --- a/src/node/services/lsp/lspClient.test.ts +++ b/src/node/services/lsp/lspClient.test.ts @@ -1,6 +1,10 @@ import { describe, expect, it, mock } from "bun:test"; import { LspClient } from "./lspClient"; -import type { CreateLspClientOptions, LspPublishDiagnosticsParams, LspServerDescriptor } from "./types"; +import type { + CreateLspClientOptions, + LspPublishDiagnosticsParams, + LspServerDescriptor, +} from "./types"; function createDescriptor(): LspServerDescriptor { return { @@ -15,7 +19,14 @@ function createDescriptor(): LspServerDescriptor { function createTransport() { return { - onmessage: undefined as ((message: { jsonrpc: "2.0"; id?: number | string; method?: string; params?: unknown }) => void) | undefined, + onmessage: undefined as + | ((message: { + jsonrpc: "2.0"; + id?: number | string; + method?: string; + params?: unknown; + }) => void) + | undefined, onclose: undefined as (() => void) | undefined, onerror: undefined as ((error: Error) => void) | undefined, start: mock(() => undefined), @@ -29,7 +40,10 @@ function createTransport() { function createClient(options?: Partial) { const transport = createTransport(); const client = new (LspClient as unknown as { - new (clientOptions: CreateLspClientOptions, transport: ReturnType): LspClient; + new ( + clientOptions: CreateLspClientOptions, + transport: ReturnType + ): LspClient; })( { descriptor: createDescriptor(), @@ -89,10 +103,33 @@ describe("LspClient publishDiagnostics handling", () => { message: "Type 'string' is not assignable to type 'number'.", }, ], + rawDiagnosticCount: 1, }); }); - it("ignores malformed publishDiagnostics notifications", () => { + it("tags explicit clears so the manager can distinguish them from malformed publishes", () => { + const onPublishDiagnostics = mock((_params: LspPublishDiagnosticsParams) => undefined); + const { transport } = createClient({ onPublishDiagnostics }); + + transport.onmessage?.({ + jsonrpc: "2.0", + method: "textDocument/publishDiagnostics", + params: { + uri: "file:///tmp/workspace/src/example.ts", + diagnostics: [], + }, + }); + + expect(onPublishDiagnostics).toHaveBeenCalledTimes(1); + expect(onPublishDiagnostics.mock.calls[0]?.[0]).toEqual({ + uri: "file:///tmp/workspace/src/example.ts", + version: undefined, + diagnostics: [], + rawDiagnosticCount: 0, + }); + }); + + it("preserves malformed inner diagnostics so the manager can ignore them", () => { const onPublishDiagnostics = mock((_params: LspPublishDiagnosticsParams) => undefined); const { transport } = createClient({ onPublishDiagnostics }); @@ -110,6 +147,7 @@ describe("LspClient publishDiagnostics handling", () => { uri: "file:///tmp/workspace/src/example.ts", version: undefined, diagnostics: [], + rawDiagnosticCount: 1, }); }); }); diff --git a/src/node/services/lsp/lspClient.ts b/src/node/services/lsp/lspClient.ts index cbcdff62c7..8b8af40dc8 100644 --- a/src/node/services/lsp/lspClient.ts +++ b/src/node/services/lsp/lspClient.ts @@ -427,6 +427,7 @@ function parsePublishDiagnosticsParams(params: unknown): LspPublishDiagnosticsPa uri: record.uri, version: typeof record.version === "number" ? record.version : undefined, diagnostics, + rawDiagnosticCount: record.diagnostics.length, }; } diff --git a/src/node/services/lsp/lspManager.test.ts b/src/node/services/lsp/lspManager.test.ts index 2eefd6da66..e1ca08a448 100644 --- a/src/node/services/lsp/lspManager.test.ts +++ b/src/node/services/lsp/lspManager.test.ts @@ -207,6 +207,7 @@ describe("LspManager", () => { uri: file.uri, version, diagnostics, + rawDiagnosticCount: diagnostics.length, }); return Promise.resolve(version); }); @@ -247,9 +248,11 @@ describe("LspManager", () => { }); expect(second).toEqual([]); - const workspaceDiagnostics = (manager as unknown as { - workspaceDiagnostics: Map>>; - }).workspaceDiagnostics; + const workspaceDiagnostics = ( + manager as unknown as { + workspaceDiagnostics: Map>>; + } + ).workspaceDiagnostics; expect(workspaceDiagnostics.get("ws-1")?.values().next().value?.size ?? 0).toBe(0); await manager.disposeWorkspace("ws-1"); @@ -257,6 +260,138 @@ describe("LspManager", () => { expect(close).toHaveBeenCalledTimes(1); }); + test("ignores malformed publishes so they do not clear cached diagnostics or satisfy waits", async () => { + const ensureFile = mock((file: Parameters[0]) => { + const version = ensureFile.mock.calls.length; + clientOptions?.onPublishDiagnostics?.({ + uri: file.uri, + version, + diagnostics: version === 1 ? [createDiagnostic("first pass")] : [], + rawDiagnosticCount: 1, + }); + return Promise.resolve(version); + }); + const close = mock(() => Promise.resolve(undefined)); + let clientOptions: CreateLspClientOptions | undefined; + const clientFactory = mock((options: CreateLspClientOptions): Promise => { + clientOptions = options; + return Promise.resolve({ + isClosed: false, + ensureFile, + query: mock(() => Promise.resolve({ operation: "hover" as const, hover: "unused" })), + close, + }); + }); + + const manager = new LspManager({ + registry: createRegistry(), + clientFactory, + }); + const runtime = new LocalRuntime(workspacePath); + + const first = await manager.collectPostMutationDiagnostics({ + workspaceId: "ws-1", + runtime, + workspacePath, + filePaths: ["src/example.ts"], + timeoutMs: 20, + }); + expect(first).toHaveLength(1); + expect(first[0]?.diagnostics[0]?.message).toBe("first pass"); + + const second = await manager.collectPostMutationDiagnostics({ + workspaceId: "ws-1", + runtime, + workspacePath, + filePaths: ["src/example.ts"], + timeoutMs: 20, + }); + expect(second).toEqual([]); + + const workspaceDiagnostics = ( + manager as unknown as { + workspaceDiagnostics: Map< + string, + Map> + >; + } + ).workspaceDiagnostics; + const cachedDiagnostics = [ + ...(workspaceDiagnostics.get("ws-1")?.values().next().value?.values() ?? []), + ]; + expect(cachedDiagnostics).toHaveLength(1); + expect(cachedDiagnostics[0]?.diagnostics[0]?.message).toBe("first pass"); + + await manager.disposeWorkspace("ws-1"); + expect(close).toHaveBeenCalledTimes(1); + }); + + test("collects post-mutation diagnostics across roots without serial waits", async () => { + const secondEnsureStarted = createDeferred(); + const publishByUri = new Map void>(); + let secondEnsureResolved = false; + const clientFactory = mock((options: CreateLspClientOptions): Promise => { + const ensureFile = mock(async (file: Parameters[0]) => { + publishByUri.set(file.uri, () => { + options.onPublishDiagnostics?.({ + uri: file.uri, + version: 1, + diagnostics: [createDiagnostic(`diagnostic for ${file.uri}`)], + rawDiagnosticCount: 1, + }); + }); + if (publishByUri.size === 2 && !secondEnsureResolved) { + secondEnsureResolved = true; + secondEnsureStarted.resolve(); + } + return 1; + }); + + return Promise.resolve({ + isClosed: false, + ensureFile, + query: mock(() => Promise.resolve({ operation: "hover" as const, hover: "unused" })), + close: mock(() => Promise.resolve(undefined)), + }); + }); + + const manager = new LspManager({ + registry: createRegistry(), + clientFactory, + }); + const runtime = new LocalRuntime(workspacePath); + + const diagnosticsPromise = manager.collectPostMutationDiagnostics({ + workspaceId: "ws-1", + runtime, + workspacePath, + filePaths: ["src/example.ts", "packages/pkg/src/nested.ts"], + timeoutMs: 200, + }); + + await Promise.race([ + secondEnsureStarted.promise, + new Promise((_, reject) => { + setTimeout( + () => + reject(new Error("Expected both files to start waiting before diagnostics publish")), + 50 + ); + }), + ]); + + for (const publish of publishByUri.values()) { + publish(); + } + + const diagnostics = await diagnosticsPromise; + expect(diagnostics).toHaveLength(2); + expect(diagnostics.map((entry) => entry.path)).toEqual([ + path.join(workspacePath, "packages", "pkg", "src", "nested.ts"), + path.join(workspacePath, "src", "example.ts"), + ]); + }); + test("keeps diagnostics isolated by root uri within the same workspace", async () => { const clientFactory = mock((options: CreateLspClientOptions): Promise => { const ensureFile = mock((file: Parameters[0]) => { @@ -264,6 +399,7 @@ describe("LspManager", () => { uri: file.uri, version: 1, diagnostics: [createDiagnostic(`diagnostic for ${options.rootPath}`)], + rawDiagnosticCount: 1, }); return Promise.resolve(1); }); diff --git a/src/node/services/lsp/lspManager.ts b/src/node/services/lsp/lspManager.ts index 51fcfcd789..f41b56dd90 100644 --- a/src/node/services/lsp/lspManager.ts +++ b/src/node/services/lsp/lspManager.ts @@ -164,57 +164,85 @@ export class LspManager { this.acquireLease(options.workspaceId); try { - const diagnostics: LspFileDiagnostics[] = []; - const seenFilePaths = new Set(); - - for (const filePath of options.filePaths) { - if (seenFilePaths.has(filePath)) { - continue; - } - seenFilePaths.add(filePath); - - try { - const context = await this.resolveClientContext({ - workspaceId: options.workspaceId, - runtime: options.runtime, - workspacePath: options.workspacePath, - filePath, - }); - const previousPublish = this.getLatestDiagnosticPublish( - options.workspaceId, - context.clientKey, - context.fileHandle.uri - ); - const expectedVersion = await context.client.ensureFile(context.fileHandle); - const freshPublish = await this.waitForFreshDiagnostics({ - workspaceId: options.workspaceId, - clientKey: context.clientKey, - uri: context.fileHandle.uri, - previousReceivedAtMs: previousPublish?.receivedAtMs, - expectedVersion, - timeoutMs: options.timeoutMs ?? LSP_POST_MUTATION_DIAGNOSTICS_TIMEOUT_MS, - }); - - if (!freshPublish) { - continue; + const fileContextsByClientKey = new Map< + string, + Array<{ filePath: string; context: ResolvedLspClientContext }> + >(); + const uniqueFilePaths = [...new Set(options.filePaths)]; + + await Promise.all( + uniqueFilePaths.map(async (filePath) => { + try { + const context = await this.resolveClientContext({ + workspaceId: options.workspaceId, + runtime: options.runtime, + workspacePath: options.workspacePath, + filePath, + }); + const entriesForClient = fileContextsByClientKey.get(context.clientKey) ?? []; + entriesForClient.push({ filePath, context }); + fileContextsByClientKey.set(context.clientKey, entriesForClient); + } catch (error) { + log.debug("Skipping post-mutation LSP diagnostics for file", { + workspaceId: options.workspaceId, + filePath, + error, + }); } + }) + ); - const snapshot = this.getCachedDiagnostics( - options.workspaceId, - context.clientKey, - context.fileHandle.uri - ); - if (snapshot && snapshot.receivedAtMs === freshPublish.receivedAtMs) { - diagnostics.push(snapshot); - } - } catch (error) { - log.debug("Skipping post-mutation LSP diagnostics for file", { - workspaceId: options.workspaceId, - filePath, - error, - }); - } - } + const diagnostics = ( + await Promise.all( + [...fileContextsByClientKey.values()].map( + async (fileContexts) => + await Promise.all( + fileContexts.map(async ({ filePath, context }) => { + try { + const previousPublish = this.getLatestDiagnosticPublish( + options.workspaceId, + context.clientKey, + context.fileHandle.uri + ); + const expectedVersion = await context.client.ensureFile(context.fileHandle); + const freshPublish = await this.waitForFreshDiagnostics({ + workspaceId: options.workspaceId, + clientKey: context.clientKey, + uri: context.fileHandle.uri, + previousReceivedAtMs: previousPublish?.receivedAtMs, + expectedVersion, + timeoutMs: options.timeoutMs ?? LSP_POST_MUTATION_DIAGNOSTICS_TIMEOUT_MS, + }); + + if (!freshPublish) { + return undefined; + } + + const snapshot = this.getCachedDiagnostics( + options.workspaceId, + context.clientKey, + context.fileHandle.uri + ); + if (snapshot && snapshot.receivedAtMs === freshPublish.receivedAtMs) { + return snapshot; + } + + return undefined; + } catch (error) { + log.debug("Skipping post-mutation LSP diagnostics for file", { + workspaceId: options.workspaceId, + filePath, + error, + }); + return undefined; + } + }) + ) + ) + ) + ) + .flat() + .filter((snapshot): snapshot is LspFileDiagnostics => snapshot !== undefined); this.markActivity(options.workspaceId); return diagnostics.sort((left, right) => left.path.localeCompare(right.path)); @@ -465,6 +493,10 @@ export class LspManager { pathMapper: LspPathMapper; params: LspPublishDiagnosticsParams; }): void { + if (isMalformedDiagnosticPublish(context.params)) { + return; + } + const receivedAtMs = Date.now(); const publishReceipt: LspDiagnosticPublishReceipt = { version: context.params.version, @@ -827,6 +859,10 @@ function isFreshDiagnosticPublish( return publish.receivedAtMs > (previousReceivedAtMs ?? 0); } +function isMalformedDiagnosticPublish(params: LspPublishDiagnosticsParams): boolean { + return params.rawDiagnosticCount > 0 && params.diagnostics.length === 0; +} + async function pathExists(runtime: Runtime, candidatePath: string): Promise { try { await runtime.stat(candidatePath); diff --git a/src/node/services/lsp/types.ts b/src/node/services/lsp/types.ts index bdc51b6430..cb275bc926 100644 --- a/src/node/services/lsp/types.ts +++ b/src/node/services/lsp/types.ts @@ -39,7 +39,11 @@ export interface LspMarkedString { } export interface LspHover { - contents: string | LspMarkupContent | LspMarkedString | Array; + contents: + | string + | LspMarkupContent + | LspMarkedString + | Array; } export interface LspDiagnostic { @@ -54,6 +58,7 @@ export interface LspPublishDiagnosticsParams { uri: string; version?: number; diagnostics: LspDiagnostic[]; + rawDiagnosticCount: number; } export interface LspFileDiagnostics { diff --git a/src/node/services/tools/task_apply_git_patch.test.ts b/src/node/services/tools/task_apply_git_patch.test.ts index f517d9baaf..ce17ff759a 100644 --- a/src/node/services/tools/task_apply_git_patch.test.ts +++ b/src/node/services/tools/task_apply_git_patch.test.ts @@ -34,8 +34,9 @@ async function commitFile( content: string, message: string ): Promise { + await fsPromises.mkdir(path.dirname(path.join(repoPath, fileName)), { recursive: true }); await fsPromises.writeFile(path.join(repoPath, fileName), content, "utf-8"); - execSync(`git add -- ${fileName}`, { cwd: repoPath, stdio: "ignore" }); + execSync(`git add -- ${JSON.stringify(fileName)}`, { cwd: repoPath, stdio: "ignore" }); execSync(`git commit -m ${JSON.stringify(message)}`, { cwd: repoPath, stdio: "ignore" }); } @@ -46,21 +47,23 @@ async function buildReadyProjectArtifact(params: { projectPath: string; projectName: string; childRepo: string; - baseSha: string; + baseSha?: string; headSha: string; + commitCount?: number; + formatPatchArgs?: string; }) { const patchPath = getSubagentGitPatchMboxPath( params.sessionDir, params.childTaskId, params.storageKey ); - const patch = execSync( - `git format-patch --stdout --binary ${params.baseSha}..${params.headSha}`, - { - cwd: params.childRepo, - encoding: "buffer", - } - ); + const formatPatchArgs = + params.formatPatchArgs ?? + (params.baseSha ? `${params.baseSha}..${params.headSha}` : `--root ${params.headSha}`); + const patch = execSync(`git format-patch --stdout --binary ${formatPatchArgs}`, { + cwd: params.childRepo, + encoding: "buffer", + }); await fsPromises.mkdir(path.dirname(patchPath), { recursive: true }); await fsPromises.writeFile(patchPath, patch); @@ -70,9 +73,9 @@ async function buildReadyProjectArtifact(params: { projectName: params.projectName, storageKey: params.storageKey, status: "ready" as const, - baseCommitSha: params.baseSha, + ...(params.baseSha ? { baseCommitSha: params.baseSha } : {}), headCommitSha: params.headSha, - commitCount: 1, + commitCount: params.commitCount ?? 1, mboxPath: patchPath, }; } @@ -766,6 +769,84 @@ describe("task_apply_git_patch tool", () => { expect(await readSubagentGitPatchArtifact(currentSessionDir, childTaskId)).toBeNull(); }, 20_000); + it("reports changed files across a root patch series, including paths with spaces", async () => { + const childRepo = path.join(rootDir, "root-series-child"); + const targetRepo = path.join(rootDir, "root-series-target"); + for (const repo of [childRepo, targetRepo]) { + await fsPromises.mkdir(repo, { recursive: true }); + initGitRepo(repo); + } + + await commitFile(childRepo, "README.md", "hello\n", "first change"); + await commitFile(childRepo, "docs/file with spaces.md", "second\n", "second change"); + const headSha = execSync("git rev-parse HEAD", { cwd: childRepo, encoding: "utf-8" }).trim(); + + const muxRoot = path.join(rootDir, "root-series-mux"); + const workspaceId = "root-series-workspace"; + const sessionDir = path.join(muxRoot, "sessions", workspaceId); + await fsPromises.mkdir(sessionDir, { recursive: true }); + await writeWorkspaceConfig({ + muxRoot, + workspaceId, + workspaceName: "root-series", + primaryProjectPath: targetRepo, + projects: [{ projectPath: targetRepo, projectName: "root-series" }], + }); + + const childTaskId = "root-series-task"; + await writePatchArtifact({ + sessionDir, + workspaceId, + childTaskId, + projectArtifacts: [ + await buildReadyProjectArtifact({ + sessionDir, + childTaskId, + storageKey: "root-series", + projectPath: targetRepo, + projectName: "root-series", + childRepo, + headSha, + commitCount: 2, + formatPatchArgs: `--root ${headSha}`, + }), + ], + }); + + const mutationCalls: Array<{ filePaths: string[] }> = []; + const tool = createTaskApplyGitPatchTool({ + ...getTestDeps(), + workspaceId, + cwd: targetRepo, + runtime: createRuntime({ type: "local", srcBaseDir: "/tmp" }), + runtimeTempDir: "/tmp", + workspaceSessionDir: sessionDir, + onFilesMutated: async (params) => { + mutationCalls.push(params); + return undefined; + }, + }); + + const result = (await tool.execute!({ task_id: childTaskId }, mockToolCallOptions)) as { + success: boolean; + appliedCommits?: Array<{ subject: string }>; + }; + + expect(result.success).toBe(true); + expect(result.appliedCommits?.map((commit) => commit.subject)).toEqual([ + "first change", + "second change", + ]); + expect(mutationCalls).toEqual([ + { + filePaths: [ + path.join(targetRepo, "README.md"), + path.join(targetRepo, "docs", "file with spaces.md"), + ], + }, + ]); + }, 20_000); + it("appends post-apply diagnostics notes for real applies", async () => { const fixture = await setupSingleProjectPatchFixture(rootDir, "diagnostics"); const mutationCalls: Array<{ filePaths: string[] }> = []; diff --git a/src/node/services/tools/task_apply_git_patch.ts b/src/node/services/tools/task_apply_git_patch.ts index a9794d36cf..5ee861368c 100644 --- a/src/node/services/tools/task_apply_git_patch.ts +++ b/src/node/services/tools/task_apply_git_patch.ts @@ -184,46 +184,51 @@ function joinRepoRelativePath(repoCwd: string, relativePath: string): string { return pathModule.join(repoCwd, ...normalizedSegments); } -async function getChangedFilesForAppliedRange(params: { +async function getChangedFilesForAppliedCommits(params: { runtime: ToolConfiguration["runtime"]; cwd: string; - beforeHeadSha: string | undefined; + appliedCommits: AppliedCommit[]; }): Promise { - const diffCommand = params.beforeHeadSha - ? `git diff --name-only ${params.beforeHeadSha}..HEAD --` - : "git diff-tree --no-commit-id --name-only -r HEAD"; + const changedFiles = new Set(); - try { - const result = await execBuffered(params.runtime, diffCommand, { - cwd: params.cwd, - timeout: 30, - }); - if (result.exitCode !== 0) { - log.debug("task_apply_git_patch: git diff --name-only failed", { + for (const commit of params.appliedCommits) { + if (!commit.sha) { + continue; + } + + try { + const result = await execBuffered( + params.runtime, + `git diff-tree --root --no-commit-id --name-only -z -r ${commit.sha} --`, + { + cwd: params.cwd, + timeout: 30, + } + ); + if (result.exitCode !== 0) { + log.debug("task_apply_git_patch: git diff-tree --name-only failed", { + cwd: params.cwd, + commitSha: commit.sha, + exitCode: result.exitCode, + stderr: result.stderr.trim(), + stdout: result.stdout.trim(), + }); + continue; + } + + for (const relativePath of result.stdout.split("\u0000").filter((line) => line.length > 0)) { + changedFiles.add(joinRepoRelativePath(params.cwd, relativePath)); + } + } catch (error) { + log.debug("task_apply_git_patch: git diff-tree --name-only threw", { cwd: params.cwd, - exitCode: result.exitCode, - stderr: result.stderr.trim(), - stdout: result.stdout.trim(), + commitSha: commit.sha, + error, }); - return []; } - - return Array.from( - new Set( - result.stdout - .split("\n") - .map((line) => line.replace(/\r$/, "").trim()) - .filter((line) => line.length > 0) - .map((line) => joinRepoRelativePath(params.cwd, line)) - ) - ); - } catch (error) { - log.debug("task_apply_git_patch: git diff --name-only threw", { - cwd: params.cwd, - error, - }); - return []; } + + return [...changedFiles]; } const MAX_PARENT_WORKSPACE_DEPTH = 32; @@ -923,10 +928,10 @@ async function applyProjectPatch(params: { let postMutationNote: string | undefined; if (params.onFilesMutated) { - const changedFiles = await getChangedFilesForAppliedRange({ + const changedFiles = await getChangedFilesForAppliedCommits({ runtime: params.runtime, cwd: params.repoCwd, - beforeHeadSha, + appliedCommits, }); if (changedFiles.length > 0) { try { From 4f4ee4304fae5b3883de756efd3c494f9cd0d2cc Mon Sep 17 00:00:00 2001 From: terion Date: Mon, 13 Apr 2026 13:54:26 +0300 Subject: [PATCH 05/46] Fail closed on task_apply_git_patch diff discovery errors --- .../tools/task_apply_git_patch.test.ts | 136 ++++++++++++++++++ .../services/tools/task_apply_git_patch.ts | 12 +- 2 files changed, 147 insertions(+), 1 deletion(-) diff --git a/src/node/services/tools/task_apply_git_patch.test.ts b/src/node/services/tools/task_apply_git_patch.test.ts index ce17ff759a..3f164dc55c 100644 --- a/src/node/services/tools/task_apply_git_patch.test.ts +++ b/src/node/services/tools/task_apply_git_patch.test.ts @@ -6,6 +6,7 @@ import { execSync } from "node:child_process"; import type { ToolExecutionOptions } from "ai"; +import type { ExecOptions, ExecStream, Runtime } from "@/node/runtime/Runtime"; import { createTaskApplyGitPatchTool } from "@/node/services/tools/task_apply_git_patch"; import { getSubagentGitPatchArtifactsFilePath, @@ -21,6 +22,51 @@ const mockToolCallOptions: ToolExecutionOptions = { messages: [], }; +function createTextStream(content: string): ReadableStream { + return new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode(content)); + controller.close(); + }, + }); +} + +function createExecStream(stdout = "", stderr = "", exitCode = 0): ExecStream { + return { + stdout: createTextStream(stdout), + stderr: createTextStream(stderr), + stdin: new WritableStream({ + write() { + return Promise.resolve(); + }, + close() { + return Promise.resolve(); + }, + }), + exitCode: Promise.resolve(exitCode), + duration: Promise.resolve(0), + }; +} + +function createRuntimeThatFailsNthDiffTree( + baseRuntime: Runtime, + failOnCall: number, + state: { diffTreeCalls: number } +): Runtime { + const runtime = Object.create(baseRuntime) as Runtime; + runtime.exec = async (command: string, options: ExecOptions) => { + if (command.startsWith("git diff-tree --root --no-commit-id --name-only -z -r ")) { + state.diffTreeCalls += 1; + if (state.diffTreeCalls === failOnCall) { + return createExecStream("", "diff-tree failed", 1); + } + } + + return await baseRuntime.exec(command, options); + }; + return runtime; +} + function initGitRepo(repoPath: string): void { execSync("git init -b main", { cwd: repoPath, stdio: "ignore" }); execSync('git config user.email "test@example.com"', { cwd: repoPath, stdio: "ignore" }); @@ -847,6 +893,96 @@ describe("task_apply_git_patch tool", () => { ]); }, 20_000); + it("fails closed for post-apply diagnostics when changed-file discovery fails mid-series", async () => { + const childRepo = path.join(rootDir, "diff-tree-child"); + const targetRepo = path.join(rootDir, "diff-tree-target"); + for (const repo of [childRepo, targetRepo]) { + await fsPromises.mkdir(repo, { recursive: true }); + initGitRepo(repo); + } + + await commitFile(childRepo, "README.md", "hello\n", "base"); + await commitFile(targetRepo, "README.md", "hello\n", "base"); + const baseSha = execSync("git rev-parse HEAD", { cwd: childRepo, encoding: "utf-8" }).trim(); + + await commitFile(childRepo, "README.md", "hello\nfirst\n", "first change"); + await commitFile(childRepo, "docs/second.md", "second\n", "second change"); + const headSha = execSync("git rev-parse HEAD", { cwd: childRepo, encoding: "utf-8" }).trim(); + + const muxRoot = path.join(rootDir, "diff-tree-mux"); + const workspaceId = "diff-tree-workspace"; + const sessionDir = path.join(muxRoot, "sessions", workspaceId); + await fsPromises.mkdir(sessionDir, { recursive: true }); + await writeWorkspaceConfig({ + muxRoot, + workspaceId, + workspaceName: "diff-tree", + primaryProjectPath: targetRepo, + projects: [{ projectPath: targetRepo, projectName: "diff-tree" }], + }); + + const childTaskId = "diff-tree-task"; + await writePatchArtifact({ + sessionDir, + workspaceId, + childTaskId, + projectArtifacts: [ + await buildReadyProjectArtifact({ + sessionDir, + childTaskId, + storageKey: "diff-tree", + projectPath: targetRepo, + projectName: "diff-tree", + childRepo, + baseSha, + headSha, + commitCount: 2, + }), + ], + }); + + const runtimeState = { diffTreeCalls: 0 }; + const runtime = createRuntimeThatFailsNthDiffTree( + createRuntime({ type: "local", srcBaseDir: "/tmp" }), + 2, + runtimeState + ); + const mutationCalls: Array<{ filePaths: string[] }> = []; + const tool = createTaskApplyGitPatchTool({ + ...getTestDeps(), + workspaceId, + cwd: targetRepo, + runtime, + runtimeTempDir: "/tmp", + workspaceSessionDir: sessionDir, + onFilesMutated: async (params) => { + mutationCalls.push(params); + return "should not run"; + }, + }); + + const result = (await tool.execute!({ task_id: childTaskId }, mockToolCallOptions)) as { + success: boolean; + note?: string; + appliedCommits?: Array<{ subject: string }>; + }; + + expect(result.success).toBe(true); + expect(result.appliedCommits?.map((commit) => commit.subject)).toEqual([ + "first change", + "second change", + ]); + expect(runtimeState.diffTreeCalls).toBe(2); + expect(mutationCalls).toEqual([]); + expect(result.note ?? "").not.toContain("should not run"); + expect(await fsPromises.readFile(path.join(targetRepo, "README.md"), "utf-8")).toBe( + "hello\nfirst\n" + ); + expect(await fsPromises.readFile(path.join(targetRepo, "docs", "second.md"), "utf-8")).toBe( + "second\n" + ); + }, 20_000); + it("appends post-apply diagnostics notes for real applies", async () => { const fixture = await setupSingleProjectPatchFixture(rootDir, "diagnostics"); const mutationCalls: Array<{ filePaths: string[] }> = []; diff --git a/src/node/services/tools/task_apply_git_patch.ts b/src/node/services/tools/task_apply_git_patch.ts index 5ee861368c..a5ec5bf831 100644 --- a/src/node/services/tools/task_apply_git_patch.ts +++ b/src/node/services/tools/task_apply_git_patch.ts @@ -190,6 +190,7 @@ async function getChangedFilesForAppliedCommits(params: { appliedCommits: AppliedCommit[]; }): Promise { const changedFiles = new Set(); + let discoveryFailed = false; for (const commit of params.appliedCommits) { if (!commit.sha) { @@ -213,7 +214,8 @@ async function getChangedFilesForAppliedCommits(params: { stderr: result.stderr.trim(), stdout: result.stdout.trim(), }); - continue; + discoveryFailed = true; + break; } for (const relativePath of result.stdout.split("\u0000").filter((line) => line.length > 0)) { @@ -225,9 +227,17 @@ async function getChangedFilesForAppliedCommits(params: { commitSha: commit.sha, error, }); + discoveryFailed = true; + break; } } + if (discoveryFailed) { + // Post-apply diagnostics must fail closed here because a partial path list could hide + // diagnostics from later commits in the applied series. + return []; + } + return [...changedFiles]; } From 21781c6d62e5daddd69269ac6b2535cba037135b Mon Sep 17 00:00:00 2001 From: terion Date: Mon, 13 Apr 2026 14:42:33 +0300 Subject: [PATCH 06/46] =?UTF-8?q?=F0=9F=A4=96=20feat:=20add=20workspace=20?= =?UTF-8?q?LSP=20diagnostics=20snapshots?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add shared workspace LSP diagnostics schemas, backend list/subscribe plumbing, browser WorkspaceStore caching/hooks, and focused coverage for manager/store/server paths. --- src/browser/stores/WorkspaceStore.test.ts | 221 ++++++++++++++++++++- src/browser/stores/WorkspaceStore.ts | 147 +++++++++++++- src/cli/server.test.ts | 100 ++++++++++ src/common/orpc/schemas.ts | 16 ++ src/common/orpc/schemas/api.ts | 11 + src/common/orpc/schemas/workspaceLsp.ts | 40 ++++ src/common/orpc/types.ts | 7 + src/node/orpc/context.ts | 2 + src/node/orpc/router.ts | 46 +++++ src/node/services/lsp/lspManager.test.ts | 133 ++++++++++++- src/node/services/lsp/lspManager.ts | 99 ++++++++- src/node/services/serviceContainer.test.ts | 3 +- src/node/services/serviceContainer.ts | 1 + 13 files changed, 813 insertions(+), 13 deletions(-) create mode 100644 src/common/orpc/schemas/workspaceLsp.ts diff --git a/src/browser/stores/WorkspaceStore.test.ts b/src/browser/stores/WorkspaceStore.test.ts index 7b88319e63..9e7941b2f2 100644 --- a/src/browser/stores/WorkspaceStore.test.ts +++ b/src/browser/stores/WorkspaceStore.test.ts @@ -1,7 +1,11 @@ import { describe, expect, it, beforeEach, afterEach, mock, type Mock } from "bun:test"; import type { FrontendWorkspaceMetadata } from "@/common/types/workspace"; import type { StreamStartEvent, ToolCallStartEvent } from "@/common/types/stream"; -import type { WorkspaceActivitySnapshot, WorkspaceChatMessage } from "@/common/orpc/types"; +import type { + WorkspaceActivitySnapshot, + WorkspaceChatMessage, + WorkspaceLspDiagnosticsSnapshot, +} from "@/common/orpc/types"; import { DEFAULT_RUNTIME_CONFIG } from "@/common/constants/workspace"; import { DEFAULT_AUTO_COMPACTION_THRESHOLD_PERCENT } from "@/common/constants/ui"; import { @@ -49,6 +53,98 @@ const mockHistoryLoadMore = mock( ); const mockActivityList = mock(() => Promise.resolve>({})); +interface ControlledAsyncIterator { + iterator: AsyncGenerator; + push: (value: T) => void; + end: () => void; +} + +function createControlledAsyncIterator(signal?: AbortSignal): ControlledAsyncIterator { + const queue: T[] = []; + let ended = false; + let resolveNext: ((result: IteratorResult) => void) | null = null; + + const push = (value: T) => { + if (ended) { + return; + } + if (resolveNext) { + const resolve = resolveNext; + resolveNext = null; + resolve({ value, done: false }); + return; + } + queue.push(value); + }; + + const end = () => { + if (ended) { + return; + } + ended = true; + if (resolveNext) { + const resolve = resolveNext; + resolveNext = null; + resolve({ value: undefined as void, done: true }); + } + }; + + signal?.addEventListener("abort", end, { once: true }); + + const iterator: AsyncGenerator = { + async next(): Promise> { + if (queue.length > 0) { + return { value: queue.shift()!, done: false }; + } + if (ended) { + return { value: undefined as void, done: true }; + } + return await new Promise>((resolve) => { + resolveNext = resolve; + }); + }, + return(): Promise> { + end(); + return Promise.resolve({ value: undefined as void, done: true }); + }, + throw(error: unknown): Promise> { + end(); + return Promise.reject(error instanceof Error ? error : new Error(String(error))); + }, + [Symbol.asyncDispose](): Promise { + end(); + return Promise.resolve(); + }, + [Symbol.asyncIterator]() { + return this; + }, + }; + + return { iterator, push, end }; +} + +const lspDiagnosticsSubscriptions: Array<{ + workspaceId: string; + signal?: AbortSignal; + push: (snapshot: WorkspaceLspDiagnosticsSnapshot) => void; + end: () => void; +}> = []; + +function emitLspDiagnosticsSnapshot(snapshot: WorkspaceLspDiagnosticsSnapshot): void { + for (const subscription of lspDiagnosticsSubscriptions) { + if (subscription.workspaceId === snapshot.workspaceId) { + subscription.push(snapshot); + } + } +} + +function clearLspDiagnosticsSubscriptions(): void { + for (const subscription of lspDiagnosticsSubscriptions) { + subscription.end(); + } + lspDiagnosticsSubscriptions.length = 0; +} + type WorkspaceActivityEvent = | { type: "activity"; @@ -99,6 +195,21 @@ const mockSetAutoCompactionThreshold = mock(() => Promise.resolve({ success: true, data: undefined }) ); const mockGetStartupAutoRetryModel = mock(() => Promise.resolve({ success: true, data: null })); +const mockListLspDiagnostics = mock(({ workspaceId }: { workspaceId: string }) => + Promise.resolve({ workspaceId, diagnostics: [] }) +); +const mockSubscribeLspDiagnostics = mock( + ({ workspaceId }: { workspaceId: string }, options?: { signal?: AbortSignal }) => { + const controlled = createControlledAsyncIterator(options?.signal); + lspDiagnosticsSubscriptions.push({ + workspaceId, + signal: options?.signal, + push: controlled.push, + end: controlled.end, + }); + return controlled.iterator; + } +); const mockClient = { workspace: { @@ -113,6 +224,10 @@ const mockClient = { }, setAutoCompactionThreshold: mockSetAutoCompactionThreshold, getStartupAutoRetryModel: mockGetStartupAutoRetryModel, + lsp: { + listDiagnostics: mockListLspDiagnostics, + subscribeDiagnostics: mockSubscribeLspDiagnostics, + }, }, terminal: { activity: { @@ -279,6 +394,9 @@ describe("WorkspaceStore", () => { mockTerminalActivitySubscribe.mockClear(); mockSetAutoCompactionThreshold.mockClear(); mockGetStartupAutoRetryModel.mockClear(); + mockListLspDiagnostics.mockClear(); + mockSubscribeLspDiagnostics.mockClear(); + clearLspDiagnosticsSubscriptions(); global.window.localStorage?.clear?.(); mockHistoryLoadMore.mockResolvedValue({ messages: [], @@ -314,6 +432,7 @@ describe("WorkspaceStore", () => { afterEach(() => { store.dispose(); + clearLspDiagnosticsSubscriptions(); }); describe("pinned todo auto-collapse", () => { @@ -4485,4 +4604,104 @@ describe("WorkspaceStore", () => { expect(usage.lastContextUsage?.model).toBe("claude-3-5-sonnet-20241022"); }); }); + + + describe("workspace LSP diagnostics snapshots", () => { + it("starts empty and replaces the cached snapshot after the first payload", async () => { + const workspaceId = "lsp-diagnostics-snapshot"; + createAndAddWorkspace(store, workspaceId); + + expect(store.getWorkspaceLspDiagnosticsSnapshot(workspaceId)).toBeNull(); + + const unsubscribe = store.subscribeLspDiagnostics(workspaceId, () => undefined); + const subscribed = await waitUntil(() => mockSubscribeLspDiagnostics.mock.calls.length === 1); + expect(subscribed).toBe(true); + + emitLspDiagnosticsSnapshot({ + workspaceId, + diagnostics: [ + { + uri: "file:///workspace/src/example.ts", + path: "/workspace/src/example.ts", + serverId: "typescript", + rootUri: "file:///workspace", + version: 1, + diagnostics: [ + { + range: { + start: { line: 1, character: 1 }, + end: { line: 1, character: 5 }, + }, + severity: 1, + source: "tsserver", + message: "first snapshot", + }, + ], + receivedAtMs: 1, + }, + ], + }); + + const received = await waitUntil( + () => + store.getWorkspaceLspDiagnosticsSnapshot(workspaceId)?.diagnostics[0]?.diagnostics[0] + ?.message === "first snapshot" + ); + expect(received).toBe(true); + + emitLspDiagnosticsSnapshot({ workspaceId, diagnostics: [] }); + const cleared = await waitUntil( + () => store.getWorkspaceLspDiagnosticsSnapshot(workspaceId)?.diagnostics.length === 0 + ); + expect(cleared).toBe(true); + + unsubscribe(); + }); + + it("shares one backend subscription across listeners and clears state on the last unsubscribe", async () => { + const workspaceId = "lsp-diagnostics-refcount"; + createAndAddWorkspace(store, workspaceId); + + const unsubscribeOne = store.subscribeLspDiagnostics(workspaceId, () => undefined); + const unsubscribeTwo = store.subscribeLspDiagnostics(workspaceId, () => undefined); + const subscribed = await waitUntil(() => mockSubscribeLspDiagnostics.mock.calls.length === 1); + expect(subscribed).toBe(true); + expect(mockSubscribeLspDiagnostics).toHaveBeenCalledTimes(1); + + emitLspDiagnosticsSnapshot({ workspaceId, diagnostics: [] }); + await waitUntil( + () => store.getWorkspaceLspDiagnosticsSnapshot(workspaceId)?.workspaceId === workspaceId + ); + + unsubscribeOne(); + expect(lspDiagnosticsSubscriptions[0]?.signal?.aborted).toBe(false); + expect(store.getWorkspaceLspDiagnosticsSnapshot(workspaceId)?.workspaceId).toBe(workspaceId); + + unsubscribeTwo(); + const aborted = await waitUntil(() => lspDiagnosticsSubscriptions[0]?.signal?.aborted === true); + expect(aborted).toBe(true); + expect(store.getWorkspaceLspDiagnosticsSnapshot(workspaceId)).toBeNull(); + }); + + it("aborts LSP diagnostics subscriptions when a workspace is removed", async () => { + const workspaceId = "lsp-diagnostics-remove"; + createAndAddWorkspace(store, workspaceId); + + store.subscribeLspDiagnostics(workspaceId, () => undefined); + const subscribed = await waitUntil(() => mockSubscribeLspDiagnostics.mock.calls.length === 1); + expect(subscribed).toBe(true); + + emitLspDiagnosticsSnapshot({ workspaceId, diagnostics: [] }); + await waitUntil( + () => store.getWorkspaceLspDiagnosticsSnapshot(workspaceId)?.workspaceId === workspaceId + ); + + store.removeWorkspace(workspaceId); + + const aborted = await waitUntil(() => lspDiagnosticsSubscriptions[0]?.signal?.aborted === true); + expect(aborted).toBe(true); + expect(store.getWorkspaceLspDiagnosticsSnapshot(workspaceId)).toBeNull(); + }); + }); + }); diff --git a/src/browser/stores/WorkspaceStore.ts b/src/browser/stores/WorkspaceStore.ts index 34800653ed..9d54e28ec4 100644 --- a/src/browser/stores/WorkspaceStore.ts +++ b/src/browser/stores/WorkspaceStore.ts @@ -4,6 +4,7 @@ import type { FrontendWorkspaceMetadata } from "@/common/types/workspace"; import type { WorkspaceActivitySnapshot, WorkspaceChatMessage, + WorkspaceLspDiagnosticsSnapshot, WorkspaceStatsSnapshot, OnChatMode, ProvidersConfigMap, @@ -526,6 +527,13 @@ export class WorkspaceStore { // Per-workspace listener refcount for useWorkspaceStatsSnapshot(). // Used to only subscribe to backend stats when something in the UI is actually reading them. private statsListenerCounts = new Map(); + + // Workspace LSP diagnostics snapshots (from workspace.lsp.subscribeDiagnostics) + private workspaceLspDiagnostics = new Map(); + private lspDiagnosticsStore = new MapStore(); + private lspDiagnosticsUnsubscribers = new Map void>(); + private lspDiagnosticsListenerCounts = new Map(); + // Cumulative session usage (from session-usage.json) private sessionUsage = new Map>(); @@ -1021,11 +1029,15 @@ export class WorkspaceStore { return; } - // Drop stats subscriptions before swapping clients so reconnects resubscribe cleanly. + // Drop lazy workspace subscriptions before swapping clients so reconnects resubscribe cleanly. for (const unsubscribe of this.statsUnsubscribers.values()) { unsubscribe(); } this.statsUnsubscribers.clear(); + for (const unsubscribe of this.lspDiagnosticsUnsubscribers.values()) { + unsubscribe(); + } + this.lspDiagnosticsUnsubscribers.clear(); this.client = client; this.clientChangeController.abort(); @@ -1333,6 +1345,58 @@ export class WorkspaceStore { }); } + private subscribeToLspDiagnostics(workspaceId: string): void { + if (!this.client) { + return; + } + + if (!this.isWorkspaceRegistered(workspaceId)) { + return; + } + if ((this.lspDiagnosticsListenerCounts.get(workspaceId) ?? 0) <= 0) { + return; + } + if (this.lspDiagnosticsUnsubscribers.has(workspaceId)) { + return; + } + + const controller = new AbortController(); + const { signal } = controller; + let iterator: AsyncIterator | null = null; + + (async () => { + try { + const subscribedIterator = await this.client!.workspace.lsp.subscribeDiagnostics( + { workspaceId }, + { signal } + ); + iterator = subscribedIterator; + + for await (const snapshot of subscribedIterator) { + if (signal.aborted) break; + queueMicrotask(() => { + if (signal.aborted) { + return; + } + this.workspaceLspDiagnostics.set(workspaceId, snapshot); + this.lspDiagnosticsStore.bump(workspaceId); + }); + } + } catch (error) { + if (signal.aborted || isAbortError(error)) return; + console.warn( + `[WorkspaceStore] Error in LSP diagnostics subscription for ${workspaceId}:`, + error + ); + } + })(); + + this.lspDiagnosticsUnsubscribers.set(workspaceId, () => { + controller.abort(); + void iterator?.return?.(); + }); + } + /** * Cancel any pending idle state bump for a workspace. * Used when immediate state visibility is needed (e.g., stream-end). @@ -1890,6 +1954,12 @@ export class WorkspaceStore { }); } + getWorkspaceLspDiagnosticsSnapshot(workspaceId: string): WorkspaceLspDiagnosticsSnapshot | null { + return this.lspDiagnosticsStore.get(workspaceId, () => { + return this.workspaceLspDiagnostics.get(workspaceId) ?? null; + }); + } + /** * Bump state for a workspace to trigger React re-renders. * Used by addEphemeralMessage for frontend-only messages. @@ -2184,6 +2254,46 @@ export class WorkspaceStore { }; } + subscribeLspDiagnostics(workspaceId: string, listener: () => void): () => void { + const unsubscribeFromStore = this.lspDiagnosticsStore.subscribeKey(workspaceId, listener); + + const previousCount = this.lspDiagnosticsListenerCounts.get(workspaceId) ?? 0; + const nextCount = previousCount + 1; + this.lspDiagnosticsListenerCounts.set(workspaceId, nextCount); + + if (previousCount === 0) { + this.subscribeToLspDiagnostics(workspaceId); + } + + return () => { + unsubscribeFromStore(); + + const currentCount = this.lspDiagnosticsListenerCounts.get(workspaceId); + if (!currentCount) { + console.warn( + `[WorkspaceStore] LSP diagnostics listener count underflow for ${workspaceId} (already 0)` + ); + return; + } + + if (currentCount === 1) { + this.lspDiagnosticsListenerCounts.delete(workspaceId); + + const unsubscribe = this.lspDiagnosticsUnsubscribers.get(workspaceId); + if (unsubscribe) { + unsubscribe(); + this.lspDiagnosticsUnsubscribers.delete(workspaceId); + } + this.workspaceLspDiagnostics.delete(workspaceId); + this.lspDiagnosticsStore.bump(workspaceId); + this.lspDiagnosticsStore.delete(workspaceId); + return; + } + + this.lspDiagnosticsListenerCounts.set(workspaceId, currentCount - 1); + }; + } + /** * Subscribe to consumer store changes for a specific workspace. */ @@ -3163,6 +3273,9 @@ export class WorkspaceStore { // Stats snapshots are subscribed lazily via subscribeStats(). this.subscribeToStats(workspaceId); + // LSP diagnostics snapshots are subscribed lazily via subscribeLspDiagnostics(). + this.subscribeToLspDiagnostics(workspaceId); + this.ensureActiveOnChatSubscription(); if (!this.client) { @@ -3190,6 +3303,12 @@ export class WorkspaceStore { this.statsUnsubscribers.delete(workspaceId); } + const lspDiagnosticsUnsubscribe = this.lspDiagnosticsUnsubscribers.get(workspaceId); + if (lspDiagnosticsUnsubscribe) { + lspDiagnosticsUnsubscribe(); + this.lspDiagnosticsUnsubscribers.delete(workspaceId); + } + const unsubscribe = this.ipcUnsubscribers.get(workspaceId); if (unsubscribe) { unsubscribe(); @@ -3218,6 +3337,9 @@ export class WorkspaceStore { this.workspaceStats.delete(workspaceId); this.statsStore.delete(workspaceId); this.statsListenerCounts.delete(workspaceId); + this.workspaceLspDiagnostics.delete(workspaceId); + this.lspDiagnosticsStore.delete(workspaceId); + this.lspDiagnosticsListenerCounts.delete(workspaceId); this.historyPagination.delete(workspaceId); this.preReplayUsageSnapshot.delete(workspaceId); this.sessionUsage.delete(workspaceId); @@ -3268,6 +3390,10 @@ export class WorkspaceStore { unsubscribe(); } this.statsUnsubscribers.clear(); + for (const unsubscribe of this.lspDiagnosticsUnsubscribers.values()) { + unsubscribe(); + } + this.lspDiagnosticsUnsubscribers.clear(); for (const unsubscribe of this.ipcUnsubscribers.values()) { unsubscribe(); @@ -3304,6 +3430,9 @@ export class WorkspaceStore { this.workspaceStats.clear(); this.statsStore.clear(); this.statsListenerCounts.clear(); + this.workspaceLspDiagnostics.clear(); + this.lspDiagnosticsStore.clear(); + this.lspDiagnosticsListenerCounts.clear(); this.historyPagination.clear(); this.preReplayUsageSnapshot.clear(); this.sessionUsage.clear(); @@ -3953,6 +4082,22 @@ export function useWorkspaceStatsSnapshot(workspaceId: string): WorkspaceStatsSn return useSyncExternalStore(subscribe, getSnapshot); } +export function useWorkspaceLspDiagnosticsSnapshot( + workspaceId: string +): WorkspaceLspDiagnosticsSnapshot | null { + const store = getStoreInstance(); + const subscribe = useCallback( + (listener: () => void) => store.subscribeLspDiagnostics(workspaceId, listener), + [store, workspaceId] + ); + const getSnapshot = useCallback( + () => store.getWorkspaceLspDiagnosticsSnapshot(workspaceId), + [store, workspaceId] + ); + + return useSyncExternalStore(subscribe, getSnapshot); +} + /** * Hook for consumer breakdown (lazy, with tokenization). * Updates after async Web Worker calculation completes. diff --git a/src/cli/server.test.ts b/src/cli/server.test.ts index bdac7bfa21..571f48a74a 100644 --- a/src/cli/server.test.ts +++ b/src/cli/server.test.ts @@ -26,12 +26,15 @@ import { ServiceContainer } from "@/node/services/serviceContainer"; import type { RouterClient } from "@orpc/server"; import { createOrpcServer, type OrpcServer } from "@/node/orpc/server"; import type { ProjectConfig } from "@/common/types/project"; +import type { WorkspaceLspDiagnosticsSnapshot } from "@/common/orpc/types"; +import type { LspFileDiagnostics } from "@/node/services/lsp/types"; import { shouldExposeLaunchProject } from "@/cli/launchProject"; // --- Test Server Factory --- interface TestServerHandle { server: OrpcServer; + services: ServiceContainer; tempDir: string; close: () => Promise; } @@ -72,9 +75,12 @@ async function createTestServer(): Promise { return { server, + services, tempDir, close: async () => { await server.close(); + await services.dispose(); + await services.shutdown(); // Cleanup temp directory await fs.rm(tempDir, { recursive: true, force: true }).catch(() => undefined); }, @@ -124,6 +130,22 @@ function createProjectConfig(projectKind?: ProjectConfig["projectKind"]): Projec }; } +function setWorkspaceLspDiagnostics( + services: ServiceContainer, + snapshot: WorkspaceLspDiagnosticsSnapshot +): void { + const manager = services.lspManager as unknown as { + workspaceDiagnostics: Map>>; + notifyWorkspaceDiagnosticsListeners: (workspaceId: string) => void; + }; + const diagnosticsByUri = new Map(); + for (const diagnostics of snapshot.diagnostics) { + diagnosticsByUri.set(diagnostics.uri, diagnostics); + } + manager.workspaceDiagnostics.set(snapshot.workspaceId, new Map([["typescript:file:///workspace", diagnosticsByUri]])); + manager.notifyWorkspaceDiagnosticsListeners(snapshot.workspaceId); +} + describe("shouldExposeLaunchProject", () => { test("returns true when only system projects exist", () => { const projects: Array<[string, ProjectConfig]> = [ @@ -271,6 +293,38 @@ describe("oRPC Server Endpoints", () => { expect(ticks).toHaveLength(1); expect(ticks[0].tick).toBe(1); }); + + test("workspace.lsp.listDiagnostics returns the current workspace snapshot", async () => { + const client = createHttpClient(serverHandle.server.baseUrl); + setWorkspaceLspDiagnostics(serverHandle.services, { + workspaceId: "lsp-http-workspace", + diagnostics: [ + { + uri: "file:///workspace/src/example.ts", + path: "/workspace/src/example.ts", + serverId: "typescript", + rootUri: "file:///workspace", + version: 1, + diagnostics: [ + { + range: { + start: { line: 1, character: 1 }, + end: { line: 1, character: 5 }, + }, + severity: 1, + source: "tsserver", + message: "http snapshot", + }, + ], + receivedAtMs: 1, + }, + ], + }); + + const result = await client.workspace.lsp.listDiagnostics({ workspaceId: "lsp-http-workspace" }); + expect(result.workspaceId).toBe("lsp-http-workspace"); + expect(result.diagnostics[0]?.diagnostics[0]?.message).toBe("http snapshot"); + }); }); describe("WebSocket endpoint (/orpc/ws)", () => { @@ -351,6 +405,52 @@ describe("oRPC Server Endpoints", () => { close(); } }); + + test("workspace.lsp.subscribeDiagnostics emits initial and subsequent snapshots", async () => { + const workspaceId = "lsp-ws-workspace"; + const { client, close } = await createWebSocketClient(serverHandle.server.wsUrl); + try { + const stream = await client.workspace.lsp.subscribeDiagnostics({ workspaceId }); + const initial = + (await stream.next()) as IteratorResult; + expect(initial.done).toBe(false); + expect(initial.value?.workspaceId).toBe(workspaceId); + expect(initial.value?.diagnostics).toEqual([]); + + setWorkspaceLspDiagnostics(serverHandle.services, { + workspaceId, + diagnostics: [ + { + uri: "file:///workspace/src/example.ts", + path: "/workspace/src/example.ts", + serverId: "typescript", + rootUri: "file:///workspace", + version: 2, + diagnostics: [ + { + range: { + start: { line: 2, character: 1 }, + end: { line: 2, character: 5 }, + }, + severity: 1, + source: "tsserver", + message: "ws snapshot", + }, + ], + receivedAtMs: 2, + }, + ], + }); + + const update = + (await stream.next()) as IteratorResult; + expect(update.done).toBe(false); + expect(update.value?.diagnostics[0]?.diagnostics[0]?.message).toBe("ws snapshot"); + await stream.return?.(); + } finally { + close(); + } + }); }); describe("Cross-transport consistency", () => { diff --git a/src/common/orpc/schemas.ts b/src/common/orpc/schemas.ts index 352110e81d..798e851b42 100644 --- a/src/common/orpc/schemas.ts +++ b/src/common/orpc/schemas.ts @@ -39,6 +39,22 @@ export { WorkspaceStatsSnapshotSchema, } from "./schemas/workspaceStats"; +// Workspace LSP diagnostics schemas +export { + LspDiagnosticSchema, + LspFileDiagnosticsSchema, + LspPositionSchema, + LspRangeSchema, + WorkspaceLspDiagnosticsSnapshotSchema, +} from "./schemas/workspaceLsp"; +export type { + LspDiagnostic, + LspFileDiagnostics, + LspPosition, + LspRange, + WorkspaceLspDiagnosticsSnapshot, +} from "./schemas/workspaceLsp"; + // Analytics schemas export { AgentCostRowSchema, diff --git a/src/common/orpc/schemas/api.ts b/src/common/orpc/schemas/api.ts index a6386d608e..8530669d7e 100644 --- a/src/common/orpc/schemas/api.ts +++ b/src/common/orpc/schemas/api.ts @@ -38,6 +38,7 @@ import { } from "./terminal"; import { BashToolResultSchema, FileTreeNodeSchema } from "./tools"; import { WorkspaceStatsSnapshotSchema } from "./workspaceStats"; +import { WorkspaceLspDiagnosticsSnapshotSchema } from "./workspaceLsp"; import { FrontendWorkspaceMetadataSchema, GitStatusSchema, @@ -1399,6 +1400,16 @@ export const workspace = { output: ResultSchema(z.void(), z.string()), }, }, + lsp: { + listDiagnostics: { + input: z.object({ workspaceId: z.string() }), + output: WorkspaceLspDiagnosticsSnapshotSchema, + }, + subscribeDiagnostics: { + input: z.object({ workspaceId: z.string() }), + output: eventIterator(WorkspaceLspDiagnosticsSnapshotSchema), + }, + }, getSessionUsage: { input: z.object({ workspaceId: z.string() }), output: SessionUsageFileSchema.optional(), diff --git a/src/common/orpc/schemas/workspaceLsp.ts b/src/common/orpc/schemas/workspaceLsp.ts new file mode 100644 index 0000000000..983c3a9d07 --- /dev/null +++ b/src/common/orpc/schemas/workspaceLsp.ts @@ -0,0 +1,40 @@ +import { z } from "zod"; + +export const LspPositionSchema = z.object({ + line: z.number(), + character: z.number(), +}); + +export const LspRangeSchema = z.object({ + start: LspPositionSchema, + end: LspPositionSchema, +}); + +export const LspDiagnosticSchema = z.object({ + range: LspRangeSchema, + severity: z.union([z.literal(1), z.literal(2), z.literal(3), z.literal(4)]).optional(), + code: z.union([z.string(), z.number()]).optional(), + source: z.string().optional(), + message: z.string(), +}); + +export const LspFileDiagnosticsSchema = z.object({ + uri: z.string(), + path: z.string(), + serverId: z.string(), + rootUri: z.string(), + version: z.number().optional(), + diagnostics: z.array(LspDiagnosticSchema), + receivedAtMs: z.number(), +}); + +export const WorkspaceLspDiagnosticsSnapshotSchema = z.object({ + workspaceId: z.string(), + diagnostics: z.array(LspFileDiagnosticsSchema), +}); + +export type LspPosition = z.infer; +export type LspRange = z.infer; +export type LspDiagnostic = z.infer; +export type LspFileDiagnostics = z.infer; +export type WorkspaceLspDiagnosticsSnapshot = z.infer; diff --git a/src/common/orpc/types.ts b/src/common/orpc/types.ts index 74442366ac..cc7dd25f0c 100644 --- a/src/common/orpc/types.ts +++ b/src/common/orpc/types.ts @@ -46,6 +46,13 @@ export type UpdateStatus = z.infer; export type DesktopPrereqStatus = z.infer; export type ChatMuxMessage = z.infer; export type WorkspaceStatsSnapshot = z.infer; +export type LspPosition = z.infer; +export type LspRange = z.infer; +export type LspDiagnostic = z.infer; +export type LspFileDiagnostics = z.infer; +export type WorkspaceLspDiagnosticsSnapshot = z.infer< + typeof schemas.WorkspaceLspDiagnosticsSnapshotSchema +>; export type WorkspaceActivitySnapshot = z.infer; export type FrontendWorkspaceMetadataSchemaType = z.infer< typeof schemas.FrontendWorkspaceMetadataSchema diff --git a/src/node/orpc/context.ts b/src/node/orpc/context.ts index eb05595838..2f1a6de79b 100644 --- a/src/node/orpc/context.ts +++ b/src/node/orpc/context.ts @@ -38,6 +38,7 @@ import type { CoderService } from "@/node/services/coderService"; import type { ServerAuthService } from "@/node/services/serverAuthService"; import type { SshPromptService } from "@/node/services/sshPromptService"; import type { AnalyticsService } from "@/node/services/analytics/analyticsService"; +import type { LspManager } from "@/node/services/lsp/lspManager"; import type { DesktopBridgeServer } from "@/node/services/desktop/DesktopBridgeServer"; import type { DesktopSessionManager } from "@/node/services/desktop/DesktopSessionManager"; import type { DesktopTokenManager } from "@/node/services/desktop/DesktopTokenManager"; @@ -70,6 +71,7 @@ export interface ORPCContext { telemetryService: TelemetryService; experimentsService: ExperimentsService; sessionUsageService: SessionUsageService; + lspManager: LspManager; devToolsService: DevToolsService; browserSessionDiscoveryService: AgentBrowserSessionDiscoveryService; browserBridgeTokenManager: BrowserBridgeTokenManager; diff --git a/src/node/orpc/router.ts b/src/node/orpc/router.ts index 545ea42287..128b9f6670 100644 --- a/src/node/orpc/router.ts +++ b/src/node/orpc/router.ts @@ -16,6 +16,7 @@ import type { UpdateStatus, WorkspaceActivitySnapshot, WorkspaceChatMessage, + WorkspaceLspDiagnosticsSnapshot, WorkspaceStatsSnapshot, FrontendWorkspaceMetadataSchemaType, } from "@/common/orpc/types"; @@ -3889,6 +3890,51 @@ export const router = (authToken?: string) => { .handler(async ({ context, input }) => { return context.sessionUsageService.getSessionUsageBatch(input.workspaceIds); }), + lsp: { + listDiagnostics: t + .input(schemas.workspace.lsp.listDiagnostics.input) + .output(schemas.workspace.lsp.listDiagnostics.output) + .handler(({ context, input }) => { + return context.lspManager.getWorkspaceDiagnosticsSnapshot(input.workspaceId); + }), + subscribeDiagnostics: t + .input(schemas.workspace.lsp.subscribeDiagnostics.input) + .output(schemas.workspace.lsp.subscribeDiagnostics.output) + .handler(async function* ({ context, input, signal }) { + const workspaceId = input.workspaceId; + + if (signal?.aborted) { + return; + } + + context.lspManager.acquireLease(workspaceId); + const queue = createAsyncEventQueue(); + const onAbort = () => { + queue.end(); + }; + + if (signal) { + signal.addEventListener("abort", onAbort, { once: true }); + } + + const unsubscribe = context.lspManager.subscribeWorkspaceDiagnostics( + workspaceId, + (snapshot) => { + queue.push(snapshot); + } + ); + + try { + yield context.lspManager.getWorkspaceDiagnosticsSnapshot(workspaceId); + yield* queue.iterate(); + } finally { + signal?.removeEventListener("abort", onAbort); + queue.end(); + unsubscribe(); + context.lspManager.releaseLease(workspaceId); + } + }), + }, stats: { subscribe: t .input(schemas.workspace.stats.subscribe.input) diff --git a/src/node/services/lsp/lspManager.test.ts b/src/node/services/lsp/lspManager.test.ts index e1ca08a448..48f295a1aa 100644 --- a/src/node/services/lsp/lspManager.test.ts +++ b/src/node/services/lsp/lspManager.test.ts @@ -2,6 +2,7 @@ import * as fs from "node:fs/promises"; import * as os from "node:os"; import * as path from "node:path"; import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test"; +import type { WorkspaceLspDiagnosticsSnapshot } from "@/common/orpc/types"; import { LocalRuntime } from "@/node/runtime/LocalRuntime"; import type { CreateLspClientOptions, @@ -68,7 +69,7 @@ describe("LspManager", () => { }); test("reuses clients for the same workspace root and forwards normalized positions", async () => { - const ensureFile = mock(async () => 1); + const ensureFile = mock(() => Promise.resolve(1)); let lastQueryRequest: Parameters[0] | undefined; const query = mock((request: Parameters[0]) => { lastQueryRequest = request; @@ -137,7 +138,7 @@ describe("LspManager", () => { }); test("deduplicates concurrent client creation for the same workspace root", async () => { - const ensureFile = mock(async () => 1); + const ensureFile = mock(() => Promise.resolve(1)); const query = mock(() => Promise.resolve({ operation: "hover" as const, @@ -331,7 +332,7 @@ describe("LspManager", () => { const publishByUri = new Map void>(); let secondEnsureResolved = false; const clientFactory = mock((options: CreateLspClientOptions): Promise => { - const ensureFile = mock(async (file: Parameters[0]) => { + const ensureFile = mock((file: Parameters[0]) => { publishByUri.set(file.uri, () => { options.onPublishDiagnostics?.({ uri: file.uri, @@ -344,7 +345,7 @@ describe("LspManager", () => { secondEnsureResolved = true; secondEnsureStarted.resolve(); } - return 1; + return Promise.resolve(1); }); return Promise.resolve({ @@ -433,4 +434,128 @@ describe("LspManager", () => { path.join(workspacePath, "src", "example.ts"), ]); }); + + test("returns cloned sorted snapshots and notifies listeners on publish and clear", async () => { + const clientFactory = mock((options: CreateLspClientOptions): Promise => { + const ensureFile = mock((file: Parameters[0]) => { + options.onPublishDiagnostics?.({ + uri: file.uri, + version: 1, + diagnostics: [createDiagnostic(`diagnostic for ${file.runtimePath}`)], + rawDiagnosticCount: 1, + }); + return Promise.resolve(1); + }); + + return Promise.resolve({ + isClosed: false, + ensureFile, + query: mock(() => Promise.resolve({ operation: "hover" as const, hover: "unused" })), + close: mock(() => Promise.resolve(undefined)), + }); + }); + + const manager = new LspManager({ + registry: createRegistry(), + clientFactory, + }); + const runtime = new LocalRuntime(workspacePath); + const snapshots: WorkspaceLspDiagnosticsSnapshot[] = []; + const unsubscribe = manager.subscribeWorkspaceDiagnostics("ws-1", (snapshot) => { + snapshots.push(snapshot); + }); + + await manager.collectPostMutationDiagnostics({ + workspaceId: "ws-1", + runtime, + workspacePath, + filePaths: ["src/example.ts", "packages/pkg/src/nested.ts"], + timeoutMs: 20, + }); + + const snapshot = manager.getWorkspaceDiagnosticsSnapshot("ws-1"); + expect(snapshot.diagnostics.map((entry) => entry.path)).toEqual([ + path.join(workspacePath, "packages", "pkg", "src", "nested.ts"), + path.join(workspacePath, "src", "example.ts"), + ]); + const firstDiagnostic = snapshot.diagnostics[0]; + expect(firstDiagnostic).toBeDefined(); + if (!firstDiagnostic) { + throw new Error("Expected the first diagnostics entry to exist"); + } + const firstMessage = firstDiagnostic.diagnostics[0]; + expect(firstMessage).toBeDefined(); + if (!firstMessage) { + throw new Error("Expected the first diagnostic message to exist"); + } + firstMessage.message = "mutated"; + const refreshedSnapshot = manager.getWorkspaceDiagnosticsSnapshot("ws-1"); + expect(refreshedSnapshot.diagnostics[0]?.diagnostics[0]?.message).toBe( + `diagnostic for ${path.join(workspacePath, "packages", "pkg", "src", "nested.ts")}` + ); + + expect(snapshots).toHaveLength(2); + expect(snapshots[0]?.diagnostics).toHaveLength(1); + expect(snapshots[1]?.diagnostics).toHaveLength(2); + + await manager.disposeWorkspace("ws-1"); + expect(manager.getWorkspaceDiagnosticsSnapshot("ws-1").diagnostics).toEqual([]); + expect(snapshots.at(-1)?.diagnostics).toEqual([]); + + unsubscribe(); + }); + + test("ignores malformed publishes in workspace snapshots", async () => { + const ensureFile = mock((file: Parameters[0]) => { + const version = ensureFile.mock.calls.length; + clientOptions?.onPublishDiagnostics?.({ + uri: file.uri, + version, + diagnostics: version === 1 ? [createDiagnostic("first pass")] : [], + rawDiagnosticCount: 1, + }); + return Promise.resolve(version); + }); + const close = mock(() => Promise.resolve(undefined)); + let clientOptions: CreateLspClientOptions | undefined; + const clientFactory = mock((options: CreateLspClientOptions): Promise => { + clientOptions = options; + return Promise.resolve({ + isClosed: false, + ensureFile, + query: mock(() => Promise.resolve({ operation: "hover" as const, hover: "unused" })), + close, + }); + }); + + const manager = new LspManager({ + registry: createRegistry(), + clientFactory, + }); + const runtime = new LocalRuntime(workspacePath); + + await manager.collectPostMutationDiagnostics({ + workspaceId: "ws-1", + runtime, + workspacePath, + filePaths: ["src/example.ts"], + timeoutMs: 20, + }); + + await manager.collectPostMutationDiagnostics({ + workspaceId: "ws-1", + runtime, + workspacePath, + filePaths: ["src/example.ts"], + timeoutMs: 20, + }); + + expect(manager.getWorkspaceDiagnosticsSnapshot("ws-1").diagnostics).toHaveLength(1); + expect( + manager.getWorkspaceDiagnosticsSnapshot("ws-1").diagnostics[0]?.diagnostics[0]?.message + ).toBe("first pass"); + + await manager.disposeWorkspace("ws-1"); + expect(close).toHaveBeenCalledTimes(1); + }); }); diff --git a/src/node/services/lsp/lspManager.ts b/src/node/services/lsp/lspManager.ts index f41b56dd90..d5ed55bd53 100644 --- a/src/node/services/lsp/lspManager.ts +++ b/src/node/services/lsp/lspManager.ts @@ -9,6 +9,7 @@ import { LSP_PREVIEW_CONTEXT_LINES, } from "@/constants/lsp"; import type { Runtime } from "@/node/runtime/Runtime"; +import type { WorkspaceLspDiagnosticsSnapshot } from "@/common/orpc/types"; import { log } from "@/node/services/log"; import { LspClient } from "./lspClient"; import { LspPathMapper } from "./lspPathMapper"; @@ -50,6 +51,7 @@ interface LspDiagnosticPublishReceipt { } type LspDiagnosticWaiter = (publish: LspDiagnosticPublishReceipt) => void; +type WorkspaceDiagnosticsListener = (snapshot: WorkspaceLspDiagnosticsSnapshot) => void; type LspClientFactory = (options: CreateLspClientOptions) => Promise; export interface LspManagerOptions { @@ -94,6 +96,7 @@ export class LspManager { string, Map>> >(); + private readonly workspaceDiagnosticListeners = new Map>(); private readonly registry: readonly LspServerDescriptor[]; private readonly clientFactory: LspClientFactory; private readonly idleTimeoutMs: number; @@ -128,6 +131,35 @@ export class LspManager { this.workspaceLeases.set(workspaceId, currentLeaseCount - 1); } + getWorkspaceDiagnosticsSnapshot(workspaceId: string): WorkspaceLspDiagnosticsSnapshot { + const diagnostics = Array.from(this.workspaceDiagnostics.get(workspaceId)?.values() ?? []) + .flatMap((diagnosticsForClient) => [...diagnosticsForClient.values()]) + .sort(compareFileDiagnostics) + .map(cloneFileDiagnostics); + + return { + workspaceId, + diagnostics, + }; + } + + subscribeWorkspaceDiagnostics( + workspaceId: string, + listener: WorkspaceDiagnosticsListener + ): () => void { + const listeners = + this.workspaceDiagnosticListeners.get(workspaceId) ?? new Set(); + listeners.add(listener); + this.workspaceDiagnosticListeners.set(workspaceId, listeners); + + return () => { + listeners.delete(listener); + if (listeners.size === 0) { + this.workspaceDiagnosticListeners.delete(workspaceId); + } + }; + } + async query(options: LspManagerQueryOptions): Promise { this.acquireLease(options.workspaceId); @@ -284,8 +316,13 @@ export class LspManager { async dispose(): Promise { clearInterval(this.idleInterval); - const workspaceIds = [...this.workspaceClients.keys()]; - await Promise.all(workspaceIds.map(async (workspaceId) => this.disposeWorkspace(workspaceId))); + const workspaceIds = new Set([ + ...this.workspaceClients.keys(), + ...this.workspaceDiagnostics.keys(), + ...this.workspaceDiagnosticListeners.keys(), + ]); + await Promise.all([...workspaceIds].map(async (workspaceId) => this.disposeWorkspace(workspaceId))); + this.workspaceDiagnosticListeners.clear(); } private async getOrCreateClient( @@ -533,13 +570,15 @@ export class LspManager { context.params.uri, publishReceipt ); + this.notifyWorkspaceDiagnosticsListeners(context.workspaceId); } private getOrCreateWorkspaceDiagnostics( workspaceId: string, clientKey: string ): Map { - const workspaceDiagnostics = this.workspaceDiagnostics.get(workspaceId) ?? new Map(); + const workspaceDiagnostics = + this.workspaceDiagnostics.get(workspaceId) ?? new Map>(); this.workspaceDiagnostics.set(workspaceId, workspaceDiagnostics); const diagnosticsForClient = workspaceDiagnostics.get(clientKey) ?? new Map(); @@ -551,7 +590,9 @@ export class LspManager { workspaceId: string, clientKey: string ): Map { - const workspacePublishes = this.workspaceDiagnosticPublishes.get(workspaceId) ?? new Map(); + const workspacePublishes = + this.workspaceDiagnosticPublishes.get(workspaceId) ?? + new Map>(); this.workspaceDiagnosticPublishes.set(workspaceId, workspacePublishes); const publishesForClient = workspacePublishes.get(clientKey) ?? new Map(); @@ -645,7 +686,9 @@ export class LspManager { clientKey: string, uri: string ): Set { - const workspaceWaiters = this.diagnosticWaiters.get(workspaceId) ?? new Map(); + const workspaceWaiters = + this.diagnosticWaiters.get(workspaceId) ?? + new Map>>(); this.diagnosticWaiters.set(workspaceId, workspaceWaiters); const clientWaiters = workspaceWaiters.get(clientKey) ?? new Map>(); @@ -666,7 +709,7 @@ export class LspManager { if (clientWaiters.size === 0) { workspaceWaiters?.delete(clientKey); } - if (workspaceWaiters && workspaceWaiters.size === 0) { + if (workspaceWaiters?.size === 0) { this.diagnosticWaiters.delete(workspaceId); } } @@ -687,10 +730,23 @@ export class LspManager { } } + private notifyWorkspaceDiagnosticsListeners(workspaceId: string): void { + const listeners = this.workspaceDiagnosticListeners.get(workspaceId); + if (!listeners || listeners.size === 0) { + return; + } + + const snapshot = this.getWorkspaceDiagnosticsSnapshot(workspaceId); + for (const listener of [...listeners]) { + listener(snapshot); + } + } + private clearWorkspaceDiagnostics(workspaceId: string): void { this.workspaceDiagnostics.delete(workspaceId); this.workspaceDiagnosticPublishes.delete(workspaceId); this.diagnosticWaiters.delete(workspaceId); + this.notifyWorkspaceDiagnosticsListeners(workspaceId); } private getClientKey(serverId: string, rootUri: string): string { @@ -843,6 +899,37 @@ function selectPathModule(filePath: string): PathModule { return path.posix; } +function compareFileDiagnostics(left: LspFileDiagnostics, right: LspFileDiagnostics): number { + return ( + left.path.localeCompare(right.path) || + left.serverId.localeCompare(right.serverId) || + left.rootUri.localeCompare(right.rootUri) || + left.uri.localeCompare(right.uri) + ); +} + +function cloneFileDiagnostics(diagnostics: LspFileDiagnostics): LspFileDiagnostics { + return { + uri: diagnostics.uri, + path: diagnostics.path, + serverId: diagnostics.serverId, + rootUri: diagnostics.rootUri, + version: diagnostics.version, + diagnostics: diagnostics.diagnostics.map((diagnostic) => ({ + ...diagnostic, + range: cloneRange(diagnostic.range), + })), + receivedAtMs: diagnostics.receivedAtMs, + }; +} + +function cloneRange(range: LspRange): LspRange { + return { + start: { ...range.start }, + end: { ...range.end }, + }; +} + function isFreshDiagnosticPublish( publish: LspDiagnosticPublishReceipt | undefined, previousReceivedAtMs: number | undefined, diff --git a/src/node/services/serviceContainer.test.ts b/src/node/services/serviceContainer.test.ts index 860113e437..ddaa4fe251 100644 --- a/src/node/services/serviceContainer.test.ts +++ b/src/node/services/serviceContainer.test.ts @@ -194,12 +194,13 @@ describe("ServiceContainer", () => { ).toBe(true); }); - it("exposes desktopSessionManager in the ORPC context", () => { + it("exposes desktopSessionManager and lspManager in the ORPC context", () => { services = new ServiceContainer(config); const context = services.toORPCContext(); expect(context.desktopSessionManager).toBe(services.desktopSessionManager); + expect(context.lspManager).toBe(services.lspManager); }); it("closes desktop sessions during shutdown", async () => { diff --git a/src/node/services/serviceContainer.ts b/src/node/services/serviceContainer.ts index 5457066a40..3fad2ed147 100644 --- a/src/node/services/serviceContainer.ts +++ b/src/node/services/serviceContainer.ts @@ -631,6 +631,7 @@ export class ServiceContainer { telemetryService: this.telemetryService, experimentsService: this.experimentsService, sessionUsageService: this.sessionUsageService, + lspManager: this.lspManager, devToolsService: this.devToolsService, browserSessionDiscoveryService: this.browserSessionDiscoveryService, browserBridgeTokenManager: this.browserBridgeTokenManager, From 7248bd6976b9419c94385e25d8a11355a0aae47a Mon Sep 17 00:00:00 2001 From: terion Date: Mon, 13 Apr 2026 15:17:59 +0300 Subject: [PATCH 07/46] =?UTF-8?q?=F0=9F=A4=96=20fix:=20harden=20LSP=20diag?= =?UTF-8?q?nostics=20recovery?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reconnect active workspace diagnostics subscriptions after client swaps, retry diagnostics streams when they end unexpectedly, and settle pending LSP diagnostic waits during workspace disposal. --- src/browser/stores/WorkspaceStore.test.ts | 212 +++++++++++++++++++++- src/browser/stores/WorkspaceStore.ts | 90 +++++++-- src/node/services/lsp/lspManager.test.ts | 75 ++++++++ src/node/services/lsp/lspManager.ts | 34 +++- 4 files changed, 382 insertions(+), 29 deletions(-) diff --git a/src/browser/stores/WorkspaceStore.test.ts b/src/browser/stores/WorkspaceStore.test.ts index 9e7941b2f2..3338594919 100644 --- a/src/browser/stores/WorkspaceStore.test.ts +++ b/src/browser/stores/WorkspaceStore.test.ts @@ -200,7 +200,9 @@ const mockListLspDiagnostics = mock(({ workspaceId }: { workspaceId: string }) = ); const mockSubscribeLspDiagnostics = mock( ({ workspaceId }: { workspaceId: string }, options?: { signal?: AbortSignal }) => { - const controlled = createControlledAsyncIterator(options?.signal); + const controlled = createControlledAsyncIterator( + options?.signal + ); lspDiagnosticsSubscriptions.push({ workspaceId, signal: options?.signal, @@ -4605,7 +4607,6 @@ describe("WorkspaceStore", () => { }); }); - describe("workspace LSP diagnostics snapshots", () => { it("starts empty and replaces the cached snapshot after the first payload", async () => { const workspaceId = "lsp-diagnostics-snapshot"; @@ -4678,11 +4679,211 @@ describe("WorkspaceStore", () => { expect(store.getWorkspaceLspDiagnosticsSnapshot(workspaceId)?.workspaceId).toBe(workspaceId); unsubscribeTwo(); - const aborted = await waitUntil(() => lspDiagnosticsSubscriptions[0]?.signal?.aborted === true); + const aborted = await waitUntil( + () => lspDiagnosticsSubscriptions[0]?.signal?.aborted === true + ); expect(aborted).toBe(true); expect(store.getWorkspaceLspDiagnosticsSnapshot(workspaceId)).toBeNull(); }); + it("re-subscribes active diagnostics listeners after the client reconnects", async () => { + const workspaceId = "lsp-diagnostics-reconnect"; + createAndAddWorkspace(store, workspaceId); + + const unsubscribe = store.subscribeLspDiagnostics(workspaceId, () => undefined); + const subscribed = await waitUntil(() => mockSubscribeLspDiagnostics.mock.calls.length === 1); + expect(subscribed).toBe(true); + + emitLspDiagnosticsSnapshot({ + workspaceId, + diagnostics: [ + { + uri: "file:///workspace/src/example.ts", + path: "/workspace/src/example.ts", + serverId: "typescript", + rootUri: "file:///workspace", + version: 1, + diagnostics: [ + { + range: { + start: { line: 1, character: 1 }, + end: { line: 1, character: 5 }, + }, + severity: 1, + source: "tsserver", + message: "before reconnect", + }, + ], + receivedAtMs: 1, + }, + ], + }); + const firstReceived = await waitUntil( + () => + store.getWorkspaceLspDiagnosticsSnapshot(workspaceId)?.diagnostics[0]?.diagnostics[0] + ?.message === "before reconnect" + ); + expect(firstReceived).toBe(true); + + store.setClient(null); + const oldSignal = lspDiagnosticsSubscriptions[0]?.signal; + const oldAborted = await waitUntil(() => oldSignal?.aborted === true); + expect(oldAborted).toBe(true); + + store.setClient(mockClient as any); + const resubscribed = await waitUntil( + () => mockSubscribeLspDiagnostics.mock.calls.length === 2 + ); + expect(resubscribed).toBe(true); + + emitLspDiagnosticsSnapshot({ + workspaceId, + diagnostics: [ + { + uri: "file:///workspace/src/example.ts", + path: "/workspace/src/example.ts", + serverId: "typescript", + rootUri: "file:///workspace", + version: 2, + diagnostics: [ + { + range: { + start: { line: 2, character: 1 }, + end: { line: 2, character: 5 }, + }, + severity: 1, + source: "tsserver", + message: "after reconnect", + }, + ], + receivedAtMs: 2, + }, + ], + }); + const secondReceived = await waitUntil( + () => + store.getWorkspaceLspDiagnosticsSnapshot(workspaceId)?.diagnostics[0]?.diagnostics[0] + ?.message === "after reconnect" + ); + expect(secondReceived).toBe(true); + + unsubscribe(); + }); + + it("retries LSP diagnostics subscriptions after an unexpected stream end", async () => { + store.setClient(null); + + const customSubscriptions: Array<{ + signal?: AbortSignal; + push: (snapshot: WorkspaceLspDiagnosticsSnapshot) => void; + end: () => void; + }> = []; + const customSubscribeDiagnostics = mock( + (_input: { workspaceId: string }, options?: { signal?: AbortSignal }) => { + const controlled = createControlledAsyncIterator( + options?.signal + ); + customSubscriptions.push({ + signal: options?.signal, + push: controlled.push, + end: controlled.end, + }); + return controlled.iterator; + } + ); + const customClient = { + workspace: { + ...mockClient.workspace, + lsp: { + ...mockClient.workspace.lsp, + subscribeDiagnostics: customSubscribeDiagnostics, + }, + }, + terminal: mockClient.terminal, + }; + store.setClient(customClient as any); + + const workspaceId = "lsp-diagnostics-retry"; + createAndAddWorkspace(store, workspaceId); + + const unsubscribe = store.subscribeLspDiagnostics(workspaceId, () => undefined); + const subscribed = await waitUntil(() => customSubscribeDiagnostics.mock.calls.length === 1); + expect(subscribed).toBe(true); + + customSubscriptions[0]?.push({ + workspaceId, + diagnostics: [ + { + uri: "file:///workspace/src/example.ts", + path: "/workspace/src/example.ts", + serverId: "typescript", + rootUri: "file:///workspace", + version: 1, + diagnostics: [ + { + range: { + start: { line: 1, character: 1 }, + end: { line: 1, character: 5 }, + }, + severity: 1, + source: "tsserver", + message: "before retry", + }, + ], + receivedAtMs: 1, + }, + ], + }); + const firstReceived = await waitUntil( + () => + store.getWorkspaceLspDiagnosticsSnapshot(workspaceId)?.diagnostics[0]?.diagnostics[0] + ?.message === "before retry" + ); + expect(firstReceived).toBe(true); + + customSubscriptions[0]?.end(); + const retried = await waitUntil( + () => customSubscribeDiagnostics.mock.calls.length === 2, + 4000 + ); + expect(retried).toBe(true); + expect(customSubscriptions[1]).toBeDefined(); + + customSubscriptions[1]?.push({ + workspaceId, + diagnostics: [ + { + uri: "file:///workspace/src/example.ts", + path: "/workspace/src/example.ts", + serverId: "typescript", + rootUri: "file:///workspace", + version: 2, + diagnostics: [ + { + range: { + start: { line: 2, character: 1 }, + end: { line: 2, character: 5 }, + }, + severity: 1, + source: "tsserver", + message: "after retry", + }, + ], + receivedAtMs: 2, + }, + ], + }); + const secondReceived = await waitUntil( + () => + store.getWorkspaceLspDiagnosticsSnapshot(workspaceId)?.diagnostics[0]?.diagnostics[0] + ?.message === "after retry" + ); + expect(secondReceived).toBe(true); + expect(customSubscriptions[1]?.signal?.aborted).toBe(false); + + unsubscribe(); + }); + it("aborts LSP diagnostics subscriptions when a workspace is removed", async () => { const workspaceId = "lsp-diagnostics-remove"; createAndAddWorkspace(store, workspaceId); @@ -4698,10 +4899,11 @@ describe("WorkspaceStore", () => { store.removeWorkspace(workspaceId); - const aborted = await waitUntil(() => lspDiagnosticsSubscriptions[0]?.signal?.aborted === true); + const aborted = await waitUntil( + () => lspDiagnosticsSubscriptions[0]?.signal?.aborted === true + ); expect(aborted).toBe(true); expect(store.getWorkspaceLspDiagnosticsSnapshot(workspaceId)).toBeNull(); }); }); - }); diff --git a/src/browser/stores/WorkspaceStore.ts b/src/browser/stores/WorkspaceStore.ts index 9d54e28ec4..15328eaaf3 100644 --- a/src/browser/stores/WorkspaceStore.ts +++ b/src/browser/stores/WorkspaceStore.ts @@ -1062,6 +1062,9 @@ export class WorkspaceStore { for (const workspaceId of this.statsListenerCounts.keys()) { this.subscribeToStats(workspaceId); } + for (const workspaceId of this.lspDiagnosticsListenerCounts.keys()) { + this.subscribeToLspDiagnostics(workspaceId); + } this.ensureActiveOnChatSubscription(); void this.refreshProvidersConfig(client); @@ -1345,8 +1348,22 @@ export class WorkspaceStore { }); } + private canContinueLspDiagnosticsSubscription( + workspaceId: string, + client: RouterClient, + signal: AbortSignal + ): boolean { + return ( + !signal.aborted && + this.client === client && + this.isWorkspaceRegistered(workspaceId) && + (this.lspDiagnosticsListenerCounts.get(workspaceId) ?? 0) > 0 + ); + } + private subscribeToLspDiagnostics(workspaceId: string): void { - if (!this.client) { + const client = this.client; + if (!client) { return; } @@ -1365,29 +1382,62 @@ export class WorkspaceStore { let iterator: AsyncIterator | null = null; (async () => { - try { - const subscribedIterator = await this.client!.workspace.lsp.subscribeDiagnostics( - { workspaceId }, - { signal } - ); - iterator = subscribedIterator; + let attempt = 0; - for await (const snapshot of subscribedIterator) { - if (signal.aborted) break; - queueMicrotask(() => { - if (signal.aborted) { + while (this.canContinueLspDiagnosticsSubscription(workspaceId, client, signal)) { + try { + const subscribedIterator = await client.workspace.lsp.subscribeDiagnostics( + { workspaceId }, + { signal } + ); + + if (!this.canContinueLspDiagnosticsSubscription(workspaceId, client, signal)) { + void subscribedIterator.return?.(); + return; + } + + iterator = subscribedIterator; + + for await (const snapshot of subscribedIterator) { + if (!this.canContinueLspDiagnosticsSubscription(workspaceId, client, signal)) { return; } - this.workspaceLspDiagnostics.set(workspaceId, snapshot); - this.lspDiagnosticsStore.bump(workspaceId); - }); + + attempt = 0; + queueMicrotask(() => { + if (!this.canContinueLspDiagnosticsSubscription(workspaceId, client, signal)) { + return; + } + this.workspaceLspDiagnostics.set(workspaceId, snapshot); + this.lspDiagnosticsStore.bump(workspaceId); + }); + } + + if (!this.canContinueLspDiagnosticsSubscription(workspaceId, client, signal)) { + return; + } + + console.warn( + `[WorkspaceStore] LSP diagnostics subscription ended unexpectedly for ${workspaceId}; retrying...` + ); + } catch (error) { + if (!this.canContinueLspDiagnosticsSubscription(workspaceId, client, signal)) { + return; + } + if (!isAbortError(error)) { + console.warn( + `[WorkspaceStore] Error in LSP diagnostics subscription for ${workspaceId}:`, + error + ); + } + } finally { + void iterator?.return?.(); + iterator = null; } - } catch (error) { - if (signal.aborted || isAbortError(error)) return; - console.warn( - `[WorkspaceStore] Error in LSP diagnostics subscription for ${workspaceId}:`, - error - ); + + const delayMs = calculateSubscriptionBackoffMs(attempt); + attempt++; + await this.sleepWithAbort(delayMs, signal); } })(); diff --git a/src/node/services/lsp/lspManager.test.ts b/src/node/services/lsp/lspManager.test.ts index 48f295a1aa..a8cb9d4ae2 100644 --- a/src/node/services/lsp/lspManager.test.ts +++ b/src/node/services/lsp/lspManager.test.ts @@ -22,6 +22,17 @@ function createDeferred() { return { promise, resolve, reject }; } +async function waitUntil(condition: () => boolean, timeoutMs = 2000): Promise { + const start = Date.now(); + while (Date.now() - start < timeoutMs) { + if (condition()) { + return true; + } + await new Promise((resolve) => setTimeout(resolve, 10)); + } + return false; +} + function createRegistry(): readonly LspServerDescriptor[] { return [ { @@ -505,6 +516,70 @@ describe("LspManager", () => { unsubscribe(); }); + test("disposeWorkspace settles pending diagnostics waits promptly", async () => { + const ensureFile = mock(() => Promise.resolve(1)); + const close = mock(() => Promise.resolve(undefined)); + const clientFactory = mock((_options: CreateLspClientOptions): Promise => { + return Promise.resolve({ + isClosed: false, + ensureFile, + query: mock(() => Promise.resolve({ operation: "hover" as const, hover: "unused" })), + close, + }); + }); + + const manager = new LspManager({ + registry: createRegistry(), + clientFactory, + }); + const runtime = new LocalRuntime(workspacePath); + + const diagnosticsPromise = manager.collectPostMutationDiagnostics({ + workspaceId: "ws-1", + runtime, + workspacePath, + filePaths: ["src/example.ts"], + timeoutMs: 10_000, + }); + + const diagnosticWaiters = ( + manager as unknown as { + diagnosticWaiters: Map>>>; + } + ).diagnosticWaiters; + const waiterRegistered = await waitUntil(() => { + const workspaceWaiters = diagnosticWaiters.get("ws-1"); + if (!workspaceWaiters) { + return false; + } + for (const clientWaiters of workspaceWaiters.values()) { + for (const waiters of clientWaiters.values()) { + if (waiters.size > 0) { + return true; + } + } + } + return false; + }); + expect(waiterRegistered).toBe(true); + + await manager.disposeWorkspace("ws-1"); + + const completion = await Promise.race([ + diagnosticsPromise.then((diagnostics) => ({ type: "resolved" as const, diagnostics })), + new Promise<{ type: "timeout" }>((resolve) => { + setTimeout(() => resolve({ type: "timeout" }), 200); + }), + ]); + expect(completion.type).toBe("resolved"); + if (completion.type !== "resolved") { + throw new Error("Expected diagnostics wait to resolve after workspace disposal"); + } + expect(completion.diagnostics).toEqual([]); + expect(diagnosticWaiters.has("ws-1")).toBe(false); + expect(close).toHaveBeenCalledTimes(1); + }); + test("ignores malformed publishes in workspace snapshots", async () => { const ensureFile = mock((file: Parameters[0]) => { const version = ensureFile.mock.calls.length; diff --git a/src/node/services/lsp/lspManager.ts b/src/node/services/lsp/lspManager.ts index d5ed55bd53..ceccdb5e01 100644 --- a/src/node/services/lsp/lspManager.ts +++ b/src/node/services/lsp/lspManager.ts @@ -50,7 +50,7 @@ interface LspDiagnosticPublishReceipt { receivedAtMs: number; } -type LspDiagnosticWaiter = (publish: LspDiagnosticPublishReceipt) => void; +type LspDiagnosticWaiter = (publish?: LspDiagnosticPublishReceipt) => void; type WorkspaceDiagnosticsListener = (snapshot: WorkspaceLspDiagnosticsSnapshot) => void; type LspClientFactory = (options: CreateLspClientOptions) => Promise; @@ -96,7 +96,10 @@ export class LspManager { string, Map>> >(); - private readonly workspaceDiagnosticListeners = new Map>(); + private readonly workspaceDiagnosticListeners = new Map< + string, + Set + >(); private readonly registry: readonly LspServerDescriptor[]; private readonly clientFactory: LspClientFactory; private readonly idleTimeoutMs: number; @@ -321,7 +324,9 @@ export class LspManager { ...this.workspaceDiagnostics.keys(), ...this.workspaceDiagnosticListeners.keys(), ]); - await Promise.all([...workspaceIds].map(async (workspaceId) => this.disposeWorkspace(workspaceId))); + await Promise.all( + [...workspaceIds].map(async (workspaceId) => this.disposeWorkspace(workspaceId)) + ); this.workspaceDiagnosticListeners.clear(); } @@ -578,7 +583,8 @@ export class LspManager { clientKey: string ): Map { const workspaceDiagnostics = - this.workspaceDiagnostics.get(workspaceId) ?? new Map>(); + this.workspaceDiagnostics.get(workspaceId) ?? + new Map>(); this.workspaceDiagnostics.set(workspaceId, workspaceDiagnostics); const diagnosticsForClient = workspaceDiagnostics.get(clientKey) ?? new Map(); @@ -657,6 +663,10 @@ export class LspManager { }; const onPublish: LspDiagnosticWaiter = (publish) => { + if (!publish) { + finish(undefined); + return; + } if ( !isFreshDiagnosticPublish(publish, params.previousReceivedAtMs, params.expectedVersion) ) { @@ -714,6 +724,21 @@ export class LspManager { } } + private settleWorkspaceDiagnosticWaiters(workspaceId: string): void { + const workspaceWaiters = this.diagnosticWaiters.get(workspaceId); + if (!workspaceWaiters) { + return; + } + + for (const clientWaiters of workspaceWaiters.values()) { + for (const waiters of clientWaiters.values()) { + for (const waiter of [...waiters]) { + waiter(undefined); + } + } + } + } + private notifyDiagnosticWaiters( workspaceId: string, clientKey: string, @@ -743,6 +768,7 @@ export class LspManager { } private clearWorkspaceDiagnostics(workspaceId: string): void { + this.settleWorkspaceDiagnosticWaiters(workspaceId); this.workspaceDiagnostics.delete(workspaceId); this.workspaceDiagnosticPublishes.delete(workspaceId); this.diagnosticWaiters.delete(workspaceId); From 7d52175d4b10a7a38d55b6c7c3ce87d79be431d3 Mon Sep 17 00:00:00 2001 From: terion Date: Mon, 13 Apr 2026 15:31:33 +0300 Subject: [PATCH 08/46] fix lsp disposal diagnostics races --- src/node/services/lsp/lspManager.test.ts | 122 +++++++++++++++++++++++ src/node/services/lsp/lspManager.ts | 42 +++++++- 2 files changed, 161 insertions(+), 3 deletions(-) diff --git a/src/node/services/lsp/lspManager.test.ts b/src/node/services/lsp/lspManager.test.ts index a8cb9d4ae2..7aff182355 100644 --- a/src/node/services/lsp/lspManager.test.ts +++ b/src/node/services/lsp/lspManager.test.ts @@ -516,6 +516,128 @@ describe("LspManager", () => { unsubscribe(); }); + test("ignores late diagnostics publishes from disposed clients", async () => { + let clientOptions: CreateLspClientOptions | undefined; + let ensuredFile: Parameters[0] | undefined; + const ensureFile = mock((file: Parameters[0]) => { + ensuredFile = file; + clientOptions?.onPublishDiagnostics?.({ + uri: file.uri, + version: 1, + diagnostics: [createDiagnostic("first pass")], + rawDiagnosticCount: 1, + }); + return Promise.resolve(1); + }); + const close = mock(() => Promise.resolve(undefined)); + const clientFactory = mock((options: CreateLspClientOptions): Promise => { + clientOptions = options; + return Promise.resolve({ + isClosed: false, + ensureFile, + query: mock(() => Promise.resolve({ operation: "hover" as const, hover: "unused" })), + close, + }); + }); + + const manager = new LspManager({ + registry: createRegistry(), + clientFactory, + }); + const runtime = new LocalRuntime(workspacePath); + + const diagnostics = await manager.collectPostMutationDiagnostics({ + workspaceId: "ws-1", + runtime, + workspacePath, + filePaths: ["src/example.ts"], + timeoutMs: 20, + }); + expect(diagnostics).toHaveLength(1); + + await manager.disposeWorkspace("ws-1"); + expect(manager.getWorkspaceDiagnosticsSnapshot("ws-1").diagnostics).toEqual([]); + + if (!clientOptions || !ensuredFile) { + throw new Error("Expected the LSP client to publish diagnostics for the test file"); + } + + clientOptions.onPublishDiagnostics?.({ + uri: ensuredFile.uri, + version: 2, + diagnostics: [createDiagnostic("late publish")], + rawDiagnosticCount: 1, + }); + + expect(manager.getWorkspaceDiagnosticsSnapshot("ws-1").diagnostics).toEqual([]); + const workspaceDiagnostics = ( + manager as unknown as { + workspaceDiagnostics: Map>>; + } + ).workspaceDiagnostics; + expect(workspaceDiagnostics.has("ws-1")).toBe(false); + expect(close).toHaveBeenCalledTimes(1); + }); + + test("collectPostMutationDiagnostics exits promptly when disposal wins before waiter registration", async () => { + const ensureFile = mock(() => Promise.resolve(1)); + const close = mock(() => Promise.resolve(undefined)); + const clientFactory = mock((_options: CreateLspClientOptions): Promise => { + return Promise.resolve({ + isClosed: false, + ensureFile, + query: mock(() => Promise.resolve({ operation: "hover" as const, hover: "unused" })), + close, + }); + }); + + const manager = new LspManager({ + registry: createRegistry(), + clientFactory, + }); + const runtime = new LocalRuntime(workspacePath); + const managerWithInternals = manager as unknown as { + diagnosticWaiters: Map>>>; + waitForFreshDiagnostics: (params: { + workspaceId: string; + workspaceGeneration: number; + clientKey: string; + uri: string; + previousReceivedAtMs?: number; + expectedVersion: number; + timeoutMs: number; + }) => Promise; + }; + const originalWaitForFreshDiagnostics = + managerWithInternals.waitForFreshDiagnostics.bind(managerWithInternals); + managerWithInternals.waitForFreshDiagnostics = async (params) => { + await manager.disposeWorkspace("ws-1"); + return await originalWaitForFreshDiagnostics(params); + }; + + const diagnosticsPromise = manager.collectPostMutationDiagnostics({ + workspaceId: "ws-1", + runtime, + workspacePath, + filePaths: ["src/example.ts"], + timeoutMs: 10_000, + }); + + const completion = await Promise.race([ + diagnosticsPromise.then((diagnostics) => ({ type: "resolved" as const, diagnostics })), + new Promise<{ type: "timeout" }>((resolve) => { + setTimeout(() => resolve({ type: "timeout" }), 200); + }), + ]); + expect(completion.type).toBe("resolved"); + if (completion.type !== "resolved") { + throw new Error("Expected diagnostics collection to resolve after workspace disposal"); + } + expect(completion.diagnostics).toEqual([]); + expect(managerWithInternals.diagnosticWaiters.has("ws-1")).toBe(false); + expect(close).toHaveBeenCalledTimes(1); + }); + test("disposeWorkspace settles pending diagnostics waits promptly", async () => { const ensureFile = mock(() => Promise.resolve(1)); const close = mock(() => Promise.resolve(undefined)); diff --git a/src/node/services/lsp/lspManager.ts b/src/node/services/lsp/lspManager.ts index ceccdb5e01..6ba541d09d 100644 --- a/src/node/services/lsp/lspManager.ts +++ b/src/node/services/lsp/lspManager.ts @@ -39,6 +39,7 @@ interface WorkspaceClients { interface ResolvedLspClientContext { client: LspClientInstance; clientKey: string; + workspaceGeneration: number; descriptor: LspServerDescriptor; fileHandle: LspClientFileHandle; pathMapper: LspPathMapper; @@ -83,6 +84,9 @@ export interface LspManagerCollectDiagnosticsOptions { export class LspManager { private readonly workspaceClients = new Map(); + // Disposal bumps the generation so late publishes from closing clients and stale + // post-mutation waits cannot recreate diagnostics for an already-cleared workspace. + private readonly workspaceGenerations = new Map(); private readonly workspaceLeases = new Map(); private readonly workspaceDiagnostics = new Map< string, @@ -242,6 +246,7 @@ export class LspManager { const expectedVersion = await context.client.ensureFile(context.fileHandle); const freshPublish = await this.waitForFreshDiagnostics({ workspaceId: options.workspaceId, + workspaceGeneration: context.workspaceGeneration, clientKey: context.clientKey, uri: context.fileHandle.uri, previousReceivedAtMs: previousPublish?.receivedAtMs, @@ -287,6 +292,8 @@ export class LspManager { } async disposeWorkspace(workspaceId: string): Promise { + this.workspaceGenerations.set(workspaceId, this.getWorkspaceGeneration(workspaceId) + 1); + const entry = this.workspaceClients.get(workspaceId); if (!entry) { this.clearWorkspaceDiagnostics(workspaceId); @@ -337,7 +344,7 @@ export class LspManager { pathMapper: LspPathMapper, rootPath: string, rootUri: string - ): Promise<{ client: LspClientInstance; clientKey: string }> { + ): Promise<{ client: LspClientInstance; clientKey: string; workspaceGeneration: number }> { const workspaceEntry = this.workspaceClients.get(workspaceId) ?? { clients: new Map(), pendingClients: new Map>(), @@ -346,10 +353,11 @@ export class LspManager { this.workspaceClients.set(workspaceId, workspaceEntry); const clientKey = this.getClientKey(descriptor.id, rootUri); + const workspaceGeneration = this.getWorkspaceGeneration(workspaceId); const existingClient = workspaceEntry.clients.get(clientKey); if (existingClient && !existingClient.isClosed) { workspaceEntry.lastActivity = Date.now(); - return { client: existingClient, clientKey }; + return { client: existingClient, clientKey, workspaceGeneration }; } if (existingClient?.isClosed) { @@ -361,6 +369,7 @@ export class LspManager { return { client: await pendingClient, clientKey, + workspaceGeneration, }; } @@ -374,6 +383,7 @@ export class LspManager { onPublishDiagnostics: (params) => this.handlePublishDiagnostics({ workspaceId, + workspaceGeneration, clientKey, serverId: descriptor.id, rootUri, @@ -413,6 +423,7 @@ export class LspManager { return { client: await clientPromise, clientKey, + workspaceGeneration, }; } @@ -463,6 +474,7 @@ export class LspManager { return { client: clientResult.client, clientKey: clientResult.clientKey, + workspaceGeneration: clientResult.workspaceGeneration, descriptor, fileHandle, pathMapper, @@ -479,6 +491,10 @@ export class LspManager { entry.lastActivity = Date.now(); } + private getWorkspaceGeneration(workspaceId: string): number { + return this.workspaceGenerations.get(workspaceId) ?? 0; + } + private async cleanupIdleWorkspaces(): Promise { const now = Date.now(); for (const [workspaceId, entry] of this.workspaceClients) { @@ -529,12 +545,17 @@ export class LspManager { private handlePublishDiagnostics(context: { workspaceId: string; + workspaceGeneration: number; clientKey: string; serverId: string; rootUri: string; pathMapper: LspPathMapper; params: LspPublishDiagnosticsParams; }): void { + if (context.workspaceGeneration !== this.getWorkspaceGeneration(context.workspaceId)) { + return; + } + if (isMalformedDiagnosticPublish(context.params)) { return; } @@ -624,12 +645,17 @@ export class LspManager { private async waitForFreshDiagnostics(params: { workspaceId: string; + workspaceGeneration: number; clientKey: string; uri: string; previousReceivedAtMs?: number; expectedVersion: number; timeoutMs: number; }): Promise { + if (params.workspaceGeneration !== this.getWorkspaceGeneration(params.workspaceId)) { + return undefined; + } + const existingPublish = this.getLatestDiagnosticPublish( params.workspaceId, params.clientKey, @@ -645,9 +671,14 @@ export class LspManager { let settled = false; const workspaceWaiters = this.getOrCreateDiagnosticWaiters( params.workspaceId, + params.workspaceGeneration, params.clientKey, params.uri ); + if (!workspaceWaiters) { + resolve(undefined); + return; + } const finish = (publish?: LspDiagnosticPublishReceipt) => { if (settled) { @@ -693,9 +724,14 @@ export class LspManager { private getOrCreateDiagnosticWaiters( workspaceId: string, + workspaceGeneration: number, clientKey: string, uri: string - ): Set { + ): Set | undefined { + if (workspaceGeneration !== this.getWorkspaceGeneration(workspaceId)) { + return undefined; + } + const workspaceWaiters = this.diagnosticWaiters.get(workspaceId) ?? new Map>>(); From b43cd801e97042035546738afba54aadfa07e930 Mon Sep 17 00:00:00 2001 From: terion Date: Mon, 13 Apr 2026 16:02:44 +0300 Subject: [PATCH 09/46] =?UTF-8?q?=F0=9F=A4=96=20feat:=20add=20stats=20diag?= =?UTF-8?q?nostics=20panel?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a Diagnostics sub-tab to the Stats sidebar and render grouped LSP diagnostics from the workspace store without subscribing while the tab is inactive. Add focused tests for the new diagnostics panel states and the Stats sub-tab list. --- .../StatsContainer.diagnostics.test.tsx | 39 +++++ .../features/RightSidebar/StatsContainer.tsx | 8 +- .../StatsTab.diagnostics.test.tsx | 127 +++++++++++++++ .../features/RightSidebar/StatsTab.tsx | 150 +++++++++++++++++- 4 files changed, 320 insertions(+), 4 deletions(-) create mode 100644 src/browser/features/RightSidebar/StatsContainer.diagnostics.test.tsx create mode 100644 src/browser/features/RightSidebar/StatsTab.diagnostics.test.tsx diff --git a/src/browser/features/RightSidebar/StatsContainer.diagnostics.test.tsx b/src/browser/features/RightSidebar/StatsContainer.diagnostics.test.tsx new file mode 100644 index 0000000000..2b8b7ee2c3 --- /dev/null +++ b/src/browser/features/RightSidebar/StatsContainer.diagnostics.test.tsx @@ -0,0 +1,39 @@ +import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test"; +import { GlobalWindow } from "happy-dom"; +import { cleanup, render } from "@testing-library/react"; + +void mock.module("./CostsTab", () => ({ + CostsTab: () =>
Costs panel
, +})); + +describe("StatsContainer diagnostics", () => { + let originalWindow: typeof globalThis.window; + let originalDocument: typeof globalThis.document; + + beforeEach(() => { + originalWindow = globalThis.window; + originalDocument = globalThis.document; + + globalThis.window = new GlobalWindow() as unknown as Window & typeof globalThis; + globalThis.document = globalThis.window.document; + globalThis.window.localStorage.clear(); + }); + + afterEach(() => { + cleanup(); + globalThis.window = originalWindow; + globalThis.document = originalDocument; + }); + + test("shows the Diagnostics sub-tab alongside the existing Stats tabs", async () => { + // eslint-disable-next-line no-restricted-syntax -- The test must import after mock.module() so StatsContainer sees the stubbed CostsTab instead of loading Mermaid-heavy dependencies. + const { StatsContainer } = await import("./StatsContainer"); + const view = render(); + + expect(view.getByText("Costs panel")).toBeTruthy(); + expect(view.getByRole("button", { name: "Cost" })).toBeTruthy(); + expect(view.getByRole("button", { name: "Timing" })).toBeTruthy(); + expect(view.getByRole("button", { name: "Models" })).toBeTruthy(); + expect(view.getByRole("button", { name: "Diagnostics" })).toBeTruthy(); + }); +}); diff --git a/src/browser/features/RightSidebar/StatsContainer.tsx b/src/browser/features/RightSidebar/StatsContainer.tsx index 8613e4f8f8..775d5d6a84 100644 --- a/src/browser/features/RightSidebar/StatsContainer.tsx +++ b/src/browser/features/RightSidebar/StatsContainer.tsx @@ -5,13 +5,14 @@ * - "Cost" — renders CostsTab * - "Timing" — renders TimingPanel from StatsTab * - "Models" — renders ModelBreakdownPanel from StatsTab + * - "Diagnostics" — renders DiagnosticsPanel from StatsTab */ import { usePersistedState } from "@/browser/hooks/usePersistedState"; import { CostsTab } from "./CostsTab"; -import { TimingPanel, ModelBreakdownPanel } from "./StatsTab"; +import { TimingPanel, ModelBreakdownPanel, DiagnosticsPanel } from "./StatsTab"; -type StatsSubTab = "cost" | "timing" | "models"; +type StatsSubTab = "cost" | "timing" | "models" | "diagnostics"; interface StatsOption { value: StatsSubTab; @@ -22,6 +23,7 @@ const OPTIONS: StatsOption[] = [ { value: "cost", label: "Cost" }, { value: "timing", label: "Timing" }, { value: "models", label: "Models" }, + { value: "diagnostics", label: "Diagnostics" }, ]; interface StatsContainerProps { @@ -59,6 +61,8 @@ export function StatsContainer(props: StatsContainerProps) { {effectiveTab === "cost" && } {effectiveTab === "timing" && } {effectiveTab === "models" && } + {/* Keep the LSP subscription inside the leaf panel so inactive Stats tabs stay unsubscribed. */} + {effectiveTab === "diagnostics" && } ); } diff --git a/src/browser/features/RightSidebar/StatsTab.diagnostics.test.tsx b/src/browser/features/RightSidebar/StatsTab.diagnostics.test.tsx new file mode 100644 index 0000000000..5a11a00c89 --- /dev/null +++ b/src/browser/features/RightSidebar/StatsTab.diagnostics.test.tsx @@ -0,0 +1,127 @@ +import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test"; +import { GlobalWindow } from "happy-dom"; +import { cleanup, render } from "@testing-library/react"; + +import type { WorkspaceLspDiagnosticsSnapshot } from "@/common/orpc/types"; + +let currentSnapshot: WorkspaceLspDiagnosticsSnapshot | null = null; + +void mock.module("@/browser/stores/WorkspaceStore", () => ({ + useWorkspaceStatsSnapshot: () => null, + useWorkspaceLspDiagnosticsSnapshot: () => currentSnapshot, +})); + +import { DiagnosticsPanel } from "./StatsTab"; + +describe("DiagnosticsPanel", () => { + let originalWindow: typeof globalThis.window; + let originalDocument: typeof globalThis.document; + + beforeEach(() => { + originalWindow = globalThis.window; + originalDocument = globalThis.document; + + globalThis.window = new GlobalWindow() as unknown as Window & typeof globalThis; + globalThis.document = globalThis.window.document; + currentSnapshot = null; + }); + + afterEach(() => { + cleanup(); + currentSnapshot = null; + globalThis.window = originalWindow; + globalThis.document = originalDocument; + }); + + test("shows a loading state before the first diagnostics snapshot arrives", () => { + const view = render(); + + expect(view.getByText("Loading diagnostics...")).toBeTruthy(); + }); + + test("shows an empty state when the diagnostics snapshot contains no entries", () => { + currentSnapshot = { + workspaceId: "workspace-1", + diagnostics: [], + }; + + const view = render(); + + expect(view.getByText("No diagnostics for this workspace.")).toBeTruthy(); + }); + + test("renders grouped diagnostics with severity labels and file metadata", () => { + currentSnapshot = { + workspaceId: "workspace-1", + diagnostics: [ + { + uri: "file:///workspace/src/example.ts", + path: "/workspace/src/example.ts", + serverId: "typescript", + rootUri: "file:///workspace", + version: 3, + diagnostics: [ + { + range: { + start: { line: 1, character: 4 }, + end: { line: 1, character: 10 }, + }, + severity: 1, + source: "tsserver", + code: 2322, + message: "Type 'string' is not assignable to type 'number'.", + }, + { + range: { + start: { line: 4, character: 2 }, + end: { line: 4, character: 8 }, + }, + severity: 2, + source: "eslint", + code: "no-console", + message: "Unexpected console statement.", + }, + ], + receivedAtMs: 1, + }, + { + uri: "file:///workspace/src/utils.ts", + path: "/workspace/src/utils.ts", + serverId: "typescript", + rootUri: "file:///workspace", + version: 1, + diagnostics: [ + { + range: { + start: { line: 9, character: 0 }, + end: { line: 9, character: 4 }, + }, + severity: 3, + source: "tsserver", + message: "'helper' is declared but its value is never read.", + }, + ], + receivedAtMs: 2, + }, + ], + }; + + const view = render(); + + expect(view.getByText("/workspace/src/example.ts")).toBeTruthy(); + expect(view.getByText("/workspace/src/utils.ts")).toBeTruthy(); + expect(view.getByText("Type 'string' is not assignable to type 'number'.")).toBeTruthy(); + expect(view.getByText("Unexpected console statement.")).toBeTruthy(); + expect(view.getByText("'helper' is declared but its value is never read.")).toBeTruthy(); + + expect(view.container.textContent).toContain( + "3 diagnostics across 2 files · 1 error · 1 warning · 1 information" + ); + expect(view.container.textContent).toContain("Server: typescript"); + expect(view.container.textContent).toContain("Error · Line 2, Column 5 · tsserver · Code 2322"); + expect(view.container.textContent).toContain( + "Warning · Line 5, Column 3 · eslint · Code no-console" + ); + expect(view.container.textContent).toContain("Information · Line 10, Column 1 · tsserver"); + }); +}); diff --git a/src/browser/features/RightSidebar/StatsTab.tsx b/src/browser/features/RightSidebar/StatsTab.tsx index ec33b4c397..2a6ecb3ee3 100644 --- a/src/browser/features/RightSidebar/StatsTab.tsx +++ b/src/browser/features/RightSidebar/StatsTab.tsx @@ -1,9 +1,12 @@ import React from "react"; -import type { WorkspaceStatsSnapshot } from "@/common/orpc/types"; +import type { WorkspaceStatsSnapshot, WorkspaceLspDiagnosticsSnapshot } from "@/common/orpc/types"; import { usePersistedState } from "@/browser/hooks/usePersistedState"; -import { useWorkspaceStatsSnapshot } from "@/browser/stores/WorkspaceStore"; +import { + useWorkspaceStatsSnapshot, + useWorkspaceLspDiagnosticsSnapshot, +} from "@/browser/stores/WorkspaceStore"; import { ToggleGroup, type ToggleOption } from "@/browser/components/ToggleGroup/ToggleGroup"; import { useTelemetry } from "@/browser/hooks/useTelemetry"; import { computeTimingPercentages } from "@/browser/utils/timingPercentages"; @@ -29,6 +32,52 @@ const VIEW_MODE_OPTIONS: Array> = [ { value: "last-request", label: "Last Request" }, ]; +const DIAGNOSTIC_SEVERITY_LABELS: Record<1 | 2 | 3 | 4, string> = { + 1: "Error", + 2: "Warning", + 3: "Information", + 4: "Hint", +}; + +function getDiagnosticSeverityLabel(severity?: 1 | 2 | 3 | 4): string { + if (severity == null) { + return "Unknown"; + } + + return DIAGNOSTIC_SEVERITY_LABELS[severity]; +} + +function formatDiagnosticLocation(line: number, character: number): string { + return `Line ${line + 1}, Column ${character + 1}`; +} + +function collectDiagnosticsSummary(snapshot: WorkspaceLspDiagnosticsSnapshot) { + const counts: Record<1 | 2 | 3 | 4, number> = { + 1: 0, + 2: 0, + 3: 0, + 4: 0, + }; + + const files = snapshot.diagnostics.filter((file) => file.diagnostics.length > 0); + let totalDiagnostics = 0; + + for (const file of files) { + totalDiagnostics += file.diagnostics.length; + for (const diagnostic of file.diagnostics) { + if (diagnostic.severity != null) { + counts[diagnostic.severity] += 1; + } + } + } + + return { + files, + counts, + totalDiagnostics, + }; +} + // Exported for unit tests. export function formatModelBreakdownLabel(entry: { model: string; @@ -781,6 +830,103 @@ export function TimingPanel(props: TimingPanelProps) { ); } +export interface DiagnosticsPanelProps { + workspaceId: string; +} + +/** + * Standalone Diagnostics panel for the Stats sub-tab. + * Subscribes lazily so inactive Stats tabs do not hold background LSP listeners open. + */ +export function DiagnosticsPanel(props: DiagnosticsPanelProps) { + const snapshot = useWorkspaceLspDiagnosticsSnapshot(props.workspaceId); + + // The workspace store uses null to mean "no snapshot yet", which is different from + // a valid empty payload after the LSP server reports zero diagnostics. + if (snapshot === null) { + return ( +
+
+

Loading diagnostics...

+
+
+ ); + } + + const { files, counts, totalDiagnostics } = collectDiagnosticsSummary(snapshot); + + if (totalDiagnostics === 0) { + return ( +
+
+

No diagnostics for this workspace.

+
+
+ ); + } + + const summaryParts = ([1, 2, 3, 4] as const) + .flatMap((severity) => { + const count = counts[severity]; + if (count === 0) { + return []; + } + + const label = getDiagnosticSeverityLabel(severity).toLowerCase(); + return [`${count} ${label}${count === 1 ? "" : "s"}`]; + }) + .join(" · "); + + return ( +
+
+ Diagnostics +

+ {totalDiagnostics} diagnostic{totalDiagnostics === 1 ? "" : "s"} across {files.length}{" "} + file + {files.length === 1 ? "" : "s"} + {summaryParts.length > 0 ? ` · ${summaryParts}` : ""} +

+
+ +
+ {files.map((file) => ( +
+
+

{file.path}

+

Server: {file.serverId}

+
+ +
    + {file.diagnostics.map((diagnostic, index) => { + const metadata = [ + getDiagnosticSeverityLabel(diagnostic.severity), + formatDiagnosticLocation( + diagnostic.range.start.line, + diagnostic.range.start.character + ), + diagnostic.source, + diagnostic.code != null ? `Code ${diagnostic.code}` : undefined, + ].filter((value): value is string => value != null); + + return ( +
  • +

    {diagnostic.message}

    +

    {metadata.join(" · ")}

    +
    • + ); + })} +
+
+ ))} +
+
+ ); +} + export interface ModelBreakdownPanelProps { workspaceId: string; } From d9dccb33fd1d32386c44db4fb5544a02a733d3f6 Mon Sep 17 00:00:00 2001 From: terion Date: Mon, 13 Apr 2026 16:21:55 +0300 Subject: [PATCH 10/46] fix diagnostics stats follow-ups --- .../StatsContainer.diagnostics.test.tsx | 20 +++- .../features/RightSidebar/StatsContainer.tsx | 1 + .../StatsTab.diagnostics.test.tsx | 47 +++++++- .../features/RightSidebar/StatsTab.tsx | 113 +++++++++++------- 4 files changed, 132 insertions(+), 49 deletions(-) diff --git a/src/browser/features/RightSidebar/StatsContainer.diagnostics.test.tsx b/src/browser/features/RightSidebar/StatsContainer.diagnostics.test.tsx index 2b8b7ee2c3..ba696bcc87 100644 --- a/src/browser/features/RightSidebar/StatsContainer.diagnostics.test.tsx +++ b/src/browser/features/RightSidebar/StatsContainer.diagnostics.test.tsx @@ -1,6 +1,6 @@ import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test"; import { GlobalWindow } from "happy-dom"; -import { cleanup, render } from "@testing-library/react"; +import { cleanup, fireEvent, render } from "@testing-library/react"; void mock.module("./CostsTab", () => ({ CostsTab: () =>
Costs panel
, @@ -25,15 +25,23 @@ describe("StatsContainer diagnostics", () => { globalThis.document = originalDocument; }); - test("shows the Diagnostics sub-tab alongside the existing Stats tabs", async () => { + test("switches to the Diagnostics sub-tab and exposes button pressed state", async () => { // eslint-disable-next-line no-restricted-syntax -- The test must import after mock.module() so StatsContainer sees the stubbed CostsTab instead of loading Mermaid-heavy dependencies. const { StatsContainer } = await import("./StatsContainer"); const view = render(); + const costButton = view.getByRole("button", { name: "Cost" }); + const diagnosticsButton = view.getByRole("button", { name: "Diagnostics" }); + expect(view.getByText("Costs panel")).toBeTruthy(); - expect(view.getByRole("button", { name: "Cost" })).toBeTruthy(); - expect(view.getByRole("button", { name: "Timing" })).toBeTruthy(); - expect(view.getByRole("button", { name: "Models" })).toBeTruthy(); - expect(view.getByRole("button", { name: "Diagnostics" })).toBeTruthy(); + expect(costButton.getAttribute("aria-pressed")).toBe("true"); + expect(diagnosticsButton.getAttribute("aria-pressed")).toBe("false"); + + fireEvent.click(diagnosticsButton); + + expect(view.queryByText("Costs panel")).toBeNull(); + expect(view.getByText("Loading diagnostics...")).toBeTruthy(); + expect(costButton.getAttribute("aria-pressed")).toBe("false"); + expect(diagnosticsButton.getAttribute("aria-pressed")).toBe("true"); }); }); diff --git a/src/browser/features/RightSidebar/StatsContainer.tsx b/src/browser/features/RightSidebar/StatsContainer.tsx index 775d5d6a84..56cddb5d1c 100644 --- a/src/browser/features/RightSidebar/StatsContainer.tsx +++ b/src/browser/features/RightSidebar/StatsContainer.tsx @@ -45,6 +45,7 @@ export function StatsContainer(props: StatsContainerProps) {