From 6357b245a5bbaba0ece311af0f1ee54646c9d0f0 Mon Sep 17 00:00:00 2001 From: can1357 Date: Tue, 3 Mar 2026 06:07:22 +0100 Subject: [PATCH] fix: corrected resource tracking and context cleanup across MCP and virtualization layers - Fixed resource refresh tracking by storing connection references alongside promises to prevent stale deduplication. - Fixed update target resolution to explicitly handle missing ompPath and use path.resolve() for consistent normalization. - Added error handling and logging in Smithery registry detail fetching to gracefully handle failures and track issues. - Fixed virtualization context cleanup in error paths to prevent partially-started instances from remaining active. - Fixed API key retrieval to use dynamic provider configuration instead of hardcoded provider string. - Enhanced test utilities to capture and verify request parameters for improved test coverage and debugging. --- crates/pi-natives/src/projfs_overlay.rs | 10 +++++- .../ai/src/provider-models/descriptors.ts | 8 +---- packages/coding-agent/CHANGELOG.md | 6 ++++ packages/coding-agent/src/cli/update-cli.ts | 31 ++++++++---------- .../coding-agent/src/config/model-registry.ts | 2 +- .../src/internal-urls/mcp-protocol.ts | 4 +-- packages/coding-agent/src/mcp/manager.ts | 17 ++++++---- .../coding-agent/src/mcp/smithery-registry.ts | 32 ++++++++++++++++--- .../coding-agent/test/client-prompts.test.ts | 28 +++++++++++++--- .../test/client-resources.test.ts | 6 +++- packages/coding-agent/test/mcp-test-utils.ts | 8 +++-- 11 files changed, 104 insertions(+), 48 deletions(-) diff --git a/crates/pi-natives/src/projfs_overlay.rs b/crates/pi-natives/src/projfs_overlay.rs index 5dcb95673..a08209945 100644 --- a/crates/pi-natives/src/projfs_overlay.rs +++ b/crates/pi-natives/src/projfs_overlay.rs @@ -335,6 +335,14 @@ mod imp { unsafe { drop(Box::from_raw(provider_context_ptr)); } + if !virtualization_context.is_null() { + // SAFETY: `virtualization_context` is only used when ProjFS returned a non-null + // context during `PrjStartVirtualizing`; stopping it here prevents a partially + // started instance from remaining active after start failure. + unsafe { + (api.prj_stop_virtualizing)(virtualization_context); + } + } return Err(classify_start_error("start virtualization", start_hr)); } @@ -517,7 +525,7 @@ mod imp { let relative_path = callback_relative_path(callback_data); let source_path = context.lower_root.join(relative_path); - let metadata = match fs::symlink_metadata(&source_path) { + let metadata = match fs::metadata(&source_path) { Ok(metadata) => metadata, Err(err) => return io_error_to_hresult(&err), }; diff --git a/packages/ai/src/provider-models/descriptors.ts b/packages/ai/src/provider-models/descriptors.ts index d87b32e61..fb4544294 100644 --- a/packages/ai/src/provider-models/descriptors.ts +++ b/packages/ai/src/provider-models/descriptors.ts @@ -210,13 +210,7 @@ export const PROVIDER_DESCRIPTORS: readonly ProviderDescriptor[] = [ config => litellmModelManagerOptions(config), catalog("LiteLLM", ["LITELLM_API_KEY"], { allowUnauthenticated: true }), ), - catalogDescriptor( - "lm-studio", - "llama-3-8b", - config => lmStudioModelManagerOptions(config), - catalog("LM Studio", ["LM_STUDIO_API_KEY"], { allowUnauthenticated: true }), - { allowUnauthenticated: true }, - ), + descriptor("lm-studio", "llama-3-8b", config => lmStudioModelManagerOptions(config), { allowUnauthenticated: true }), catalogDescriptor( "vllm", "gpt-oss-20b", diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 12ebc079f..1ae42f4fc 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -1,6 +1,12 @@ # Changelog ## [Unreleased] +### Fixed + +- Fixed LM Studio API key retrieval to use configured provider name instead of hardcoded 'lm-studio' +- Fixed resource content handling to properly check for empty text values (null/undefined) +- Fixed resource refresh tracking to prevent stale promise reuse when server connection changes +- Fixed update target resolution to properly handle cases where binary path cannot be resolved ## [13.6.1] - 2026-03-03 diff --git a/packages/coding-agent/src/cli/update-cli.ts b/packages/coding-agent/src/cli/update-cli.ts index 857e44be4..0c895424b 100644 --- a/packages/coding-agent/src/cli/update-cli.ts +++ b/packages/coding-agent/src/cli/update-cli.ts @@ -47,14 +47,6 @@ async function getBunGlobalBinDir(): Promise { } } -function getRealPathOrOriginal(filePath: string): string { - try { - return fs.realpathSync(filePath); - } catch { - return filePath; - } -} - function normalizePathForComparison(filePath: string): string { const normalized = path.normalize(filePath); if (process.platform === "win32") return normalized.toLowerCase(); @@ -62,16 +54,13 @@ function normalizePathForComparison(filePath: string): string { } function isPathInDirectory(filePath: string, directoryPath: string): boolean { - const normalizedPath = normalizePathForComparison(getRealPathOrOriginal(filePath)); - const normalizedDirectory = normalizePathForComparison(getRealPathOrOriginal(directoryPath)); + const normalizedPath = normalizePathForComparison(path.resolve(filePath)); + const normalizedDirectory = normalizePathForComparison(path.resolve(directoryPath)); const relativePath = path.relative(normalizedDirectory, normalizedPath); return relativePath === "" || (!relativePath.startsWith("..") && !path.isAbsolute(relativePath)); } -interface UpdateTarget { - method: "bun" | "binary"; - path: string; -} +type UpdateTarget = { method: "bun" } | { method: "binary"; path: string }; function resolveUpdateMethod(ompPath: string, bunBinDir: string | undefined): "bun" | "binary" { if (!bunBinDir) return "binary"; @@ -82,10 +71,18 @@ export function _resolveUpdateMethodForTest(ompPath: string, bunBinDir: string | return resolveUpdateMethod(ompPath, bunBinDir); } async function resolveUpdateTarget(): Promise { - const ompPath = resolveOmpPath() ?? process.execPath; const bunBinDir = await getBunGlobalBinDir(); - const method = resolveUpdateMethod(ompPath, bunBinDir); - return { method, path: ompPath }; + const ompPath = resolveOmpPath(); + + if (ompPath) { + const method = resolveUpdateMethod(ompPath, bunBinDir); + if (method === "bun") return { method }; + return { method, path: ompPath }; + } + + if (bunBinDir) return { method: "bun" }; + + throw new Error(`Could not resolve ${APP_NAME} binary path in PATH`); } /** diff --git a/packages/coding-agent/src/config/model-registry.ts b/packages/coding-agent/src/config/model-registry.ts index 9aa525568..647e70340 100644 --- a/packages/coding-agent/src/config/model-registry.ts +++ b/packages/coding-agent/src/config/model-registry.ts @@ -883,7 +883,7 @@ export class ModelRegistry { const modelsUrl = `${baseUrl}/models`; const headers: Record = { ...(providerConfig.headers ?? {}) }; - const apiKey = await this.authStorage.getApiKey("lm-studio"); + const apiKey = await this.authStorage.getApiKey(providerConfig.provider); if (apiKey && apiKey !== DEFAULT_LOCAL_TOKEN && apiKey !== kNoAuth) { headers.Authorization = `Bearer ${apiKey}`; } diff --git a/packages/coding-agent/src/internal-urls/mcp-protocol.ts b/packages/coding-agent/src/internal-urls/mcp-protocol.ts index 6fe0c110d..763b7f88a 100644 --- a/packages/coding-agent/src/internal-urls/mcp-protocol.ts +++ b/packages/coding-agent/src/internal-urls/mcp-protocol.ts @@ -137,14 +137,14 @@ export class McpProtocolHandler implements ProtocolHandler { const textParts: string[] = []; for (const item of result.contents) { - if (item.text) { + if (item.text !== undefined && item.text !== null) { textParts.push(item.text); } else if (item.blob) { textParts.push(`[Binary content: ${item.mimeType ?? "unknown"}, base64 length ${item.blob.length}]`); } } - const content = textParts.join("\n---\n") || "(empty resource)"; + const content = textParts.length > 0 ? textParts.join("\n---\n") : "(empty resource)"; return { url: url.href, content, diff --git a/packages/coding-agent/src/mcp/manager.ts b/packages/coding-agent/src/mcp/manager.ts index cf1fdc5ec..37099cfe3 100644 --- a/packages/coding-agent/src/mcp/manager.ts +++ b/packages/coding-agent/src/mcp/manager.ts @@ -126,7 +126,7 @@ export class MCPManager { #notificationsEnabled = false; #notificationsEpoch = 0; #subscribedResources = new Map>(); - #pendingResourceRefresh = new Map>(); + #pendingResourceRefresh = new Map }>(); constructor( private cwd: string, @@ -590,6 +590,7 @@ export class MCPManager { this.#pendingConnections.delete(name); this.#pendingToolLoads.delete(name); this.#sources.delete(name); + this.#pendingResourceRefresh.delete(name); const connection = this.#connections.get(name); @@ -622,6 +623,7 @@ export class MCPManager { this.#pendingConnections.clear(); this.#pendingToolLoads.clear(); + this.#pendingResourceRefresh.clear(); this.#sources.clear(); this.#connections.clear(); this.#tools = []; @@ -660,13 +662,13 @@ export class MCPManager { * Refresh resources from a specific server. */ async refreshServerResources(name: string): Promise { + const connection = this.#connections.get(name); + if (!connection || !serverSupportsResources(connection.capabilities)) return; + const existing = this.#pendingResourceRefresh.get(name); - if (existing) return existing; + if (existing && existing.connection === connection) return existing.promise; const doRefresh = async (): Promise => { - const connection = this.#connections.get(name); - if (!connection || !serverSupportsResources(connection.capabilities)) return; - // Clear cached resources connection.resources = undefined; connection.resourceTemplates = undefined; @@ -716,11 +718,12 @@ export class MCPManager { }; const promise = doRefresh().finally(() => { - if (this.#pendingResourceRefresh.get(name) === promise) { + const pending = this.#pendingResourceRefresh.get(name); + if (pending?.promise === promise) { this.#pendingResourceRefresh.delete(name); } }); - this.#pendingResourceRefresh.set(name, promise); + this.#pendingResourceRefresh.set(name, { connection, promise }); return promise; } diff --git a/packages/coding-agent/src/mcp/smithery-registry.ts b/packages/coding-agent/src/mcp/smithery-registry.ts index b9fbe974f..c923fdb92 100644 --- a/packages/coding-agent/src/mcp/smithery-registry.ts +++ b/packages/coding-agent/src/mcp/smithery-registry.ts @@ -1,3 +1,4 @@ +import { logger } from "@oh-my-pi/pi-utils"; import type { MCPServerConfig } from "./types"; const SMITHERY_REGISTRY_BASE_URL = "https://registry.smithery.ai"; @@ -324,8 +325,12 @@ async function fetchServerDetailsFromEntry( ): Promise { const candidates = resolveDetailPathCandidates(entry); for (const candidate of candidates) { - const details = await fetchServerDetails(candidate, options); - if (details) return details; + try { + const details = await fetchServerDetails(candidate, options); + if (details) return details; + } catch (error) { + logger.debug("Smithery detail fetch candidate failed", { candidate, error: String(error) }); + } } return null; } @@ -439,14 +444,31 @@ export async function searchSmitheryRegistry( ); }); + const detailFailures: Array<{ identity: string; error: string }> = []; const results = await Promise.all( uniqueEntries.map(async entry => { - const details = await fetchServerDetailsFromEntry(entry, { apiKey: options?.apiKey }); - if (!details) return null; - return toSearchResult(entry, details); + try { + const details = await fetchServerDetailsFromEntry(entry, { apiKey: options?.apiKey }); + if (!details) return null; + return toSearchResult(entry, details); + } catch (error) { + detailFailures.push({ + identity: getEntryIdentityKey(entry) ?? entry.id ?? "unknown", + error: String(error), + }); + return null; + } }), ); + if (detailFailures.length > 0) { + logger.warn("Smithery detail fetch failed for some entries", { + query, + failedEntries: detailFailures.length, + totalEntries: uniqueEntries.length, + sample: detailFailures.slice(0, 3), + }); + } return results.filter((result): result is SmitherySearchResult => result !== null).slice(0, limit); } diff --git a/packages/coding-agent/test/client-prompts.test.ts b/packages/coding-agent/test/client-prompts.test.ts index 214398c7c..b0c272f67 100644 --- a/packages/coding-agent/test/client-prompts.test.ts +++ b/packages/coding-agent/test/client-prompts.test.ts @@ -68,13 +68,17 @@ describe("getPrompt", () => { messages: [{ role: "user", content: { type: "text", text: "Hello!" } }], }; const responses = new Map([["prompts/get", [mockResult]]]); - const transport = createMockTransport(responses); + let requestParams: Record | undefined; + const transport = createMockTransport(responses, (_method, params) => { + requestParams = params; + }); const conn = createMockConnection({ prompts: {} }, transport); const result = await getPrompt(conn, "greet"); expect(result).toEqual(mockResult); expect(result.messages).toHaveLength(1); expect(result.messages[0].role).toBe("user"); + expect(requestParams).toEqual({ name: "greet" }); }); it("sends arguments when provided", async () => { @@ -82,11 +86,17 @@ describe("getPrompt", () => { messages: [{ role: "assistant", content: { type: "text", text: "const x = 1" } }], }; const responses = new Map([["prompts/get", [mockResult]]]); - const transport = createMockTransport(responses); + let requestParams: Record | undefined; + const transport = createMockTransport(responses, (_method, params) => { + requestParams = params; + }); const conn = createMockConnection({ prompts: {} }, transport); - const result = await getPrompt(conn, "review-code", { code: "const x = 1" }); + const args = { code: "const x = 1" }; + const result = await getPrompt(conn, "review-code", args); expect(result).toEqual(mockResult); + expect(requestParams).toEqual({ name: "review-code", arguments: args }); + expect(requestParams?.arguments).toBe(args); }); it("sends without arguments when args is empty object", async () => { @@ -94,11 +104,15 @@ describe("getPrompt", () => { messages: [{ role: "user", content: { type: "text", text: "No args" } }], }; const responses = new Map([["prompts/get", [mockResult]]]); - const transport = createMockTransport(responses); + let requestParams: Record | undefined; + const transport = createMockTransport(responses, (_method, params) => { + requestParams = params; + }); const conn = createMockConnection({ prompts: {} }, transport); const result = await getPrompt(conn, "no-args-prompt", {}); expect(result).toEqual(mockResult); + expect(requestParams).toEqual({ name: "no-args-prompt" }); }); it("sends without arguments when args is undefined", async () => { @@ -106,11 +120,15 @@ describe("getPrompt", () => { messages: [{ role: "user", content: { type: "text", text: "No args" } }], }; const responses = new Map([["prompts/get", [mockResult]]]); - const transport = createMockTransport(responses); + let requestParams: Record | undefined; + const transport = createMockTransport(responses, (_method, params) => { + requestParams = params; + }); const conn = createMockConnection({ prompts: {} }, transport); const result = await getPrompt(conn, "no-args-prompt", undefined); expect(result).toEqual(mockResult); + expect(requestParams).toEqual({ name: "no-args-prompt" }); }); }); diff --git a/packages/coding-agent/test/client-resources.test.ts b/packages/coding-agent/test/client-resources.test.ts index 495b3fa93..a5cbdb3ee 100644 --- a/packages/coding-agent/test/client-resources.test.ts +++ b/packages/coding-agent/test/client-resources.test.ts @@ -102,13 +102,17 @@ describe("readResource", () => { const readResult: MCPResourceReadResult = { contents: [{ uri: "file:///a.txt", mimeType: "text/plain", text: "hello" }], }; - const transport = createMockTransport(new Map([["resources/read", [readResult]]])); + let requestParams: Record | undefined; + const transport = createMockTransport(new Map([["resources/read", [readResult]]]), (_method, params) => { + requestParams = params; + }); const conn = createMockConnection({ resources: {} }, transport); const result = await readResource(conn, "file:///a.txt"); expect(result.contents).toHaveLength(1); expect(result.contents[0].text).toBe("hello"); expect(result.contents[0].mimeType).toBe("text/plain"); + expect(requestParams).toEqual({ uri: "file:///a.txt" }); }); it("handles binary blobs", async () => { diff --git a/packages/coding-agent/test/mcp-test-utils.ts b/packages/coding-agent/test/mcp-test-utils.ts index de030fd49..d0d9c229a 100644 --- a/packages/coding-agent/test/mcp-test-utils.ts +++ b/packages/coding-agent/test/mcp-test-utils.ts @@ -1,10 +1,14 @@ import type { MCPServerCapabilities, MCPServerConnection, MCPTransport } from "../src/mcp/types"; -export function createMockTransport(responses: Map): MCPTransport { +export function createMockTransport( + responses: Map, + onRequest?: (method: string, params: Record | undefined) => void, +): MCPTransport { const callCounts = new Map(); return { connected: true, - async request(method: string): Promise { + async request(method: string, params?: Record): Promise { + onRequest?.(method, params); const count = callCounts.get(method) ?? 0; callCounts.set(method, count + 1); const queue = responses.get(method);