feat(muse-proxy): warm sessions on a persistent muse serve (MSP) host (#480)

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');
+});