test(pi-claude-bridge): cover protocol and session recovery

pi/extensions/claude-bridge/tests/mcp.test.ts

@@ -0,0 +1,49 @@
+import assert from "node:assert/strict";
+import { test } from "node:test";
+import { startToolServer, TOOL_USE_ID_META } from "../mcp.ts";
+
+test("MCP initialization, notifications, validation, pairing, and tool errors", async () => {
+  const calls: Array<{ name: string; id: string }> = [];
+  let toolName = "read";
+  const server = await startToolServer({
+    listTools: () => [{ name: toolName, description: "Test tool", inputSchema: { type: "object" } }],
+    callTool: async (name, id) => {
+      calls.push({ name, id });
+      if (name === "fail") throw new Error("EXPECTED_MCP_ERROR");
+      return { content: [{ type: "text", text: id }] };
+    },
+  });
+  async function rpc(method: string, params: Record<string, unknown> = {}) {
+    const response = await fetch(server.url, {
+      method: "POST", body: JSON.stringify({ jsonrpc: "2.0", id: 1, method, params }),
+    });
+    assert.equal(response.status, 200);
+    return await response.json() as any;
+  }
+  try {
+    assert.match(server.url, /^http:\/\/127\.0\.0\.1:/);
+    assert.equal((await fetch(server.url)).status, 405);
+    assert.equal((await fetch(server.url, { method: "POST", body: "not-json" })).status, 400);
+    const notification = await fetch(server.url, {
+      method: "POST", body: JSON.stringify({ jsonrpc: "2.0", method: "notifications/initialized" }),
+    });
+    assert.equal(notification.status, 202);
+    assert.equal(await notification.text(), "");
+    assert.equal((await rpc("initialize")).result.protocolVersion, "2025-06-18");
+    assert.equal((await rpc("tools/list")).result.tools[0].name, "read");
+    toolName = "changed";
+    assert.equal((await rpc("tools/list")).result.tools[0].name, "changed");
+    assert.equal((await rpc("server/discover")).error.code, -32601);
+    assert.equal((await rpc("tools/call", {})).error.code, -32602);
+    assert.equal((await rpc("tools/call", { name: "read" })).error.code, -32602);
+    assert.equal(calls.length, 0);
+    const result = await rpc("tools/call", { name: "read", _meta: { [TOOL_USE_ID_META]: "pair-73921" } });
+    assert.equal(result.result.content[0].text, "pair-73921");
+    assert.deepEqual(calls, [{ name: "read", id: "pair-73921" }]);
+    const error = await rpc("tools/call", { name: "fail", _meta: { [TOOL_USE_ID_META]: "error-id" } });
+    assert.equal(error.result.isError, true);
+    assert.equal(error.result.content[0].text, "EXPECTED_MCP_ERROR");
+  } finally {
+    await server.close();
+  }
+});

pi/extensions/claude-bridge/tests/mock-claude.cjs

@@ -0,0 +1,81 @@
+// Offline CLI fixture. The regression runner installs it as `claude` in a temporary PATH.
+const fs = require('node:fs');
+const args = process.argv.slice(2);
+const flag = (name) => args[args.indexOf(name) + 1];
+const model = flag('--model');
+const mcp = args.includes('--mcp-config') ? JSON.parse(fs.readFileSync(flag('--mcp-config'), 'utf8')).mcpServers.pi.url : null;
+let count = 0, memory = '', buffer = '';
+const emit = (message) => process.stdout.write(JSON.stringify(message) + '\n');
+const event = (event) => emit({ type: 'stream_event', event });
+const delay = (ms) => new Promise((resolve) => setTimeout(resolve, ms));
+async function answer(text) {
+  const id = `text-${++count}`;
+  event({ type: 'message_start', message: { id, usage: { input_tokens: 10, output_tokens: 0 } } });
+  event({ type: 'content_block_start', index: 0, content_block: { type: 'text' } });
+  event({ type: 'content_block_delta', index: 0, delta: { type: 'text_delta', text } });
+  event({ type: 'content_block_stop', index: 0 });
+  event({ type: 'message_delta', delta: { stop_reason: 'end_turn' }, usage: { output_tokens: 5 } });
+  event({ type: 'message_stop' });
+  emit({ type: 'assistant', message: { id, content: [{ type: 'text', text }], usage: { input_tokens: 10, output_tokens: 5 } } });
+  // Expose the original bug: a new Pi prompt must not start before this result.
+  await delay(150);
+  emit({ type: 'result', subtype: 'success', result: text, usage: { input_tokens: 100, output_tokens: 50 } });
+}
+async function tool(name, input, nonStreaming = false, inlineInput = false) {
+  const id = `tool-${++count}`;
+  const block = { type: 'tool_use', id, name: `mcp__pi__${name}`, input };
+  const message = { id: `assistant-${id}`, content: [block], usage: { input_tokens: 10, output_tokens: 5 } };
+  if (nonStreaming) emit({ type: 'assistant', message });
+  else {
+    event({ type: 'message_start', message });
+    event({ type: 'content_block_start', index: 0, content_block: block });
+    if (!inlineInput) event({ type: 'content_block_delta', index: 0, delta: { type: 'input_json_delta', partial_json: JSON.stringify(input) } });
+    event({ type: 'content_block_stop', index: 0 });
+    event({ type: 'message_delta', delta: { stop_reason: 'tool_use' }, usage: { output_tokens: 5 } });
+    event({ type: 'message_stop' });
+    // The full assistant record can arrive after Pi opens its tool-result stream.
+    await delay(80);
+    emit({ type: 'assistant', message });
+  }
+  const response = await fetch(mcp, { method: 'POST', body: JSON.stringify({ jsonrpc: '2.0', id, method: 'tools/call', params: { name, arguments: input, _meta: { 'claudecode/toolUseId': id } } }) });
+  const result = (await response.json()).result;
+  await answer(result?.isError ? 'TOOL_ERROR_HANDLED' : 'TOOL_OK');
+}
+async function handle(record) {
+  const content = record.message.content;
+  const text = content.filter((block) => block.type === 'text').map((block) => block.text).join('\n');
+  const prompt = content.at(-1)?.text || '';
+  if (text.includes('MEMORY_73921')) memory = 'MEMORY_73921';
+  fs.appendFileSync(process.env.BRIDGE_TEST_LOG, JSON.stringify({ pid: process.pid, model, effort: args.includes('--effort') ? flag('--effort') : null, text, apiKey: Boolean(process.env.ANTHROPIC_API_KEY) }) + '\n');
+  if (prompt === 'MOCK_ABORT') {
+    event({ type: 'message_start', message: { id: 'abort', usage: {} } });
+    event({ type: 'content_block_start', index: 0, content_block: { type: 'text' } });
+    event({ type: 'content_block_delta', index: 0, delta: { type: 'text_delta', text: 'ABORT_STARTED' } });
+    return;
+  }
+  if (prompt === 'MOCK_EXIT') return process.exit(7);
+  if (prompt === 'MOCK_INLINE_TOOL') return tool('read', { path: process.env.BRIDGE_TEST_FILE }, false, true);
+  if (prompt === 'MOCK_TOOL') return tool('read', { path: process.env.BRIDGE_TEST_FILE });
+  if (prompt === 'MOCK_FALLBACK_TOOL') return tool('read', { path: process.env.BRIDGE_TEST_FILE }, true);
+  if (prompt === 'MOCK_TOOL_ERROR') return tool('read', { path: process.env.BRIDGE_TEST_FILE + '.missing' });
+  if (prompt === 'MOCK_SLOW_TOOL') return tool('bridge_wait', {});
+  if (prompt === 'MOCK_MEMORY') return answer(memory || 'UNKNOWN');
+  if (prompt === 'MOCK_MODEL') return answer(model);
+  if (prompt === 'MOCK_CLI_ERROR') {
+    emit({ type: 'result', is_error: true, errors: ['MOCK_PROVIDER_FAILURE'] });
+    return;
+  }
+  return answer(prompt);
+}
+let queue = Promise.resolve();
+process.stdin.on('data', (chunk) => {
+  buffer += chunk;
+  let newline;
+  while ((newline = buffer.indexOf('\n')) !== -1) {
+    const line = buffer.slice(0, newline); buffer = buffer.slice(newline + 1);
+    if (line) queue = queue.then(() => handle(JSON.parse(line))).catch((error) => { console.error(error); process.exit(1); });
+  }
+});
+process.stdin.on('end', () => { queue.finally(() => process.exit(0)); });
+// Exercise old-process cleanup after the replacement request has already started.
+process.on('SIGTERM', () => setTimeout(() => process.exit(0), 250));

pi/extensions/claude-bridge/tests/provider.test.cjs

@@ -0,0 +1,147 @@
+const assert = require('node:assert/strict');
+const { spawn } = require('node:child_process');
+const fs = require('node:fs');
+const os = require('node:os');
+const path = require('node:path');
+const { test } = require('node:test');
+
+// Uses the installed Pi loader and runtime. No npm installation or API request is needed.
+test('provider turn boundaries, restarts, fallback tools, errors, and steering', { timeout: 90000 }, async () => {
+  const dir = fs.mkdtempSync(path.join(os.tmpdir(), 'pi-bridge-regression-'));
+  const cli = path.join(dir, 'claude');
+  const log = path.join(dir, 'cli.jsonl');
+  const fixture = path.join(dir, 'fixture.txt');
+  const helper = path.join(dir, 'helper.ts');
+  fs.writeFileSync(cli, `#!/bin/sh\nexec ${JSON.stringify(process.execPath)} ${JSON.stringify(path.join(__dirname, 'mock-claude.cjs'))} "$@"\n`, { mode: 0o700 });
+  fs.writeFileSync(fixture, 'FIXTURE_CONTENT');
+  fs.writeFileSync(path.join(dir, 'settings.json'), JSON.stringify({ compaction: { enabled: true, keepRecentTokens: 100, reserveTokens: 1000 } }));
+  fs.writeFileSync(helper, `export default function(pi) {
+    let suffix = '';
+    pi.on('before_agent_start', event => suffix ? { systemPrompt: event.systemPrompt + suffix } : undefined);
+    pi.registerCommand('bridge_prompt', { description: 'Change test prompt', handler: async () => { suffix = '\\nTEST_PROMPT_CHANGED'; } });
+    pi.registerCommand('bridge_tools', { description: 'Change test tools', handler: async () => { pi.setActiveTools(['read']); } });
+    pi.registerTool({ name: 'bridge_wait', label: 'Wait', description: 'Test wait', parameters: { type: 'object', properties: {} },
+      execute: async () => { await new Promise(r => setTimeout(r, 400)); return { content: [{ type: 'text', text: 'WAIT_OK' }], details: undefined }; } });
+  }`);
+  const child = spawn('pi', [
+    '--offline', '--no-extensions', '--no-mcp', '--no-skills', '--no-prompt-templates', '--no-context-files', '--no-session',
+    '-e', path.resolve(__dirname, '../index.ts'), '-e', helper,
+    '--model', 'claude-bridge/claude-haiku-5-5', '--thinking', 'off', '--tools', 'read,bridge_wait', '--mode', 'rpc',
+  ], { env: { ...process.env, PATH: `${dir}:${process.env.PATH}`, ANTHROPIC_API_KEY: 'must-not-reach-child', BRIDGE_TEST_LOG: log, BRIDGE_TEST_FILE: fixture, PI_CODING_AGENT_DIR: dir } });
+  let buffer = '', stderr = '', seq = 0;
+  const events = [], waiters = new Set();
+  child.stderr.on('data', (chunk) => { stderr += chunk; });
+  child.stdout.on('data', (chunk) => {
+    buffer += chunk;
+    let newline;
+    while ((newline = buffer.indexOf('\n')) !== -1) {
+      const line = buffer.slice(0, newline); buffer = buffer.slice(newline + 1);
+      let record;
+      try { record = JSON.parse(line); } catch { continue; }
+      events.push(record);
+      for (const waiter of [...waiters]) if (waiter.predicate(record)) {
+        waiters.delete(waiter); clearTimeout(waiter.timer); waiter.resolve(record);
+      }
+    }
+  });
+  function wait(predicate) {
+    return new Promise((resolve, reject) => {
+      const waiter = { predicate, resolve, timer: setTimeout(() => {
+        waiters.delete(waiter); reject(new Error(`RPC timeout. stderr: ${stderr}`));
+      }, 10000) };
+      waiters.add(waiter);
+    });
+  }
+  async function command(type, fields = {}) {
+    const id = String(++seq);
+    const response = wait((record) => record.type === 'response' && record.id === id);
+    child.stdin.write(JSON.stringify({ id, type, ...fields }) + '\n');
+    const record = await response;
+    assert.equal(record.success, true, JSON.stringify(record));
+    return record;
+  }
+  function answersSince(start) {
+    return events.slice(start).filter((record) => record.type === 'message_end' && record.message?.role === 'assistant');
+  }
+  async function prompt(message) {
+    const start = events.length;
+    const settled = wait((record) => record.type === 'agent_settled');
+    await command('prompt', { message }); await settled;
+    const answers = answersSince(start);
+    const last = answers.at(-1)?.message;
+    return { start, answers, last, text: last?.content.filter((block) => block.type === 'text').map((block) => block.text).join('') };
+  }
+  const cliRecords = () => fs.readFileSync(log, 'utf8').trim().split('\n').map(JSON.parse);
+  try {
+    await command('set_auto_retry', { enabled: false });
+    for (let i = 0; i < 5; i++) {
+      const result = await prompt(`RAPID_${i}`);
+      assert.equal(result.text, `RAPID_${i}`);
+      assert.equal(result.answers.length, 1);
+      assert.equal(result.last.usage.input, 10, 'Do not overwrite message usage with aggregate CLI result usage');
+    }
+    assert.equal(new Set(cliRecords().map((record) => record.pid)).size, 1, 'Reuse the process on normal turns');
+    assert.ok(cliRecords().every((record) => record.apiKey === false));
+    assert.equal((await prompt('Remember MEMORY_73921')).last.stopReason, 'stop');
+    assert.equal((await prompt('MOCK_MEMORY')).text, 'MEMORY_73921');
+    for (const name of ['MOCK_TOOL', 'MOCK_INLINE_TOOL', 'MOCK_FALLBACK_TOOL', 'MOCK_TOOL_ERROR']) {
+      const result = await prompt(name);
+      assert.equal(result.text, name === 'MOCK_TOOL_ERROR' ? 'TOOL_ERROR_HANDLED' : 'TOOL_OK');
+      const tools = events.slice(result.start).filter((record) => record.type === 'tool_execution_end');
+      assert.equal(tools.length, 1, 'Do not execute the late full assistant record twice');
+      assert.equal(tools[0].isError, name === 'MOCK_TOOL_ERROR');
+    }
+    const error = await prompt('MOCK_CLI_ERROR');
+    assert.equal(error.last.stopReason, 'error');
+    assert.match(error.last.errorMessage, /MOCK_PROVIDER_FAILURE/);
+    assert.equal((await prompt('AFTER_ERROR_OK')).text, 'AFTER_ERROR_OK');
+    const exited = await prompt('MOCK_EXIT');
+    assert.equal(exited.last.stopReason, 'error');
+    assert.match(exited.last.errorMessage, /exited with code 7/);
+    assert.equal((await prompt('MOCK_MEMORY')).text, 'MEMORY_73921', 'Restore after an unexpected CLI exit');
+    await command('set_thinking_level', { level: 'high' });
+    assert.equal((await prompt('MOCK_MEMORY')).text, 'MEMORY_73921');
+    assert.equal(cliRecords().at(-1).effort, 'high');
+    await command('set_thinking_level', { level: 'off' });
+    const started = wait((record) => record.type === 'message_update' && record.assistantMessageEvent?.type === 'text_delta');
+    await command('prompt', { message: 'MOCK_ABORT' }); await started;
+    await command('abort');
+    assert.equal((await prompt('MOCK_MEMORY')).text, 'MEMORY_73921', 'Restore memory after abort');
+    assert.match(cliRecords().at(-1).text, /restore the active Pi conversation/);
+    assert.equal((await prompt('MOCK_TOOL')).text, 'TOOL_OK', 'Old-process cleanup must not clear new pending calls');
+    const toolStart = wait((record) => record.type === 'tool_execution_start' && record.toolName === 'bridge_wait');
+    const settled = wait((record) => record.type === 'agent_settled');
+    const start = events.length;
+    await command('prompt', { message: 'MOCK_SLOW_TOOL' }); await toolStart;
+    await command('steer', { message: 'STEER_OK' }); await settled;
+    const steered = answersSince(start).at(-1).message.content.filter((block) => block.type === 'text').map((block) => block.text).join('');
+    assert.equal(steered, 'STEER_OK');
+    const beforePromptChange = cliRecords().at(-1).pid;
+    await command('prompt', { message: '/bridge_prompt' });
+    assert.equal((await prompt('MOCK_MEMORY')).text, 'MEMORY_73921');
+    assert.notEqual(cliRecords().at(-1).pid, beforePromptChange, 'Changed prompt restarts the CLI');
+    const beforeToolChange = cliRecords().at(-1).pid;
+    await command('prompt', { message: '/bridge_tools' });
+    assert.equal((await prompt('MOCK_MEMORY')).text, 'MEMORY_73921');
+    assert.notEqual(cliRecords().at(-1).pid, beforeToolChange, 'Changed tools restart the CLI');
+    await command('set_model', { provider: 'claude-bridge', modelId: 'claude-haiku-4-5' });
+    assert.equal((await prompt('MOCK_MODEL')).text, 'claude-haiku-4-5');
+    await prompt('Compaction context: ' + 'Previous test data. '.repeat(500));
+    await command('compact', { customInstructions: 'Preserve MEMORY_73921.' });
+    assert.equal((await prompt('MOCK_MEMORY')).text, 'MEMORY_73921', 'Restore the compacted active transcript');
+    assert.match(cliRecords().at(-1).text, /restore the active Pi conversation/);
+    const forks = await command('get_fork_messages');
+    await command('fork', { entryId: forks.data.messages.find((message) => message.text === 'RAPID_0').entryId });
+    assert.equal((await prompt('MOCK_MEMORY')).text, 'UNKNOWN', 'An earlier branch must not retain abandoned history');
+    await command('new_session');
+    assert.equal((await prompt('MOCK_MEMORY')).text, 'UNKNOWN');
+    assert.ok(!cliRecords().at(-1).text.includes('MEMORY_73921'));
+  } finally {
+    for (const waiter of waiters) clearTimeout(waiter.timer);
+    const closed = new Promise((resolve) => child.once('close', resolve));
+    child.stdin.end();
+    const timer = setTimeout(() => child.kill('SIGKILL'), 3000);
+    await closed; clearTimeout(timer);
+    fs.rmSync(dir, { recursive: true, force: true });
+  }
+});

pi/extensions/claude-bridge/tests/transcript.test.ts

@@ -0,0 +1,56 @@
+import assert from "node:assert/strict";
+import { test } from "node:test";
+import { extendsTranscript, restoredContent, transcriptKeys, type TranscriptMessage } from "../transcript.ts";
+
+const user = { role: "user", content: "Remember MEMORY_73921" };
+const assistant = { role: "assistant", content: [{ type: "text", text: "ACK", index: 0 }], timestamp: 1 };
+
+test("fingerprints ignore runtime fields but detect branch and content changes", () => {
+  const previous = transcriptKeys([user, assistant]);
+  const current = transcriptKeys([
+    { role: "system", content: "Changed prompt" }, user,
+    { ...assistant, timestamp: 2, usage: { input: 10 }, content: [{ type: "text", text: "ACK" }] },
+    { role: "user", content: "Next" },
+  ]);
+  assert.ok(extendsTranscript(previous, current));
+  assert.ok(!extendsTranscript(previous, transcriptKeys([user])));
+  assert.ok(!extendsTranscript(previous, transcriptKeys([user, { ...assistant, content: "Different" }])));
+});
+
+test("tool argument key order does not force a restart", () => {
+  const message = (args: unknown): TranscriptMessage => ({ role: "assistant", content: [{ type: "toolCall", id: "id", name: "read", arguments: args }] });
+  assert.deepEqual(transcriptKeys([message({ a: 1, b: 2 })]), transcriptKeys([message({ b: 2, a: 1 })]));
+});
+
+test("restoration includes prior messages, tool results, and current input without native calls", () => {
+  const blocks = restoredContent([
+    { role: "system", content: "SYSTEM_PROMPT_NOT_REPLAYED" }, user, assistant,
+    { role: "toolResult", toolCallId: "old-id", toolName: "read", content: "ERROR_73921", isError: true },
+    { role: "user", content: "Current request" },
+  ]);
+  const text = blocks.filter((block) => block.type === "text").map((block) => block.text).join("\n");
+  assert.match(text, /MEMORY_73921/);
+  assert.match(text, /ERROR_73921/);
+  assert.match(text, /"isError":true/);
+  assert.ok(!text.includes("SYSTEM_PROMPT_NOT_REPLAYED"));
+  assert.equal(blocks.at(-1)?.type, "text");
+  assert.equal((blocks.at(-1) as { text: string }).text, "Current request");
+  assert.ok(blocks.every((block) => block.type === "text" || block.type === "image"));
+});
+
+test("historical images stay image blocks; thinking signatures are not replayed", () => {
+  const blocks = restoredContent([
+    { role: "user", content: [{ type: "image", data: "aW1hZ2U=", mimeType: "image/png" }] },
+    { role: "assistant", content: [{ type: "thinking", thinking: "private", thinkingSignature: "secret-signature" }] },
+    { role: "user", content: "Next" },
+  ]);
+  assert.equal(blocks.filter((block) => block.type === "image").length, 1);
+  const text = JSON.stringify(blocks);
+  assert.ok(!text.includes("secret-signature"));
+  assert.ok(!text.includes("private"));
+});
+
+test("tool-result tail gets a continuation request instead of replaying tool calls", () => {
+  const blocks = restoredContent([{ role: "toolResult", toolCallId: "id", content: "DONE" }]);
+  assert.match((blocks.at(-1) as { text: string }).text, /Continue the latest user request/);
+});