From a61e39bb235f477dd87e0cc0cce8a6e0814900ff Mon Sep 17 00:00:00 2001 From: Muse Sweep Date: Sat, 24 Oct 2026 16:13:09 +0000 Subject: [PATCH] feat(muse-proxy): warm sessions on a persistent muse serve (MSP) host (#480) One persistent 'muse serve' process hosts durable MSP sessions keyed by sha256(model + system + first user message), persisted to warm-sessions.json in the workdir (LRU 64, 2h TTL) and resumed after proxy/host restarts. Keyed hits send only the new user turn; the provider prompt cache then serves the ~26k scaffolding tokens as cache reads (measured: turn 2 input 26,455 tokens, 26,225 of them cacheReadTokens). Diverged histories re-bootstrap as a fresh session whose first turn is the flattened transcript (today's exec behaviour); any warm failure before the response is committed falls back to the spawn-per-request exec path, and MUSE_PROXY_WARM=0 disables the warm path. Raw MSP (NDJSON JSON-RPC 2.0) over stdlib only - no SDK runtime dependency. Also publishes endpoint.json like the sibling proxies and pins NO_PROXY loopback for child processes. --- bin/muse-proxy.mjs | 628 +++++++++++++++++++++++++++++++++++++++- bin/muse-proxy.test.mjs | 176 ++++++++--- 2 files changed, 759 insertions(+), 45 deletions(-) diff --git a/bin/muse-proxy.mjs b/bin/muse-proxy.mjs index 1a0531c72..a2d504778 100644 --- a/bin/muse-proxy.mjs +++ b/bin/muse-proxy.mjs @@ -5,6 +5,11 @@ // subprocess per chat-completions request, prompt delivered via --prompt-file, // streamed back out of the CLI's JSONL event stream. // +// Warm path (#480): chat completions are served by one persistent +// `muse serve` (MSP) host with durable per-conversation sessions whenever +// possible — see the warm-sessions section below. The exec path remains as +// the fallback for any warm failure and when MUSE_PROXY_WARM=0. +// // Auth is the CLI's own OAuth login (~/.config/muse/auth.json). META_API_KEY // must NOT be set for the child: the CLI binary states it takes priority over // the account login, and the Model API key is a separate pay-as-you-go wallet @@ -147,9 +152,578 @@ const childEnv = () => { env.https_proxy = EGRESS_PROXY; env.http_proxy = EGRESS_PROXY; env.SSL_CERT_FILE = EGRESS_CA; + env.NO_PROXY = [env.NO_PROXY, '127.0.0.1', 'localhost', '::1'].filter(Boolean).join(','); + env.no_proxy = env.NO_PROXY; return env; }; +// --------------------------------------------------------------------------- +// Warm sessions (#480): one persistent `muse serve` (MSP over stdio) host +// --------------------------------------------------------------------------- +// The exec path above pays a fresh CLI spawn (+ ~26k tokens of session +// scaffolding) for every request. `muse serve` instead hosts durable MSP +// sessions in one long-lived process: a conversation's first turn creates a +// session, later turns send only the new user message, and the provider +// serves the scaffolding from its prompt cache (verified on the sandbox: +// turn 2 of a served session reported 26,225 cacheReadTokens of 26,455 +// input tokens — ~230 uncached — vs 26,708 fully re-read input tokens for +// the same turn through the exec path). +// +// MSP is newline-delimited JSON-RPC 2.0 on the host's stdio (schema: +// `muse schema generate-json-schema`), so this is driven with the standard +// library only — no @muse-code/sdk runtime dependency. +// +// Conversations are keyed by sha256(model + system + first user message). +// The key -> session map persists to warm-sessions.json in the workdir +// (LRU 64 entries, 2h idle TTL), so sessions survive proxy restarts: the +// sessions themselves are durable in the CLI's own store and are re-attached +// with session/resume on the new host. A request whose history no longer +// matches the mapped session (diverged prefix, extra user turns, expired or +// lost session) is re-bootstrapped as a fresh session whose first turn is +// the full flattened transcript — exactly what the exec path sends on every +// request. Any warm-path failure before the response is committed falls +// back to the exec path; the warm path can be disabled with +// MUSE_PROXY_WARM=0. + +const WARM_ENABLED = process.env.MUSE_PROXY_WARM !== '0'; +const WARM_STATE_FILE = + process.env.MUSE_PROXY_WARM_STATE || path.join(WORKDIR, 'warm-sessions.json'); +const WARM_TTL_MS = 2 * 60 * 60 * 1000; +const WARM_MAX_SESSIONS = 64; +const WARM_TURN_TIMEOUT_MS = 300 * 1000; + +// MSP command ids must be UUIDv7 (the host rejects v4: "expected UUIDv7"). +let uuidV7LastMs = -1; +let uuidV7Seq = 0; +const uuidv7 = () => { + let now = Date.now(); + if (now <= uuidV7LastMs) { + now = uuidV7LastMs; + uuidV7Seq += 1; + if (uuidV7Seq > 0xfff) { now += 1; uuidV7Seq = 0; } + } else { + uuidV7Seq = 0; + } + uuidV7LastMs = now; + const bytes = crypto.randomBytes(16); + let t = now; + for (let i = 5; i >= 0; i -= 1) { bytes[i] = t & 0xff; t = Math.floor(t / 256); } + bytes[6] = 0x70 | ((uuidV7Seq >> 8) & 0x0f); + bytes[7] = uuidV7Seq & 0xff; + bytes[8] = (bytes[8] & 0x3f) | 0x80; + const hex = bytes.toString('hex'); + return `${hex.slice(0, 8)}-${hex.slice(8, 12)}-${hex.slice(12, 16)}-${hex.slice(16, 20)}-${hex.slice(20)}`; +}; + +const messageText = (m) => { + if (!m || typeof m !== 'object') return ''; + if (typeof m.content === 'string') return m.content; + if (Array.isArray(m.content)) { + return m.content.map(c => (c && typeof c.text === 'string') ? c.text : '').join('\n'); + } + return m.content ? JSON.stringify(m.content) : ''; +}; + +const messageHash = (m) => + crypto.createHash('sha256').update(`${m.role}\x00${messageText(m)}`).digest('hex'); + +const conversationKey = (messages, baseModel) => { + const system = messages.filter(m => m.role === 'system').map(messageText).join('\n\n'); + const firstUser = messages.find(m => m.role === 'user'); + return crypto.createHash('sha256') + .update(`${baseModel}\x00${system}\x00${firstUser ? messageText(firstUser) : ''}`) + .digest('hex'); +}; + +// --- persistent key -> session map ----------------------------------------- +const warmStore = { + map: new Map(), + loaded: false, + load() { + if (this.loaded) return; + this.loaded = true; + try { + const raw = JSON.parse(fs.readFileSync(WARM_STATE_FILE, 'utf-8')); + const now = Date.now(); + for (const e of raw.entries || []) { + if (e && e.key && e.sessionId && now - (e.lastUsedAt || 0) <= WARM_TTL_MS) { + this.map.set(e.key, e); + } + } + } catch { /* no state yet */ } + }, + save() { + try { + const entries = [...this.map.values()].map(e => ({ + key: e.key, sessionId: e.sessionId, model: e.model, + hashes: e.hashes, sentCount: e.sentCount, lastUsedAt: e.lastUsedAt + })); + const tmp = `${WARM_STATE_FILE}.tmp`; + fs.writeFileSync(tmp, JSON.stringify({ entries }, null, 2)); + fs.renameSync(tmp, WARM_STATE_FILE); + } catch (err) { + console.warn(`[muse-proxy warm] state save failed: ${err.message}`); + } + }, + get(key) { + this.load(); + const entry = this.map.get(key); + if (!entry) return null; + if (Date.now() - (entry.lastUsedAt || 0) > WARM_TTL_MS) { + this.map.delete(key); + this.save(); + return null; + } + return entry; + }, + set(key, entry) { + this.load(); + this.map.set(key, entry); + if (this.map.size > WARM_MAX_SESSIONS) { + const byAge = [...this.map.values()].sort((a, b) => (a.lastUsedAt || 0) - (b.lastUsedAt || 0)); + for (const old of byAge.slice(0, this.map.size - WARM_MAX_SESSIONS)) { + this.map.delete(old.key); + warmHost.leave(old.sessionId); + } + } + this.save(); + }, + delete(key) { + this.map.delete(key); + this.save(); + } +}; + +// --- the one persistent serve host ------------------------------------------ +const warmHost = { + child: null, + readyPromise: null, + gen: 0, + seq: 0, + pending: new Map(), + listeners: new Set(), + stderrTail: '', + + ensure() { + if (this.readyPromise) return this.readyPromise; + this.readyPromise = this._start().catch(err => { + this.readyPromise = null; + throw err; + }); + return this.readyPromise; + }, + + async _start() { + const args = ['serve', '--provider', 'meta', '--disable-request-user-input']; + console.log(`[muse-proxy warm] Spawning serve host: ${MUSE_BIN} ${args.join(' ')}`); + const child = spawn(MUSE_BIN, args, { + cwd: WORKDIR, + env: childEnv(), + stdio: ['pipe', 'pipe', 'pipe'] + }); + this.child = child; + this.stderrTail = ''; + let lineBuffer = ''; + child.stdout.on('data', chunk => { + lineBuffer += chunk.toString('utf-8'); + const lines = lineBuffer.split('\n'); + lineBuffer = lines.pop(); + for (const line of lines) this._onLine(line); + }); + child.stderr.on('data', chunk => { + const msg = chunk.toString('utf-8'); + this.stderrTail = (this.stderrTail + msg).slice(-2000); + if (msg.trim().length > 0) console.warn(`[muse-proxy warm stderr] ${msg.trim()}`); + }); + child.on('error', err => this._onExit(`spawn error: ${err.message}`)); + child.on('close', (code, signal) => this._onExit(`serve host exited (code ${code}, signal ${signal})`)); + let result; + try { + result = await this.rpc('initialize', { + clientInfo: { name: 'muse_proxy', version: '1.0.0' }, + capabilities: { userInputDialogs: false } + }, 20000); + this._write({ jsonrpc: '2.0', method: 'initialized' }); + } catch (err) { + try { child.kill('SIGKILL'); } catch { /* already gone */ } + if (this.child === child) this.child = null; + throw err; + } + this.gen += 1; + console.log(`[muse-proxy warm] serve host ready (gen ${this.gen}, server ${result?.serverInfo?.name} ${result?.serverInfo?.version})`); + }, + + _onExit(reason) { + if (!this.child && this.pending.size === 0) return; + this.child = null; + this.readyPromise = null; + const err = new Error(reason); + for (const p of this.pending.values()) p.reject(err); + this.pending.clear(); + for (const fn of [...this.listeners]) { + try { fn({ method: 'host/died', params: { reason } }); } catch { /* listener */ } + } + }, + + _write(obj) { + if (!this.child || !this.child.stdin.writable) throw new Error('serve host stdin is not writable'); + this.child.stdin.write(JSON.stringify(obj) + '\n'); + }, + + _onLine(line) { + const trimmed = line.trim(); + if (!trimmed || !trimmed.startsWith('{')) return; + let msg; + try { msg = JSON.parse(trimmed); } catch { return; } + if (msg.id !== undefined && msg.method === undefined && this.pending.has(msg.id)) { + const p = this.pending.get(msg.id); + this.pending.delete(msg.id); + if (msg.error) p.reject(new Error(`MSP ${msg.error.message || 'error'} (code ${msg.error.code})`)); + else p.resolve(msg.result); + return; + } + if (msg.id !== undefined && typeof msg.method === 'string') { + // Server-initiated request (approval/request, userInput/request): + // acknowledge presentation only, as the SDK's receipt does. Sessions + // run with approvalMode allowAll and request-user-input withheld, so + // no decision is expected to follow. + try { this._write({ jsonrpc: '2.0', id: msg.id, result: {} }); } catch { /* host gone */ } + console.warn(`[muse-proxy warm] server request ${msg.method} acknowledged`); + return; + } + if (typeof msg.method === 'string') { + for (const fn of [...this.listeners]) { + try { fn(msg); } catch { /* listener */ } + } + } + }, + + rpc(method, params, timeoutMs = 30000) { + return new Promise((resolve, reject) => { + const id = ++this.seq; + const timer = setTimeout(() => { + this.pending.delete(id); + reject(new Error(`MSP ${method} timed out after ${timeoutMs}ms`)); + }, timeoutMs); + this.pending.set(id, { + resolve: v => { clearTimeout(timer); resolve(v); }, + reject: e => { clearTimeout(timer); reject(e); } + }); + try { + this._write({ jsonrpc: '2.0', id, method, params }); + } catch (err) { + this.pending.delete(id); + clearTimeout(timer); + reject(err); + } + }); + }, + + onNotification(fn) { + this.listeners.add(fn); + return () => this.listeners.delete(fn); + }, + + leave(sessionId) { + if (!sessionId || !this.readyPromise) return; + this.rpc('session/leave', { commandId: uuidv7(), sessionId }, 5000).catch(() => {}); + }, + + async startSession(baseModel) { + const result = await this.rpc('session/start', { + commandId: uuidv7(), + modelId: baseModel, + approvalMode: 'allowAll', + workspaceRoots: [WORKDIR] + }); + const sessionId = result?.session?.sessionId; + if (!sessionId) throw new Error('session/start returned no sessionId'); + return sessionId; + }, + + async resumeSession(sessionId) { + const result = await this.rpc('session/resume', { + commandId: uuidv7(), + sessionId, + excludeItems: true + }); + const resumed = result?.session?.sessionId; + if (!resumed) throw new Error('session/resume returned no session'); + return resumed; + }, + + // Run one turn on a loaded session. Streams agentMessage text deltas to + // onDelta (when given) and resolves with the turn's full text + usage. + runTurn({ sessionId, text, effort, onDelta }) { + return new Promise((resolve, reject) => { + let turnId = null; + let deltasText = ''; + let lastAgentText = null; + let lastFinalText = null; + let usage = null; + const itemKinds = new Map(); + let settled = false; + const settle = (fn, value) => { + if (settled) return; + settled = true; + clearTimeout(timer); + off(); + fn(value); + }; + const timer = setTimeout(() => { + settle(reject, new Error(`warm turn timed out after ${WARM_TURN_TIMEOUT_MS}ms`)); + }, WARM_TURN_TIMEOUT_MS); + const off = this.onNotification(msg => { + const params = msg.params || {}; + if (params.sessionId !== sessionId && msg.method !== 'host/died') return; + switch (msg.method) { + case 'host/died': + settle(reject, new Error(`serve host died mid-turn: ${params.reason}`)); + break; + case 'item/started': + if (params.item && params.item.itemId) itemKinds.set(params.item.itemId, params.item.kind); + break; + case 'item/delta': { + const kind = itemKinds.get(params.itemId); + if ((params.field === undefined || params.field === 'text') && + (kind === undefined || kind === 'agentMessage') && + typeof params.delta === 'string') { + deltasText += params.delta; + if (onDelta) { + try { onDelta(params.delta); } catch { /* response gone */ } + } + } + break; + } + case 'item/completed': + if (params.item && params.item.kind === 'agentMessage' && typeof params.item.text === 'string') { + lastAgentText = params.item.text; + if (params.item.messagePhase === 'finalAnswer') lastFinalText = params.item.text; + } + break; + case 'session/tokenUsage': + if (params.usage) { + usage = { + promptTokens: params.usage.inputTokens, + completionTokens: params.usage.outputTokens + }; + } + break; + case 'turn/completed': + if (turnId && params.turnId !== turnId) break; + if (params.terminal !== 'completed') { + settle(reject, new Error(`warm turn terminated with state '${params.terminal}'${params.reason ? `: ${params.reason}` : ''}`)); + } else { + settle(resolve, { + deltasText, + fullText: lastFinalText ?? lastAgentText ?? deltasText, + promptTokens: usage?.promptTokens, + completionTokens: usage?.completionTokens + }); + } + break; + default: + break; + } + }); + this.rpc('turn/start', { + commandId: uuidv7(), + sessionId, + input: [{ type: 'text', text }], + reasoningEffort: effort + }).then(ack => { + turnId = ack?.turnId || turnId; + }).catch(err => settle(reject, err)); + }); + } +}; + +// Per-conversation serialization: turns on one session must not interleave, +// and the delta decision must be made against the latest committed state. +const warmChains = new Map(); +const enqueueWarm = (key, job) => { + const prev = warmChains.get(key) || Promise.resolve(); + const run = prev.then(job, job); + warmChains.set(key, run.catch(() => {})); + return run; +}; + +// Attempt the chat completion on the warm serve path. Returns true when the +// response has been fully handled (success, or a stream already committed +// that had to be terminated in place); false when nothing was written and +// the caller should fall back to the exec path. +const tryWarmChat = async ({ messages, baseModel, effort, fullName, stream, res, chatId, created }) => { + const key = conversationKey(messages, baseModel); + let headed = false; + let streamedText = ''; + const writeChunk = (content) => { + if (!stream || res.writableEnded) return; + if (!headed) { + headed = true; + res.writeHead(200, { + 'Content-Type': 'text/event-stream; charset=utf-8', + 'Cache-Control': 'no-cache, no-transform', + 'Connection': 'keep-alive', + 'X-Accel-Buffering': 'no' + }); + } + const chunk = { + id: chatId, + object: 'chat.completion.chunk', + created, + model: fullName, + choices: [{ index: 0, delta: { content }, finish_reason: null }] + }; + res.write(`data: ${JSON.stringify(chunk)}\n\n`); + streamedText += content; + }; + + const job = async () => { + await warmHost.ensure(); + let entry = warmStore.get(key); + let sessionId = null; + let turnText = null; + const hashes = messages.map(messageHash); + + if (entry && entry.sessionId && entry.model === baseModel) { + if (entry.hostGen === warmHost.gen) { + sessionId = entry.sessionId; + } else { + try { + sessionId = await warmHost.resumeSession(entry.sessionId); + entry.hostGen = warmHost.gen; + console.log(`[muse-proxy warm] resumed session ${sessionId} for key ${key.slice(0, 8)}`); + } catch (err) { + console.warn(`[muse-proxy warm] resume failed (${err.message}); re-bootstrapping`); + sessionId = null; + } + } + } + + if (sessionId) { + const last = messages[messages.length - 1]; + const suffix = messages.slice(entry.sentCount); + const prefixOk = entry.sentCount > 0 && + entry.hashes.length >= entry.sentCount && + entry.hashes.slice(0, entry.sentCount).every((h, i) => h === hashes[i]); + const midUsers = suffix.slice(0, -1).filter(m => m.role === 'user').length; + if (prefixOk && suffix.length > 0 && last.role === 'user' && midUsers === 0) { + turnText = messageText(last); + } else { + console.log(`[muse-proxy warm] history diverged for key ${key.slice(0, 8)}; re-bootstrapping`); + sessionId = null; + } + } + + if (!sessionId) { + const replacedSessionId = entry && entry.sessionId; + sessionId = await warmHost.startSession(baseModel); + if (replacedSessionId && replacedSessionId !== sessionId) warmHost.leave(replacedSessionId); + turnText = formatPrompt(messages); + entry = { + key, sessionId, model: baseModel, + hashes: [], sentCount: 0, lastUsedAt: Date.now(), hostGen: warmHost.gen + }; + warmStore.set(key, entry); + console.log(`[muse-proxy warm] new session ${sessionId} for key ${key.slice(0, 8)}`); + } + + const result = await warmHost.runTurn({ + sessionId, + text: turnText, + effort, + onDelta: stream ? writeChunk : null + }); + + entry.sessionId = sessionId; + entry.hashes = hashes; + entry.sentCount = messages.length; + entry.lastUsedAt = Date.now(); + entry.hostGen = warmHost.gen; + warmStore.set(key, entry); + return result; + }; + + let result; + try { + result = await enqueueWarm(key, job); + } catch (err) { + console.warn(`[muse-proxy warm] warm path failed (${err.message}); ${headed ? 'stream already committed, terminating in place' : 'falling back to exec'}`); + if (!headed) return false; + if (!res.writableEnded) { + res.write(`data: ${JSON.stringify({ error: { message: `Muse execution failed: ${err.message}` } })}\n\n`); + res.write('data: [DONE]\n\n'); + } + res.end(); + return true; + } + + // Reconcile deltas with the authoritative final text (same rule as the + // exec path, #484): only a missing tail can still be appended to a stream. + let fullContent = result.deltasText; + let unstreamedSuffix = ''; + if (typeof result.fullText === 'string' && result.fullText.length >= fullContent.length) { + if (result.fullText.startsWith(fullContent)) unstreamedSuffix = result.fullText.slice(fullContent.length); + fullContent = result.fullText; + } + const promptTokens = (typeof result.promptTokens === 'number' && result.promptTokens > 0) + ? result.promptTokens + : Math.max(1, Math.ceil(formatPrompt(messages).length / 4)); + let completionTokens = (typeof result.completionTokens === 'number' && result.completionTokens > 0) + ? result.completionTokens + : 0; + if (completionTokens === 0 && fullContent.length > 0) { + completionTokens = Math.max(1, Math.ceil(fullContent.length / 4)); + } + + if (stream) { + if (unstreamedSuffix.length > 0) writeChunk(unstreamedSuffix); + if (!headed) { + headed = true; + res.writeHead(200, { + 'Content-Type': 'text/event-stream; charset=utf-8', + 'Cache-Control': 'no-cache, no-transform', + 'Connection': 'keep-alive', + 'X-Accel-Buffering': 'no' + }); + } + const endChunk = { + id: chatId, + object: 'chat.completion.chunk', + created, + model: fullName, + choices: [{ index: 0, delta: {}, finish_reason: 'stop' }], + usage: { + prompt_tokens: promptTokens, + completion_tokens: completionTokens, + total_tokens: promptTokens + completionTokens + } + }; + res.write(`data: ${JSON.stringify(endChunk)}\n\n`); + res.write('data: [DONE]\n\n'); + res.end(); + } else { + res.writeHead(200, { 'Content-Type': 'application/json' }); + res.end(JSON.stringify({ + id: chatId, + object: 'chat.completion', + created, + model: fullName, + choices: [{ + index: 0, + message: { role: 'assistant', content: fullContent }, + finish_reason: 'stop' + }], + usage: { + prompt_tokens: promptTokens, + completion_tokens: completionTokens, + total_tokens: promptTokens + completionTokens + } + }, null, 2)); + } + return true; +}; + const server = http.createServer((req, res) => { res.setHeader('Access-Control-Allow-Origin', '*'); res.setHeader('Access-Control-Allow-Methods', 'GET, POST, OPTIONS'); @@ -172,7 +746,13 @@ const server = http.createServer((req, res) => { port: PORT, binary: MUSE_BIN, defaultModel: 'muse-spark-1.3-high', - endpoint: '/v1/chat/completions' + endpoint: '/v1/chat/completions', + warm: { + enabled: WARM_ENABLED, + hostGen: warmHost.gen, + hostUp: Boolean(warmHost.child), + sessions: warmStore.map.size + } }, null, 2)); return; } @@ -216,7 +796,7 @@ const server = http.createServer((req, res) => { } }); - req.on('end', () => { + req.on('end', async () => { let body; try { body = JSON.parse(bodyText); @@ -240,6 +820,23 @@ const server = http.createServer((req, res) => { const chatId = `chatcmpl-muse-${crypto.randomUUID()}`; const created = Math.floor(Date.now() / 1000); + // Warm path first (#480): a persistent serve host + MSP session + // answers when it can; any failure before the response is committed + // returns false and the exec path below takes over unchanged. + if (WARM_ENABLED) { + let warmHandled = false; + try { + warmHandled = await tryWarmChat({ + messages: body.messages, baseModel, effort, fullName, stream, res, chatId, created + }); + } catch (err) { + console.warn(`[muse-proxy warm] unexpected warm error (${err.message}); falling back to exec`); + warmHandled = false; + } + if (warmHandled) return; + if (res.writableEnded) return; + } + // Prompt goes in via --prompt-file (no argv length/escaping concerns) const promptFile = path.join(WORKDIR, `prompt-${crypto.randomUUID()}.txt`); fs.writeFileSync(promptFile, prompt, 'utf-8'); @@ -510,7 +1107,34 @@ const server = http.createServer((req, res) => { res.end(JSON.stringify({ error: { message: `Route not found: ${req.method} ${pathname}` } })); }); + +// Publish the bound endpoint so other tools can find this proxy without +// knowing its port in advance: $XDG_STATE_HOME//endpoint.json +// (default ~/.local/state//endpoint.json), written on listen, +// removed on exit. +const publishEndpoint = () => { + try { + const stateHome = + process.env.XDG_STATE_HOME || path.join(os.homedir(), '.local', 'state'); + const dir = path.join(stateHome, 'muse-proxy'); + fs.mkdirSync(dir, { recursive: true }); + const file = path.join(dir, 'endpoint.json'); + fs.writeFileSync( + `${file}.tmp`, + JSON.stringify({ host: HOST, port: PORT, pid: process.pid, url: `http://${HOST}:${PORT}` }) + '\n', + ); + fs.renameSync(`${file}.tmp`, file); + const cleanup = () => { + try { fs.unlinkSync(file); } catch {} + }; + process.on('exit', cleanup); + process.on('SIGTERM', () => { cleanup(); process.exit(0); }); + process.on('SIGINT', () => { cleanup(); process.exit(0); }); + } catch {} +}; + server.listen(PORT, HOST, () => { + publishEndpoint(); console.log(`[muse-proxy] OpenAI-compatible bridge listening on http://${HOST}:${PORT}`); console.log(`[muse-proxy] Resolved muse binary: ${MUSE_BIN}`); console.log(`[muse-proxy] Workdir: ${WORKDIR}`); diff --git a/bin/muse-proxy.test.mjs b/bin/muse-proxy.test.mjs index b19bf4132..281c6bc8e 100644 --- a/bin/muse-proxy.test.mjs +++ b/bin/muse-proxy.test.mjs @@ -1,4 +1,6 @@ -// Tests for muse-proxy's streaming/terminal-text reconciliation (#484). +// Tests for muse-proxy's streaming/terminal-text reconciliation (#484) +// and its warm serve/MSP session path (#480). The exec-path tests pin +// MUSE_PROXY_WARM=0; the warm tests drive a stub `muse serve` host. // Run: node --test bin/muse-proxy.test.mjs // // The proxy spawns one `muse exec --json` child per request. These tests @@ -17,45 +19,82 @@ import { fileURLToPath } from 'node:url'; const binDir = path.dirname(fileURLToPath(import.meta.url)); -const STUB = `#!/bin/sh -# The prompt file path is the argument after --prompt-file. -prompt="" -prev="" -for a in "$@"; do - if [ "$prev" = "--prompt-file" ]; then prompt="$a"; fi - prev="$a" -done -if grep -q FAILME "$prompt" 2>/dev/null; then - echo "stub muse: auth expired" >&2 - exit 3 -fi -if grep -q FAILPARTIAL "$prompt" 2>/dev/null; then - printf '%s\\n' '{"payload_type":"run.output.delta","payload":{"text":"partial answer"}}' - printf '%s\\n' '{"payload_type":"run.terminal.failed","payload":{"terminal":"failed","text":"partial answer"}}' - echo "stub muse: run failed midway" >&2 - exit 0 -fi -if grep -q CRASHPARTIAL "$prompt" 2>/dev/null; then - printf '%s\\n' '{"payload_type":"run.output.delta","payload":{"text":"partial answer"}}' - echo "stub muse: crashed midway" >&2 - exit 3 -fi -if grep -q FAILSTATE "$prompt" 2>/dev/null; then - printf '%s\\n' '{"payload_type":"run.terminal.failed","payload":{"terminal":"failed"}}' - exit 0 -fi -if grep -q SIGKILLME "$prompt" 2>/dev/null; then - kill -9 $$ -fi -if grep -q DIVERGE "$prompt" 2>/dev/null; then - printf '%s\\n' '{"payload_type":"run.output.delta","payload":{"text":"abc"}}' - printf '%s\\n' '{"payload_type":"run.terminal.completed","payload":{"terminal":"completed","text":"xyz-longer"}}' -else - printf '%s\\n' '{"payload_type":"run.output.delta","payload":{"text":"Hel"}}' - printf '%s\\n' '{"payload_type":"run.usage","payload":{"models":[{"input_tokens":100,"output_tokens":3}]}}' - printf '%s\\n' '{"payload_type":"run.terminal.completed","payload":{"terminal":"completed","text":"Hello world"}}' -fi -`; +const STUB = `#!/usr/bin/env node +// Stub muse: 'exec' reproduces the fixed JSONL behaviours the exec-path +// tests assert on; 'serve' speaks just enough MSP (NDJSON JSON-RPC) for +// the warm-path tests, logging session/turn traffic to $STUB_LOG. +const fs = require('fs'); +const args = process.argv.slice(2); +if (args[0] === 'serve') { + if (process.env.STUB_SERVE_FAIL === '1') process.exit(3); + const log = (line) => { if (process.env.STUB_LOG) fs.appendFileSync(process.env.STUB_LOG, line + '\\n'); }; + let buf = ''; + process.stdin.on('data', (d) => { + buf += d.toString('utf-8'); + let i; + while ((i = buf.indexOf('\\n')) >= 0) { + const line = buf.slice(0, i); buf = buf.slice(i + 1); + if (!line.trim()) continue; + let msg; try { msg = JSON.parse(line); } catch { continue; } + if (msg.id === undefined || !msg.method) continue; + const reply = (result) => process.stdout.write(JSON.stringify({ jsonrpc: '2.0', id: msg.id, result }) + '\\n'); + const notify = (method, params) => process.stdout.write(JSON.stringify({ jsonrpc: '2.0', method, params }) + '\\n'); + if (msg.method === 'initialize') { + reply({ serverInfo: { name: 'stub-muse', version: '0' } }); + } else if (msg.method === 'session/start') { + log('START ' + msg.params.modelId); + reply({ session: { sessionId: 'stub-session-1' } }); + } else if (msg.method === 'session/resume') { + log('RESUME ' + msg.params.sessionId); + reply({ session: { sessionId: msg.params.sessionId } }); + } else if (msg.method === 'turn/start') { + const text = (msg.params.input || []).map(p => p.text || '').join(''); + log('TURN ' + text); + const sid = msg.params.sessionId; + const turnId = msg.params.commandId; + reply({ commandId: msg.params.commandId, turnId, disposition: 'started', startedNewTurn: true, status: 'accepted' }); + notify('item/started', { sessionId: sid, item: { itemId: 'a1', kind: 'agentMessage' } }); + notify('item/delta', { sessionId: sid, itemId: 'a1', delta: 'RE:' + text }); + notify('item/completed', { sessionId: sid, item: { itemId: 'a1', kind: 'agentMessage', messagePhase: 'finalAnswer', text: 'RE:' + text } }); + notify('session/tokenUsage', { sessionId: sid, turnId, usage: { inputTokens: 42, outputTokens: 7 } }); + notify('turn/completed', { sessionId: sid, turnId, terminal: 'completed' }); + } else { + reply({}); + } + } + }); +} else { + let promptFile = null; + for (let i = 0; i < args.length; i++) if (args[i] === '--prompt-file') promptFile = args[i + 1]; + let prompt = ''; + try { prompt = fs.readFileSync(promptFile, 'utf-8'); } catch {} + const line = (o) => console.log(JSON.stringify(o)); + if (prompt.includes('FAILME')) { console.error('stub muse: auth expired'); process.exit(3); } + if (prompt.includes('FAILPARTIAL')) { + line({ payload_type: 'run.output.delta', payload: { text: 'partial answer' } }); + line({ payload_type: 'run.terminal.failed', payload: { terminal: 'failed', text: 'partial answer' } }); + console.error('stub muse: run failed midway'); + process.exit(0); + } + if (prompt.includes('CRASHPARTIAL')) { + line({ payload_type: 'run.output.delta', payload: { text: 'partial answer' } }); + console.error('stub muse: crashed midway'); + process.exit(3); + } + if (prompt.includes('FAILSTATE')) { + line({ payload_type: 'run.terminal.failed', payload: { terminal: 'failed' } }); + process.exit(0); + } + if (prompt.includes('SIGKILLME')) { process.kill(process.pid, 'SIGKILL'); } + if (prompt.includes('DIVERGE')) { + line({ payload_type: 'run.output.delta', payload: { text: 'abc' } }); + line({ payload_type: 'run.terminal.completed', payload: { terminal: 'completed', text: 'xyz-longer' } }); + } else { + line({ payload_type: 'run.output.delta', payload: { text: 'Hel' } }); + line({ payload_type: 'run.usage', payload: { models: [{ input_tokens: 100, output_tokens: 3 }] } }); + line({ payload_type: 'run.terminal.completed', payload: { terminal: 'completed', text: 'Hello world' } }); + } +}`; const freePort = () => new Promise((resolve, reject) => { const srv = net.createServer(); @@ -66,20 +105,24 @@ const freePort = () => new Promise((resolve, reject) => { }); }); -const startProxy = async () => { +const startProxy = async (extraEnv = {}) => { const tmp = fs.mkdtempSync(path.join(os.tmpdir(), 'muse-proxy-test-')); const stubDir = path.join(tmp, 'stub'); fs.mkdirSync(stubDir); fs.mkdirSync(path.join(tmp, 'home')); fs.writeFileSync(path.join(stubDir, 'muse'), STUB, { mode: 0o755 }); const port = await freePort(); + const logFile = path.join(tmp, 'stub.log'); const child = spawn(process.execPath, [path.join(binDir, 'muse-proxy.mjs')], { env: { ...process.env, HOME: path.join(tmp, 'home'), PATH: `${stubDir}:${process.env.PATH}`, MUSE_PROXY_PORT: String(port), - MUSE_PROXY_WORKDIR: path.join(tmp, 'work') + MUSE_PROXY_WORKDIR: path.join(tmp, 'work'), + MUSE_PROXY_WARM: '0', + STUB_LOG: logFile, + ...extraEnv }, stdio: ['ignore', 'pipe', 'pipe'] }); @@ -95,7 +138,7 @@ const startProxy = async () => { if (Date.now() > deadline) throw new Error(`proxy did not start:\n${log}`); await new Promise(r => setTimeout(r, 50)); } - return { port, stop: () => child.kill('SIGKILL') }; + return { port, stop: () => child.kill('SIGKILL'), logFile }; }; const chat = async (port, body) => { @@ -251,3 +294,50 @@ test('divergent terminal text: non-stream uses it, stream keeps the deltas alrea assert.equal(content, 'abc'); assert.equal(finish, 'stop'); }); + +test('warm serve path: second turn sends only the new user message (#480)', async (t) => { + const { port, stop, logFile } = await startProxy({ MUSE_PROXY_WARM: '1' }); + t.after(stop); + + const turn1 = JSON.parse(await chat(port, { + model: 'muse-spark-1.3', + messages: [{ role: 'user', content: 'alpha one' }] + })); + assert.equal(turn1.choices[0].message.content, 'RE:alpha one'); + assert.equal(turn1.usage.prompt_tokens, 42, 'usage must come from the MSP tokenUsage notification'); + + const turn2 = JSON.parse(await chat(port, { + model: 'muse-spark-1.3', + messages: [ + { role: 'user', content: 'alpha one' }, + { role: 'assistant', content: 'RE:alpha one' }, + { role: 'user', content: 'beta two' } + ] + })); + assert.equal(turn2.choices[0].message.content, 'RE:beta two'); + + const log = fs.readFileSync(logFile, 'utf-8').trim().split('\n'); + assert.equal(log.filter(l => l.startsWith('START')).length, 1, 'one MSP session for the whole conversation'); + const turns = log.filter(l => l.startsWith('TURN')); + assert.equal(turns.length, 2); + assert.equal(turns[0], 'TURN alpha one'); + assert.equal(turns[1], 'TURN beta two', 'turn 2 must send only the new message, not the re-flattened history'); +}); + +test('warm serve path: streaming turn streams the MSP deltas (#480)', async (t) => { + const { port, stop } = await startProxy({ MUSE_PROXY_WARM: '1' }); + t.after(stop); + const { content, finish } = streamedContent(await chat(port, { + messages: [{ role: 'user', content: 'alpha one' }], + stream: true + })); + assert.equal(content, 'RE:alpha one'); + assert.equal(finish, 'stop'); +}); + +test('warm serve failure falls back to the exec path (#480)', async (t) => { + const { port, stop } = await startProxy({ MUSE_PROXY_WARM: '1', STUB_SERVE_FAIL: '1' }); + t.after(stop); + const plain = JSON.parse(await chat(port, { messages: [{ role: 'user', content: 'hi' }] })); + assert.equal(plain.choices[0].message.content, 'Hello world', 'exec stub answer proves the fallback ran'); +});