From d74d0765d7bb480812dc55891161c934bc337327 Mon Sep 17 00:00:00 2001 From: YueZengwu Date: Fri, 17 Jul 2026 20:39:46 +0800 Subject: [PATCH 1/5] feat(runtime): host-local model catalogs for Cursor and Kimi Discover model ids from the bound computer (agent models / ~/.kimi-code) via a correlated daemon RPC, and render DEFAULT + catalog + custom id fallbacks in the agent Model picker. Co-authored-by: Cursor --- apps/cli/src/commands/daemon/start.ts | 23 ++ apps/cli/src/core/client-runtime.ts | 16 ++ .../src/__tests__/discover-models.test.ts | 42 ++++ packages/client/src/client-connection.ts | 26 ++ packages/client/src/index.ts | 8 + .../runtime/capabilities/discover-models.ts | 201 +++++++++++++++ .../__tests__/clients-provider-models.test.ts | 66 +++++ packages/server/src/api/agent/ws-client.ts | 23 ++ packages/server/src/api/clients.ts | 44 +++- .../server/src/services/connection-manager.ts | 62 +++++ packages/shared/src/index.ts | 14 ++ .../shared/src/schemas/provider-models.ts | 69 ++++++ packages/web/src/api/activity.ts | 11 + .../src/pages/agent-detail/model-section.tsx | 233 ++++++++++++++++-- .../src/pages/agent-detail/runtime-tab.tsx | 1 + .../cursor-provider-surfaces.test.tsx | 9 +- .../kimi-code-provider-surfaces.test.tsx | 6 +- 17 files changed, 821 insertions(+), 33 deletions(-) create mode 100644 packages/client/src/__tests__/discover-models.test.ts create mode 100644 packages/client/src/runtime/capabilities/discover-models.ts create mode 100644 packages/server/src/__tests__/clients-provider-models.test.ts create mode 100644 packages/shared/src/schemas/provider-models.ts diff --git a/apps/cli/src/commands/daemon/start.ts b/apps/cli/src/commands/daemon/start.ts index cd2bb9bc8..6e1584a79 100644 --- a/apps/cli/src/commands/daemon/start.ts +++ b/apps/cli/src/commands/daemon/start.ts @@ -8,6 +8,7 @@ import { configureClientLoggerForService, createLogger, discoverClaudeCodeSkills, + discoverProviderModels, flushClientSentry, initClientSentry, } from "@first-tree/client"; @@ -366,6 +367,28 @@ export function registerDaemonStartCommand(daemon: Command): void { }).finally(() => capabilityRefresher.endInteractive(command.provider)); }); + // Host-local model catalog: web opens Model settings → server asks this + // daemon → we discover from the real provider and reply on the WS. + runtime.onProviderModelsList((command) => { + void (async () => { + try { + const catalog = await discoverProviderModels(command.provider); + runtime.sendProviderModelsResult(command.ref, catalog); + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + writeStatus("⚠️", `provider-models:list failed for ${command.provider}: ${message}`); + runtime.sendProviderModelsResult(command.ref, { + provider: command.provider, + models: [], + defaultModelId: null, + fetchedAt: new Date().toISOString(), + source: "unavailable", + error: message, + }); + } + })(); + }); + await runtime.start(); // Post-register capabilities upload + arm the background poll — the diff --git a/apps/cli/src/core/client-runtime.ts b/apps/cli/src/core/client-runtime.ts index e437fefaf..36c0bffb9 100644 --- a/apps/cli/src/core/client-runtime.ts +++ b/apps/cli/src/core/client-runtime.ts @@ -8,6 +8,7 @@ import { getChildProcessRegistry, getHandlerFactory, hasHandler, + type ProviderModelsListCommand, type RuntimeAuthCommand, registerBuiltinHandlers, type UpdateHooks, @@ -17,6 +18,7 @@ import { AGENT_BIND_REJECT_REASONS, type AgentPinnedMessage, type ClientPausedReason, + type ProviderModelCatalog, type RuntimeProvider, runtimeProviderSchema, } from "@first-tree/shared"; @@ -284,6 +286,20 @@ export class ClientRuntime { this.connection.on("runtime-auth:start", callback); } + /** + * Register a handler for the server→client `provider-models:list` command. + * The daemon discovers models from the host-local provider and replies with + * `provider-models:result` on the same connection. + */ + onProviderModelsList(callback: (command: ProviderModelsListCommand) => void): void { + this.connection.on("provider-models:list", callback); + } + + /** Reply to a correlated `provider-models:list` with the discovered catalog. */ + sendProviderModelsResult(ref: string, catalog: ProviderModelCatalog): void { + this.connection.sendProviderModelsResult(ref, catalog); + } + addAgent(name: string, config: AgentConfig): void { if (this.agentNames.has(name)) return; // The runtime provider is a valid enum value, but this client build may not diff --git a/packages/client/src/__tests__/discover-models.test.ts b/packages/client/src/__tests__/discover-models.test.ts new file mode 100644 index 000000000..a54b96689 --- /dev/null +++ b/packages/client/src/__tests__/discover-models.test.ts @@ -0,0 +1,42 @@ +import { describe, expect, it } from "vitest"; +import { parseCursorModelsOutput, parseKimiConfigModels } from "../runtime/capabilities/discover-models.js"; + +describe("parseCursorModelsOutput", () => { + it("parses id/label rows and marks the default", () => { + const parsed = parseCursorModelsOutput(`Available models + +auto - Auto (default) +gpt-5.2 - GPT-5.2 +composer-2.5 - Composer 2.5 +`); + expect(parsed.defaultModelId).toBe("auto"); + expect(parsed.models).toEqual([ + { id: "auto", label: "Auto", isDefault: true, hint: "default" }, + { id: "gpt-5.2", label: "GPT-5.2" }, + { id: "composer-2.5", label: "Composer 2.5" }, + ]); + }); +}); + +describe("parseKimiConfigModels", () => { + it("reads default_model and [models.\".\"] sections", () => { + const parsed = parseKimiConfigModels(` +default_model = "kimi-code/k3" + +[models."kimi-code/kimi-for-coding"] +provider = "managed:kimi-code" +model = "kimi-for-coding" +display_name = "K2.7 Coding" + +[models."kimi-code/k3"] +provider = "managed:kimi-code" +model = "k3" +display_name = "K3" +`); + expect(parsed.defaultModelId).toBe("kimi-code/k3"); + expect(parsed.models).toEqual([ + { id: "kimi-code/kimi-for-coding", label: "K2.7 Coding" }, + { id: "kimi-code/k3", label: "K3", isDefault: true, hint: "default" }, + ]); + }); +}); diff --git a/packages/client/src/client-connection.ts b/packages/client/src/client-connection.ts index bef230b8e..7c7dcc988 100644 --- a/packages/client/src/client-connection.ts +++ b/packages/client/src/client-connection.ts @@ -16,10 +16,14 @@ import { inboxDeliverFrameSchema, inboxRecoverAcceptedFrameSchema, inboxRecoverRejectedFrameSchema, + PROVIDER_MODELS_LIST_TYPE, + PROVIDER_MODELS_RESULT_TYPE, + type ProviderModelCatalog, RUNTIME_AUTH_START_TYPE, type RuntimeAuthMethod, type RuntimeProvider, type RuntimeState, + providerModelsListCommandSchema, runtimeAuthStartCommandSchema, type ServerWelcomeFrame, type SessionEvent, @@ -164,6 +168,12 @@ export type RuntimeAuthCommand = { ref: string; }; +/** Server→client command to discover host-local provider models. */ +export type ProviderModelsListCommand = { + provider: RuntimeProvider; + ref: string; +}; + /** * Welcome frame received after `auth:ok`. `isReconnect` is true for every * occurrence after the first welcome in the lifetime of this `ClientConnection` @@ -199,6 +209,7 @@ type ClientConnectionEvents = { "agent:pinned": [message: AgentPinnedMessage]; "session:command": [command: SessionCommand]; "runtime-auth:start": [command: RuntimeAuthCommand]; + "provider-models:list": [command: ProviderModelsListCommand]; "session:reconcile:result": [result: SessionReconcileResult]; "auth:expired": []; /** @@ -1082,6 +1093,12 @@ export class ClientConnection extends EventEmitter { this.ws.send(JSON.stringify({ type: "session:reconcile", agentId, chatIds })); } + /** Reply to a `provider-models:list` reverse command with the host catalog. */ + sendProviderModelsResult(ref: string, catalog: ProviderModelCatalog): void { + if (!this.ws || this.ws.readyState !== WebSocket.OPEN) return; + this.ws.send(JSON.stringify({ type: PROVIDER_MODELS_RESULT_TYPE, ref, catalog })); + } + async disconnect(): Promise { this.closing = true; this.connectAbort?.abort(); @@ -1579,6 +1596,15 @@ export class ClientConnection extends EventEmitter { return; } + if (type === PROVIDER_MODELS_LIST_TYPE) { + const parsed = providerModelsListCommandSchema.safeParse(msg); + if (parsed.success) { + const { provider, ref } = parsed.data; + this.emit("provider-models:list", { provider, ref }); + } + return; + } + if (type === "session:reconcile:result") { const agentId = msg.agentId as string; const staleChatIds = Array.isArray(msg.staleChatIds) ? (msg.staleChatIds as string[]) : null; diff --git a/packages/client/src/index.ts b/packages/client/src/index.ts index 9afefb6fb..8e8af2381 100644 --- a/packages/client/src/index.ts +++ b/packages/client/src/index.ts @@ -1,6 +1,7 @@ export type { BoundAgent, ClientConnectionConfig, + ProviderModelsListCommand, RuntimeAuthCommand, ServerWelcome, SessionCommand, @@ -42,6 +43,12 @@ export { resolveCodexRuntimeBinary, } from "./runtime/capabilities/codex.js"; export { probeCursorCapability } from "./runtime/capabilities/cursor.js"; +export { + type DiscoverModelsDeps, + discoverProviderModels, + parseCursorModelsOutput, + parseKimiConfigModels, +} from "./runtime/capabilities/discover-models.js"; export { CAPABILITY_REFRESH_BASE_MS, CAPABILITY_REFRESH_MAX_MS, @@ -54,6 +61,7 @@ export { revalidateCapabilities, shouldFullReprobe, } from "./runtime/capabilities/index.js"; + export type { AdoptOptions, ChildCategory, diff --git a/packages/client/src/runtime/capabilities/discover-models.ts b/packages/client/src/runtime/capabilities/discover-models.ts new file mode 100644 index 000000000..6d11b58a8 --- /dev/null +++ b/packages/client/src/runtime/capabilities/discover-models.ts @@ -0,0 +1,201 @@ +import { homedir } from "node:os"; +import { join } from "node:path"; +import { readFile } from "node:fs/promises"; +import type { ProviderModelCatalog, ProviderModelOption, RuntimeProvider } from "@first-tree/shared"; +import { findCursorExecutableOnPath } from "../cursor-binary.js"; +import { runCommand } from "./launch-probe.js"; + +/** Ceiling for `agent models` — account catalog fetch can be network-bound. */ +const CURSOR_MODELS_TIMEOUT_MS = 20_000; + +export type DiscoverModelsDeps = { + env?: NodeJS.ProcessEnv; + now?: () => Date; + findCursorBinary?: (env?: Record) => string | null; + runCursorModels?: (binary: string, env: NodeJS.ProcessEnv) => Promise<{ ok: boolean; stdout: string; stderr: string }>; + readKimiConfig?: () => Promise; + kimiConfigPath?: string; +}; + +function fetchedAt(deps: DiscoverModelsDeps): string { + return (deps.now ?? (() => new Date()))().toISOString(); +} + +function unavailable(provider: RuntimeProvider, error: string, deps: DiscoverModelsDeps): ProviderModelCatalog { + return { + provider, + models: [], + defaultModelId: null, + fetchedAt: fetchedAt(deps), + source: "unavailable", + error, + }; +} + +/** + * Parse `agent models` / `agent --list-models` text: + * Available models + * auto - Auto (default) + * gpt-5.2 - GPT-5.2 + */ +export function parseCursorModelsOutput(stdout: string): { + models: ProviderModelOption[]; + defaultModelId: string | null; +} { + const models: ProviderModelOption[] = []; + let defaultModelId: string | null = null; + for (const rawLine of stdout.split(/\r?\n/)) { + const line = rawLine.trim(); + if (!line || /^available models$/i.test(line)) continue; + const match = /^(\S+)\s+-\s+(.+)$/.exec(line); + if (!match) continue; + const id = match[1]!; + const label = match[2]!.trim(); + const isDefault = /\(default\)/i.test(label); + if (isDefault) defaultModelId = id; + models.push({ + id, + label: label.replace(/\s*\(default\)\s*/i, "").trim() || id, + ...(isDefault ? { isDefault: true, hint: "default" } : {}), + }); + } + return { models, defaultModelId }; +} + +/** + * Minimal TOML extract for Kimi's `~/.kimi-code/config.toml`: + * default_model = "kimi-code/k3" + * [models."kimi-code/k3"] + * display_name = "K3" + * model = "k3" + * + * Avoids adding a TOML dependency for a narrow, stable config shape. + */ +export function parseKimiConfigModels(toml: string): { + models: ProviderModelOption[]; + defaultModelId: string | null; +} { + const defaultMatch = /^default_model\s*=\s*"([^"]+)"/m.exec(toml); + const defaultModelId = defaultMatch?.[1] ?? null; + + const models: ProviderModelOption[] = []; + const sectionRe = /^\[models\."([^"]+)"\]\s*$/gm; + const sections: Array<{ id: string; headerStart: number; bodyStart: number }> = []; + for (const match of toml.matchAll(sectionRe)) { + sections.push({ + id: match[1]!, + headerStart: match.index!, + bodyStart: match.index! + match[0].length, + }); + } + for (let i = 0; i < sections.length; i++) { + const section = sections[i]!; + const bodyEnd = i + 1 < sections.length ? sections[i + 1]!.headerStart : toml.length; + const body = toml.slice(section.bodyStart, bodyEnd); + const displayName = /^display_name\s*=\s*"([^"]+)"/m.exec(body)?.[1]; + const isDefault = defaultModelId === section.id; + models.push({ + id: section.id, + ...(displayName ? { label: displayName } : {}), + ...(isDefault ? { isDefault: true, hint: "default" } : {}), + }); + } + return { models, defaultModelId }; +} + +async function discoverCursorModels(deps: DiscoverModelsDeps): Promise { + const env = deps.env ?? process.env; + const findBinary = deps.findCursorBinary ?? findCursorExecutableOnPath; + const binary = findBinary(env); + if (!binary) { + return unavailable("cursor", "cursor-agent / agent binary not found on this host", deps); + } + const run = + deps.runCursorModels ?? + (async (bin, processEnv) => { + const result = await runCommand(bin, ["models"], { timeoutMs: CURSOR_MODELS_TIMEOUT_MS, env: processEnv }); + return { ok: result.ok, stdout: result.stdout, stderr: result.stderr }; + }); + const result = await run(binary, env); + if (!result.ok) { + const detail = (result.stderr || result.stdout || "agent models failed").trim(); + return unavailable("cursor", detail.slice(0, 500), deps); + } + const parsed = parseCursorModelsOutput(result.stdout); + if (parsed.models.length === 0) { + return unavailable("cursor", "agent models returned no parseable model rows", deps); + } + return { + provider: "cursor", + models: parsed.models, + defaultModelId: parsed.defaultModelId, + fetchedAt: fetchedAt(deps), + source: "provider-cli", + error: null, + }; +} + +async function discoverKimiModels(deps: DiscoverModelsDeps): Promise { + const path = deps.kimiConfigPath ?? join(homedir(), ".kimi-code", "config.toml"); + const read = + deps.readKimiConfig ?? + (async () => { + try { + return await readFile(path, "utf8"); + } catch (err) { + const code = err && typeof err === "object" && "code" in err ? String((err as { code: unknown }).code) : ""; + if (code === "ENOENT") return null; + throw err; + } + }); + let toml: string | null; + try { + toml = await read(); + } catch (err) { + return unavailable("kimi-code", err instanceof Error ? err.message : String(err), deps); + } + if (toml == null) { + return unavailable("kimi-code", `Kimi config not found at ${path}`, deps); + } + const parsed = parseKimiConfigModels(toml); + if (parsed.models.length === 0) { + return unavailable("kimi-code", "Kimi config has no [models.\".\"] entries", deps); + } + return { + provider: "kimi-code", + models: parsed.models, + defaultModelId: parsed.defaultModelId, + fetchedAt: fetchedAt(deps), + source: "provider-config", + error: null, + }; +} + +/** + * Discover the model catalog for a runtime provider from the host-local + * provider. Phase 1 implements Cursor + Kimi; other providers return + * `source: "unavailable"` so the web can keep its curated/fallback UI. + */ +export async function discoverProviderModels( + provider: RuntimeProvider, + deps: DiscoverModelsDeps = {}, +): Promise { + switch (provider) { + case "cursor": + return discoverCursorModels(deps); + case "kimi-code": + return discoverKimiModels(deps); + case "claude-code": + case "claude-code-tui": + case "codex": + return unavailable( + provider, + `Host-local model discovery for ${provider} lands in a later phase`, + deps, + ); + default: { + const _exhaustive: never = provider; + return unavailable(_exhaustive, `Unknown provider: ${String(provider)}`, deps); + } + } +} diff --git a/packages/server/src/__tests__/clients-provider-models.test.ts b/packages/server/src/__tests__/clients-provider-models.test.ts new file mode 100644 index 000000000..a159ae012 --- /dev/null +++ b/packages/server/src/__tests__/clients-provider-models.test.ts @@ -0,0 +1,66 @@ +import crypto from "node:crypto"; +import { describe, expect, it, vi } from "vitest"; +import type { WebSocket } from "ws"; +import { + removeClientConnection, + resolveClientReply, + setClientConnection, +} from "../services/connection-manager.js"; +import { createAdminContext, useTestApp } from "./helpers.js"; + +/** + * `GET /api/v1/clients/:clientId/providers/:provider/models` asks the connected + * daemon for a host-local model catalog and waits for the correlated reply. + */ +describe("GET /clients/:clientId/providers/:provider/models", () => { + const getApp = useTestApp(); + + it("forwards provider-models:list and returns the daemon catalog", async () => { + const app = getApp(); + const admin = await createAdminContext(app, { username: `pm-${crypto.randomUUID().slice(0, 6)}` }); + const ws = { readyState: 1, send: vi.fn(), close: vi.fn() }; + setClientConnection(admin.clientId, ws as unknown as WebSocket); + try { + const pending = app.inject({ + method: "GET", + url: `/api/v1/clients/${admin.clientId}/providers/cursor/models`, + headers: { authorization: `Bearer ${admin.accessToken}` }, + }); + + // Wait briefly for the reverse command to be sent, then resolve the waiter. + await vi.waitFor(() => { + expect(ws.send).toHaveBeenCalled(); + }); + const frame = JSON.parse(String(ws.send.mock.calls[0]?.[0])); + expect(frame).toMatchObject({ type: "provider-models:list", provider: "cursor" }); + expect(typeof frame.ref).toBe("string"); + + const catalog = { + provider: "cursor" as const, + models: [{ id: "auto", label: "Auto", isDefault: true }], + defaultModelId: "auto", + fetchedAt: new Date().toISOString(), + source: "provider-cli" as const, + error: null, + }; + expect(resolveClientReply(admin.clientId, frame.ref, catalog)).toBe(true); + + const res = await pending; + expect(res.statusCode).toBe(200); + expect(res.json()).toMatchObject(catalog); + } finally { + removeClientConnection(admin.clientId, ws as unknown as WebSocket); + } + }); + + it("returns 503 when the daemon is not connected", async () => { + const app = getApp(); + const admin = await createAdminContext(app, { username: `pm-${crypto.randomUUID().slice(0, 6)}` }); + const res = await app.inject({ + method: "GET", + url: `/api/v1/clients/${admin.clientId}/providers/kimi-code/models`, + headers: { authorization: `Bearer ${admin.accessToken}` }, + }); + expect(res.statusCode).toBe(503); + }); +}); diff --git a/packages/server/src/api/agent/ws-client.ts b/packages/server/src/api/agent/ws-client.ts index ccb1e2275..be115a2f4 100644 --- a/packages/server/src/api/agent/ws-client.ts +++ b/packages/server/src/api/agent/ws-client.ts @@ -5,6 +5,7 @@ import { AUTH_RETRYABLE_CODES, type AuthRejectedCode, type AuthRetryableCode, + PROVIDER_MODELS_RESULT_TYPE, agentBindRequestSchema, agentPinnedMessageSchema, clientRegisterSchema, @@ -12,6 +13,7 @@ import { inboxAckFrameSchema, inboxDeliverFrameSchema, inboxRecoverFrameSchema, + providerModelsResultFrameSchema, runtimeStateMessageSchema, sessionEventMessageSchema, sessionEventRejectedReasonSchema, @@ -1757,6 +1759,27 @@ export function clientWsRoutes(notifier: Notifier, instanceId: string) { await reconcilePinnedAgentsForClient(); } socket.send(JSON.stringify({ type: "heartbeat:ack" })); + } else if (type === PROVIDER_MODELS_RESULT_TYPE) { + if (!clientId) { + socket.send(JSON.stringify({ type: "error", message: "Must register client first" })); + return; + } + const result = providerModelsResultFrameSchema.safeParse(msg); + if (!result.success) { + socket.send(JSON.stringify({ type: "error", message: "Malformed provider-models:result frame" })); + return; + } + const resolved = connectionManager.resolveClientReply( + clientId, + result.data.ref, + result.data.catalog, + ); + if (!resolved) { + app.log.debug( + { clientId, ref: result.data.ref }, + "provider-models:result matched no pending HTTP waiter", + ); + } } } catch (err) { const message = err instanceof Error ? err.message : "Internal error"; diff --git a/packages/server/src/api/clients.ts b/packages/server/src/api/clients.ts index a2baf7c83..0a621b4ea 100644 --- a/packages/server/src/api/clients.ts +++ b/packages/server/src/api/clients.ts @@ -1,7 +1,10 @@ import { randomUUID } from "node:crypto"; import { + PROVIDER_MODELS_LIST_TYPE, RUNTIME_AUTH_START_TYPE, + providerModelCatalogSchema, runtimeAuthStartRequestSchema, + runtimeProviderSchema, updateClientCapabilitiesSchema, } from "@first-tree/shared"; import { getChannelConfig } from "@first-tree/shared/channel"; @@ -11,7 +14,7 @@ import { stampClientResource } from "../observability/request-context.js"; import { requireUser } from "../scope/require-user.js"; import { expiryToSeconds } from "../services/auth.js"; import * as clientService from "../services/client.js"; -import { forceDisconnectClient, sendToClient } from "../services/connection-manager.js"; +import { forceDisconnectClient, rejectPendingRepliesForClient, sendToClient, waitForClientReply } from "../services/connection-manager.js"; import { serializeDate } from "../utils.js"; import { clientCommandVersionHint } from "./client-command-version.js"; @@ -91,6 +94,45 @@ export async function clientRoutes(app: FastifyInstance): Promise { return { ref, started: true as const }; }); + // Host-local model catalog: ask the connected daemon to discover models from + // the real provider on that computer, wait for the correlated reply, and + // return the catalog to the web. 503 when the computer is offline or times out. + app.get<{ Params: { clientId: string; provider: string } }>( + "/:clientId/providers/:provider/models", + async (request) => { + const { userId } = requireUser(request); + const { clientId, provider: rawProvider } = request.params; + stampClientResource(request, clientId); + await clientService.assertClientOwner(app.db, clientId, { userId }); + await clientService.assertClientNotRetired(app.db, clientId); + const provider = runtimeProviderSchema.parse(rawProvider); + const ref = randomUUID(); + const replyPromise = waitForClientReply(clientId, ref); + const delivered = sendToClient(clientId, { + type: PROVIDER_MODELS_LIST_TYPE, + provider, + ref, + }); + if (!delivered) { + rejectPendingRepliesForClient(clientId, new Error("Computer not connected")); + await replyPromise.catch(() => undefined); + throw new ServiceUnavailableError( + "Could not list models because this computer is not connected. Make sure the daemon is running, then retry.", + ); + } + try { + const raw = await replyPromise; + return providerModelCatalogSchema.parse(raw); + } catch (err) { + throw new ServiceUnavailableError( + err instanceof Error + ? err.message + : "Could not list models from this computer. Retry after the daemon is connected.", + ); + } + }, + ); + app.post<{ Params: { clientId: string } }>("/:clientId/disconnect", async (request) => { const { userId } = requireUser(request); const { clientId } = request.params; diff --git a/packages/server/src/services/connection-manager.ts b/packages/server/src/services/connection-manager.ts index f5924a505..734f9b409 100644 --- a/packages/server/src/services/connection-manager.ts +++ b/packages/server/src/services/connection-manager.ts @@ -168,6 +168,7 @@ export function removeClientConnection(clientId: string, ws: WebSocket): string[ activeConnections.delete(agentId); } clientConnections.delete(clientId); + rejectPendingRepliesForClient(clientId, new Error("Client disconnected")); return agentIds; } @@ -222,5 +223,66 @@ export function forceDisconnectClient(clientId: string): string[] { activeConnections.delete(agentId); } clientConnections.delete(clientId); + rejectPendingRepliesForClient(clientId, new Error("Client disconnected")); return agentIds; } + +/** + * HTTP→daemon request/response correlation. Used by host-local discovery + * (provider model catalogs) where the web needs a synchronous answer from the + * connected computer rather than fire-and-forget + poll. + */ +type PendingClientReply = { + clientId: string; + resolve: (value: unknown) => void; + reject: (error: Error) => void; + timer: ReturnType; +}; + +const pendingClientReplies = new Map(); + +const DEFAULT_CLIENT_REPLY_TIMEOUT_MS = 25_000; + +export function waitForClientReply( + clientId: string, + ref: string, + timeoutMs: number = DEFAULT_CLIENT_REPLY_TIMEOUT_MS, +): Promise { + return new Promise((resolve, reject) => { + if (pendingClientReplies.has(ref)) { + reject(new Error(`Duplicate pending client reply ref: ${ref}`)); + return; + } + const timer = setTimeout(() => { + pendingClientReplies.delete(ref); + reject(new Error("Timed out waiting for the computer to reply")); + }, timeoutMs); + pendingClientReplies.set(ref, { clientId, resolve, reject, timer }); + }); +} + +export function resolveClientReply(clientId: string, ref: string, value: unknown): boolean { + const pending = pendingClientReplies.get(ref); + if (!pending || pending.clientId !== clientId) return false; + clearTimeout(pending.timer); + pendingClientReplies.delete(ref); + pending.resolve(value); + return true; +} + +export function rejectPendingRepliesForClient(clientId: string, error: Error): void { + for (const [ref, pending] of pendingClientReplies) { + if (pending.clientId !== clientId) continue; + clearTimeout(pending.timer); + pendingClientReplies.delete(ref); + pending.reject(error); + } +} + +/** Test seam: drop every pending reply without resolving. */ +export function clearPendingClientRepliesForTests(): void { + for (const pending of pendingClientReplies.values()) { + clearTimeout(pending.timer); + } + pendingClientReplies.clear(); +} diff --git a/packages/shared/src/index.ts b/packages/shared/src/index.ts index 2f012a0d0..6134c6f4a 100644 --- a/packages/shared/src/index.ts +++ b/packages/shared/src/index.ts @@ -974,6 +974,20 @@ export { runtimeAuthStartRequestSchema, runtimeAuthStartResponseSchema, } from "./schemas/runtime-auth.js"; +export { + PROVIDER_MODELS_LIST_TYPE, + PROVIDER_MODELS_RESULT_TYPE, + type ProviderModelCatalog, + type ProviderModelCatalogSource, + type ProviderModelOption, + type ProviderModelsListCommand, + type ProviderModelsResultFrame, + providerModelCatalogSchema, + providerModelCatalogSourceSchema, + providerModelOptionSchema, + providerModelsListCommandSchema, + providerModelsResultFrameSchema, +} from "./schemas/provider-models.js"; export { DEFAULT_RUNTIME_PROVIDER, DISABLED_RUNTIME_PROVIDERS, diff --git a/packages/shared/src/schemas/provider-models.ts b/packages/shared/src/schemas/provider-models.ts new file mode 100644 index 000000000..c5c8bf244 --- /dev/null +++ b/packages/shared/src/schemas/provider-models.ts @@ -0,0 +1,69 @@ +import { z } from "zod"; +import { runtimeProviderSchema } from "./runtime-provider.js"; + +/** + * Host-local model catalog discovery: the daemon asks the real provider on the + * computer for the models the operator can pick, and the web renders that list. + * First Tree never maintains a global model matrix and never proxies provider + * credentials — discovery runs only on the daemon host. + * + * Wire shape mirrors runtime-auth: + * - Web → Server HTTP `GET /clients/:clientId/providers/:provider/models` + * - Server → daemon reverse command `provider-models:list` (ref-correlated) + * - Daemon → Server reply frame `provider-models:result` + * - Server resolves the pending HTTP request with the catalog + */ + +/** Server→client command: discover models for one runtime provider. */ +export const PROVIDER_MODELS_LIST_TYPE = "provider-models:list" as const; + +/** Client→server reply carrying the discovered catalog (or an unavailable stub). */ +export const PROVIDER_MODELS_RESULT_TYPE = "provider-models:result" as const; + +export const providerModelOptionSchema = z.object({ + /** Exact value written to `config.payload.model`. */ + id: z.string().min(1), + /** Human-facing label; defaults to `id` in the UI when absent. */ + label: z.string().optional(), + /** Secondary hint (e.g. "default", "flagship"). */ + hint: z.string().optional(), + /** Provider-reported default; informational — Web still uses empty string for DEFAULT. */ + isDefault: z.boolean().optional(), +}); +export type ProviderModelOption = z.infer; + +export const providerModelCatalogSourceSchema = z.enum([ + "provider-cli", + "provider-config", + "provider-rpc", + "provider-cache", + "unavailable", +]); +export type ProviderModelCatalogSource = z.infer; + +export const providerModelCatalogSchema = z.object({ + provider: runtimeProviderSchema, + models: z.array(providerModelOptionSchema), + /** Provider's own default model id when known (Kimi `default_model`, Cursor `auto`, …). */ + defaultModelId: z.string().nullable().optional(), + fetchedAt: z.string(), + source: providerModelCatalogSourceSchema, + /** Present when discovery failed or the provider is not supported yet. */ + error: z.string().nullable().optional(), +}); +export type ProviderModelCatalog = z.infer; + +export const providerModelsListCommandSchema = z.object({ + type: z.literal(PROVIDER_MODELS_LIST_TYPE), + provider: runtimeProviderSchema, + /** Correlation id tying command → result → HTTP response. */ + ref: z.string().min(1), +}); +export type ProviderModelsListCommand = z.infer; + +export const providerModelsResultFrameSchema = z.object({ + type: z.literal(PROVIDER_MODELS_RESULT_TYPE), + ref: z.string().min(1), + catalog: providerModelCatalogSchema, +}); +export type ProviderModelsResultFrame = z.infer; diff --git a/packages/web/src/api/activity.ts b/packages/web/src/api/activity.ts index 8b709b6e6..1a9a7873f 100644 --- a/packages/web/src/api/activity.ts +++ b/packages/web/src/api/activity.ts @@ -2,6 +2,7 @@ import type { AgentType, ClientCapabilities, ConnectTokenResponse, + ProviderModelCatalog, RuntimeAuthMethod, RuntimeProvider, UpdateAttempt, @@ -130,6 +131,16 @@ export function startRuntimeAuth( return api.post(`/clients/${encodeURIComponent(clientId)}/runtime-auth/start`, body); } +/** + * Ask the connected computer's daemon for the host-local model catalog for + * this provider (Cursor `agent models`, Kimi `~/.kimi-code/config.toml`, …). + */ +export function getProviderModels(clientId: string, provider: RuntimeProvider): Promise { + return api.get( + `/clients/${encodeURIComponent(clientId)}/providers/${encodeURIComponent(provider)}/models`, + ); +} + export function resetAgentActivity(agentId: string): Promise<{ reset: boolean }> { // The agent uuid in the path is enough for the server to resolve the // owning org — no `withOrg` needed. diff --git a/packages/web/src/pages/agent-detail/model-section.tsx b/packages/web/src/pages/agent-detail/model-section.tsx index 49c36c93a..b6f735037 100644 --- a/packages/web/src/pages/agent-detail/model-section.tsx +++ b/packages/web/src/pages/agent-detail/model-section.tsx @@ -1,5 +1,6 @@ -import type { RuntimeProvider } from "@first-tree/shared"; +import type { ProviderModelCatalog, RuntimeProvider } from "@first-tree/shared"; import { useEffect, useMemo, useRef, useState } from "react"; +import { getProviderModels } from "../../api/activity.js"; import { Input } from "../../components/ui/input.js"; import { Select, type SelectOption } from "../../components/ui/select.js"; import { ConfigRow } from "./flat-section.js"; @@ -7,6 +8,11 @@ import { ConfigRow } from "./flat-section.js"; /** * Model — an inline dropdown. Selecting a value saves it immediately (onChange), * consistent with every other control on the page; there is no draft / Save Bar. + * + * Phase 1 host-local catalogs (Cursor / Kimi): options come from the bound + * computer's real provider. Fallbacks required by product: + * 1. DEFAULT (empty string) — inherit the provider's local default config + * 2. Custom — free-form exact model id the operator types */ export type ModelOption = SelectOption; @@ -56,15 +62,15 @@ export const CODEX_MODEL_OPTIONS: Array = { value: CODEX_MODEL_IDS.GPT_5_2, label: CODEX_MODEL_IDS.GPT_5_2 }, ]; -const MODEL_OPTIONS_BY_PROVIDER: Record = { +/** Sentinel Select value for the custom free-form path — never written to config. */ +export const CUSTOM_MODEL_SENTINEL = "__custom__"; + +const HOST_LOCAL_MODEL_PROVIDERS: ReadonlySet = new Set(["cursor", "kimi-code"]); + +const CURATED_OPTIONS_BY_PROVIDER: Partial> = { "claude-code": CLAUDE_MODEL_OPTIONS, "claude-code-tui": CLAUDE_MODEL_OPTIONS, codex: CODEX_MODEL_OPTIONS, - // Cursor deliberately has NO curated list: the account-dependent SKU catalog - // is large and shifting, so the control is a free-form exact-id input (see - // ModelSection) and First Tree does not maintain a model matrix. - cursor: [], - "kimi-code": [], }; const MODEL_HELP_BY_PROVIDER: Record = { @@ -72,8 +78,9 @@ const MODEL_HELP_BY_PROVIDER: Record = { "claude-code-tui": "Applies to new sessions immediately. Model swap restarts the tmux session (~2–4s).", codex: "Applies to new sessions immediately. Unset lets the CLI pick by auth mode.", cursor: - "Exact Cursor model id, passed through verbatim on the next turn. Leave empty for the Cursor default (auto). An id your account can't use fails visibly — no silent fallback.", - "kimi-code": "Exact Kimi model id, passed to new sessions. Leave empty to use the model configured in ~/.kimi-code.", + "Pick a model from this computer's Cursor account catalog, leave DEFAULT to inherit Cursor's local default (auto), or enter a custom exact model id. An id your account can't use fails visibly — no silent fallback.", + "kimi-code": + "Pick a model from this computer's ~/.kimi-code config, leave DEFAULT to inherit the local default_model, or enter a custom exact model id.", }; export type ModelSectionProps = { @@ -82,51 +89,227 @@ export type ModelSectionProps = { disabled?: boolean; /** Runtime this agent runs on — drives the option list and help copy. */ provider?: RuntimeProvider; + /** Bound computer id — required to fetch host-local catalogs (Cursor / Kimi). */ + clientId?: string | null; }; -const UNSET_OPTION: ModelOption = { value: "", label: "(unset — inherits local)" }; - -export function ModelSection({ value, onChange, disabled, provider = "claude-code" }: ModelSectionProps) { - const presetOptions = MODEL_OPTIONS_BY_PROVIDER[provider]; +const DEFAULT_OPTION: ModelOption = { value: "", label: "DEFAULT", hint: "inherit local" }; - const items = useMemo(() => { - const list: ModelOption[] = [UNSET_OPTION, ...presetOptions]; +export function ModelSection({ + value, + onChange, + disabled, + provider = "claude-code", + clientId = null, +}: ModelSectionProps) { + const usesHostCatalog = HOST_LOCAL_MODEL_PROVIDERS.has(provider); + const presetOptions = CURATED_OPTIONS_BY_PROVIDER[provider] ?? []; + const curatedItems = useMemo(() => { + const list: ModelOption[] = [DEFAULT_OPTION, ...presetOptions]; if (value !== "" && !presetOptions.some((o) => o.value === value)) { list.push({ value, label: value, hint: "custom" }); } return list; }, [presetOptions, value]); - if (provider === "cursor" || provider === "kimi-code") { + if (usesHostCatalog) { return ( - + ); } return ( - ); } +function HostLocalModelPicker({ + value, + onChange, + disabled, + provider, + clientId, +}: { + value: string; + onChange: (v: string) => void; + disabled?: boolean; + provider: RuntimeProvider; + clientId: string | null; +}) { + const [catalog, setCatalog] = useState(null); + const [loading, setLoading] = useState(false); + const [loadError, setLoadError] = useState(null); + const [customMode, setCustomMode] = useState(false); + + useEffect(() => { + if (!clientId) { + setCatalog(null); + setLoadError(null); + setLoading(false); + return; + } + let cancelled = false; + setLoading(true); + setLoadError(null); + void getProviderModels(clientId, provider) + .then((next) => { + if (cancelled) return; + setCatalog(next); + if (next.source === "unavailable") { + setLoadError(next.error ?? "Could not load models from this computer"); + } + }) + .catch((err: unknown) => { + if (cancelled) return; + setCatalog(null); + setLoadError(err instanceof Error ? err.message : String(err)); + }) + .finally(() => { + if (!cancelled) setLoading(false); + }); + return () => { + cancelled = true; + }; + }, [clientId, provider]); + + const catalogOptions = useMemo(() => { + if (!catalog || catalog.source === "unavailable") return []; + return catalog.models.map((m) => ({ + value: m.id, + label: m.label ?? m.id, + hint: m.hint, + })); + }, [catalog]); + + const valueInCatalog = value !== "" && catalogOptions.some((o) => o.value === value); + const showCustomInput = customMode || (value !== "" && !valueInCatalog); + + useEffect(() => { + if (value !== "" && !valueInCatalog && catalogOptions.length > 0) { + setCustomMode(true); + } + if (value === "") setCustomMode(false); + }, [value, valueInCatalog, catalogOptions.length]); + + const selectValue = showCustomInput ? CUSTOM_MODEL_SENTINEL : value; + + const items = useMemo(() => { + const list: ModelOption[] = [DEFAULT_OPTION, ...catalogOptions]; + list.push({ value: CUSTOM_MODEL_SENTINEL, label: "Custom…", hint: "type exact id" }); + if (value !== "" && !catalogOptions.some((o) => o.value === value) && value !== CUSTOM_MODEL_SENTINEL) { + // Keep the current custom id visible in the list while editing. + list.splice(list.length - 1, 0, { value, label: value, hint: "custom" }); + } + return list; + }, [catalogOptions, value]); + + const onSelectChange = (next: string) => { + if (next === CUSTOM_MODEL_SENTINEL) { + setCustomMode(true); + return; + } + setCustomMode(false); + onChange(next); + }; + + // No computer bound / catalog unavailable → free-form with DEFAULT empty via clearing. + if (!clientId || (catalog?.source === "unavailable" && catalogOptions.length === 0 && !loading)) { + return ( +
+ 8} + mono + aria-label="Model" + placeholder={loading ? "Loading models…" : undefined} + /> + {showCustomInput ? ( + + ) : null} + {loadError && catalogOptions.length > 0 ? ( +

+ Catalog warning: {loadError} +

+ ) : null} +
+ ); +} + /** - * Free-form exact-model input (cursor). Commits on blur / Enter rather than - * per keystroke — every other control on this page saves on change, but a - * text field saving each character would spam config writes. + * Free-form exact-model input. Commits on blur / Enter rather than per + * keystroke — every other control on this page saves on change, but a text + * field saving each character would spam config writes. */ function FreeFormModelInput({ value, onChange, disabled, provider, + placeholder, }: { value: string; onChange: (v: string) => void; disabled?: boolean; - provider: "cursor" | "kimi-code"; + provider: RuntimeProvider; + placeholder?: string; }) { const [draft, setDraft] = useState(value); // Latch scoped to ONE Enter→blur sequence: Enter followed by the immediate @@ -160,9 +343,11 @@ function FreeFormModelInput({ if (e.key === "Enter") commit(); }} disabled={disabled} - placeholder={provider === "kimi-code" ? "local Kimi default" : "auto (Cursor default)"} + placeholder={ + placeholder ?? (provider === "kimi-code" ? "local Kimi default" : "auto (Cursor default)") + } className="font-mono" - aria-label="Model" + aria-label="Custom model id" /> ); } diff --git a/packages/web/src/pages/agent-detail/runtime-tab.tsx b/packages/web/src/pages/agent-detail/runtime-tab.tsx index 7eb19f6e9..91029bb6c 100644 --- a/packages/web/src/pages/agent-detail/runtime-tab.tsx +++ b/packages/web/src/pages/agent-detail/runtime-tab.tsx @@ -78,6 +78,7 @@ export function RuntimeTab() { onChange={(v) => configSave.save({ model: v }, { field: "model" })} disabled={editsDisabled} provider={ctx.setupRuntimeProvider} + clientId={ctx.clientStatus?.clientId ?? ctx.agent.clientId} /> {/* Cursor has no separate reasoning-effort channel — effort/fast variants live in the provider-native model id, so the control diff --git a/packages/web/src/pages/clients/__tests__/cursor-provider-surfaces.test.tsx b/packages/web/src/pages/clients/__tests__/cursor-provider-surfaces.test.tsx index a73957915..24b3a6af1 100644 --- a/packages/web/src/pages/clients/__tests__/cursor-provider-surfaces.test.tsx +++ b/packages/web/src/pages/clients/__tests__/cursor-provider-surfaces.test.tsx @@ -70,15 +70,14 @@ describe("cursor provider — setup card surfaces", () => { }); }); -describe("cursor provider — free-form model input", () => { - it("renders a text input with the auto hint and commits the exact id on blur", async () => { +describe("cursor provider — DEFAULT + custom model id fallback", () => { + it("renders a custom model id input and commits the exact id on blur", async () => { const saved: string[] = []; const el = await render( saved.push(v)} provider="cursor" />); - const input = el.querySelector('input[aria-label="Model"]'); + const input = el.querySelector('input[aria-label="Custom model id"]'); expect(input).not.toBeNull(); if (!input) throw new Error("unreachable"); - expect(input.placeholder).toContain("auto"); await act(async () => { // React tracks the value setter — go through the native prototype setter @@ -120,6 +119,6 @@ describe("cursor provider — free-form model input", () => { it("keeps the dropdown for claude/codex providers (no free-form regression)", async () => { const el = await render( {}} provider="claude-code" />); - expect(el.querySelector('input[aria-label="Model"]')).toBeNull(); + expect(el.querySelector('input[aria-label="Custom model id"]')).toBeNull(); }); }); diff --git a/packages/web/src/pages/clients/__tests__/kimi-code-provider-surfaces.test.tsx b/packages/web/src/pages/clients/__tests__/kimi-code-provider-surfaces.test.tsx index fd6c7be87..eeaa11cef 100644 --- a/packages/web/src/pages/clients/__tests__/kimi-code-provider-surfaces.test.tsx +++ b/packages/web/src/pages/clients/__tests__/kimi-code-provider-surfaces.test.tsx @@ -42,13 +42,13 @@ describe("Kimi Code provider surfaces", () => { expect(providerInstallHint("kimi-code", "darwin")).toContain("bundled with First Tree"); }); - it("uses the free-form model field with a local-Kimi default hint", async () => { + it("uses the free-form custom model field when no computer catalog is bound", async () => { const saved: string[] = []; const element = await render( saved.push(value)} provider="kimi-code" />, ); - const input = element.querySelector('input[aria-label="Model"]'); - expect(input?.placeholder).toContain("local Kimi default"); + const input = element.querySelector('input[aria-label="Custom model id"]'); + expect(input).not.toBeNull(); if (!input) throw new Error("model input missing"); await act(async () => { From c6ed8e2ca9f811c8408f1f6b9277b475db0bb6cb Mon Sep 17 00:00:00 2001 From: YueZengwu Date: Fri, 17 Jul 2026 20:43:33 +0800 Subject: [PATCH 2/5] fix(runtime): satisfy CodeQL and branch/lint CI for model catalogs Rewrite Cursor/Kimi parsers without polynomial regexes, fix Biome export order/formatting, and prepare a feat/-prefixed branch for CI. Co-authored-by: Cursor --- packages/client/src/client-connection.ts | 2 +- .../runtime/capabilities/discover-models.ts | 103 +++++++++++++----- packages/server/src/api/agent/ws-client.ts | 8 +- packages/server/src/api/clients.ts | 9 +- packages/shared/src/index.ts | 28 ++--- packages/web/src/api/activity.ts | 4 +- .../src/pages/agent-detail/model-section.tsx | 13 +-- 7 files changed, 101 insertions(+), 66 deletions(-) diff --git a/packages/client/src/client-connection.ts b/packages/client/src/client-connection.ts index 7c7dcc988..0a51c0d52 100644 --- a/packages/client/src/client-connection.ts +++ b/packages/client/src/client-connection.ts @@ -19,11 +19,11 @@ import { PROVIDER_MODELS_LIST_TYPE, PROVIDER_MODELS_RESULT_TYPE, type ProviderModelCatalog, + providerModelsListCommandSchema, RUNTIME_AUTH_START_TYPE, type RuntimeAuthMethod, type RuntimeProvider, type RuntimeState, - providerModelsListCommandSchema, runtimeAuthStartCommandSchema, type ServerWelcomeFrame, type SessionEvent, diff --git a/packages/client/src/runtime/capabilities/discover-models.ts b/packages/client/src/runtime/capabilities/discover-models.ts index 6d11b58a8..e616c2e4e 100644 --- a/packages/client/src/runtime/capabilities/discover-models.ts +++ b/packages/client/src/runtime/capabilities/discover-models.ts @@ -1,6 +1,6 @@ +import { readFile } from "node:fs/promises"; import { homedir } from "node:os"; import { join } from "node:path"; -import { readFile } from "node:fs/promises"; import type { ProviderModelCatalog, ProviderModelOption, RuntimeProvider } from "@first-tree/shared"; import { findCursorExecutableOnPath } from "../cursor-binary.js"; import { runCommand } from "./launch-probe.js"; @@ -12,7 +12,10 @@ export type DiscoverModelsDeps = { env?: NodeJS.ProcessEnv; now?: () => Date; findCursorBinary?: (env?: Record) => string | null; - runCursorModels?: (binary: string, env: NodeJS.ProcessEnv) => Promise<{ ok: boolean; stdout: string; stderr: string }>; + runCursorModels?: ( + binary: string, + env: NodeJS.ProcessEnv, + ) => Promise<{ ok: boolean; stdout: string; stderr: string }>; readKimiConfig?: () => Promise; kimiConfigPath?: string; }; @@ -37,6 +40,9 @@ function unavailable(provider: RuntimeProvider, error: string, deps: DiscoverMod * Available models * auto - Auto (default) * gpt-5.2 - GPT-5.2 + * + * Uses indexOf/slice instead of `\s+` / `.+` regexes so CodeQL does not flag + * polynomial-time matching on CLI stdout. */ export function parseCursorModelsOutput(stdout: string): { models: ProviderModelOption[]; @@ -46,16 +52,24 @@ export function parseCursorModelsOutput(stdout: string): { let defaultModelId: string | null = null; for (const rawLine of stdout.split(/\r?\n/)) { const line = rawLine.trim(); - if (!line || /^available models$/i.test(line)) continue; - const match = /^(\S+)\s+-\s+(.+)$/.exec(line); - if (!match) continue; - const id = match[1]!; - const label = match[2]!.trim(); - const isDefault = /\(default\)/i.test(label); - if (isDefault) defaultModelId = id; + if (!line) continue; + if (line.toLowerCase() === "available models") continue; + const sep = line.indexOf(" - "); + if (sep <= 0) continue; + const id = line.slice(0, sep); + // Model ids are single tokens (`auto`, `gpt-5.2`); reject spaced ids. + if (!id || id.includes(" ") || id.includes("\t")) continue; + let label = line.slice(sep + 3).trim(); + const defaultMarker = "(default)"; + const defaultAt = label.toLowerCase().indexOf(defaultMarker); + const isDefault = defaultAt >= 0; + if (isDefault) { + defaultModelId = id; + label = `${label.slice(0, defaultAt)}${label.slice(defaultAt + defaultMarker.length)}`.trim(); + } models.push({ id, - label: label.replace(/\s*\(default\)\s*/i, "").trim() || id, + label: label || id, ...(isDefault ? { isDefault: true, hint: "default" } : {}), }); } @@ -75,24 +89,59 @@ export function parseKimiConfigModels(toml: string): { models: ProviderModelOption[]; defaultModelId: string | null; } { - const defaultMatch = /^default_model\s*=\s*"([^"]+)"/m.exec(toml); - const defaultModelId = defaultMatch?.[1] ?? null; - + let defaultModelId: string | null = null; const models: ProviderModelOption[] = []; - const sectionRe = /^\[models\."([^"]+)"\]\s*$/gm; const sections: Array<{ id: string; headerStart: number; bodyStart: number }> = []; - for (const match of toml.matchAll(sectionRe)) { - sections.push({ - id: match[1]!, - headerStart: match.index!, - bodyStart: match.index! + match[0].length, - }); + + // Line-oriented scan avoids `\s*` / `.+` regexes that CodeQL flags as + // polynomial on untrusted host config text. + let cursor = 0; + while (cursor <= toml.length) { + const nextNl = toml.indexOf("\n", cursor); + const lineEnd = nextNl === -1 ? toml.length : nextNl; + const line = toml.slice(cursor, lineEnd).replace(/\r$/, ""); + const trimmed = line.trim(); + + if (trimmed.startsWith("default_model")) { + const eq = trimmed.indexOf("="); + if (eq > 0) { + const raw = trimmed.slice(eq + 1).trim(); + if (raw.startsWith('"') && raw.endsWith('"') && raw.length >= 2) { + defaultModelId = raw.slice(1, -1); + } + } + } + + if (trimmed.startsWith('[models."') && trimmed.endsWith('"]')) { + const id = trimmed.slice('[models."'.length, -'"]'.length); + if (id.length > 0) { + const bodyStart = nextNl === -1 ? toml.length : nextNl + 1; + sections.push({ id, headerStart: cursor, bodyStart }); + } + } + + if (nextNl === -1) break; + cursor = nextNl + 1; } + for (let i = 0; i < sections.length; i++) { - const section = sections[i]!; - const bodyEnd = i + 1 < sections.length ? sections[i + 1]!.headerStart : toml.length; + const section = sections[i]; + if (!section) continue; + const next = sections[i + 1]; + const bodyEnd = next ? next.headerStart : toml.length; const body = toml.slice(section.bodyStart, bodyEnd); - const displayName = /^display_name\s*=\s*"([^"]+)"/m.exec(body)?.[1]; + let displayName: string | undefined; + for (const bodyLine of body.split("\n")) { + const t = bodyLine.replace(/\r$/, "").trim(); + if (!t.startsWith("display_name")) continue; + const eq = t.indexOf("="); + if (eq <= 0) continue; + const raw = t.slice(eq + 1).trim(); + if (raw.startsWith('"') && raw.endsWith('"') && raw.length >= 2) { + displayName = raw.slice(1, -1); + } + break; + } const isDefault = defaultModelId === section.id; models.push({ id: section.id, @@ -159,7 +208,7 @@ async function discoverKimiModels(deps: DiscoverModelsDeps): Promise { - return api.get( - `/clients/${encodeURIComponent(clientId)}/providers/${encodeURIComponent(provider)}/models`, - ); + return api.get(`/clients/${encodeURIComponent(clientId)}/providers/${encodeURIComponent(provider)}/models`); } export function resetAgentActivity(agentId: string): Promise<{ reset: boolean }> { diff --git a/packages/web/src/pages/agent-detail/model-section.tsx b/packages/web/src/pages/agent-detail/model-section.tsx index b6f735037..c1702ebfc 100644 --- a/packages/web/src/pages/agent-detail/model-section.tsx +++ b/packages/web/src/pages/agent-detail/model-section.tsx @@ -128,14 +128,7 @@ export function ModelSection({ return ( - ); } @@ -343,9 +336,7 @@ function FreeFormModelInput({ if (e.key === "Enter") commit(); }} disabled={disabled} - placeholder={ - placeholder ?? (provider === "kimi-code" ? "local Kimi default" : "auto (Cursor default)") - } + placeholder={placeholder ?? (provider === "kimi-code" ? "local Kimi default" : "auto (Cursor default)")} className="font-mono" aria-label="Custom model id" /> From 655dcf786199378fdf01765f4ae85e9cad55868e Mon Sep 17 00:00:00 2001 From: YueZengwu Date: Fri, 17 Jul 2026 20:44:29 +0800 Subject: [PATCH 3/5] refactor: leave Web Model picker to kiven-fte Keep host-local discovery + GET models API in this PR; revert web UI so frontend ownership stays with kiven against the shared catalog contract. Co-authored-by: Cursor --- packages/web/src/api/activity.ts | 9 - .../src/pages/agent-detail/model-section.tsx | 224 ++---------------- .../src/pages/agent-detail/runtime-tab.tsx | 1 - .../cursor-provider-surfaces.test.tsx | 9 +- .../kimi-code-provider-surfaces.test.tsx | 6 +- 5 files changed, 32 insertions(+), 217 deletions(-) diff --git a/packages/web/src/api/activity.ts b/packages/web/src/api/activity.ts index c1a8f60e9..8b709b6e6 100644 --- a/packages/web/src/api/activity.ts +++ b/packages/web/src/api/activity.ts @@ -2,7 +2,6 @@ import type { AgentType, ClientCapabilities, ConnectTokenResponse, - ProviderModelCatalog, RuntimeAuthMethod, RuntimeProvider, UpdateAttempt, @@ -131,14 +130,6 @@ export function startRuntimeAuth( return api.post(`/clients/${encodeURIComponent(clientId)}/runtime-auth/start`, body); } -/** - * Ask the connected computer's daemon for the host-local model catalog for - * this provider (Cursor `agent models`, Kimi `~/.kimi-code/config.toml`, …). - */ -export function getProviderModels(clientId: string, provider: RuntimeProvider): Promise { - return api.get(`/clients/${encodeURIComponent(clientId)}/providers/${encodeURIComponent(provider)}/models`); -} - export function resetAgentActivity(agentId: string): Promise<{ reset: boolean }> { // The agent uuid in the path is enough for the server to resolve the // owning org — no `withOrg` needed. diff --git a/packages/web/src/pages/agent-detail/model-section.tsx b/packages/web/src/pages/agent-detail/model-section.tsx index c1702ebfc..49c36c93a 100644 --- a/packages/web/src/pages/agent-detail/model-section.tsx +++ b/packages/web/src/pages/agent-detail/model-section.tsx @@ -1,6 +1,5 @@ -import type { ProviderModelCatalog, RuntimeProvider } from "@first-tree/shared"; +import type { RuntimeProvider } from "@first-tree/shared"; import { useEffect, useMemo, useRef, useState } from "react"; -import { getProviderModels } from "../../api/activity.js"; import { Input } from "../../components/ui/input.js"; import { Select, type SelectOption } from "../../components/ui/select.js"; import { ConfigRow } from "./flat-section.js"; @@ -8,11 +7,6 @@ import { ConfigRow } from "./flat-section.js"; /** * Model — an inline dropdown. Selecting a value saves it immediately (onChange), * consistent with every other control on the page; there is no draft / Save Bar. - * - * Phase 1 host-local catalogs (Cursor / Kimi): options come from the bound - * computer's real provider. Fallbacks required by product: - * 1. DEFAULT (empty string) — inherit the provider's local default config - * 2. Custom — free-form exact model id the operator types */ export type ModelOption = SelectOption; @@ -62,15 +56,15 @@ export const CODEX_MODEL_OPTIONS: Array = { value: CODEX_MODEL_IDS.GPT_5_2, label: CODEX_MODEL_IDS.GPT_5_2 }, ]; -/** Sentinel Select value for the custom free-form path — never written to config. */ -export const CUSTOM_MODEL_SENTINEL = "__custom__"; - -const HOST_LOCAL_MODEL_PROVIDERS: ReadonlySet = new Set(["cursor", "kimi-code"]); - -const CURATED_OPTIONS_BY_PROVIDER: Partial> = { +const MODEL_OPTIONS_BY_PROVIDER: Record = { "claude-code": CLAUDE_MODEL_OPTIONS, "claude-code-tui": CLAUDE_MODEL_OPTIONS, codex: CODEX_MODEL_OPTIONS, + // Cursor deliberately has NO curated list: the account-dependent SKU catalog + // is large and shifting, so the control is a free-form exact-id input (see + // ModelSection) and First Tree does not maintain a model matrix. + cursor: [], + "kimi-code": [], }; const MODEL_HELP_BY_PROVIDER: Record = { @@ -78,9 +72,8 @@ const MODEL_HELP_BY_PROVIDER: Record = { "claude-code-tui": "Applies to new sessions immediately. Model swap restarts the tmux session (~2–4s).", codex: "Applies to new sessions immediately. Unset lets the CLI pick by auth mode.", cursor: - "Pick a model from this computer's Cursor account catalog, leave DEFAULT to inherit Cursor's local default (auto), or enter a custom exact model id. An id your account can't use fails visibly — no silent fallback.", - "kimi-code": - "Pick a model from this computer's ~/.kimi-code config, leave DEFAULT to inherit the local default_model, or enter a custom exact model id.", + "Exact Cursor model id, passed through verbatim on the next turn. Leave empty for the Cursor default (auto). An id your account can't use fails visibly — no silent fallback.", + "kimi-code": "Exact Kimi model id, passed to new sessions. Leave empty to use the model configured in ~/.kimi-code.", }; export type ModelSectionProps = { @@ -89,220 +82,51 @@ export type ModelSectionProps = { disabled?: boolean; /** Runtime this agent runs on — drives the option list and help copy. */ provider?: RuntimeProvider; - /** Bound computer id — required to fetch host-local catalogs (Cursor / Kimi). */ - clientId?: string | null; }; -const DEFAULT_OPTION: ModelOption = { value: "", label: "DEFAULT", hint: "inherit local" }; +const UNSET_OPTION: ModelOption = { value: "", label: "(unset — inherits local)" }; -export function ModelSection({ - value, - onChange, - disabled, - provider = "claude-code", - clientId = null, -}: ModelSectionProps) { - const usesHostCatalog = HOST_LOCAL_MODEL_PROVIDERS.has(provider); - const presetOptions = CURATED_OPTIONS_BY_PROVIDER[provider] ?? []; - const curatedItems = useMemo(() => { - const list: ModelOption[] = [DEFAULT_OPTION, ...presetOptions]; +export function ModelSection({ value, onChange, disabled, provider = "claude-code" }: ModelSectionProps) { + const presetOptions = MODEL_OPTIONS_BY_PROVIDER[provider]; + + const items = useMemo(() => { + const list: ModelOption[] = [UNSET_OPTION, ...presetOptions]; if (value !== "" && !presetOptions.some((o) => o.value === value)) { list.push({ value, label: value, hint: "custom" }); } return list; }, [presetOptions, value]); - if (usesHostCatalog) { + if (provider === "cursor" || provider === "kimi-code") { return ( - + ); } return ( - ); } -function HostLocalModelPicker({ - value, - onChange, - disabled, - provider, - clientId, -}: { - value: string; - onChange: (v: string) => void; - disabled?: boolean; - provider: RuntimeProvider; - clientId: string | null; -}) { - const [catalog, setCatalog] = useState(null); - const [loading, setLoading] = useState(false); - const [loadError, setLoadError] = useState(null); - const [customMode, setCustomMode] = useState(false); - - useEffect(() => { - if (!clientId) { - setCatalog(null); - setLoadError(null); - setLoading(false); - return; - } - let cancelled = false; - setLoading(true); - setLoadError(null); - void getProviderModels(clientId, provider) - .then((next) => { - if (cancelled) return; - setCatalog(next); - if (next.source === "unavailable") { - setLoadError(next.error ?? "Could not load models from this computer"); - } - }) - .catch((err: unknown) => { - if (cancelled) return; - setCatalog(null); - setLoadError(err instanceof Error ? err.message : String(err)); - }) - .finally(() => { - if (!cancelled) setLoading(false); - }); - return () => { - cancelled = true; - }; - }, [clientId, provider]); - - const catalogOptions = useMemo(() => { - if (!catalog || catalog.source === "unavailable") return []; - return catalog.models.map((m) => ({ - value: m.id, - label: m.label ?? m.id, - hint: m.hint, - })); - }, [catalog]); - - const valueInCatalog = value !== "" && catalogOptions.some((o) => o.value === value); - const showCustomInput = customMode || (value !== "" && !valueInCatalog); - - useEffect(() => { - if (value !== "" && !valueInCatalog && catalogOptions.length > 0) { - setCustomMode(true); - } - if (value === "") setCustomMode(false); - }, [value, valueInCatalog, catalogOptions.length]); - - const selectValue = showCustomInput ? CUSTOM_MODEL_SENTINEL : value; - - const items = useMemo(() => { - const list: ModelOption[] = [DEFAULT_OPTION, ...catalogOptions]; - list.push({ value: CUSTOM_MODEL_SENTINEL, label: "Custom…", hint: "type exact id" }); - if (value !== "" && !catalogOptions.some((o) => o.value === value) && value !== CUSTOM_MODEL_SENTINEL) { - // Keep the current custom id visible in the list while editing. - list.splice(list.length - 1, 0, { value, label: value, hint: "custom" }); - } - return list; - }, [catalogOptions, value]); - - const onSelectChange = (next: string) => { - if (next === CUSTOM_MODEL_SENTINEL) { - setCustomMode(true); - return; - } - setCustomMode(false); - onChange(next); - }; - - // No computer bound / catalog unavailable → free-form with DEFAULT empty via clearing. - if (!clientId || (catalog?.source === "unavailable" && catalogOptions.length === 0 && !loading)) { - return ( -
- 8} - mono - aria-label="Model" - placeholder={loading ? "Loading models…" : undefined} - /> - {showCustomInput ? ( - - ) : null} - {loadError && catalogOptions.length > 0 ? ( -

- Catalog warning: {loadError} -

- ) : null} -
- ); -} - /** - * Free-form exact-model input. Commits on blur / Enter rather than per - * keystroke — every other control on this page saves on change, but a text - * field saving each character would spam config writes. + * Free-form exact-model input (cursor). Commits on blur / Enter rather than + * per keystroke — every other control on this page saves on change, but a + * text field saving each character would spam config writes. */ function FreeFormModelInput({ value, onChange, disabled, provider, - placeholder, }: { value: string; onChange: (v: string) => void; disabled?: boolean; - provider: RuntimeProvider; - placeholder?: string; + provider: "cursor" | "kimi-code"; }) { const [draft, setDraft] = useState(value); // Latch scoped to ONE Enter→blur sequence: Enter followed by the immediate @@ -336,9 +160,9 @@ function FreeFormModelInput({ if (e.key === "Enter") commit(); }} disabled={disabled} - placeholder={placeholder ?? (provider === "kimi-code" ? "local Kimi default" : "auto (Cursor default)")} + placeholder={provider === "kimi-code" ? "local Kimi default" : "auto (Cursor default)"} className="font-mono" - aria-label="Custom model id" + aria-label="Model" /> ); } diff --git a/packages/web/src/pages/agent-detail/runtime-tab.tsx b/packages/web/src/pages/agent-detail/runtime-tab.tsx index 91029bb6c..7eb19f6e9 100644 --- a/packages/web/src/pages/agent-detail/runtime-tab.tsx +++ b/packages/web/src/pages/agent-detail/runtime-tab.tsx @@ -78,7 +78,6 @@ export function RuntimeTab() { onChange={(v) => configSave.save({ model: v }, { field: "model" })} disabled={editsDisabled} provider={ctx.setupRuntimeProvider} - clientId={ctx.clientStatus?.clientId ?? ctx.agent.clientId} /> {/* Cursor has no separate reasoning-effort channel — effort/fast variants live in the provider-native model id, so the control diff --git a/packages/web/src/pages/clients/__tests__/cursor-provider-surfaces.test.tsx b/packages/web/src/pages/clients/__tests__/cursor-provider-surfaces.test.tsx index 24b3a6af1..a73957915 100644 --- a/packages/web/src/pages/clients/__tests__/cursor-provider-surfaces.test.tsx +++ b/packages/web/src/pages/clients/__tests__/cursor-provider-surfaces.test.tsx @@ -70,14 +70,15 @@ describe("cursor provider — setup card surfaces", () => { }); }); -describe("cursor provider — DEFAULT + custom model id fallback", () => { - it("renders a custom model id input and commits the exact id on blur", async () => { +describe("cursor provider — free-form model input", () => { + it("renders a text input with the auto hint and commits the exact id on blur", async () => { const saved: string[] = []; const el = await render( saved.push(v)} provider="cursor" />); - const input = el.querySelector('input[aria-label="Custom model id"]'); + const input = el.querySelector('input[aria-label="Model"]'); expect(input).not.toBeNull(); if (!input) throw new Error("unreachable"); + expect(input.placeholder).toContain("auto"); await act(async () => { // React tracks the value setter — go through the native prototype setter @@ -119,6 +120,6 @@ describe("cursor provider — DEFAULT + custom model id fallback", () => { it("keeps the dropdown for claude/codex providers (no free-form regression)", async () => { const el = await render( {}} provider="claude-code" />); - expect(el.querySelector('input[aria-label="Custom model id"]')).toBeNull(); + expect(el.querySelector('input[aria-label="Model"]')).toBeNull(); }); }); diff --git a/packages/web/src/pages/clients/__tests__/kimi-code-provider-surfaces.test.tsx b/packages/web/src/pages/clients/__tests__/kimi-code-provider-surfaces.test.tsx index eeaa11cef..fd6c7be87 100644 --- a/packages/web/src/pages/clients/__tests__/kimi-code-provider-surfaces.test.tsx +++ b/packages/web/src/pages/clients/__tests__/kimi-code-provider-surfaces.test.tsx @@ -42,13 +42,13 @@ describe("Kimi Code provider surfaces", () => { expect(providerInstallHint("kimi-code", "darwin")).toContain("bundled with First Tree"); }); - it("uses the free-form custom model field when no computer catalog is bound", async () => { + it("uses the free-form model field with a local-Kimi default hint", async () => { const saved: string[] = []; const element = await render( saved.push(value)} provider="kimi-code" />, ); - const input = element.querySelector('input[aria-label="Custom model id"]'); - expect(input).not.toBeNull(); + const input = element.querySelector('input[aria-label="Model"]'); + expect(input?.placeholder).toContain("local Kimi default"); if (!input) throw new Error("model input missing"); await act(async () => { From d988e2bfa3ca0b352f51b728074cdd1b19042087 Mon Sep 17 00:00:00 2001 From: YueZengwu Date: Fri, 17 Jul 2026 20:52:09 +0800 Subject: [PATCH 4/5] fix(runtime): cross-replica model catalog RPC and Kimi discovery Route provider-models commands/results across server replicas via PG NOTIFY plus durable clients.metadata rendezvous, and parse Kimi catalogs with smol-toml including KIMI_CODE_HOME and bare model aliases. Co-authored-by: Cursor --- packages/client/package.json | 1 + .../src/__tests__/discover-models.test.ts | 49 ++++++- .../runtime/capabilities/discover-models.ts | 92 +++++-------- .../__tests__/clients-provider-models.test.ts | 73 ++++++++++ .../src/__tests__/notifier-extra.test.ts | 56 +++++++- packages/server/src/api/agent/ws-client.ts | 31 ++++- packages/server/src/api/clients.ts | 32 +++-- packages/server/src/services/notifier.ts | 127 ++++++++++++++++++ .../src/services/provider-models-rpc.ts | 101 ++++++++++++++ pnpm-lock.yaml | 23 +++- 10 files changed, 508 insertions(+), 77 deletions(-) create mode 100644 packages/server/src/services/provider-models-rpc.ts diff --git a/packages/client/package.json b/packages/client/package.json index d85e429e2..1891db3d3 100644 --- a/packages/client/package.json +++ b/packages/client/package.json @@ -37,6 +37,7 @@ "mdast-util-from-markdown": "^2.0.3", "pino": "^9.5.0", "semver": "^7.6.3", + "smol-toml": "^1.7.0", "ws": "^8.20.0", "yaml": "^2.8.3", "zod": "^4.0.0" diff --git a/packages/client/src/__tests__/discover-models.test.ts b/packages/client/src/__tests__/discover-models.test.ts index a54b96689..efd9b5235 100644 --- a/packages/client/src/__tests__/discover-models.test.ts +++ b/packages/client/src/__tests__/discover-models.test.ts @@ -1,5 +1,9 @@ import { describe, expect, it } from "vitest"; -import { parseCursorModelsOutput, parseKimiConfigModels } from "../runtime/capabilities/discover-models.js"; +import { + parseCursorModelsOutput, + parseKimiConfigModels, + resolveKimiConfigPath, +} from "../runtime/capabilities/discover-models.js"; describe("parseCursorModelsOutput", () => { it("parses id/label rows and marks the default", () => { @@ -18,8 +22,25 @@ composer-2.5 - Composer 2.5 }); }); +describe("resolveKimiConfigPath", () => { + it("defaults to ~/.kimi-code/config.toml", () => { + expect(resolveKimiConfigPath({}, "/home/user")).toBe("/home/user/.kimi-code/config.toml"); + }); + + it("honors KIMI_CODE_HOME relocation", () => { + expect(resolveKimiConfigPath({ KIMI_CODE_HOME: "/opt/kimi" }, "/home/user")).toBe("/opt/kimi/config.toml"); + }); + + it("trims KIMI_CODE_HOME and ignores empty", () => { + expect(resolveKimiConfigPath({ KIMI_CODE_HOME: " /custom/kimi " }, "/home/user")).toBe( + "/custom/kimi/config.toml", + ); + expect(resolveKimiConfigPath({ KIMI_CODE_HOME: " " }, "/home/user")).toBe("/home/user/.kimi-code/config.toml"); + }); +}); + describe("parseKimiConfigModels", () => { - it("reads default_model and [models.\".\"] sections", () => { + it('reads default_model and quoted [models."."] sections', () => { const parsed = parseKimiConfigModels(` default_model = "kimi-code/k3" @@ -39,4 +60,28 @@ display_name = "K3" { id: "kimi-code/k3", label: "K3", isDefault: true, hint: "default" }, ]); }); + + it("accepts bare model aliases and marks default_model", () => { + const parsed = parseKimiConfigModels(` +default_model = "gemini-3-pro-preview" + +[models.gemini-3-pro-preview] +provider = "openai" +model = "gemini-3-pro-preview" +display_name = "Gemini 3 Pro" + +[models."kimi-code/k3"] +provider = "managed:kimi-code" +model = "k3" +display_name = "K3" +`); + expect(parsed.defaultModelId).toBe("gemini-3-pro-preview"); + expect(parsed.models).toEqual( + expect.arrayContaining([ + { id: "gemini-3-pro-preview", label: "Gemini 3 Pro", isDefault: true, hint: "default" }, + { id: "kimi-code/k3", label: "K3" }, + ]), + ); + expect(parsed.models).toHaveLength(2); + }); }); diff --git a/packages/client/src/runtime/capabilities/discover-models.ts b/packages/client/src/runtime/capabilities/discover-models.ts index e616c2e4e..199b9407a 100644 --- a/packages/client/src/runtime/capabilities/discover-models.ts +++ b/packages/client/src/runtime/capabilities/discover-models.ts @@ -2,6 +2,7 @@ import { readFile } from "node:fs/promises"; import { homedir } from "node:os"; import { join } from "node:path"; import type { ProviderModelCatalog, ProviderModelOption, RuntimeProvider } from "@first-tree/shared"; +import { parse as parseToml } from "smol-toml"; import { findCursorExecutableOnPath } from "../cursor-binary.js"; import { runCommand } from "./launch-probe.js"; @@ -77,74 +78,35 @@ export function parseCursorModelsOutput(stdout: string): { } /** - * Minimal TOML extract for Kimi's `~/.kimi-code/config.toml`: - * default_model = "kimi-code/k3" - * [models."kimi-code/k3"] - * display_name = "K3" - * model = "k3" - * - * Avoids adding a TOML dependency for a narrow, stable config shape. + * Parse Kimi Code `config.toml` model tables via a real TOML parser so we + * accept both quoted headers (`[models."kimi-code/k3"]`) and bare aliases + * (`[models.gemini-3-pro-preview]`) documented by Kimi. */ export function parseKimiConfigModels(toml: string): { models: ProviderModelOption[]; defaultModelId: string | null; } { - let defaultModelId: string | null = null; - const models: ProviderModelOption[] = []; - const sections: Array<{ id: string; headerStart: number; bodyStart: number }> = []; - - // Line-oriented scan avoids `\s*` / `.+` regexes that CodeQL flags as - // polynomial on untrusted host config text. - let cursor = 0; - while (cursor <= toml.length) { - const nextNl = toml.indexOf("\n", cursor); - const lineEnd = nextNl === -1 ? toml.length : nextNl; - const line = toml.slice(cursor, lineEnd).replace(/\r$/, ""); - const trimmed = line.trim(); - - if (trimmed.startsWith("default_model")) { - const eq = trimmed.indexOf("="); - if (eq > 0) { - const raw = trimmed.slice(eq + 1).trim(); - if (raw.startsWith('"') && raw.endsWith('"') && raw.length >= 2) { - defaultModelId = raw.slice(1, -1); - } - } - } - - if (trimmed.startsWith('[models."') && trimmed.endsWith('"]')) { - const id = trimmed.slice('[models."'.length, -'"]'.length); - if (id.length > 0) { - const bodyStart = nextNl === -1 ? toml.length : nextNl + 1; - sections.push({ id, headerStart: cursor, bodyStart }); - } - } + let data: Record; + try { + data = parseToml(toml) as Record; + } catch { + return { models: [], defaultModelId: null }; + } - if (nextNl === -1) break; - cursor = nextNl + 1; + const defaultModelId = typeof data.default_model === "string" ? data.default_model : null; + const modelsRaw = data.models; + if (!modelsRaw || typeof modelsRaw !== "object" || Array.isArray(modelsRaw)) { + return { models: [], defaultModelId }; } - for (let i = 0; i < sections.length; i++) { - const section = sections[i]; - if (!section) continue; - const next = sections[i + 1]; - const bodyEnd = next ? next.headerStart : toml.length; - const body = toml.slice(section.bodyStart, bodyEnd); - let displayName: string | undefined; - for (const bodyLine of body.split("\n")) { - const t = bodyLine.replace(/\r$/, "").trim(); - if (!t.startsWith("display_name")) continue; - const eq = t.indexOf("="); - if (eq <= 0) continue; - const raw = t.slice(eq + 1).trim(); - if (raw.startsWith('"') && raw.endsWith('"') && raw.length >= 2) { - displayName = raw.slice(1, -1); - } - break; - } - const isDefault = defaultModelId === section.id; + const models: ProviderModelOption[] = []; + for (const [id, value] of Object.entries(modelsRaw as Record)) { + if (!id || typeof value !== "object" || value === null || Array.isArray(value)) continue; + const row = value as Record; + const displayName = typeof row.display_name === "string" ? row.display_name : undefined; + const isDefault = defaultModelId === id; models.push({ - id: section.id, + id, ...(displayName ? { label: displayName } : {}), ...(isDefault ? { isDefault: true, hint: "default" } : {}), }); @@ -152,6 +114,13 @@ export function parseKimiConfigModels(toml: string): { return { models, defaultModelId }; } +/** Effective Kimi config path: `$KIMI_CODE_HOME/config.toml` or `~/.kimi-code/config.toml`. */ +export function resolveKimiConfigPath(env: NodeJS.ProcessEnv = process.env, home: string = homedir()): string { + const custom = env.KIMI_CODE_HOME?.trim(); + const root = custom && custom.length > 0 ? custom : join(home, ".kimi-code"); + return join(root, "config.toml"); +} + async function discoverCursorModels(deps: DiscoverModelsDeps): Promise { const env = deps.env ?? process.env; const findBinary = deps.findCursorBinary ?? findCursorExecutableOnPath; @@ -185,7 +154,8 @@ async function discoverCursorModels(deps: DiscoverModelsDeps): Promise { - const path = deps.kimiConfigPath ?? join(homedir(), ".kimi-code", "config.toml"); + const env = deps.env ?? process.env; + const path = deps.kimiConfigPath ?? resolveKimiConfigPath(env); const read = deps.readKimiConfig ?? (async () => { @@ -208,7 +178,7 @@ async function discoverKimiModels(deps: DiscoverModelsDeps): Promise { }); expect(res.statusCode).toBe(503); }); + + it("fans the reverse command across replicas when the socket is remote", async () => { + const app = getApp(); + const admin = await createAdminContext(app, { username: `pm-${crypto.randomUUID().slice(0, 6)}` }); + await app.db + .update(clients) + .set({ status: "connected", instanceId: "replica-other" }) + .where(eq(clients.id, admin.clientId)); + + const notifyCommand = vi.spyOn(app.notifier, "notifyDaemonClientCommand").mockResolvedValue(); + + const pending = app.inject({ + method: "GET", + url: `/api/v1/clients/${admin.clientId}/providers/cursor/models`, + headers: { authorization: `Bearer ${admin.accessToken}` }, + }); + + await vi.waitFor(() => { + expect(notifyCommand).toHaveBeenCalled(); + }); + const command = notifyCommand.mock.calls[0]?.[0]; + expect(command).toMatchObject({ + type: "provider-models:list", + clientId: admin.clientId, + provider: "cursor", + }); + expect(typeof command?.ref).toBe("string"); + const ref = command?.ref; + if (!ref) throw new Error("expected notifyDaemonClientCommand ref"); + + const catalog = { + provider: "cursor" as const, + models: [{ id: "auto", label: "Auto", isDefault: true }], + defaultModelId: "auto", + fetchedAt: new Date().toISOString(), + source: "provider-cli" as const, + error: null, + }; + // Simulate the socket-owning replica: durable store + local miss + result wake. + await storeModelCatalogRpcResult(app.db, admin.clientId, ref, catalog); + expect(await readModelCatalogRpcResult(app.db, admin.clientId, ref)).toMatchObject(catalog); + // Direct resolve mirrors what onDaemonClientCommandResult does after read. + expect(resolveClientReply(admin.clientId, ref, catalog)).toBe(true); + + const res = await pending; + expect(res.statusCode).toBe(200); + expect(res.json()).toMatchObject(catalog); + notifyCommand.mockRestore(); + }); + + it("resolves a remote waiter from metadata after a result wake", async () => { + const app = getApp(); + const admin = await createAdminContext(app, { username: `pm-${crypto.randomUUID().slice(0, 6)}` }); + const ref = crypto.randomUUID(); + const catalog = { + provider: "kimi-code" as const, + models: [{ id: "gemini-3-pro-preview", label: "Gemini 3 Pro", isDefault: true }], + defaultModelId: "gemini-3-pro-preview", + fetchedAt: new Date().toISOString(), + source: "provider-config" as const, + error: null, + }; + + const replyPromise = waitForClientReply(admin.clientId, ref); + await storeModelCatalogRpcResult(app.db, admin.clientId, ref, catalog); + await app.notifier.notifyDaemonClientCommandResult({ clientId: admin.clientId, ref }); + + await expect(replyPromise).resolves.toMatchObject(catalog); + }); }); diff --git a/packages/server/src/__tests__/notifier-extra.test.ts b/packages/server/src/__tests__/notifier-extra.test.ts index 5a20be671..2da7ff44a 100644 --- a/packages/server/src/__tests__/notifier-extra.test.ts +++ b/packages/server/src/__tests__/notifier-extra.test.ts @@ -102,8 +102,15 @@ describe("createNotifier", () => { await notifier.notifyChatAudience("chat_1"); await notifier.notifyChatUpdated("chat_1"); await notifier.notifyAgentRouteChange(payload); + await notifier.notifyDaemonClientCommand({ + type: "provider-models:list", + clientId: "client_1", + provider: "cursor", + ref: "ref_1", + }); + await notifier.notifyDaemonClientCommandResult({ clientId: "client_1", ref: "ref_1" }); - expect(ok.calls).toHaveLength(10); + expect(ok.calls).toHaveLength(12); expect(ok.calls.map((values) => values[0])).toEqual([ "inbox_notifications", "config_changes", @@ -115,6 +122,8 @@ describe("createNotifier", () => { "chat_audience_events", "chat_updated_events", "agent_route_events", + "daemon_client_commands", + "daemon_client_command_results", ]); const failing = createNotifier(makeListenClient(true).client as never); @@ -128,6 +137,17 @@ describe("createNotifier", () => { await expect(failing.notifyChatAudience("chat_1")).resolves.toBeUndefined(); await expect(failing.notifyChatUpdated("chat_1")).resolves.toBeUndefined(); await expect(failing.notifyAgentRouteChange(payload)).resolves.toBeUndefined(); + await expect( + failing.notifyDaemonClientCommand({ + type: "provider-models:list", + clientId: "client_1", + provider: "cursor", + ref: "ref_1", + }), + ).resolves.toBeUndefined(); + await expect( + failing.notifyDaemonClientCommandResult({ clientId: "client_1", ref: "ref_1" }), + ).resolves.toBeUndefined(); }); it("parses LISTEN payloads, ignores malformed data, swallows handler errors, and stops idempotently", async () => { @@ -160,6 +180,14 @@ describe("createNotifier", () => { throw new Error("consumer failed"); }); const agentRouteSecond = vi.fn(); + const daemonCommand = vi.fn(() => { + throw new Error("consumer failed"); + }); + const daemonCommandSecond = vi.fn(); + const daemonResult = vi.fn(() => { + throw new Error("consumer failed"); + }); + const daemonResultSecond = vi.fn(); notifier.onConfigChange(config); notifier.onSessionStateChange(sessionState); @@ -176,6 +204,10 @@ describe("createNotifier", () => { notifier.onMeChatsChanged(meChatsChanged); notifier.onAgentRouteChange(agentRoute); notifier.onAgentRouteChange(agentRouteSecond); + notifier.onDaemonClientCommand(daemonCommand); + notifier.onDaemonClientCommand(daemonCommandSecond); + notifier.onDaemonClientCommandResult(daemonResult); + notifier.onDaemonClientCommandResult(daemonResultSecond); await notifier.start(); listeners.get("config_changes")?.("agent"); @@ -206,6 +238,19 @@ describe("createNotifier", () => { ); listeners.get("agent_route_events")?.("{not json"); listeners.get("agent_route_events")?.(JSON.stringify({ agentId: "agent_1" })); + listeners.get("daemon_client_commands")?.( + JSON.stringify({ + type: "provider-models:list", + clientId: "client_1", + provider: "cursor", + ref: "ref_1", + }), + ); + listeners.get("daemon_client_commands")?.("{not json"); + listeners.get("daemon_client_commands")?.(JSON.stringify({ type: "provider-models:list" })); + listeners.get("daemon_client_command_results")?.(JSON.stringify({ clientId: "client_1", ref: "ref_1" })); + listeners.get("daemon_client_command_results")?.("{not json"); + listeners.get("daemon_client_command_results")?.(JSON.stringify({ clientId: "client_1" })); expect(config).toHaveBeenCalledWith("agent"); expect(sessionState).toHaveBeenCalledWith({ @@ -243,10 +288,17 @@ describe("createNotifier", () => { runtimeProvider: "codex", targetClientId: "client_1", }); + expect(daemonCommandSecond).toHaveBeenCalledWith({ + type: "provider-models:list", + clientId: "client_1", + provider: "cursor", + ref: "ref_1", + }); + expect(daemonResultSecond).toHaveBeenCalledWith({ clientId: "client_1", ref: "ref_1" }); await notifier.stop(); await notifier.stop(); - expect(unlisteners).toHaveLength(11); + expect(unlisteners).toHaveLength(13); for (const unlisten of unlisteners) { expect(unlisten).toHaveBeenCalledTimes(1); } diff --git a/packages/server/src/api/agent/ws-client.ts b/packages/server/src/api/agent/ws-client.ts index a0b27bb0c..db9aa9902 100644 --- a/packages/server/src/api/agent/ws-client.ts +++ b/packages/server/src/api/agent/ws-client.ts @@ -12,6 +12,7 @@ import { inboxAckFrameSchema, inboxDeliverFrameSchema, inboxRecoverFrameSchema, + PROVIDER_MODELS_LIST_TYPE, PROVIDER_MODELS_RESULT_TYPE, providerModelsResultFrameSchema, runtimeStateMessageSchema, @@ -56,6 +57,7 @@ import * as landingCampaignChatStateService from "../../services/landing-campaig import * as notificationService from "../../services/notification.js"; import type { InboxPushHandler, Notifier } from "../../services/notifier.js"; import * as presenceService from "../../services/presence.js"; +import { readModelCatalogRpcResult, storeModelCatalogRpcResult } from "../../services/provider-models-rpc.js"; import * as runtimeLivenessService from "../../services/runtime-liveness.js"; import * as sessionEventService from "../../services/session-event.js"; @@ -295,6 +297,29 @@ export function clientWsRoutes(notifier: Notifier, instanceId: string) { connectionManager.sendToClient(payload.targetClientId, frame.data); }); + // Cross-replica reverse commands: only the replica that owns the daemon + // WebSocket can deliver; others no-op on sendToClient. + notifier.onDaemonClientCommand((payload) => { + if (payload.type !== PROVIDER_MODELS_LIST_TYPE) return; + connectionManager.sendToClient(payload.clientId, { + type: PROVIDER_MODELS_LIST_TYPE, + provider: payload.provider, + ref: payload.ref, + }); + }); + + // Cross-replica result wake: catalog is in clients.metadata; resolve any + // local HTTP waiter that registered waitForClientReply for this ref. + notifier.onDaemonClientCommandResult((payload) => { + void (async () => { + const catalog = await readModelCatalogRpcResult(app.db, payload.clientId, payload.ref); + if (!catalog) return; + connectionManager.resolveClientReply(payload.clientId, payload.ref, catalog); + })().catch((err) => { + app.log.debug({ err, clientId: payload.clientId, ref: payload.ref }, "provider-models result wake failed"); + }); + }); + // WS upgrade is excluded from HTTP tracing in app.ts via the autotelic // plugin's `ignoreRoutes` — fastify hijacks the reply on upgrade, so a // `onResponse`-terminated HTTP span would never end. The connection's @@ -1769,11 +1794,15 @@ export function clientWsRoutes(notifier: Notifier, instanceId: string) { socket.send(JSON.stringify({ type: "error", message: "Malformed provider-models:result frame" })); return; } + // Durable rendezvous first so another replica's HTTP waiter can + // load the catalog after the tiny result-wake NOTIFY. + await storeModelCatalogRpcResult(app.db, clientId, result.data.ref, result.data.catalog); const resolved = connectionManager.resolveClientReply(clientId, result.data.ref, result.data.catalog); + await notifier.notifyDaemonClientCommandResult({ clientId, ref: result.data.ref }); if (!resolved) { app.log.debug( { clientId, ref: result.data.ref }, - "provider-models:result matched no pending HTTP waiter", + "provider-models:result matched no pending HTTP waiter on this replica", ); } } diff --git a/packages/server/src/api/clients.ts b/packages/server/src/api/clients.ts index dc0785726..9ac049393 100644 --- a/packages/server/src/api/clients.ts +++ b/packages/server/src/api/clients.ts @@ -20,6 +20,7 @@ import { sendToClient, waitForClientReply, } from "../services/connection-manager.js"; +import { isClientConnectedSomewhere } from "../services/provider-models-rpc.js"; import { serializeDate } from "../utils.js"; import { clientCommandVersionHint } from "./client-command-version.js"; @@ -101,7 +102,10 @@ export async function clientRoutes(app: FastifyInstance): Promise { // Host-local model catalog: ask the connected daemon to discover models from // the real provider on that computer, wait for the correlated reply, and - // return the catalog to the web. 503 when the computer is offline or times out. + // return the catalog to the web. When the daemon socket lives on another + // replica, fan the reverse command via PG NOTIFY; the result is stored in + // clients.metadata and woken with a tiny result NOTIFY. 503 when offline + // or timed out. app.get<{ Params: { clientId: string; provider: string } }>( "/:clientId/providers/:provider/models", async (request) => { @@ -113,17 +117,27 @@ export async function clientRoutes(app: FastifyInstance): Promise { const provider = runtimeProviderSchema.parse(rawProvider); const ref = randomUUID(); const replyPromise = waitForClientReply(clientId, ref); - const delivered = sendToClient(clientId, { + const command = { type: PROVIDER_MODELS_LIST_TYPE, provider, ref, - }); - if (!delivered) { - rejectPendingRepliesForClient(clientId, new Error("Computer not connected")); - await replyPromise.catch(() => undefined); - throw new ServiceUnavailableError( - "Could not list models because this computer is not connected. Make sure the daemon is running, then retry.", - ); + }; + const deliveredLocally = sendToClient(clientId, command); + if (!deliveredLocally) { + const client = await clientService.getClient(app.db, clientId); + if (!client || !isClientConnectedSomewhere(client)) { + rejectPendingRepliesForClient(clientId, new Error("Computer not connected")); + await replyPromise.catch(() => undefined); + throw new ServiceUnavailableError( + "Could not list models because this computer is not connected. Make sure the daemon is running, then retry.", + ); + } + await app.notifier.notifyDaemonClientCommand({ + type: PROVIDER_MODELS_LIST_TYPE, + clientId, + provider, + ref, + }); } try { const raw = await replyPromise; diff --git a/packages/server/src/services/notifier.ts b/packages/server/src/services/notifier.ts index 28f244351..9ce4f48a8 100644 --- a/packages/server/src/services/notifier.ts +++ b/packages/server/src/services/notifier.ts @@ -38,6 +38,18 @@ const CHAT_AUDIENCE_CHANNEL = "chat_audience_events"; */ const CHAT_UPDATED_CHANNEL = "chat_updated_events"; const AGENT_ROUTE_CHANNEL = "agent_route_events"; +/** + * Cross-replica reverse command to a connected daemon (e.g. provider-models:list). + * Payload is small JSON: `{ type, clientId, provider, ref }`. The replica that + * owns the client's WebSocket delivers it via `sendToClient`; others no-op. + */ +const DAEMON_CLIENT_COMMAND_CHANNEL = "daemon_client_commands"; +/** + * Cross-replica wake that a daemon command result is ready. Payload is + * `{ clientId, ref }` only — the catalog lives in `clients.metadata` so large + * Cursor lists stay under the PG NOTIFY 8KB limit. + */ +const DAEMON_CLIENT_COMMAND_RESULT_CHANNEL = "daemon_client_command_results"; /** * A viewer's PRIVATE me-chats projection changed (currently: they pinned or * unpinned a chat). Carries `:` so the WS layer @@ -97,6 +109,22 @@ export type AgentRouteChangePayload = { reason: string; }; export type AgentRouteChangeHandler = (payload: AgentRouteChangePayload) => void; + +/** Small reverse-command frame fan-out for host-local daemon RPCs. */ +export type DaemonClientCommandPayload = { + type: string; + clientId: string; + provider: string; + ref: string; +}; +export type DaemonClientCommandHandler = (payload: DaemonClientCommandPayload) => void; + +/** Wake waiters that a correlated daemon RPC result is stored in client metadata. */ +export type DaemonClientCommandResultPayload = { + clientId: string; + ref: string; +}; +export type DaemonClientCommandResultHandler = (payload: DaemonClientCommandResultPayload) => void; export type MeChatsChangedHandler = (payload: { humanAgentId: string; organizationId: string }) => void; /** @@ -146,6 +174,17 @@ export type Notifier = { notifyMeChatsChanged(humanAgentId: string, organizationId: string): Promise; /** Agent runtime route changed: fan local WS detach/pin handling to every server replica. */ notifyAgentRouteChange(payload: AgentRouteChangePayload): Promise; + /** + * Fan a small reverse-command frame to every replica so the process that + * owns the daemon WebSocket can `sendToClient`. Payload must stay tiny + * (no catalog bodies). + */ + notifyDaemonClientCommand(payload: DaemonClientCommandPayload): Promise; + /** + * Wake waiters that a correlated daemon RPC result is durable in + * `clients.metadata` (catalog bodies are too large for NOTIFY). + */ + notifyDaemonClientCommandResult(payload: DaemonClientCommandResultPayload): Promise; /** * Push a raw JSON frame to every socket currently subscribed to `inboxId` * on **this server instance only**. Unlike `notify`, does not fan out @@ -174,6 +213,10 @@ export type Notifier = { onMeChatsChanged(handler: MeChatsChangedHandler): void; /** Register a handler for agent runtime route changes. */ onAgentRouteChange(handler: AgentRouteChangeHandler): void; + /** Register a handler for cross-replica daemon reverse commands. */ + onDaemonClientCommand(handler: DaemonClientCommandHandler): void; + /** Register a handler for cross-replica daemon RPC result wakes. */ + onDaemonClientCommandResult(handler: DaemonClientCommandResultHandler): void; /** Start listening for PG notifications */ start(): Promise; /** Stop listening */ @@ -192,6 +235,8 @@ export function createNotifier(listenClient: postgres.Sql): Notifier { const chatUpdatedHandlers: ChatUpdatedChangeHandler[] = []; const meChatsChangedHandlers: MeChatsChangedHandler[] = []; const agentRouteHandlers: AgentRouteChangeHandler[] = []; + const daemonClientCommandHandlers: DaemonClientCommandHandler[] = []; + const daemonClientCommandResultHandlers: DaemonClientCommandResultHandler[] = []; let unlistenInboxFn: (() => Promise) | null = null; let unlistenConfigFn: (() => Promise) | null = null; let unlistenSessionStateFn: (() => Promise) | null = null; @@ -203,6 +248,8 @@ export function createNotifier(listenClient: postgres.Sql): Notifier { let unlistenChatUpdatedFn: (() => Promise) | null = null; let unlistenMeChatsChangedFn: (() => Promise) | null = null; let unlistenAgentRouteFn: (() => Promise) | null = null; + let unlistenDaemonClientCommandFn: (() => Promise) | null = null; + let unlistenDaemonClientCommandResultFn: (() => Promise) | null = null; function handleNotification(payload: string) { // payload format: "inboxId:messageId" @@ -344,6 +391,22 @@ export function createNotifier(listenClient: postgres.Sql): Notifier { } }, + async notifyDaemonClientCommand(payload: DaemonClientCommandPayload) { + try { + await listenClient`SELECT pg_notify(${DAEMON_CLIENT_COMMAND_CHANNEL}, ${JSON.stringify(payload)})`; + } catch { + // fire-and-forget — HTTP waiter timeout is the durable fallback. + } + }, + + async notifyDaemonClientCommandResult(payload: DaemonClientCommandResultPayload) { + try { + await listenClient`SELECT pg_notify(${DAEMON_CLIENT_COMMAND_RESULT_CHANNEL}, ${JSON.stringify(payload)})`; + } catch { + // fire-and-forget — HTTP waiter timeout is the durable fallback. + } + }, + async pushFrameToInbox(inboxId: string, frame: string): Promise { const map = subscriptions.get(inboxId); if (!map) return 0; @@ -404,6 +467,14 @@ export function createNotifier(listenClient: postgres.Sql): Notifier { agentRouteHandlers.push(handler); }, + onDaemonClientCommand(handler: DaemonClientCommandHandler) { + daemonClientCommandHandlers.push(handler); + }, + + onDaemonClientCommandResult(handler: DaemonClientCommandResultHandler) { + daemonClientCommandResultHandlers.push(handler); + }, + async start() { const inboxResult = await listenClient.listen(INBOX_CHANNEL, (payload) => { if (payload) handleNotification(payload); @@ -595,6 +666,54 @@ export function createNotifier(listenClient: postgres.Sql): Notifier { } }); unlistenAgentRouteFn = agentRouteResult.unlisten; + + const daemonClientCommandResult = await listenClient.listen(DAEMON_CLIENT_COMMAND_CHANNEL, (payload) => { + if (!payload) return; + try { + const parsed = JSON.parse(payload) as Partial; + if ( + typeof parsed.type !== "string" || + typeof parsed.clientId !== "string" || + typeof parsed.provider !== "string" || + typeof parsed.ref !== "string" + ) { + return; + } + for (const handler of daemonClientCommandHandlers) { + try { + handler(parsed as DaemonClientCommandPayload); + } catch { + // swallow — handler errors must not poison fan-out + } + } + } catch { + // ignore malformed payloads + } + }); + unlistenDaemonClientCommandFn = daemonClientCommandResult.unlisten; + + const daemonClientCommandResultWake = await listenClient.listen( + DAEMON_CLIENT_COMMAND_RESULT_CHANNEL, + (payload) => { + if (!payload) return; + try { + const parsed = JSON.parse(payload) as Partial; + if (typeof parsed.clientId !== "string" || typeof parsed.ref !== "string") { + return; + } + for (const handler of daemonClientCommandResultHandlers) { + try { + handler(parsed as DaemonClientCommandResultPayload); + } catch { + // swallow — handler errors must not poison fan-out + } + } + } catch { + // ignore malformed payloads + } + }, + ); + unlistenDaemonClientCommandResultFn = daemonClientCommandResultWake.unlisten; }, async stop() { @@ -642,6 +761,14 @@ export function createNotifier(listenClient: postgres.Sql): Notifier { await unlistenAgentRouteFn(); unlistenAgentRouteFn = null; } + if (unlistenDaemonClientCommandFn) { + await unlistenDaemonClientCommandFn(); + unlistenDaemonClientCommandFn = null; + } + if (unlistenDaemonClientCommandResultFn) { + await unlistenDaemonClientCommandResultFn(); + unlistenDaemonClientCommandResultFn = null; + } }, }; } diff --git a/packages/server/src/services/provider-models-rpc.ts b/packages/server/src/services/provider-models-rpc.ts new file mode 100644 index 000000000..c504d4a3b --- /dev/null +++ b/packages/server/src/services/provider-models-rpc.ts @@ -0,0 +1,101 @@ +import { type ProviderModelCatalog, providerModelCatalogSchema } from "@first-tree/shared"; +import { eq } from "drizzle-orm"; +import type { Database } from "../db/connection.js"; +import { clients } from "../db/schema/clients.js"; + +/** + * Durable rendezvous for host-local model-catalog RPC. + * + * PG NOTIFY payloads must stay small (≈8KB). Cursor catalogs can exceed that, + * so the socket-owning replica stores the catalog under + * `clients.metadata.modelCatalogRpc[ref]` and fans a tiny `{ clientId, ref }` + * wake. The HTTP-serving replica (possibly a different process) loads the + * catalog from metadata and resolves its local waiter. + */ + +const RPC_METADATA_KEY = "modelCatalogRpc"; +const MAX_AGE_MS = 120_000; +const MAX_ENTRIES = 20; + +type RpcEntry = { + catalog: ProviderModelCatalog; + storedAt: string; +}; + +function asRpcMap(raw: unknown): Record { + if (!raw || typeof raw !== "object" || Array.isArray(raw)) return {}; + const out: Record = {}; + for (const [ref, value] of Object.entries(raw as Record)) { + if (!value || typeof value !== "object" || Array.isArray(value)) continue; + const row = value as Record; + const parsed = providerModelCatalogSchema.safeParse(row.catalog); + if (!parsed.success || typeof row.storedAt !== "string") continue; + out[ref] = { catalog: parsed.data, storedAt: row.storedAt }; + } + return out; +} + +function pruneRpcMap(map: Record, nowMs: number): Record { + const fresh = Object.entries(map).filter(([, entry]) => { + const t = Date.parse(entry.storedAt); + return Number.isFinite(t) && nowMs - t < MAX_AGE_MS; + }); + fresh.sort((a, b) => Date.parse(b[1].storedAt) - Date.parse(a[1].storedAt)); + return Object.fromEntries(fresh.slice(0, MAX_ENTRIES)); +} + +/** Persist a catalog so any replica can resolve the correlated HTTP waiter. */ +export async function storeModelCatalogRpcResult( + db: Database, + clientId: string, + ref: string, + catalog: ProviderModelCatalog, +): Promise { + const [client] = await db + .select({ metadata: clients.metadata }) + .from(clients) + .where(eq(clients.id, clientId)) + .limit(1); + if (!client) return; + + const base = (client.metadata ?? {}) as Record; + const existing = asRpcMap(base[RPC_METADATA_KEY]); + const now = new Date(); + const next = pruneRpcMap( + { + ...existing, + [ref]: { catalog, storedAt: now.toISOString() }, + }, + now.getTime(), + ); + await db + .update(clients) + .set({ metadata: { ...base, [RPC_METADATA_KEY]: next } }) + .where(eq(clients.id, clientId)); +} + +/** Load a previously stored catalog (does not delete — prune happens on write). */ +export async function readModelCatalogRpcResult( + db: Database, + clientId: string, + ref: string, +): Promise { + const [client] = await db + .select({ metadata: clients.metadata }) + .from(clients) + .where(eq(clients.id, clientId)) + .limit(1); + if (!client) return null; + const base = (client.metadata ?? {}) as Record; + const entry = asRpcMap(base[RPC_METADATA_KEY])[ref]; + return entry?.catalog ?? null; +} + +/** + * True when the DB says a daemon WebSocket is live somewhere (this process or + * another replica). Used after a local `sendToClient` miss to decide between + * cross-replica fan-out and a hard 503. + */ +export function isClientConnectedSomewhere(client: { status: string; instanceId: string | null }): boolean { + return client.status === "connected" && client.instanceId != null; +} diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index fcb9d1774..7e2fffb35 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -126,6 +126,9 @@ importers: semver: specifier: ^7.6.3 version: 7.7.4 + smol-toml: + specifier: ^1.7.0 + version: 1.7.0 ws: specifier: ^8.20.0 version: 8.20.0 @@ -7657,6 +7660,22 @@ snapshots: chai: 5.3.3 tinyrainbow: 2.0.0 + '@vitest/mocker@3.2.4(vite@6.4.1(@types/node@22.19.15)(jiti@2.6.1)(lightningcss@1.32.0)(tsx@4.21.0)(yaml@2.8.2))': + dependencies: + '@vitest/spy': 3.2.4 + estree-walker: 3.0.3 + magic-string: 0.30.21 + optionalDependencies: + vite: 6.4.1(@types/node@22.19.15)(jiti@2.6.1)(lightningcss@1.32.0)(tsx@4.21.0)(yaml@2.8.2) + + '@vitest/mocker@3.2.4(vite@6.4.1(@types/node@22.19.15)(jiti@2.6.1)(lightningcss@1.32.0)(tsx@4.21.0)(yaml@2.8.3))': + dependencies: + '@vitest/spy': 3.2.4 + estree-walker: 3.0.3 + magic-string: 0.30.21 + optionalDependencies: + vite: 6.4.1(@types/node@22.19.15)(jiti@2.6.1)(lightningcss@1.32.0)(tsx@4.21.0)(yaml@2.8.3) + '@vitest/mocker@3.2.4(vite@6.4.1(@types/node@22.19.19)(jiti@2.6.1)(lightningcss@1.32.0)(tsx@4.21.0)(yaml@2.8.3))': dependencies: '@vitest/spy': 3.2.4 @@ -10323,7 +10342,7 @@ snapshots: dependencies: '@types/chai': 5.2.3 '@vitest/expect': 3.2.4 - '@vitest/mocker': 3.2.4(vite@6.4.1(@types/node@22.19.19)(jiti@2.6.1)(lightningcss@1.32.0)(tsx@4.21.0)(yaml@2.8.3)) + '@vitest/mocker': 3.2.4(vite@6.4.1(@types/node@22.19.15)(jiti@2.6.1)(lightningcss@1.32.0)(tsx@4.21.0)(yaml@2.8.2)) '@vitest/pretty-format': 3.2.4 '@vitest/runner': 3.2.4 '@vitest/snapshot': 3.2.4 @@ -10366,7 +10385,7 @@ snapshots: dependencies: '@types/chai': 5.2.3 '@vitest/expect': 3.2.4 - '@vitest/mocker': 3.2.4(vite@6.4.1(@types/node@22.19.19)(jiti@2.6.1)(lightningcss@1.32.0)(tsx@4.21.0)(yaml@2.8.3)) + '@vitest/mocker': 3.2.4(vite@6.4.1(@types/node@22.19.15)(jiti@2.6.1)(lightningcss@1.32.0)(tsx@4.21.0)(yaml@2.8.3)) '@vitest/pretty-format': 3.2.4 '@vitest/runner': 3.2.4 '@vitest/snapshot': 3.2.4 From 278be4fba804696d0501a053bcb29504ecf0848d Mon Sep 17 00:00:00 2001 From: YueZengwu Date: Fri, 17 Jul 2026 21:03:08 +0800 Subject: [PATCH 5/5] fix(runtime): harden model-catalog rendezvous under concurrency Atomic jsonb_set for catalog refs and capabilities, DB fallback when the result wake is lost, and instance_id-scoped command/result delivery so stale sockets after takeover cannot win the RPC. Co-authored-by: Cursor --- .../__tests__/clients-provider-models.test.ts | 157 +++++++++++++++++- .../src/__tests__/notifier-extra.test.ts | 4 + packages/server/src/api/agent/ws-client.ts | 20 ++- packages/server/src/api/clients.ts | 31 ++-- packages/server/src/services/client.ts | 15 +- .../server/src/services/connection-manager.ts | 11 +- packages/server/src/services/notifier.ts | 5 +- .../src/services/provider-models-rpc.ts | 92 +++++----- 8 files changed, 261 insertions(+), 74 deletions(-) diff --git a/packages/server/src/__tests__/clients-provider-models.test.ts b/packages/server/src/__tests__/clients-provider-models.test.ts index 5e96ace5a..beaae76ee 100644 --- a/packages/server/src/__tests__/clients-provider-models.test.ts +++ b/packages/server/src/__tests__/clients-provider-models.test.ts @@ -1,12 +1,14 @@ import crypto from "node:crypto"; import { eq } from "drizzle-orm"; -import { describe, expect, it, vi } from "vitest"; +import { afterEach, describe, expect, it, vi } from "vitest"; import type { WebSocket } from "ws"; import { clients } from "../db/schema/clients.js"; +import * as clientService from "../services/client.js"; import { removeClientConnection, resolveClientReply, setClientConnection, + setClientReplyTimeoutMsForTests, waitForClientReply, } from "../services/connection-manager.js"; import { readModelCatalogRpcResult, storeModelCatalogRpcResult } from "../services/provider-models-rpc.js"; @@ -19,9 +21,18 @@ import { createAdminContext, useTestApp } from "./helpers.js"; describe("GET /clients/:clientId/providers/:provider/models", () => { const getApp = useTestApp(); + afterEach(() => { + setClientReplyTimeoutMsForTests(null); + }); + + async function markClientOnInstance(app: ReturnType, clientId: string, instanceId: string) { + await app.db.update(clients).set({ status: "connected", instanceId }).where(eq(clients.id, clientId)); + } + it("forwards provider-models:list and returns the daemon catalog", async () => { const app = getApp(); const admin = await createAdminContext(app, { username: `pm-${crypto.randomUUID().slice(0, 6)}` }); + await markClientOnInstance(app, admin.clientId, app.config.instanceId); const ws = { readyState: 1, send: vi.fn(), close: vi.fn() }; setClientConnection(admin.clientId, ws as unknown as WebSocket); try { @@ -31,7 +42,6 @@ describe("GET /clients/:clientId/providers/:provider/models", () => { headers: { authorization: `Bearer ${admin.accessToken}` }, }); - // Wait briefly for the reverse command to be sent, then resolve the waiter. await vi.waitFor(() => { expect(ws.send).toHaveBeenCalled(); }); @@ -68,13 +78,10 @@ describe("GET /clients/:clientId/providers/:provider/models", () => { expect(res.statusCode).toBe(503); }); - it("fans the reverse command across replicas when the socket is remote", async () => { + it("fans the reverse command to the DB-authoritative instance when remote", async () => { const app = getApp(); const admin = await createAdminContext(app, { username: `pm-${crypto.randomUUID().slice(0, 6)}` }); - await app.db - .update(clients) - .set({ status: "connected", instanceId: "replica-other" }) - .where(eq(clients.id, admin.clientId)); + await markClientOnInstance(app, admin.clientId, "replica-other"); const notifyCommand = vi.spyOn(app.notifier, "notifyDaemonClientCommand").mockResolvedValue(); @@ -92,6 +99,7 @@ describe("GET /clients/:clientId/providers/:provider/models", () => { type: "provider-models:list", clientId: admin.clientId, provider: "cursor", + targetInstanceId: "replica-other", }); expect(typeof command?.ref).toBe("string"); const ref = command?.ref; @@ -105,10 +113,8 @@ describe("GET /clients/:clientId/providers/:provider/models", () => { source: "provider-cli" as const, error: null, }; - // Simulate the socket-owning replica: durable store + local miss + result wake. await storeModelCatalogRpcResult(app.db, admin.clientId, ref, catalog); expect(await readModelCatalogRpcResult(app.db, admin.clientId, ref)).toMatchObject(catalog); - // Direct resolve mirrors what onDaemonClientCommandResult does after read. expect(resolveClientReply(admin.clientId, ref, catalog)).toBe(true); const res = await pending; @@ -117,6 +123,50 @@ describe("GET /clients/:clientId/providers/:provider/models", () => { notifyCommand.mockRestore(); }); + it("does not deliver on a stale local socket after instance takeover", async () => { + const app = getApp(); + const admin = await createAdminContext(app, { username: `pm-${crypto.randomUUID().slice(0, 6)}` }); + // DB says another replica owns the connection, but this process still has a socket. + await markClientOnInstance(app, admin.clientId, "replica-other"); + const staleWs = { readyState: 1, send: vi.fn(), close: vi.fn() }; + setClientConnection(admin.clientId, staleWs as unknown as WebSocket); + + const notifyCommand = vi.spyOn(app.notifier, "notifyDaemonClientCommand").mockResolvedValue(); + try { + const pending = app.inject({ + method: "GET", + url: `/api/v1/clients/${admin.clientId}/providers/cursor/models`, + headers: { authorization: `Bearer ${admin.accessToken}` }, + }); + + await vi.waitFor(() => { + expect(notifyCommand).toHaveBeenCalled(); + }); + expect(staleWs.send).not.toHaveBeenCalled(); + expect(notifyCommand.mock.calls[0]?.[0]).toMatchObject({ + targetInstanceId: "replica-other", + clientId: admin.clientId, + }); + + const ref = notifyCommand.mock.calls[0]?.[0]?.ref; + if (!ref) throw new Error("expected ref"); + const catalog = { + provider: "cursor" as const, + models: [{ id: "auto", label: "Auto" }], + defaultModelId: "auto", + fetchedAt: new Date().toISOString(), + source: "provider-cli" as const, + error: null, + }; + expect(resolveClientReply(admin.clientId, ref, catalog)).toBe(true); + const res = await pending; + expect(res.statusCode).toBe(200); + } finally { + notifyCommand.mockRestore(); + removeClientConnection(admin.clientId, staleWs as unknown as WebSocket); + } + }); + it("resolves a remote waiter from metadata after a result wake", async () => { const app = getApp(); const admin = await createAdminContext(app, { username: `pm-${crypto.randomUUID().slice(0, 6)}` }); @@ -136,4 +186,93 @@ describe("GET /clients/:clientId/providers/:provider/models", () => { await expect(replyPromise).resolves.toMatchObject(catalog); }); + + it("returns a stored catalog when the result wake is lost", async () => { + const app = getApp(); + const admin = await createAdminContext(app, { username: `pm-${crypto.randomUUID().slice(0, 6)}` }); + await markClientOnInstance(app, admin.clientId, app.config.instanceId); + setClientReplyTimeoutMsForTests(80); + + const ws = { readyState: 1, send: vi.fn(), close: vi.fn() }; + setClientConnection(admin.clientId, ws as unknown as WebSocket); + try { + const pending = app.inject({ + method: "GET", + url: `/api/v1/clients/${admin.clientId}/providers/cursor/models`, + headers: { authorization: `Bearer ${admin.accessToken}` }, + }); + + await vi.waitFor(() => { + expect(ws.send).toHaveBeenCalled(); + }); + const frame = JSON.parse(String(ws.send.mock.calls[0]?.[0])); + const catalog = { + provider: "cursor" as const, + models: [{ id: "auto", label: "Auto", isDefault: true }], + defaultModelId: "auto", + fetchedAt: new Date().toISOString(), + source: "provider-cli" as const, + error: null, + }; + // Durable store arrives, but no resolveClientReply / result NOTIFY (lost wake). + await storeModelCatalogRpcResult(app.db, admin.clientId, frame.ref, catalog); + + const res = await pending; + expect(res.statusCode).toBe(200); + expect(res.json()).toMatchObject(catalog); + } finally { + removeClientConnection(admin.clientId, ws as unknown as WebSocket); + } + }); + + it("stores concurrent refs without clobbering sibling metadata", async () => { + const app = getApp(); + const admin = await createAdminContext(app, { username: `pm-${crypto.randomUUID().slice(0, 6)}` }); + + const detectedAt = new Date().toISOString(); + await clientService.updateClientCapabilities(app.db, admin.clientId, { + cursor: { state: "ok", available: true, detectedAt, sdkVersion: "1.0.0" }, + }); + + const ref1 = crypto.randomUUID(); + const ref2 = crypto.randomUUID(); + const cat1 = { + provider: "cursor" as const, + models: [{ id: "auto", label: "Auto" }], + defaultModelId: "auto", + fetchedAt: new Date().toISOString(), + source: "provider-cli" as const, + error: null, + }; + const cat2 = { + provider: "kimi-code" as const, + models: [{ id: "k3", label: "K3" }], + defaultModelId: "k3", + fetchedAt: new Date().toISOString(), + source: "provider-config" as const, + error: null, + }; + + await Promise.all([ + storeModelCatalogRpcResult(app.db, admin.clientId, ref1, cat1), + storeModelCatalogRpcResult(app.db, admin.clientId, ref2, cat2), + ]); + + expect(await readModelCatalogRpcResult(app.db, admin.clientId, ref1)).toMatchObject(cat1); + expect(await readModelCatalogRpcResult(app.db, admin.clientId, ref2)).toMatchObject(cat2); + + await clientService.updateClientCapabilities(app.db, admin.clientId, { + cursor: { state: "ok", available: true, detectedAt, sdkVersion: "1.0.1" }, + "kimi-code": { state: "ok", available: true, detectedAt }, + }); + + expect(await readModelCatalogRpcResult(app.db, admin.clientId, ref1)).toMatchObject(cat1); + expect(await readModelCatalogRpcResult(app.db, admin.clientId, ref2)).toMatchObject(cat2); + + const row = await clientService.getClient(app.db, admin.clientId); + expect(clientService.extractCapabilities(row?.metadata)).toMatchObject({ + cursor: { available: true, sdkVersion: "1.0.1" }, + "kimi-code": { available: true }, + }); + }); }); diff --git a/packages/server/src/__tests__/notifier-extra.test.ts b/packages/server/src/__tests__/notifier-extra.test.ts index 2da7ff44a..9f7b0d4c6 100644 --- a/packages/server/src/__tests__/notifier-extra.test.ts +++ b/packages/server/src/__tests__/notifier-extra.test.ts @@ -107,6 +107,7 @@ describe("createNotifier", () => { clientId: "client_1", provider: "cursor", ref: "ref_1", + targetInstanceId: "instance_1", }); await notifier.notifyDaemonClientCommandResult({ clientId: "client_1", ref: "ref_1" }); @@ -143,6 +144,7 @@ describe("createNotifier", () => { clientId: "client_1", provider: "cursor", ref: "ref_1", + targetInstanceId: "instance_1", }), ).resolves.toBeUndefined(); await expect( @@ -244,6 +246,7 @@ describe("createNotifier", () => { clientId: "client_1", provider: "cursor", ref: "ref_1", + targetInstanceId: "instance_1", }), ); listeners.get("daemon_client_commands")?.("{not json"); @@ -293,6 +296,7 @@ describe("createNotifier", () => { clientId: "client_1", provider: "cursor", ref: "ref_1", + targetInstanceId: "instance_1", }); expect(daemonResultSecond).toHaveBeenCalledWith({ clientId: "client_1", ref: "ref_1" }); diff --git a/packages/server/src/api/agent/ws-client.ts b/packages/server/src/api/agent/ws-client.ts index db9aa9902..5fa6533dc 100644 --- a/packages/server/src/api/agent/ws-client.ts +++ b/packages/server/src/api/agent/ws-client.ts @@ -297,10 +297,12 @@ export function clientWsRoutes(notifier: Notifier, instanceId: string) { connectionManager.sendToClient(payload.targetClientId, frame.data); }); - // Cross-replica reverse commands: only the replica that owns the daemon - // WebSocket can deliver; others no-op on sendToClient. + // Cross-replica reverse commands: only the DB-authoritative instance may + // deliver. A stale open socket on a previous replica must not receive the + // same ref after reconnect/takeover. notifier.onDaemonClientCommand((payload) => { if (payload.type !== PROVIDER_MODELS_LIST_TYPE) return; + if (payload.targetInstanceId !== instanceId) return; connectionManager.sendToClient(payload.clientId, { type: PROVIDER_MODELS_LIST_TYPE, provider: payload.provider, @@ -1794,6 +1796,20 @@ export function clientWsRoutes(notifier: Notifier, instanceId: string) { socket.send(JSON.stringify({ type: "error", message: "Malformed provider-models:result frame" })); return; } + // Only the DB-authoritative instance accepts results — a replaced + // connection on a prior replica must not win the rendezvous. + const [owner] = await app.db + .select({ instanceId: clients.instanceId }) + .from(clients) + .where(eq(clients.id, clientId)) + .limit(1); + if (!owner || owner.instanceId !== instanceId) { + app.log.debug( + { clientId, ref: result.data.ref, instanceId, ownerInstanceId: owner?.instanceId ?? null }, + "ignoring provider-models:result from non-authoritative instance", + ); + return; + } // Durable rendezvous first so another replica's HTTP waiter can // load the catalog after the tiny result-wake NOTIFY. await storeModelCatalogRpcResult(app.db, clientId, result.data.ref, result.data.catalog); diff --git a/packages/server/src/api/clients.ts b/packages/server/src/api/clients.ts index 9ac049393..30a8d4020 100644 --- a/packages/server/src/api/clients.ts +++ b/packages/server/src/api/clients.ts @@ -20,7 +20,7 @@ import { sendToClient, waitForClientReply, } from "../services/connection-manager.js"; -import { isClientConnectedSomewhere } from "../services/provider-models-rpc.js"; +import { isClientConnectedSomewhere, readModelCatalogRpcResult } from "../services/provider-models-rpc.js"; import { serializeDate } from "../utils.js"; import { clientCommandVersionHint } from "./client-command-version.js"; @@ -102,10 +102,10 @@ export async function clientRoutes(app: FastifyInstance): Promise { // Host-local model catalog: ask the connected daemon to discover models from // the real provider on that computer, wait for the correlated reply, and - // return the catalog to the web. When the daemon socket lives on another - // replica, fan the reverse command via PG NOTIFY; the result is stored in - // clients.metadata and woken with a tiny result NOTIFY. 503 when offline - // or timed out. + // return the catalog to the web. Delivery is scoped to the DB-authoritative + // `clients.instance_id` (local send or PG NOTIFY fan-out). Results are stored + // in clients.metadata; on waiter timeout we still read that durable copy so a + // lost NOTIFY does not false-503. Hard 503 only when offline / truly missing. app.get<{ Params: { clientId: string; provider: string } }>( "/:clientId/providers/:provider/models", async (request) => { @@ -115,34 +115,45 @@ export async function clientRoutes(app: FastifyInstance): Promise { await clientService.assertClientOwner(app.db, clientId, { userId }); await clientService.assertClientNotRetired(app.db, clientId); const provider = runtimeProviderSchema.parse(rawProvider); + const client = await clientService.getClient(app.db, clientId); + if (!client || !isClientConnectedSomewhere(client) || !client.instanceId) { + throw new ServiceUnavailableError( + "Could not list models because this computer is not connected. Make sure the daemon is running, then retry.", + ); + } + const targetInstanceId = client.instanceId; const ref = randomUUID(); const replyPromise = waitForClientReply(clientId, ref); - const command = { + const daemonFrame = { type: PROVIDER_MODELS_LIST_TYPE, provider, ref, }; - const deliveredLocally = sendToClient(clientId, command); - if (!deliveredLocally) { - const client = await clientService.getClient(app.db, clientId); - if (!client || !isClientConnectedSomewhere(client)) { + if (targetInstanceId === app.config.instanceId) { + const delivered = sendToClient(clientId, daemonFrame); + if (!delivered) { rejectPendingRepliesForClient(clientId, new Error("Computer not connected")); await replyPromise.catch(() => undefined); throw new ServiceUnavailableError( "Could not list models because this computer is not connected. Make sure the daemon is running, then retry.", ); } + } else { await app.notifier.notifyDaemonClientCommand({ type: PROVIDER_MODELS_LIST_TYPE, clientId, provider, ref, + targetInstanceId, }); } try { const raw = await replyPromise; return providerModelCatalogSchema.parse(raw); } catch (err) { + // Race-safe fallback: catalog may already be durable while the wake was lost. + const stored = await readModelCatalogRpcResult(app.db, clientId, ref); + if (stored) return stored; throw new ServiceUnavailableError( err instanceof Error ? err.message diff --git a/packages/server/src/services/client.ts b/packages/server/src/services/client.ts index 96acf94cc..ba52efe7c 100644 --- a/packages/server/src/services/client.ts +++ b/packages/server/src/services/client.ts @@ -346,9 +346,20 @@ export async function updateClientCapabilities( if (existingCapabilities.success && stableJson(existingCapabilities.data) === stableJson(parsed.data)) { return; } - const merged = { ...baseMetadata, capabilities: parsed.data }; - await db.update(clients).set({ metadata: merged }).where(eq(clients.id, clientId)); + // Atomic key update so concurrent modelCatalogRpc refs (and other metadata + // writers) are not erased by a whole-object metadata replace. + await db + .update(clients) + .set({ + metadata: sql`jsonb_set( + COALESCE(${clients.metadata}, '{}'::jsonb), + '{capabilities}', + ${JSON.stringify(parsed.data)}::jsonb, + true + )`, + }) + .where(eq(clients.id, clientId)); } /** diff --git a/packages/server/src/services/connection-manager.ts b/packages/server/src/services/connection-manager.ts index 734f9b409..7b100923a 100644 --- a/packages/server/src/services/connection-manager.ts +++ b/packages/server/src/services/connection-manager.ts @@ -241,12 +241,19 @@ type PendingClientReply = { const pendingClientReplies = new Map(); -const DEFAULT_CLIENT_REPLY_TIMEOUT_MS = 25_000; +export const DEFAULT_CLIENT_REPLY_TIMEOUT_MS = 25_000; + +let clientReplyTimeoutMsForTests: number | null = null; + +/** Test seam: shorten HTTP↔daemon reply waits without changing production default. */ +export function setClientReplyTimeoutMsForTests(timeoutMs: number | null): void { + clientReplyTimeoutMsForTests = timeoutMs; +} export function waitForClientReply( clientId: string, ref: string, - timeoutMs: number = DEFAULT_CLIENT_REPLY_TIMEOUT_MS, + timeoutMs: number = clientReplyTimeoutMsForTests ?? DEFAULT_CLIENT_REPLY_TIMEOUT_MS, ): Promise { return new Promise((resolve, reject) => { if (pendingClientReplies.has(ref)) { diff --git a/packages/server/src/services/notifier.ts b/packages/server/src/services/notifier.ts index 9ce4f48a8..d6a744c02 100644 --- a/packages/server/src/services/notifier.ts +++ b/packages/server/src/services/notifier.ts @@ -116,6 +116,8 @@ export type DaemonClientCommandPayload = { clientId: string; provider: string; ref: string; + /** DB-authoritative `clients.instance_id` — only that replica may deliver. */ + targetInstanceId: string; }; export type DaemonClientCommandHandler = (payload: DaemonClientCommandPayload) => void; @@ -675,7 +677,8 @@ export function createNotifier(listenClient: postgres.Sql): Notifier { typeof parsed.type !== "string" || typeof parsed.clientId !== "string" || typeof parsed.provider !== "string" || - typeof parsed.ref !== "string" + typeof parsed.ref !== "string" || + typeof parsed.targetInstanceId !== "string" ) { return; } diff --git a/packages/server/src/services/provider-models-rpc.ts b/packages/server/src/services/provider-models-rpc.ts index c504d4a3b..aa340aa06 100644 --- a/packages/server/src/services/provider-models-rpc.ts +++ b/packages/server/src/services/provider-models-rpc.ts @@ -1,5 +1,5 @@ import { type ProviderModelCatalog, providerModelCatalogSchema } from "@first-tree/shared"; -import { eq } from "drizzle-orm"; +import { eq, sql } from "drizzle-orm"; import type { Database } from "../db/connection.js"; import { clients } from "../db/schema/clients.js"; @@ -8,73 +8,68 @@ import { clients } from "../db/schema/clients.js"; * * PG NOTIFY payloads must stay small (≈8KB). Cursor catalogs can exceed that, * so the socket-owning replica stores the catalog under - * `clients.metadata.modelCatalogRpc[ref]` and fans a tiny `{ clientId, ref }` - * wake. The HTTP-serving replica (possibly a different process) loads the - * catalog from metadata and resolves its local waiter. + * `clients.metadata.modelCatalogRpc[ref]` with an atomic top-level `jsonb_set` + * (sibling keys like `capabilities` stay intact) and a nested `||` merge for + * the ref (concurrent UPDATEs on the same client row serialize under the row + * lock and re-read the latest map). A tiny `{ clientId, ref }` wake fans out + * after the durable write. */ const RPC_METADATA_KEY = "modelCatalogRpc"; +/** Ignore durable entries older than this (logical TTL; keys are not rewritten). */ const MAX_AGE_MS = 120_000; -const MAX_ENTRIES = 20; +const REF_RE = /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i; type RpcEntry = { catalog: ProviderModelCatalog; storedAt: string; }; -function asRpcMap(raw: unknown): Record { - if (!raw || typeof raw !== "object" || Array.isArray(raw)) return {}; - const out: Record = {}; - for (const [ref, value] of Object.entries(raw as Record)) { - if (!value || typeof value !== "object" || Array.isArray(value)) continue; - const row = value as Record; - const parsed = providerModelCatalogSchema.safeParse(row.catalog); - if (!parsed.success || typeof row.storedAt !== "string") continue; - out[ref] = { catalog: parsed.data, storedAt: row.storedAt }; - } - return out; -} - -function pruneRpcMap(map: Record, nowMs: number): Record { - const fresh = Object.entries(map).filter(([, entry]) => { - const t = Date.parse(entry.storedAt); - return Number.isFinite(t) && nowMs - t < MAX_AGE_MS; - }); - fresh.sort((a, b) => Date.parse(b[1].storedAt) - Date.parse(a[1].storedAt)); - return Object.fromEntries(fresh.slice(0, MAX_ENTRIES)); +function asRpcEntry(raw: unknown): RpcEntry | null { + if (!raw || typeof raw !== "object" || Array.isArray(raw)) return null; + const row = raw as Record; + const parsed = providerModelCatalogSchema.safeParse(row.catalog); + if (!parsed.success || typeof row.storedAt !== "string") return null; + const storedMs = Date.parse(row.storedAt); + if (!Number.isFinite(storedMs) || Date.now() - storedMs >= MAX_AGE_MS) return null; + return { catalog: parsed.data, storedAt: row.storedAt }; } -/** Persist a catalog so any replica can resolve the correlated HTTP waiter. */ +/** + * Persist one catalog ref without replacing the whole `clients.metadata` object. + * Concurrent refs and sibling metadata writers (capabilities) are preserved. + */ export async function storeModelCatalogRpcResult( db: Database, clientId: string, ref: string, catalog: ProviderModelCatalog, ): Promise { - const [client] = await db - .select({ metadata: clients.metadata }) - .from(clients) - .where(eq(clients.id, clientId)) - .limit(1); - if (!client) return; - - const base = (client.metadata ?? {}) as Record; - const existing = asRpcMap(base[RPC_METADATA_KEY]); - const now = new Date(); - const next = pruneRpcMap( - { - ...existing, - [ref]: { catalog, storedAt: now.toISOString() }, - }, - now.getTime(), - ); + if (!REF_RE.test(ref)) { + throw new Error(`Invalid model-catalog RPC ref: ${ref}`); + } + const entry = { + catalog, + storedAt: new Date().toISOString(), + }; + // Top-level jsonb_set keeps capabilities / lastUpdateAttempt. Nested || merges + // one ref; concurrent UPDATEs on this row serialize and re-evaluate against + // the latest map under READ COMMITTED. await db .update(clients) - .set({ metadata: { ...base, [RPC_METADATA_KEY]: next } }) + .set({ + metadata: sql`jsonb_set( + COALESCE(${clients.metadata}, '{}'::jsonb), + '{modelCatalogRpc}', + COALESCE(${clients.metadata} -> 'modelCatalogRpc', '{}'::jsonb) + || jsonb_build_object(${ref}::text, ${JSON.stringify(entry)}::jsonb), + true + )`, + }) .where(eq(clients.id, clientId)); } -/** Load a previously stored catalog (does not delete — prune happens on write). */ +/** Load a previously stored catalog when still within the logical TTL. */ export async function readModelCatalogRpcResult( db: Database, clientId: string, @@ -87,14 +82,15 @@ export async function readModelCatalogRpcResult( .limit(1); if (!client) return null; const base = (client.metadata ?? {}) as Record; - const entry = asRpcMap(base[RPC_METADATA_KEY])[ref]; + const map = base[RPC_METADATA_KEY]; + if (!map || typeof map !== "object" || Array.isArray(map)) return null; + const entry = asRpcEntry((map as Record)[ref]); return entry?.catalog ?? null; } /** * True when the DB says a daemon WebSocket is live somewhere (this process or - * another replica). Used after a local `sendToClient` miss to decide between - * cross-replica fan-out and a hard 503. + * another replica). Used to decide between cross-replica fan-out and a hard 503. */ export function isClientConnectedSomewhere(client: { status: string; instanceId: string | null }): boolean { return client.status === "connected" && client.instanceId != null;