diff --git a/scripts/bug-bash-logging.sh b/scripts/bug-bash-logging.sh new file mode 100755 index 0000000..2cfe4d7 --- /dev/null +++ b/scripts/bug-bash-logging.sh @@ -0,0 +1,217 @@ +#!/usr/bin/env bash +# +# Bug-bash scenario for the event-logging system. +# +# Walks through a realistic project lifecycle with STITCH_MCP_LOG=1, then +# inspects .stitch-mcp/log/ to confirm events were captured correctly. +# +# Requirements: +# - bun (project's dev runtime) +# - jq (brew install jq) +# - Authenticated Stitch session (env vars or gcloud ADC already set up) +# +# Usage: +# ./scripts/bug-bash-logging.sh +# ./scripts/bug-bash-logging.sh --keep # don't wipe .stitch-mcp/log/ first +# +set -uo pipefail + +KEEP=0 +[[ "${1:-}" == "--keep" ]] && KEEP=1 + +LOG_DIR=".stitch-mcp/log" +EVENTS="$LOG_DIR/events.jsonl" +BLOBS="$LOG_DIR/blobs" + +export STITCH_MCP_LOG=1 + +# --- helpers --------------------------------------------------------------- + +# All UI output goes to stderr so $(tool ...) captures clean JSON on stdout. +bold() { printf '\033[1m%s\033[0m\n' "$*" >&2; } +green() { printf '\033[32m%s\033[0m\n' "$*" >&2; } +red() { printf '\033[31m%s\033[0m\n' "$*" >&2; } +gray() { printf '\033[90m%s\033[0m\n' "$*" >&2; } + +step() { + echo >&2 + bold "▶ $*" +} + +# Invoke a tool via the CLI. Echoes raw JSON on stdout, status on stderr. +# Usage: tool +tool() { + local name=$1 + local args=$2 + gray " args: $args" + local start=$(date +%s) + local out + if ! out=$(bun run src/cli.ts tool "$name" -d "$args" -o json 2>/tmp/bug-bash.err); then + red " ✖ tool $name failed (exit $?)" + cat /tmp/bug-bash.err >&2 + echo "$out" + return 1 + fi + local end=$(date +%s) + gray " duration: $((end - start))s" + echo "$out" +} + +require() { + command -v "$1" >/dev/null 2>&1 || { red "✖ missing dependency: $1"; exit 1; } +} + +# --- preflight ------------------------------------------------------------- + +require bun +require jq + +if [[ $KEEP -eq 0 ]]; then + bold "Wiping $LOG_DIR (use --keep to preserve)" + rm -rf "$LOG_DIR" +fi + +# --- step 1: create project ------------------------------------------------ + +step "create_project" +PROJECT_RESPONSE=$(tool create_project '{"title":"bug-bash logging '"$(date +%s)"'"}') +PROJECT_ID=$(echo "$PROJECT_RESPONSE" | jq -r '.. | objects | .name? // empty | select(test("^projects/"))' | head -1 | sed 's|^projects/||') +if [[ -z "$PROJECT_ID" ]]; then + red "✖ failed to extract project ID" + echo "$PROJECT_RESPONSE" | jq '.' >&2 + exit 1 +fi +green " ✓ project_id: $PROJECT_ID" + +# --- step 2: generate first screen ----------------------------------------- + +step "generate_screen_from_text (cold — no previous screens)" +echo " this can take 1–3 min..." +SCREEN1_RESPONSE=$(tool generate_screen_from_text "$(jq -n --arg p "$PROJECT_ID" '{ + projectId: $p, + prompt: "A clean mobile login screen with email and password inputs and a primary CTA button" +}')") +SCREEN1_ID=$(echo "$SCREEN1_RESPONSE" | jq -r '.. | objects | .id? // empty | select(test("^[a-f0-9]+$"))' | head -1) +if [[ -z "$SCREEN1_ID" ]]; then + red "✖ failed to extract screen ID" + echo "$SCREEN1_RESPONSE" | jq '.structuredContent | keys' >&2 || true + exit 1 +fi +green " ✓ screen_id: $SCREEN1_ID" + +# --- step 3: edit that screen ---------------------------------------------- + +step "edit_screens (modify the login screen)" +echo " this can take 1–3 min..." +EDIT_RESPONSE=$(tool edit_screens "$(jq -n --arg p "$PROJECT_ID" --arg s "$SCREEN1_ID" '{ + projectId: $p, + selectedScreenIds: [$s], + prompt: "Add a Sign in with Google button above the email input" +}')") +EDITED_ID=$(echo "$EDIT_RESPONSE" | jq -r '.. | objects | .id? // empty | select(test("^[a-f0-9]+$"))' | head -1) +green " ✓ edited screen id: ${EDITED_ID:-}" + +# --- step 4: generate variants --------------------------------------------- + +step "generate_variants (3 variants of the login screen)" +echo " this can take 2–5 min..." +VARIANTS_RESPONSE=$(tool generate_variants "$(jq -n --arg p "$PROJECT_ID" --arg s "$SCREEN1_ID" '{ + projectId: $p, + selectedScreenIds: [$s], + prompt: "Try different visual styles — minimal, playful, and corporate", + variantOptions: { variantCount: 3, creativeRange: "EXPLORE" } +}')") +VARIANT_COUNT=$(echo "$VARIANTS_RESPONSE" | jq '[.. | objects | .id? // empty | select(test("^[a-f0-9]+$"))] | unique | length') +VARIANT_API_ERR=$(echo "$VARIANTS_RESPONSE" | jq -r '.error // empty') +if [[ -n "$VARIANT_API_ERR" ]]; then + red " ⚠ generate_variants API error: $VARIANT_API_ERR (logging system should still capture as call.failed)" +else + green " ✓ produced screens: $VARIANT_COUNT" +fi + +# --- step 5: read-only call (sanity-check the read path) ------------------- + +step "list_screens (read path — should log without asset capture)" +LIST_RESPONSE=$(tool list_screens "$(jq -n --arg p "$PROJECT_ID" '{projectId: $p}')") +LIST_COUNT=$(echo "$LIST_RESPONSE" | jq '[.. | objects | .name? // empty | select(test("^screens/"))] | unique | length') +green " ✓ list_screens returned $LIST_COUNT screens" + +# --- step 6: an unclassified tool (should NOT crash the user call) --------- + +step "list_design_systems (unclassified tool — capture should fail silently)" +DS_RESPONSE=$(tool list_design_systems '{}') +if echo "$DS_RESPONSE" | jq -e '.isError == true' >/dev/null 2>&1; then + red " ⚠ user-visible error returned — capture leaked into the response" +else + green " ✓ user-visible response is intact (capture failure was swallowed)" +fi + +# --- inspection ------------------------------------------------------------ + +echo +bold "═══════════════════════════════════════════════════════════════" +bold " Inspection" +bold "═══════════════════════════════════════════════════════════════" + +if [[ ! -f "$EVENTS" ]]; then + red "✖ events.jsonl was not created — logging may not be wired through" + exit 1 +fi + +EVENT_COUNT=$(wc -l <"$EVENTS" | tr -d ' ') +BLOB_COUNT=$(find "$BLOBS" -type f 2>/dev/null | wc -l | tr -d ' ') +BLOB_SIZE=$(du -sh "$BLOBS" 2>/dev/null | awk '{print $1}') + +echo +bold "Counts" +echo " events: $EVENT_COUNT lines" +echo " blobs: $BLOB_COUNT files ($BLOB_SIZE)" + +echo +bold "Event types (call.requested vs call.completed vs call.failed)" +jq -r '.type' "$EVENTS" | sort | uniq -c + +echo +bold "Tools called (from call.requested events)" +jq -r 'select(.type == "call.requested") | .payload.tool' "$EVENTS" | sort | uniq -c + +echo +bold "Trace pairs (each tool call should appear twice — requested + completed/failed)" +jq -r '.trace_id' "$EVENTS" | sort | uniq -c | awk ' + $1 == 2 { paired++ } + $1 == 1 { unpaired++; print " ⚠ unpaired trace: " $2 } + $1 > 2 { extra++; print " ⚠ trace with " $1 " events: " $2 } + END { + print "" + print " paired traces: " (paired+0) + print " unpaired traces: " (unpaired+0) + print " >2-event traces: " (extra+0) + } +' + +echo +bold "Generative-call asset capture" +jq -c 'select(.type == "call.completed" and .payload.kind == "generative")' "$EVENTS" | while read -r line; do + tool=$(echo "$line" | jq -r '.payload.tool') + produced=$(echo "$line" | jq '.payload.produced_screens | length') + has_html=$(echo "$line" | jq '[.payload.produced_screens[]?.html_blob // empty] | length') + has_screenshot=$(echo "$line" | jq '[.payload.produced_screens[]?.screenshot_blob // empty] | length') + echo " $tool — $produced screens, $has_html html blobs, $has_screenshot screenshots" +done + +echo +bold "Failures (call.failed)" +FAIL_COUNT=$(jq -c 'select(.type == "call.failed")' "$EVENTS" | wc -l | tr -d ' ') +if [[ "$FAIL_COUNT" == "0" ]]; then + green " ✓ none" +else + red " ✖ $FAIL_COUNT failures:" + jq -c 'select(.type == "call.failed") | {tool: .payload.tool, is_error: .payload.is_error, err: .payload.error_text}' "$EVENTS" +fi + +echo +bold "Summary" +echo " project_id: $PROJECT_ID" +echo " log dir: $LOG_DIR" +echo " re-inspect: jq -c . $EVENTS" +echo " blob dir: $BLOBS" diff --git a/scripts/log-probe-2.ts b/scripts/log-probe-2.ts new file mode 100644 index 0000000..595988e --- /dev/null +++ b/scripts/log-probe-2.ts @@ -0,0 +1,70 @@ +/** + * Follow-up probe: variants + edit + a successful get_screen. + * Reuses the project + screen created in 3-generate-screen.json. + */ +import 'dotenv/config'; +import { Client } from '@modelcontextprotocol/sdk/client/index.js'; +import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js'; +import { CallToolResultSchema } from '@modelcontextprotocol/sdk/types.js'; +import { mkdir, readFile, writeFile } from 'node:fs/promises'; +import { join } from 'node:path'; + +const OUT = '.stitch-mcp/log-probe'; +const URL2 = process.env.STITCH_MCP_URL ?? 'https://stitch.googleapis.com/mcp'; + +const prior = JSON.parse(await readFile(`${OUT}/3-generate-screen.json`, 'utf-8')); +const sc = prior.result.structuredContent; +const projectId: string = sc.projectId; +// Find the design component, take first screen +let screenId: string | undefined; +for (const c of sc.outputComponents) { + if (c.design?.screens?.length) { screenId = c.design.screens[0].id; break; } +} +if (!projectId || !screenId) { + throw new Error(`missing prior ids; projectId=${projectId} screenId=${screenId}`); +} +console.error(`reusing projectId=${projectId} screenId=${screenId}`); + +await mkdir(OUT, { recursive: true }); + +const transport = new StreamableHTTPClientTransport(new URL(URL2), { + requestInit: { headers: { 'X-Goog-Api-Key': process.env.STITCH_API_KEY!, Accept: 'application/json, text/event-stream' } }, +}); +const client = new Client({ name: 'log-probe-2', version: '0.0.0' }, { capabilities: {} }); +await client.connect(transport); + +const progressEvents: any[] = []; +async function call(label: string, name: string, args: any) { + console.error(`\n[${label}] ${name}`); + const t0 = Date.now(); + const myProgress: any[] = []; + try { + const result = await client.callTool({ name, arguments: args }, CallToolResultSchema, { + timeout: 600_000, + onprogress: (p: any) => { const ev = { ts: Date.now()-t0, ...p }; myProgress.push(ev); progressEvents.push({ label, ...ev }); console.error(` progress`, p); }, + }); + const ms = Date.now() - t0; + console.error(` ok in ${ms}ms`); + await writeFile(join(OUT, `${label}.json`), JSON.stringify({ duration_ms: ms, progress: myProgress, result }, null, 2)); + } catch (e: any) { + const ms = Date.now() - t0; + console.error(` THREW in ${ms}ms:`, e?.message); + await writeFile(join(OUT, `${label}.json`), JSON.stringify({ duration_ms: ms, progress: myProgress, threw: { message: e?.message, code: e?.code } }, null, 2)); + } +} + +await call('4-generate-variants', 'generate_variants', { + projectId, + selectedScreenIds: [screenId], + prompt: 'Vary the color scheme and the typographic layout', + variantOptions: { variantCount: 2, creativeRange: 'EXPLORE', aspects: ['COLOR_SCHEME', 'LAYOUT'] }, +}); + +await call('5-edit-screens', 'edit_screens', { + projectId, selectedScreenIds: [screenId], prompt: 'Make the headline larger and serif', +}); + +await call('6-get-screen', 'get_screen', { projectId, screenId }); + +console.error(`\nProgress events captured: ${progressEvents.length}`); +await client.close(); diff --git a/scripts/log-probe.ts b/scripts/log-probe.ts new file mode 100644 index 0000000..a1f6bc6 --- /dev/null +++ b/scripts/log-probe.ts @@ -0,0 +1,141 @@ +/** + * One-off probe to capture raw wire shapes for the logging-system design. + * Saves verbatim envelopes to .stitch-mcp/log-probe/. + * + * Run: STITCH_API_KEY=... bun scripts/log-probe.ts + */ +import 'dotenv/config'; +import { Client } from '@modelcontextprotocol/sdk/client/index.js'; +import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js'; +import { CallToolResultSchema } from '@modelcontextprotocol/sdk/types.js'; +import { mkdir, writeFile } from 'node:fs/promises'; +import { join } from 'node:path'; + +const OUT = '.stitch-mcp/log-probe'; +const STITCH_URL = process.env.STITCH_MCP_URL ?? 'https://stitch.googleapis.com/mcp'; + +if (!process.env.STITCH_API_KEY) { + console.error('STITCH_API_KEY required'); + process.exit(1); +} + +await mkdir(OUT, { recursive: true }); + +const transport = new StreamableHTTPClientTransport(new URL(STITCH_URL), { + requestInit: { + headers: { + 'X-Goog-Api-Key': process.env.STITCH_API_KEY!, + Accept: 'application/json, text/event-stream', + }, + }, +}); + +const client = new Client({ name: 'log-probe', version: '0.0.0' }, { capabilities: {} }); +await client.connect(transport); + +const progressEvents: any[] = []; + +async function dump(label: string, data: unknown) { + const file = join(OUT, `${label}.json`); + await writeFile(file, JSON.stringify(data, null, 2)); + console.error(` -> ${file} (${JSON.stringify(data).length} bytes)`); +} + +async function call(label: string, name: string, args: Record, opts: { progress?: boolean } = {}) { + console.error(`\n[${label}] ${name}`); + const t0 = Date.now(); + const myProgress: any[] = []; + try { + const result = await client.callTool( + { name, arguments: args }, + CallToolResultSchema, + { + timeout: 600_000, + onprogress: opts.progress + ? (p: any) => { + const ev = { ts: Date.now() - t0, ...p }; + myProgress.push(ev); + progressEvents.push({ label, ...ev }); + console.error(` progress @${ev.ts}ms`, p); + } + : undefined, + }, + ); + const ms = Date.now() - t0; + console.error(` ok in ${ms}ms`); + await dump(label, { duration_ms: ms, progress: myProgress, result }); + return result as any; + } catch (e: any) { + const ms = Date.now() - t0; + console.error(` THREW after ${ms}ms:`, e?.message ?? e); + await dump(label, { duration_ms: ms, progress: myProgress, threw: { message: e?.message, code: e?.code, data: e?.data } }); + return null; + } +} + +// 1) cheap baseline +await call('1-list-projects', 'list_projects', {}); + +// 2) create throwaway project +const created = await call('2-create-project', 'create_project', { title: `log-probe-${Date.now()}` }); +const createdSC: any = created?.structuredContent ?? (() => { + try { return JSON.parse(created?.content?.[0]?.text ?? '{}'); } catch { return {}; } +})(); +const projectName: string | undefined = createdSC?.name; +const projectId = projectName?.replace(/^projects\//, ''); +console.error(` projectId = ${projectId}`); + +if (projectId) { + // 3) generate one screen, request progress + const gen = await call( + '3-generate-screen', + 'generate_screen_from_text', + { projectId, prompt: 'A minimal landing page for a typography-focused blog. One headline, one paragraph, one CTA.' }, + { progress: true }, + ); + const genSC: any = gen?.structuredContent ?? (() => { + try { return JSON.parse(gen?.content?.[0]?.text ?? '{}'); } catch { return {}; } + })(); + const firstScreenId = genSC?.outputComponents?.[0]?.design?.screens?.[0]?.id; + console.error(` firstScreenId = ${firstScreenId}`); + + // 4) variants on that screen, capture aspect attribution + if (firstScreenId) { + await call( + '4-generate-variants', + 'generate_variants', + { + projectId, + selectedScreenIds: [firstScreenId], + prompt: 'Vary the color scheme and the layout', + variantOptions: { variantCount: 2, creativeRange: 'EXPLORE', aspects: ['COLOR_SCHEME', 'LAYOUT'] }, + }, + { progress: true }, + ); + + // 5) edit screens + await call( + '5-edit-screens', + 'edit_screens', + { projectId, selectedScreenIds: [firstScreenId], prompt: 'Make the headline larger and serif' }, + { progress: true }, + ); + + // 6) get_screen, look at envelope of read-only call + await call('6-get-screen', 'get_screen', { projectId, screenId: firstScreenId }); + } + + // 7) post-state of project (designMd may have grown) + await call('7-get-project-after', 'get_project', { name: `projects/${projectId}` }); +} + +// 8) deliberate failure: bogus screen id +await call('8-fail-bogus-screen', 'get_screen', { projectId: '0', screenId: 'doesnotexist' }); + +// 9) deliberate failure: missing required field (server-side, since SDK won't block raw call) +await call('9-fail-missing-required', 'edit_screens', { projectId: projectId ?? '0' } as any); + +await writeFile(join(OUT, 'progress-summary.json'), JSON.stringify(progressEvents, null, 2)); +console.error(`\nProgress events captured: ${progressEvents.length}`); + +await client.close(); diff --git a/scripts/log-validate.ts b/scripts/log-validate.ts new file mode 100644 index 0000000..3051f40 --- /dev/null +++ b/scripts/log-validate.ts @@ -0,0 +1,582 @@ +/** + * Validation prototype for the logging system. + * + * Replays every probe artifact (real Stitch responses) through the proposed + * file structure to verify mechanical/structural correctness, then runs ONE + * live get_screen call to exercise the end-to-end live path. Prints stats and + * sample queries so we can assess the structure on real data. + * + * Output goes to .stitch-mcp/log/ (events.jsonl + blobs/ + index/). + * Reset between runs by deleting that directory. + */ +import 'dotenv/config'; +import { createHash } from 'node:crypto'; +import { mkdir, readFile, readdir, writeFile, appendFile, rename, stat } from 'node:fs/promises'; +import { join, dirname } from 'node:path'; +import { z } from 'zod'; +import { Client } from '@modelcontextprotocol/sdk/client/index.js'; +import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js'; +import { CallToolResultSchema } from '@modelcontextprotocol/sdk/types.js'; + +const ROOT = '.stitch-mcp/log'; +const PROBE = '.stitch-mcp/log-probe'; + +// -------------------------------------------------------------------------- +// Schemas (locked from earlier conversation) +// -------------------------------------------------------------------------- + +const BlobRef = z.object({ sha256: z.string(), size: z.number(), mime: z.string() }); +type BlobRef = z.infer; + +const ProducedScreen = z.object({ + project_id: z.string(), + screen_id: z.string(), + name: z.string(), + title: z.string().optional(), + device_type: z.string().optional(), + generated_by: z.string().optional(), + agent_type: z.string().optional(), + status: z.string().optional(), + width: z.number().optional(), + height: z.number().optional(), + parent_screen_id: z.string().nullable(), + sibling_screen_ids: z.array(z.string()), + effective_prompt: z.string().optional(), + html_file_name: z.string().optional(), + screenshot_file_name: z.string().optional(), + html_blob: BlobRef.optional(), + screenshot_blob: BlobRef.optional(), + theme_blob: BlobRef.optional(), + design_system_blob: BlobRef.optional(), +}); + +const ReadSummary = z.object({ + project_id: z.string().optional(), + screen_id: z.string().optional(), + title: z.string().optional(), + device_type: z.string().optional(), + html_file_name: z.string().optional(), + screenshot_file_name: z.string().optional(), + project_count: z.number().optional(), + project_ids: z.array(z.string()).optional(), + screen_count: z.number().optional(), + screen_ids: z.array(z.string()).optional(), +}); + +const Envelope = z.object({ + id: z.string(), + time: z.string(), + trace_id: z.string(), + schema_version: z.literal(1), + type: z.enum(['call.requested', 'call.completed', 'call.failed', 'observation']), +}); + +const RequestedPayload = z.object({ + tool: z.string(), + project_id: z.string().optional(), + selected_screen_ids: z.array(z.string()).optional(), + user_prompt: z.string().optional(), + variant_options: z.record(z.string(), z.unknown()).optional(), + device_type: z.string().optional(), + model_id: z.string().optional(), + args_blob: BlobRef, +}); + +const CompletedPayload = z.object({ + tool: z.string(), + duration_ms: z.number(), + kind: z.enum(['generative', 'read']), + stitch_session_id: z.string().optional(), + structured_content_blob: BlobRef.optional(), + produced_screens: z.array(ProducedScreen).optional(), + text_components: z.array(z.string()).optional(), + suggestions: z.array(z.string()).optional(), + design_system_asset_name: z.string().optional(), + read_summary: ReadSummary.optional(), +}); + +const FailedPayload = z.object({ + tool: z.string(), + duration_ms: z.number(), + is_error: z.union([z.literal(true), z.literal('empty')]), + error_text: z.string().optional(), + raw_blob: BlobRef.optional(), +}); + +const READ_TOOLS = new Set(['get_screen', 'list_screens', 'list_projects', 'get_project', 'create_project']); +const GEN_TOOLS = new Set(['generate_screen_from_text', 'edit_screens', 'generate_variants']); + +// -------------------------------------------------------------------------- +// Tiny helpers +// -------------------------------------------------------------------------- + +let counter = 0; +function nextId(): string { + // sortable: unix-ms-base36 + counter + random + const t = Date.now().toString(36).padStart(8, '0'); + const c = (counter++).toString(36).padStart(3, '0'); + const r = Math.floor(Math.random() * 1e9).toString(36); + return `${t}${c}${r}`; +} + +function nowIso(): string { + return new Date().toISOString(); +} + +function extForMime(m: string): string { + if (m.startsWith('text/html')) return 'html'; + if (m.startsWith('image/png')) return 'png'; + if (m.startsWith('image/webp')) return 'webp'; + if (m.startsWith('image/jpeg')) return 'jpg'; + if (m.startsWith('application/json') || m.includes('json')) return 'json'; + return 'bin'; +} + +async function ensureDir(p: string) { await mkdir(p, { recursive: true }); } + +// -------------------------------------------------------------------------- +// Blob store +// -------------------------------------------------------------------------- + +class BlobStore { + private root: string; + constructor(root: string) { this.root = root; } + + async putBuffer(buf: Buffer, mime: string): Promise { + const sha = createHash('sha256').update(buf).digest('hex'); + const ext = extForMime(mime); + const dir = join(this.root, sha.slice(0, 2)); + await ensureDir(dir); + const final = join(dir, `${sha}.${ext}`); + try { + const s = await stat(final); + // already exists — dedupped + return { sha256: sha, size: s.size, mime }; + } catch { /* not present */ } + const tmp = join(this.root, '..', 'tmp', `${sha}.${ext}.partial`); + await ensureDir(dirname(tmp)); + await writeFile(tmp, buf); + await rename(tmp, final); + return { sha256: sha, size: buf.length, mime }; + } + + async putJson(value: unknown): Promise { + const buf = Buffer.from(JSON.stringify(value)); + return this.putBuffer(buf, 'application/json'); + } + + async fetchAndPut(url: string, fallbackMime?: string): Promise<{ ref: BlobRef; ok: boolean }> { + try { + const res = await fetch(url, { + redirect: 'follow', + headers: { 'User-Agent': 'Mozilla/5.0 Chrome/120 stitch-mcp-log/0.1' }, + }); + if (!res.ok) return { ref: { sha256: '', size: 0, mime: fallbackMime ?? '' }, ok: false }; + const mime = res.headers.get('content-type')?.split(';')[0] ?? fallbackMime ?? 'application/octet-stream'; + const buf = Buffer.from(await res.arrayBuffer()); + const ref = await this.putBuffer(buf, mime); + return { ref, ok: true }; + } catch { + return { ref: { sha256: '', size: 0, mime: fallbackMime ?? '' }, ok: false }; + } + } +} + +// -------------------------------------------------------------------------- +// Event store +// -------------------------------------------------------------------------- + +class EventStore { + constructor(public root: string, public blobs: BlobStore) {} + + async append(event: unknown): Promise { + Envelope.parse(event); // validate envelope shape early + const line = JSON.stringify(event) + '\n'; + await appendFile(join(this.root, 'events.jsonl'), line, 'utf8'); + } +} + +// -------------------------------------------------------------------------- +// Capture pipeline — turn (request, response, duration) into events +// -------------------------------------------------------------------------- + +interface CaptureCtx { + store: EventStore; + blobs: BlobStore; +} + +function pickScreenComponents(structured: any): any[] { + const oc: any[] = structured?.outputComponents ?? []; + const screens: any[] = []; + for (const c of oc) if (c?.design?.screens) screens.push(...c.design.screens); + return screens; +} + +function pickTextComponents(structured: any): string[] { + return (structured?.outputComponents ?? []).flatMap((c: any) => (typeof c?.text === 'string' ? [c.text] : [])); +} + +function pickSuggestions(structured: any): string[] { + return (structured?.outputComponents ?? []).flatMap((c: any) => (typeof c?.suggestion === 'string' ? [c.suggestion] : [])); +} + +function pickDesignSystemComponent(structured: any): any | null { + const c = (structured?.outputComponents ?? []).find((x: any) => x?.designSystem); + return c?.designSystem ?? null; +} + +async function captureCall( + ctx: CaptureCtx, + args: Record, + tool: string, + result: any, + durationMs: number, + startedAt: string, + finishedAt: string, +): Promise<{ trace_id: string; producedScreenIds: string[] }> { + const trace_id = nextId(); + const argsBlob = await ctx.blobs.putJson(args); + + // requested + const requestedPayload: z.infer = { + tool, + args_blob: argsBlob, + project_id: typeof args.projectId === 'string' ? args.projectId : undefined, + selected_screen_ids: Array.isArray(args.selectedScreenIds) ? args.selectedScreenIds as string[] : undefined, + user_prompt: typeof args.prompt === 'string' ? args.prompt : undefined, + variant_options: typeof args.variantOptions === 'object' && args.variantOptions ? args.variantOptions as any : undefined, + device_type: typeof args.deviceType === 'string' ? args.deviceType : undefined, + model_id: typeof args.modelId === 'string' ? args.modelId : undefined, + }; + await ctx.store.append({ + id: nextId(), time: startedAt, trace_id, schema_version: 1, + type: 'call.requested', payload: RequestedPayload.parse(requestedPayload), + }); + + // failure: explicit + if (result?.isError) { + const errorText = result?.content?.[0]?.text ?? ''; + const rawBlob = await ctx.blobs.putJson(result); + await ctx.store.append({ + id: nextId(), time: finishedAt, trace_id, schema_version: 1, + type: 'call.failed', payload: FailedPayload.parse({ tool, duration_ms: durationMs, is_error: true, error_text: errorText, raw_blob: rawBlob }), + }); + return { trace_id, producedScreenIds: [] }; + } + + const sc = result?.structuredContent ?? null; + const kind: 'generative' | 'read' = GEN_TOOLS.has(tool) ? 'generative' : 'read'; + + if (kind === 'read') { + const summary: z.infer = {}; + if (tool === 'get_screen') { + summary.project_id = (args.projectId as string) ?? undefined; + summary.screen_id = (args.screenId as string) ?? undefined; + summary.title = sc?.title; + summary.device_type = sc?.deviceType; + summary.html_file_name = sc?.htmlCode?.name; + summary.screenshot_file_name = sc?.screenshot?.name; + } else if (tool === 'list_screens') { + summary.project_id = (args.projectId as string) ?? undefined; + const screens: any[] = sc?.screens ?? []; + summary.screen_count = screens.length; + summary.screen_ids = screens.map(s => s?.id ?? s?.name?.split('/screens/')[1]).filter(Boolean); + } else if (tool === 'list_projects') { + const projects: any[] = sc?.projects ?? []; + summary.project_count = projects.length; + summary.project_ids = projects.map(p => p?.name?.replace(/^projects\//, '')).filter(Boolean); + } else if (tool === 'get_project') { + summary.project_id = sc?.name?.replace(/^projects\//, ''); + summary.title = sc?.title; + } else if (tool === 'create_project') { + summary.project_id = sc?.name?.replace(/^projects\//, ''); + summary.title = sc?.title; + } + await ctx.store.append({ + id: nextId(), time: finishedAt, trace_id, schema_version: 1, + type: 'call.completed', + payload: CompletedPayload.parse({ tool, duration_ms: durationMs, kind: 'read', read_summary: summary }), + }); + return { trace_id, producedScreenIds: [] }; + } + + // generative + const screens = pickScreenComponents(sc); + const ds = pickDesignSystemComponent(sc); + const designSystemBlob = ds ? await ctx.blobs.putJson(ds) : undefined; + const dsAssetName: string | undefined = ds?.name; + const stitchSessionId: string | undefined = sc?.sessionId ? String(sc.sessionId) : undefined; + + const selectedParents: string[] = (args.selectedScreenIds as string[] | undefined) ?? []; + const parentScreenId = selectedParents[0] ?? null; + const allIds = screens.map(s => s.id).filter(Boolean); + + const produced: z.infer[] = []; + for (const s of screens) { + const sibs = allIds.filter(id => id !== s.id); + const themeBlob = s.theme && Object.keys(s.theme).length > 0 ? await ctx.blobs.putJson(s.theme) : undefined; + let htmlBlob: BlobRef | undefined; + let screenshotBlob: BlobRef | undefined; + if (s.htmlCode?.downloadUrl) { + const r = await ctx.blobs.fetchAndPut(s.htmlCode.downloadUrl, s.htmlCode.mimeType); + if (r.ok) htmlBlob = r.ref; + } + if (s.screenshot?.downloadUrl) { + const r = await ctx.blobs.fetchAndPut(s.screenshot.downloadUrl); + if (r.ok) screenshotBlob = r.ref; + } + produced.push(ProducedScreen.parse({ + project_id: (args.projectId as string) ?? sc?.projectId ?? '', + screen_id: s.id, + name: s.name, + title: s.title, + device_type: s.deviceType, + generated_by: s.generatedBy, + agent_type: s.screenMetadata?.agentType, + status: s.screenMetadata?.status, + width: typeof s.width === 'string' ? parseInt(s.width) : s.width, + height: typeof s.height === 'string' ? parseInt(s.height) : s.height, + parent_screen_id: parentScreenId, + sibling_screen_ids: sibs, + effective_prompt: s.prompt, + html_file_name: s.htmlCode?.name, + screenshot_file_name: s.screenshot?.name, + html_blob: htmlBlob, + screenshot_blob: screenshotBlob, + theme_blob: themeBlob, + design_system_blob: designSystemBlob, + })); + } + + // implicit failure: zero screens + if (produced.length === 0) { + const rawBlob = await ctx.blobs.putJson(result); + await ctx.store.append({ + id: nextId(), time: finishedAt, trace_id, schema_version: 1, + type: 'call.failed', payload: FailedPayload.parse({ tool, duration_ms: durationMs, is_error: 'empty', raw_blob: rawBlob }), + }); + return { trace_id, producedScreenIds: [] }; + } + + const structuredBlob = sc ? await ctx.blobs.putJson(sc) : undefined; + await ctx.store.append({ + id: nextId(), time: finishedAt, trace_id, schema_version: 1, + type: 'call.completed', + payload: CompletedPayload.parse({ + tool, duration_ms: durationMs, kind: 'generative', + stitch_session_id: stitchSessionId, + structured_content_blob: structuredBlob, + produced_screens: produced, + text_components: pickTextComponents(sc), + suggestions: pickSuggestions(sc), + design_system_asset_name: dsAssetName, + }), + }); + return { trace_id, producedScreenIds: produced.map(p => p.screen_id) }; +} + +// -------------------------------------------------------------------------- +// Index builder +// -------------------------------------------------------------------------- + +async function buildIndex(root: string) { + const indexDir = join(root, 'index'); + await ensureDir(indexDir); + await ensureDir(join(indexDir, 'by-project')); + + const lines = (await readFile(join(root, 'events.jsonl'), 'utf8')).split('\n').filter(Boolean); + const events = lines.map(l => JSON.parse(l)); + + const screens = new Map(); + const prompts: any[] = []; + const lineage: any[] = []; + const byProject = new Map(); + const requestedById = new Map(); + + for (const ev of events) { + if (ev.type === 'call.requested') requestedById.set(ev.trace_id, ev); + if (ev.type === 'call.completed') { + const p = ev.payload; + const projId = p.read_summary?.project_id ?? p.produced_screens?.[0]?.project_id; + if (projId) { + const arr = byProject.get(projId) ?? []; + arr.push({ event_id: ev.id, time: ev.time, type: ev.type, tool: p.tool, screen_ids: (p.produced_screens ?? []).map((s: any) => s.screen_id) }); + byProject.set(projId, arr); + } + if (p.kind === 'generative' && p.produced_screens) { + const req = requestedById.get(ev.trace_id); + const userPrompt = req?.payload?.user_prompt; + for (const s of p.produced_screens) { + screens.set(s.screen_id, { + screen_id: s.screen_id, + project_id: s.project_id, + title: s.title, + device_type: s.device_type, + root: s.parent_screen_id == null, + parent_screen_id: s.parent_screen_id, + created_by_event_id: ev.id, + created_at: ev.time, + user_prompt: userPrompt, + effective_prompt: s.effective_prompt, + html_blob: s.html_blob, + screenshot_blob: s.screenshot_blob, + theme_blob: s.theme_blob, + signals: [], + }); + prompts.push({ + screen_id: s.screen_id, + tool: p.tool, + trace_id: ev.trace_id, + user_prompt: userPrompt, + effective_prompt: s.effective_prompt, + design_system_asset_name: p.design_system_asset_name ?? null, + variant_options: req?.payload?.variant_options ?? null, + signal: null, + }); + if (s.parent_screen_id) { + lineage.push({ + parent_screen_id: s.parent_screen_id, + child_screen_id: s.screen_id, + relation: p.tool === 'edit_screens' ? 'edited_into' : p.tool === 'generate_variants' ? 'variant_of' : 'derived', + sibling_screen_ids: s.sibling_screen_ids ?? [], + trace_id: ev.trace_id, + tool: p.tool, + time: ev.time, + }); + } + } + } + } + } + + // Derived signals: a screen that was viewed (read on it) after an edit/generation = "viewed" + const signals: any[] = []; + const byScreenViews = new Map(); + for (const ev of events) { + if (ev.type === 'call.completed' && ev.payload.kind === 'read' && ev.payload.tool === 'get_screen' && ev.payload.read_summary?.screen_id) { + const sid = ev.payload.read_summary.screen_id; + const arr = byScreenViews.get(sid) ?? []; + arr.push(ev.id); + byScreenViews.set(sid, arr); + } + } + for (const [sid, viewIds] of byScreenViews) { + const sc = screens.get(sid); + if (sc) { + sc.signals.push('viewed'); + signals.push({ screen_id: sid, signal: 'viewed', reason: `get_screen called ${viewIds.length}x`, source_event_ids: viewIds }); + } + } + // Variant survivor / abandoned: among siblings, ones with any downstream activity = survivor + for (const ev of events) { + if (ev.type === 'call.completed' && ev.payload.kind === 'generative' && ev.payload.tool === 'generate_variants' && ev.payload.produced_screens?.length) { + for (const s of ev.payload.produced_screens) { + const hasDownstream = lineage.some(l => l.parent_screen_id === s.screen_id) || byScreenViews.has(s.screen_id); + const sig = hasDownstream ? 'variant_survivor' : 'abandoned'; + signals.push({ screen_id: s.screen_id, signal: sig, reason: 'variant outcome inferred from downstream activity', source_event_id: ev.id }); + const sc = screens.get(s.screen_id); + if (sc) sc.signals.push(sig); + } + } + } + + // Write + await writeFile(join(indexDir, 'screens.jsonl'), [...screens.values()].map(s => JSON.stringify(s)).join('\n') + '\n'); + await writeFile(join(indexDir, 'prompts.jsonl'), prompts.map(p => JSON.stringify(p)).join('\n') + '\n'); + await writeFile(join(indexDir, 'lineage.jsonl'), lineage.map(l => JSON.stringify(l)).join('\n') + '\n'); + await writeFile(join(indexDir, 'signals.jsonl'), signals.map(s => JSON.stringify(s)).join('\n') + '\n'); + for (const [pid, arr] of byProject) { + await writeFile(join(indexDir, 'by-project', `${pid}.jsonl`), arr.map(a => JSON.stringify(a)).join('\n') + '\n'); + } + await writeFile(join(indexDir, 'meta.json'), JSON.stringify({ + schema_version: 1, + built_through_event_id: events[events.length - 1]?.id ?? null, + built_at: nowIso(), + counts: { events: events.length, screens: screens.size, prompts: prompts.length, lineage: lineage.length, signals: signals.length }, + }, null, 2)); +} + +// -------------------------------------------------------------------------- +// Driver +// -------------------------------------------------------------------------- + +interface ProbeReplay { + label: string; + tool: string; + args: Record; + resultFile: string; +} + +// Reconstruct the args we sent in scripts/log-probe{,-2}.ts +function probeReplays(projectId: string, screenId: string): ProbeReplay[] { + return [ + { label: '1', tool: 'list_projects', args: {}, resultFile: '1-list-projects.json' }, + { label: '2', tool: 'create_project', args: { title: `log-probe-replay` }, resultFile: '2-create-project.json' }, + { label: '3', tool: 'generate_screen_from_text', args: { projectId, prompt: 'A minimal landing page for a typography-focused blog. One headline, one paragraph, one CTA.' }, resultFile: '3-generate-screen.json' }, + { label: '4', tool: 'generate_variants', args: { projectId, selectedScreenIds: [screenId], prompt: 'Vary the color scheme and the typographic layout', variantOptions: { variantCount: 2, creativeRange: 'EXPLORE', aspects: ['COLOR_SCHEME', 'LAYOUT'] } }, resultFile: '4-generate-variants.json' }, + { label: '5', tool: 'edit_screens', args: { projectId, selectedScreenIds: [screenId], prompt: 'Make the headline larger and serif' }, resultFile: '5-edit-screens.json' }, + { label: '6', tool: 'get_screen', args: { projectId, screenId }, resultFile: '6-get-screen.json' }, + { label: '7', tool: 'get_project', args: { name: `projects/${projectId}` }, resultFile: '7-get-project-after.json' }, + { label: '8', tool: 'get_screen', args: { projectId: '0', screenId: 'doesnotexist' }, resultFile: '8-fail-bogus-screen.json' }, + { label: '9', tool: 'edit_screens', args: { projectId }, resultFile: '9-fail-missing-required.json' }, + ]; +} + +async function main() { + console.error('=== reset .stitch-mcp/log/ ==='); + // intentionally do not delete — we want to be append-only; user can `rm -rf` if needed + await ensureDir(ROOT); + await ensureDir(join(ROOT, 'blobs')); + await ensureDir(join(ROOT, 'index')); + await ensureDir(join(ROOT, 'tmp')); + + const blobs = new BlobStore(join(ROOT, 'blobs')); + const store = new EventStore(ROOT, blobs); + + // discover the projectId / screenId from probe artifact + const gen = JSON.parse(await readFile(join(PROBE, '3-generate-screen.json'), 'utf8')); + const sc = gen.result.structuredContent; + const projectId: string = sc.projectId; + const screenId: string = pickScreenComponents(sc)[0].id; + + console.error(`projectId=${projectId} firstScreenId=${screenId}`); + + // 1) Replay probes + const replays = probeReplays(projectId, screenId); + for (const r of replays) { + const path = join(PROBE, r.resultFile); + let probe: any; + try { probe = JSON.parse(await readFile(path, 'utf8')); } catch { console.error(`SKIP ${r.label} (missing ${path})`); continue; } + const result = probe.result ?? null; + const duration = probe.duration_ms ?? 0; + if (result == null) { console.error(`SKIP ${r.label} (probe THREW: ${JSON.stringify(probe.threw)})`); continue; } + const finishedAt = nowIso(); + const startedAt = new Date(Date.now() - duration).toISOString(); + const { trace_id, producedScreenIds } = await captureCall({ store, blobs }, r.args, r.tool, result, duration, startedAt, finishedAt); + console.error(` [${r.label}] ${r.tool} trace=${trace_id.slice(0,12)} produced=${producedScreenIds.length}`); + } + + // 2) One LIVE call to validate the live path end-to-end + console.error('\n=== live get_screen ==='); + const transport = new StreamableHTTPClientTransport(new URL(process.env.STITCH_MCP_URL ?? 'https://stitch.googleapis.com/mcp'), { + requestInit: { headers: { 'X-Goog-Api-Key': process.env.STITCH_API_KEY!, Accept: 'application/json, text/event-stream' } }, + }); + const client = new Client({ name: 'log-validate', version: '0.0.0' }, { capabilities: {} }); + await client.connect(transport); + const t0 = Date.now(); + const startedAt = nowIso(); + const liveArgs = { projectId, screenId }; + const liveResult = await client.callTool({ name: 'get_screen', arguments: liveArgs }, CallToolResultSchema); + const dur = Date.now() - t0; + await captureCall({ store, blobs }, liveArgs, 'get_screen', liveResult, dur, startedAt, nowIso()); + console.error(` live get_screen ok in ${dur}ms`); + await client.close(); + + // 3) Build index + console.error('\n=== build index ==='); + const t1 = Date.now(); + await buildIndex(ROOT); + console.error(` index built in ${Date.now() - t1}ms`); +} + +await main(); diff --git a/src/commands/proxy/LoggingCallToolHandler.test.ts b/src/commands/proxy/LoggingCallToolHandler.test.ts new file mode 100644 index 0000000..298d7ed --- /dev/null +++ b/src/commands/proxy/LoggingCallToolHandler.test.ts @@ -0,0 +1,121 @@ +import { describe, test, expect, beforeEach, afterEach } from 'bun:test'; +import { installLoggingCallToolHandler } from './LoggingCallToolHandler.js'; +import type { CaptureInput, CaptureResult, CaptureSpec } from '../../lib/log/capture/spec.js'; + +class RecordingCapture implements CaptureSpec { + inputs: CaptureInput[] = []; + async capture(input: CaptureInput): Promise { + this.inputs.push(input); + return { success: true, data: { trace_id: 't', produced_screen_ids: [], warnings: [] } }; + } +} + +interface FakeProxy { + registered: { schema: any; handler: any } | null; + server: { server: { setRequestHandler: (schema: any, handler: any) => void } }; +} + +function makeFakeProxy(): FakeProxy { + const proxy: FakeProxy = { + registered: null, + server: { + server: { + setRequestHandler: (schema: any, handler: any) => { + proxy.registered = { schema, handler }; + }, + }, + }, + }; + return proxy; +} + +let originalFetch: typeof globalThis.fetch; +beforeEach(() => { originalFetch = globalThis.fetch; }); +afterEach(() => { globalThis.fetch = originalFetch; }); + +describe('installLoggingCallToolHandler', () => { + test('registers a handler that forwards to Stitch and captures the result', async () => { + const cap = new RecordingCapture(); + const proxy = makeFakeProxy(); + let postedUrl = ''; + let postedHeaders: Record = {}; + let postedBody: any = null; + globalThis.fetch = (async (url: any, init?: any) => { + postedUrl = String(url); + postedHeaders = (init?.headers as any) ?? {}; + postedBody = JSON.parse(init?.body as string); + return new Response(JSON.stringify({ jsonrpc: '2.0', id: postedBody.id, result: { structuredContent: { ok: true } } }), { + status: 200, + headers: { 'content-type': 'application/json' }, + }); + }) as any; + + installLoggingCallToolHandler(proxy, cap, { apiKey: 'KEY', url: 'https://stitch.test/mcp' }); + expect(proxy.registered).not.toBeNull(); + + const result = await proxy.registered!.handler({ params: { name: 'list_projects', arguments: {} } }); + expect(result).toEqual({ structuredContent: { ok: true } }); + + // forwarded with right headers + expect(postedUrl).toBe('https://stitch.test/mcp'); + expect(postedHeaders['X-Goog-Api-Key']).toBe('KEY'); + expect(postedBody).toMatchObject({ jsonrpc: '2.0', method: 'tools/call', params: { name: 'list_projects', arguments: {} } }); + + // capture invoked + expect(cap.inputs).toHaveLength(1); + expect(cap.inputs[0]!.tool).toBe('list_projects'); + expect(cap.inputs[0]!.result).toEqual({ structuredContent: { ok: true } }); + }); + + test('on Stitch HTTP error: returns isError envelope to client + captures the failure', async () => { + const cap = new RecordingCapture(); + const proxy = makeFakeProxy(); + globalThis.fetch = (async () => new Response('upstream boom', { status: 500 })) as any; + + installLoggingCallToolHandler(proxy, cap, { apiKey: 'KEY', url: 'https://stitch.test/mcp' }); + const result = await proxy.registered!.handler({ params: { name: 'get_screen', arguments: { projectId: 'p', screenId: 's' } } }); + + expect(result).toMatchObject({ isError: true }); + expect(result.content?.[0]?.text).toContain('Error calling get_screen'); + expect(cap.inputs).toHaveLength(1); + expect((cap.inputs[0]!.result as any).isError).toBe(true); + }); + + test('on Stitch RPC-level error: same — surfaces as isError + captures', async () => { + const cap = new RecordingCapture(); + const proxy = makeFakeProxy(); + globalThis.fetch = (async () => new Response(JSON.stringify({ jsonrpc: '2.0', id: 1, error: { code: -1, message: 'permission denied' } }), { + status: 200, + headers: { 'content-type': 'application/json' }, + })) as any; + + installLoggingCallToolHandler(proxy, cap, { apiKey: 'KEY', url: 'https://stitch.test/mcp' }); + const result = await proxy.registered!.handler({ params: { name: 'list_projects', arguments: {} } }); + expect(result.isError).toBe(true); + expect(result.content?.[0]?.text).toContain('permission denied'); + expect(cap.inputs).toHaveLength(1); + }); + + test('throws if proxy.server.server.setRequestHandler is missing', () => { + const cap = new RecordingCapture(); + expect(() => installLoggingCallToolHandler({}, cap, { apiKey: 'K', url: 'u' })).toThrow(); + }); + + test('throws if no apiKey available', () => { + const cap = new RecordingCapture(); + const proxy = makeFakeProxy(); + expect(() => installLoggingCallToolHandler(proxy, cap, { apiKey: '', url: 'u' })).toThrow(/STITCH_API_KEY/); + }); + + test('capture failure does not propagate; client still gets the result', async () => { + const proxy = makeFakeProxy(); + globalThis.fetch = (async () => new Response(JSON.stringify({ jsonrpc: '2.0', id: 1, result: { ok: 1 } }), { + status: 200, headers: { 'content-type': 'application/json' }, + })) as any; + const flakey: CaptureSpec = { capture: async () => { throw new Error('flake'); } }; + + installLoggingCallToolHandler(proxy, flakey, { apiKey: 'K', url: 'u' }); + const r = await proxy.registered!.handler({ params: { name: 'list_projects', arguments: {} } }); + expect(r).toEqual({ ok: 1 }); + }); +}); diff --git a/src/commands/proxy/LoggingCallToolHandler.ts b/src/commands/proxy/LoggingCallToolHandler.ts new file mode 100644 index 0000000..e867be7 --- /dev/null +++ b/src/commands/proxy/LoggingCallToolHandler.ts @@ -0,0 +1,82 @@ +import { CallToolRequestSchema } from '@modelcontextprotocol/sdk/types.js'; +import type { CaptureSpec } from '../../lib/log/capture/spec.js'; + +const DEFAULT_STITCH_MCP_URL = 'https://stitch.googleapis.com/mcp'; + +interface ForwardOptions { + apiKey: string; + url: string; +} + +/** + * Replace the SDK proxy's tools/call handler with a capture-wrapped variant. + * Must be called AFTER {@link StitchProxy.start} (which registers the original). + */ +export function installLoggingCallToolHandler(proxy: any, capture: CaptureSpec, opts?: Partial): void { + const apiKey = opts?.apiKey ?? process.env.STITCH_API_KEY; + const url = opts?.url ?? process.env.STITCH_MCP_URL ?? DEFAULT_STITCH_MCP_URL; + if (!apiKey) throw new Error('logging proxy requires STITCH_API_KEY'); + + // The SDK exposes McpServer at proxy.server, whose underlying Server is at .server. + const server = proxy?.server?.server; + if (!server || typeof server.setRequestHandler !== 'function') { + throw new Error('cannot install logging handler: proxy.server.server.setRequestHandler missing'); + } + + server.setRequestHandler(CallToolRequestSchema, async (request: any) => { + const { name, arguments: args } = request.params; + const startedAt = new Date().toISOString(); + const t0 = Date.now(); + + let result: any; + try { + result = await forwardToStitch({ apiKey, url }, name, args ?? {}); + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + result = { content: [{ type: 'text', text: `Error calling ${name}: ${message}` }], isError: true }; + } + const finishedAt = new Date().toISOString(); + const durationMs = Date.now() - t0; + + // best-effort capture — never propagate failures to the MCP client + try { + await capture.capture({ + tool: name, + args: (args ?? {}) as Record, + result, + duration_ms: durationMs, + started_at: startedAt, + finished_at: finishedAt, + }); + } catch (err) { + const msg = err instanceof Error ? err.message : String(err); + // eslint-disable-next-line no-console + console.error(`[stitch-mcp log] capture failed: ${msg}`); + } + + return result; + }); +} + +/** Minimal JSON-RPC forwarder, mirroring the SDK's private forwardToStitch. */ +async function forwardToStitch(opts: ForwardOptions, name: string, args: Record): Promise { + const body = { + jsonrpc: '2.0', + method: 'tools/call', + params: { name, arguments: args }, + id: Date.now(), + }; + const res = await globalThis.fetch(opts.url, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + Accept: 'application/json', + 'X-Goog-Api-Key': opts.apiKey, + }, + body: JSON.stringify(body), + }); + if (!res.ok) throw new Error(`Stitch API error (${res.status}): ${await res.text()}`); + const json = (await res.json()) as { result?: unknown; error?: { message: string } }; + if (json.error) throw new Error(`Stitch RPC error: ${json.error.message}`); + return json.result; +} diff --git a/src/commands/proxy/handler.ts b/src/commands/proxy/handler.ts index 59403ee..7db848b 100644 --- a/src/commands/proxy/handler.ts +++ b/src/commands/proxy/handler.ts @@ -1,6 +1,8 @@ import { StitchProxy } from '@google/stitch-sdk'; import type { StitchProxy as StitchProxyType } from '@google/stitch-sdk'; import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js'; +import { createCaptureHandler, isLogEnabled } from '../../lib/log/factory.js'; +import { installLoggingCallToolHandler } from './LoggingCallToolHandler.js'; interface ProxyCommandInput { port?: number; @@ -32,6 +34,19 @@ export class ProxyCommandHandler { }); const transport = this.createTransport(); await proxy.start(transport); + + // Override the SDK's tools/call handler with a capture-wrapped variant. + // Must run AFTER proxy.start() (which registers the original handler). + if (isLogEnabled()) { + try { + installLoggingCallToolHandler(proxy, createCaptureHandler()); + } catch (err) { + const msg = err instanceof Error ? err.message : String(err); + // eslint-disable-next-line no-console + console.error(`[stitch-mcp log] failed to install capture handler: ${msg}`); + } + } + await transport.onclose; return { success: true, data: { status: 'running' } }; } catch (e: any) { diff --git a/src/commands/tool/handler.ts b/src/commands/tool/handler.ts index 12b37b5..b743ebb 100644 --- a/src/commands/tool/handler.ts +++ b/src/commands/tool/handler.ts @@ -11,6 +11,8 @@ import { ShowSchemaStep } from './steps/ShowSchemaStep.js'; import { ParseArgsStep } from './steps/ParseArgsStep.js'; import { ValidateToolStep } from './steps/ValidateToolStep.js'; import { ExecuteToolStep } from './steps/ExecuteToolStep.js'; +import { LogExecuteToolStep } from './steps/LogExecuteToolStep.js'; +import { createCaptureHandler, isLogEnabled } from '../../lib/log/factory.js'; export const deps = { runSteps, @@ -31,12 +33,15 @@ export class ToolCommandHandler { this.client = client || new StitchToolClient(); this.stitchInstance = stitchInstance || defaultStitch; this.tools = tools || defaultVirtualTools; + const executeStep: CommandStep = isLogEnabled() + ? new LogExecuteToolStep(createCaptureHandler()) + : new deps.ExecuteToolStep(); this.steps = [ new deps.ListToolsStep(), new deps.ShowSchemaStep(), new deps.ParseArgsStep(), new deps.ValidateToolStep(), - new deps.ExecuteToolStep(), + executeStep, ]; } diff --git a/src/commands/tool/steps/LogExecuteToolStep.test.ts b/src/commands/tool/steps/LogExecuteToolStep.test.ts new file mode 100644 index 0000000..e2ff4ee --- /dev/null +++ b/src/commands/tool/steps/LogExecuteToolStep.test.ts @@ -0,0 +1,129 @@ +import { describe, test, expect } from 'bun:test'; +import { LogExecuteToolStep } from './LogExecuteToolStep.js'; +import type { CaptureInput, CaptureResult, CaptureSpec } from '../../../lib/log/capture/spec.js'; +import type { ToolContext } from '../context.js'; + +class RecordingCapture implements CaptureSpec { + inputs: CaptureInput[] = []; + async capture(input: CaptureInput): Promise { + this.inputs.push(input); + return { success: true, data: { trace_id: 't', produced_screen_ids: [], warnings: [] } }; + } +} + +function makeContext(opts: { + toolName: string; + parsedArgs: Record; + rawResult?: any; + rawThrows?: Error; + isConnected?: boolean; +}): ToolContext { + // Build a fake StitchToolClient-shaped object whose private `client` calls our stub. + const inner = { + callTool: async () => { + if (opts.rawThrows) throw opts.rawThrows; + return opts.rawResult; + }, + }; + const fakeStitchClient: any = { + isConnected: opts.isConnected ?? true, + connect: async () => { fakeStitchClient.isConnected = true; }, + client: inner, + callTool: async () => { throw new Error('should not be called by LogExecuteToolStep'); }, + listTools: async () => ({ tools: [] }), + close: async () => {}, + }; + return { + input: { toolName: opts.toolName, output: 'json', showSchema: false }, + client: fakeStitchClient, + stitch: {} as any, + virtualTools: [], + parsedArgs: opts.parsedArgs, + } as unknown as ToolContext; +} + +describe('LogExecuteToolStep', () => { + test('makes the raw call, captures the envelope, and sets context.result from structuredContent', async () => { + const cap = new RecordingCapture(); + const step = new LogExecuteToolStep(cap); + const raw = { structuredContent: { hello: 'world', sessionId: 'sess1' } }; + const ctx = makeContext({ toolName: 'list_projects', parsedArgs: {}, rawResult: raw }); + + const r = await step.run(ctx); + expect(r.success).toBe(true); + expect(cap.inputs).toHaveLength(1); + expect(cap.inputs[0]!.tool).toBe('list_projects'); + expect(cap.inputs[0]!.result).toEqual(raw); + expect(typeof cap.inputs[0]!.duration_ms).toBe('number'); + expect(ctx.result).toEqual({ success: true, data: { hello: 'world', sessionId: 'sess1' } }); + }); + + test('on isError envelope: sets context.result.success=false and still captures', async () => { + const cap = new RecordingCapture(); + const step = new LogExecuteToolStep(cap); + const raw = { isError: true, content: [{ type: 'text', text: 'Requested entity was not found.' }] }; + const ctx = makeContext({ toolName: 'get_screen', parsedArgs: { projectId: '0', screenId: 'x' }, rawResult: raw }); + + const r = await step.run(ctx); + expect(r.success).toBe(true); + expect(ctx.result).toMatchObject({ success: false }); + if (ctx.result?.success) throw new Error(); + expect(ctx.result?.error).toContain('Requested entity was not found.'); + expect(cap.inputs).toHaveLength(1); + }); + + test('connect() is called when not yet connected', async () => { + const cap = new RecordingCapture(); + const step = new LogExecuteToolStep(cap); + const raw = { structuredContent: { ok: true } }; + const ctx = makeContext({ toolName: 'list_projects', parsedArgs: {}, rawResult: raw, isConnected: false }); + + expect((ctx.client as any).isConnected).toBe(false); + await step.run(ctx); + expect((ctx.client as any).isConnected).toBe(true); + }); + + test('virtual tools bypass capture and use the original execute path', async () => { + const cap = new RecordingCapture(); + const step = new LogExecuteToolStep(cap); + const ctx = makeContext({ toolName: 'list-tools-virtual', parsedArgs: {} }); + (ctx as any).virtualTools = [{ + name: 'list-tools-virtual', + description: '', + execute: async () => ({ stub: true }), + }]; + + const r = await step.run(ctx); + expect(r.success).toBe(true); + expect(ctx.result).toEqual({ success: true, data: { stub: true } }); + expect(cap.inputs).toHaveLength(0); // no capture for virtual tools + }); + + test('parseToolResponse fallback: parses content[0].text as JSON', async () => { + const cap = new RecordingCapture(); + const step = new LogExecuteToolStep(cap); + const raw = { content: [{ type: 'text', text: '{"a":1}' }] }; + const ctx = makeContext({ toolName: 'list_projects', parsedArgs: {}, rawResult: raw }); + + await step.run(ctx); + expect(ctx.result).toEqual({ success: true, data: { a: 1 } }); + expect(cap.inputs).toHaveLength(1); + }); + + test('underlying call throws → step returns failure AND still captures as isError envelope', async () => { + const cap = new RecordingCapture(); + const step = new LogExecuteToolStep(cap); + const ctx = makeContext({ toolName: 'list_projects', parsedArgs: {}, rawThrows: new Error('boom') }); + + const r = await step.run(ctx); + expect(r.success).toBe(false); + expect(ctx.result).toMatchObject({ success: false, error: 'boom' }); + + // Per Bug #2: transport-level throws must still capture so we don't lose the attempt. + expect(cap.inputs).toHaveLength(1); + const passed = cap.inputs[0]!; + expect(passed.tool).toBe('list_projects'); + expect((passed.result as any).isError).toBe(true); + expect((passed.result as any).content[0].text).toBe('boom'); + }); +}); diff --git a/src/commands/tool/steps/LogExecuteToolStep.ts b/src/commands/tool/steps/LogExecuteToolStep.ts new file mode 100644 index 0000000..19b2dd2 --- /dev/null +++ b/src/commands/tool/steps/LogExecuteToolStep.ts @@ -0,0 +1,107 @@ +import { CallToolResultSchema } from '@modelcontextprotocol/sdk/types.js'; +import type { CommandStep, StepResult } from '../../../framework/CommandStep.js'; +import type { CaptureSpec } from '../../../lib/log/capture/spec.js'; +import type { ToolContext } from '../context.js'; + +/** + * Execute step variant that captures the raw MCP envelope before parsing. + * Used in place of {@link ExecuteToolStep} when STITCH_MCP_LOG=1. + * + * Reaches through the {@link StitchToolClient} into its underlying MCP `Client` + * so the raw {@link CallToolResult} is observable (the public callTool() throws + * on isError, losing the envelope). + */ +export class LogExecuteToolStep implements CommandStep { + id = 'execute-tool'; + name = 'Execute tool (with capture)'; + + constructor(private readonly capture: CaptureSpec) {} + + async shouldRun(context: ToolContext): Promise { + return context.parsedArgs !== undefined; + } + + async run(context: ToolContext): Promise { + const tool = context.input.toolName!; + const args = context.parsedArgs!; + + // Virtual tools have no MCP envelope to capture; behave exactly as ExecuteToolStep. + const virtualTool = context.virtualTools.find((t) => t.name === tool); + if (virtualTool) { + try { + const result = await virtualTool.execute(context.client, args, context.stitch); + context.result = { success: true, data: result }; + return { success: true }; + } catch (e: any) { + context.result = { success: false, error: `Virtual tool execution failed: ${e.message || String(e)}` }; + return { success: false, error: e }; + } + } + + // Ensure connected via the public surface (StitchToolClient handles auth + transport). + const stitch = context.client as any; + if (!stitch.isConnected) { + await stitch.connect(); + } + const rawClient = stitch.client; // the underlying MCP SDK Client + if (!rawClient) { + context.result = { success: false, error: 'logging path requires a connected StitchToolClient' }; + return { success: false }; + } + + const startedAt = new Date().toISOString(); + const t0 = Date.now(); + let raw: any = null; + let threw: Error | null = null; + try { + // Stitch generations can take minutes; the MCP SDK default (60s) is too short. + raw = await rawClient.callTool({ name: tool, arguments: args }, CallToolResultSchema, { timeout: 600_000 }); + } catch (e: any) { + threw = e instanceof Error ? e : new Error(String(e)); + } + const finishedAt = new Date().toISOString(); + const durationMs = Date.now() - t0; + + // Always capture — including transport-level throws — so we don't lose the attempt. + try { + await this.capture.capture({ + tool, args, + result: threw + ? { isError: true, content: [{ type: 'text', text: threw.message }] } + : raw, + duration_ms: durationMs, + started_at: startedAt, finished_at: finishedAt, + }); + } catch (e) { + const msg = e instanceof Error ? e.message : String(e); + // eslint-disable-next-line no-console + console.error(`[stitch-mcp log] capture failed: ${msg}`); + } + + if (threw) { + context.result = { success: false, error: threw.message }; + return { success: false, error: threw }; + } + context.result = parseToolResponse(raw); + return { success: true }; + } +} + +/** Mirrors {@link StitchToolClient}'s parseToolResponse for context.result shape. */ +function parseToolResponse(raw: any): { success: boolean; data?: any; error?: string } { + if (raw?.isError) { + const errorText = (raw.content ?? []) + .map((c: any) => (c.type === 'text' ? c.text : '')) + .join(''); + return { success: false, error: errorText }; + } + if (raw?.structuredContent) { + return { success: true, data: raw.structuredContent }; + } + const textContent = raw?.content?.find((c: any) => c.type === 'text'); + if (textContent?.text) { + try { return { success: true, data: JSON.parse(textContent.text) }; } + catch { return { success: true, data: textContent.text }; } + } + return { success: true, data: raw }; +} diff --git a/src/lib/log/append.test.ts b/src/lib/log/append.test.ts new file mode 100644 index 0000000..b793056 --- /dev/null +++ b/src/lib/log/append.test.ts @@ -0,0 +1,78 @@ +import { describe, test, expect, beforeEach, afterEach } from 'bun:test'; +import { mkdtemp, readFile, rm } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; +import { appendEvent } from './append.js'; + +let root: string; +let path: string; + +beforeEach(async () => { + root = await mkdtemp(join(tmpdir(), 'append-')); + path = join(root, 'events.jsonl'); +}); +afterEach(async () => { + await rm(root, { recursive: true, force: true }); +}); + +const validEvent = { + id: '01abc', + time: '2026-04-25T18:00:00Z', + trace_id: '01trace', + schema_version: 1 as const, + type: 'call.requested', + payload: { tool: 'list_projects' }, +}; + +describe('appendEvent', () => { + test('writes exactly one line ending in newline', async () => { + const r = await appendEvent(path, validEvent); + expect(r.success).toBe(true); + + const content = await readFile(path, 'utf8'); + expect(content.endsWith('\n')).toBe(true); + const lines = content.split('\n').filter(Boolean); + expect(lines).toHaveLength(1); + expect(JSON.parse(lines[0]!)).toMatchObject({ id: '01abc', type: 'call.requested' }); + }); + + test('appends successive events on separate lines', async () => { + await appendEvent(path, validEvent); + await appendEvent(path, { ...validEvent, id: '01def', type: 'call.completed' }); + const content = await readFile(path, 'utf8'); + const lines = content.split('\n').filter(Boolean); + expect(lines).toHaveLength(2); + expect(JSON.parse(lines[1]!).type).toBe('call.completed'); + }); + + test('returns EVENT_VALIDATION_FAILED on bad envelope', async () => { + const r = await appendEvent(path, { ...validEvent, schema_version: 2 }); + expect(r.success).toBe(false); + if (r.success) return; + expect(r.error.code).toBe('EVENT_VALIDATION_FAILED'); + }); + + test('returns EVENT_VALIDATION_FAILED when type is missing', async () => { + const { type: _t, ...without } = validEvent; + const r = await appendEvent(path, without); + expect(r.success).toBe(false); + if (r.success) return; + expect(r.error.code).toBe('EVENT_VALIDATION_FAILED'); + }); + + test('returns EVENT_WRITE_FAILED when destination cannot be written', async () => { + const r = await appendEvent('/proc/cannot-write-here/x.jsonl', validEvent); + expect(r.success).toBe(false); + if (r.success) return; + expect(r.error.code).toBe('EVENT_WRITE_FAILED'); + }); + + test('30 concurrent appends produce 30 well-formed lines (line atomicity)', async () => { + const events = Array.from({ length: 30 }, (_, i) => ({ ...validEvent, id: `e${i}` })); + await Promise.all(events.map((e) => appendEvent(path, e))); + const content = await readFile(path, 'utf8'); + const lines = content.split('\n').filter(Boolean); + expect(lines).toHaveLength(30); + for (const l of lines) expect(() => JSON.parse(l)).not.toThrow(); + }); +}); diff --git a/src/lib/log/append.ts b/src/lib/log/append.ts new file mode 100644 index 0000000..c81c52e --- /dev/null +++ b/src/lib/log/append.ts @@ -0,0 +1,52 @@ +import { appendFile, mkdir } from 'node:fs/promises'; +import { dirname } from 'node:path'; +import { z } from 'zod'; + +const EnvelopeSchema = z.object({ + id: z.string().min(1), + time: z.string().min(1), + trace_id: z.string().min(1), + schema_version: z.literal(1), + type: z.string().min(1), + payload: z.unknown(), +}); + +export type AppendResult = + | { success: true } + | { + success: false; + error: { + code: 'EVENT_VALIDATION_FAILED' | 'EVENT_WRITE_FAILED'; + message: string; + recoverable: boolean; + }; + }; + +/** Validate envelope shape and append exactly one JSON line, ending in `\n`. */ +export async function appendEvent(eventsPath: string, event: unknown): Promise { + const parsed = EnvelopeSchema.safeParse(event); + if (!parsed.success) { + return { + success: false, + error: { + code: 'EVENT_VALIDATION_FAILED', + message: parsed.error.issues.map((i) => `${i.path.join('.')}: ${i.message}`).join('; '), + recoverable: false, + }, + }; + } + try { + await mkdir(dirname(eventsPath), { recursive: true }); + await appendFile(eventsPath, JSON.stringify(parsed.data) + '\n', 'utf8'); + return { success: true }; + } catch (e) { + return { + success: false, + error: { + code: 'EVENT_WRITE_FAILED', + message: e instanceof Error ? e.message : String(e), + recoverable: false, + }, + }; + } +} diff --git a/src/lib/log/blob-store/handler.test.ts b/src/lib/log/blob-store/handler.test.ts new file mode 100644 index 0000000..ccd3867 --- /dev/null +++ b/src/lib/log/blob-store/handler.test.ts @@ -0,0 +1,177 @@ +import { describe, test, expect, beforeEach, afterEach } from 'bun:test'; +import { mkdtemp, rm, readdir, stat } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; +import { createHash } from 'node:crypto'; + +import { BlobStoreHandler } from './handler.js'; +import { BlobRefSchema } from './spec.js'; + +let root: string; + +beforeEach(async () => { + root = await mkdtemp(join(tmpdir(), 'blob-store-')); +}); + +afterEach(async () => { + await rm(root, { recursive: true, force: true }); +}); + +describe('put()', () => { + test('is idempotent — same content does not write twice', async () => { + const store = new BlobStoreHandler(root); + const buf = Buffer.from('hello'); + + const r1 = await store.put(buf, 'text/plain'); + expect(r1.success).toBe(true); + if (!r1.success) return; + const path = join(root, r1.data.sha256.slice(0, 2), `${r1.data.sha256}.bin`); + const mtime1 = (await stat(path)).mtimeMs; + + // wait a tick so a re-write would have a different mtime + await new Promise((res) => setTimeout(res, 10)); + + const r2 = await store.put(buf, 'text/plain'); + expect(r2.success).toBe(true); + if (!r2.success) return; + expect(r2.data.sha256).toBe(r1.data.sha256); + + const mtime2 = (await stat(path)).mtimeMs; + expect(mtime2).toBe(mtime1); // not re-written + }); +}); + +describe('has() / get()', () => { + test('has() reflects whether a sha is present', async () => { + const store = new BlobStoreHandler(root); + const buf = Buffer.from('payload'); + const put = await store.put(buf, 'application/json'); + expect(put.success).toBe(true); + if (!put.success) return; + + const present = await store.has(put.data.sha256); + expect(present.success).toBe(true); + if (!present.success) return; + expect(present.data).toBe(true); + + const absent = await store.has('0'.repeat(64)); + expect(absent.success).toBe(true); + if (!absent.success) return; + expect(absent.data).toBe(false); + }); + + test('get() round-trips bytes; null when absent', async () => { + const store = new BlobStoreHandler(root); + const buf = Buffer.from([1, 2, 3, 4, 5]); + const put = await store.put(buf, 'application/json'); + expect(put.success).toBe(true); + if (!put.success) return; + + const present = await store.get(put.data.sha256); + expect(present.success).toBe(true); + if (!present.success) return; + expect(present.data).not.toBeNull(); + expect(Buffer.compare(present.data!, buf)).toBe(0); + + const absent = await store.get('0'.repeat(64)); + expect(absent.success).toBe(true); + if (!absent.success) return; + expect(absent.data).toBeNull(); + }); +}); + +describe('put() — error mapping', () => { + test('returns Failure(BLOB_WRITE_FAILED) when the root path is not writable', async () => { + // Use a path that contains a non-directory ancestor (a regular file). mkdir will EEXIST/ENOTDIR. + const buf = Buffer.from('x'); + const fakeFile = join(root, 'this-is-a-file'); + await Bun.write(fakeFile, 'sentinel'); + const store = new BlobStoreHandler(join(fakeFile, 'cannot-be-a-dir')); + + const r = await store.put(buf, 'application/json'); + expect(r.success).toBe(false); + if (r.success) return; + expect(r.error.code).toBe('BLOB_WRITE_FAILED'); + expect(r.error.message).toBeDefined(); + }); +}); + +describe('fetch()', () => { + let originalFetch: typeof globalThis.fetch; + + beforeEach(() => { originalFetch = globalThis.fetch; }); + afterEach(() => { globalThis.fetch = originalFetch; }); + + test('downloads bytes and stores them as a blob', async () => { + const html = Buffer.from('

hi

'); + let capturedUrl: string | undefined; + let capturedInit: RequestInit | undefined; + globalThis.fetch = (async (input: any, init?: any) => { + capturedUrl = String(input); + capturedInit = init; + return new Response(html, { + status: 200, + headers: { 'content-type': 'text/html; charset=utf-8' }, + }); + }) as any; + + const store = new BlobStoreHandler(root); + const r = await store.fetch('https://example.com/page'); + expect(r.success).toBe(true); + if (!r.success) return; + + expect(r.data.mime).toBe('text/html'); + expect(r.data.size).toBe(html.length); + expect(capturedUrl).toBe('https://example.com/page'); + expect((capturedInit as any)?.redirect).toBe('follow'); + + // round-trip via get() + const got = await store.get(r.data.sha256); + expect(got.success).toBe(true); + if (!got.success) return; + expect(Buffer.compare(got.data!, html)).toBe(0); + }); + + test('returns Failure(BLOB_FETCH_HTTP_ERROR) on non-2xx', async () => { + globalThis.fetch = (async () => new Response('nope', { status: 404 })) as any; + const store = new BlobStoreHandler(root); + const r = await store.fetch('https://example.com/missing'); + expect(r.success).toBe(false); + if (r.success) return; + expect(r.error.code).toBe('BLOB_FETCH_HTTP_ERROR'); + expect(r.error.message).toContain('404'); + }); + + test('returns Failure(BLOB_FETCH_NETWORK) when fetch throws', async () => { + globalThis.fetch = (async () => { throw new TypeError('connection refused'); }) as any; + const store = new BlobStoreHandler(root); + const r = await store.fetch('https://example.com/down'); + expect(r.success).toBe(false); + if (r.success) return; + expect(r.error.code).toBe('BLOB_FETCH_NETWORK'); + expect(r.error.recoverable).toBe(true); + }); +}); + +describe('put() — extra cases', () => { + test('writes one file under /. and returns a valid BlobRef', async () => { + const store = new BlobStoreHandler(root); + const buf = Buffer.from(JSON.stringify({ hello: 'world' })); + const expectedSha = createHash('sha256').update(buf).digest('hex'); + + const result = await store.put(buf, 'application/json'); + + expect(result.success).toBe(true); + if (!result.success) return; + + expect(BlobRefSchema.safeParse(result.data).success).toBe(true); + expect(result.data.sha256).toBe(expectedSha); + expect(result.data.size).toBe(buf.length); + expect(result.data.mime).toBe('application/json'); + + const expectedPath = join(root, expectedSha.slice(0, 2), `${expectedSha}.json`); + const s = await stat(expectedPath); + expect(s.isFile()).toBe(true); + expect(s.size).toBe(buf.length); + }); +}); diff --git a/src/lib/log/blob-store/handler.ts b/src/lib/log/blob-store/handler.ts new file mode 100644 index 0000000..ad7d07b --- /dev/null +++ b/src/lib/log/blob-store/handler.ts @@ -0,0 +1,94 @@ +import { createHash } from 'node:crypto'; +import { mkdir, readdir, readFile, stat, writeFile } from 'node:fs/promises'; +import { dirname, join } from 'node:path'; +import type { BlobStoreSpec, PutResult, HasResult, GetResult } from './spec.js'; + +const EXT_BY_MIME: Record = { + 'application/json': 'json', + 'text/html': 'html', + 'image/png': 'png', + 'image/webp': 'webp', + 'image/jpeg': 'jpg', +}; + +function extForMime(mime: string): string { + return EXT_BY_MIME[mime] ?? 'bin'; +} + +export class BlobStoreHandler implements BlobStoreSpec { + constructor(private readonly root: string) {} + + async put(buffer: Buffer, mime: string): Promise { + try { + const sha256 = createHash('sha256').update(buffer).digest('hex'); + const existing = await this.findBySha(sha256); + if (existing) { + const s = await stat(existing); + return { success: true, data: { sha256, size: s.size, mime } }; + } + const path = join(this.root, sha256.slice(0, 2), `${sha256}.${extForMime(mime)}`); + await mkdir(dirname(path), { recursive: true }); + await writeFile(path, buffer); + return { success: true, data: { sha256, size: buffer.length, mime } }; + } catch (e) { + return { + success: false, + error: { + code: 'BLOB_WRITE_FAILED', + message: e instanceof Error ? e.message : String(e), + recoverable: false, + }, + }; + } + } + + private async findBySha(sha256: string): Promise { + const dir = join(this.root, sha256.slice(0, 2)); + let entries: string[]; + try { entries = await readdir(dir); } catch { return null; } + const match = entries.find((e) => e.startsWith(`${sha256}.`)); + return match ? join(dir, match) : null; + } + + async fetch(url: string, mimeHint?: string): Promise { + let response: Response; + try { + response = await globalThis.fetch(url, { + redirect: 'follow', + headers: { 'User-Agent': 'stitch-mcp-log/0.1 (Mozilla/5.0)' }, + }); + } catch (e) { + return { + success: false, + error: { + code: 'BLOB_FETCH_NETWORK', + message: e instanceof Error ? e.message : String(e), + recoverable: true, + }, + }; + } + if (!response.ok) { + return { + success: false, + error: { + code: 'BLOB_FETCH_HTTP_ERROR', + message: `HTTP ${response.status} for ${url}`, + recoverable: false, + }, + }; + } + const mime = (response.headers.get('content-type') ?? mimeHint ?? 'application/octet-stream').split(';')[0]!.trim(); + const buffer = Buffer.from(await response.arrayBuffer()); + return this.put(buffer, mime); + } + async has(sha256: string): Promise { + const path = await this.findBySha(sha256); + return { success: true, data: path != null }; + } + async get(sha256: string): Promise { + const path = await this.findBySha(sha256); + if (!path) return { success: true, data: null }; + const buf = await readFile(path); + return { success: true, data: buf }; + } +} diff --git a/src/lib/log/blob-store/spec.test.ts b/src/lib/log/blob-store/spec.test.ts new file mode 100644 index 0000000..8a2e5c1 --- /dev/null +++ b/src/lib/log/blob-store/spec.test.ts @@ -0,0 +1,94 @@ +import { describe, test, expect } from 'bun:test'; +import { + BlobRefSchema, + BlobStoreErrorSchema, + Sha256Schema, + PutResultSchema, + GetResultSchema, + HasResultSchema, +} from './spec.js'; + +describe('Sha256Schema', () => { + const validShas = [ + 'a'.repeat(64), + '0123456789abcdef'.repeat(4), + ]; + for (const v of validShas) { + test(`accepts valid sha (len=${v.length})`, () => { + expect(Sha256Schema.safeParse(v).success).toBe(true); + }); + } + + const invalidCases: { val: unknown; reason: string }[] = [ + { val: '', reason: 'empty' }, + { val: 'A'.repeat(64), reason: 'uppercase rejected' }, + { val: 'g'.repeat(64), reason: 'non-hex char' }, + { val: 'a'.repeat(63), reason: 'too short' }, + { val: 'a'.repeat(65), reason: 'too long' }, + { val: 12345, reason: 'not a string' }, + ]; + for (const c of invalidCases) { + test(`rejects ${c.reason}`, () => { + expect(Sha256Schema.safeParse(c.val).success).toBe(false); + }); + } +}); + +describe('BlobRefSchema', () => { + test('accepts a well-formed ref', () => { + const r = BlobRefSchema.safeParse({ sha256: 'a'.repeat(64), size: 0, mime: 'application/json' }); + expect(r.success).toBe(true); + }); + + test('rejects negative size', () => { + const r = BlobRefSchema.safeParse({ sha256: 'a'.repeat(64), size: -1, mime: 'text/html' }); + expect(r.success).toBe(false); + }); + + test('rejects empty mime', () => { + const r = BlobRefSchema.safeParse({ sha256: 'a'.repeat(64), size: 1, mime: '' }); + expect(r.success).toBe(false); + }); +}); + +describe('BlobStoreErrorSchema', () => { + test('accepts a known error code', () => { + const r = BlobStoreErrorSchema.safeParse({ code: 'BLOB_WRITE_FAILED', message: 'disk full', recoverable: false }); + expect(r.success).toBe(true); + }); + + test('rejects unknown error code', () => { + const r = BlobStoreErrorSchema.safeParse({ code: 'NOPE', message: 'x', recoverable: true }); + expect(r.success).toBe(false); + }); +}); + +describe('Result discriminated unions', () => { + test('PutResult: success branch parses', () => { + const r = PutResultSchema.safeParse({ + success: true, + data: { sha256: 'a'.repeat(64), size: 10, mime: 'image/png' }, + }); + expect(r.success).toBe(true); + }); + + test('PutResult: failure branch parses', () => { + const r = PutResultSchema.safeParse({ + success: false, + error: { code: 'BLOB_FETCH_HTTP_ERROR', message: '404', recoverable: false }, + }); + expect(r.success).toBe(true); + }); + + test('GetResult: data null is valid (absent)', () => { + const r = GetResultSchema.safeParse({ success: true, data: null }); + expect(r.success).toBe(true); + }); + + test('HasResult: data must be boolean', () => { + const ok = HasResultSchema.safeParse({ success: true, data: true }); + expect(ok.success).toBe(true); + const bad = HasResultSchema.safeParse({ success: true, data: 'yes' }); + expect(bad.success).toBe(false); + }); +}); diff --git a/src/lib/log/blob-store/spec.ts b/src/lib/log/blob-store/spec.ts new file mode 100644 index 0000000..baa9235 --- /dev/null +++ b/src/lib/log/blob-store/spec.ts @@ -0,0 +1,77 @@ +import { z } from 'zod'; + +// --- shared schemas --------------------------------------------------------- + +export const Sha256Schema = z.string().regex(/^[0-9a-f]{64}$/, 'must be 64-char lowercase hex'); +export const MimeSchema = z.string().min(1); + +export const BlobRefSchema = z.object({ + sha256: Sha256Schema, + size: z.number().int().nonnegative(), + mime: MimeSchema, +}); +export type BlobRef = z.infer; + +// --- error union ------------------------------------------------------------ + +export const BlobStoreErrorCodeSchema = z.enum([ + 'BLOB_FETCH_NETWORK', + 'BLOB_FETCH_HTTP_ERROR', + 'BLOB_WRITE_FAILED', + 'BLOB_READ_FAILED', + 'BLOB_INVALID_INPUT', +]); +export type BlobStoreErrorCode = z.infer; + +export const BlobStoreErrorSchema = z.object({ + code: BlobStoreErrorCodeSchema, + message: z.string(), + suggestion: z.string().optional(), + recoverable: z.boolean(), +}); +export type BlobStoreError = z.infer; + +const FailureSchema = z.object({ + success: z.literal(false), + error: BlobStoreErrorSchema, +}); + +// --- per-method result schemas --------------------------------------------- + +export const PutSuccessSchema = z.object({ + success: z.literal(true), + data: BlobRefSchema, +}); +export const PutResultSchema = z.union([PutSuccessSchema, FailureSchema]); +export type PutResult = z.infer; + +export const HasSuccessSchema = z.object({ + success: z.literal(true), + data: z.boolean(), +}); +export const HasResultSchema = z.union([HasSuccessSchema, FailureSchema]); +export type HasResult = z.infer; + +export const GetSuccessSchema = z.object({ + success: z.literal(true), + // null = not found; success=true because "absent" is a valid answer, not an error + data: z.instanceof(Buffer).nullable(), +}); +export const GetResultSchema = z.union([GetSuccessSchema, FailureSchema]); +export type GetResult = z.infer; + +// --- capability ------------------------------------------------------------- + +export interface BlobStoreSpec { + /** Hash, dedupe, and persist a buffer. Returns a content-addressed BlobRef. */ + put(buffer: Buffer, mime: string): Promise; + + /** Fetch a URL (following redirects) and persist the bytes. */ + fetch(url: string, mimeHint?: string): Promise; + + /** Cheap existence check by sha256. */ + has(sha256: string): Promise; + + /** Read bytes by sha256. data === null when absent. */ + get(sha256: string): Promise; +} diff --git a/src/lib/log/capture/handler.test.ts b/src/lib/log/capture/handler.test.ts new file mode 100644 index 0000000..3d08a14 --- /dev/null +++ b/src/lib/log/capture/handler.test.ts @@ -0,0 +1,337 @@ +import { describe, test, expect, beforeEach } from 'bun:test'; +import { readFile } from 'node:fs/promises'; +import { existsSync } from 'node:fs'; +import { join } from 'node:path'; +import type { BlobRef, BlobStoreSpec, GetResult, HasResult, PutResult } from '../blob-store/spec.js'; +import type { AppendResult } from '../append.js'; +import type { AppendFn, Event } from './spec.js'; +import { CaptureHandler } from './handler.js'; + +// --- in-memory test doubles ------------------------------------------------- + +class FakeBlobs implements BlobStoreSpec { + store = new Map(); // sha → bytes + putCount = 0; + fetchCount = 0; + /** url → response. If the response has `.fail`, fetch returns failure. */ + urls = new Map(); + + async put(buffer: Buffer, mime: string): Promise { + this.putCount++; + const sha = await sha256(buffer); + this.store.set(sha, buffer); + return { success: true, data: { sha256: sha, size: buffer.length, mime } }; + } + async fetch(url: string, mimeHint?: string): Promise { + this.fetchCount++; + const r = this.urls.get(url); + if (!r) return { success: false, error: { code: 'BLOB_FETCH_HTTP_ERROR', message: `no mock for ${url}`, recoverable: false } }; + if (r.fail) return { success: false, error: { code: r.fail.code, message: `mocked ${r.fail.status ?? ''}`, recoverable: false } }; + const buf = r.body ?? Buffer.from(''); + return this.put(buf, r.mime ?? mimeHint ?? 'application/octet-stream'); + } + async has(sha256: string): Promise { return { success: true, data: this.store.has(sha256) }; } + async get(sha256: string): Promise { return { success: true, data: this.store.get(sha256) ?? null }; } +} + +async function sha256(buf: Buffer): Promise { + const { createHash } = await import('node:crypto'); + return createHash('sha256').update(buf).digest('hex'); +} + +function makeAppend(): { fn: AppendFn; events: Event[] } { + const events: Event[] = []; + const fn: AppendFn = async (e) => { events.push(e); return { success: true } as AppendResult; }; + return { fn, events }; +} + +const PROBE = '.stitch-mcp/log-probe'; +async function probe(name: string): Promise<{ result: any; duration_ms: number }> { + const raw = JSON.parse(await readFile(join(PROBE, name), 'utf8')); + return { result: raw.result, duration_ms: raw.duration_ms }; +} + +// Probe artifacts are real Stitch responses captured by `bun scripts/log-probe.ts`. +// They live under .stitch-mcp/log-probe/ which is gitignored, so they're absent in +// CI. Suites that load them are skipped automatically when the directory is +// missing — run them locally after `bun scripts/log-probe.ts` populates fixtures. +// Follow-up: commit a scrubbed subset to tests/fixtures/ so CI can run these too. +const PROBE_AVAILABLE = existsSync(PROBE); +const describeWithProbe = PROBE_AVAILABLE ? describe : describe.skip; + +// reusable counter-based newId +function makeIdGen() { + let i = 0; return () => `id${i++}`; +} + +let blobs: FakeBlobs; +let appendCtl: ReturnType; +let h: CaptureHandler; + +beforeEach(() => { + blobs = new FakeBlobs(); + appendCtl = makeAppend(); + h = new CaptureHandler({ blobs, append: appendCtl.fn, newId: makeIdGen() }); +}); + +// --- 1) generate_screen_from_text ------------------------------------------- + +describeWithProbe('generate_screen_from_text', () => { + test('writes call.requested + call.completed; produces 1 screen with all blobs filled', async () => { + const { result, duration_ms } = await probe('3-generate-screen.json'); + // mock the asset URLs that the screen references + const sc = result.structuredContent; + const screen = sc.outputComponents.find((c: any) => c.design)?.design.screens[0]; + blobs.urls.set(screen.htmlCode.downloadUrl, { mime: 'text/html', body: Buffer.from('real') }); + blobs.urls.set(screen.screenshot.downloadUrl, { mime: 'image/png', body: Buffer.from('PNG-bytes') }); + + const r = await h.capture({ + tool: 'generate_screen_from_text', + args: { projectId: '7240244230308968338', prompt: 'A minimal landing page for a typography-focused blog. One headline, one paragraph, one CTA.' }, + result, duration_ms, started_at: '2026-04-25T18:07:47Z', finished_at: '2026-04-25T18:09:12Z', + }); + + expect(r.success).toBe(true); + if (!r.success) return; + expect(r.data.produced_screen_ids).toEqual(['871df1dad8e54bb0ad9a3322ceee6260']); + expect(r.data.warnings).toEqual([]); + + expect(appendCtl.events.map((e) => e.type)).toEqual(['call.requested', 'call.completed']); + + const completed = appendCtl.events[1]!; + expect(completed.payload).toMatchObject({ tool: 'generate_screen_from_text', kind: 'generative' }); + if (completed.type !== 'call.completed' || completed.payload.kind !== 'generative') throw new Error(); + expect(completed.payload.stitch_session_id).toBe('6904021969822429940'); + expect(completed.payload.produced_screens).toHaveLength(1); + + const ps = completed.payload.produced_screens[0]!; + expect(ps.screen_id).toBe('871df1dad8e54bb0ad9a3322ceee6260'); + expect(ps.parent_screen_id).toBeNull(); + expect(ps.sibling_screen_ids).toEqual([]); + expect(ps.effective_prompt.startsWith('A minimal landing page for a typography-focused blog.')).toBe(true); + expect(ps.html_blob).not.toBeNull(); + expect(ps.screenshot_blob).not.toBeNull(); + expect(ps.theme_blob).not.toBeNull(); + expect(ps.design_system_blob).not.toBeNull(); + }); + + test('effective_prompt comes from screen.prompt, NOT args.prompt', async () => { + const { result, duration_ms } = await probe('3-generate-screen.json'); + const sc = result.structuredContent; + const screen = sc.outputComponents.find((c: any) => c.design)?.design.screens[0]; + blobs.urls.set(screen.htmlCode.downloadUrl, { mime: 'text/html', body: Buffer.from('') }); + blobs.urls.set(screen.screenshot.downloadUrl, { mime: 'image/png', body: Buffer.from('png') }); + + const userInput = 'Short prompt'; // deliberately different from screen.prompt + await h.capture({ + tool: 'generate_screen_from_text', + args: { projectId: '7240244230308968338', prompt: userInput }, + result, duration_ms, started_at: 't0', finished_at: 't1', + }); + + const requested = appendCtl.events[0]!; + if (requested.type !== 'call.requested') throw new Error(); + expect(requested.payload.user_prompt).toBe(userInput); + + const completed = appendCtl.events[1]!; + if (completed.type !== 'call.completed' || completed.payload.kind !== 'generative') throw new Error(); + const ep = completed.payload.produced_screens[0]!.effective_prompt; + expect(ep).not.toBe(userInput); + expect(ep.length).toBeGreaterThan(userInput.length); + }); +}); + +// --- 2) generate_variants ---------------------------------------------------- + +describeWithProbe('generate_variants', () => { + test('produces N screens with shared parent and sibling cross-refs', async () => { + const { result, duration_ms } = await probe('4-generate-variants.json'); + const screens = result.structuredContent.outputComponents.find((c: any) => c.design)?.design.screens; + for (const s of screens) { + blobs.urls.set(s.htmlCode.downloadUrl, { mime: 'text/html', body: Buffer.from('') }); + blobs.urls.set(s.screenshot.downloadUrl, { mime: 'image/png', body: Buffer.from('vpng') }); + } + + const parentId = '871df1dad8e54bb0ad9a3322ceee6260'; + const r = await h.capture({ + tool: 'generate_variants', + args: { + projectId: '7240244230308968338', selectedScreenIds: [parentId], + prompt: 'Vary it', variantOptions: { variantCount: 2, creativeRange: 'EXPLORE', aspects: ['COLOR_SCHEME', 'LAYOUT'] }, + }, + result, duration_ms, started_at: 't0', finished_at: 't1', + }); + expect(r.success).toBe(true); + if (!r.success) return; + expect(r.data.produced_screen_ids).toHaveLength(2); + + const completed = appendCtl.events[1]!; + if (completed.type !== 'call.completed' || completed.payload.kind !== 'generative') throw new Error(); + + for (const ps of completed.payload.produced_screens) { + expect(ps.parent_screen_id).toBe(parentId); + expect(ps.sibling_screen_ids).toHaveLength(1); + const otherId = completed.payload.produced_screens.find((p) => p.screen_id !== ps.screen_id)!.screen_id; + expect(ps.sibling_screen_ids).toEqual([otherId]); + // variants come back with empty theme → theme_blob null is expected + expect(ps.theme_blob).toBeNull(); + } + }); +}); + +// --- 3) edit_screens --------------------------------------------------------- + +describeWithProbe('edit_screens', () => { + test('produces 1 child screen with parent set and theme/designSystem captured', async () => { + const { result, duration_ms } = await probe('5-edit-screens.json'); + const screen = result.structuredContent.outputComponents.find((c: any) => c.design)?.design.screens[0]; + blobs.urls.set(screen.htmlCode.downloadUrl, { mime: 'text/html', body: Buffer.from('') }); + blobs.urls.set(screen.screenshot.downloadUrl, { mime: 'image/png', body: Buffer.from('epng') }); + + const parentId = '871df1dad8e54bb0ad9a3322ceee6260'; + const r = await h.capture({ + tool: 'edit_screens', + args: { projectId: '7240244230308968338', selectedScreenIds: [parentId], prompt: 'Make it bigger' }, + result, duration_ms, started_at: 't0', finished_at: 't1', + }); + expect(r.success).toBe(true); + const completed = appendCtl.events[1]!; + if (completed.type !== 'call.completed' || completed.payload.kind !== 'generative') throw new Error(); + expect(completed.payload.produced_screens).toHaveLength(1); + const ps = completed.payload.produced_screens[0]!; + expect(ps.parent_screen_id).toBe(parentId); + expect(ps.sibling_screen_ids).toEqual([]); + expect(ps.theme_blob).not.toBeNull(); // edit returns full theme + }); +}); + +// --- 4) read calls ----------------------------------------------------------- + +describeWithProbe('read calls', () => { + test('get_screen logs read summary with project_id+screen_ids+result_blob; NO blob fetches', async () => { + const { result, duration_ms } = await probe('6-get-screen.json'); + const r = await h.capture({ + tool: 'get_screen', + args: { projectId: '7240244230308968338', screenId: '871df1dad8e54bb0ad9a3322ceee6260' }, + result, duration_ms, started_at: 't0', finished_at: 't1', + }); + expect(r.success).toBe(true); + expect(blobs.fetchCount).toBe(0); // no eager downloads on read + + const completed = appendCtl.events[1]!; + if (completed.type !== 'call.completed' || completed.payload.kind !== 'read') throw new Error(); + expect(completed.payload.project_id).toBe('7240244230308968338'); + expect(completed.payload.screen_ids).toEqual(['871df1dad8e54bb0ad9a3322ceee6260']); + expect(completed.payload.result_blob).toBeDefined(); + expect(completed.payload.result_blob.mime).toBe('application/json'); + }); + + test('list_projects logs as read; result is blobbed and returned_project_ids extracted', async () => { + const { result, duration_ms } = await probe('1-list-projects.json'); + const beforePuts = blobs.putCount; + const r = await h.capture({ + tool: 'list_projects', args: {}, result, duration_ms, started_at: 't0', finished_at: 't1', + }); + expect(r.success).toBe(true); + // args_blob + result_blob now (response IS captured for replay) + expect(blobs.putCount - beforePuts).toBe(2); + expect(blobs.fetchCount).toBe(0); + + const completed = appendCtl.events[1]!; + if (completed.type !== 'call.completed' || completed.payload.kind !== 'read') throw new Error(); + expect(completed.payload.tool).toBe('list_projects'); + expect(completed.payload.result_blob).toBeDefined(); + // list_projects response is full of name="projects/" entries + expect((completed.payload.returned_project_ids ?? []).length).toBeGreaterThan(0); + }); +}); + +// --- 5) failures ------------------------------------------------------------ + +describeWithProbe('failures', () => { + test('explicit isError → call.failed with error_text', async () => { + const { result, duration_ms } = await probe('8-fail-bogus-screen.json'); + const r = await h.capture({ + tool: 'get_screen', + args: { projectId: '0', screenId: 'doesnotexist' }, + result, duration_ms, started_at: 't0', finished_at: 't1', + }); + expect(r.success).toBe(true); + expect(appendCtl.events.map((e) => e.type)).toEqual(['call.requested', 'call.failed']); + const failed = appendCtl.events[1]!; + if (failed.type !== 'call.failed') throw new Error(); + expect(failed.payload.is_error).toBe(true); + expect(failed.payload.error_text).toContain('not found'); + }); + + test('zero produced screens on a generative call → call.failed with is_error="empty"', async () => { + const { result, duration_ms } = await probe('9-fail-missing-required.json'); + const r = await h.capture({ + tool: 'edit_screens', + args: { projectId: '7240244230308968338' }, // missing selectedScreenIds + prompt — server returns no-op + result, duration_ms, started_at: 't0', finished_at: 't1', + }); + expect(r.success).toBe(true); + const failed = appendCtl.events[1]!; + if (failed.type !== 'call.failed') throw new Error(); + expect(failed.payload.is_error).toBe('empty'); + }); +}); + +// --- 6) robustness ---------------------------------------------------------- + +describeWithProbe('robustness', () => { + test('a single failed asset fetch does NOT abort capture; warning recorded', async () => { + const { result, duration_ms } = await probe('3-generate-screen.json'); + const sc = result.structuredContent; + const screen = sc.outputComponents.find((c: any) => c.design)?.design.screens[0]; + // HTML succeeds, screenshot URL is mocked to FAIL + blobs.urls.set(screen.htmlCode.downloadUrl, { mime: 'text/html', body: Buffer.from('') }); + blobs.urls.set(screen.screenshot.downloadUrl, { fail: { code: 'BLOB_FETCH_HTTP_ERROR', status: 410 } }); + + const r = await h.capture({ + tool: 'generate_screen_from_text', + args: { projectId: '7240244230308968338', prompt: 'p' }, + result, duration_ms, started_at: 't0', finished_at: 't1', + }); + expect(r.success).toBe(true); + if (!r.success) return; + expect(r.data.produced_screen_ids).toHaveLength(1); + expect(r.data.warnings.length).toBeGreaterThanOrEqual(1); + expect(r.data.warnings[0]).toContain('screenshot'); + + const completed = appendCtl.events[1]!; + if (completed.type !== 'call.completed' || completed.payload.kind !== 'generative') throw new Error(); + const ps = completed.payload.produced_screens[0]!; + expect(ps.html_blob).not.toBeNull(); + expect(ps.screenshot_blob).toBeNull(); // failed; null preserved + }); +}); + +// Inline-only — runs in CI without probe fixtures. +describe('unknown tools', () => { + test('unknown tool is logged with kind="unknown" and a result_blob', async () => { + const r = await h.capture({ + tool: 'list_design_systems', + args: { projectId: 'p1' }, + result: { structuredContent: { designSystems: [] } }, + duration_ms: 5, + started_at: 't0', + finished_at: 't1', + }); + expect(r.success).toBe(true); + if (!r.success) return; + + expect(appendCtl.events).toHaveLength(2); + const requested = appendCtl.events[0]!; + if (requested.type !== 'call.requested') throw new Error(); + expect(requested.payload.tool).toBe('list_design_systems'); + + const completed = appendCtl.events[1]!; + if (completed.type !== 'call.completed' || completed.payload.kind !== 'unknown') throw new Error(); + expect(completed.payload.tool).toBe('list_design_systems'); + expect(completed.payload.project_id).toBe('p1'); + expect(completed.payload.result_blob).toBeDefined(); + expect(completed.payload.result_blob.mime).toBe('application/json'); + }); +}); diff --git a/src/lib/log/capture/handler.ts b/src/lib/log/capture/handler.ts new file mode 100644 index 0000000..5881fef --- /dev/null +++ b/src/lib/log/capture/handler.ts @@ -0,0 +1,328 @@ +import { randomUUID } from 'node:crypto'; +import type { BlobRef } from '../blob-store/spec.js'; +import { + type AppendFn, + type CaptureDeps, + type CaptureInput, + type CaptureResult, + type CaptureSpec, + type Event, + type ProducedScreen, + type ToolKind, + CaptureInputSchema, + kindOf, +} from './spec.js'; + +interface ScreenComponent { + id?: string; + name?: string; + prompt?: string; + theme?: Record; + designSystem?: Record | null; + htmlCode?: { downloadUrl?: string; mimeType?: string }; + screenshot?: { downloadUrl?: string }; +} + +export class CaptureHandler implements CaptureSpec { + private readonly blobs: CaptureDeps['blobs']; + private readonly append: AppendFn; + private readonly now: () => Date; + private readonly newId: () => string; + + constructor(deps: CaptureDeps) { + this.blobs = deps.blobs; + this.append = deps.append; + this.now = deps.now ?? (() => new Date()); + this.newId = deps.newId ?? (() => randomUUID()); + } + + async capture(input: CaptureInput): Promise { + const parsed = CaptureInputSchema.safeParse(input); + if (!parsed.success) { + return this.fail('CAPTURE_INVALID_INPUT', parsed.error.message, false); + } + const kind: ToolKind = kindOf(input.tool); + + const trace_id = this.newId(); + const warnings: string[] = []; + + // 1) requested event — always + const argsBuf = Buffer.from(JSON.stringify(input.args)); + const argsBlob = await this.blobs.put(argsBuf, 'application/json'); + if (!argsBlob.success) { + return this.fail('CAPTURE_BLOB_FATAL', `args_blob: ${argsBlob.error.message}`, false); + } + + const requested: Event = { + id: this.newId(), + time: input.started_at, + trace_id, + schema_version: 1, + type: 'call.requested', + payload: { + tool: input.tool, + project_id: typeof input.args.projectId === 'string' ? input.args.projectId : undefined, + selected_screen_ids: Array.isArray(input.args.selectedScreenIds) + ? (input.args.selectedScreenIds as string[]) + : undefined, + user_prompt: typeof input.args.prompt === 'string' ? input.args.prompt : undefined, + variant_options: + input.args.variantOptions && typeof input.args.variantOptions === 'object' + ? (input.args.variantOptions as Record) + : undefined, + device_type: typeof input.args.deviceType === 'string' ? input.args.deviceType : undefined, + model_id: typeof input.args.modelId === 'string' ? input.args.modelId : undefined, + args_blob: argsBlob.data, + }, + }; + const ar = await this.append(requested); + if (!ar.success) return this.fail('CAPTURE_APPEND_FAILED', ar.error.message, true); + + // 2) explicit failure + const r = input.result as { isError?: boolean; structuredContent?: any; content?: { type: string; text?: string }[] } | null; + if (r && r.isError === true) { + const errorText = r.content?.find((c) => c.type === 'text')?.text ?? ''; + const rawBlob = await this.blobs.put(Buffer.from(JSON.stringify(r)), 'application/json'); + const failed: Event = { + id: this.newId(), + time: input.finished_at, + trace_id, + schema_version: 1, + type: 'call.failed', + payload: { + tool: input.tool, + duration_ms: input.duration_ms, + is_error: true, + error_text: errorText, + raw_blob: rawBlob.success ? rawBlob.data : undefined, + }, + }; + const fr = await this.append(failed); + if (!fr.success) return this.fail('CAPTURE_APPEND_FAILED', fr.error.message, true); + return { success: true, data: { trace_id, produced_screen_ids: [], warnings } }; + } + + if (kind === 'read') { + const resultBlob = await this.blobs.put(Buffer.from(JSON.stringify(r ?? {})), 'application/json'); + if (!resultBlob.success) { + return this.fail('CAPTURE_BLOB_FATAL', `result_blob: ${resultBlob.error.message}`, false); + } + const returned = extractReturnedIds(r); + const completed: Event = { + id: this.newId(), + time: input.finished_at, + trace_id, + schema_version: 1, + type: 'call.completed', + payload: { + tool: input.tool, + duration_ms: input.duration_ms, + kind: 'read', + project_id: typeof input.args.projectId === 'string' ? input.args.projectId : undefined, + screen_ids: + typeof input.args.screenId === 'string' + ? [input.args.screenId] + : Array.isArray(input.args.selectedScreenIds) + ? (input.args.selectedScreenIds as string[]) + : undefined, + returned_project_ids: returned.projects.length > 0 ? returned.projects : undefined, + returned_screen_ids: returned.screens.length > 0 ? returned.screens : undefined, + result_blob: resultBlob.data, + }, + }; + const cr = await this.append(completed); + if (!cr.success) return this.fail('CAPTURE_APPEND_FAILED', cr.error.message, true); + return { success: true, data: { trace_id, produced_screen_ids: [], warnings } }; + } + + if (kind === 'unknown') { + const resultBlob = await this.blobs.put(Buffer.from(JSON.stringify(r ?? {})), 'application/json'); + if (!resultBlob.success) { + return this.fail('CAPTURE_BLOB_FATAL', `result_blob: ${resultBlob.error.message}`, false); + } + const completed: Event = { + id: this.newId(), + time: input.finished_at, + trace_id, + schema_version: 1, + type: 'call.completed', + payload: { + tool: input.tool, + duration_ms: input.duration_ms, + kind: 'unknown', + project_id: typeof input.args.projectId === 'string' ? input.args.projectId : undefined, + result_blob: resultBlob.data, + }, + }; + const cr = await this.append(completed); + if (!cr.success) return this.fail('CAPTURE_APPEND_FAILED', cr.error.message, true); + return { success: true, data: { trace_id, produced_screen_ids: [], warnings } }; + } + + // 3) generative path + const sc = (r as any)?.structuredContent ?? null; + const screens: ScreenComponent[] = pickScreens(sc); + + // implicit failure: zero produced screens on a generative call + if (screens.length === 0) { + const rawBlob = await this.blobs.put(Buffer.from(JSON.stringify(r ?? {})), 'application/json'); + const failed: Event = { + id: this.newId(), + time: input.finished_at, + trace_id, + schema_version: 1, + type: 'call.failed', + payload: { + tool: input.tool, + duration_ms: input.duration_ms, + is_error: 'empty', + raw_blob: rawBlob.success ? rawBlob.data : undefined, + }, + }; + const fr = await this.append(failed); + if (!fr.success) return this.fail('CAPTURE_APPEND_FAILED', fr.error.message, true); + return { success: true, data: { trace_id, produced_screen_ids: [], warnings } }; + } + + const structuredBlob = await this.blobs.put(Buffer.from(JSON.stringify(sc)), 'application/json'); + if (!structuredBlob.success) { + return this.fail('CAPTURE_BLOB_FATAL', `structured_content_blob: ${structuredBlob.error.message}`, false); + } + + // top-level designSystem component (only on generate_screen_from_text in practice) + const dsAsset = pickDesignSystemComponent(sc); + let dsAssetBlob: BlobRef | null = null; + if (dsAsset) { + const r2 = await this.blobs.put(Buffer.from(JSON.stringify(dsAsset)), 'application/json'); + if (r2.success) dsAssetBlob = r2.data; + else warnings.push(`design_system_asset put failed: ${r2.error.message}`); + } + + const allScreenIds = screens.map((s) => s.id ?? '').filter(Boolean); + const selectedParents = (input.args.selectedScreenIds as string[] | undefined) ?? []; + const parent = selectedParents[0] ?? null; + + const produced: ProducedScreen[] = []; + for (const s of screens) { + const screenId = s.id ?? ''; + const siblings = allScreenIds.filter((id) => id !== screenId); + + // theme blob (populated for generate/edit; empty {} for variants) + let themeBlob: BlobRef | null = null; + if (s.theme && Object.keys(s.theme).length > 0) { + const tr = await this.blobs.put(Buffer.from(JSON.stringify(s.theme)), 'application/json'); + if (tr.success) themeBlob = tr.data; + else warnings.push(`screen ${screenId} theme: ${tr.error.message}`); + } + + // per-screen design system (generate only) + let dsBlob: BlobRef | null = null; + if (s.designSystem) { + const dr = await this.blobs.put(Buffer.from(JSON.stringify(s.designSystem)), 'application/json'); + if (dr.success) dsBlob = dr.data; + else warnings.push(`screen ${screenId} design_system: ${dr.error.message}`); + } else if (dsAssetBlob) { + // fall back to the top-level designSystem component blob when per-screen is absent + dsBlob = dsAssetBlob; + } + + // eager downloads + let htmlBlob: BlobRef | null = null; + if (s.htmlCode?.downloadUrl) { + const fr = await this.blobs.fetch(s.htmlCode.downloadUrl, s.htmlCode.mimeType ?? 'text/html'); + if (fr.success) htmlBlob = fr.data; + else warnings.push(`screen ${screenId} html: ${fr.error.code} ${fr.error.message}`); + } + let shotBlob: BlobRef | null = null; + if (s.screenshot?.downloadUrl) { + const fr = await this.blobs.fetch(s.screenshot.downloadUrl); + if (fr.success) shotBlob = fr.data; + else warnings.push(`screen ${screenId} screenshot: ${fr.error.code} ${fr.error.message}`); + } + + produced.push({ + project_id: (input.args.projectId as string) ?? sc?.projectId ?? '', + screen_id: screenId, + name: s.name ?? '', + parent_screen_id: parent, + sibling_screen_ids: siblings, + effective_prompt: s.prompt ?? '', + html_blob: htmlBlob, + screenshot_blob: shotBlob, + theme_blob: themeBlob, + design_system_blob: dsBlob, + }); + } + + const completed: Event = { + id: this.newId(), + time: input.finished_at, + trace_id, + schema_version: 1, + type: 'call.completed', + payload: { + tool: input.tool, + duration_ms: input.duration_ms, + kind: 'generative', + stitch_session_id: typeof sc?.sessionId !== 'undefined' ? String(sc.sessionId) : undefined, + structured_content_blob: structuredBlob.data, + produced_screens: produced, + }, + }; + const cr = await this.append(completed); + if (!cr.success) return this.fail('CAPTURE_APPEND_FAILED', cr.error.message, true); + + return { + success: true, + data: { trace_id, produced_screen_ids: produced.map((p) => p.screen_id), warnings }, + }; + } + + private fail(code: any, message: string, recoverable: boolean): CaptureResult { + return { success: false, error: { code, message, recoverable } }; + } +} + +// --- helpers (extraction over heterogeneous outputComponents) --------------- + +function pickScreens(sc: any): ScreenComponent[] { + const out: ScreenComponent[] = []; + for (const c of sc?.outputComponents ?? []) { + if (c?.design?.screens) for (const s of c.design.screens) out.push(s as ScreenComponent); + } + return out; +} + +function pickDesignSystemComponent(sc: any): Record | null { + const c = (sc?.outputComponents ?? []).find((x: any) => x?.designSystem); + return c?.designSystem ?? null; +} + +/** + * Walk an MCP result and pull out any `name` fields shaped like + * "projects/" or "screens/". Used to record which entities a + * read-tool response actually returned so the log is queryable without + * re-reading the result blob. + */ +function extractReturnedIds(result: unknown): { projects: string[]; screens: string[] } { + const projects = new Set(); + const screens = new Set(); + const seen = new WeakSet(); + const visit = (node: unknown) => { + if (!node || typeof node !== 'object') return; + if (seen.has(node as object)) return; + seen.add(node as object); + const name = (node as { name?: unknown }).name; + if (typeof name === 'string') { + if (name.startsWith('projects/')) projects.add(name.slice('projects/'.length)); + else if (name.startsWith('screens/')) screens.add(name.slice('screens/'.length)); + } + if (Array.isArray(node)) { + for (const v of node) visit(v); + } else { + for (const v of Object.values(node as Record)) visit(v); + } + }; + visit(result); + return { projects: Array.from(projects), screens: Array.from(screens) }; +} diff --git a/src/lib/log/capture/spec.test.ts b/src/lib/log/capture/spec.test.ts new file mode 100644 index 0000000..9fb9cd1 --- /dev/null +++ b/src/lib/log/capture/spec.test.ts @@ -0,0 +1,135 @@ +import { describe, test, expect } from 'bun:test'; +import { + CaptureInputSchema, + CompletedPayloadSchema, + EventSchema, + FailedPayloadSchema, + ProducedScreenSchema, + RequestedPayloadSchema, + kindOf, +} from './spec.js'; + +const validBlob = { sha256: 'a'.repeat(64), size: 10, mime: 'application/json' }; + +describe('kindOf()', () => { + test.each([ + ['generate_screen_from_text', 'generative'], + ['edit_screens', 'generative'], + ['generate_variants', 'generative'], + ['get_screen', 'read'], + ['list_screens', 'read'], + ['list_projects', 'read'], + ['get_project', 'read'], + ['create_project', 'read'], + ['unknown_tool', 'unknown'], + ['list_design_systems', 'unknown'], + ] as const)('kindOf(%s) === %s', (tool, expected) => { + expect(kindOf(tool)).toBe(expected); + }); +}); + +describe('ProducedScreenSchema', () => { + test('accepts a minimal screen with all blobs null', () => { + const r = ProducedScreenSchema.safeParse({ + project_id: 'p', screen_id: 's', name: 'projects/p/screens/s', + parent_screen_id: null, sibling_screen_ids: [], + effective_prompt: '', html_blob: null, screenshot_blob: null, + theme_blob: null, design_system_blob: null, + }); + expect(r.success).toBe(true); + }); + + test('rejects when sibling_screen_ids missing', () => { + const r = ProducedScreenSchema.safeParse({ + project_id: 'p', screen_id: 's', name: 'n', + parent_screen_id: null, effective_prompt: '', + html_blob: null, screenshot_blob: null, theme_blob: null, design_system_blob: null, + }); + expect(r.success).toBe(false); + }); +}); + +describe('Payload schemas', () => { + test('Requested requires args_blob', () => { + expect(RequestedPayloadSchema.safeParse({ tool: 'list_projects' }).success).toBe(false); + expect(RequestedPayloadSchema.safeParse({ tool: 'list_projects', args_blob: validBlob }).success).toBe(true); + }); + + test('Completed-generative requires produced_screens + structured_content_blob', () => { + const r = CompletedPayloadSchema.safeParse({ + tool: 'generate_screen_from_text', duration_ms: 100, kind: 'generative', + structured_content_blob: validBlob, produced_screens: [], + }); + expect(r.success).toBe(true); + }); + + test('Completed-read requires result_blob', () => { + const missing = CompletedPayloadSchema.safeParse({ + tool: 'get_screen', duration_ms: 1, kind: 'read', + project_id: 'p', screen_ids: ['s'], + }); + expect(missing.success).toBe(false); + + const ok = CompletedPayloadSchema.safeParse({ + tool: 'get_screen', duration_ms: 1, kind: 'read', + project_id: 'p', screen_ids: ['s'], result_blob: validBlob, + }); + expect(ok.success).toBe(true); + if (ok.success) expect((ok.data as any).kind).toBe('read'); + }); + + test('Completed-unknown requires result_blob', () => { + const missing = CompletedPayloadSchema.safeParse({ + tool: 'list_design_systems', duration_ms: 1, kind: 'unknown', + }); + expect(missing.success).toBe(false); + + const ok = CompletedPayloadSchema.safeParse({ + tool: 'list_design_systems', duration_ms: 1, kind: 'unknown', + result_blob: validBlob, + }); + expect(ok.success).toBe(true); + if (ok.success) expect((ok.data as any).kind).toBe('unknown'); + }); + + test('Failed accepts is_error true | "empty"', () => { + expect(FailedPayloadSchema.safeParse({ tool: 'x', duration_ms: 1, is_error: true }).success).toBe(true); + expect(FailedPayloadSchema.safeParse({ tool: 'x', duration_ms: 1, is_error: 'empty' }).success).toBe(true); + expect(FailedPayloadSchema.safeParse({ tool: 'x', duration_ms: 1, is_error: false }).success).toBe(false); + }); +}); + +describe('EventSchema discrimination', () => { + test('routes by type', () => { + const ok = EventSchema.safeParse({ + id: 'a', time: 't', trace_id: 'r', schema_version: 1, + type: 'call.requested', payload: { tool: 'list_projects', args_blob: validBlob }, + }); + expect(ok.success).toBe(true); + }); + + test('rejects unknown type', () => { + const r = EventSchema.safeParse({ + id: 'a', time: 't', trace_id: 'r', schema_version: 1, + type: 'observation', payload: {}, + }); + expect(r.success).toBe(false); + }); +}); + +describe('CaptureInputSchema', () => { + test('accepts minimal valid input', () => { + const r = CaptureInputSchema.safeParse({ + tool: 'list_projects', args: {}, result: {}, duration_ms: 0, + started_at: 't0', finished_at: 't1', + }); + expect(r.success).toBe(true); + }); + + test('rejects negative duration', () => { + const r = CaptureInputSchema.safeParse({ + tool: 'x', args: {}, result: {}, duration_ms: -1, started_at: 't', finished_at: 't', + }); + expect(r.success).toBe(false); + }); +}); diff --git a/src/lib/log/capture/spec.ts b/src/lib/log/capture/spec.ts new file mode 100644 index 0000000..d37b3ef --- /dev/null +++ b/src/lib/log/capture/spec.ts @@ -0,0 +1,170 @@ +import { z } from 'zod'; +import { BlobRefSchema } from '../blob-store/spec.js'; + +// --- tool taxonomy ---------------------------------------------------------- + +export const GENERATIVE_TOOLS = new Set([ + 'generate_screen_from_text', + 'edit_screens', + 'generate_variants', +] as const); + +export const READ_TOOLS = new Set([ + 'get_screen', + 'list_screens', + 'list_projects', + 'get_project', + 'create_project', +] as const); + +export type ToolKind = 'generative' | 'read' | 'unknown'; + +export function kindOf(tool: string): ToolKind { + if (GENERATIVE_TOOLS.has(tool as any)) return 'generative'; + if (READ_TOOLS.has(tool as any)) return 'read'; + return 'unknown'; +} + +// --- per-screen artifact ----------------------------------------------------- + +export const ProducedScreenSchema = z.object({ + project_id: z.string(), + screen_id: z.string(), + name: z.string(), + parent_screen_id: z.string().nullable(), + sibling_screen_ids: z.array(z.string()), + effective_prompt: z.string(), + // null when a fetch failed (warning recorded) or absent on the response (variants) + html_blob: BlobRefSchema.nullable(), + screenshot_blob: BlobRefSchema.nullable(), + theme_blob: BlobRefSchema.nullable(), + design_system_blob: BlobRefSchema.nullable(), +}); +export type ProducedScreen = z.infer; + +// --- event payloads --------------------------------------------------------- + +export const RequestedPayloadSchema = z.object({ + tool: z.string(), + project_id: z.string().optional(), + selected_screen_ids: z.array(z.string()).optional(), + user_prompt: z.string().optional(), + variant_options: z.record(z.string(), z.unknown()).optional(), + device_type: z.string().optional(), + model_id: z.string().optional(), + args_blob: BlobRefSchema, +}); + +export const CompletedGenerativePayloadSchema = z.object({ + tool: z.string(), + duration_ms: z.number().int().nonnegative(), + kind: z.literal('generative'), + stitch_session_id: z.string().optional(), + structured_content_blob: BlobRefSchema, + produced_screens: z.array(ProducedScreenSchema), +}); + +export const CompletedReadPayloadSchema = z.object({ + tool: z.string(), + duration_ms: z.number().int().nonnegative(), + kind: z.literal('read'), + project_id: z.string().optional(), + screen_ids: z.array(z.string()).optional(), + returned_project_ids: z.array(z.string()).optional(), + returned_screen_ids: z.array(z.string()).optional(), + result_blob: BlobRefSchema, +}); + +export const CompletedUnknownPayloadSchema = z.object({ + tool: z.string(), + duration_ms: z.number().int().nonnegative(), + kind: z.literal('unknown'), + project_id: z.string().optional(), + result_blob: BlobRefSchema, +}); + +export const CompletedPayloadSchema = z.discriminatedUnion('kind', [ + CompletedGenerativePayloadSchema, + CompletedReadPayloadSchema, + CompletedUnknownPayloadSchema, +]); + +export const FailedPayloadSchema = z.object({ + tool: z.string(), + duration_ms: z.number().int().nonnegative(), + is_error: z.union([z.literal(true), z.literal('empty')]), + error_text: z.string().optional(), + raw_blob: BlobRefSchema.optional(), +}); + +// --- envelope + event union -------------------------------------------------- + +const baseEnvelope = { + id: z.string().min(1), + time: z.string().min(1), + trace_id: z.string().min(1), + schema_version: z.literal(1), +}; + +export const EventSchema = z.discriminatedUnion('type', [ + z.object({ ...baseEnvelope, type: z.literal('call.requested'), payload: RequestedPayloadSchema }), + z.object({ ...baseEnvelope, type: z.literal('call.completed'), payload: CompletedPayloadSchema }), + z.object({ ...baseEnvelope, type: z.literal('call.failed'), payload: FailedPayloadSchema }), +]); +export type Event = z.infer; + +// --- capture I/O ------------------------------------------------------------ + +export const CaptureInputSchema = z.object({ + tool: z.string().min(1), + args: z.record(z.string(), z.unknown()), + result: z.unknown(), // raw MCP CallToolResult + duration_ms: z.number().int().nonnegative(), + started_at: z.string().min(1), + finished_at: z.string().min(1), +}); +export type CaptureInput = z.infer; + +export const CaptureErrorCodeSchema = z.enum([ + 'CAPTURE_UNKNOWN_TOOL', // tool isn't in either taxonomy set + 'CAPTURE_APPEND_FAILED', // appendEvent rejected + 'CAPTURE_BLOB_FATAL', // critical blob (args/result) failed; can't proceed + 'CAPTURE_INVALID_INPUT', +]); + +const CaptureFailure = z.object({ + success: z.literal(false), + error: z.object({ + code: CaptureErrorCodeSchema, + message: z.string(), + recoverable: z.boolean(), + }), +}); +const CaptureSuccess = z.object({ + success: z.literal(true), + data: z.object({ + trace_id: z.string(), + produced_screen_ids: z.array(z.string()), + warnings: z.array(z.string()), // soft-failures (one blob fetch died, etc.) + }), +}); +export const CaptureResultSchema = z.union([CaptureSuccess, CaptureFailure]); +export type CaptureResult = z.infer; + +// --- dependency contracts (for testability) --------------------------------- + +import type { BlobStoreSpec } from '../blob-store/spec.js'; +import type { AppendResult } from '../append.js'; + +export type AppendFn = (event: Event) => Promise; + +export interface CaptureSpec { + capture(input: CaptureInput): Promise; +} + +export interface CaptureDeps { + blobs: BlobStoreSpec; + append: AppendFn; + now?: () => Date; // injectable for deterministic tests + newId?: () => string; +} diff --git a/src/lib/log/factory.ts b/src/lib/log/factory.ts new file mode 100644 index 0000000..7c7051c --- /dev/null +++ b/src/lib/log/factory.ts @@ -0,0 +1,20 @@ +import { join } from 'node:path'; +import { appendEvent } from './append.js'; +import { BlobStoreHandler } from './blob-store/handler.js'; +import { CaptureHandler } from './capture/handler.js'; +import type { CaptureSpec } from './capture/spec.js'; + +export const DEFAULT_LOG_ROOT = '.stitch-mcp/log'; + +export function isLogEnabled(): boolean { + return process.env.STITCH_MCP_LOG === '1'; +} + +export function createCaptureHandler(root: string = DEFAULT_LOG_ROOT): CaptureSpec { + const blobs = new BlobStoreHandler(join(root, 'blobs')); + const eventsPath = join(root, 'events.jsonl'); + return new CaptureHandler({ + blobs, + append: (event) => appendEvent(eventsPath, event), + }); +}