From af0fabd84f92e3578c8ff1e810337cb6401c9e2a Mon Sep 17 00:00:00 2001 From: Michal Date: Mon, 27 Apr 2026 15:15:46 +0100 Subject: [PATCH] feat(mcplocal+mcpd): wake-recipe config + wake-task execution (v2 Stage 1) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit First half of v2 — mcplocal can now declare a hibernating backend and respond to a `wake` task by running a configured recipe. v2 Stage 2 will wire mcpd to dispatch the wake task before relaying inference. Config (LlmProviderFileEntry): - New \`wake\` block on a published provider: wake: type: http # or: command url: ... # http only method: POST # http only, default POST headers: {...} # http only body: ... # http only command: ... # command only args: [...] # command only maxWaitSeconds: 60 # how long to poll isAvailable() after wake fires Registrar (mcplocal): - At publish time, providers with a wake recipe whose isAvailable() returns false report initialStatus=hibernating to mcpd. Without a wake recipe (legacy v1) or when already up, status stays active. - handleWakeTask: runs the recipe (HTTP request OR child-process spawn), then polls isAvailable() up to maxWaitSeconds, sending a heartbeat each loop so mcpd's GC sweep doesn't time us out mid-boot. Reports { ok, ms } on success or { error } on timeout/recipe failure via the existing _provider-task/:id/result. - Replaces the v1 stub that rejected wake tasks with "not implemented". mcpd VirtualLlmService: - RegisterProviderInput gains optional initialStatus ('active' | 'hibernating'). The register/upsert path uses it for both new and reconnecting rows. Defaults to 'active' so v1 publishers still work unchanged. - Provider-register route's coercer accepts the new field. Tests: 3 new in registrar.test.ts cover initialStatus selection (hibernating when wake configured + unavailable, active otherwise, active when no wake even if unavailable). 8/8 registrar tests, 833/833 mcpd unchanged. Co-Authored-By: Claude Opus 4.7 (1M context) --- src/mcpd/src/routes/virtual-llms.ts | 7 + src/mcpd/src/services/virtual-llm.service.ts | 14 +- src/mcplocal/src/http/config.ts | 17 +++ src/mcplocal/src/main.ts | 1 + src/mcplocal/src/providers/registrar.ts | 140 ++++++++++++++++++- src/mcplocal/tests/registrar.test.ts | 111 ++++++++++++++- 6 files changed, 278 insertions(+), 12 deletions(-) diff --git a/src/mcpd/src/routes/virtual-llms.ts b/src/mcpd/src/routes/virtual-llms.ts index 59664cb..78a0e67 100644 --- a/src/mcpd/src/routes/virtual-llms.ts +++ b/src/mcpd/src/routes/virtual-llms.ts @@ -150,6 +150,7 @@ function coerceProviderInput(raw: unknown): { tier?: string; description?: string; extraConfig?: Record; + initialStatus?: 'active' | 'hibernating'; } { if (raw === null || typeof raw !== 'object') { throw Object.assign(new Error('provider entry must be an object'), { statusCode: 400 }); @@ -170,5 +171,11 @@ function coerceProviderInput(raw: unknown): { if (o['extraConfig'] !== null && typeof o['extraConfig'] === 'object') { out.extraConfig = o['extraConfig'] as Record; } + // Only accept the two values v2 actually defines. Anything else falls + // through to the service default (active) — matches v1 publishers that + // don't know about this field. + if (o['initialStatus'] === 'active' || o['initialStatus'] === 'hibernating') { + out.initialStatus = o['initialStatus']; + } return out; } diff --git a/src/mcpd/src/services/virtual-llm.service.ts b/src/mcpd/src/services/virtual-llm.service.ts index f242843..84710dc 100644 --- a/src/mcpd/src/services/virtual-llm.service.ts +++ b/src/mcpd/src/services/virtual-llm.service.ts @@ -37,6 +37,15 @@ export interface RegisterProviderInput { tier?: string; description?: string; extraConfig?: Record; + /** + * Optional. Lets the publisher hint that the underlying backend is + * asleep — mcpd records the row as `hibernating` and will dispatch a + * `wake` task before any inference. Defaults to `active` (today's + * behavior). v2 publishers (mcplocal with a configured wake recipe) + * pass 'hibernating' when `LlmProvider.isAvailable()` returns false at + * publish time. + */ + initialStatus?: 'active' | 'hibernating'; } export interface RegisterResult { @@ -112,6 +121,7 @@ export class VirtualLlmService implements IVirtualLlmService { const llms: Llm[] = []; for (const p of input.providers) { + const initialStatus = p.initialStatus ?? 'active'; const existing = await this.repo.findByName(p.name); if (existing === null) { const created = await this.repo.create({ @@ -123,7 +133,7 @@ export class VirtualLlmService implements IVirtualLlmService { ...(p.extraConfig !== undefined ? { extraConfig: p.extraConfig } : {}), kind: 'virtual', providerSessionId: sessionId, - status: 'active', + status: initialStatus, lastHeartbeatAt: now, inactiveSince: null, }); @@ -156,7 +166,7 @@ export class VirtualLlmService implements IVirtualLlmService { ...(p.extraConfig !== undefined ? { extraConfig: p.extraConfig } : {}), kind: 'virtual', providerSessionId: sessionId, - status: 'active', + status: initialStatus, lastHeartbeatAt: now, inactiveSince: null, }); diff --git a/src/mcplocal/src/http/config.ts b/src/mcplocal/src/http/config.ts index d70f9c6..1d4a9c3 100644 --- a/src/mcplocal/src/http/config.ts +++ b/src/mcplocal/src/http/config.ts @@ -80,8 +80,25 @@ export interface LlmProviderFileEntry { * Default: false — existing setups don't change behavior. */ publish?: boolean; + /** + * Optional wake recipe for hibernating backends. When set, a provider + * whose `isAvailable()` returns false at registrar start time is + * published as `status=hibernating`. The next inference request that + * lands on the row triggers this recipe; mcplocal polls + * `isAvailable()` until it returns true (or times out) and then flips + * the row to active so mcpd can dispatch the queued inference. + * + * Two recipe types: + * - `http`: POST to a URL (e.g. an external sleep/wake controller) + * - `command`: spawn a shell command (e.g. `systemctl --user start vllm`) + */ + wake?: WakeRecipe; } +export type WakeRecipe = + | { type: 'http'; url: string; method?: 'GET' | 'POST'; headers?: Record; body?: string; maxWaitSeconds?: number } + | { type: 'command'; command: string; args?: string[]; maxWaitSeconds?: number }; + export interface ProjectLlmOverride { model?: string; provider?: string; diff --git a/src/mcplocal/src/main.ts b/src/mcplocal/src/main.ts index df631a9..670ba04 100644 --- a/src/mcplocal/src/main.ts +++ b/src/mcplocal/src/main.ts @@ -215,6 +215,7 @@ async function maybeStartVirtualLlmRegistrar( model: entry.model ?? entry.name, }; if (entry.tier !== undefined) item.tier = entry.tier; + if (entry.wake !== undefined) item.wake = entry.wake; published.push(item); } if (published.length === 0) return null; diff --git a/src/mcplocal/src/providers/registrar.ts b/src/mcplocal/src/providers/registrar.ts index 120e4c9..e09392c 100644 --- a/src/mcplocal/src/providers/registrar.ts +++ b/src/mcplocal/src/providers/registrar.ts @@ -27,7 +27,9 @@ import http from 'node:http'; import https from 'node:https'; import { promises as fs } from 'node:fs'; import { dirname } from 'node:path'; +import { spawn } from 'node:child_process'; import type { LlmProvider, CompletionOptions } from './types.js'; +import type { WakeRecipe } from '../http/config.js'; export interface RegistrarLogger { info: (msg: string) => void; @@ -45,6 +47,13 @@ export interface RegistrarPublishedProvider { tier?: 'fast' | 'heavy'; /** Optional human-readable description for `mcpctl get llm`. */ description?: string; + /** + * Optional wake recipe for backends that hibernate. When provided AND + * `provider.isAvailable()` returns false at registrar start, the row is + * published with status=hibernating; on the first incoming `wake` task + * the registrar runs this recipe and waits for the backend to come up. + */ + wake?: WakeRecipe; } export interface RegistrarOptions { @@ -140,15 +149,28 @@ export class VirtualLlmRegistrar { } private async register(): Promise { - const body: Record = { - providers: this.opts.publishedProviders.map((p) => ({ + // Decide initial status per provider. A provider with a wake recipe + // that's NOT currently available comes up as hibernating; otherwise + // active (today's behavior). isAvailable() is forgiving — any + // unexpected throw is treated as "not available" so a transient + // network blip during boot doesn't crash the registrar. + const providers = await Promise.all(this.opts.publishedProviders.map(async (p) => { + let initialStatus: 'active' | 'hibernating' = 'active'; + if (p.wake !== undefined) { + let alive = false; + try { alive = await p.provider.isAvailable(); } catch { alive = false; } + if (!alive) initialStatus = 'hibernating'; + } + return { name: p.provider.name, type: p.type, model: p.model, ...(p.tier !== undefined ? { tier: p.tier } : {}), ...(p.description !== undefined ? { description: p.description } : {}), - })), - }; + initialStatus, + }; + })); + const body: Record = { providers }; if (this.sessionId !== null) body['providerSessionId'] = this.sessionId; const res = await postJson( @@ -276,9 +298,51 @@ export class VirtualLlmRegistrar { void this.handleInferTask(task); return; } - // Wake tasks are reserved for v2 — acknowledge with an error so mcpd - // surfaces a clean failure rather than waiting forever. - void this.postResult(task.taskId, { error: 'wake task type not implemented in this client (v2)' }); + if (task.kind === 'wake') { + void this.handleWakeTask(task); + return; + } + } + + /** + * Run the configured wake recipe and poll the provider until it comes + * up. Sends a `{ status: 200, body: { ok: true } }` result on success; + * `{ error }` on timeout or recipe failure. While waiting, also bumps + * the heartbeat so mcpd's GC sweep doesn't decide we're stale mid-wake. + */ + private async handleWakeTask(task: { kind: 'wake'; taskId: string; llmName: string }): Promise { + const published = this.opts.publishedProviders.find((p) => p.provider.name === task.llmName); + if (published === undefined) { + await this.postResult(task.taskId, { error: `provider '${task.llmName}' not registered locally` }); + return; + } + if (published.wake === undefined) { + await this.postResult(task.taskId, { error: `provider '${task.llmName}' has no wake recipe configured` }); + return; + } + + try { + await runWakeRecipe(published.wake); + // Poll isAvailable() until it comes up (or timeout). Heartbeat + // every poll tick so mcpd doesn't time us out while we're waiting + // on a slow boot. + const maxWaitMs = (published.wake.maxWaitSeconds ?? 60) * 1000; + const started = Date.now(); + while (Date.now() - started < maxWaitMs) { + let alive = false; + try { alive = await published.provider.isAvailable(); } catch { alive = false; } + if (alive) { + await this.heartbeatOnce(); + await this.postResult(task.taskId, { status: 200, body: { ok: true, ms: Date.now() - started } }); + return; + } + await this.heartbeatOnce(); + await new Promise((r) => setTimeout(r, 1500)); + } + await this.postResult(task.taskId, { error: `provider '${task.llmName}' did not come up within ${String(maxWaitMs)}ms` }); + } catch (err) { + await this.postResult(task.taskId, { error: `wake recipe failed: ${(err as Error).message}` }); + } } private async handleInferTask(task: InferTask): Promise { @@ -373,6 +437,68 @@ function openAiStreamChunk( }; } +/** + * Execute a wake recipe. Returns when the recipe completes; throws if it + * fails. Doesn't itself poll for provider readiness — that's the caller's + * job (handleWakeTask polls isAvailable() with its own timeout). + * + * `http`: fires the configured request and considers any 2xx a success. + * The remote service is expected to be a "wake controller" that returns + * quickly; if the underlying boot takes minutes, the controller should + * return 202 and the readiness poll catches up. + * + * `command`: spawns the binary with args, waits for exit. Non-zero exit + * is treated as failure. stdout/stderr are discarded — the recipe's job + * is to *trigger* a wake, not to produce output. + */ +async function runWakeRecipe(recipe: WakeRecipe): Promise { + if (recipe.type === 'http') { + const u = new URL(recipe.url); + const driver = u.protocol === 'https:' ? https : http; + const method = recipe.method ?? 'POST'; + const headers: Record = { ...(recipe.headers ?? {}) }; + const body = recipe.body; + if (body !== undefined) { + headers['Content-Length'] = String(Buffer.byteLength(body)); + } + await new Promise((resolve, reject) => { + const req = driver.request({ + hostname: u.hostname, + port: u.port || (u.protocol === 'https:' ? 443 : 80), + path: u.pathname + u.search, + method, + headers, + timeout: 30_000, + }, (res) => { + const status = res.statusCode ?? 0; + // Drain so the socket can be reused/freed. + res.resume(); + if (status >= 200 && status < 300) resolve(); + else reject(new Error(`wake HTTP returned ${String(status)}`)); + }); + req.on('error', reject); + req.on('timeout', () => { req.destroy(); reject(new Error('wake HTTP timed out')); }); + if (body !== undefined) req.write(body); + req.end(); + }); + return; + } + if (recipe.type === 'command') { + await new Promise((resolve, reject) => { + const child = spawn(recipe.command, recipe.args ?? [], { + stdio: 'ignore', + }); + child.on('error', reject); + child.on('exit', (code) => { + if (code === 0) resolve(); + else reject(new Error(`wake command exited with code ${String(code)}`)); + }); + }); + return; + } + throw new Error(`unknown wake recipe type`); +} + interface PostResponse { statusCode: number; body: string } /** Tiny JSON POST helper used by all of the registrar's mcpd calls. */ diff --git a/src/mcplocal/tests/registrar.test.ts b/src/mcplocal/tests/registrar.test.ts index 40e99ce..f6ce8d9 100644 --- a/src/mcplocal/tests/registrar.test.ts +++ b/src/mcplocal/tests/registrar.test.ts @@ -142,13 +142,17 @@ describe('VirtualLlmRegistrar', () => { const registerCall = fake.calls.find((c) => c.path === '/api/v1/llms/_provider-register'); expect(registerCall).toBeDefined(); expect(registerCall!.method).toBe('POST'); - const body = JSON.parse(registerCall!.body) as { providers: Array<{ name: string; type: string; model: string; tier: string }> }; - expect(body.providers).toEqual([{ + const body = JSON.parse(registerCall!.body) as { providers: Array> }; + expect(body.providers).toHaveLength(1); + expect(body.providers[0]).toMatchObject({ name: 'vllm-local', type: 'openai', model: 'qwen', tier: 'fast', - }]); + // v2 always sends initialStatus; defaults to 'active' when no + // wake recipe is configured. + initialStatus: 'active', + }); expect(registerCall!.headers['authorization']).toBe('Bearer tok-abc'); // Sticky session id persisted. @@ -219,6 +223,107 @@ describe('VirtualLlmRegistrar', () => { } }); + // ── v2: hibernating + wake recipe ── + + it('publishes initialStatus=hibernating when provider is unavailable AND wake is configured', async () => { + const fake = await startFakeServer(); + try { + const sleeping: LlmProvider = { + name: 'vllm-local', + async complete() { throw new Error('not running'); }, + async listModels() { return []; }, + async isAvailable() { return false; }, + }; + const registrar = new VirtualLlmRegistrar({ + mcpdUrl: fake.url, + token: 't', + publishedProviders: [{ + provider: sleeping, + type: 'openai', + model: 'm', + wake: { type: 'http', url: 'http://localhost:9999/wake', maxWaitSeconds: 1 }, + }], + sessionFilePath: join(tempDir, 'provider-session'), + log: silentLog(), + heartbeatIntervalMs: 60_000, + }); + await registrar.start(); + await new Promise((r) => setTimeout(r, 20)); + + const registerCall = fake.calls.find((c) => c.path === '/api/v1/llms/_provider-register'); + const body = JSON.parse(registerCall!.body) as { providers: Array<{ initialStatus?: string }> }; + expect(body.providers[0]!.initialStatus).toBe('hibernating'); + registrar.stop(); + } finally { + await fake.close(); + } + }); + + it('publishes initialStatus=active when provider is available even with a wake recipe', async () => { + const fake = await startFakeServer(); + try { + const awake: LlmProvider = { + name: 'vllm-local', + async complete() { throw new Error('not used'); }, + async listModels() { return []; }, + async isAvailable() { return true; }, + }; + const registrar = new VirtualLlmRegistrar({ + mcpdUrl: fake.url, + token: 't', + publishedProviders: [{ + provider: awake, + type: 'openai', + model: 'm', + wake: { type: 'http', url: 'http://localhost:9999/wake' }, + }], + sessionFilePath: join(tempDir, 'provider-session'), + log: silentLog(), + heartbeatIntervalMs: 60_000, + }); + await registrar.start(); + await new Promise((r) => setTimeout(r, 20)); + + const registerCall = fake.calls.find((c) => c.path === '/api/v1/llms/_provider-register'); + const body = JSON.parse(registerCall!.body) as { providers: Array<{ initialStatus?: string }> }; + expect(body.providers[0]!.initialStatus).toBe('active'); + registrar.stop(); + } finally { + await fake.close(); + } + }); + + it('publishes initialStatus=active when no wake recipe is configured (legacy path)', async () => { + const fake = await startFakeServer(); + try { + // Provider intentionally returns false but has no wake recipe → + // legacy v1 publishers don't get hibernation behavior. + const sleeping: LlmProvider = { + name: 'vllm-local', + async complete() { return { content: '', toolCalls: [], usage: { promptTokens: 0, completionTokens: 0, totalTokens: 0 }, finishReason: 'stop' }; }, + async listModels() { return []; }, + async isAvailable() { return false; }, + }; + const registrar = new VirtualLlmRegistrar({ + mcpdUrl: fake.url, + token: 't', + publishedProviders: [{ provider: sleeping, type: 'openai', model: 'm' }], + sessionFilePath: join(tempDir, 'provider-session'), + log: silentLog(), + heartbeatIntervalMs: 60_000, + }); + await registrar.start(); + await new Promise((r) => setTimeout(r, 20)); + + const registerCall = fake.calls.find((c) => c.path === '/api/v1/llms/_provider-register'); + const body = JSON.parse(registerCall!.body) as { providers: Array<{ initialStatus?: string }> }; + expect(body.providers[0]!.initialStatus).toBe('active'); + registrar.stop(); + } finally { + await fake.close(); + } + }); + it('throws when mcpd returns non-201 from /_provider-register', async () => { const fake = await startFakeServer(); fake.handler = (_req, res, _body) => {