From 6b764639414fec462afad661f7d9972ad956fbf0 Mon Sep 17 00:00:00 2001 From: can1357 Date: Wed, 10 Jun 2026 10:41:04 +0200 Subject: [PATCH 01/27] fix: fixed stats dashboard asset serving for bundled CLI distributions - Fixed npm CLI bundling by invoking `generate-client-bundle.ts --generate` before build and `--reset` after. - Added `PI_BUNDLED` handling in `packages/stats/src/server.ts` and switched bundled distributions to embedded assets. - Decoded embedded archives via `decodeEmbeddedClientArchive` and treated legacy placeholder content as missing. - Handled missing dashboard source/build directories as ENOENT during mtime scans to avoid startup crashes. - Extended `--smoke-test` to start the stats server and validate dashboard HTML contains expected asset markers. --- packages/coding-agent/CHANGELOG.md | 7 ++- packages/coding-agent/scripts/bundle-dist.ts | 47 +++++++++------ packages/coding-agent/src/cli.ts | 26 ++++++-- .../test/stats-dashboard-bundle.test.ts | 30 ++++++++++ packages/stats/CHANGELOG.md | 8 +++ .../stats/scripts/generate-client-bundle.ts | 15 ++--- .../stats/src/embedded-client.generated.txt | 7 --- packages/stats/src/embedded-client.ts | 26 ++++++++ packages/stats/src/server.ts | 59 ++++++++++++------- scripts/install-tests/run-ci.sh | 9 +-- 10 files changed, 162 insertions(+), 72 deletions(-) create mode 100644 packages/coding-agent/test/stats-dashboard-bundle.test.ts create mode 100644 packages/stats/src/embedded-client.ts diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index f488065de..d25ba3994 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -1,6 +1,11 @@ # Changelog ## [Unreleased] +### Fixed + +- Fixed npm CLI distribution bundles by embedding the stats dashboard client bundle so dashboard assets are served in prebuilt installs +- Fixed the CLI smoke-test command to start the stats server and verify dashboard HTML is served, catching bundled-asset regressions +- Added verification of a `
` and `index.js` in smoke-test dashboard responses ## [15.10.12] - 2026-06-10 @@ -9971,4 +9976,4 @@ Initial public release. - Git branch display in footer - Message queueing during streaming responses - OAuth integration for Gmail and Google Calendar access -- HTML export with syntax highlighting and collapsible sections +- HTML export with syntax highlighting and collapsible sections \ No newline at end of file diff --git a/packages/coding-agent/scripts/bundle-dist.ts b/packages/coding-agent/scripts/bundle-dist.ts index ad835aeb9..df754cf57 100755 --- a/packages/coding-agent/scripts/bundle-dist.ts +++ b/packages/coding-agent/scripts/bundle-dist.ts @@ -51,25 +51,34 @@ async function cleanBundleOutputs(): Promise { async function main(): Promise { const start = Bun.nanoseconds(); await cleanBundleOutputs(); - await runCommand([ - "bun", - "build", - "--target=bun", - "--outdir", - "dist", - "--minify-whitespace", - "--minify-syntax", - "--keep-names", - "--external", - "mupdf", - "--external", - "@oh-my-pi/pi-natives", - "--external", - "@huggingface/transformers", - "--define", - 'process.env.PI_BUNDLED="true"', - "./src/cli.ts", - ]); + // The npm bundle ships no stats dashboard sources or prebuilt dist/client, + // so embed the dashboard archive the same way compiled binaries do + // (scripts/build-binary.ts). Reset afterwards to keep the checked-in + // placeholder empty. + await runCommand(["bun", "--cwd=../stats", "scripts/generate-client-bundle.ts", "--generate"]); + try { + await runCommand([ + "bun", + "build", + "--target=bun", + "--outdir", + "dist", + "--minify-whitespace", + "--minify-syntax", + "--keep-names", + "--external", + "mupdf", + "--external", + "@oh-my-pi/pi-natives", + "--external", + "@huggingface/transformers", + "--define", + 'process.env.PI_BUNDLED="true"', + "./src/cli.ts", + ]); + } finally { + await runCommand(["bun", "--cwd=../stats", "scripts/generate-client-bundle.ts", "--reset"]); + } await ensureShebang(); const stat = await fs.stat(cliPath); const elapsedMs = (Bun.nanoseconds() - start) / 1_000_000; diff --git a/packages/coding-agent/src/cli.ts b/packages/coding-agent/src/cli.ts index 337efd411..c51b070cc 100755 --- a/packages/coding-agent/src/cli.ts +++ b/packages/coding-agent/src/cli.ts @@ -43,19 +43,33 @@ async function showHelp(config: CliConfig): Promise { } } /** - * Smoke-test entry. Spawns bundled workers, pings them, exits. + * Smoke-test entry. Spawns bundled workers, serves the stats dashboard once, + * pings everything, then exits. * - * Purpose: catch the silent worker-load regressions that hit compiled - * binaries (issues #1011 and #1027). Version/help paths do not spawn worker - * modules on a fresh install, so this probe is the minimal end-to-end test - * that proves `new Worker(...)` resolves and bundled worker modules evaluate. + * Purpose: catch the silent worker-load and bundled-asset regressions that hit + * compiled binaries and the npm CLI bundle. Version/help paths do not spawn + * worker modules or serve dashboard assets on a fresh install, so this probe is + * the minimal end-to-end test that proves those distribution-only paths work. * Wired into `scripts/install-tests/run-ci.sh` so binary / source-link / * tarball installs all exercise it on every CI run. */ async function runSmokeTest(): Promise { - const { smokeTestSyncWorker } = await import("@oh-my-pi/omp-stats"); + const { smokeTestSyncWorker, startServer } = await import("@oh-my-pi/omp-stats"); const { smokeTestTinyTitleWorker } = await import("./tiny/title-client"); await smokeTestSyncWorker(); + + const statsServer = await startServer(0); + try { + const response = await fetch(`http://127.0.0.1:${statsServer.port}/`); + if (!response.ok) throw new Error(`stats dashboard smoke failed: HTTP ${response.status}`); + const html = await response.text(); + if (!html.includes('
') || !html.includes("index.js")) { + throw new Error("stats dashboard smoke failed: dashboard HTML was not served"); + } + } finally { + statsServer.stop(); + } + await smokeTestTinyTitleWorker(); process.stdout.write("smoke-test: ok\n"); } diff --git a/packages/coding-agent/test/stats-dashboard-bundle.test.ts b/packages/coding-agent/test/stats-dashboard-bundle.test.ts new file mode 100644 index 000000000..308763721 --- /dev/null +++ b/packages/coding-agent/test/stats-dashboard-bundle.test.ts @@ -0,0 +1,30 @@ +import { describe, expect, it } from "bun:test"; +import * as path from "node:path"; + +describe("stats dashboard assets in distributed CLI builds", () => { + const repoRoot = path.resolve(import.meta.dir, "../../.."); + const bundleScriptPath = path.join(repoRoot, "packages/coding-agent/scripts/bundle-dist.ts"); + const cliPath = path.join(repoRoot, "packages/coding-agent/src/cli.ts"); + const statsServerPath = path.join(repoRoot, "packages/stats/src/server.ts"); + + it("embeds the stats client archive while building the npm CLI bundle", async () => { + const bundleScript = await Bun.file(bundleScriptPath).text(); + expect(bundleScript).toContain(`"scripts/generate-client-bundle.ts", "--generate"`); + expect(bundleScript).toContain(`"scripts/generate-client-bundle.ts", "--reset"`); + expect(bundleScript).toContain(`process.env.PI_BUNDLED="true"`); + }); + + it("uses embedded stats assets for prebuilt CLI distributions", async () => { + const statsServer = await Bun.file(statsServerPath).text(); + expect(statsServer).toContain("process.env.PI_BUNDLED"); + expect(statsServer).toContain("USE_EMBEDDED_CLIENT"); + expect(statsServer).toContain("Embedded stats client bundle missing"); + }); + + it("probes dashboard static assets in the install smoke test path", async () => { + const cliSource = await Bun.file(cliPath).text(); + expect(cliSource).toContain("startServer(0)"); + expect(cliSource).toContain("127.0.0.1"); + expect(cliSource).toContain("dashboard HTML was not served"); + }); +}); diff --git a/packages/stats/CHANGELOG.md b/packages/stats/CHANGELOG.md index f9d81a3ef..12e98ce90 100644 --- a/packages/stats/CHANGELOG.md +++ b/packages/stats/CHANGELOG.md @@ -1,6 +1,14 @@ # Changelog ## [Unreleased] +### Added + +- Added support for prebuilt npm bundle mode via `PI_BUNDLED`, allowing the stats server to use an embedded dashboard bundle in packaged CLI distributions + +### Fixed + +- Fixed handling of legacy `embedded-client.generated.txt` placeholder content so it is treated as missing archive instead of being decoded into invalid bytes +- Fixed ENOENT handling while scanning dashboard source/build directories so missing `client/` or `dist/client` trees no longer crash startup ## [15.10.11] - 2026-06-10 diff --git a/packages/stats/scripts/generate-client-bundle.ts b/packages/stats/scripts/generate-client-bundle.ts index ef0b9eebb..a5383a09f 100644 --- a/packages/stats/scripts/generate-client-bundle.ts +++ b/packages/stats/scripts/generate-client-bundle.ts @@ -11,16 +11,9 @@ const DIST_CLIENT_DIR = path.join("dist", "client"); const GENERATE_FLAG = "--generate"; const RESET_FLAG = "--reset"; -function placeholderContent(): string { - return `/** - * Embedded stats dashboard bundle for compiled binaries. - * - * This file is generated by \`bun --cwd=packages/stats scripts/generate-client-bundle.ts --generate\` during - * binary builds. The checked-in value is intentionally empty. - */ -export const EMBEDDED_CLIENT_ARCHIVE_TAR_GZ_BASE64 = ""; -`; -} +// `--reset` restores the checked-in state: an empty file. The runtime treats +// blank (or any non-base64) content as "no archive embedded" and builds the +// dashboard from source instead; see src/embedded-client.ts. async function collectFiles(dir: string): Promise { const entries = await fs.readdir(dir, { withFileTypes: true }); @@ -60,7 +53,7 @@ async function buildArchiveBase64(dir: string): Promise { async function main(): Promise { if (process.argv.includes(RESET_FLAG)) { - await Bun.write(GENERATED_FILE, placeholderContent()); + await Bun.write(GENERATED_FILE, ""); console.log(`Reset ${GENERATED_FILE}`); return; } diff --git a/packages/stats/src/embedded-client.generated.txt b/packages/stats/src/embedded-client.generated.txt index 962b4d37c..e69de29bb 100644 --- a/packages/stats/src/embedded-client.generated.txt +++ b/packages/stats/src/embedded-client.generated.txt @@ -1,7 +0,0 @@ -/** - * Embedded stats dashboard bundle for compiled binaries. - * - * This file is generated by `bun --cwd=packages/stats scripts/generate-client-bundle.ts --generate` during - * binary builds. The checked-in value is intentionally empty. - */ -export const EMBEDDED_CLIENT_ARCHIVE_TAR_GZ_BASE64 = ""; diff --git a/packages/stats/src/embedded-client.ts b/packages/stats/src/embedded-client.ts new file mode 100644 index 000000000..32ce10934 --- /dev/null +++ b/packages/stats/src/embedded-client.ts @@ -0,0 +1,26 @@ +/** + * Embedded stats dashboard archive handling. + * + * `embedded-client.generated.txt` holds the base64 of a gzipped tar of the + * built dashboard (`dist/client`). It is populated by + * `scripts/generate-client-bundle.ts --generate` for compiled binaries and the + * prepacked npm bundle, and reset to an empty file afterwards so the dev tree + * keeps building the dashboard from source. + */ + +/** + * Decode the generated archive text. + * + * Returns `null` when the content is blank or not a raw gzip archive encoded as + * base64 — notably the legacy placeholder that contained a TypeScript + * `export const … = "";` stub, which must be treated as "no archive embedded" + * rather than decoded into garbage bytes. + */ +export function decodeEmbeddedClientArchive(txt: string): Buffer | null { + const normalized = txt.replaceAll(/\s+/g, ""); + if (!normalized) return null; + if (!/^[A-Za-z0-9+/]+={0,2}$/.test(normalized)) return null; + const archiveBytes = Buffer.from(normalized, "base64"); + if (archiveBytes[0] !== 0x1f || archiveBytes[1] !== 0x8b) return null; + return archiveBytes; +} diff --git a/packages/stats/src/server.ts b/packages/stats/src/server.ts index c06c591b2..c3bdcba1b 100644 --- a/packages/stats/src/server.ts +++ b/packages/stats/src/server.ts @@ -1,6 +1,8 @@ +import type { Dirent } from "node:fs"; import * as fs from "node:fs/promises"; import * as os from "node:os"; import * as path from "node:path"; +import { isEnoent } from "@oh-my-pi/pi-utils"; import { $ } from "bun"; import { getBehaviorDashboardStats, @@ -14,24 +16,27 @@ import { getTotalMessageCount, syncAllSessions, } from "./aggregator"; +import { decodeEmbeddedClientArchive } from "./embedded-client"; import embeddedClientArchiveTxt from "./embedded-client.generated.txt"; -const getEmbeddedClientArchive = (() => { - const txt = embeddedClientArchiveTxt.replaceAll(/[\s\r\n]/g, "").trim(); - if (!txt) return null; - return () => Buffer.from(txt, "base64"); -})(); +const EMBEDDED_CLIENT_ARCHIVE = decodeEmbeddedClientArchive(embeddedClientArchiveTxt); const CLIENT_DIR = path.join(import.meta.dir, "client"); const STATIC_DIR = path.join(import.meta.dir, "..", "dist", "client"); const IS_BUN_COMPILED = - Bun.env.PI_COMPILED || + Boolean(process.env.PI_COMPILED || Bun.env.PI_COMPILED) || import.meta.url.includes("$bunfs") || import.meta.url.includes("~BUN") || import.meta.url.includes("%7EBUN"); +// The prepacked npm bundle (coding-agent dist/cli.js) constant-folds +// process.env.PI_BUNDLED at build time. Like compiled binaries, it ships no +// dashboard sources or prebuilt dist/client next to the bundle, so the +// embedded archive is the only viable asset source. +const IS_PREBUILT = IS_BUN_COMPILED || Boolean(process.env.PI_BUNDLED || Bun.env.PI_BUNDLED); +const USE_EMBEDDED_CLIENT = EMBEDDED_CLIENT_ARCHIVE !== null || IS_PREBUILT; -const COMPILED_CLIENT_DIR_ROOT = path.join(os.tmpdir(), "omp-stats-client"); -let compiledClientDirPromise: Promise | null = null; +const EMBEDDED_CLIENT_DIR_ROOT = path.join(os.tmpdir(), "omp-stats-client"); +let embeddedClientDirPromise: Promise | null = null; function sanitizeArchivePath(archivePath: string): string | null { const normalized = archivePath.replaceAll("\\", "/").replace(/^\.\//, ""); @@ -56,18 +61,19 @@ async function extractEmbeddedClientArchive(archiveBytes: Buffer, outputDir: str } } -async function getCompiledClientDir(): Promise { - if (!IS_BUN_COMPILED) return STATIC_DIR; - if (compiledClientDirPromise) return compiledClientDirPromise; +async function getEmbeddedClientDir(): Promise { + if (!USE_EMBEDDED_CLIENT) return STATIC_DIR; + if (embeddedClientDirPromise) return embeddedClientDirPromise; - const archiveBytes = getEmbeddedClientArchive?.(); - if (!archiveBytes) { - throw new Error("Compiled stats client bundle missing. Rebuild binary with embedded stats assets."); + if (!EMBEDDED_CLIENT_ARCHIVE) { + throw new Error( + "Embedded stats client bundle missing. Rebuild the omp binary or npm bundle with embedded stats assets.", + ); } - compiledClientDirPromise = (async () => { - const bundleHash = Bun.hash(archiveBytes).toString(16); - const outputDir = path.join(COMPILED_CLIENT_DIR_ROOT, bundleHash); + embeddedClientDirPromise = (async () => { + const bundleHash = Bun.hash(EMBEDDED_CLIENT_ARCHIVE).toString(16); + const outputDir = path.join(EMBEDDED_CLIENT_DIR_ROOT, bundleHash); const markerPath = path.join(outputDir, "index.html"); try { const marker = await fs.stat(markerPath); @@ -76,15 +82,24 @@ async function getCompiledClientDir(): Promise { await fs.rm(outputDir, { recursive: true, force: true }); await fs.mkdir(outputDir, { recursive: true }); - await extractEmbeddedClientArchive(archiveBytes, outputDir); + await extractEmbeddedClientArchive(EMBEDDED_CLIENT_ARCHIVE, outputDir); return outputDir; })(); - return compiledClientDirPromise; + return embeddedClientDirPromise; } async function getLatestMtime(dir: string): Promise { - const entries = await fs.readdir(dir, { withFileTypes: true }); + let entries: Dirent[]; + try { + entries = await fs.readdir(dir, { withFileTypes: true }); + } catch (err) { + // Tolerate missing source trees (e.g. installs without the dashboard + // sources); the caller falls back to prebuilt assets or a clear build + // failure instead of crashing on the scan. + if (isEnoent(err)) return 0; + throw err; + } const promises = []; for (const entry of entries) { @@ -108,7 +123,7 @@ async function getLatestMtime(dir: string): Promise { } const ensureClientBuild = async () => { - if (IS_BUN_COMPILED) return; + if (USE_EMBEDDED_CLIENT) return; const indexPath = path.join(STATIC_DIR, "index.html"); const cssPath = path.join(STATIC_DIR, "styles.css"); const clientSourceMtime = await getLatestMtime(CLIENT_DIR); @@ -247,7 +262,7 @@ async function handleApi(req: Request): Promise { * Handle static file requests. */ async function handleStatic(requestPath: string): Promise { - const staticDir = IS_BUN_COMPILED ? await getCompiledClientDir() : STATIC_DIR; + const staticDir = await getEmbeddedClientDir(); const filePath = requestPath === "/" ? "/index.html" : requestPath; const fullPath = path.join(staticDir, filePath); diff --git a/scripts/install-tests/run-ci.sh b/scripts/install-tests/run-ci.sh index a27736e69..62f50b140 100755 --- a/scripts/install-tests/run-ci.sh +++ b/scripts/install-tests/run-ci.sh @@ -21,12 +21,9 @@ smoke_cli() { XDG_DATA_HOME="$runtime_dir/xdg" HOME="$runtime_dir/home" "$omp_bin" --version XDG_DATA_HOME="$runtime_dir/xdg" HOME="$runtime_dir/home" "$omp_bin" --help >/dev/null XDG_DATA_HOME="$runtime_dir/xdg" HOME="$runtime_dir/home" "$omp_bin" stats --summary >/dev/null - # Spawns the stats sync worker via `new Worker(...)` and waits for a pong. - # Regression probe for #1011 (browser tab worker) and #1027 (stats sync - # worker) — both broke silently in compiled binaries because the `with - # { type: "file" }` import pattern only copies the worker as a raw asset - # without bundling its imports. `stats --summary` doesn't catch this on a - # fresh install (no session files = no Worker spawn). + # Spawns bundled workers and serves the stats dashboard once. Regression + # probe for #1011/#1027 worker loading and for npm/compiled distributions + # missing the dashboard assets that `stats --summary` never touches. XDG_DATA_HOME="$runtime_dir/xdg" HOME="$runtime_dir/home" "$omp_bin" --smoke-test } From fd06c70ddeb180501094e5369df515d0acc1a8f8 Mon Sep 17 00:00:00 2001 From: can1357 Date: Wed, 10 Jun 2026 10:55:13 +0200 Subject: [PATCH 02/27] fix(ask): restored checkmark glyph on custom answers and done option - Replaced ask-tool icon with success glyph swapped during status sweep. - Updated tests to expect a single success glyph on custom input. --- packages/coding-agent/CHANGELOG.md | 1 + packages/coding-agent/src/tools/ask.ts | 6 ++++-- packages/coding-agent/test/tools/ask.test.ts | 15 +++++++-------- 3 files changed, 12 insertions(+), 10 deletions(-) diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index d25ba3994..0183d097f 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -6,6 +6,7 @@ - Fixed npm CLI distribution bundles by embedding the stats dashboard client bundle so dashboard assets are served in prebuilt installs - Fixed the CLI smoke-test command to start the stats server and verify dashboard HTML is served, catching bundled-asset regressions - Added verification of a `
` and `index.js` in smoke-test dashboard responses +- Restored the checkmark glyph on ask-tool custom answers and the multi-select "Done selecting" option, which a status-glyph sweep had swapped for the ask tool icon ## [15.10.12] - 2026-06-10 diff --git a/packages/coding-agent/src/tools/ask.ts b/packages/coding-agent/src/tools/ask.ts index 411d0a03a..bc63c1c62 100644 --- a/packages/coding-agent/src/tools/ask.ts +++ b/packages/coding-agent/src/tools/ask.ts @@ -104,7 +104,7 @@ const RECOMMENDED_SUFFIX = " (Recommended)"; const TIMEOUT_DETECTION_TOLERANCE_MS = 1_000; function getDoneOptionLabel(): string { - return `${theme.symbol("tool.ask")} Done selecting`; + return `${theme.status.success} Done selecting`; } /** Add "(Recommended)" suffix to the option at the given index if not already present */ @@ -694,7 +694,9 @@ function normalizeRenderQuestions(raw: unknown): NonNullable { expect(renderedText).toContain("second line"); expect(renderedText).toContain("third line"); - // Count tool.ask glyphs — should be exactly one for the custom input block, - // plus one for the question status icon (if present). The key contract is that - // continuation lines do NOT get their own glyph. - const askGlyph = theme!.symbol("tool.ask"); + // Count success glyphs — should be exactly one for the custom input block. + // The key contract is that continuation lines do NOT get their own glyph. + const successGlyph = theme!.symbol("status.success"); const successIconCount = ( - renderedText.match(new RegExp(askGlyph.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"), "g")) || [] + renderedText.match(new RegExp(successGlyph.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"), "g")) || [] ).length; - // One icon on the status line header + one on the custom input first line = 2 max - expect(successIconCount).toBeLessThanOrEqual(2); + // One glyph on the custom input first line; header uses the tool.ask icon. + expect(successIconCount).toBe(1); // Ensure "second line" and "third line" are NOT preceded by a success icon on their own line const lines = renderedText.split("\n"); @@ -844,7 +843,7 @@ describe("AskTool multiline custom input rendering", () => { const trimmed = line.trim(); if (trimmed.includes("second line") || trimmed.includes("third line")) { // These continuation lines must NOT start with a success icon - expect(trimmed.startsWith(askGlyph)).toBe(false); + expect(trimmed.startsWith(successGlyph)).toBe(false); } } }); From 6f81403b351248e90f8f18518f3f031c604711ad Mon Sep 17 00:00:00 2001 From: can1357 Date: Wed, 10 Jun 2026 10:50:48 +0200 Subject: [PATCH 03/27] feat(session-stats): added LLM-assisted token-usage audit - Added standalone Bun script reading session JSONL with real per-request usage. - Scanned main-vs-subagent split, residency, repeated reads, and edit churn. - Classified costliest sessions via pi-ai for hygiene, spawn, and waste verdicts. - Added stats:audit script, prompt, and scanFile unit tests. --- package.json | 1 + scripts/session-stats/README.md | 42 +- scripts/session-stats/audit-prompt.md | 41 + scripts/session-stats/audit.test.ts | 174 +++ scripts/session-stats/audit.ts | 1438 +++++++++++++++++++++++++ 5 files changed, 1694 insertions(+), 2 deletions(-) create mode 100644 scripts/session-stats/audit-prompt.md create mode 100644 scripts/session-stats/audit.test.ts create mode 100755 scripts/session-stats/audit.ts diff --git a/package.json b/package.json index ff6594257..173cd305c 100644 --- a/package.json +++ b/package.json @@ -130,6 +130,7 @@ "stats:tools": "python3 scripts/session-stats/analyze.py tools", "stats:edits": "python3 scripts/session-stats/analyze.py edits", "stats:followups": "python3 scripts/session-stats/analyze.py followups", + "stats:audit": "bun scripts/session-stats/audit.ts", "test:py": "python3 -m pytest -x python/omp-rpc/tests && python3 -m pytest -x python/robomp/tests", "robomp:install": "pip install -e 'python/robomp[dev]'", "robomp:serve": "python3 -m robomp serve", diff --git a/scripts/session-stats/README.md b/scripts/session-stats/README.md index 35b39c400..62d0d2766 100644 --- a/scripts/session-stats/README.md +++ b/scripts/session-stats/README.md @@ -9,8 +9,10 @@ scripts. ``` scripts/session-stats/ - sync.py # walks ~/.omp/agent/sessions/ and populates ss_* tables - analyze.py # tools | edits | followups subcommands over the synced db + sync.py # walks ~/.omp/agent/sessions/ and populates ss_* tables + analyze.py # tools | edits | followups subcommands over the synced db + audit.ts # LLM-assisted token-usage audit (no sync needed) + audit-prompt.md # system prompt for the audit classifier ``` ## One-time prep @@ -70,6 +72,42 @@ audit reads each call's `is_error` flag as the authoritative success/failure signal and decodes hashline op kinds (`replace`, `insert after`, `delete`, `replace block`, …) into the verb distribution. +## Usage audit (`audit.ts`) + +Standalone Bun script — reads session JSONL directly (no `sync.py` / tiktoken +needed) and uses the *real* per-request usage recorded in each assistant +message (input/output/cacheRead/cacheWrite + nominal cost) instead of +re-tokenizing. + +```sh +bun run stats:audit # last week, scan + LLM analysis +bun run stats:audit -- --no-llm # scan-only report +bun run stats:audit -- --since 3d --folder Projects-pi +bun run stats:audit -- --min-cost 5 --max-llm 8 --json /tmp/audit.json +bun run stats:audit -- --digest-dir /tmp/digests # inspect classifier inputs +bun run stats:audit -- --session parser # classify sessions matching id/title (ignores --min-cost) +bun run stats:audit -- --no-cache # force fresh LLM verdicts +``` + +The scan phase reports the main-vs-subagent usage split, per-folder and +per-session cost, per-tool traffic (estimated arg/result tokens plus a +*context-residency* metric: result tokens × subsequent requests), repeated +reads of the same file, the largest single tool results, compactions, and +edit-failure churn. + +The LLM phase (default `anthropic/claude-sonnet-4-6` via `@oh-my-pi/pi-ai`, +credentials resolved through omp's auth storage — stored key, OAuth, or env +var) classifies the costliest sessions: multi-topic sessions that should have +been split or handed off, task spawns that were wasteful or failed to transfer +context, and the biggest waste sources with concrete fixes. A final aggregate +call distills systemic findings and quick wins across sessions. + +Verdicts are cached in `~/.omp/stats-audit-cache.json` (keyed by session id + +digest hash + model + prompt hash, so any change to the transcript, digest +format, or `audit-prompt.md` invalidates the entry automatically). Re-runs +reuse cached verdicts for free; `--no-cache` bypasses reads but still writes +fresh results. + The Rust crate that previously lived here was retired in favor of this SQLite-backed flow. The schema persists everything the analyses used to recompute on every run (token counts, hashline parse output, success flags), diff --git a/scripts/session-stats/audit-prompt.md b/scripts/session-stats/audit-prompt.md new file mode 100644 index 000000000..88f3af214 --- /dev/null +++ b/scripts/session-stats/audit-prompt.md @@ -0,0 +1,41 @@ +# Role + +You are a token-efficiency auditor for **omp**, a terminal coding agent. You receive a digest of one recorded session (or an aggregate of per-session verdicts) and return a structured analysis by calling the `respond` tool. Never reply with prose; always call the tool. + +# How omp sessions spend tokens + +- A session is a conversation with a main agent. Context is append-only: every tool result, user message, and assistant message stays in context and is re-sent on **every subsequent request** (cached prefixes are re-billed at ~10% of input price as `cache-read`). +- Therefore a large tool result early in a long session costs far more than its own size. The digest's `residency` metric approximates this: result tokens × number of later requests. +- `task` spawns subagents: isolated contexts that do work and return only a final report to the main context. Subagents are the cheap way to do exploration/bulk edits — their intermediate tool traffic never lands in the main context. A spawn is wasted when the child re-discovers context the parent already had (a thin assignment prompt → the child burns tokens re-exploring), when the work was trivial enough to do inline, or when the child fails/errors and the parent redoes the work. +- `compaction` events mean the context grew past its limit and was summarized — a strong sign the session ran too long or accumulated bloat. +- Users can start fresh sessions, use `/handoff` (summarize + continue in a new session), or delegate to subagents. Switching topics inside one long session drags the entire prior topic's context into every request of the new topic. + +# Digest format notes + +- Token counts labelled `~` are estimates (chars/4). Usage totals (`billed-in`, `out`, `cost`) are real numbers recorded from the API. +- `cache-read N%` is the fraction of input that was cache-hits. A low ratio in a long session means cache churn (model switches, branch edits, parallel branches) — expensive. +- Turn flow lists each user message with the work it triggered. `[synthetic/steering]` turns were injected by the system, not typed by the user. +- `Repeated reads` lists files read ≥3 times in the same context — usually re-reads after edits or forgetting earlier reads. Each line carries measured figures: ` ×N (~Xtok total, ~Y residency)` — X is the summed result size, Y is the measured residency cost. Quote these numbers; do not derive your own residency estimates for repeated reads. +- Spawn entries show the assignment prompt size, the child's own spend, and how the child ended. `ended: (no final text; last tool: X)` is NOT a failure — many subagents deliver their report through the task result channel and never emit trailing prose. Judge spawn failure from `[ERRORED]` flags, `(no output)`, or a useless merged result — not from the absence of final text. +- `merged result ~N` is the task report as it sits in the parent context NOW; `[Output truncated - N tokens]` in a snippet means the result was later pruned from context (the prune is a context-saving feature working as intended, not data loss). + +# Your judgments + +1. **Session hygiene** (`score`, `multiTopic`, `topics`, `shouldHaveSplit`, `handoffOpportunities`) + - Identify the distinct *unrelated* topics. Sequential phases of one task (implement → test → docs) are ONE topic. Unrelated bugfix dropped into a feature session IS a second topic. + - `shouldHaveSplit` only when a split/handoff would have plausibly saved real money: e.g. topic B started after context already held 100k+ tokens of topic A. + - `handoffOpportunities`: name the specific turn ("T7: new topic 'fix CI' while 180k of refactor context was loaded — fresh session would have started at ~10k"). + - Score 0–10 for token efficiency only (not task success). 8–10 lean sessions, well-delegated; 4–7 noticeable waste; 0–3 heavy waste (multiple compactions, repeated giant results, redundant re-reads, dead spawns). + +2. **Spawn quality** (`spawnVerdicts`) — judge each spawn group worth judging (skip trivial ones, cap 10): + - `good`: meaningful work isolated from main context, reasonable prompt, useful report. + - `unnecessary`: work small enough to do inline; spawn overhead (system prompt + exploration) exceeded the savings. + - `wrong-granularity`: should have been more/fewer parallel tasks, or sequenced (children redid each other's discovery). + - `context-transfer-failure`: assignment too thin — child visibly re-explored what the parent knew (big child spend on discovery, prompt under ~1k tokens, child asking-the-codebase questions the parent had answered). Cite the evidence. + - `failed`: child errored/died/produced nothing useful; parent paid for it anyway. + +3. **Waste sources** (`waste`) — the biggest concrete token sinks, largest first, with a practical fix each. Ground them in the digest: residency-heavy tools, repeated reads, giant single results (full-file reads where a range would do, unfiltered test output), edit retry churn, low cache-read ratio, synthetic auto-continue loops, model choice (e.g. expensive model on mechanical work). Each item carries `estTokens` (tokens attributable to the waste) and `estUsd` (realistic dollars a leaner workflow would have saved) — keep the two consistent. Distinguish *residency* tokens from *billed* tokens: residency is re-paid on each later request at ~10% of input price (cache-read), so a residency-derived figure must be discounted accordingly — never price it as if it were full-rate input tokens. + +Be specific and quantitative: cite turn numbers, file paths, spawn labels, and token figures from the digest. No generic advice ("use tools efficiently"); every claim must trace to a line of the digest. `headline` MUST be a non-empty single sentence; a response with an empty or missing headline is invalid and will be retried. + +For aggregate requests (input starts with `# AGGREGATE`): the per-session verdicts arrive as one compact JSON object per line under `Per-session data (JSON, one per line):`. Synthesize *recurring* patterns across sessions into `systemicIssues` (issue + evidence + fix), `quickWins` (one-line habit changes ranked by savings), and a 2–4 sentence `summary` addressed directly to the user. Cite only sessions and figures present in that data, and refer to sessions by their title. Never re-attribute main-context turns as subagent spend (prior failure: a 160-request main-context debugging phase was misreported as a "runaway child"). Do not restate per-session findings verbatim; find the patterns. diff --git a/scripts/session-stats/audit.test.ts b/scripts/session-stats/audit.test.ts new file mode 100644 index 000000000..9df7ec2e7 --- /dev/null +++ b/scripts/session-stats/audit.test.ts @@ -0,0 +1,174 @@ +import { afterAll, describe, expect, it } from "bun:test"; +import * as fs from "node:fs/promises"; +import * as os from "node:os"; +import * as path from "node:path"; +import { normalizeReadPath, parseSince, scanFile } from "./audit"; + +describe("parseSince", () => { + it("maps window suffixes to milliseconds", () => { + expect(parseSince("12h")).toBe(12 * 3_600_000); + expect(parseSince("3d")).toBe(3 * 24 * 3_600_000); + expect(parseSince("1w")).toBe(7 * 24 * 3_600_000); + expect(parseSince("w")).toBe(7 * 24 * 3_600_000); + expect(parseSince("2mo")).toBe(60 * 24 * 3_600_000); + }); + + it("rejects unparseable windows", () => { + expect(() => parseSince("soon")).toThrow(); + expect(() => parseSince("5x")).toThrow(); + }); +}); + +describe("normalizeReadPath", () => { + it("strips line/raw selectors from plain paths", () => { + expect(normalizeReadPath("src/foo.ts:50-200")).toBe("src/foo.ts"); + expect(normalizeReadPath("src/foo.ts:50+10")).toBe("src/foo.ts"); + expect(normalizeReadPath("src/foo.ts:raw")).toBe("src/foo.ts"); + expect(normalizeReadPath("src/foo.ts:5-16,960-973")).toBe("src/foo.ts"); + expect(normalizeReadPath("src/foo.ts:2-4:raw")).toBe("src/foo.ts"); + }); + + it("keeps internal URL schemes distinct instead of collapsing to the scheme", () => { + expect(normalizeReadPath("artifact://37")).toBe("artifact://37"); + expect(normalizeReadPath("agent://h0qbtw5y/report")).toBe("agent://h0qbtw5y/report"); + expect(normalizeReadPath("artifact://37:50-100")).toBe("artifact://37"); + }); + + it("leaves selector-free paths untouched", () => { + expect(normalizeReadPath("docs/readme.md")).toBe("docs/readme.md"); + }); +}); + +// --------------------------------------------------------------------------- +// scanFile contract on a synthetic session transcript + +const tmpDir = await fs.mkdtemp(path.join(os.tmpdir(), "audit-test-")); +afterAll(() => fs.rm(tmpDir, { recursive: true, force: true })); + +function asst(opts: { + ts: number; + usage: { input: number; output: number; cacheRead?: number; cacheWrite?: number; cost: number }; + content: unknown[]; + stopReason?: string; +}): string { + return JSON.stringify({ + type: "message", + message: { + role: "assistant", + model: "test-model", + stopReason: opts.stopReason ?? "toolUse", + timestamp: opts.ts, + usage: { + input: opts.usage.input, + output: opts.usage.output, + cacheRead: opts.usage.cacheRead ?? 0, + cacheWrite: opts.usage.cacheWrite ?? 0, + cost: { total: opts.usage.cost }, + }, + content: opts.content, + }, + }); +} + +function toolResult(callId: string, text: string, ts: number, isError = false): string { + return JSON.stringify({ + type: "message", + message: { + role: "toolResult", + toolCallId: callId, + toolName: "read", + isError, + timestamp: ts, + content: [{ type: "text", text }], + }, + }); +} + +it("scanFile recovers usage, turns, spawns, residency, and pruned result sizes", async () => { + const lines = [ + JSON.stringify({ type: "session", id: "s1", timestamp: "2026-06-09T10:00:00.000Z", title: "Test session" }), + JSON.stringify({ type: "message", message: { role: "user", content: "fix the bug", timestamp: 1000 } }), + // req 1: reads a file (result re-paid by 2 later requests) + asst({ + ts: 2000, + usage: { input: 100, output: 10, cacheWrite: 50, cost: 0.5 }, + content: [{ type: "toolCall", id: "c1", name: "read", arguments: { path: "src/a.ts:10-20" } }], + }), + toolResult("c1", "x".repeat(400), 2500), // ~100 tokens + // req 2: spawns a task; its result was pruned and stores a placeholder + asst({ + ts: 3000, + usage: { input: 200, output: 20, cacheRead: 300, cost: 1.0 }, + content: [ + { + type: "toolCall", + id: "c2", + name: "task", + arguments: { agent: "task", tasks: [{ id: "FixParser", description: "Fix the parser" }] }, + }, + ], + }), + JSON.stringify({ + type: "message", + message: { + role: "toolResult", + toolCallId: "c2", + toolName: "task", + isError: false, + timestamp: 3500, + content: [{ type: "text", text: "[Output truncated - 1993 tokens]" }], + }, + }), + // second user turn + final answer with no tools + JSON.stringify({ type: "message", message: { role: "user", content: "thanks, also re-read it", timestamp: 4000 } }), + asst({ + ts: 5000, + usage: { input: 400, output: 40, cacheRead: 600, cost: 2.0 }, + content: [ + { type: "text", text: "All done." }, + { type: "toolCall", id: "c3", name: "read", arguments: { path: "src/a.ts:raw" } }, + ], + stopReason: "stop", + }), + toolResult("c3", "y".repeat(200), 5500), + ]; + const file = path.join(tmpDir, "synthetic.jsonl"); + await Bun.write(file, `${lines.join("\n")}\n`); + + const scan = await scanFile(file); + if (!scan) throw new Error("scanFile returned undefined"); + + // Real usage sums (not estimates). + expect(scan.usage.requests).toBe(3); + expect(scan.usage.input).toBe(700); + expect(scan.usage.output).toBe(70); + expect(scan.usage.cacheRead).toBe(900); + expect(scan.usage.cacheWrite).toBe(50); + expect(scan.usage.cost).toBeCloseTo(3.5); + // Context peak = max single-request input+cacheRead+cacheWrite. + expect(scan.contextPeak).toBe(1000); + + // Turn segmentation: 2 user turns; first turn carries 2 requests. + expect(scan.turns.length).toBe(2); + expect(scan.turns[0].requests).toBe(2); + expect(scan.turns[0].cost).toBeCloseTo(1.5); + expect(scan.turns[1].requests).toBe(1); + + // Both reads of src/a.ts group together despite different selectors, and + // per-path residency weights each result by later requests: c1 (100 tok, + // request 1 of 3) → ×2 = 200; c3 (50 tok, request 3 of 3) → ×0. + expect(scan.readCounts.get("src/a.ts")?.count).toBe(2); + expect(scan.readCounts.get("src/a.ts")?.residency).toBe(200); + + // Pruned task result recovers its true size from the placeholder. + expect(scan.spawns.length).toBe(1); + expect(scan.spawns[0].labels).toEqual(["FixParser"]); + expect(scan.spawns[0].resultToks).toBe(1993); + + // Residency: c1 result (100 tok) lands at request 1 → re-paid by 2 later + // requests; c3 result (50 tok) lands at request 3 → no later requests. + expect(scan.toolAgg.get("read")?.residency).toBe(200); + + expect(scan.lastAssistantText).toBe("All done."); + expect(scan.title).toBe("Test session"); +}); diff --git a/scripts/session-stats/audit.ts b/scripts/session-stats/audit.ts new file mode 100755 index 000000000..b012c356d --- /dev/null +++ b/scripts/session-stats/audit.ts @@ -0,0 +1,1438 @@ +#!/usr/bin/env bun +/** + * Token-usage audit over the local omp session corpus (~/.omp/agent/sessions/). + * + * Phase 1 (scan, no LLM): walks recent sessions, sums *real* per-request usage + * (input/output/cacheRead/cacheWrite + nominal cost recorded in each assistant + * message), splits main-context vs subagent usage, and aggregates tool traffic + * (estimated tokens in args/results, context residency, repeated reads, edit + * failures, compactions). + * + * Phase 2 (classify): for the costliest sessions, builds a compact digest and + * asks a small model (default: anthropic/claude-sonnet-4-6 via @oh-my-pi/pi-ai) + * to judge: + * a) session hygiene — multiple topics in one chat, missed handoff points, + * b) task-spawn quality — wasteful spawns, context-transfer failures, + * c) the biggest sources of waste given the tool traffic. + * A final aggregate call distills systemic findings across sessions. + * + * Usage: + * bun scripts/session-stats/audit.ts # last week, scan + LLM + * bun scripts/session-stats/audit.ts --since 3d --no-llm # scan only + * bun scripts/session-stats/audit.ts --folder Projects-pi --max-llm 6 + * bun scripts/session-stats/audit.ts --json out.json + * + * Auth: resolves an API key for the classifier provider through omp's auth + * storage (~/.omp/agent/agent.db: stored key, OAuth, or env var fallback). + */ + +import type { Dirent } from "node:fs"; +import * as fs from "node:fs/promises"; +import * as os from "node:os"; +import * as path from "node:path"; +import { parseArgs } from "node:util"; +import { + type Api, + AuthStorage, + completeSimple, + type Model, + SqliteAuthCredentialStore, + type Tool, + type ToolCall, +} from "@oh-my-pi/pi-ai"; +import { type GeneratedProvider, getBundledModel } from "@oh-my-pi/pi-catalog/models"; +import { getAgentDbPath, isEnoent } from "@oh-my-pi/pi-utils"; +import SYSTEM_PROMPT from "./audit-prompt.md" with { type: "text" }; + +const SESSIONS_ROOT = path.join(os.homedir(), ".omp", "agent", "sessions"); +const DEFAULT_MODEL = "anthropic/claude-sonnet-4-6"; +const CACHE_PATH = path.join(os.homedir(), ".omp", "stats-audit-cache.json"); + +// -------------------------------------------------------------------------- +// CLI + +interface CliOptions { + since: number; // ms window + folder?: string; + exclude?: string; + model: string; + maxLlm: number; + minCost: number; + concurrency: number; + json?: string; + noLlm: boolean; + limit?: number; + digestDir?: string; + session?: string; + noCache: boolean; +} + +export function parseSince(raw: string): number { + const m = /^(\d+)?\s*(h|d|w|mo|m)$/.exec(raw.trim()); + if (!m) throw new Error(`invalid --since "${raw}" (use e.g. 12h, 3d, 1w, 1mo)`); + const n = m[1] ? Number.parseInt(m[1], 10) : 1; + const HOUR = 3_600_000; + switch (m[2]) { + case "h": + return n * HOUR; + case "d": + return n * 24 * HOUR; + case "w": + return n * 7 * 24 * HOUR; + default: // m | mo + return n * 30 * 24 * HOUR; + } +} + +function parseCli(argv: string[]): CliOptions { + const { values } = parseArgs({ + args: argv, + options: { + since: { type: "string", default: "1w" }, + folder: { type: "string" }, + exclude: { type: "string" }, + model: { type: "string", default: DEFAULT_MODEL }, + "max-llm": { type: "string", default: "12" }, + "min-cost": { type: "string", default: "1" }, + concurrency: { type: "string", default: "4" }, + json: { type: "string" }, + "no-llm": { type: "boolean", default: false }, + limit: { type: "string" }, + "digest-dir": { type: "string" }, + session: { type: "string" }, + "no-cache": { type: "boolean", default: false }, + help: { type: "boolean", default: false }, + }, + }); + if (values.help) { + console.log( + `session audit — token usage analysis over ~/.omp/agent/sessions\n\n` + + ` --since <12h|3d|1w|1mo> window by session mtime (default 1w)\n` + + ` --folder only folders containing substring\n` + + ` --exclude drop folders containing substring\n` + + ` --model classifier model (default ${DEFAULT_MODEL})\n` + + ` --max-llm classify at most n sessions (default 12)\n` + + ` --min-cost classify only sessions >= cost (default 1)\n` + + ` --concurrency parallel classifier calls (default 4)\n` + + ` --json write full machine-readable results\n` + + ` --digest-dir dump per-session digests fed to the model\n` + + ` --session classify sessions whose id/title matches (ignores --min-cost)\n` + + ` --no-cache skip the verdict cache (~/.omp/stats-audit-cache.json)\n` + + ` --no-llm scan + report only\n` + + ` --limit scan at most n session groups (debug)`, + ); + process.exit(0); + } + return { + since: parseSince(values.since), + folder: values.folder, + exclude: values.exclude, + model: values.model, + maxLlm: Number.parseInt(values["max-llm"], 10), + minCost: Number.parseFloat(values["min-cost"]), + concurrency: Math.max(1, Number.parseInt(values.concurrency, 10)), + json: values.json, + noLlm: values["no-llm"], + limit: values.limit ? Number.parseInt(values.limit, 10) : undefined, + digestDir: values["digest-dir"], + session: values.session, + noCache: values["no-cache"], + }; +} + +// -------------------------------------------------------------------------- +// Usage accounting + +interface UsageTotals { + input: number; + output: number; + cacheRead: number; + cacheWrite: number; + cost: number; + requests: number; +} + +function emptyUsage(): UsageTotals { + return { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0, requests: 0 }; +} + +function addUsage(into: UsageTotals, from: UsageTotals): void { + into.input += from.input; + into.output += from.output; + into.cacheRead += from.cacheRead; + into.cacheWrite += from.cacheWrite; + into.cost += from.cost; + into.requests += from.requests; +} + +function billedTokens(u: UsageTotals): number { + return u.input + u.output + u.cacheRead + u.cacheWrite; +} + +// -------------------------------------------------------------------------- +// Per-file scan + +interface ToolAgg { + calls: number; + argToks: number; + resultToks: number; + errors: number; + /** Σ resultToks × (requests issued after the result landed) — how heavily + * the result sat in context for the rest of the session. */ + residency: number; +} + +interface SpawnCall { + callId: string; + agent: string; + labels: string[]; + descriptions: string[]; + argToks: number; + ts: number; + resultToks: number; + isError: boolean; + resultSnippet: string; +} + +interface TurnInfo { + ts: number; + text: string; + tokens: number; + synthetic: boolean; + requests: number; + outToks: number; + cost: number; + tools: Map; + spawnAgents: string[]; +} + +interface TopResult { + tool: string; + summary: string; + toks: number; +} + +interface FileScan { + path: string; + stem: string; + title?: string; + usage: UsageTotals; + models: Map; + turns: TurnInfo[]; + toolAgg: Map; + spawns: SpawnCall[]; + readCounts: Map; + editErrors: number; + editCalls: number; + compactions: number; + asstErrors: number; + contextPeak: number; + firstTs: number; + lastTs: number; + topResults: TopResult[]; + lastAssistantText: string; + lastToolName: string; +} + +function estTokens(text: string): number { + return Math.ceil(text.length / 4); +} + +/** Matches the placeholder the session writer stores when a tool result was + * pruned from context (`[Output truncated - N tokens]`); N is the true size. */ +const TRUNCATED_RESULT_RE = /\[Output truncated - (\d+) tokens?\]/; + +/** Group key for repeated-read detection: keep `scheme://` URLs intact and + * strip trailing line/raw selectors (`:50-200`, `:raw`, `:2-4:raw`, …). */ +export function normalizeReadPath(p: string): string { + let out = p; + for (;;) { + const next = out.replace(/:(?:raw|conflicts|[0-9][0-9+\-,]*)$/i, ""); + if (next === out) return out; + out = next; + } +} + +function contentText(content: unknown): string { + if (typeof content === "string") return content; + if (!Array.isArray(content)) return ""; + let out = ""; + for (const item of content) { + if (item && typeof item === "object" && (item as { type?: string }).type === "text") { + out += (item as { text?: string }).text ?? ""; + } + } + return out; +} + +function clip(text: string, max: number): string { + const flat = text.replace(/\s+/g, " ").trim(); + return flat.length > max ? `${flat.slice(0, max)}…` : flat; +} + +/** Human-meaningful one-liner for a tool call's arguments. */ +function argSummary(name: string, args: Record | undefined): string { + if (!args) return ""; + const pick = (...keys: string[]): string => { + for (const k of keys) { + const v = args[k]; + if (typeof v === "string" && v) return v; + } + return ""; + }; + switch (name) { + case "read": + case "write": + case "edit": + return clip(pick("path", "file_path", "input"), 80); + case "bash": + return clip(pick("command", "cmd"), 80); + case "search": + return clip(pick("pattern"), 60); + case "find": + return clip(JSON.stringify(args.paths ?? args.pattern ?? ""), 60); + case "task": { + const tasks = Array.isArray(args.tasks) ? args.tasks : []; + const ids = tasks.map(t => (t as { id?: string }).id ?? "?").join(","); + return clip(`agent=${args.agent ?? "?"} tasks=[${ids}]`, 100); + } + default: + return clip(JSON.stringify(args), 70); + } +} + +interface PendingCall { + name: string; + args: Record | undefined; + requestIndex: number; +} + +export async function scanFile(filePath: string): Promise { + let text: string; + try { + text = await Bun.file(filePath).text(); + } catch { + return undefined; + } + const scan: FileScan = { + path: filePath, + stem: path.basename(filePath, ".jsonl"), + usage: emptyUsage(), + models: new Map(), + turns: [], + toolAgg: new Map(), + spawns: [], + readCounts: new Map(), + editErrors: 0, + editCalls: 0, + compactions: 0, + asstErrors: 0, + contextPeak: 0, + firstTs: 0, + lastTs: 0, + topResults: [], + lastAssistantText: "", + lastToolName: "", + }; + const pending = new Map(); + const resultLog: { tool: string; toks: number; requestIndex: number }[] = []; + const readLog: { path: string; toks: number; requestIndex: number }[] = []; + const spawnByCallId = new Map(); + let requestCount = 0; + + const tool = (name: string): ToolAgg => { + let agg = scan.toolAgg.get(name); + if (!agg) { + agg = { calls: 0, argToks: 0, resultToks: 0, errors: 0, residency: 0 }; + scan.toolAgg.set(name, agg); + } + return agg; + }; + + for (const line of text.split("\n")) { + if (!line) continue; + let entry: Record; + try { + entry = JSON.parse(line); + } catch { + continue; // torn tail line from a crashed writer + } + const type = entry.type; + if (type === "session") { + scan.title = (entry.title as string) ?? undefined; + const ts = Date.parse((entry.timestamp as string) ?? ""); + if (Number.isFinite(ts)) scan.firstTs = ts; + continue; + } + if (type === "compaction") { + scan.compactions++; + continue; + } + if (type !== "message") continue; + const msg = entry.message as Record | undefined; + if (!msg) continue; + const ts = typeof msg.timestamp === "number" ? msg.timestamp : 0; + if (ts) { + if (!scan.firstTs) scan.firstTs = ts; + scan.lastTs = Math.max(scan.lastTs, ts); + } + const role = msg.role; + + if (role === "user") { + const textBlob = contentText(msg.content); + scan.turns.push({ + ts, + text: clip(textBlob, 400), + tokens: estTokens(textBlob), + synthetic: msg.synthetic === true || msg.steering === true, + requests: 0, + outToks: 0, + cost: 0, + tools: new Map(), + spawnAgents: [], + }); + continue; + } + + if (role === "assistant") { + requestCount++; + const usage = msg.usage as Record | undefined; + const u: UsageTotals = { + input: (usage?.input as number) || 0, + output: (usage?.output as number) || 0, + cacheRead: (usage?.cacheRead as number) || 0, + cacheWrite: (usage?.cacheWrite as number) || 0, + cost: ((usage?.cost as Record | undefined)?.total as number) || 0, + requests: 1, + }; + addUsage(scan.usage, u); + scan.contextPeak = Math.max(scan.contextPeak, u.input + u.cacheRead + u.cacheWrite); + if (typeof msg.model === "string") { + scan.models.set(msg.model, (scan.models.get(msg.model) ?? 0) + 1); + } + if (msg.stopReason === "error") scan.asstErrors++; + + const turn = scan.turns[scan.turns.length - 1]; + if (turn) { + turn.requests++; + turn.outToks += u.output; + turn.cost += u.cost; + } + + const content = Array.isArray(msg.content) ? msg.content : []; + for (const block of content) { + const b = block as Record; + if (b.type === "text" && typeof b.text === "string" && b.text.trim()) { + scan.lastAssistantText = clip(b.text, 300); + } + if (b.type !== "toolCall") continue; + const name = (b.name as string) ?? "?"; + const args = b.arguments as Record | undefined; + const argToks = estTokens(JSON.stringify(args ?? {})); + const agg = tool(name); + agg.calls++; + agg.argToks += argToks; + const callId = (b.id as string) ?? ""; + pending.set(callId, { name, args, requestIndex: requestCount }); + scan.lastToolName = name; + if (turn) turn.tools.set(name, (turn.tools.get(name) ?? 0) + 1); + if (name === "edit") scan.editCalls++; + if (name === "read") { + const p = typeof args?.path === "string" ? normalizeReadPath(args.path as string) : ""; + if (p) { + const rec = scan.readCounts.get(p) ?? { count: 0, toks: 0, residency: 0 }; + rec.count++; + scan.readCounts.set(p, rec); + } + } + if (name === "task") { + const tasks = Array.isArray(args?.tasks) ? (args?.tasks as Record[]) : []; + const spawn: SpawnCall = { + callId, + agent: typeof args?.agent === "string" ? (args.agent as string) : "?", + labels: tasks.map(t => (typeof t.id === "string" ? t.id : "?")), + descriptions: tasks.map(t => clip(String(t.description ?? t.assignment ?? ""), 90)), + argToks, + ts, + resultToks: 0, + isError: false, + resultSnippet: "", + }; + scan.spawns.push(spawn); + spawnByCallId.set(callId, spawn); + if (turn) turn.spawnAgents.push(spawn.agent); + } + } + continue; + } + + if (role === "toolResult") { + const callId = (msg.toolCallId as string) ?? ""; + const call = pending.get(callId); + const name = call?.name ?? (msg.toolName as string) ?? "?"; + const textBlob = contentText(msg.content); + const truncated = TRUNCATED_RESULT_RE.exec(textBlob); + const toks = truncated ? Math.max(estTokens(textBlob), Number.parseInt(truncated[1], 10)) : estTokens(textBlob); + const agg = tool(name); + agg.resultToks += toks; + if (msg.isError === true) { + agg.errors++; + if (name === "edit") scan.editErrors++; + } + resultLog.push({ tool: name, toks, requestIndex: call?.requestIndex ?? requestCount }); + if (name === "read" && call?.args && typeof call.args.path === "string") { + const p = normalizeReadPath(call.args.path as string); + const rec = scan.readCounts.get(p); + if (rec) rec.toks += toks; + readLog.push({ path: p, toks, requestIndex: call.requestIndex }); + } + const spawn = spawnByCallId.get(callId); + if (spawn) { + spawn.resultToks = toks; + spawn.isError = msg.isError === true; + spawn.resultSnippet = clip(textBlob, 240); + } + if (toks > 2000) { + scan.topResults.push({ tool: name, summary: argSummary(name, call?.args), toks }); + if (scan.topResults.length > 24) { + scan.topResults.sort((a, b) => b.toks - a.toks); + scan.topResults.length = 12; + } + } + pending.delete(callId); + } + } + + // Context residency: result tokens weighted by how many later requests re-paid them. + for (const r of resultLog) { + const later = Math.max(0, requestCount - r.requestIndex); + const agg = scan.toolAgg.get(r.tool); + if (agg) agg.residency += r.toks * later; + } + // Per-path read residency: same weighting, attributed to the normalized path. + for (const r of readLog) { + const rec = scan.readCounts.get(r.path); + if (rec) rec.residency += r.toks * Math.max(0, requestCount - r.requestIndex); + } + scan.topResults.sort((a, b) => b.toks - a.toks); + scan.topResults.length = Math.min(scan.topResults.length, 12); + return scan; +} + +// -------------------------------------------------------------------------- +// Session groups (main + subagent files) + +interface SessionGroup { + folder: string; + id: string; + mtime: number; + main: FileScan; + children: FileScan[]; + usage: UsageTotals; // main + children + subUsage: UsageTotals; // children only +} + +interface DiscoveredGroup { + folder: string; + id: string; + mainPath: string; + childPaths: string[]; + mtime: number; +} + +async function discoverGroups(opts: CliOptions): Promise { + const cutoff = Date.now() - opts.since; + const groups: DiscoveredGroup[] = []; + let folders: string[]; + try { + folders = await fs.readdir(SESSIONS_ROOT); + } catch { + throw new Error(`sessions root not found: ${SESSIONS_ROOT}`); + } + for (const folder of folders) { + if (opts.folder && !folder.includes(opts.folder)) continue; + if (opts.exclude && folder.includes(opts.exclude)) continue; + const folderPath = path.join(SESSIONS_ROOT, folder); + let entries: Dirent[]; + try { + entries = await fs.readdir(folderPath, { withFileTypes: true }); + } catch { + continue; + } + const subdirs = new Set(); + const mains = new Map(); + for (const e of entries) { + if (e.isDirectory()) { + subdirs.add(e.name); + } else if (e.name.endsWith(".jsonl")) { + const p = path.join(folderPath, e.name); + const stat = await fs.stat(p); + mains.set(e.name.slice(0, -6), { path: p, mtime: stat.mtimeMs }); + } + } + for (const [id, main] of mains) { + let childPaths: string[] = []; + let mtime = main.mtime; + if (subdirs.has(id)) { + const dirPath = path.join(folderPath, id); + const nested = await fs.readdir(dirPath, { withFileTypes: true, recursive: true }); + for (const e of nested) { + if (!e.isFile() || !e.name.endsWith(".jsonl")) continue; + const p = path.join(e.parentPath, e.name); + childPaths.push(p); + } + childPaths.sort(); + for (const p of childPaths) { + const stat = await fs.stat(p); + mtime = Math.max(mtime, stat.mtimeMs); + } + } + if (mtime < cutoff) continue; + groups.push({ folder, id, mainPath: main.path, childPaths, mtime }); + } + } + groups.sort((a, b) => b.mtime - a.mtime); + if (opts.limit !== undefined) groups.length = Math.min(groups.length, opts.limit); + return groups; +} + +async function scanGroup(d: DiscoveredGroup): Promise { + const main = await scanFile(d.mainPath); + if (!main) return undefined; + const children: FileScan[] = []; + for (const p of d.childPaths) { + const child = await scanFile(p); + if (child) children.push(child); + } + const usage = emptyUsage(); + const subUsage = emptyUsage(); + addUsage(usage, main.usage); + for (const c of children) { + addUsage(usage, c.usage); + addUsage(subUsage, c.usage); + } + if (usage.requests === 0) return undefined; // header-only session, never used + return { folder: d.folder, id: d.id, mtime: d.mtime, main, children, usage, subUsage }; +} + +/** Run `fn` over `items` with bounded concurrency, preserving order. */ +async function mapPool(items: readonly T[], limit: number, fn: (item: T, index: number) => Promise): Promise { + const out = new Array(items.length); + let next = 0; + const workers = Array.from({ length: Math.min(limit, items.length) }, async () => { + while (next < items.length) { + const i = next++; + out[i] = await fn(items[i], i); + } + }); + await Promise.all(workers); + return out; +} + +// -------------------------------------------------------------------------- +// Formatting helpers + +function fmtTok(n: number): string { + if (n >= 1e9) return `${(n / 1e9).toFixed(2)}B`; + if (n >= 1e6) return `${(n / 1e6).toFixed(1)}M`; + if (n >= 1e3) return `${(n / 1e3).toFixed(1)}k`; + return String(Math.round(n)); +} + +function fmtMoney(n: number): string { + return `$${n.toFixed(2)}`; +} + +function fmtPct(part: number, whole: number): string { + if (whole <= 0) return "0%"; + return `${((part / whole) * 100).toFixed(1)}%`; +} + +function fmtDur(ms: number): string { + if (ms <= 0) return "0m"; + const m = Math.round(ms / 60000); + if (m < 60) return `${m}m`; + return `${Math.floor(m / 60)}h${m % 60 ? `${m % 60}m` : ""}`; +} + +function pad(s: string, w: number): string { + return s.length >= w ? s : s + " ".repeat(w - s.length); +} + +function padl(s: string, w: number): string { + return s.length >= w ? s : " ".repeat(w - s.length) + s; +} + +// -------------------------------------------------------------------------- +// Digest builder (classifier input) + +function toolLine(name: string, agg: ToolAgg): string { + const err = agg.errors ? ` errors=${agg.errors}` : ""; + return `${name}: ${agg.calls} calls, args~${fmtTok(agg.argToks)}, results~${fmtTok(agg.resultToks)}, residency~${fmtTok(agg.residency)}${err}`; +} + +function mergeToolAggs(scans: FileScan[]): Map { + const merged = new Map(); + for (const s of scans) { + for (const [name, agg] of s.toolAgg) { + const m = merged.get(name); + if (m) { + m.calls += agg.calls; + m.argToks += agg.argToks; + m.resultToks += agg.resultToks; + m.errors += agg.errors; + m.residency += agg.residency; + } else { + merged.set(name, { ...agg }); + } + } + } + return merged; +} + +/** Strip a `-2`/`-3` retry suffix from a subagent file stem. */ +function baseLabel(stem: string): string { + return stem.replace(/-\d+$/, ""); +} + +/** How a (sub)agent transcript ended, for digests. A child ending on a tool + * call is normal — its report flows back through the task result channel. */ +function endedStr(s: FileScan): string { + if (s.lastAssistantText) return `"${s.lastAssistantText}"`; + if (s.lastToolName) return `(no final text; last tool: ${s.lastToolName})`; + return `(no output)`; +} + +function buildDigest(g: SessionGroup): string { + const lines: string[] = []; + const m = g.main; + const models = [...m.models.entries()].map(([id, n]) => `${id}×${n}`).join(", "); + const cacheable = m.usage.input + m.usage.cacheRead; + const wall = m.lastTs - m.firstTs; + lines.push(`# SESSION ${g.id}`); + lines.push(`title: ${m.title ?? "(untitled)"}`); + lines.push(`project folder: ${g.folder}`); + lines.push(`models: ${models || "?"}`); + lines.push(`wall time: ${fmtDur(wall)}; user turns: ${m.turns.filter(t => !t.synthetic).length}`); + lines.push( + `MAIN context totals: ${m.usage.requests} requests, billed-in ${fmtTok(m.usage.input + m.usage.cacheRead + m.usage.cacheWrite)} ` + + `(cache-read ${fmtPct(m.usage.cacheRead, cacheable)}), out ${fmtTok(m.usage.output)}, cost ${fmtMoney(m.usage.cost)}`, + ); + lines.push( + `context peak: ${fmtTok(m.contextPeak)} tok; compactions: ${m.compactions}; assistant errors: ${m.asstErrors}`, + ); + lines.push( + `SUBAGENTS: ${g.children.length} runs, cost ${fmtMoney(g.subUsage.cost)} (${fmtPct(g.subUsage.cost, g.usage.cost)} of session), ` + + `billed ${fmtTok(billedTokens(g.subUsage))} tok`, + ); + + lines.push(`\n## Turn flow (main context)`); + const t0 = m.turns[0]?.ts ?? m.firstTs; + const shown = m.turns.slice(0, 40); + for (let i = 0; i < shown.length; i++) { + const t = shown[i]; + const toolStr = + [...t.tools.entries()] + .sort((a, b) => b[1] - a[1]) + .slice(0, 6) + .map(([n, c]) => `${n}×${c}`) + .join(" ") || "none"; + const spawnStr = t.spawnAgents.length ? ` | spawns: ${t.spawnAgents.join(",")}` : ""; + const syn = t.synthetic ? " [synthetic/steering]" : ""; + lines.push( + `T${i + 1} +${fmtDur(t.ts - t0)}${syn} [${fmtTok(t.tokens)}t] "${t.text}"` + + `\n → ${t.requests} req | tools: ${toolStr} | out ${fmtTok(t.outToks)} | ${fmtMoney(t.cost)}${spawnStr}`, + ); + } + if (m.turns.length > shown.length) lines.push(`… ${m.turns.length - shown.length} more turns`); + + lines.push(`\n## Tool traffic in main context (token counts are ~estimates)`); + const sortedTools = [...m.toolAgg.entries()].sort((a, b) => b[1].resultToks - a[1].resultToks); + for (const [name, agg] of sortedTools.slice(0, 14)) lines.push(toolLine(name, agg)); + + const repeats = [...m.readCounts.entries()] + .filter(([, r]) => r.count >= 3) + .sort((a, b) => b[1].residency - a[1].residency) + .slice(0, 8); + if (repeats.length) { + lines.push(`\n## Repeated reads of the same file (waste signal)`); + for (const [p, r] of repeats) lines.push(`${p} ×${r.count} (~${fmtTok(r.toks)}tok total, ~${fmtTok(r.residency)} residency)`); + } + + if (m.topResults.length) { + lines.push(`\n## Largest single tool results in main context`); + for (const r of m.topResults.slice(0, 8)) lines.push(`${r.tool} "${r.summary}" → ~${fmtTok(r.toks)} tok`); + } + + if (m.editCalls) { + lines.push(`\n## Edits: ${m.editCalls} calls, ${m.editErrors} failed`); + } + + // Spawn ↔ child linkage + const childByLabel = new Map(); + for (const c of g.children) { + const key = baseLabel(c.stem); + const list = childByLabel.get(key) ?? []; + list.push(c); + childByLabel.set(key, list); + } + const linked = new Set(); + if (m.spawns.length) { + lines.push(`\n## Task spawns from main context`); + for (const spawn of m.spawns.slice(0, 24)) { + const head = `task(agent=${spawn.agent}) prompt~${fmtTok(spawn.argToks)} → merged result~${fmtTok(spawn.resultToks)}${spawn.isError ? " [ERRORED]" : ""}`; + lines.push(head); + for (let i = 0; i < spawn.labels.length; i++) { + const label = spawn.labels[i]; + const kids = childByLabel.get(label) ?? []; + const kid = kids.find(k => !linked.has(k)) ?? kids[0]; + let childStr = "child log missing"; + if (kid) { + linked.add(kid); + childStr = + `child: ${kid.usage.requests} req, billed ${fmtTok(billedTokens(kid.usage))}, ${fmtMoney(kid.usage.cost)}, ` + + `${fmtDur(kid.lastTs - kid.firstTs)}, ended: ${endedStr(kid)}`; + } + lines.push(` - ${label}: "${spawn.descriptions[i] ?? ""}" | ${childStr}`); + } + if (spawn.resultSnippet) lines.push(` merged result snippet: "${spawn.resultSnippet}"`); + } + if (m.spawns.length > 24) lines.push(`… ${m.spawns.length - 24} more spawn calls`); + } + const unlinked = g.children.filter(c => !linked.has(c)); + if (unlinked.length) { + lines.push(`\n## Other subagent runs (eval agent()/irc/etc., not matched to a task call)`); + for (const c of unlinked.slice(0, 16)) { + lines.push( + `${c.stem}: ${c.usage.requests} req, billed ${fmtTok(billedTokens(c.usage))}, ${fmtMoney(c.usage.cost)}, ended: ${endedStr(c)}`, + ); + } + if (unlinked.length > 16) lines.push(`… ${unlinked.length - 16} more`); + } + + let digest = lines.join("\n"); + if (digest.length > 26000) digest = `${digest.slice(0, 26000)}\n…[digest truncated]`; + return digest; +} + +// -------------------------------------------------------------------------- +// Classifier + +interface SpawnVerdict { + label: string; + verdict: "good" | "unnecessary" | "wrong-granularity" | "context-transfer-failure" | "failed"; + why: string; +} + +interface WasteItem { + source: string; + estTokens: number; + estUsd: number; + fix: string; +} + +interface SessionVerdict { + score: number; + multiTopic: boolean; + topics: string[]; + shouldHaveSplit: boolean; + handoffOpportunities: string[]; + spawnVerdicts: SpawnVerdict[]; + waste: WasteItem[]; + headline: string; +} + +const SESSION_SCHEMA = { + type: "object", + additionalProperties: false, + properties: { + score: { type: "integer", minimum: 0, maximum: 10, description: "token-efficiency score for this session" }, + multiTopic: { type: "boolean" }, + topics: { type: "array", maxItems: 5, items: { type: "string" } }, + shouldHaveSplit: { type: "boolean", description: "true when separate chats/handoff would have saved tokens" }, + handoffOpportunities: { + type: "array", + maxItems: 4, + items: { type: "string" }, + description: "specific turns/moments where a fresh session, /handoff, or a subagent would have been cheaper", + }, + spawnVerdicts: { + type: "array", + maxItems: 10, + items: { + type: "object", + additionalProperties: false, + properties: { + label: { type: "string" }, + verdict: { + type: "string", + enum: ["good", "unnecessary", "wrong-granularity", "context-transfer-failure", "failed"], + }, + why: { type: "string" }, + }, + required: ["label", "verdict", "why"], + }, + }, + waste: { + type: "array", + maxItems: 5, + items: { + type: "object", + additionalProperties: false, + properties: { + source: { type: "string" }, + estTokens: { type: "integer", description: "rough wasted tokens attributable to this source" }, + estUsd: { + type: "number", + description: "realistic dollars this waste cost — what a leaner workflow would have saved", + }, + fix: { type: "string" }, + }, + required: ["source", "estTokens", "estUsd", "fix"], + }, + description: "biggest sources of waste, largest first", + }, + headline: { type: "string", description: "one-sentence takeaway for this session" }, + }, + required: ["score", "multiTopic", "topics", "shouldHaveSplit", "handoffOpportunities", "spawnVerdicts", "waste", "headline"], +} as const; + +interface AggregateFindings { + systemicIssues: { issue: string; evidence: string; fix: string }[]; + quickWins: string[]; + summary: string; +} + +const AGGREGATE_SCHEMA = { + type: "object", + additionalProperties: false, + properties: { + systemicIssues: { + type: "array", + maxItems: 6, + items: { + type: "object", + additionalProperties: false, + properties: { + issue: { type: "string" }, + evidence: { type: "string", description: "which sessions/numbers support this" }, + fix: { type: "string", description: "concrete habit or workflow change" }, + }, + required: ["issue", "evidence", "fix"], + }, + }, + quickWins: { type: "array", maxItems: 5, items: { type: "string" } }, + summary: { type: "string" }, + }, + required: ["systemicIssues", "quickWins", "summary"], +} as const; + +function validateSessionVerdict(v: SessionVerdict): string | undefined { + if (typeof v.headline !== "string" || !v.headline.trim()) return "headline missing or empty"; + if (typeof v.score !== "number" || !Number.isFinite(v.score)) return "score is not a finite number"; + if (!Array.isArray(v.waste)) return "waste is not an array"; + if (!Array.isArray(v.spawnVerdicts)) return "spawnVerdicts is not an array"; + if (!Array.isArray(v.topics)) return "topics is not an array"; + if (!Array.isArray(v.handoffOpportunities)) return "handoffOpportunities is not an array"; + return undefined; +} + +function validateAggregate(a: AggregateFindings): string | undefined { + if (typeof a.summary !== "string" || !a.summary.trim()) return "summary missing or empty"; + if (!Array.isArray(a.systemicIssues)) return "systemicIssues is not an array"; + return undefined; +} + +function finiteNumber(n: unknown, fallback: number): number { + return typeof n === "number" && Number.isFinite(n) ? n : fallback; +} + +/** Clamp/default every field the renderer touches so `undefined` can never + * reach the report, and sort waste by dollars desc (tokens desc tie-break). + * Also applied to cached verdicts, which may predate schema changes. */ +function normalizeVerdict(v: SessionVerdict): SessionVerdict { + const waste = (Array.isArray(v.waste) ? v.waste : []) + .map(w => ({ + source: typeof w.source === "string" && w.source ? w.source : "(unspecified)", + estTokens: finiteNumber(w.estTokens, 0), + estUsd: finiteNumber(w.estUsd, 0), + fix: typeof w.fix === "string" ? w.fix : "", + })) + .sort((a, b) => b.estUsd - a.estUsd || b.estTokens - a.estTokens); + const spawnVerdicts = (Array.isArray(v.spawnVerdicts) ? v.spawnVerdicts : []).map(s => ({ + label: typeof s.label === "string" ? s.label : "?", + verdict: s.verdict, + why: typeof s.why === "string" ? s.why : "", + })); + return { + score: Math.min(10, Math.max(0, Math.round(finiteNumber(v.score, 0)))), + multiTopic: v.multiTopic === true, + topics: (Array.isArray(v.topics) ? v.topics : []).map(String), + shouldHaveSplit: v.shouldHaveSplit === true, + handoffOpportunities: (Array.isArray(v.handoffOpportunities) ? v.handoffOpportunities : []).map(String), + spawnVerdicts, + waste, + headline: typeof v.headline === "string" ? v.headline : "", + }; +} + +interface Classifier { + model: Model; + apiKey: string; +} + +async function openClassifier(modelSpec: string): Promise { + const slash = modelSpec.indexOf("/"); + if (slash <= 0) throw new Error(`--model must be /, got "${modelSpec}"`); + const provider = modelSpec.slice(0, slash); + const modelId = modelSpec.slice(slash + 1); + const model = getBundledModel(provider as GeneratedProvider, modelId); + if (!model) throw new Error(`unknown model "${modelSpec}" (not in bundled catalog)`); + const store = await SqliteAuthCredentialStore.open(getAgentDbPath()); + const storage = new AuthStorage(store); + await storage.reload(); + const apiKey = await storage.getApiKey(provider); + if (!apiKey) { + throw new Error(`no credentials for provider "${provider}" (omp login or env var required)`); + } + return { model, apiKey }; +} + +async function completeStructured( + cls: Classifier, + prompt: string, + schema: Record, + validate: (value: T) => string | undefined, +): Promise<{ value: T; usage: UsageTotals }> { + const respond: Tool = { + name: "respond", + description: "Return your analysis by calling this tool with the requested structured fields.", + parameters: schema as Tool["parameters"], + strict: false, + }; + let lastError = ""; + for (let attempt = 0; attempt < 3; attempt++) { + const response = await completeSimple( + cls.model, + { + systemPrompt: [SYSTEM_PROMPT], + messages: [{ role: "user", content: [{ type: "text", text: prompt }], timestamp: Date.now() }], + tools: [respond], + }, + { + apiKey: cls.apiKey, + toolChoice: { type: "tool", name: "respond" }, + disableReasoning: true, + }, + ); + if (response.stopReason === "error" || response.stopReason === "aborted") { + lastError = response.errorMessage ?? response.stopReason; + await Bun.sleep(1500 * (attempt + 1)); + continue; + } + const call = response.content.find((c): c is ToolCall => c.type === "toolCall" && c.name === "respond"); + if (!call) { + lastError = "model returned no structured tool call"; + continue; + } + const value = call.arguments as T; + const problem = validate(value); + if (problem !== undefined) { + lastError = `invalid structured response: ${problem}`; + continue; + } + return { value, usage: usageOf(response) }; + } + throw new Error(`classifier call failed: ${lastError}`); +} + +function usageOf(response: AssistantMessageLike): UsageTotals { + const u = response.usage; + return { + input: u.input, + output: u.output, + cacheRead: u.cacheRead, + cacheWrite: u.cacheWrite, + cost: u.cost.total, + requests: 1, + }; +} + +/** Narrow view of pi-ai's AssistantMessage used here (content + usage). */ +interface AssistantMessageLike { + content: (ToolCall | { type: string })[]; + usage: { input: number; output: number; cacheRead: number; cacheWrite: number; cost: { total: number } }; + stopReason: string; + errorMessage?: string; +} + +// -------------------------------------------------------------------------- +// Verdict cache + +interface VerdictCacheEntry { + verdict: SessionVerdict; + model: string; + ts: number; +} + +interface VerdictCache { + entries: Record; +} + +async function loadVerdictCache(): Promise { + try { + const parsed = (await Bun.file(CACHE_PATH).json()) as Partial | null; + if (parsed && typeof parsed === "object" && parsed.entries && typeof parsed.entries === "object") { + return { entries: parsed.entries }; + } + } catch (err) { + if (!isEnoent(err)) process.stderr.write(`verdict cache unreadable, starting fresh (${CACHE_PATH})\n`); + } + return { entries: {} }; +} + +/** Persist the cache, pruned to the newest 500 entries by timestamp. */ +async function saveVerdictCache(cache: VerdictCache): Promise { + const newest = Object.entries(cache.entries) + .sort((a, b) => b[1].ts - a[1].ts) + .slice(0, 500); + await Bun.write(CACHE_PATH, JSON.stringify({ entries: Object.fromEntries(newest) })); +} + +/** Digest + system-prompt hashes make staleness automatic: any change to the + * session transcript, digest format, model, or prompt misses the cache. */ +function verdictCacheKey(groupId: string, digest: string, model: string): string { + return `${groupId}:${Bun.hash(digest).toString(16)}:${model}:${Bun.hash(SYSTEM_PROMPT).toString(16)}`; +} + +// -------------------------------------------------------------------------- +// Report + +interface AuditResult { + windowMs: number; + groups: SessionGroup[]; + verdicts: Map; + aggregate?: AggregateFindings; + classifierUsage: UsageTotals; + classifierModel?: string; +} + +function printScanReport(res: AuditResult): void { + const { groups } = res; + const total = emptyUsage(); + const sub = emptyUsage(); + let files = 0; + let compactions = 0; + for (const g of groups) { + addUsage(total, g.usage); + addUsage(sub, g.subUsage); + files += 1 + g.children.length; + compactions += g.main.compactions; + } + const days = res.windowMs / 86_400_000; + console.log(`\nSESSION AUDIT — last ${days >= 1 ? `${days.toFixed(0)}d` : fmtDur(res.windowMs)}`); + console.log(`corpus: ${groups.length} sessions (${files} jsonl files)`); + console.log( + `spend (nominal): ${fmtMoney(total.cost)} | billed ${fmtTok(billedTokens(total))} tok ` + + `(in ${fmtTok(total.input)}, cache-read ${fmtTok(total.cacheRead)}, cache-write ${fmtTok(total.cacheWrite)}, out ${fmtTok(total.output)})`, + ); + console.log( + `subagent share: ${fmtPct(sub.cost, total.cost)} of cost (${fmtMoney(sub.cost)}), ` + + `${fmtPct(billedTokens(sub), billedTokens(total))} of tokens, ${fmtPct(sub.requests, total.requests)} of requests`, + ); + console.log(`compactions in main contexts: ${compactions}`); + + // Folder split + const byFolder = new Map(); + for (const g of groups) { + let rec = byFolder.get(g.folder); + if (!rec) { + rec = { usage: emptyUsage(), sub: emptyUsage(), n: 0 }; + byFolder.set(g.folder, rec); + } + addUsage(rec.usage, g.usage); + addUsage(rec.sub, g.subUsage); + rec.n++; + } + console.log(`\nby project folder (top 12 by cost):`); + const folders = [...byFolder.entries()].sort((a, b) => b[1].usage.cost - a[1].usage.cost).slice(0, 12); + for (const [folder, rec] of folders) { + console.log( + ` ${pad(clip(folder, 44), 46)} ${padl(fmtMoney(rec.usage.cost), 9)} ${padl(fmtTok(billedTokens(rec.usage)), 8)} tok ` + + `${padl(String(rec.n), 4)} sess sub ${fmtPct(rec.sub.cost, rec.usage.cost)}`, + ); + } + + // Tool traffic across everything + const allScans: FileScan[] = []; + for (const g of groups) { + allScans.push(g.main, ...g.children); + } + const tools = mergeToolAggs(allScans); + console.log(`\ntool traffic, all contexts (arg/result tokens are ~estimates):`); + console.log( + ` ${pad("tool", 16)} ${padl("calls", 7)} ${padl("argTok", 9)} ${padl("resTok", 9)} ${padl("res/call", 9)} ${padl("errs", 6)} ${padl("residency", 11)}`, + ); + const toolRows = [...tools.entries()].sort((a, b) => b[1].resultToks - a[1].resultToks).slice(0, 16); + for (const [name, agg] of toolRows) { + console.log( + ` ${pad(clip(name, 15), 16)} ${padl(String(agg.calls), 7)} ${padl(fmtTok(agg.argToks), 9)} ${padl(fmtTok(agg.resultToks), 9)} ` + + `${padl(fmtTok(agg.calls ? agg.resultToks / agg.calls : 0), 9)} ${padl(String(agg.errors), 6)} ${padl(fmtTok(agg.residency), 11)}`, + ); + } + + // Biggest single results anywhere + const allTop: TopResult[] = []; + for (const s of allScans) allTop.push(...s.topResults); + allTop.sort((a, b) => b.toks - a.toks); + if (allTop.length) { + console.log(`\nlargest single tool results (corpus-wide):`); + for (const r of allTop.slice(0, 10)) { + console.log(` ~${padl(fmtTok(r.toks), 7)} ${r.tool} ${r.summary}`); + } + } + + // Top sessions + console.log(`\ntop sessions by cost:`); + const top = [...groups].sort((a, b) => b.usage.cost - a.usage.cost).slice(0, 15); + for (const g of top) { + const flags: string[] = []; + if (g.main.compactions) flags.push(`${g.main.compactions} compactions`); + if (g.main.asstErrors) flags.push(`${g.main.asstErrors} errors`); + const flagStr = flags.length ? ` [${flags.join(", ")}]` : ""; + console.log( + ` ${padl(fmtMoney(g.usage.cost), 8)} ${pad(new Date(g.main.firstTs).toISOString().slice(0, 16), 17)} ` + + `${pad(clip(g.folder, 26), 27)} "${clip(g.main.title ?? g.id, 46)}" sub ${fmtPct(g.subUsage.cost, g.usage.cost)}${flagStr}`, + ); + } +} + +function printVerdicts(res: AuditResult): void { + if (!res.verdicts.size) return; + console.log(`\n${"─".repeat(72)}`); + console.log(`LLM analysis (${res.classifierModel}) — ${res.verdicts.size} sessions`); + const ordered = res.groups.filter(g => res.verdicts.has(g.id)).sort((a, b) => b.usage.cost - a.usage.cost); + for (const g of ordered) { + const v = res.verdicts.get(g.id); + if (!v) continue; + console.log(`\n[${fmtMoney(g.usage.cost)}] "${clip(g.main.title ?? g.id, 60)}" (${g.folder}) — score ${v.score}/10`); + console.log(` ${v.headline}`); + if (v.multiTopic) console.log(` topics: ${v.topics.join(" | ")}${v.shouldHaveSplit ? " → should have split" : ""}`); + for (const h of v.handoffOpportunities) console.log(` handoff: ${h}`); + for (const s of v.spawnVerdicts) { + if (s.verdict === "good") continue; + console.log(` spawn ${s.label}: ${s.verdict} — ${s.why}`); + } + for (const w of v.waste) { + console.log(` waste ~${fmtMoney(w.estUsd)} (~${fmtTok(w.estTokens)}tok): ${w.source} → ${w.fix}`); + } + } +} + +function printAggregate(res: AuditResult): void { + const agg = res.aggregate; + if (!agg) return; + console.log(`\n${"─".repeat(72)}`); + console.log(`SYSTEMIC FINDINGS`); + console.log(`\n${agg.summary}`); + for (let i = 0; i < agg.systemicIssues.length; i++) { + const s = agg.systemicIssues[i]; + console.log(`\n${i + 1}. ${s.issue}`); + console.log(` evidence: ${s.evidence}`); + console.log(` fix: ${s.fix}`); + } + if (agg.quickWins.length) { + console.log(`\nquick wins:`); + for (const q of agg.quickWins) console.log(` - ${q}`); + } + console.log( + `\nclassifier spend: ${fmtMoney(res.classifierUsage.cost)} (${res.classifierUsage.requests} calls, ` + + `in ${fmtTok(res.classifierUsage.input + res.classifierUsage.cacheRead + res.classifierUsage.cacheWrite)}, out ${fmtTok(res.classifierUsage.output)})`, + ); +} + +// -------------------------------------------------------------------------- +// JSON export + +function exportJson(res: AuditResult): Record { + return { + windowMs: res.windowMs, + classifierModel: res.classifierModel, + classifierUsage: res.classifierUsage, + aggregate: res.aggregate, + sessions: res.groups.map(g => ({ + id: g.id, + folder: g.folder, + title: g.main.title, + startedAt: g.main.firstTs, + usage: g.usage, + subUsage: g.subUsage, + contextPeak: g.main.contextPeak, + compactions: g.main.compactions, + turns: g.main.turns.length, + spawns: g.main.spawns.map(s => ({ + agent: s.agent, + labels: s.labels, + resultToks: s.resultToks, + isError: s.isError, + })), + children: g.children.map(c => ({ + label: c.stem, + usage: c.usage, + requests: c.usage.requests, + })), + tools: Object.fromEntries(g.main.toolAgg), + verdict: res.verdicts.get(g.id), + })), + }; +} + +// -------------------------------------------------------------------------- +// Main + +async function main(): Promise { + const opts = parseCli(process.argv.slice(2)); + + process.stderr.write(`discovering sessions under ${SESSIONS_ROOT} …\n`); + const discovered = await discoverGroups(opts); + process.stderr.write(`scanning ${discovered.length} session groups …\n`); + + let done = 0; + const scanned = await mapPool(discovered, 8, async d => { + const g = await scanGroup(d); + done++; + if (done % 50 === 0) process.stderr.write(` ${done}/${discovered.length}\n`); + return g; + }); + const groups = scanned.filter((g): g is SessionGroup => g !== undefined); + + const res: AuditResult = { + windowMs: opts.since, + groups, + verdicts: new Map(), + classifierUsage: emptyUsage(), + }; + + printScanReport(res); + + if (!opts.noLlm && groups.length) { + const sessionFilter = opts.session?.toLowerCase(); + const matched = sessionFilter + ? groups.filter( + g => g.id.toLowerCase().includes(sessionFilter) || (g.main.title ?? "").toLowerCase().includes(sessionFilter), + ) + : groups.filter(g => g.usage.cost >= opts.minCost); + const candidates = matched.sort((a, b) => b.usage.cost - a.usage.cost).slice(0, opts.maxLlm); + if (!candidates.length) { + console.log( + sessionFilter + ? `\n(no sessions matching "${opts.session}"; skipping LLM analysis)` + : `\n(no sessions ≥ ${fmtMoney(opts.minCost)}; skipping LLM analysis)`, + ); + } else { + const cls = await openClassifier(opts.model); + res.classifierModel = `${cls.model.provider}/${cls.model.id}`; + process.stderr.write(`\nclassifying ${candidates.length} sessions with ${res.classifierModel} …\n`); + if (opts.digestDir) await fs.mkdir(opts.digestDir, { recursive: true }); + + const cache = await loadVerdictCache(); + let cacheHits = 0; + await mapPool(candidates, opts.concurrency, async g => { + const digest = buildDigest(g); + if (opts.digestDir) { + await Bun.write(path.join(opts.digestDir, `${g.id}.md`), digest); + } + const key = verdictCacheKey(g.id, digest, res.classifierModel ?? opts.model); + if (!opts.noCache) { + const hit = cache.entries[key]; + if (hit) { + res.verdicts.set(g.id, normalizeVerdict(hit.verdict)); + cacheHits++; + process.stderr.write(` ✓ ${clip(g.main.title ?? g.id, 50)} (cached)\n`); + return; + } + } + try { + const { value, usage } = await completeStructured( + cls, + digest, + SESSION_SCHEMA, + validateSessionVerdict, + ); + const verdict = normalizeVerdict(value); + res.verdicts.set(g.id, verdict); + cache.entries[key] = { verdict, model: res.classifierModel ?? opts.model, ts: Date.now() }; + addUsage(res.classifierUsage, usage); + process.stderr.write(` ✓ ${clip(g.main.title ?? g.id, 50)} (score ${verdict.score})\n`); + } catch (err) { + process.stderr.write(` ✗ ${clip(g.main.title ?? g.id, 50)}: ${err instanceof Error ? err.message : err}\n`); + } + }); + await saveVerdictCache(cache); + process.stderr.write(`${res.verdicts.size}/${candidates.length} verdicts (${cacheHits} from cache)\n`); + + printVerdicts(res); + + if (res.verdicts.size >= 2) { + const total = emptyUsage(); + const sub = emptyUsage(); + for (const g of groups) { + addUsage(total, g.usage); + addUsage(sub, g.subUsage); + } + const parts: string[] = [ + `# AGGREGATE across ${groups.length} sessions, window ${fmtDur(opts.since)}`, + `total nominal cost ${fmtMoney(total.cost)}; subagent share ${fmtPct(sub.cost, total.cost)}`, + `\nPer-session data (JSON, one per line):`, + ]; + const round2 = (n: number): number => Math.round(n * 100) / 100; + for (const g of groups) { + const v = res.verdicts.get(g.id); + if (!v) continue; + parts.push( + JSON.stringify({ + id: g.id, + title: g.main.title ?? g.id, + costUsd: round2(g.usage.cost), + subagentPct: round2(g.usage.cost > 0 ? (g.subUsage.cost / g.usage.cost) * 100 : 0), + requests: g.usage.requests, + contextPeak: g.main.contextPeak, + compactions: g.main.compactions, + score: v.score, + headline: v.headline, + topics: v.topics, + shouldHaveSplit: v.shouldHaveSplit, + spawnIssues: v.spawnVerdicts.filter(s => s.verdict !== "good").map(s => ({ label: s.label, verdict: s.verdict })), + waste: v.waste.map(w => ({ source: w.source, estTokens: w.estTokens, estUsd: round2(w.estUsd) })), + }), + ); + } + parts.push( + `\nProduce the cross-session aggregate: systemic issues (recurring patterns, not one-offs), quick wins, and a short summary addressed to the user. Cite only sessions and figures present in the data above.`, + ); + try { + const { value, usage } = await completeStructured( + cls, + parts.join("\n"), + AGGREGATE_SCHEMA, + validateAggregate, + ); + res.aggregate = value; + addUsage(res.classifierUsage, usage); + } catch (err) { + process.stderr.write(`aggregate pass failed: ${err instanceof Error ? err.message : err}\n`); + } + printAggregate(res); + } + } + } + + if (opts.json) { + await Bun.write(opts.json, JSON.stringify(exportJson(res), null, 1)); + console.log(`\nwrote ${opts.json}`); + } +} + +if (import.meta.main) { + await main(); +} From 70fc147d3f2926fd2cfc29cc1ab3359c29848479 Mon Sep 17 00:00:00 2001 From: can1357 Date: Wed, 10 Jun 2026 11:20:26 +0200 Subject: [PATCH 04/27] feat(coding-agent): enabled startup input capture with queued replay and async submit handling - Added pre-TUI startup input capture that queued non-empty Enter submissions. - Replayed queued startup submissions after initialization and injected startup editor text into InteractiveMode. - Handled startup Ctrl+C to clear typed text and Ctrl+D to exit, with detach restoring terminal state. - Extended TUI Editor `onSubmit` to accept `void | Promise` for async handlers. --- packages/coding-agent/CHANGELOG.md | 12 ++ packages/coding-agent/src/main.ts | 76 ++++--- .../src/modes/interactive-mode.ts | 5 +- .../coding-agent/src/modes/startup-input.ts | 192 ++++++++++++++++++ .../coding-agent/test/startup-input.test.ts | 146 +++++++++++++ packages/tui/CHANGELOG.md | 5 +- packages/tui/src/components/editor.ts | 2 +- 7 files changed, 409 insertions(+), 29 deletions(-) create mode 100644 packages/coding-agent/src/modes/startup-input.ts create mode 100644 packages/coding-agent/test/startup-input.test.ts diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 0183d097f..3de17230b 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -1,8 +1,20 @@ # Changelog ## [Unreleased] + +### Added + +- Added pre-TUI startup input capture so users could type while interactive sessions initialize and keep their draft while the application loads +- Added queued submission replay so Enter presses made before startup completion are submitted automatically once interactive mode begins + +### Changed + +- Changed interactive startup to carry the startup editor state into the live prompt so text entered during splash is preserved in the editor when the TUI takes over +- Changed model-scope display during startup to appear as an in-UI information notification instead of a direct stdout line + ### Fixed +- Fixed startup Ctrl+C handling in pre-TUI mode so it now clears typed text before exiting on a second press - Fixed npm CLI distribution bundles by embedding the stats dashboard client bundle so dashboard assets are served in prebuilt installs - Fixed the CLI smoke-test command to start the stats server and verify dashboard HTML is served, catching bundled-asset regressions - Added verification of a `
` and `index.js` in smoke-test dashboard responses diff --git a/packages/coding-agent/src/main.ts b/packages/coding-agent/src/main.ts index 89f60ca4a..b4b17195b 100644 --- a/packages/coding-agent/src/main.ts +++ b/packages/coding-agent/src/main.ts @@ -51,10 +51,10 @@ import { ExtensionRunner } from "./extensibility/extensions/runner"; import type { ExtensionUIContext } from "./extensibility/extensions/types"; import { scheduleMarketplaceAutoUpdate } from "./extensibility/plugins/marketplace-auto-update"; import type { MCPManager } from "./mcp"; -import { WelcomeComponent } from "./modes/components/welcome"; import { InteractiveMode } from "./modes/interactive-mode"; import type { PrintModeOptions } from "./modes/print-mode"; import { CURRENT_SETUP_VERSION } from "./modes/setup-version"; +import { StartupInput } from "./modes/startup-input"; import { initTheme, stopThemeWatcher } from "./modes/theme/theme"; import type { SubmittedUserInput } from "./modes/types"; import { @@ -97,33 +97,31 @@ function maybeShowStartupSplash(options: { modelName?: string; providerName?: string; lspServers?: LspStartupServerInfo[]; -}): void { - if (!options.isInteractive) return; - if (options.resuming || options.quiet) return; - if ($env.PI_TIMING) return; - if (!process.stdin.isTTY || !process.stdout.isTTY) return; +}): StartupInput | undefined { + if (!options.isInteractive) return undefined; + if (options.resuming || options.quiet) return undefined; + if ($env.PI_TIMING) return undefined; + if (!process.stdin.isTTY || !process.stdout.isTTY) return undefined; // First-run launches go straight into the setup wizard, which paints its own // splash — keep the minimal two-line notice there. if (options.setupPending) { process.stdout.write(`${chalk.dim(`omp ${options.version}`)}\n${chalk.dim("Initializing session…")}\n`); - return; + return undefined; } - // Render the same welcome box the TUI paints first: recent sessions as a - // loading placeholder (the fixed slot count keeps the box height stable) and - // the logo held on the intro animation's first frame so the in-TUI intro - // continues from the frame shown here. Clearing the screen first puts the - // box at the same origin the TUI's first full paint (clearScrollback) uses, - // so the live welcome replaces this frame in place without shifting. - const welcome = new WelcomeComponent( - options.version, - options.modelName ?? "", - options.providerName ?? "", - null, - options.lspServers ?? [], - ); - welcome.holdIntroFirstFrame(); - const lines = welcome.render(process.stdout.columns || 80); - process.stdout.write(`\x1b[2J\x1b[H\x1b[3J\n${lines.join("\n")}\n`); + // Paint the same welcome box the TUI paints first (recent sessions as a + // loading placeholder, logo held on the intro animation's first frame) plus + // a live editor, and start capturing raw stdin so the user can type — and + // even submit — while the session loads in the background. The TUI's first + // full paint (clearScrollback) replaces this frame in place; the editor + // instance itself is handed to InteractiveMode so nothing typed is lost. + const startupInput = new StartupInput({ + version: options.version, + modelName: options.modelName ?? "", + providerName: options.providerName ?? "", + lspServers: options.lspServers ?? [], + }); + startupInput.start(); + return startupInput; } async function checkForNewVersion(currentVersion: string): Promise { @@ -372,6 +370,7 @@ async function runInteractiveMode( initialMessage?: string, initialImages?: ImageContent[], titleSystemPrompt?: string, + startupInput?: StartupInput, ): Promise { const mode = new InteractiveMode( session, @@ -382,6 +381,7 @@ async function runInteractiveMode( mcpManager, eventBus, titleSystemPrompt, + startupInput?.editor, ); // Cold-launch gate: the full setup wizard (every scene + the overlay and @@ -400,6 +400,12 @@ async function runInteractiveMode( }) : []; + // Hand the terminal over: stop the pre-TUI capture (restoring cooked mode so + // ProcessTerminal records the correct prior raw state) right before the TUI + // grabs stdin. Keystrokes typed during init's awaits stay OS-buffered and + // flow into the TUI once it resumes stdin. + startupInput?.detach(); + await mode.init({ suppressWelcomeIntro: resuming || setupScenes.length > 0, clearInitialTerminalHistory: true, @@ -460,8 +466,19 @@ async function runInteractiveMode( } } + const startupSubmissions = [...(startupInput?.queuedSubmissions ?? [])]; while (true) { - const input = await mode.getUserInput(); + const inputPromise = mode.getUserInput(); + const queuedText = startupSubmissions.shift(); + if (queuedText !== undefined) { + // Replay through the real submit pipeline (slash commands, bash/python + // modes, title generation) exactly as if Enter were pressed now. + await mode.editor.onSubmit?.(queuedText); + // Handled inline (e.g. a slash command) without consuming the pending + // input wait — replay the next queued item on the next iteration. + if (mode.onInputCallback) continue; + } + const input = await inputPromise; await submitInteractiveInput(mode, session, input); } } @@ -1264,7 +1281,7 @@ export async function runRootCommand( settingsInstance.get("lsp.lazy") ? "available" : "connecting", ) : []; - maybeShowStartupSplash({ + const startupInput = maybeShowStartupSplash({ isInteractive, resuming: Boolean(parsedArgs.continue || parsedArgs.resume || parsedArgs.fork), quiet: settingsInstance.get("startup.quiet"), @@ -1275,6 +1292,9 @@ export async function runRootCommand( lspServers: splashLspServers, }); + // TEMP-SMOKE: stretch the startup gap for VHS verification. REMOVE. + if (process.env.PI_DEBUG_SLOW_START) await Bun.sleep(Number(process.env.PI_DEBUG_SLOW_START)); + const { session, setToolUIContext, modelFallbackMessage, lspServers, mcpManager } = await createSession({ ...sessionOptions, eventBus, @@ -1322,7 +1342,10 @@ export async function runRootCommand( return `${scopedModel.model.id}${thinkingStr}`; }) .join(", "); - process.stdout.write(`${chalk.dim(`Model scope: ${modelList} ${chalk.gray("(Ctrl+P to cycle)")}`)}\n`); + // Routed through the TUI (not stdout): the startup capture owns the + // terminal in raw mode here, and the TUI's first clearScrollback paint + // would wipe a pre-TUI line anyway. + notifs.push({ kind: "info", message: `Model scope: ${modelList} (Ctrl+P to cycle)` }); } if ($env.PI_TIMING) { @@ -1350,6 +1373,7 @@ export async function runRootCommand( initialMessage, initialImages, titleSystemPrompt, + startupInput, ); } else { // Branch-only single-shot runner: keep print-mode code out of normal interactive startup. diff --git a/packages/coding-agent/src/modes/interactive-mode.ts b/packages/coding-agent/src/modes/interactive-mode.ts index b1cfec78d..238634188 100644 --- a/packages/coding-agent/src/modes/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive-mode.ts @@ -388,6 +388,7 @@ export class InteractiveMode implements InteractiveModeContext { mcpManager?: MCPManager, eventBus?: EventBus, titleSystemPrompt?: string, + startupEditor?: CustomEditor, ) { this.session = session; this.sessionManager = session.sessionManager; @@ -422,7 +423,9 @@ export class InteractiveMode implements InteractiveModeContext { this.btwContainer = new Container(); this.omfgContainer = new Container(); this.errorBannerContainer = new Container(); - this.editor = new CustomEditor(getEditorTheme()); + // Adopt the pre-TUI startup editor when provided: typed text, cursor, + // paste buffers, and undo history carry over from the splash phase. + this.editor = startupEditor ?? new CustomEditor(getEditorTheme()); this.editor.setUseTerminalCursor(this.ui.getShowHardwareCursor()); this.editor.setAutocompleteMaxVisible(settings.get("autocompleteMaxVisible")); this.editor.onAutocompleteCancel = () => { diff --git a/packages/coding-agent/src/modes/startup-input.ts b/packages/coding-agent/src/modes/startup-input.ts new file mode 100644 index 000000000..1bba19721 --- /dev/null +++ b/packages/coding-agent/src/modes/startup-input.ts @@ -0,0 +1,192 @@ +import { StdinBuffer, truncateToWidth } from "@oh-my-pi/pi-tui"; +import { postmortem } from "@oh-my-pi/pi-utils"; +import { CustomEditor } from "./components/custom-editor"; +import { type LspServerInfo, WelcomeComponent } from "./components/welcome"; +import { getEditorTheme, theme } from "./theme/theme"; + +/** Synchronized-output guards (DEC 2026); unsupported terminals ignore them. */ +const SYNC_BEGIN = "\x1b[?2026h"; +const SYNC_END = "\x1b[?2026l"; + +export interface StartupInputOptions { + version: string; + modelName: string; + providerName: string; + lspServers: LspServerInfo[]; +} + +/** + * Pre-TUI live input phase. Paints the same frame the TUI's first full paint + * will produce — welcome box held on the intro's first frame, blank chat area, + * editor box — and runs a real {@link CustomEditor} against raw stdin while + * session creation continues in the background. + * + * The editor instance is handed to InteractiveMode at construction, so typed + * text, cursor position, paste buffers, and undo history carry seamlessly into + * the live UI. Enter submissions made before the session is ready are queued + * (rendered dimmed in the chat area, where the real transcript will appear) + * and replayed through the real submit pipeline once the input loop starts. + * + * Handoff contract: {@link detach} must run before `ProcessTerminal.start()` + * grabs stdin — it restores cooked mode (so the terminal records the correct + * prior raw state) and pauses stdin, leaving OS-buffered keystrokes to flow + * into the TUI once it resumes. + */ +export class StartupInput { + readonly editor: CustomEditor; + readonly #welcome: WelcomeComponent; + #queued: string[] = []; + #stdinBuffer: StdinBuffer | undefined; + #dataListener: ((chunk: string) => void) | undefined; + #resizeListener: (() => void) | undefined; + #unregisterCleanup: (() => void) | undefined; + #started = false; + #detached = false; + #wasRaw = false; + + constructor(options: StartupInputOptions) { + this.#welcome = new WelcomeComponent( + options.version, + options.modelName, + options.providerName, + null, + options.lspServers, + ); + // Freeze the logo on the intro's first frame so the in-TUI intro picks up + // exactly where this frame leaves off. + this.#welcome.holdIntroFirstFrame(); + + this.editor = new CustomEditor(getEditorTheme()); + // `Editor.#submitValue` expands paste markers, trims, and clears the + // buffer before invoking onSubmit, so `text` is final plain text. + this.editor.onSubmit = text => { + if (text) this.#queued.push(text); + this.#paintLiveRegion(); + }; + // Ctrl+C: clear typed text; on an empty editor abort startup (pre-TUI raw + // mode swallows SIGINT, so this is the muscle-memory escape hatch). + this.editor.onClear = () => { + if (this.editor.getText()) { + this.editor.setText(""); + this.#paintLiveRegion(); + } else { + this.#exit(130); + } + }; + // Ctrl+D: same exit semantics as the live UI on an idle session. + this.editor.onExit = () => this.#exit(0); + } + + /** Enter submissions captured before the session was ready, in order. */ + get queuedSubmissions(): readonly string[] { + return this.#queued; + } + + /** Grab stdin (raw mode), paint the initial frame, and start echoing input. */ + start(): void { + if (this.#started) return; + this.#started = true; + this.#wasRaw = process.stdin.isRaw === true; + process.stdin.setRawMode?.(true); + process.stdin.setEncoding("utf8"); + process.stdin.resume(); + + // Same sequence-splitting pipeline ProcessTerminal uses, so the editor + // receives single key events and bracketed pastes arrive re-wrapped. + const buffer = new StdinBuffer({ timeout: 50 }); + buffer.on("data", sequence => this.feedInput(sequence)); + buffer.on("paste", content => this.feedInput(`\x1b[200~${content}\x1b[201~`)); + this.#stdinBuffer = buffer; + this.#dataListener = chunk => buffer.process(chunk); + process.stdin.on("data", this.#dataListener); + this.#resizeListener = () => this.#paintFull(); + process.stdout.on("resize", this.#resizeListener); + // Crash safety: a fatal error before handoff must not leave the user's + // terminal in raw mode with a hidden cursor. + this.#unregisterCleanup = postmortem.register("startup-input-restore", () => this.#restoreTerminal()); + + // Bracketed paste on; hardware cursor off (the editor draws its own). + process.stdout.write("\x1b[?2004h\x1b[?25l"); + this.#paintFull(); + } + + /** Route one complete input sequence into the editor and refresh the frame. */ + feedInput(sequence: string): void { + this.editor.handleInput(sequence); + this.#paintLiveRegion(); + } + + /** + * Stop capturing and hand the terminal to the TUI. Restores cooked mode so + * `ProcessTerminal.start()` records the correct prior state, and pauses + * stdin so keystrokes typed during the remaining init await flow into the + * TUI once it resumes. The painted frame is left in place — the TUI's first + * full paint replaces it at the same origin. + */ + detach(): void { + this.#unregisterCleanup?.(); + this.#unregisterCleanup = undefined; + this.editor.onSubmit = undefined; + this.editor.onClear = undefined; + this.editor.onExit = undefined; + this.#restoreTerminal(); + } + + #restoreTerminal(): void { + if (!this.#started || this.#detached) return; + this.#detached = true; + if (this.#dataListener) process.stdin.off("data", this.#dataListener); + if (this.#resizeListener) process.stdout.off("resize", this.#resizeListener); + this.#stdinBuffer?.removeAllListeners(); + process.stdin.pause(); + process.stdin.setRawMode?.(this.#wasRaw); + process.stdout.write("\x1b[?2004l\x1b[?25h"); + } + + #exit(code: number): void { + this.detach(); + process.stdout.write("\r\n"); + void postmortem.quit(code); + } + + /** + * Compose the full frame, mirroring the TUI's first-paint layout: Spacer, + * welcome box, Spacer, chat area (queued submissions), hook Spacer, editor. + * `liveRegionIndex` marks the first row that changes with input; everything + * above it is the stable welcome prefix. + */ + #frameRows(width: number): { rows: string[]; liveRegionIndex: number } { + const rows: string[] = ["", ...this.#welcome.render(width), ""]; + const liveRegionIndex = rows.length; + for (const text of this.#queued) { + rows.push(truncateToWidth(theme.fg("dim", ` › ${text.replace(/\s+/g, " ")}`), Math.max(0, width - 1))); + } + rows.push(""); + const terminalRows = process.stdout.rows || 24; + this.editor.setMaxHeight(Math.max(3, Math.min(10, terminalRows - rows.length - 2))); + rows.push(...this.editor.render(width)); + return { rows, liveRegionIndex }; + } + + #paintFull(): void { + if (!this.#started || this.#detached) return; + const width = process.stdout.columns || 80; + const { rows } = this.#frameRows(width); + // Raw mode disables ONLCR; emit explicit CR+LF between rows. + process.stdout.write(`${SYNC_BEGIN}\x1b[2J\x1b[H\x1b[3J${rows.join("\r\n")}${SYNC_END}`); + } + + #paintLiveRegion(): void { + if (!this.#started || this.#detached) return; + const width = process.stdout.columns || 80; + const { rows, liveRegionIndex } = this.#frameRows(width); + if (rows.length >= (process.stdout.rows || 24)) { + // Frame taller than the viewport: the initial write scrolled, so + // absolute row addressing no longer maps to the frame. Repaint all. + this.#paintFull(); + return; + } + const region = rows.slice(liveRegionIndex).join("\r\n"); + process.stdout.write(`${SYNC_BEGIN}\x1b[${liveRegionIndex + 1};1H\x1b[0J${region}${SYNC_END}`); + } +} diff --git a/packages/coding-agent/test/startup-input.test.ts b/packages/coding-agent/test/startup-input.test.ts new file mode 100644 index 000000000..e4e2f0d16 --- /dev/null +++ b/packages/coding-agent/test/startup-input.test.ts @@ -0,0 +1,146 @@ +import { afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "bun:test"; +import * as path from "node:path"; +import { Agent } from "@oh-my-pi/pi-agent-core"; +import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry"; +import { resetSettingsForTest, Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; +import { InteractiveMode } from "@oh-my-pi/pi-coding-agent/modes/interactive-mode"; +import { StartupInput } from "@oh-my-pi/pi-coding-agent/modes/startup-input"; +import { initTheme } from "@oh-my-pi/pi-coding-agent/modes/theme/theme"; +import { AgentSession } from "@oh-my-pi/pi-coding-agent/session/agent-session"; +import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage"; +import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager"; +import { TempDir } from "@oh-my-pi/pi-utils"; + +function createStartupInput(): StartupInput { + return new StartupInput({ + version: "0.0.0-test", + modelName: "Test Model", + providerName: "test", + lspServers: [], + }); +} + +beforeAll(() => { + initTheme(); +}); + +// `start()` is never called: the capture stays detached from stdin/stdout, so +// these tests exercise the pure input → editor → queue contract. +describe("StartupInput pre-TUI capture", () => { + it("queues Enter submissions in order and clears the editor between them", () => { + const startup = createStartupInput(); + + for (const ch of "hello") startup.feedInput(ch); + expect(startup.editor.getText()).toBe("hello"); + + startup.feedInput("\r"); + expect([...startup.queuedSubmissions]).toEqual(["hello"]); + expect(startup.editor.getText()).toBe(""); + + for (const ch of "world") startup.feedInput(ch); + startup.feedInput("\r"); + expect([...startup.queuedSubmissions]).toEqual(["hello", "world"]); + }); + + it("ignores Enter on an empty editor", () => { + const startup = createStartupInput(); + startup.feedInput("\r"); + expect(startup.queuedSubmissions).toHaveLength(0); + }); + + it("queues paste placeholders fully expanded, ready for the real submit pipeline", () => { + const startup = createStartupInput(); + const pasted = Array.from({ length: 12 }, (_, i) => `line ${i + 1}`).join("\n"); + + startup.feedInput(`\x1b[200~${pasted}\x1b[201~`); + // Large pastes collapse to an atomic placeholder in the editor buffer... + expect(startup.editor.getText()).toContain("[Paste #1"); + + startup.feedInput("\r"); + // ...but the queued submission must be final plain text: the splash + // editor's paste map is gone by the time the queue is replayed. + expect(startup.queuedSubmissions[0]).toBe(pasted); + }); + + it("clears typed text on Ctrl+C without queueing it", () => { + const startup = createStartupInput(); + for (const ch of "draft") startup.feedInput(ch); + + startup.feedInput("\x03"); + expect(startup.editor.getText()).toBe(""); + expect(startup.queuedSubmissions).toHaveLength(0); + }); + + it("releases its editor handlers on detach so the live UI owns submission", () => { + const startup = createStartupInput(); + expect(startup.editor.onSubmit).toBeDefined(); + + startup.detach(); + expect(startup.editor.onSubmit).toBeUndefined(); + expect(startup.editor.onClear).toBeUndefined(); + expect(startup.editor.onExit).toBeUndefined(); + }); +}); + +describe("InteractiveMode startup editor adoption", () => { + let tempDir: TempDir; + let authStorage: AuthStorage; + let session: AgentSession; + let mode: InteractiveMode; + + beforeEach(async () => { + resetSettingsForTest(); + tempDir = TempDir.createSync("@pi-startup-input-"); + await Settings.init({ inMemory: true, cwd: tempDir.path() }); + authStorage = await AuthStorage.create(path.join(tempDir.path(), "testauth.db")); + const modelRegistry = new ModelRegistry(authStorage); + const model = modelRegistry.find("anthropic", "claude-sonnet-4-5"); + if (!model) { + throw new Error("Expected claude-sonnet-4-5 to exist in registry"); + } + + session = new AgentSession({ + agent: new Agent({ + initialState: { + model, + systemPrompt: ["Test"], + tools: [], + messages: [], + }, + }), + sessionManager: SessionManager.create(tempDir.path(), tempDir.path()), + settings: Settings.isolated(), + modelRegistry, + }); + }); + + afterEach(async () => { + vi.restoreAllMocks(); + mode?.stop(); + await session?.dispose(); + authStorage?.close(); + tempDir?.removeSync(); + resetSettingsForTest(); + }); + + it("adopts the splash editor instance so text typed during startup carries over", () => { + const startup = createStartupInput(); + for (const ch of "typed while loading") startup.feedInput(ch); + startup.detach(); + + mode = new InteractiveMode( + session, + "test", + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + startup.editor, + ); + + expect(mode.editor).toBe(startup.editor); + expect(mode.editor.getText()).toBe("typed while loading"); + }); +}); diff --git a/packages/tui/CHANGELOG.md b/packages/tui/CHANGELOG.md index 391c03a3c..b34b92467 100644 --- a/packages/tui/CHANGELOG.md +++ b/packages/tui/CHANGELOG.md @@ -1,6 +1,9 @@ # Changelog ## [Unreleased] +### Added + +- Added support for asynchronous `onSubmit` handlers by allowing the callback to return a `Promise` ## [15.10.11] - 2026-06-10 @@ -1273,4 +1276,4 @@ Initial release under @oh-my-pi scope. See previous releases at [badlogic/pi-mon ### Fixed -- **Readline-style Ctrl+W**: Now skips trailing whitespace before deleting the preceding word, matching standard readline behavior. ([#306](https://github.com/badlogic/pi-mono/pull/306) by [@kim0](https://github.com/kim0)) +- **Readline-style Ctrl+W**: Now skips trailing whitespace before deleting the preceding word, matching standard readline behavior. ([#306](https://github.com/badlogic/pi-mono/pull/306) by [@kim0](https://github.com/kim0)) \ No newline at end of file diff --git a/packages/tui/src/components/editor.ts b/packages/tui/src/components/editor.ts index 35a5e40fb..cb8cd5d61 100644 --- a/packages/tui/src/components/editor.ts +++ b/packages/tui/src/components/editor.ts @@ -450,7 +450,7 @@ export class Editor implements Component, Focusable { // Debounce timer for autocomplete updates #autocompleteTimeout?: NodeJS.Timeout; - onSubmit?: (text: string) => void; + onSubmit?: (text: string) => void | Promise; onAltEnter?: (text: string) => void; onChange?: (text: string) => void; onAutocompleteCancel?: () => void; From 1c885471bcb500f092ca466ad7636e31150b76ba Mon Sep 17 00:00:00 2001 From: can1357 Date: Wed, 10 Jun 2026 17:40:50 +0200 Subject: [PATCH 05/27] fix(utils): unescaped \n/\t in {{join}} prompt helper separator --- packages/utils/CHANGELOG.md | 4 ++++ packages/utils/src/prompt.ts | 4 +++- packages/utils/test/prompt.test.ts | 14 ++++++++++++++ 3 files changed, 21 insertions(+), 1 deletion(-) diff --git a/packages/utils/CHANGELOG.md b/packages/utils/CHANGELOG.md index 5bb5266df..0326fafa8 100644 --- a/packages/utils/CHANGELOG.md +++ b/packages/utils/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- Fixed the `{{join}}` prompt helper joining with a literal two-character `\n` when templates pass `"\n"` as the separator — Handlebars string literals carry no escape processing. The separator now unescapes `\n`/`\t`, matching the `{{#list}}` helper's documented convention (visible as literal `\n` between paths in compaction `` lists). + ## [15.10.11] - 2026-06-10 ### Added diff --git a/packages/utils/src/prompt.ts b/packages/utils/src/prompt.ts index c175c5e78..d0bd02ad8 100644 --- a/packages/utils/src/prompt.ts +++ b/packages/utils/src/prompt.ts @@ -342,10 +342,12 @@ handlebars.registerHelper( /** * {{join array ", "}} * Joins an array with a separator (default: ", "). + * Note: Use \n/\t in the separator for newlines/tabs (unescaped automatically, + * same convention as {{#list}} — Handlebars string literals carry no escapes). */ handlebars.registerHelper("join", (context: unknown[], separator?: unknown): string => { if (!Array.isArray(context)) return ""; - const sep = typeof separator === "string" ? separator : ", "; + const sep = typeof separator === "string" ? separator.replace(/\\n/g, "\n").replace(/\\t/g, "\t") : ", "; return context.join(sep); }); diff --git a/packages/utils/test/prompt.test.ts b/packages/utils/test/prompt.test.ts index 494b8d323..7dd5e2126 100644 --- a/packages/utils/test/prompt.test.ts +++ b/packages/utils/test/prompt.test.ts @@ -94,3 +94,17 @@ describe("compile cache", () => { expect(prompt.render("{{#if a}}{ {{b}}}{{/if}}", { a: true, b: "v" })).toBe("{ v}"); }); }); + +describe("helpers: join", () => { + it('unescapes \\n and \\t in the separator (Handlebars string literals carry no escapes)', () => { + // Regression: `{{join files "\n"}}` used to emit the literal two-char `\n` + // between entries (visible in compaction lists). + expect(prompt.render('{{join files "\\n"}}', { files: ["a.ts", "b.ts"] })).toBe("a.ts\nb.ts"); + expect(prompt.render('{{join files "\\t"}}', { files: ["a.ts", "b.ts"] })).toBe("a.ts\tb.ts"); + }); + + it("defaults to comma-space and tolerates non-arrays", () => { + expect(prompt.render("{{join files}}", { files: ["a", "b"] })).toBe("a, b"); + expect(prompt.render("{{join files}}", { files: "not-an-array" })).toBe(""); + }); +}); From 67ddc7a17f53b47584d444d5bbe137b5930f6010 Mon Sep 17 00:00:00 2001 From: can1357 Date: Wed, 10 Jun 2026 17:41:12 +0200 Subject: [PATCH 06/27] feat(natives): added renderSnapcompactPng with bundled 5x8 BDF font --- Cargo.lock | 1 + Cargo.toml | 1 + crates/pi-natives/Cargo.toml | 1 + crates/pi-natives/src/fonts/5x8.bdf | 2897 ++++++++++++++++++++++++++ crates/pi-natives/src/lib.rs | 1 + crates/pi-natives/src/snapcompact.rs | 244 +++ packages/natives/CHANGELOG.md | 4 + packages/natives/native/index.d.ts | 10 + packages/natives/native/index.js | 1 + types/assets/index.d.ts | 5 + 10 files changed, 3165 insertions(+) create mode 100644 crates/pi-natives/src/fonts/5x8.bdf create mode 100644 crates/pi-natives/src/snapcompact.rs diff --git a/Cargo.lock b/Cargo.lock index 0891c360b..73899fa6c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2436,6 +2436,7 @@ dependencies = [ "pi-ast", "pi-iso", "pi-shell", + "png", "portable-pty", "rayon", "regex", diff --git a/Cargo.toml b/Cargo.toml index 4a0095746..b2cca1af0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -274,6 +274,7 @@ image = { version = "0.25", default-features = false, features = [ "gif", "webp", ] } +png = "0.18" inferno = { version = "0.12", default-features = false } syntect = { version = "5.3", default-features = false, features = [ "default-syntaxes", diff --git a/crates/pi-natives/Cargo.toml b/crates/pi-natives/Cargo.toml index 1137ff346..5a570185d 100644 --- a/crates/pi-natives/Cargo.toml +++ b/crates/pi-natives/Cargo.toml @@ -36,6 +36,7 @@ pi-ast.workspace = true pi-iso.workspace = true pi-shell.workspace = true portable-pty.workspace = true +png.workspace = true rayon.workspace = true regex.workspace = true serde.workspace = true diff --git a/crates/pi-natives/src/fonts/5x8.bdf b/crates/pi-natives/src/fonts/5x8.bdf new file mode 100644 index 000000000..0fd5096d6 --- /dev/null +++ b/crates/pi-natives/src/fonts/5x8.bdf @@ -0,0 +1,2897 @@ +STARTFONT 2.1 +COMMENT $ucs-fonts: 5x8.bdf,v 1.32 2006-01-05 20:03:17+00 mgk25 Rel $ +COMMENT Send bug reports to Markus Kuhn +FONT -Misc-Fixed-Medium-R-Normal--8-80-75-75-C-50-ISO10646-1 +SIZE 11 75 75 +FONTBOUNDINGBOX 5 8 0 -1 +STARTPROPERTIES 22 +FONTNAME_REGISTRY "" +FOUNDRY "Misc" +FAMILY_NAME "Fixed" +WEIGHT_NAME "Medium" +SLANT "R" +SETWIDTH_NAME "Normal" +ADD_STYLE_NAME "" +PIXEL_SIZE 8 +POINT_SIZE 80 +RESOLUTION_X 75 +RESOLUTION_Y 75 +SPACING "C" +AVERAGE_WIDTH 50 +CHARSET_REGISTRY "ISO10646" +CHARSET_ENCODING "1" +FONT_DESCENT 1 +FONT_ASCENT 7 +COPYRIGHT "Public domain font. Share and enjoy." +DEFAULT_CHAR 0 +_XMBDFED_INFO "Edited with xmbdfed 4.5." +CAP_HEIGHT 6 +X_HEIGHT 4 +ENDPROPERTIES +CHARS 191 +STARTCHAR space +ENCODING 32 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +00 +00 +00 +00 +00 +00 +ENDCHAR +STARTCHAR exclam +ENCODING 33 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +20 +20 +20 +20 +00 +20 +00 +ENDCHAR +STARTCHAR quotedbl +ENCODING 34 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +50 +50 +50 +00 +00 +00 +00 +ENDCHAR +STARTCHAR numbersign +ENCODING 35 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +50 +50 +F8 +50 +F8 +50 +50 +00 +ENDCHAR +STARTCHAR dollar +ENCODING 36 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +20 +70 +A0 +70 +28 +70 +20 +00 +ENDCHAR +STARTCHAR percent +ENCODING 37 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +40 +50 +20 +50 +10 +00 +00 +ENDCHAR +STARTCHAR ampersand +ENCODING 38 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +40 +A0 +A0 +40 +A0 +A0 +50 +00 +ENDCHAR +STARTCHAR quotesingle +ENCODING 39 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +20 +20 +20 +00 +00 +00 +00 +ENDCHAR +STARTCHAR parenleft +ENCODING 40 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +20 +40 +40 +40 +40 +20 +00 +ENDCHAR +STARTCHAR parenright +ENCODING 41 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +40 +20 +20 +20 +20 +40 +00 +ENDCHAR +STARTCHAR asterisk +ENCODING 42 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +90 +60 +F0 +60 +90 +00 +ENDCHAR +STARTCHAR plus +ENCODING 43 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +20 +20 +F8 +20 +20 +00 +ENDCHAR +STARTCHAR comma +ENCODING 44 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +00 +00 +00 +30 +20 +40 +ENDCHAR +STARTCHAR hyphen +ENCODING 45 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +00 +00 +F0 +00 +00 +00 +ENDCHAR +STARTCHAR period +ENCODING 46 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +00 +00 +00 +20 +70 +20 +ENDCHAR +STARTCHAR slash +ENCODING 47 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +10 +10 +20 +40 +80 +80 +00 +ENDCHAR +STARTCHAR zero +ENCODING 48 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +20 +50 +50 +50 +50 +20 +00 +ENDCHAR +STARTCHAR one +ENCODING 49 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +20 +60 +20 +20 +20 +70 +00 +ENDCHAR +STARTCHAR two +ENCODING 50 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +60 +90 +10 +60 +80 +F0 +00 +ENDCHAR +STARTCHAR three +ENCODING 51 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +F0 +20 +60 +10 +90 +60 +00 +ENDCHAR +STARTCHAR four +ENCODING 52 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +20 +60 +A0 +F0 +20 +20 +00 +ENDCHAR +STARTCHAR five +ENCODING 53 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +F0 +80 +E0 +10 +90 +60 +00 +ENDCHAR +STARTCHAR six +ENCODING 54 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +60 +80 +E0 +90 +90 +60 +00 +ENDCHAR +STARTCHAR seven +ENCODING 55 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +F0 +10 +20 +20 +40 +40 +00 +ENDCHAR +STARTCHAR eight +ENCODING 56 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +60 +90 +60 +90 +90 +60 +00 +ENDCHAR +STARTCHAR nine +ENCODING 57 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +60 +90 +90 +70 +10 +60 +00 +ENDCHAR +STARTCHAR colon +ENCODING 58 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +60 +60 +00 +60 +60 +00 +ENDCHAR +STARTCHAR semicolon +ENCODING 59 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +30 +30 +00 +30 +20 +40 +ENDCHAR +STARTCHAR less +ENCODING 60 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +10 +20 +40 +40 +20 +10 +00 +ENDCHAR +STARTCHAR equal +ENCODING 61 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +00 +F0 +00 +F0 +00 +00 +ENDCHAR +STARTCHAR greater +ENCODING 62 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +40 +20 +10 +10 +20 +40 +00 +ENDCHAR +STARTCHAR question +ENCODING 63 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +20 +50 +10 +20 +00 +20 +00 +ENDCHAR +STARTCHAR at +ENCODING 64 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +30 +48 +98 +A8 +A8 +90 +40 +30 +ENDCHAR +STARTCHAR A +ENCODING 65 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +60 +90 +90 +F0 +90 +90 +00 +ENDCHAR +STARTCHAR B +ENCODING 66 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +E0 +90 +E0 +90 +90 +E0 +00 +ENDCHAR +STARTCHAR C +ENCODING 67 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +60 +90 +80 +80 +90 +60 +00 +ENDCHAR +STARTCHAR D +ENCODING 68 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +E0 +90 +90 +90 +90 +E0 +00 +ENDCHAR +STARTCHAR E +ENCODING 69 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +F0 +80 +E0 +80 +80 +F0 +00 +ENDCHAR +STARTCHAR F +ENCODING 70 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +F0 +80 +E0 +80 +80 +80 +00 +ENDCHAR +STARTCHAR G +ENCODING 71 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +60 +90 +80 +B0 +90 +60 +00 +ENDCHAR +STARTCHAR H +ENCODING 72 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +90 +90 +F0 +90 +90 +90 +00 +ENDCHAR +STARTCHAR I +ENCODING 73 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +70 +20 +20 +20 +20 +70 +00 +ENDCHAR +STARTCHAR J +ENCODING 74 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +70 +20 +20 +20 +A0 +40 +00 +ENDCHAR +STARTCHAR K +ENCODING 75 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +90 +A0 +C0 +A0 +A0 +90 +00 +ENDCHAR +STARTCHAR L +ENCODING 76 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +80 +80 +80 +80 +80 +F0 +00 +ENDCHAR +STARTCHAR M +ENCODING 77 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +90 +F0 +F0 +90 +90 +90 +00 +ENDCHAR +STARTCHAR N +ENCODING 78 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +90 +D0 +F0 +B0 +B0 +90 +00 +ENDCHAR +STARTCHAR O +ENCODING 79 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +60 +90 +90 +90 +90 +60 +00 +ENDCHAR +STARTCHAR P +ENCODING 80 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +E0 +90 +90 +E0 +80 +80 +00 +ENDCHAR +STARTCHAR Q +ENCODING 81 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +60 +90 +90 +D0 +B0 +60 +10 +ENDCHAR +STARTCHAR R +ENCODING 82 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +E0 +90 +90 +E0 +90 +90 +00 +ENDCHAR +STARTCHAR S +ENCODING 83 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +60 +90 +40 +20 +90 +60 +00 +ENDCHAR +STARTCHAR T +ENCODING 84 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +70 +20 +20 +20 +20 +20 +00 +ENDCHAR +STARTCHAR U +ENCODING 85 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +90 +90 +90 +90 +90 +60 +00 +ENDCHAR +STARTCHAR V +ENCODING 86 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +90 +90 +90 +90 +60 +60 +00 +ENDCHAR +STARTCHAR W +ENCODING 87 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +90 +90 +90 +F0 +F0 +90 +00 +ENDCHAR +STARTCHAR X +ENCODING 88 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +90 +90 +60 +60 +90 +90 +00 +ENDCHAR +STARTCHAR Y +ENCODING 89 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +88 +88 +50 +20 +20 +20 +00 +ENDCHAR +STARTCHAR Z +ENCODING 90 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +F0 +10 +20 +40 +80 +F0 +00 +ENDCHAR +STARTCHAR bracketleft +ENCODING 91 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +70 +40 +40 +40 +40 +70 +00 +ENDCHAR +STARTCHAR backslash +ENCODING 92 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +80 +80 +40 +20 +10 +10 +00 +ENDCHAR +STARTCHAR bracketright +ENCODING 93 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +70 +10 +10 +10 +10 +70 +00 +ENDCHAR +STARTCHAR asciicircum +ENCODING 94 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +20 +50 +00 +00 +00 +00 +00 +ENDCHAR +STARTCHAR underscore +ENCODING 95 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +00 +00 +00 +00 +00 +F0 +ENDCHAR +STARTCHAR grave +ENCODING 96 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +40 +20 +00 +00 +00 +00 +00 +ENDCHAR +STARTCHAR a +ENCODING 97 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +00 +70 +90 +90 +70 +00 +ENDCHAR +STARTCHAR b +ENCODING 98 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +80 +80 +E0 +90 +90 +E0 +00 +ENDCHAR +STARTCHAR c +ENCODING 99 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +00 +30 +40 +40 +30 +00 +ENDCHAR +STARTCHAR d +ENCODING 100 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +10 +10 +70 +90 +90 +70 +00 +ENDCHAR +STARTCHAR e +ENCODING 101 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +00 +60 +B0 +C0 +60 +00 +ENDCHAR +STARTCHAR f +ENCODING 102 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +20 +50 +40 +E0 +40 +40 +00 +ENDCHAR +STARTCHAR g +ENCODING 103 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +00 +60 +90 +70 +10 +60 +ENDCHAR +STARTCHAR h +ENCODING 104 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +80 +80 +E0 +90 +90 +90 +00 +ENDCHAR +STARTCHAR i +ENCODING 105 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +20 +00 +60 +20 +20 +70 +00 +ENDCHAR +STARTCHAR j +ENCODING 106 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +10 +00 +10 +10 +10 +50 +20 +ENDCHAR +STARTCHAR k +ENCODING 107 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +80 +80 +90 +E0 +90 +90 +00 +ENDCHAR +STARTCHAR l +ENCODING 108 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +60 +20 +20 +20 +20 +70 +00 +ENDCHAR +STARTCHAR m +ENCODING 109 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +00 +D0 +A8 +A8 +A8 +00 +ENDCHAR +STARTCHAR n +ENCODING 110 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +00 +E0 +90 +90 +90 +00 +ENDCHAR +STARTCHAR o +ENCODING 111 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +00 +60 +90 +90 +60 +00 +ENDCHAR +STARTCHAR p +ENCODING 112 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +00 +E0 +90 +E0 +80 +80 +ENDCHAR +STARTCHAR q +ENCODING 113 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +00 +70 +90 +70 +10 +10 +ENDCHAR +STARTCHAR r +ENCODING 114 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +00 +A0 +D0 +80 +80 +00 +ENDCHAR +STARTCHAR s +ENCODING 115 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +00 +30 +60 +10 +60 +00 +ENDCHAR +STARTCHAR t +ENCODING 116 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +40 +40 +E0 +40 +50 +20 +00 +ENDCHAR +STARTCHAR u +ENCODING 117 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +00 +90 +90 +90 +70 +00 +ENDCHAR +STARTCHAR v +ENCODING 118 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +00 +50 +50 +50 +20 +00 +ENDCHAR +STARTCHAR w +ENCODING 119 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +00 +88 +A8 +A8 +50 +00 +ENDCHAR +STARTCHAR x +ENCODING 120 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +00 +90 +60 +60 +90 +00 +ENDCHAR +STARTCHAR y +ENCODING 121 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +00 +90 +90 +70 +90 +60 +ENDCHAR +STARTCHAR z +ENCODING 122 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +00 +F0 +20 +40 +F0 +00 +ENDCHAR +STARTCHAR braceleft +ENCODING 123 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +30 +40 +20 +C0 +20 +40 +30 +00 +ENDCHAR +STARTCHAR bar +ENCODING 124 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +20 +20 +20 +20 +20 +20 +00 +ENDCHAR +STARTCHAR braceright +ENCODING 125 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +C0 +20 +40 +30 +40 +20 +C0 +00 +ENDCHAR +STARTCHAR asciitilde +ENCODING 126 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +50 +A0 +00 +00 +00 +00 +00 +ENDCHAR +STARTCHAR space +ENCODING 160 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +00 +00 +00 +00 +00 +00 +ENDCHAR +STARTCHAR exclamdown +ENCODING 161 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +20 +00 +20 +20 +20 +20 +00 +ENDCHAR +STARTCHAR cent +ENCODING 162 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +20 +70 +A0 +A0 +70 +20 +ENDCHAR +STARTCHAR sterling +ENCODING 163 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +20 +50 +E0 +40 +50 +A0 +00 +ENDCHAR +STARTCHAR currency +ENCODING 164 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +88 +70 +50 +70 +88 +00 +ENDCHAR +STARTCHAR yen +ENCODING 165 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +88 +50 +F8 +20 +F8 +20 +00 +ENDCHAR +STARTCHAR brokenbar +ENCODING 166 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +20 +20 +20 +00 +20 +20 +20 +00 +ENDCHAR +STARTCHAR section +ENCODING 167 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +70 +80 +E0 +90 +70 +10 +E0 +00 +ENDCHAR +STARTCHAR dieresis +ENCODING 168 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +50 +00 +00 +00 +00 +00 +00 +ENDCHAR +STARTCHAR copyright +ENCODING 169 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +70 +A8 +C8 +C8 +A8 +70 +00 +ENDCHAR +STARTCHAR ordfeminine +ENCODING 170 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +30 +50 +30 +00 +70 +00 +00 +00 +ENDCHAR +STARTCHAR guillemotleft +ENCODING 171 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +00 +50 +A0 +50 +00 +00 +ENDCHAR +STARTCHAR logicalnot +ENCODING 172 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +00 +00 +70 +10 +10 +00 +ENDCHAR +STARTCHAR hyphen +ENCODING 173 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +00 +00 +70 +00 +00 +00 +ENDCHAR +STARTCHAR registered +ENCODING 174 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +70 +E8 +D8 +E8 +D8 +70 +00 +ENDCHAR +STARTCHAR macron +ENCODING 175 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +70 +00 +00 +00 +00 +00 +00 +ENDCHAR +STARTCHAR degree +ENCODING 176 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +20 +50 +20 +00 +00 +00 +00 +ENDCHAR +STARTCHAR plusminus +ENCODING 177 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +20 +70 +20 +00 +70 +00 +ENDCHAR +STARTCHAR twosuperior +ENCODING 178 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +20 +50 +10 +20 +70 +00 +00 +00 +ENDCHAR +STARTCHAR threesuperior +ENCODING 179 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +60 +10 +60 +10 +60 +00 +00 +00 +ENDCHAR +STARTCHAR acute +ENCODING 180 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +20 +40 +00 +00 +00 +00 +00 +ENDCHAR +STARTCHAR mu +ENCODING 181 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +00 +90 +90 +90 +E0 +80 +ENDCHAR +STARTCHAR paragraph +ENCODING 182 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +78 +E8 +E8 +68 +28 +28 +00 +ENDCHAR +STARTCHAR periodcentered +ENCODING 183 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +00 +00 +20 +00 +00 +00 +ENDCHAR +STARTCHAR cedilla +ENCODING 184 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +00 +00 +00 +00 +20 +40 +ENDCHAR +STARTCHAR onesuperior +ENCODING 185 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +20 +60 +20 +20 +70 +00 +00 +00 +ENDCHAR +STARTCHAR ordmasculine +ENCODING 186 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +20 +50 +20 +00 +70 +00 +00 +00 +ENDCHAR +STARTCHAR guillemotright +ENCODING 187 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +00 +A0 +50 +A0 +00 +00 +ENDCHAR +STARTCHAR onequarter +ENCODING 188 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +80 +80 +80 +A0 +60 +F0 +20 +00 +ENDCHAR +STARTCHAR onehalf +ENCODING 189 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +80 +80 +A0 +D0 +10 +20 +70 +00 +ENDCHAR +STARTCHAR threequarters +ENCODING 190 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +80 +40 +80 +60 +A0 +F0 +20 +00 +ENDCHAR +STARTCHAR questiondown +ENCODING 191 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +20 +00 +20 +40 +50 +20 +00 +ENDCHAR +STARTCHAR Agrave +ENCODING 192 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +40 +20 +60 +90 +F0 +90 +90 +00 +ENDCHAR +STARTCHAR Aacute +ENCODING 193 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +20 +40 +60 +90 +F0 +90 +90 +00 +ENDCHAR +STARTCHAR Acircumflex +ENCODING 194 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +60 +90 +60 +90 +F0 +90 +90 +00 +ENDCHAR +STARTCHAR Atilde +ENCODING 195 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +50 +A0 +60 +90 +F0 +90 +90 +00 +ENDCHAR +STARTCHAR Adieresis +ENCODING 196 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +90 +00 +60 +90 +F0 +90 +90 +00 +ENDCHAR +STARTCHAR Aring +ENCODING 197 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +60 +90 +60 +90 +F0 +90 +90 +00 +ENDCHAR +STARTCHAR AE +ENCODING 198 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +70 +A0 +A0 +F0 +A0 +B0 +00 +ENDCHAR +STARTCHAR Ccedilla +ENCODING 199 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +60 +90 +80 +80 +90 +60 +40 +ENDCHAR +STARTCHAR Egrave +ENCODING 200 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +40 +20 +F0 +80 +E0 +80 +F0 +00 +ENDCHAR +STARTCHAR Eacute +ENCODING 201 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +20 +40 +F0 +80 +E0 +80 +F0 +00 +ENDCHAR +STARTCHAR Ecircumflex +ENCODING 202 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +60 +90 +F0 +80 +E0 +80 +F0 +00 +ENDCHAR +STARTCHAR Edieresis +ENCODING 203 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +90 +00 +F0 +80 +E0 +80 +F0 +00 +ENDCHAR +STARTCHAR Igrave +ENCODING 204 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +40 +20 +70 +20 +20 +20 +70 +00 +ENDCHAR +STARTCHAR Iacute +ENCODING 205 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +10 +20 +70 +20 +20 +20 +70 +00 +ENDCHAR +STARTCHAR Icircumflex +ENCODING 206 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +20 +50 +70 +20 +20 +20 +70 +00 +ENDCHAR +STARTCHAR Idieresis +ENCODING 207 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +50 +00 +70 +20 +20 +20 +70 +00 +ENDCHAR +STARTCHAR Eth +ENCODING 208 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +70 +48 +E8 +48 +48 +70 +00 +ENDCHAR +STARTCHAR Ntilde +ENCODING 209 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +50 +A0 +90 +D0 +B0 +90 +90 +00 +ENDCHAR +STARTCHAR Ograve +ENCODING 210 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +40 +20 +60 +90 +90 +90 +60 +00 +ENDCHAR +STARTCHAR Oacute +ENCODING 211 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +20 +40 +60 +90 +90 +90 +60 +00 +ENDCHAR +STARTCHAR Ocircumflex +ENCODING 212 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +60 +90 +60 +90 +90 +90 +60 +00 +ENDCHAR +STARTCHAR Otilde +ENCODING 213 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +50 +A0 +60 +90 +90 +90 +60 +00 +ENDCHAR +STARTCHAR Odieresis +ENCODING 214 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +90 +00 +60 +90 +90 +90 +60 +00 +ENDCHAR +STARTCHAR multiply +ENCODING 215 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +00 +00 +50 +20 +50 +00 +ENDCHAR +STARTCHAR Oslash +ENCODING 216 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +70 +B0 +B0 +D0 +D0 +E0 +00 +ENDCHAR +STARTCHAR Ugrave +ENCODING 217 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +40 +20 +90 +90 +90 +90 +60 +00 +ENDCHAR +STARTCHAR Uacute +ENCODING 218 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +20 +40 +90 +90 +90 +90 +60 +00 +ENDCHAR +STARTCHAR Ucircumflex +ENCODING 219 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +60 +90 +90 +90 +90 +90 +60 +00 +ENDCHAR +STARTCHAR Udieresis +ENCODING 220 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +90 +00 +90 +90 +90 +90 +60 +00 +ENDCHAR +STARTCHAR Yacute +ENCODING 221 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +10 +20 +88 +50 +20 +20 +20 +00 +ENDCHAR +STARTCHAR Thorn +ENCODING 222 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +80 +E0 +90 +90 +E0 +80 +00 +ENDCHAR +STARTCHAR germandbls +ENCODING 223 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +60 +90 +A0 +A0 +90 +A0 +00 +ENDCHAR +STARTCHAR agrave +ENCODING 224 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +40 +20 +00 +70 +90 +90 +70 +00 +ENDCHAR +STARTCHAR aacute +ENCODING 225 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +20 +40 +00 +70 +90 +90 +70 +00 +ENDCHAR +STARTCHAR acircumflex +ENCODING 226 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +20 +50 +00 +70 +90 +90 +70 +00 +ENDCHAR +STARTCHAR atilde +ENCODING 227 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +50 +A0 +00 +70 +90 +90 +70 +00 +ENDCHAR +STARTCHAR adieresis +ENCODING 228 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +50 +00 +70 +90 +90 +70 +00 +ENDCHAR +STARTCHAR aring +ENCODING 229 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +60 +90 +60 +70 +90 +90 +70 +00 +ENDCHAR +STARTCHAR ae +ENCODING 230 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +00 +F0 +68 +B0 +78 +00 +ENDCHAR +STARTCHAR ccedilla +ENCODING 231 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +00 +30 +40 +40 +30 +20 +ENDCHAR +STARTCHAR egrave +ENCODING 232 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +40 +20 +00 +60 +B0 +C0 +60 +00 +ENDCHAR +STARTCHAR eacute +ENCODING 233 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +20 +40 +00 +60 +B0 +C0 +60 +00 +ENDCHAR +STARTCHAR ecircumflex +ENCODING 234 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +60 +90 +00 +60 +B0 +C0 +60 +00 +ENDCHAR +STARTCHAR edieresis +ENCODING 235 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +50 +00 +60 +B0 +C0 +60 +00 +ENDCHAR +STARTCHAR igrave +ENCODING 236 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +40 +20 +00 +60 +20 +20 +70 +00 +ENDCHAR +STARTCHAR iacute +ENCODING 237 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +10 +20 +00 +60 +20 +20 +70 +00 +ENDCHAR +STARTCHAR icircumflex +ENCODING 238 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +20 +50 +00 +60 +20 +20 +70 +00 +ENDCHAR +STARTCHAR idieresis +ENCODING 239 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +50 +00 +60 +20 +20 +70 +00 +ENDCHAR +STARTCHAR eth +ENCODING 240 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +A0 +40 +A0 +10 +70 +90 +60 +00 +ENDCHAR +STARTCHAR ntilde +ENCODING 241 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +50 +A0 +00 +E0 +90 +90 +90 +00 +ENDCHAR +STARTCHAR ograve +ENCODING 242 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +40 +20 +00 +60 +90 +90 +60 +00 +ENDCHAR +STARTCHAR oacute +ENCODING 243 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +20 +40 +00 +60 +90 +90 +60 +00 +ENDCHAR +STARTCHAR ocircumflex +ENCODING 244 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +60 +90 +00 +60 +90 +90 +60 +00 +ENDCHAR +STARTCHAR otilde +ENCODING 245 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +50 +A0 +00 +60 +90 +90 +60 +00 +ENDCHAR +STARTCHAR odieresis +ENCODING 246 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +90 +00 +60 +90 +90 +60 +00 +ENDCHAR +STARTCHAR divide +ENCODING 247 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +20 +00 +70 +00 +20 +00 +ENDCHAR +STARTCHAR oslash +ENCODING 248 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +00 +00 +70 +B0 +D0 +E0 +00 +ENDCHAR +STARTCHAR ugrave +ENCODING 249 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +40 +20 +00 +90 +90 +90 +70 +00 +ENDCHAR +STARTCHAR uacute +ENCODING 250 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +20 +40 +00 +90 +90 +90 +70 +00 +ENDCHAR +STARTCHAR ucircumflex +ENCODING 251 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +60 +90 +00 +90 +90 +90 +70 +00 +ENDCHAR +STARTCHAR udieresis +ENCODING 252 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +90 +00 +90 +90 +90 +70 +00 +ENDCHAR +STARTCHAR yacute +ENCODING 253 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +20 +40 +00 +90 +90 +70 +90 +60 +ENDCHAR +STARTCHAR thorn +ENCODING 254 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +80 +80 +E0 +90 +E0 +80 +80 +ENDCHAR +STARTCHAR ydieresis +ENCODING 255 +SWIDTH 436 0 +DWIDTH 5 0 +BBX 5 8 0 -1 +BITMAP +00 +90 +00 +90 +90 +70 +90 +60 +ENDCHAR +ENDFONT diff --git a/crates/pi-natives/src/lib.rs b/crates/pi-natives/src/lib.rs index e7c1b7124..de32b2792 100644 --- a/crates/pi-natives/src/lib.rs +++ b/crates/pi-natives/src/lib.rs @@ -36,6 +36,7 @@ pub mod highlight; pub mod html; pub mod keys; pub mod sixel; +pub mod snapcompact; pub use pi_ast::language; pub mod power; diff --git a/crates/pi-natives/src/snapcompact.rs b/crates/pi-natives/src/snapcompact.rs new file mode 100644 index 000000000..e1da346a4 --- /dev/null +++ b/crates/pi-natives/src/snapcompact.rs @@ -0,0 +1,244 @@ +//! Snapcompact frame rendering. +//! +//! Rasterizes pre-normalized conversation text onto a square 8-bit indexed +//! bitmap using the bundled public-domain X.org `5x8` BDF font, with glyph +//! ink cycling through six hues at sentence boundaries, then encodes the +//! bitmap as an indexed PNG. +//! +//! Text normalization, frame chunking, and archive management live in +//! `packages/agent/src/compaction/snapcompact.ts`; this module is only the +//! hot `text -> PNG bytes` path. + +use std::{borrow::Cow, collections::HashMap, sync::LazyLock}; + +use napi::bindgen_prelude::*; +use napi_derive::napi; + +/// Horizontal advance per glyph cell in pixels. +const GLYPH_ADVANCE_X: usize = 5; +/// Vertical pitch per text row in pixels. +const GLYPH_PITCH_Y: usize = 8; + +/// Upper bound on the frame edge: a hard stop against absurd allocations +/// (`size * size` pixel buffer), far above the 2576px production frame. +const MAX_FRAME_SIZE: u32 = 16384; + +/// Sentence ink palette: index 0 is the white background, 1-6 are the six +/// dark hues from the eval renderer (HLS l=0.22 s=0.95, h ∈ {0, .08, .3, +/// .5, .62, .78}), pre-baked to sRGB. +const PALETTE: [[u8; 3]; 7] = [ + [255, 255, 255], + [109, 2, 2], // red + [109, 53, 2], // amber + [24, 109, 2], // green + [2, 109, 109], // teal + [2, 32, 109], // blue + [75, 2, 109], // violet +]; +const INK_COLORS: usize = PALETTE.len() - 1; + +static FONT: LazyLock = LazyLock::new(|| parse_bdf(include_str!("fonts/5x8.bdf"))); + +struct Glyph { + /// Glyph width in pixels (≤ 8 for this font). + w: u8, + /// Glyph height in pixels. + h: i32, + xoff: i32, + yoff: i32, + /// One bitmask per bitmap row, MSB-leftmost. + rows: Vec, +} + +struct Font { + /// Glyphs keyed by Unicode code point (ASCII + Latin-1 coverage). + glyphs: HashMap, + ascent: i32, +} + +fn parse_bdf(text: &str) -> Font { + let mut glyphs = HashMap::new(); + let mut ascent = 0i32; + let mut enc = -1i64; + let mut bbx = [0i32; 4]; + let mut lines = text.lines(); + while let Some(line) = lines.next() { + if let Some(rest) = line.strip_prefix("FONT_ASCENT") { + ascent = rest.trim().parse().unwrap_or(0); + } else if let Some(rest) = line.strip_prefix("ENCODING") { + enc = rest.trim().parse().unwrap_or(-1); + } else if let Some(rest) = line.strip_prefix("BBX") { + let mut parts = rest.split_ascii_whitespace(); + for slot in &mut bbx { + *slot = parts.next().and_then(|part| part.parse().ok()).unwrap_or(0); + } + } else if line.starts_with("BITMAP") { + let mut rows = Vec::new(); + for row in lines.by_ref() { + if row.starts_with("ENDCHAR") { + break; + } + rows.push(u8::from_str_radix(row.trim(), 16).unwrap_or(0)); + } + if enc >= 0 { + glyphs.insert(enc as u32, Glyph { + w: bbx[0].clamp(0, 8) as u8, + h: bbx[1], + xoff: bbx[2], + yoff: bbx[3], + rows, + }); + } + } + } + Font { glyphs, ascent } +} + +/// Rasterize `text` onto a `size` x `size` palette-indexed bitmap, row-major +/// with no word wrap. Ink color advances at sentence boundaries (terminator +/// in `.!?` followed by a space). Characters beyond the grid capacity are +/// ignored; code points missing from the font leave their cell blank. +fn render_bitmap(text: &str, size: usize, font: &Font) -> Vec { + let cols = size / GLYPH_ADVANCE_X; + let grid_rows = size / GLYPH_PITCH_Y; + let capacity = cols * grid_rows; + let mut pixels = vec![0u8; size * size]; // 0 = white background + if capacity == 0 { + return pixels; + } + let codes: Vec = text.chars().map(|ch| ch as u32).collect(); + let count = codes.len().min(capacity); + let mut sentence = 0usize; + for i in 0..count { + let code = codes[i]; + let ink = (1 + sentence % INK_COLORS) as u8; + if matches!(code, 0x2e | 0x21 | 0x3f) && codes.get(i + 1) == Some(&0x20) { + sentence += 1; + } + let Some(glyph) = font.glyphs.get(&code) else { + continue; + }; + if glyph.rows.is_empty() { + continue; + } + let row = i / cols; + let col = i - row * cols; + let top = (row * GLYPH_PITCH_Y) as i32 + font.ascent - glyph.h - glyph.yoff; + let left = (col * GLYPH_ADVANCE_X) as i32 + glyph.xoff; + for (r, &bits) in glyph.rows.iter().enumerate() { + if bits == 0 { + continue; + } + let y = top + r as i32; + if y < 0 || y >= size as i32 { + continue; + } + let row_base = y as usize * size; + for b in 0..glyph.w { + if bits & (0x80u8 >> b) != 0 { + let x = left + i32::from(b); + if x >= 0 && (x as usize) < size { + pixels[row_base + x as usize] = ink; + } + } + } + } + } + pixels +} + +/// Pack one-byte-per-pixel palette indices into 4-bit PNG scanline data +/// (two pixels per byte, high nibble first). With only 7 palette entries, +/// 4-bit depth halves the pre-deflate stream vs 8-bit. +fn pack_nibbles(pixels: &[u8], size: usize) -> Vec { + let row_bytes = size.div_ceil(2); + let mut packed = vec![0u8; row_bytes * size]; + for y in 0..size { + let src = &pixels[y * size..(y + 1) * size]; + let dst = &mut packed[y * row_bytes..(y + 1) * row_bytes]; + for (x, &px) in src.iter().enumerate() { + dst[x / 2] |= px << (4 * (1 - x % 2)); + } + } + packed +} + +/// Encode a palette-indexed bitmap as a 4-bit indexed PNG with `None` row +/// filtering (the glyph bitmap is already minimal-entropy; filtering costs +/// encode time without helping deflate). +fn encode_indexed_png( + pixels: &[u8], + size: usize, + compression: png::Compression, +) -> Result> { + let mut palette = Vec::with_capacity(PALETTE.len() * 3); + for rgb in PALETTE { + palette.extend_from_slice(&rgb); + } + let mut out = Vec::new(); + let mut encoder = png::Encoder::new(&mut out, size as u32, size as u32); + encoder.set_color(png::ColorType::Indexed); + encoder.set_depth(png::BitDepth::Four); + encoder.set_palette(Cow::Owned(palette)); + encoder.set_compression(compression); + // MUST come after `set_compression`, which resets the filter to the + // compression level's default (`Adaptive` for `Balanced`). + encoder.set_filter(png::Filter::NoFilter); + let mut writer = encoder + .write_header() + .map_err(|err| Error::from_reason(format!("Failed to write PNG header: {err}")))?; + writer + .write_image_data(&pack_nibbles(pixels, size)) + .map_err(|err| Error::from_reason(format!("Failed to write PNG data: {err}")))?; + writer + .finish() + .map_err(|err| Error::from_reason(format!("Failed to finish PNG stream: {err}")))?; + Ok(out) +} + +/// Render one snapcompact frame: print pre-normalized text onto a +/// `size` x `size` 4-bit indexed bitmap and encode it as a PNG. +/// +/// The glyph grid holds `floor(size/5) * floor(size/8)` characters; input +/// beyond that is ignored (the caller chunks text to capacity). Returns the +/// PNG bytes. +#[napi] +pub fn render_snapcompact_png(text: String, size: u32) -> Result { + if size == 0 || size > MAX_FRAME_SIZE { + return Err(Error::from_reason(format!( + "Invalid frame size {size}: expected 1..={MAX_FRAME_SIZE}" + ))); + } + let pixels = render_bitmap(&text, size as usize, &FONT); + Ok(encode_indexed_png(&pixels, size as usize, png::Compression::Balanced)?.into()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn font_parses_ascii_coverage() { + let font = &*FONT; + assert_eq!(font.ascent, 7); + // Every printable ASCII char must have a glyph. + for cp in 0x20u32..0x7f { + assert!(font.glyphs.contains_key(&cp), "missing glyph for U+{cp:04X}"); + } + } + + #[test] + fn bitmap_inks_sentences_and_caps_capacity() { + let font = &*FONT; + // 40px -> 8 cols x 5 rows = 40 cells. + let pixels = render_bitmap("Hi. Ok.", 40, font); + let inks: Vec = pixels.iter().copied().filter(|&p| p != 0).collect(); + assert!(inks.contains(&1), "first sentence should use ink 1"); + assert!(inks.contains(&2), "second sentence should use ink 2"); + assert!(!inks.contains(&3), "no third sentence ink expected"); + + // Overflow input renders without panicking and stays in-bounds. + let overflow = render_bitmap(&"x".repeat(100), 40, font); + assert_eq!(overflow.len(), 40 * 40); + } +} diff --git a/packages/natives/CHANGELOG.md b/packages/natives/CHANGELOG.md index 9f31c6081..1184a6341 100644 --- a/packages/natives/CHANGELOG.md +++ b/packages/natives/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Added + +- Added `renderSnapcompactPng(text, size)`: rasterizes pre-normalized text onto a square 4-bit indexed PNG using the bundled public-domain X.org `5x8` BDF font, with glyph ink cycling through six hues at sentence boundaries. Replaces the JS rasterizer/PNG writer previously in `@oh-my-pi/pi-agent-core`; the font ships inside the crate (`crates/pi-natives/src/fonts/5x8.bdf`). + ## [15.10.12] - 2026-06-10 ### Added diff --git a/packages/natives/native/index.d.ts b/packages/natives/native/index.d.ts index 71aab8604..a61406901 100644 --- a/packages/natives/native/index.d.ts +++ b/packages/natives/native/index.d.ts @@ -1284,6 +1284,16 @@ export interface PtyStartOptions { */ export declare function readImageFromClipboard(): Promise +/** + * Render one snapcompact frame: print pre-normalized text onto a + * `size` x `size` 4-bit indexed bitmap and encode it as a PNG. + * + * The glyph grid holds `floor(size/5) * floor(size/8)` characters; input + * beyond that is ignored (the caller chunks text to capacity). Returns the + * PNG bytes. + */ +export declare function renderSnapcompactPng(text: string, size: number): Uint8Array + /** * Search content for a pattern (one-shot, compiles pattern each time). * For repeated searches with the same pattern, use [`grep`] with file filters. diff --git a/packages/natives/native/index.js b/packages/natives/native/index.js index ce26c56fa..3750385a0 100644 --- a/packages/natives/native/index.js +++ b/packages/natives/native/index.js @@ -59,6 +59,7 @@ export const matchesLegacySequence = nativeBindings.matchesLegacySequence; export const parseKey = nativeBindings.parseKey; export const parseKittySequence = nativeBindings.parseKittySequence; export const readImageFromClipboard = nativeBindings.readImageFromClipboard; +export const renderSnapcompactPng = nativeBindings.renderSnapcompactPng; export const search = nativeBindings.search; export const sliceWithWidth = nativeBindings.sliceWithWidth; export const summarizeCode = nativeBindings.summarizeCode; diff --git a/types/assets/index.d.ts b/types/assets/index.d.ts index 1e0a34718..05a576125 100644 --- a/types/assets/index.d.ts +++ b/types/assets/index.d.ts @@ -18,6 +18,11 @@ declare module "*.lark" { export default content; } +declare module "*.bdf" { + const content: string; + export default content; +} + // turndown-plugin-gfm has no published types declare module "turndown-plugin-gfm" { import type TurndownService from "turndown"; From bc5af8d9bef3dee36cd302df0a7c85a296dadeee Mon Sep 17 00:00:00 2001 From: can1357 Date: Wed, 10 Jun 2026 17:42:48 +0200 Subject: [PATCH 07/27] refactor(agent): extracted convertMessageToLlm and removed convertToLlm alias Splits the array-form defaultConvertToLlm into a single-message convertMessageToLlm that embedders can delegate every core role to instead of duplicating the conversion. Adds the optional images field on CompactionSummaryMessage so the converter attaches snapcompact frames after the summary text (snapcompact strategy lands next). Renames every convertToLlm call site in compaction.ts and branch-summarization.ts to the canonical defaultConvertToLlm.\n\nNote: the bundled CHANGELOG entries also cover the supersede-reads, snapcompact, and steering-queue fixes that follow in this batch (the entries land in directly adjacent lines and cannot be split by diff). --- packages/agent/CHANGELOG.md | 16 ++ .../src/compaction/branch-summarization.ts | 4 +- packages/agent/src/compaction/compaction.ts | 12 +- packages/agent/src/compaction/messages.ts | 142 ++++++++++-------- 4 files changed, 102 insertions(+), 72 deletions(-) diff --git a/packages/agent/CHANGELOG.md b/packages/agent/CHANGELOG.md index bff557c20..607bfc6ee 100644 --- a/packages/agent/CHANGELOG.md +++ b/packages/agent/CHANGELOG.md @@ -2,6 +2,22 @@ ## [Unreleased] +### Breaking Changes + +- Removed the `convertToLlm` alias export from `compaction/messages` — it duplicated `defaultConvertToLlm` under a second name. Import `defaultConvertToLlm` (array form) or the new `convertMessageToLlm` (single-message form) instead + +### Added + +- Added `convertMessageToLlm()`: the single-message core transformer behind `defaultConvertToLlm()`. Embedders with app-specific message roles should handle their own roles and delegate every core role (`user`/`developer`/`assistant`/`toolResult`/`custom`/`hookMessage`/`branchSummary`/`compactionSummary`) to it instead of duplicating the conversion — a duplicated `compactionSummary` case is how snapcompact frames once silently dropped off provider requests +- Added `pruneSupersededToolResults()` and the opt-in `PruneConfig.supersedeKey` hook so harnesses can prune stale tool results superseded by a newer read of the same file; superseded results are pruned ahead of age-based victims during overflow pruning and replaced with a `[Superseded by a newer read of this file]` placeholder. Without the new config, `pruneToolOutputs()` behavior is unchanged. +- Added `readToolSupersedeKey()` implementing the read-tool path/selector grammar (selector-free reads supersede range reads of the same file; URL-scheme paths exempt). Pruning honors prompt-cache economics: per-turn prunes only fire when the post-candidate suffix is small or the cache is cold (idle gap). +- Added the `snapcompact` compaction strategy (`snapcompactCompact()` in `compaction/snapcompact.ts`): instead of an LLM summary, discarded history is printed onto dense 2576px PNG frames with the public-domain X.org `5x8` pixel font (ink cycles per sentence) and re-attached to the compaction summary message as image blocks. Fully local — no model call; ~7x cheaper than raw text at near-parity recall. `CompactionSummaryMessage` gains an optional `images` field, `estimateTokens()` charges per attached frame, and frames persist under `preserveData.snapcompact` with an 8-frame budget that evicts middle-out: the session-head frame is pinned, the oldest unpinned frames drop first, so head and tail both survive. Rasterization and PNG encoding run in native code (`renderSnapcompactPng()` from `@oh-my-pi/pi-natives`), emitting 4-bit indexed PNGs + +### Fixed + +- Fixed queued steering messages being drained into an externally aborted run: interrupting mid-tool execution (e.g. Enter with a pending steer) dequeued the steer into the dying run — it landed in history without a response and the post-abort resume saw an empty queue, so the agent stopped instead of continuing. Steering/follow-up/aside queue polls are now skipped once the run's abort signal fires, leaving the queue intact for `Agent.continue()`. +- Fixed `` compaction lists recording the same file once per line-range/raw selector (`src/foo.ts:50-200`, `:raw`, `:1-50:raw`, …): read-tool selectors are now stripped before tracking, so reads dedupe to the base path and match their write/edit path when splitting read-only vs modified lists. Selector-polluted lists stored by earlier compactions self-heal on the next compaction. `readToolSupersedeKey()` now shares the same splitter (`splitReadSelector()`), gaining the `..` range alias and `L`-prefix forms it previously missed. + ## [15.10.12] - 2026-06-10 ### Added diff --git a/packages/agent/src/compaction/branch-summarization.ts b/packages/agent/src/compaction/branch-summarization.ts index 3cb73de02..c2d10b630 100644 --- a/packages/agent/src/compaction/branch-summarization.ts +++ b/packages/agent/src/compaction/branch-summarization.ts @@ -13,10 +13,10 @@ import { estimateTokens } from "./compaction"; import type { ReadonlySessionManager, SessionEntry } from "./entries"; import { type ConvertToLlm, - convertToLlm, createBranchSummaryMessage, createCompactionSummaryMessage, createCustomMessage, + defaultConvertToLlm, } from "./messages"; import branchSummaryPrompt from "./prompts/branch-summary.md" with { type: "text" }; import branchSummaryPreamble from "./prompts/branch-summary-preamble.md" with { type: "text" }; @@ -288,7 +288,7 @@ export async function generateBranchSummary( // Transform to LLM-compatible messages, then serialize to text // Serialization prevents the model from treating it as a conversation to continue - const llmMessages = (options.convertToLlm ?? convertToLlm)(messages); + const llmMessages = (options.convertToLlm ?? defaultConvertToLlm)(messages); const conversationText = serializeConversation(llmMessages); // Build prompt diff --git a/packages/agent/src/compaction/compaction.ts b/packages/agent/src/compaction/compaction.ts index 53d0ad464..1b23c2669 100644 --- a/packages/agent/src/compaction/compaction.ts +++ b/packages/agent/src/compaction/compaction.ts @@ -22,7 +22,7 @@ import { type AgentTelemetry, instrumentedCompleteSimple } from "../telemetry"; import { ThinkingLevel } from "../thinking"; import type { AgentMessage } from "../types"; import type { CompactionEntry, SessionEntry } from "./entries"; -import { type ConvertToLlm, convertToLlm, createBranchSummaryMessage, createCustomMessage } from "./messages"; +import { type ConvertToLlm, createBranchSummaryMessage, createCustomMessage, defaultConvertToLlm } from "./messages"; import { buildOpenAiNativeHistory, getPreservedOpenAiRemoteCompactionData, @@ -625,7 +625,7 @@ export async function generateSummary( // Serialize conversation to text so model doesn't try to continue it // Convert to LLM messages first (handles custom app messages when caller provides a transformer). - const llmMessages = (options?.convertToLlm ?? convertToLlm)(currentMessages); + const llmMessages = (options?.convertToLlm ?? defaultConvertToLlm)(currentMessages); const conversationText = serializeConversation(llmMessages); // Build the prompt with conversation wrapped in tags @@ -724,7 +724,7 @@ export async function generateHandoff( options: HandoffOptions, signal?: AbortSignal, ): Promise { - const llmMessages = (options.convertToLlm ?? convertToLlm)(messages); + const llmMessages = (options.convertToLlm ?? defaultConvertToLlm)(messages); const requestMessages: Message[] = [ ...llmMessages, { @@ -773,7 +773,7 @@ async function generateShortSummary( options?: SummaryOptions, ): Promise { const maxTokens = Math.min(512, Math.floor(0.2 * reserveTokens)); - const llmMessages = (options?.convertToLlm ?? convertToLlm)(recentMessages); + const llmMessages = (options?.convertToLlm ?? defaultConvertToLlm)(recentMessages); const conversationText = serializeConversation(llmMessages); let promptText = `\n${conversationText}\n\n\n`; @@ -1010,7 +1010,7 @@ export async function compact( ? previousRemoteCompaction.replacementHistory : undefined; const remoteHistory = buildOpenAiNativeHistory( - (summaryOptions.convertToLlm ?? convertToLlm)(remoteMessages), + (summaryOptions.convertToLlm ?? defaultConvertToLlm)(remoteMessages), model, previousReplacementHistory, ); @@ -1127,7 +1127,7 @@ async function generateTurnPrefixSummary( ): Promise { const maxTokens = Math.floor(0.5 * reserveTokens); // Smaller budget for turn prefix - const llmMessages = (options?.convertToLlm ?? convertToLlm)(messages); + const llmMessages = (options?.convertToLlm ?? defaultConvertToLlm)(messages); const conversationText = serializeConversation(llmMessages); const promptText = `\n${conversationText}\n\n\n${TURN_PREFIX_SUMMARIZATION_PROMPT}`; const summarizationMessages = [ diff --git a/packages/agent/src/compaction/messages.ts b/packages/agent/src/compaction/messages.ts index 93ae21b4f..149996e79 100644 --- a/packages/agent/src/compaction/messages.ts +++ b/packages/agent/src/compaction/messages.ts @@ -51,6 +51,8 @@ export interface CompactionSummaryMessage { shortSummary?: string; tokensBefore: number; providerPayload?: ProviderPayload; + /** Snapcompact frames archived by this compaction; appended as image blocks after the summary text. */ + images?: ImageContent[]; timestamp: number; } @@ -98,6 +100,7 @@ export function createCompactionSummaryMessage( timestamp: string, shortSummary?: string, providerPayload?: ProviderPayload, + images?: ImageContent[], ): CompactionSummaryMessage { return { role: "compactionSummary", @@ -105,6 +108,7 @@ export function createCompactionSummaryMessage( shortSummary, tokensBefore, providerPayload, + images: images && images.length > 0 ? images : undefined, timestamp: new Date(timestamp).getTime(), }; } @@ -137,6 +141,79 @@ function isCoreCompactionMessage(message: AgentMessage): message is AgentMessage ); } +/** + * Transform a single core-domain agent message to its LLM form; `undefined` + * drops it from the provider request. + * + * Single source of truth for the core roles (user/developer/assistant/ + * toolResult) and the compaction messages owned by this package. Embedders + * with their own app messages (e.g. the coding agent) handle their custom + * roles and delegate every core role here — duplicating these cases is how + * snapcompact frames once silently fell off the provider request. + */ +export function convertMessageToLlm(message: AgentMessage): Message | undefined { + if (isCoreCompactionMessage(message)) { + switch (message.role) { + case "custom": + case "hookMessage": { + const content = + typeof message.content === "string" + ? [{ type: "text" as const, text: message.content }] + : message.content; + return { + role: "developer", + content, + attribution: message.attribution, + timestamp: message.timestamp, + }; + } + case "branchSummary": + return { + role: "user", + content: [ + { + type: "text" as const, + text: renderBranchSummaryContext(message.summary), + }, + ], + attribution: "agent", + timestamp: message.timestamp, + }; + case "compactionSummary": + return { + role: "user", + content: [ + { + type: "text" as const, + text: renderCompactionSummaryContext(message.summary), + }, + ...(message.images ?? []), + ], + attribution: "agent", + providerPayload: message.providerPayload, + timestamp: message.timestamp, + }; + } + } + + switch (message.role) { + case "user": + return { ...message, attribution: message.attribution ?? "user" }; + case "developer": + return { ...message, attribution: message.attribution ?? "agent" }; + case "assistant": + return message as AssistantMessage; + case "toolResult": + return { + ...message, + content: getPrunedToolResultContent(message as ToolResultMessage), + attribution: message.attribution ?? "agent", + }; + default: + return undefined; + } +} + /** * Default compaction-domain transformer. * @@ -145,68 +222,5 @@ function isCoreCompactionMessage(message: AgentMessage): message is AgentMessage * core LLM roles and the compaction messages owned by this package. */ export function defaultConvertToLlm(messages: AgentMessage[]): Message[] { - return messages - .map((message): Message | undefined => { - if (isCoreCompactionMessage(message)) { - switch (message.role) { - case "custom": - case "hookMessage": { - const content = - typeof message.content === "string" - ? [{ type: "text" as const, text: message.content }] - : message.content; - return { - role: "developer", - content, - attribution: message.attribution, - timestamp: message.timestamp, - }; - } - case "branchSummary": - return { - role: "user", - content: [ - { - type: "text" as const, - text: renderBranchSummaryContext(message.summary), - }, - ], - attribution: "agent", - timestamp: message.timestamp, - }; - case "compactionSummary": - return { - role: "user", - content: [ - { - type: "text" as const, - text: renderCompactionSummaryContext(message.summary), - }, - ], - attribution: "agent", - providerPayload: message.providerPayload, - timestamp: message.timestamp, - }; - } - } - - switch (message.role) { - case "user": - return { ...message, attribution: message.attribution ?? "user" }; - case "developer": - return { ...message, attribution: message.attribution ?? "agent" }; - case "assistant": - return message as AssistantMessage; - case "toolResult": - return { - ...message, - content: getPrunedToolResultContent(message as ToolResultMessage), - attribution: message.attribution ?? "agent", - }; - default: - return undefined; - } - }) - .filter(message => message !== undefined); + return messages.map(convertMessageToLlm).filter(message => message !== undefined); } -export const convertToLlm = defaultConvertToLlm; From 03b5c48827a9c13aa26c2362f30afc4b228ccf84 Mon Sep 17 00:00:00 2001 From: can1357 Date: Wed, 10 Jun 2026 17:43:17 +0200 Subject: [PATCH 08/27] feat(agent): added supersedeReads pruning for redundant tool results Adds pruneSupersededToolResults() and the opt-in PruneConfig.supersedeKey hook: when a tool call shares a key with a newer one (e.g. a re-read of the same file), the older result is pruned even inside the protectTokens window and replaced with a [Superseded by a newer read of this file] placeholder. Adds readToolSupersedeKey() and the shared splitReadSelector() implementing the read-tool path/selector grammar (including the .. range alias and L-prefix forms) so selector-free reads supersede range reads of the same file and URL-scheme paths are exempt. Strips selectors before tracking in compaction lists, so reads dedupe to the base path and match write/edit paths when splitting read-only vs modified lists (selector-polluted lists from earlier compactions self-heal on the next pass). Gated by the new compaction.supersedeReads setting (default on). --- .../src/compaction/branch-summarization.ts | 3 +- packages/agent/src/compaction/compaction.ts | 3 +- packages/agent/src/compaction/pruning.ts | 182 ++++++++- packages/agent/src/compaction/utils.ts | 51 ++- .../agent/test/compaction-file-ops.test.ts | 76 ++++ packages/agent/test/supersede-prune.test.ts | 346 ++++++++++++++++++ .../src/config/settings-schema.ts | 12 + 7 files changed, 662 insertions(+), 11 deletions(-) create mode 100644 packages/agent/test/compaction-file-ops.test.ts create mode 100644 packages/agent/test/supersede-prune.test.ts diff --git a/packages/agent/src/compaction/branch-summarization.ts b/packages/agent/src/compaction/branch-summarization.ts index c2d10b630..3c1976cfb 100644 --- a/packages/agent/src/compaction/branch-summarization.ts +++ b/packages/agent/src/compaction/branch-summarization.ts @@ -27,6 +27,7 @@ import { type FileOperations, SUMMARIZATION_SYSTEM_PROMPT, serializeConversation, + stripReadSelector, upsertFileOperations, } from "./utils"; @@ -214,7 +215,7 @@ export function prepareBranchEntries(entries: SessionEntry[], tokenBudget: numbe if (entry.type === "branch_summary" && !entry.fromExtension && entry.details) { const details = entry.details as BranchSummaryDetails; if (Array.isArray(details.readFiles)) { - for (const f of details.readFiles) fileOps.read.add(f); + for (const f of details.readFiles) fileOps.read.add(stripReadSelector(f)); } if (Array.isArray(details.modifiedFiles)) { // Modified files go into both edited and written for proper deduplication diff --git a/packages/agent/src/compaction/compaction.ts b/packages/agent/src/compaction/compaction.ts index 1b23c2669..66b8dcef1 100644 --- a/packages/agent/src/compaction/compaction.ts +++ b/packages/agent/src/compaction/compaction.ts @@ -45,6 +45,7 @@ import { type FileOperations, SUMMARIZATION_SYSTEM_PROMPT, serializeConversation, + stripReadSelector, upsertFileOperations, } from "./utils"; @@ -74,7 +75,7 @@ function extractFileOperations( if (!prevCompaction.fromExtension && prevCompaction.details) { const details = prevCompaction.details as CompactionDetails; if (Array.isArray(details.readFiles)) { - for (const f of details.readFiles) fileOps.read.add(f); + for (const f of details.readFiles) fileOps.read.add(stripReadSelector(f)); } if (Array.isArray(details.modifiedFiles)) { for (const f of details.modifiedFiles) fileOps.edited.add(f); diff --git a/packages/agent/src/compaction/pruning.ts b/packages/agent/src/compaction/pruning.ts index 9aee37131..b4a546024 100644 --- a/packages/agent/src/compaction/pruning.ts +++ b/packages/agent/src/compaction/pruning.ts @@ -3,7 +3,7 @@ */ import type { ToolResultMessage } from "@oh-my-pi/pi-ai"; -import type { AgentMessage } from "../types"; +import type { AgentMessage, AgentToolCall } from "../types"; import { estimateTokens } from "./compaction"; import type { SessionEntry, SessionMessageEntry } from "./entries"; import { @@ -12,6 +12,7 @@ import { isSkillReadToolResult, type ProtectedToolMatcher, } from "./tool-protection"; +import { splitReadSelector } from "./utils"; export interface PruneConfig { /** Keep the most recent tool output tokens intact. */ @@ -20,6 +21,13 @@ export interface PruneConfig { minimumSavings: number; /** Tool-result protection matchers. String entries protect every result from that tool; predicates may inspect the paired tool call. */ protectedTools: ProtectedToolMatcher[]; + /** + * Optional supersede key function (see {@link SupersedePruneConfig.supersedeKey}). + * When provided, superseded tool results are pruned first — even inside the + * `protectTokens` window — before age-based victims. Absent, behavior is + * unchanged. + */ + supersedeKey?: SupersedeKeyFn; } export const DEFAULT_PRUNE_CONFIG: PruneConfig = { @@ -33,6 +41,34 @@ export interface PruneResult { tokensSaved: number; } +/** Exact placeholder written over a superseded tool result. */ +export const SUPERSEDED_NOTICE = "[Superseded by a newer read of this file]"; + +/** + * Maps a tool call to a supersede key. Results sharing a key form a group in + * which every result except the newest is a supersede candidate. A key `K` + * additionally supersedes keys with prefix `K + "\u0000"` (selector-free read + * supersedes selector-carrying reads of the same base path). Return + * `undefined` to exempt a call from supersede grouping. + */ +export type SupersedeKeyFn = (toolName: string, args: Record) => string | undefined; + +export interface SupersedePruneConfig { + /** Supersede key function; results sharing a key supersede older ones. */ + supersedeKey: SupersedeKeyFn; + /** Prune a candidate now when all messages after it total at most this many estimated tokens. Default 8 000. */ + suffixTokenLimit?: number; + /** Prune all candidates when the last message is at least this old (prompt cache is cold anyway). Default 30 min. */ + idleFlushMs?: number; + /** Clock override for tests. */ + now?: number; + /** Tool-result protection matchers (same contract as {@link PruneConfig.protectedTools}). */ + protectedTools: ProtectedToolMatcher[]; +} + +const DEFAULT_SUFFIX_TOKEN_LIMIT = 8_000; +const DEFAULT_IDLE_FLUSH_MS = 30 * 60_000; + function createPrunedNotice(tokens: number): string { return `[Output truncated - ${tokens} tokens]`; } @@ -44,18 +80,121 @@ function getToolResultMessage(entry: SessionEntry): ToolResultMessage | undefine return message as ToolResultMessage; } -function estimatePrunedSavings(tokens: number): number { - const noticeTokens = Math.ceil(createPrunedNotice(tokens).length / 4); +function estimatePrunedSavings(tokens: number, notice: string): number { + const noticeTokens = Math.ceil(notice.length / 4); return Math.max(0, tokens - noticeTokens); } +interface SupersedeCandidate { + entry: SessionMessageEntry; + message: ToolResultMessage; + /** Index of the entry within the `entries` array. */ + index: number; + tokens: number; +} + +/** + * Collect superseded tool results: for every unpruned, unprotected tool result + * whose paired call resolves a supersede key, a LATER result with the same key + * — or with a key that is the `"\u0000"`-prefix parent of this one — marks it + * superseded. Returned in message order. + */ +function collectSupersededResults( + entries: readonly SessionEntry[], + toolCallsById: ReadonlyMap, + supersedeKey: SupersedeKeyFn, + protectedTools: readonly ProtectedToolMatcher[], +): SupersedeCandidate[] { + const candidates: SupersedeCandidate[] = []; + const seenKeys = new Set(); + for (let i = entries.length - 1; i >= 0; i--) { + const entry = entries[i]; + const message = getToolResultMessage(entry); + if (!message || message.prunedAt !== undefined) continue; + const toolCall = toolCallsById.get(message.toolCallId); + if (!toolCall) continue; + if (isProtectedToolResult(message, toolCall, protectedTools)) continue; + const key = supersedeKey(toolCall.name, toolCall.arguments as Record); + if (key === undefined) continue; + const separator = key.indexOf("\u0000"); + const superseded = seenKeys.has(key) || (separator >= 0 && seenKeys.has(key.slice(0, separator))); + seenKeys.add(key); + if (!superseded) continue; + candidates.push({ + entry: entry as SessionMessageEntry, + message, + index: i, + tokens: estimateTokens(message as AgentMessage), + }); + } + return candidates.reverse(); +} + +/** + * Prune superseded tool results (e.g. stale `read` outputs replaced by a newer + * read of the same file). Cheap, incremental, and prompt-cache-aware: a + * candidate is pruned now only when the suffix after it is small (tail case — + * the read→edit→read loop) or when the context has been idle long enough that + * the provider cache is cold anyway (then ALL candidates flush). + */ +export function pruneSupersededToolResults(entries: SessionEntry[], config: SupersedePruneConfig): PruneResult { + const toolCallsById = collectToolCallsById(entries); + const candidates = collectSupersededResults(entries, toolCallsById, config.supersedeKey, config.protectedTools); + if (candidates.length === 0) return { prunedCount: 0, tokensSaved: 0 }; + + const now = config.now ?? Date.now(); + let lastMessageTimestamp: number | undefined; + for (let i = entries.length - 1; i >= 0; i--) { + const entry = entries[i]; + if (entry.type !== "message") continue; + const timestamp = (entry.message as AgentMessage).timestamp; + if (typeof timestamp === "number") lastMessageTimestamp = timestamp; + break; + } + const idle = + lastMessageTimestamp !== undefined && now - lastMessageTimestamp >= (config.idleFlushMs ?? DEFAULT_IDLE_FLUSH_MS); + + let toPrune: SupersedeCandidate[]; + if (idle) { + toPrune = candidates; + } else { + const suffixTokenLimit = config.suffixTokenLimit ?? DEFAULT_SUFFIX_TOKEN_LIMIT; + // suffixTokens[i] = estimated tokens of all messages strictly after entry i. + const suffixTokens = new Array(entries.length); + let accumulated = 0; + for (let i = entries.length - 1; i >= 0; i--) { + suffixTokens[i] = accumulated; + const entry = entries[i]; + if (entry.type === "message") accumulated += estimateTokens(entry.message as AgentMessage); + } + toPrune = candidates.filter(candidate => suffixTokens[candidate.index] <= suffixTokenLimit); + } + if (toPrune.length === 0) return { prunedCount: 0, tokensSaved: 0 }; + + const prunedAt = Date.now(); + let tokensSaved = 0; + for (const candidate of toPrune) { + candidate.message.content = [{ type: "text", text: SUPERSEDED_NOTICE }]; + candidate.message.prunedAt = prunedAt; + tokensSaved += estimatePrunedSavings(candidate.tokens, SUPERSEDED_NOTICE); + } + return { prunedCount: toPrune.length, tokensSaved }; +} + export function pruneToolOutputs(entries: SessionEntry[], config: PruneConfig = DEFAULT_PRUNE_CONFIG): PruneResult { let accumulatedTokens = 0; let tokensSaved = 0; let prunedCount = 0; - const candidates: Array<{ entry: SessionMessageEntry; tokens: number }> = []; + const candidates: Array<{ entry: SessionMessageEntry; tokens: number; superseded: boolean }> = []; const toolCallsById = collectToolCallsById(entries); + const supersededMessages = config.supersedeKey + ? new Set( + collectSupersededResults(entries, toolCallsById, config.supersedeKey, config.protectedTools).map( + candidate => candidate.message, + ), + ) + : undefined; for (let i = entries.length - 1; i >= 0; i--) { const entry = entries[i]; @@ -70,17 +209,23 @@ export function pruneToolOutputs(entries: SessionEntry[], config: PruneConfig = continue; } - if (accumulatedTokens < config.protectTokens || isProtected) { + // Superseded results are pruned first: they bypass the protect window + // (a stale copy of re-read content is dead weight at any age). + const superseded = supersededMessages?.has(message) ?? false; + if (!superseded && (accumulatedTokens < config.protectTokens || isProtected)) { accumulatedTokens += tokens; continue; } - candidates.push({ entry: entry as SessionMessageEntry, tokens }); + candidates.push({ entry: entry as SessionMessageEntry, tokens, superseded }); accumulatedTokens += tokens; } for (const candidate of candidates) { - tokensSaved += estimatePrunedSavings(candidate.tokens); + tokensSaved += estimatePrunedSavings( + candidate.tokens, + candidate.superseded ? SUPERSEDED_NOTICE : createPrunedNotice(candidate.tokens), + ); } if (tokensSaved < config.minimumSavings || candidates.length === 0) { @@ -90,10 +235,31 @@ export function pruneToolOutputs(entries: SessionEntry[], config: PruneConfig = const prunedAt = Date.now(); for (const candidate of candidates) { const message = candidate.entry.message as ToolResultMessage; - message.content = [{ type: "text", text: createPrunedNotice(candidate.tokens) }]; + message.content = [ + { type: "text", text: candidate.superseded ? SUPERSEDED_NOTICE : createPrunedNotice(candidate.tokens) }, + ]; message.prunedAt = prunedAt; prunedCount++; } return { prunedCount, tokensSaved }; } + +/** + * Supersede key for the `read` tool: the file path with the trailing line/raw + * selector stripped (the read tool's own splitter grammar via + * {@link splitReadSelector}, e.g. `src/foo.ts:50-200`, `:2-4:raw`). + * Internal/URL-scheme paths (`skill://…`, `https://…`) are exempt. + * Selector-free reads key on the bare path; selector-carrying reads key on + * `path + "\u0000" + selector`, so two reads collide only when the newer is + * selector-free or the selectors are identical (the pass's prefix rule lets a + * bare-path read supersede selector-carrying reads of the same file). + */ +export function readToolSupersedeKey(toolName: string, args: Record): string | undefined { + if (toolName !== "read") return undefined; + const path = args.path; + if (typeof path !== "string" || path.length === 0) return undefined; + if (path.includes("://")) return undefined; + const { path: base, sel } = splitReadSelector(path); + return sel === undefined ? base : `${base}\u0000${sel}`; +} diff --git a/packages/agent/src/compaction/utils.ts b/packages/agent/src/compaction/utils.ts index 7375952da..cee780ce4 100644 --- a/packages/agent/src/compaction/utils.ts +++ b/packages/agent/src/compaction/utils.ts @@ -26,6 +26,55 @@ export function createFileOps(): FileOperations { }; } +// Read-tool selector grammar, mirrored from the conservative filesystem splitter in +// packages/coding-agent/src/tools/path-utils.ts (splitPathAndSel). Keep in sync. +// A trailing `:chunk` is a selector only when it is a line-range list +// (`50`, `50-200`, `50+10`, `5-16,960-973`, `..` alias), `raw`, or `conflicts` — +// alone or as a `range:raw` / `raw:range` compound. +const RANGE_CHUNK_SRC = String.raw`L?\d+(?:(?:[-+]|\.\.)L?\d+|-|\.\.)?`; +const RANGE_LIST_SRC = `${RANGE_CHUNK_SRC}(?:,${RANGE_CHUNK_SRC})*`; +const READ_SELECTOR_RE = new RegExp(`^(?:${RANGE_LIST_SRC}|raw|conflicts)$`, "i"); +const READ_RANGE_ONLY_RE = new RegExp(`^${RANGE_LIST_SRC}$`, "i"); +const READ_RAW_ONLY_RE = /^raw$/i; + +/** + * Split a read-tool path into its base path and trailing selector, mirroring the + * read tool's own splitter. Single source of the grammar in this package: the + * file-operations list strips selectors via {@link stripReadSelector}, and the + * supersede-prune pass keys on both parts via `readToolSupersedeKey`. + */ +export function splitReadSelector(path: string): { path: string; sel?: string } { + const colon = path.lastIndexOf(":"); + if (colon <= 0) return { path }; + const candidate = path.slice(colon + 1); + if (!READ_SELECTOR_RE.test(candidate)) return { path }; + let base = path.slice(0, colon); + let sel = candidate; + // Compound trailing selector: `path:1-50:raw` or `path:raw:1-50`. + const inner = base.lastIndexOf(":"); + if (inner > 0) { + const innerCandidate = base.slice(inner + 1); + const innerIsRaw = READ_RAW_ONLY_RE.test(innerCandidate); + const outerIsRaw = READ_RAW_ONLY_RE.test(candidate); + const innerIsRange = READ_RANGE_ONLY_RE.test(innerCandidate); + const outerIsRange = READ_RANGE_ONLY_RE.test(candidate); + if ((innerIsRaw && outerIsRange) || (innerIsRange && outerIsRaw)) { + sel = `${innerCandidate}:${candidate}`; + base = base.slice(0, inner); + } + } + return { path: base, sel }; +} + +/** + * Strip a trailing read-tool selector (`:50-200`, `:raw`, `:1-50:raw`, `:conflicts`, …) + * so the same file read with different line ranges dedupes to one `` entry + * and matches its write/edit path when computing read-only vs modified lists. + */ +export function stripReadSelector(path: string): string { + return splitReadSelector(path).path; +} + /** * Extract file operations from tool calls in an assistant message. */ @@ -46,7 +95,7 @@ export function extractFileOpsFromMessage(message: AgentMessage, fileOps: FileOp switch (block.name) { case "read": - fileOps.read.add(path); + fileOps.read.add(stripReadSelector(path)); break; case "write": fileOps.written.add(path); diff --git a/packages/agent/test/compaction-file-ops.test.ts b/packages/agent/test/compaction-file-ops.test.ts new file mode 100644 index 000000000..fa8a0e116 --- /dev/null +++ b/packages/agent/test/compaction-file-ops.test.ts @@ -0,0 +1,76 @@ +import { describe, expect, it } from "bun:test"; +import { + computeFileLists, + createFileOps, + extractFileOpsFromMessage, + formatFileOperations, + stripReadSelector, +} from "../src/compaction/utils"; +import { createAssistantMessage } from "./helpers"; + +function readCall(id: string, path: string) { + return { type: "toolCall" as const, id, name: "read", arguments: { path } }; +} + +describe("stripReadSelector", () => { + it("strips line-range and raw selectors in every supported shape", () => { + expect(stripReadSelector("src/foo.ts:50")).toBe("src/foo.ts"); + expect(stripReadSelector("src/foo.ts:50-")).toBe("src/foo.ts"); + expect(stripReadSelector("src/foo.ts:50-200")).toBe("src/foo.ts"); + expect(stripReadSelector("src/foo.ts:50+150")).toBe("src/foo.ts"); + expect(stripReadSelector("src/foo.ts:5-16,960-973")).toBe("src/foo.ts"); + expect(stripReadSelector("src/foo.ts:2724..2727")).toBe("src/foo.ts"); + expect(stripReadSelector("src/foo.ts:raw")).toBe("src/foo.ts"); + expect(stripReadSelector("src/foo.ts:conflicts")).toBe("src/foo.ts"); + // Compound raw+range, either order. + expect(stripReadSelector("src/foo.ts:100-170:raw")).toBe("src/foo.ts"); + expect(stripReadSelector("src/foo.ts:raw:2-4")).toBe("src/foo.ts"); + }); + + it("keeps archive member paths, stripping only the trailing selector", () => { + expect(stripReadSelector("archive.zip:dir/file.ts:50-60")).toBe("archive.zip:dir/file.ts"); + expect(stripReadSelector("archive.zip:dir/file.ts")).toBe("archive.zip:dir/file.ts"); + }); + + it("leaves non-selector colons untouched", () => { + expect(stripReadSelector("db.sqlite:users")).toBe("db.sqlite:users"); + expect(stripReadSelector("local://ctx.md")).toBe("local://ctx.md"); + expect(stripReadSelector("https://example.com/page")).toBe("https://example.com/page"); + expect(stripReadSelector("src/foo.ts")).toBe("src/foo.ts"); + }); +}); + +describe("extractFileOpsFromMessage", () => { + it("dedupes the same file read through different selectors to one entry", () => { + const fileOps = createFileOps(); + const message = createAssistantMessage([ + readCall("r1", "docs/compaction.md:100-170:raw"), + readCall("r2", "docs/compaction.md:8-16,128-139,384-388"), + readCall("r3", "docs/compaction.md:raw"), + readCall("r4", "docs/compaction.md"), + ]); + extractFileOpsFromMessage(message, fileOps); + expect([...fileOps.read]).toEqual(["docs/compaction.md"]); + }); + + it("matches selector-suffixed reads against modified paths", () => { + const fileOps = createFileOps(); + const message = createAssistantMessage([ + readCall("r1", "src/login.ts:30-80"), + { type: "toolCall" as const, id: "w1", name: "write", arguments: { path: "src/login.ts" } }, + ]); + extractFileOpsFromMessage(message, fileOps); + const { readFiles, modifiedFiles } = computeFileLists(fileOps); + expect(readFiles).toEqual([]); + expect(modifiedFiles).toEqual(["src/login.ts"]); + }); +}); + +describe("formatFileOperations", () => { + it("renders one path per line, not literal \\n separators", () => { + const rendered = formatFileOperations(["a.ts", "b.ts"], ["c.ts"]); + expect(rendered).toContain("\na.ts\nb.ts\n"); + expect(rendered).toContain("\nc.ts\n"); + expect(rendered).not.toContain("\\n"); + }); +}); diff --git a/packages/agent/test/supersede-prune.test.ts b/packages/agent/test/supersede-prune.test.ts new file mode 100644 index 000000000..02fdada5c --- /dev/null +++ b/packages/agent/test/supersede-prune.test.ts @@ -0,0 +1,346 @@ +import { describe, expect, test } from "bun:test"; +import type { AgentMessage } from "@oh-my-pi/pi-agent-core"; +import type { SessionEntry, SessionMessageEntry } from "@oh-my-pi/pi-agent-core/compaction"; +import { + DEFAULT_PRUNE_CONFIG, + pruneSupersededToolResults, + pruneToolOutputs, + readToolSupersedeKey, + SUPERSEDED_NOTICE, + type SupersedePruneConfig, +} from "@oh-my-pi/pi-agent-core/compaction"; +import type { ProtectedToolContext } from "@oh-my-pi/pi-agent-core/compaction/tool-protection"; +import type { AssistantMessage, TextContent, ToolResultMessage } from "@oh-my-pi/pi-ai"; + +let idCounter = 0; +function nextId(): string { + return `entry-${idCounter++}`; +} + +function messageEntry(message: AgentMessage, timestamp: number): SessionMessageEntry { + return { type: "message", id: nextId(), parentId: null, timestamp: new Date(timestamp).toISOString(), message }; +} + +function assistantMessage(content: AssistantMessage["content"], timestamp: number): AssistantMessage { + return { + role: "assistant", + content, + timestamp, + provider: "mock", + model: "mock", + api: "mock", + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "stop", + }; +} + +function toolResultMessage(toolName: string, toolCallId: string, text: string, timestamp: number): ToolResultMessage { + return { + role: "toolResult", + toolCallId, + toolName, + content: [{ type: "text", text }], + isError: false, + timestamp, + }; +} + +/** Assistant toolCall entry + paired toolResult entry for one read. */ +function readPair(path: string, text: string, timestamp: number): [SessionMessageEntry, SessionMessageEntry] { + const callId = `call-${idCounter++}`; + return [ + messageEntry( + assistantMessage([{ type: "toolCall", id: callId, name: "read", arguments: { path } }], timestamp), + timestamp, + ), + messageEntry(toolResultMessage("read", callId, text, timestamp), timestamp), + ]; +} + +function textEntry(text: string, timestamp: number): SessionMessageEntry { + return messageEntry(assistantMessage([{ type: "text", text }], timestamp), timestamp); +} + +function resultText(entry: SessionEntry): string { + const message = (entry as SessionMessageEntry).message as ToolResultMessage; + return (message.content[0] as TextContent).text; +} + +function resultMessage(entry: SessionEntry): ToolResultMessage { + return (entry as SessionMessageEntry).message as ToolResultMessage; +} + +function cfg(over: Partial = {}): SupersedePruneConfig { + return { supersedeKey: readToolSupersedeKey, protectedTools: [], ...over }; +} + +const T0 = Date.UTC(2026, 5, 10, 12, 0, 0); +const FILE_CONTENT = "export function alpha() { return 1; }\n".repeat(50); +// Comfortably above any small suffixTokenLimit used below. +const BIG_TEXT = "const value = computeSomething(12345);\n".repeat(500); + +describe("readToolSupersedeKey", () => { + test("bare path keys on itself; non-read and non-string paths are exempt", () => { + expect(readToolSupersedeKey("read", { path: "src/foo.ts" })).toBe("src/foo.ts"); + expect(readToolSupersedeKey("bash", { path: "src/foo.ts" })).toBeUndefined(); + expect(readToolSupersedeKey("read", { path: 42 })).toBeUndefined(); + expect(readToolSupersedeKey("read", {})).toBeUndefined(); + }); + + test("URL/internal schemes are exempt", () => { + expect(readToolSupersedeKey("read", { path: "skill://react" })).toBeUndefined(); + expect(readToolSupersedeKey("read", { path: "https://example.com/page" })).toBeUndefined(); + }); + + test("strips trailing selectors into a \\u0000-separated key", () => { + expect(readToolSupersedeKey("read", { path: "src/foo.ts:50-200" })).toBe("src/foo.ts\u000050-200"); + expect(readToolSupersedeKey("read", { path: "src/foo.ts:raw" })).toBe("src/foo.ts\u0000raw"); + expect(readToolSupersedeKey("read", { path: "src/foo.ts:conflicts" })).toBe("src/foo.ts\u0000conflicts"); + expect(readToolSupersedeKey("read", { path: "src/foo.ts:2-4:raw" })).toBe("src/foo.ts\u00002-4:raw"); + expect(readToolSupersedeKey("read", { path: "src/foo.ts:5-16,960-973" })).toBe("src/foo.ts\u00005-16,960-973"); + expect(readToolSupersedeKey("read", { path: "src/foo.ts:50+150" })).toBe("src/foo.ts\u000050+150"); + }); + + test("does not strip non-selector colon segments", () => { + expect(readToolSupersedeKey("read", { path: "db.sqlite:users" })).toBe("db.sqlite:users"); + expect(readToolSupersedeKey("read", { path: "db.sqlite:users:42" })).toBe("db.sqlite:users\u000042"); + }); +}); + +describe("pruneSupersededToolResults — tail case", () => { + test("(a) older identical-path read pruned with exact placeholder when suffix small", () => { + const [call1, result1] = readPair("src/foo.ts", FILE_CONTENT, T0); + const [call2, result2] = readPair("src/foo.ts", FILE_CONTENT, T0 + 1_000); + const entries: SessionEntry[] = [call1, result1, call2, result2]; + + const result = pruneSupersededToolResults(entries, cfg({ now: T0 + 1_000 })); + + expect(result.prunedCount).toBe(1); + expect(result.tokensSaved).toBeGreaterThan(0); + expect(resultText(result1)).toBe("[Superseded by a newer read of this file]"); + expect(resultText(result1)).toBe(SUPERSEDED_NOTICE); + expect(resultMessage(result1).prunedAt).toBeDefined(); + // Latest read untouched. + expect(resultText(result2)).toBe(FILE_CONTENT); + expect(resultMessage(result2).prunedAt).toBeUndefined(); + }); + + test("(b) NOT pruned when suffix exceeds limit and no idle gap", () => { + const [call1, result1] = readPair("src/foo.ts", FILE_CONTENT, T0); + const [call2, result2] = readPair("src/foo.ts", FILE_CONTENT, T0 + 1_000); + const big = textEntry(BIG_TEXT, T0 + 2_000); + const entries: SessionEntry[] = [call1, result1, call2, result2, big]; + + const result = pruneSupersededToolResults(entries, cfg({ suffixTokenLimit: 200, now: T0 + 2_000 })); + + expect(result.prunedCount).toBe(0); + expect(result.tokensSaved).toBe(0); + expect(resultText(result1)).toBe(FILE_CONTENT); + expect(resultMessage(result1).prunedAt).toBeUndefined(); + expect(resultText(result2)).toBe(FILE_CONTENT); + }); + + test("(c) idle gap prunes all candidates regardless of suffix", () => { + const [call1, result1] = readPair("src/foo.ts", FILE_CONTENT, T0); + const [call2, result2] = readPair("src/bar.ts", FILE_CONTENT, T0 + 1_000); + const [call3, result3] = readPair("src/foo.ts", FILE_CONTENT, T0 + 2_000); + const [call4, result4] = readPair("src/bar.ts", FILE_CONTENT, T0 + 3_000); + const big = textEntry(BIG_TEXT, T0 + 4_000); + const entries: SessionEntry[] = [call1, result1, call2, result2, call3, result3, call4, result4, big]; + + // Suffix limit 0 would block every candidate; only the idle gap fires. + const result = pruneSupersededToolResults( + entries, + cfg({ suffixTokenLimit: 0, idleFlushMs: 30 * 60_000, now: T0 + 4_000 + 30 * 60_000 }), + ); + + expect(result.prunedCount).toBe(2); + expect(resultText(result1)).toBe(SUPERSEDED_NOTICE); + expect(resultText(result2)).toBe(SUPERSEDED_NOTICE); + expect(resultText(result3)).toBe(FILE_CONTENT); + expect(resultText(result4)).toBe(FILE_CONTENT); + }); + + test("no idle flush when gap is below the threshold", () => { + const [call1, result1] = readPair("src/foo.ts", FILE_CONTENT, T0); + const [call2, result2] = readPair("src/foo.ts", FILE_CONTENT, T0 + 1_000); + const big = textEntry(BIG_TEXT, T0 + 2_000); + const entries: SessionEntry[] = [call1, result1, call2, result2, big]; + + const result = pruneSupersededToolResults( + entries, + cfg({ suffixTokenLimit: 0, idleFlushMs: 30 * 60_000, now: T0 + 2_000 + 29 * 60_000 }), + ); + + expect(result.prunedCount).toBe(0); + expect(resultText(result1)).toBe(FILE_CONTENT); + expect(resultText(result2)).toBe(FILE_CONTENT); + }); +}); + +describe("pruneSupersededToolResults — selectors", () => { + test("(d) different range selectors do not supersede each other; a later selector-free read supersedes them", () => { + const [callA, resultA] = readPair("src/foo.ts:50-200", FILE_CONTENT, T0); + const [callB, resultB] = readPair("src/foo.ts:10-20", FILE_CONTENT, T0 + 1_000); + let entries: SessionEntry[] = [callA, resultA, callB, resultB]; + + // Different selectors: no candidates. + let result = pruneSupersededToolResults(entries, cfg({ now: T0 + 1_000 })); + expect(result.prunedCount).toBe(0); + expect(resultText(resultA)).toBe(FILE_CONTENT); + expect(resultText(resultB)).toBe(FILE_CONTENT); + + // Identical selector strings DO supersede. + const [callA2, resultA2] = readPair("src/foo.ts:50-200", FILE_CONTENT, T0 + 2_000); + entries = [...entries, callA2, resultA2]; + result = pruneSupersededToolResults(entries, cfg({ now: T0 + 2_000 })); + expect(result.prunedCount).toBe(1); + expect(resultText(resultA)).toBe(SUPERSEDED_NOTICE); + expect(resultText(resultB)).toBe(FILE_CONTENT); + expect(resultText(resultA2)).toBe(FILE_CONTENT); + + // A later selector-free read supersedes every selector-carrying read of the base path. + const [callFull, resultFull] = readPair("src/foo.ts", FILE_CONTENT, T0 + 3_000); + entries = [...entries, callFull, resultFull]; + result = pruneSupersededToolResults(entries, cfg({ now: T0 + 3_000 })); + expect(result.prunedCount).toBe(2); + expect(resultText(resultB)).toBe(SUPERSEDED_NOTICE); + expect(resultText(resultA2)).toBe(SUPERSEDED_NOTICE); + expect(resultText(resultFull)).toBe(FILE_CONTENT); + }); + + test("a selector-carrying read does NOT supersede an earlier selector-free read", () => { + const [callFull, resultFull] = readPair("src/foo.ts", FILE_CONTENT, T0); + const [callRange, resultRange] = readPair("src/foo.ts:50-200", FILE_CONTENT, T0 + 1_000); + const entries: SessionEntry[] = [callFull, resultFull, callRange, resultRange]; + + const result = pruneSupersededToolResults(entries, cfg({ now: T0 + 1_000 })); + + expect(result.prunedCount).toBe(0); + expect(resultText(resultFull)).toBe(FILE_CONTENT); + expect(resultText(resultRange)).toBe(FILE_CONTENT); + }); +}); + +describe("pruneSupersededToolResults — protection & latest", () => { + test("(e) latest read never pruned, even with idle flush", () => { + const [call1, result1] = readPair("src/foo.ts", FILE_CONTENT, T0); + const [call2, result2] = readPair("src/foo.ts", FILE_CONTENT, T0 + 1_000); + const [call3, result3] = readPair("src/foo.ts", FILE_CONTENT, T0 + 2_000); + const entries: SessionEntry[] = [call1, result1, call2, result2, call3, result3]; + + const result = pruneSupersededToolResults(entries, cfg({ now: T0 + 2_000 + 60 * 60_000 })); + + expect(result.prunedCount).toBe(2); + expect(resultText(result1)).toBe(SUPERSEDED_NOTICE); + expect(resultText(result2)).toBe(SUPERSEDED_NOTICE); + expect(resultText(result3)).toBe(FILE_CONTENT); + expect(resultMessage(result3).prunedAt).toBeUndefined(); + }); + + test("(f) protected tool results never pruned", () => { + const protectPlan = ({ toolCall }: ProtectedToolContext): boolean => + (toolCall?.arguments as Record | undefined)?.path === "plan.md"; + const [planCall1, planResult1] = readPair("plan.md", FILE_CONTENT, T0); + const [fooCall1, fooResult1] = readPair("src/foo.ts", FILE_CONTENT, T0 + 1_000); + const [planCall2, planResult2] = readPair("plan.md", FILE_CONTENT, T0 + 2_000); + const [fooCall2, fooResult2] = readPair("src/foo.ts", FILE_CONTENT, T0 + 3_000); + const entries: SessionEntry[] = [ + planCall1, + planResult1, + fooCall1, + fooResult1, + planCall2, + planResult2, + fooCall2, + fooResult2, + ]; + + const result = pruneSupersededToolResults(entries, cfg({ protectedTools: [protectPlan], now: T0 + 3_000 })); + + expect(result.prunedCount).toBe(1); + expect(resultText(planResult1)).toBe(FILE_CONTENT); + expect(resultText(planResult2)).toBe(FILE_CONTENT); + expect(resultText(fooResult1)).toBe(SUPERSEDED_NOTICE); + expect(resultText(fooResult2)).toBe(FILE_CONTENT); + }); + + test("already-pruned results are ignored as candidates and as superseders", () => { + const [call1, result1] = readPair("src/foo.ts", FILE_CONTENT, T0); + const [call2, result2] = readPair("src/foo.ts", FILE_CONTENT, T0 + 1_000); + resultMessage(result2).prunedAt = T0 + 1_500; + const entries: SessionEntry[] = [call1, result1, call2, result2]; + + // The only newer same-key read is itself pruned -> result1 has no live superseder. + const result = pruneSupersededToolResults(entries, cfg({ now: T0 + 2_000 })); + + expect(result.prunedCount).toBe(0); + expect(resultText(result1)).toBe(FILE_CONTENT); + }); +}); + +describe("pruneToolOutputs — supersede priority fold", () => { + test("with supersedeKey, superseded results bypass the protect window and get the supersede placeholder", () => { + const [call1, result1] = readPair("src/foo.ts", FILE_CONTENT, T0); + const [call2, result2] = readPair("src/foo.ts", FILE_CONTENT, T0 + 1_000); + const entries: SessionEntry[] = [call1, result1, call2, result2]; + + const result = pruneToolOutputs(entries, { + protectTokens: 1_000_000, // everything inside the protect window + minimumSavings: 0, + protectedTools: [], + supersedeKey: readToolSupersedeKey, + }); + + expect(result.prunedCount).toBe(1); + expect(resultText(result1)).toBe(SUPERSEDED_NOTICE); + expect(resultText(result2)).toBe(FILE_CONTENT); + }); + + test("(g) without supersedeKey, behavior is unchanged (regression guard)", () => { + const buildEntries = (): { + entries: SessionEntry[]; + oldResult: SessionMessageEntry; + newResult: SessionMessageEntry; + } => { + const [call1, result1] = readPair("src/foo.ts", FILE_CONTENT, T0); + const [call2, result2] = readPair("src/foo.ts", FILE_CONTENT, T0 + 1_000); + return { entries: [call1, result1, call2, result2], oldResult: result1, newResult: result2 }; + }; + + // Protect window covers everything: nothing pruned, superseded reads included. + const protectedFixture = buildEntries(); + const protectedRun = pruneToolOutputs(protectedFixture.entries, { + protectTokens: 1_000_000, + minimumSavings: 0, + protectedTools: [], + }); + expect(protectedRun).toEqual({ prunedCount: 0, tokensSaved: 0 }); + expect(resultText(protectedFixture.oldResult)).toBe(FILE_CONTENT); + expect(resultText(protectedFixture.newResult)).toBe(FILE_CONTENT); + + // Protect window empty: every result past it pruned with the legacy + // truncation placeholder — never the supersede placeholder. + const unprotectedFixture = buildEntries(); + const unprotectedRun = pruneToolOutputs(unprotectedFixture.entries, { + protectTokens: 0, + minimumSavings: 0, + protectedTools: [], + }); + expect(unprotectedRun.prunedCount).toBe(2); + expect(resultText(unprotectedFixture.oldResult)).toMatch(/^\[Output truncated - \d+ tokens\]$/); + expect(resultText(unprotectedFixture.newResult)).toMatch(/^\[Output truncated - \d+ tokens\]$/); + + // Default config shape is untouched. + expect(DEFAULT_PRUNE_CONFIG.supersedeKey).toBeUndefined(); + expect(DEFAULT_PRUNE_CONFIG.protectTokens).toBe(40_000); + expect(DEFAULT_PRUNE_CONFIG.minimumSavings).toBe(20_000); + }); +}); diff --git a/packages/coding-agent/src/config/settings-schema.ts b/packages/coding-agent/src/config/settings-schema.ts index 17c87bd5a..b0f266fae 100644 --- a/packages/coding-agent/src/config/settings-schema.ts +++ b/packages/coding-agent/src/config/settings-schema.ts @@ -1326,6 +1326,17 @@ export const SETTINGS_SCHEMA = { ], }, }, + + "compaction.supersedeReads": { + type: "boolean", + default: true, + ui: { + tab: "context", + label: "Supersede Stale Reads", + description: "Prune older read results when the same file is read again (cache-aware, runs every turn)", + }, + }, + // Branch summaries "branchSummary.enabled": { type: "boolean", @@ -3365,6 +3376,7 @@ export interface CompactionSettings { idleEnabled: boolean; idleThresholdTokens: number; idleTimeoutSeconds: number; + supersedeReads: boolean; } export interface ContextPromotionSettings { From 8baeb062ec7218511c4b364335e32278006629fb Mon Sep 17 00:00:00 2001 From: can1357 Date: Wed, 10 Jun 2026 17:43:33 +0200 Subject: [PATCH 09/27] feat(agent): added snapcompact compaction strategy MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds snapcompactCompact() in compaction/snapcompact.ts: instead of an LLM-generated summary, discarded history is printed onto dense 2576px PNG frames with the public-domain X.org 5x8 pixel font and re-attached to the compaction summary message as image blocks. Fully local — no model call; ~7x cheaper than raw text at near-parity recall. CompactionSummaryMessage now charges per attached frame in estimateTokens(), frames persist under preserveData.snapcompact with an 8-frame budget that evicts middle-out (session-head frame pinned so head and tail both survive). Rasterization and PNG encoding run in native code via renderSnapcompactPng(). --- packages/agent/src/compaction/compaction.ts | 7 +- packages/agent/src/compaction/index.ts | 1 + .../compaction/prompts/snapcompact-summary.md | 13 + packages/agent/src/compaction/snapcompact.ts | 316 ++++++++++++++++++ packages/agent/test/snapcompact.test.ts | 291 ++++++++++++++++ 5 files changed, 627 insertions(+), 1 deletion(-) create mode 100644 packages/agent/src/compaction/prompts/snapcompact-summary.md create mode 100644 packages/agent/src/compaction/snapcompact.ts create mode 100644 packages/agent/test/snapcompact.test.ts diff --git a/packages/agent/src/compaction/compaction.ts b/packages/agent/src/compaction/compaction.ts index 66b8dcef1..565f801dc 100644 --- a/packages/agent/src/compaction/compaction.ts +++ b/packages/agent/src/compaction/compaction.ts @@ -37,6 +37,7 @@ import compactionSummaryPrompt from "./prompts/compaction-summary.md" with { typ import compactionTurnPrefixPrompt from "./prompts/compaction-turn-prefix.md" with { type: "text" }; import compactionUpdateSummaryPrompt from "./prompts/compaction-update-summary.md" with { type: "text" }; import handoffDocumentPrompt from "./prompts/handoff-document.md" with { type: "text" }; +import { SNAPCOMPACT_FRAME_TOKEN_ESTIMATE } from "./snapcompact"; import { computeFileLists, @@ -138,7 +139,7 @@ export interface CompactionResult { export interface CompactionSettings { enabled: boolean; - strategy?: "context-full" | "handoff" | "shake" | "off"; + strategy?: "context-full" | "handoff" | "shake" | "snapcompact" | "off"; thresholdPercent?: number; thresholdTokens?: number; reserveTokens: number; @@ -311,6 +312,10 @@ export function estimateTokens(message: AgentMessage): number { case "branchSummary": case "compactionSummary": { fragments.push(message.summary); + if (message.role === "compactionSummary" && message.images) { + // Snapcompact frames render at ≥1568px; providers bill the downscaled cap. + extra += message.images.length * SNAPCOMPACT_FRAME_TOKEN_ESTIMATE; + } break; } default: diff --git a/packages/agent/src/compaction/index.ts b/packages/agent/src/compaction/index.ts index 401215724..1425b2976 100644 --- a/packages/agent/src/compaction/index.ts +++ b/packages/agent/src/compaction/index.ts @@ -10,4 +10,5 @@ export * from "./messages"; export * from "./openai"; export * from "./pruning"; export * from "./shake"; +export * from "./snapcompact"; export * from "./utils"; diff --git a/packages/agent/src/compaction/prompts/snapcompact-summary.md b/packages/agent/src/compaction/prompts/snapcompact-summary.md new file mode 100644 index 000000000..3e25c2dcb --- /dev/null +++ b/packages/agent/src/compaction/prompts/snapcompact-summary.md @@ -0,0 +1,13 @@ +Prior conversation history has been archived verbatim onto {{frameCount}} snapcompact frame{{#if multipleFrames}}s{{/if}} — the bitmap image{{#if multipleFrames}}s{{/if}} attached below{{#if multipleFrames}}, ordered oldest to newest{{/if}}. + +Reading a frame: monospace 5x8 pixel font on a white background, {{cols}} characters per row, {{rows}} rows per frame; read left to right, top to bottom. Text flows continuously with no word wrap, so words may break across row ends. Whitespace runs (including newlines) were collapsed to single spaces. Ink color cycles through six colors, advancing at sentence boundaries — a color change marks a new sentence. Roles are tagged inline as [User]:, [Assistant]:, [Assistant thinking]:, [Assistant tool calls]:, and [Tool result]:. +{{#if includedPreviousSummary}} + +The earliest frame begins with "[Summary of earlier history]" — a condensed digest of context that predates the archived conversation. +{{/if}} +{{#if truncatedChars}} + +{{truncatedChars}} characters of older history were dropped to respect the frame budget. The first frame (session start) is always kept, so the missing span sits between the first frame and the next. +{{/if}} + +Total archived: {{totalChars}} characters. Consult the frames whenever you need exact earlier details (user wording, decisions, file paths, tool output). If a region is hard to read, re-derive the fact from the workspace (re-read files, re-run commands) rather than guessing. diff --git a/packages/agent/src/compaction/snapcompact.ts b/packages/agent/src/compaction/snapcompact.ts new file mode 100644 index 000000000..5236a26eb --- /dev/null +++ b/packages/agent/src/compaction/snapcompact.ts @@ -0,0 +1,316 @@ +/** + * Snapcompact compaction: archive conversation history as dense bitmap images. + * + * Instead of asking an LLM to summarize discarded history, the serialized + * conversation is rendered into square PNG frames using the X.org `5x8` + * pixel font (public domain) — one character per 5x8 cell, row-major, glyph + * ink cycling through six hues at sentence boundaries. Vision models read + * the frames back directly, like an archivist at a snapcompact frame reader. + * + * Validated by the imageee SQuAD eval (`img-5x8-sent`, 2576px frames): + * ~0.88 F1 recall vs ~0.90 for raw text, at roughly 7x fewer input tokens. + * The provider downscales frames to its image cap (1568px for Anthropic), + * so one frame costs ~3.3k tokens while carrying ~165k characters. + * + * The whole pass is local and deterministic — no LLM call, no API key, no + * latency beyond rendering. Rasterization and PNG encoding happen in native + * code (`renderSnapcompactPng` in `crates/pi-natives/src/snapcompact.rs`). + * Frames persist in the compaction entry's `preserveData` and are + * re-attached to the compaction summary message on every context rebuild. + */ + +import type { ImageContent } from "@oh-my-pi/pi-ai"; +import { renderSnapcompactPng } from "@oh-my-pi/pi-natives"; +import { prompt } from "@oh-my-pi/pi-utils"; +import type { CompactionDetails, CompactionPreparation, CompactionResult } from "./compaction"; +import { type ConvertToLlm, defaultConvertToLlm } from "./messages"; +import { withOpenAiRemoteCompactionPreserveData } from "./openai"; +import snapcompactSummaryPrompt from "./prompts/snapcompact-summary.md" with { type: "text" }; +import { computeFileLists, serializeConversation, upsertFileOperations } from "./utils"; + +// ============================================================================ +// Constants +// ============================================================================ + +/** Frame edge in pixels. 2576px is the eval-validated sweet spot: the provider + * downscale to 1568px anti-aliases the 1px glyph strokes instead of shearing + * them, which reads *better* than rendering at 1568 directly. */ +export const SNAPCOMPACT_FRAME_SIZE = 2576; + +/** Glyph cell geometry of the bundled `5x8` BDF font. */ +const GLYPH_ADVANCE_X = 5; +const GLYPH_PITCH_Y = 8; + +/** Maximum frames carried on a compaction entry. Oldest frames are dropped + * first once the budget is exceeded (mirrors how iterative text summaries + * fade the oldest detail). 8 frames ≈ 26k image tokens ≈ 1.3M chars. */ +export const SNAPCOMPACT_MAX_FRAMES = 8; + +/** Token cost estimate per frame. Frames render at ≥1568px, so providers bill + * the downscaled long-edge cap: 1568*1568/750 ≈ 3,278 tokens (Anthropic). */ +export const SNAPCOMPACT_FRAME_TOKEN_ESTIMATE = 3300; + +/** Key under `CompactionEntry.preserveData` holding the frame archive. */ +export const SNAPCOMPACT_PRESERVE_KEY = "snapcompact"; + +// ============================================================================ +// Types +// ============================================================================ + +/** One developed snapcompact frame: a base64 PNG plus its reading geometry. */ +export interface SnapcompactFrame { + /** Base64-encoded PNG. */ + data: string; + mimeType: string; + /** Characters per row in the frame grid. */ + cols: number; + /** Rows in the frame grid. */ + rows: number; + /** Characters actually printed onto this frame. */ + chars: number; +} + +/** Frame archive persisted under `preserveData[SNAPCOMPACT_PRESERVE_KEY]`. */ +export interface SnapcompactArchive { + /** Frames ordered oldest to newest. */ + frames: SnapcompactFrame[]; + /** Characters currently readable across all frames. */ + totalChars: number; + /** Characters dropped so far to respect the frame budget. */ + truncatedChars: number; +} + +export interface SnapcompactGeometry { + cols: number; + rows: number; + /** Characters that fit one frame (cols * rows). */ + capacity: number; +} + +export interface SnapcompactOptions { + /** App-level message transformer (same contract as `SummaryOptions.convertToLlm`). */ + convertToLlm?: ConvertToLlm; + /** Frame edge in pixels. Defaults to {@link SNAPCOMPACT_FRAME_SIZE}. */ + frameSize?: number; + /** Frame budget. Defaults to {@link SNAPCOMPACT_MAX_FRAMES}. */ + maxFrames?: number; +} + +/** Result of rendering one frame, before base64 packing. */ +export interface RenderedFrame { + png: Uint8Array; + cols: number; + rows: number; + /** Characters printed (input may be shorter than capacity). */ + chars: number; +} + +// ============================================================================ +// Text normalization +// ============================================================================ + +/** Folds for common non-Latin-1 characters the 5x8 subset cannot draw. */ +const CHAR_FOLD: Record = { + "\u2018": "'", + "\u2019": "'", + "\u201a": "'", + "\u201b": "'", + "\u201c": '"', + "\u201d": '"', + "\u201e": '"', + "\u2013": "-", + "\u2014": "-", + "\u2015": "-", + "\u2212": "-", + "\u2026": "...", + "\u2022": "*", + "\u25cf": "*", + "\u25a0": "*", + "\u25aa": "*", + "\u2190": "<-", + "\u2192": "->", + "\u21d2": "=>", + "\u2713": "v", + "\u2714": "v", + "\u2717": "x", + "\u2718": "x", +}; + +/** + * Prepare text for printing: collapse whitespace runs (incl. newlines) to + * single spaces — the eval's "paragraph breaks collapsed to spaces" format — + * then fold everything outside the font's ASCII + Latin-1 coverage to ASCII + * approximations (`?` as the last resort). + */ +export function normalizeForSnapcompact(text: string): string { + const collapsed = text.replace(/\s+/g, " ").trim(); + let out = ""; + for (const ch of collapsed) { + const cp = ch.codePointAt(0) as number; + if (cp < 0x7f || (cp >= 0xa0 && cp <= 0xff)) { + out += ch; + continue; + } + const fold = CHAR_FOLD[ch]; + if (fold !== undefined) { + out += fold; + } else if (cp >= 0x2500 && cp <= 0x257f) { + // Box drawing: keep table skeletons legible. + out += cp === 0x2502 || cp === 0x2503 ? "|" : cp === 0x2500 || cp === 0x2501 ? "-" : "+"; + } else { + out += "?"; + } + } + return out; +} + +// ============================================================================ +// Rendering +// ============================================================================ +export function snapcompactGeometry(size: number = SNAPCOMPACT_FRAME_SIZE): SnapcompactGeometry { + const cols = Math.floor(size / GLYPH_ADVANCE_X); + const rows = Math.floor(size / GLYPH_PITCH_Y); + return { cols, rows, capacity: cols * rows }; +} + +/** Render one snapcompact frame from already-normalized text. */ +export function renderSnapcompactFrame(text: string, size: number = SNAPCOMPACT_FRAME_SIZE): RenderedFrame { + const { cols, rows, capacity } = snapcompactGeometry(size); + const chars = Math.min(text.length, capacity); + return { png: renderSnapcompactPng(text, size), cols, rows, chars }; +} + +// ============================================================================ +// Archive helpers +// ============================================================================ + +/** Validate and extract a persisted frame archive from `preserveData`. */ +export function getPreservedSnapcompactArchive( + preserveData: Record | undefined, +): SnapcompactArchive | undefined { + const candidate = preserveData?.[SNAPCOMPACT_PRESERVE_KEY]; + if (!candidate || typeof candidate !== "object") return undefined; + const archive = candidate as SnapcompactArchive; + if (!Array.isArray(archive.frames)) return undefined; + const frames = archive.frames.filter( + frame => + !!frame && + typeof frame.data === "string" && + frame.data.length > 0 && + typeof frame.mimeType === "string" && + typeof frame.cols === "number" && + typeof frame.rows === "number" && + typeof frame.chars === "number", + ); + if (frames.length === 0) return undefined; + return { + frames, + totalChars: typeof archive.totalChars === "number" ? archive.totalChars : 0, + truncatedChars: typeof archive.truncatedChars === "number" ? archive.truncatedChars : 0, + }; +} + +/** Convert archive frames into LLM image blocks (oldest first). */ +export function snapcompactImages(archive: SnapcompactArchive): ImageContent[] { + return archive.frames.map(frame => ({ type: "image", data: frame.data, mimeType: frame.mimeType })); +} + +// ============================================================================ +// Compaction entry point +// ============================================================================ + +/** + * Run a snapcompact compaction over prepared messages. Fully local: serializes + * the discarded history, prints it onto PNG frames, merges previously + * archived frames (oldest dropped beyond the budget), and produces a + * deterministic summary explaining how to read the frames. + * + * If the previous compaction was text-based, its summary is printed at the + * head of the frame archive as `[Summary of earlier history]` so no continuity is lost. + */ +export async function snapcompactCompact( + preparation: CompactionPreparation, + options?: SnapcompactOptions, +): Promise { + const { firstKeptEntryId, tokensBefore, previousSummary, previousPreserveData, fileOps } = preparation; + if (!firstKeptEntryId) { + throw new Error("First kept entry has no ID - session may need migration"); + } + const frameSize = options?.frameSize ?? SNAPCOMPACT_FRAME_SIZE; + const maxFrames = Math.max(1, options?.maxFrames ?? SNAPCOMPACT_MAX_FRAMES); + const geometry = snapcompactGeometry(frameSize); + + const messages = preparation.messagesToSummarize.concat(preparation.turnPrefixMessages); + const llmMessages = (options?.convertToLlm ?? defaultConvertToLlm)(messages); + let archiveText = normalizeForSnapcompact(serializeConversation(llmMessages)); + + const previousArchive = getPreservedSnapcompactArchive(previousPreserveData); + const includedPreviousSummary = !previousArchive && !!previousSummary; + if (includedPreviousSummary && previousSummary) { + const head = `[Summary of earlier history] ${normalizeForSnapcompact(previousSummary)}`; + archiveText = archiveText.length > 0 ? `${head} [Recent conversation] ${archiveText}` : head; + } + + let truncatedChars = previousArchive?.truncatedChars ?? 0; + + const newFrames: SnapcompactFrame[] = []; + for (let offset = 0; offset < archiveText.length; offset += geometry.capacity) { + const chunk = archiveText.slice(offset, offset + geometry.capacity); + const rendered = renderSnapcompactFrame(chunk, frameSize); + newFrames.push({ + data: Buffer.from(rendered.png).toBase64(), + mimeType: "image/png", + cols: rendered.cols, + rows: rendered.rows, + chars: rendered.chars, + }); + // Keep the event loop responsive between native render passes. + await Bun.sleep(0); + } + + const frames = [...(previousArchive?.frames ?? []), ...newFrames]; + if (frames.length > maxFrames) { + // Pin the earliest frame: it anchors the session head (the original + // request, or the filmed summary of even older history) the way the + // LLM-summary strategies keep the original goal alive across rounds. + // Eviction removes the oldest *unpinned* frames, so the archive fades + // from the middle out — head and tail survive. With a budget of one + // frame the pin is moot; keep the newest frame instead. + const evictStart = maxFrames >= 2 ? 1 : 0; + const dropped = frames.splice(evictStart, frames.length - maxFrames); + for (const frame of dropped) truncatedChars += frame.chars; + } + const totalChars = frames.reduce((sum, frame) => sum + frame.chars, 0); + + let summary: string; + if (frames.length === 0) { + summary = "No prior history."; + } else { + summary = prompt.render(snapcompactSummaryPrompt, { + frameCount: frames.length, + multipleFrames: frames.length > 1, + cols: geometry.cols, + rows: geometry.rows, + totalChars, + truncatedChars, + includedPreviousSummary, + }); + } + const { readFiles, modifiedFiles } = computeFileLists(fileOps); + summary = upsertFileOperations(summary, readFiles, modifiedFiles); + + // A snapcompact pass replaces any provider-side replacement history; strip the + // OpenAI remote-compaction payload like the default summarizer path does. + // OpenAI remote-compaction payload like the default summarizer path does. + const basePreserve = withOpenAiRemoteCompactionPreserveData(previousPreserveData, undefined) ?? {}; + const archive: SnapcompactArchive = { frames, totalChars, truncatedChars }; + + return { + summary, + shortSummary: `Archived ${totalChars.toLocaleString()} chars of history onto ${frames.length} snapcompact frame${frames.length === 1 ? "" : "s"}`, + firstKeptEntryId, + tokensBefore, + details: { readFiles, modifiedFiles } as CompactionDetails, + preserveData: { ...basePreserve, [SNAPCOMPACT_PRESERVE_KEY]: archive }, + }; +} diff --git a/packages/agent/test/snapcompact.test.ts b/packages/agent/test/snapcompact.test.ts new file mode 100644 index 000000000..fe224304c --- /dev/null +++ b/packages/agent/test/snapcompact.test.ts @@ -0,0 +1,291 @@ +import { describe, expect, it } from "bun:test"; +import type { ImageContent } from "@oh-my-pi/pi-ai"; +import { type CompactionPreparation, type CompactionResult, estimateTokens } from "../src/compaction/compaction"; +import { createCompactionSummaryMessage, defaultConvertToLlm } from "../src/compaction/messages"; +import { + getPreservedSnapcompactArchive, + normalizeForSnapcompact, + renderSnapcompactFrame, + SNAPCOMPACT_FRAME_TOKEN_ESTIMATE, + SNAPCOMPACT_PRESERVE_KEY, + type SnapcompactArchive, + snapcompactCompact, + snapcompactGeometry, +} from "../src/compaction/snapcompact"; +import { createFileOps } from "../src/compaction/utils"; +import { createAssistantMessage, createUserMessage } from "./helpers"; + +// Small frames keep render time negligible: 320px → 64 cols x 40 rows = 2560 chars. +const TEST_FRAME_SIZE = 320; + +function makePreparation(overrides: Partial = {}): CompactionPreparation { + return { + firstKeptEntryId: "kept-1", + messagesToSummarize: [ + createUserMessage("Fix the login bug. The token expires too early!"), + createAssistantMessage([{ type: "text", text: "Fixed the TTL comparison in src/login.ts." }]), + ], + turnPrefixMessages: [], + recentMessages: [], + isSplitTurn: false, + tokensBefore: 99000, + previousSummary: undefined, + previousPreserveData: undefined, + fileOps: createFileOps(), + settings: { enabled: true, reserveTokens: 16384, keepRecentTokens: 20000 }, + ...overrides, + }; +} + +interface DecodedPng { + width: number; + height: number; + colorType: number; + /** Palette indices, one byte per pixel (filter bytes stripped). */ + pixels: Uint8Array; +} + +/** Minimal PNG reader for the encoder's own output (indexed, filter None). */ +function decodePng(png: Uint8Array): DecodedPng { + expect(Array.from(png.subarray(0, 8))).toEqual([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]); + const view = new DataView(png.buffer, png.byteOffset, png.byteLength); + let pos = 8; + let width = 0; + let height = 0; + let colorType = -1; + let depth = 0; + const idatParts: Uint8Array[] = []; + while (pos < png.length) { + const length = view.getUint32(pos); + const type = String.fromCharCode(png[pos + 4], png[pos + 5], png[pos + 6], png[pos + 7]); + const data = png.subarray(pos + 8, pos + 8 + length); + if (type === "IHDR") { + width = view.getUint32(pos + 8); + height = view.getUint32(pos + 12); + depth = data[8]; + colorType = data[9]; + } else if (type === "IDAT") { + idatParts.push(data); + } + pos += 12 + length; + } + let idatLength = 0; + for (const part of idatParts) idatLength += part.length; + const idat = new Uint8Array(idatLength); + let offset = 0; + for (const part of idatParts) { + idat.set(part, offset); + offset += part.length; + } + // Strip the zlib envelope (2-byte header + trailing Adler-32). + const raw = Bun.inflateSync(idat.subarray(2, idat.length - 4)); + const rowBytes = depth === 4 ? Math.ceil(width / 2) : width; + expect(raw.length).toBe(height * (rowBytes + 1)); + const pixels = new Uint8Array(width * height); + for (let y = 0; y < height; y++) { + expect(raw[y * (rowBytes + 1)]).toBe(0); // filter byte: None + const row = raw.subarray(y * (rowBytes + 1) + 1, (y + 1) * (rowBytes + 1)); + if (depth === 4) { + for (let x = 0; x < width; x++) { + const byte = row[x >> 1]; + pixels[y * width + x] = x % 2 === 0 ? byte >> 4 : byte & 0xf; + } + } else { + pixels.set(row, y * width); + } + } + return { width, height, colorType, pixels }; +} + +describe("normalizeForSnapcompact", () => { + it("collapses whitespace runs and folds non-Latin-1 to ASCII", () => { + expect(normalizeForSnapcompact("a\n\n\tb c\r\nd")).toBe("a b c d"); + expect(normalizeForSnapcompact("x → y ✓ “quoted” — em…")).toBe(`x -> y v "quoted" - em...`); + expect(normalizeForSnapcompact("café größe")).toBe("café größe"); // Latin-1 has glyphs + expect(normalizeForSnapcompact("box │─┌ emoji 🎞")).toBe("box |-+ emoji ?"); + }); +}); + +describe("renderSnapcompactFrame", () => { + it("produces an indexed PNG of the declared geometry with sentence-cycled ink", () => { + const geometry = snapcompactGeometry(TEST_FRAME_SIZE); + expect(geometry).toEqual({ cols: 64, rows: 40, capacity: 2560 }); + + const frame = renderSnapcompactFrame("First sentence here. Second one differs.", TEST_FRAME_SIZE); + expect(frame.cols).toBe(64); + expect(frame.rows).toBe(40); + expect(frame.chars).toBe(40); + + const decoded = decodePng(frame.png); + expect(decoded.width).toBe(TEST_FRAME_SIZE); + expect(decoded.height).toBe(TEST_FRAME_SIZE); + expect(decoded.colorType).toBe(3); // indexed color + + // Two sentences → glyphs printed in ink 1 then ink 2; background stays 0. + const used = new Set(decoded.pixels); + expect(used.has(1)).toBe(true); + expect(used.has(2)).toBe(true); + expect(used.has(3)).toBe(false); + }); + + it("caps printed characters at frame capacity", () => { + const { capacity } = snapcompactGeometry(TEST_FRAME_SIZE); + const frame = renderSnapcompactFrame("x".repeat(capacity + 500), TEST_FRAME_SIZE); + expect(frame.chars).toBe(capacity); + }); +}); + +describe("snapcompactCompact", () => { + it("archives history onto frames with a self-describing summary", async () => { + const fileOps = createFileOps(); + fileOps.read.add("src/auth.ts"); + fileOps.edited.add("src/login.ts"); + const result = await snapcompactCompact(makePreparation({ fileOps }), { frameSize: TEST_FRAME_SIZE }); + + expect(result.firstKeptEntryId).toBe("kept-1"); + expect(result.tokensBefore).toBe(99000); + // Reading instructions reflect the actual grid geometry. + expect(result.summary).toContain("64 characters per row"); + expect(result.summary).toContain("snapcompact frame"); + // File operations are upserted like every other compaction summary. + expect(result.summary).toContain(""); + expect(result.summary).toContain("src/login.ts"); + expect(result.shortSummary).toContain("snapcompact frame"); + + const archive = getPreservedSnapcompactArchive(result.preserveData); + expect(archive).toBeDefined(); + expect(archive?.frames.length).toBe(1); + expect(archive?.frames[0].mimeType).toBe("image/png"); + expect(archive?.frames[0].chars).toBe(archive?.totalChars); + expect(archive?.truncatedChars).toBe(0); + // Frame data round-trips as a decodable PNG. + const decoded = decodePng(Buffer.from(archive?.frames[0].data ?? "", "base64")); + expect(decoded.width).toBe(TEST_FRAME_SIZE); + }); + + it("splits oversized history across frames and evicts beyond the budget", async () => { + const { capacity } = snapcompactGeometry(TEST_FRAME_SIZE); + // Sentences avoid whitespace collapse shrinking the payload below 2.5 frames. + const longText = "Important fact number one. ".repeat(Math.ceil((capacity * 2.5) / 28)); + const result = await snapcompactCompact(makePreparation({ messagesToSummarize: [createUserMessage(longText)] }), { + frameSize: TEST_FRAME_SIZE, + maxFrames: 2, + }); + const archive = getPreservedSnapcompactArchive(result.preserveData); + expect(archive?.frames.length).toBe(2); + expect(archive?.truncatedChars).toBeGreaterThan(0); + expect(result.summary).toContain("dropped"); + }); + + it("evicts the oldest unpinned frames, keeping the session-head frame alive", async () => { + let previous: CompactionResult | undefined; + let headFrameData = ""; + let secondFrameData = ""; + for (let pass = 1; pass <= 4; pass++) { + previous = await snapcompactCompact( + makePreparation({ + messagesToSummarize: [createUserMessage(`Distinct turn number ${pass}.`)], + previousSummary: previous?.summary, + previousPreserveData: previous?.preserveData, + }), + { frameSize: TEST_FRAME_SIZE, maxFrames: 3 }, + ); + const archive = getPreservedSnapcompactArchive(previous.preserveData); + if (pass === 1) headFrameData = archive?.frames[0].data ?? ""; + if (pass === 2) secondFrameData = archive?.frames[1].data ?? ""; + } + const final = getPreservedSnapcompactArchive(previous?.preserveData); + expect(final?.frames.length).toBe(3); + // The head frame (original request) is pinned through every eviction; + // the archive fades from the middle out. + expect(final?.frames[0].data).toBe(headFrameData); + expect(final?.frames.some(frame => frame.data === secondFrameData)).toBe(false); + expect(final?.truncatedChars).toBeGreaterThan(0); + }); + + it("includes the previous text summary when the prior compaction was not snapcompact", async () => { + const result = await snapcompactCompact( + makePreparation({ previousSummary: "Older context: project scaffolding done." }), + { frameSize: TEST_FRAME_SIZE }, + ); + expect(result.summary).toContain("[Summary of earlier history]"); + }); + + it("carries previous frames forward and strips the OpenAI remote payload", async () => { + const first = await snapcompactCompact(makePreparation(), { frameSize: TEST_FRAME_SIZE }); + const firstArchive = getPreservedSnapcompactArchive(first.preserveData); + + const second = await snapcompactCompact( + makePreparation({ + messagesToSummarize: [createUserMessage("A new turn happened after the first compaction.")], + previousSummary: first.summary, + previousPreserveData: { + ...first.preserveData, + openaiRemoteCompaction: { provider: "openai", replacementHistory: [] }, + appKey: "kept", + }, + }), + { frameSize: TEST_FRAME_SIZE }, + ); + + const archive = getPreservedSnapcompactArchive(second.preserveData); + expect(archive?.frames.length).toBe(2); + // Oldest frame rides along unchanged, new frame appended after it. + expect(archive?.frames[0].data).toBe(firstArchive?.frames[0].data ?? ""); + // Previous archive present → previous summary is snapcompact boilerplate, not re-archived. + expect(second.summary).not.toContain("[Summary of earlier history]"); + expect(second.preserveData?.openaiRemoteCompaction).toBeUndefined(); + expect(second.preserveData?.appKey).toBe("kept"); + }); +}); + +describe("compaction summary message with snapcompact frames", () => { + const images: ImageContent[] = [ + { type: "image", data: "ZmFrZQ==", mimeType: "image/png" }, + { type: "image", data: "ZmFrZTI=", mimeType: "image/png" }, + ]; + + it("estimateTokens charges per attached frame", () => { + const bare = createCompactionSummaryMessage("summary text", 1000, new Date().toISOString()); + const withFrames = createCompactionSummaryMessage( + "summary text", + 1000, + new Date().toISOString(), + undefined, + undefined, + images, + ); + expect(estimateTokens(withFrames) - estimateTokens(bare)).toBe(2 * SNAPCOMPACT_FRAME_TOKEN_ESTIMATE); + }); + + it("defaultConvertToLlm appends frames as image blocks after the summary text", () => { + const message = createCompactionSummaryMessage( + "the snapcompact archive", + 1000, + new Date().toISOString(), + undefined, + undefined, + images, + ); + const [converted] = defaultConvertToLlm([message]); + expect(converted.role).toBe("user"); + const content = converted.content as Array<{ type: string; text?: string; data?: string }>; + expect(content.length).toBe(3); + expect(content[0].type).toBe("text"); + expect(content[0].text).toContain("the snapcompact archive"); + expect(content[1]).toEqual(images[0]); + expect(content[2]).toEqual(images[1]); + }); + + it("getPreservedSnapcompactArchive rejects malformed payloads", () => { + expect(getPreservedSnapcompactArchive(undefined)).toBeUndefined(); + expect(getPreservedSnapcompactArchive({ [SNAPCOMPACT_PRESERVE_KEY]: "nope" })).toBeUndefined(); + expect(getPreservedSnapcompactArchive({ [SNAPCOMPACT_PRESERVE_KEY]: { frames: [] } })).toBeUndefined(); + const valid: SnapcompactArchive = { + frames: [{ data: "ZmFrZQ==", mimeType: "image/png", cols: 64, rows: 40, chars: 10 }], + totalChars: 10, + truncatedChars: 0, + }; + expect(getPreservedSnapcompactArchive({ [SNAPCOMPACT_PRESERVE_KEY]: valid })).toEqual(valid); + }); +}); From 84175ce4b2103582c23fc54df19b191f6ef85117 Mon Sep 17 00:00:00 2001 From: can1357 Date: Wed, 10 Jun 2026 17:43:51 +0200 Subject: [PATCH 10/27] fix(agent): preserved queued steering across externally aborted runs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Interrupting mid-tool execution (e.g. Enter with a pending steer) drained the steering queue into the dying run — it landed in history without a response and the post-abort resume saw an empty queue, so the agent stopped instead of continuing. Steering/follow-up/aside queue polls in runLoopBody and the post-tool-call check in executeToolCalls are now skipped once the run's abort signal fires, leaving the queue intact for Agent.continue(). --- packages/agent/src/agent-loop.ts | 27 +++-- .../test/agent-session-concurrent.test.ts | 113 ++++++++++++++++++ 2 files changed, 133 insertions(+), 7 deletions(-) diff --git a/packages/agent/src/agent-loop.ts b/packages/agent/src/agent-loop.ts index d31fae0a0..1ad693207 100644 --- a/packages/agent/src/agent-loop.ts +++ b/packages/agent/src/agent-loop.ts @@ -564,8 +564,10 @@ async function runLoopBody( streamFn?: StreamFn, ): Promise { let firstTurn = true; - // Check for steering messages at start (user may have typed while waiting) - let pendingMessages: AgentMessage[] = (await config.getSteeringMessages?.()) || []; + // Check for steering messages at start (user may have typed while waiting). + // Skip when the run is already externally aborted — dequeuing would strand + // the messages in a run that is about to die. + let pendingMessages: AgentMessage[] = signal?.aborted ? [] : (await config.getSteeringMessages?.()) || []; let harmonyRetryAttempt = 0; let harmonyTruncateResumeCount = 0; @@ -743,7 +745,12 @@ async function runLoopBody( stream.push({ type: "turn_end", message, toolResults }); - const steering = steeringMessagesFromExecution ?? ((await config.getSteeringMessages?.()) || []); + // On external abort (user interrupt), leave the steering queue intact: the + // session aborts then continues, delivering the queue into a fresh run. + // Draining it here would inject the messages right before a model call that + // instantly aborts — message lands in history, agent never responds. + const steering = + steeringMessagesFromExecution ?? (signal?.aborted ? [] : (await config.getSteeringMessages?.()) || []); if (hasMoreToolCalls) { // Mid-work: fold any non-interrupting asides into the next turn alongside steering. const asides = resolveAsides(await config.getAsideMessages?.()); @@ -758,8 +765,9 @@ async function runLoopBody( // Agent would stop here. Drain non-interrupting asides + follow-up messages. await config.onBeforeYield?.(); - const asideMessages = resolveAsides(await config.getAsideMessages?.()); - const followUpMessages = (await config.getFollowUpMessages?.()) || []; + // Skip queue drains when externally aborted (same stranding hazard as above). + const asideMessages = signal?.aborted ? [] : resolveAsides(await config.getAsideMessages?.()); + const followUpMessages = signal?.aborted ? [] : (await config.getFollowUpMessages?.()) || []; if (asideMessages.length > 0 || followUpMessages.length > 0) { // Set as pending so the inner loop processes them before stopping. pendingMessages = [...asideMessages, ...followUpMessages]; @@ -1253,11 +1261,16 @@ async function executeToolCalls( })); const checkSteering = async (): Promise => { - if (!shouldInterruptImmediately || !getSteeringMessages || interruptState.triggered) { + // `signal` (external/user abort) is checked separately from the internal + // steeringAbortController: once the run is externally aborted it is + // unwinding, and draining the steering queue here would strand the + // messages in the dying run instead of leaving them for the post-abort + // continue (interruptAndFlushQueuedMessages → Agent.continue()). + if (!shouldInterruptImmediately || !getSteeringMessages || interruptState.triggered || signal?.aborted) { return; } const check = steeringCheckTail.then(async () => { - if (interruptState.triggered) return; + if (interruptState.triggered || signal?.aborted) return; const steering = await getSteeringMessages(); if (steering.length > 0) { steeringMessages = steering; diff --git a/packages/coding-agent/test/agent-session-concurrent.test.ts b/packages/coding-agent/test/agent-session-concurrent.test.ts index b03aa25eb..08efb1804 100644 --- a/packages/coding-agent/test/agent-session-concurrent.test.ts +++ b/packages/coding-agent/test/agent-session-concurrent.test.ts @@ -227,6 +227,119 @@ describe("AgentSession concurrent prompt guard", () => { expect(session.getQueuedMessages()).toEqual({ steering: [], followUp: [] }); }); + it("delivers queued steering after interrupting mid-tool execution (queue survives external abort)", async () => { + // Regression: pressing Enter with a queued steer while a tool was running + // aborted the run, but the post-abort steering poll inside executeToolCalls + // drained the queue into the dying run — the message landed in history and + // interruptAndFlushQueuedMessages saw an empty queue, so it never resumed. + const model = getBundledModel("anthropic", "claude-sonnet-4-5")!; + const callMessages: Message[][] = []; + let toolStarted = false; + + const blockingTool: AgentTool = { + name: "mock_blocker", + label: "Mock Blocker", + description: "Blocks until aborted", + parameters: z.object({}), + execute: async (_id, _args, signal) => { + toolStarted = true; + const { promise, resolve } = Promise.withResolvers(); + if (signal?.aborted) resolve(); + else signal?.addEventListener("abort", () => resolve(), { once: true }); + await promise; + return { content: [{ type: "text" as const, text: "tool aborted" }] }; + }, + }; + + const toolCallContent: ToolCall = { + type: "toolCall", + id: "call_steer_flush_001", + name: "mock_blocker", + arguments: {}, + }; + + const agent = new Agent({ + getApiKey: () => "test-key", + initialState: { model, systemPrompt: ["Test"], tools: [blockingTool] }, + convertToLlm, + streamFn: (_model, context, options) => { + const callIndex = callMessages.length; + callMessages.push([...context.messages]); + const stream = new AssistantMessageEventStream(); + const signal = options?.signal; + queueMicrotask(() => { + if (signal?.aborted) { + // Post-abort model call inside the dying run. + stream.push({ type: "error", reason: "aborted", error: createAssistantMessage("Interrupted") }); + return; + } + if (callIndex === 0) { + const partial: AssistantMessage = { + ...createAssistantMessage(""), + content: [toolCallContent], + stopReason: "toolUse", + }; + stream.push({ type: "start", partial }); + stream.push({ type: "toolcall_start", contentIndex: 0, partial }); + stream.push({ type: "toolcall_end", contentIndex: 0, toolCall: toolCallContent, partial }); + stream.push({ type: "done", reason: "toolUse", message: partial }); + return; + } + const done = createAssistantMessage("Handled steer"); + stream.push({ type: "start", partial: done }); + stream.push({ type: "done", reason: "stop", message: done }); + signal?.addEventListener( + "abort", + () => { + stream.push({ type: "error", reason: "aborted", error: createAssistantMessage("Interrupted") }); + }, + { once: true }, + ); + }); + return stream; + }, + }); + + const sessionManager = SessionManager.inMemory(); + const settings = Settings.isolated(); + const authStorage = await AuthStorage.create(path.join(tempDir, "testauth-steer-tool-abort.db")); + authStorages.push(authStorage); + const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models-steer-tool-abort.yml")); + authStorage.setRuntimeApiKey("anthropic", "test-key"); + + session = new AgentSession({ + agent, + sessionManager, + settings, + modelRegistry, + }); + + const firstPrompt = session.prompt("First message").catch(() => {}); + await waitFor(() => toolStarted); + + await session.steer("Send this now"); + expect(session.getQueuedMessages().steering).toEqual(["Send this now"]); + + await session.interruptAndFlushQueuedMessages({ reason: "Interrupted by user" }); + await firstPrompt; + + // The resumed run's model call must carry the steer message. + const lastCall = callMessages[callMessages.length - 1]; + expect( + lastCall?.some(message => { + if (typeof message.content === "string") return message.content.includes("Send this now"); + return message.content.some(content => content.type === "text" && content.text.includes("Send this now")); + }), + ).toBe(true); + + // The agent actually resumed and produced a response after the steer. + const lastAssistant = [...agent.state.messages] + .reverse() + .find((m): m is AssistantMessage => m.role === "assistant"); + expect(lastAssistant?.content).toEqual([{ type: "text", text: "Handled steer" }]); + expect(session.getQueuedMessages()).toEqual({ steering: [], followUp: [] }); + }); + it("should allow followUp() while streaming", async () => { await createSession(); From 5406edeed0a9499ade9074f32503d9c0515c7bc6 Mon Sep 17 00:00:00 2001 From: can1357 Date: Wed, 10 Jun 2026 17:44:31 +0200 Subject: [PATCH 11/27] refactor(coding-agent): made the job tool always available and dropped async-job gating MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Removes isBackgroundJobSupportEnabled and JobTool.createIf; the tool is now registered unconditionally via `new JobTool(s)`. `async.enabled` now gates async bash commands only — the task tool runs asynchronously regardless. Deletes the async/support module and its barrel re-export. --- packages/coding-agent/src/async/index.ts | 1 - packages/coding-agent/src/async/support.ts | 5 ---- .../src/config/settings-schema.ts | 2 +- packages/coding-agent/src/tools/index.ts | 2 +- packages/coding-agent/src/tools/job.ts | 7 +---- packages/coding-agent/test/tools.test.ts | 11 +------- .../test/tools/task-async-fallback.test.ts | 28 ++++++++++--------- 7 files changed, 19 insertions(+), 37 deletions(-) delete mode 100644 packages/coding-agent/src/async/support.ts diff --git a/packages/coding-agent/src/async/index.ts b/packages/coding-agent/src/async/index.ts index f04070903..7d036b77e 100644 --- a/packages/coding-agent/src/async/index.ts +++ b/packages/coding-agent/src/async/index.ts @@ -1,2 +1 @@ export * from "./job-manager"; -export * from "./support"; diff --git a/packages/coding-agent/src/async/support.ts b/packages/coding-agent/src/async/support.ts deleted file mode 100644 index 10184393a..000000000 --- a/packages/coding-agent/src/async/support.ts +++ /dev/null @@ -1,5 +0,0 @@ -import type { Settings } from "../config/settings"; - -export function isBackgroundJobSupportEnabled(settings: Pick): boolean { - return settings.get("async.enabled") || settings.get("bash.autoBackground.enabled"); -} diff --git a/packages/coding-agent/src/config/settings-schema.ts b/packages/coding-agent/src/config/settings-schema.ts index b0f266fae..66c691233 100644 --- a/packages/coding-agent/src/config/settings-schema.ts +++ b/packages/coding-agent/src/config/settings-schema.ts @@ -2512,7 +2512,7 @@ export const SETTINGS_SCHEMA = { ui: { tab: "tools", label: "Async Execution", - description: "Enable async bash commands and background task execution", + description: "Enable async bash commands", }, }, diff --git a/packages/coding-agent/src/tools/index.ts b/packages/coding-agent/src/tools/index.ts index 2a9ba22cd..750af3f65 100644 --- a/packages/coding-agent/src/tools/index.ts +++ b/packages/coding-agent/src/tools/index.ts @@ -420,7 +420,7 @@ export const BUILTIN_TOOLS: Record = { checkpoint: CheckpointTool.createIf, rewind: RewindTool.createIf, task: s => TaskTool.create(s), - job: JobTool.createIf, + job: s => new JobTool(s), irc: IrcTool.createIf, todo: s => new TodoTool(s), web_search: s => new WebSearchTool(s), diff --git a/packages/coding-agent/src/tools/job.ts b/packages/coding-agent/src/tools/job.ts index 61760e7e9..b69dc7290 100644 --- a/packages/coding-agent/src/tools/job.ts +++ b/packages/coding-agent/src/tools/job.ts @@ -3,7 +3,7 @@ import type { Component } from "@oh-my-pi/pi-tui"; import { Text } from "@oh-my-pi/pi-tui"; import { prompt } from "@oh-my-pi/pi-utils"; import * as z from "zod/v4"; -import { type AsyncJob, type AsyncJobManager, isBackgroundJobSupportEnabled } from "../async"; +import type { AsyncJob, AsyncJobManager } from "../async"; import type { RenderResultOptions } from "../extensibility/custom-tools/types"; import type { Theme } from "../modes/theme/theme"; import jobDescription from "../prompts/tools/job.md" with { type: "text" }; @@ -78,11 +78,6 @@ export class JobTool implements AgentTool { this.description = prompt.render(jobDescription); } - static createIf(session: ToolSession): JobTool | null { - if (!isBackgroundJobSupportEnabled(session.settings)) return null; - return new JobTool(session); - } - async execute( _toolCallId: string, params: JobParams, diff --git a/packages/coding-agent/test/tools.test.ts b/packages/coding-agent/test/tools.test.ts index c633afe66..81610b74d 100644 --- a/packages/coding-agent/test/tools.test.ts +++ b/packages/coding-agent/test/tools.test.ts @@ -1501,15 +1501,6 @@ function b() { expect(output).toContain("second"); expect(output).toContain("third"); }); - - it("should expose background-job tools when bash auto-background is enabled", () => { - const autoBackgroundSession = createTestToolSession( - testDir, - Settings.isolated({ "bash.autoBackground.enabled": true }), - ); - - expect(JobTool.createIf(autoBackgroundSession)).not.toBeNull(); - }); }); describe("JobTool", () => { @@ -1520,7 +1511,7 @@ function b() { const session = createTestToolSession(testDir, Settings.isolated({ "bash.autoBackground.enabled": true }), { asyncJobManager: manager, }); - const jobTool = JobTool.createIf(session)!; + const jobTool = new JobTool(session); const jobId = manager.register("bash", "test job", async () => "success"); diff --git a/packages/coding-agent/test/tools/task-async-fallback.test.ts b/packages/coding-agent/test/tools/task-async-fallback.test.ts index 7af16035d..b6a9e9d58 100644 --- a/packages/coding-agent/test/tools/task-async-fallback.test.ts +++ b/packages/coding-agent/test/tools/task-async-fallback.test.ts @@ -1,5 +1,4 @@ import { afterEach, describe, expect, it, vi } from "bun:test"; -import { AsyncJobManager } from "@oh-my-pi/pi-coding-agent/async"; import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; import { TaskTool } from "@oh-my-pi/pi-coding-agent/task"; import * as discoveryModule from "@oh-my-pi/pi-coding-agent/task/discovery"; @@ -24,15 +23,16 @@ function getFirstText(result: { content: Array<{ type: string; text?: string }> describe("task.async-fallback", () => { afterEach(() => { vi.restoreAllMocks(); - AsyncJobManager.resetForTests(); }); - it("falls back to sync execution when async is enabled but no manager is registered", async () => { + it("falls back to sync execution when the session has no job manager", async () => { // Two-stage spy: the initial discovery during `TaskTool.create` advertises - // `task` so the tool builds; the executor's later call (inside - // `#executeSync`) advertises *nothing*, forcing the unique "Unknown agent" - // message — which is only reachable from the sync codepath. Hitting it - // proves we fell back instead of returning the old hard error. + // `task` so the tool builds; the executor's later call (inside the sync + // `#runSpawn`) advertises *nothing*, forcing the unique "Unknown agent" + // message. That re-discovery only happens on the sync codepath — the + // async path resolves agents from the create-time snapshot and returns a + // job stub immediately — so hitting it proves the missing + // `session.asyncJobManager` routed us through the sync fallback. const discoverSpy = vi.spyOn(discoveryModule, "discoverAgents"); discoverSpy.mockResolvedValueOnce({ agents: [ @@ -47,18 +47,20 @@ describe("task.async-fallback", () => { }); discoverSpy.mockResolvedValue({ agents: [], projectAgentsDir: null }); - AsyncJobManager.resetForTests(); - expect(AsyncJobManager.instance()).toBeUndefined(); - - const tool = await TaskTool.create(createSession({ "async.enabled": true })); + // createSession never wires `asyncJobManager`, which is the fallback trigger. + const tool = await TaskTool.create(createSession()); const result = await tool.execute("tool-1", { agent: "task", - tasks: [{ id: "One", description: "label", assignment: "Do the thing." }], + id: "One", + description: "label", + assignment: "Do the thing.", } as TaskParams); const text = getFirstText(result); expect(text).toContain('Unknown agent "task"'); - expect(text).not.toContain("no async job manager is available"); + expect(text).toContain("Available: none"); + // create + sync-path re-discovery; the async path would have stopped at one. + expect(discoverSpy).toHaveBeenCalledTimes(2); }); }); From 1654c759ca9a4250c7a03ca534b527560c1cce88 Mon Sep 17 00:00:00 2001 From: can1357 Date: Wed, 10 Jun 2026 17:46:02 +0200 Subject: [PATCH 12/27] feat(coding-agent): added AgentLifecycleManager for idle/parked subagents MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduces AgentLifecycleManager: when the task executor adopts a finished subagent the manager arms a TTL timer on idle, parks the agent on expiry (disposes the live session, keeps the AgentRef + sessionFile), and revives it on demand via an injected reviver. Only this manager flips parked ↔ idle. AgentRegistry now annotates session: AgentRef["session"] as null exactly when parked/aborted, and sdk.ts wires the lifecycle dispose into main-session teardown, derives agentKind once, and gates ref unregistration on parking so a parked agent stays addressable (history://, revive). --- .../src/registry/agent-lifecycle.ts | 218 ++++++++++++++++ .../src/registry/agent-registry.ts | 21 +- packages/coding-agent/src/sdk.ts | 38 ++- .../test/registry/agent-lifecycle.test.ts | 234 ++++++++++++++++++ 4 files changed, 497 insertions(+), 14 deletions(-) create mode 100644 packages/coding-agent/src/registry/agent-lifecycle.ts create mode 100644 packages/coding-agent/test/registry/agent-lifecycle.test.ts diff --git a/packages/coding-agent/src/registry/agent-lifecycle.ts b/packages/coding-agent/src/registry/agent-lifecycle.ts new file mode 100644 index 000000000..ad2ef4f16 --- /dev/null +++ b/packages/coding-agent/src/registry/agent-lifecycle.ts @@ -0,0 +1,218 @@ +/** + * AgentLifecycleManager - Owns the idle → parked → revived lifecycle of + * adopted subagents. + * + * The task executor hands a finished agent over via {@link AgentLifecycleManager.adopt}; + * from then on the manager arms a TTL timer whenever the agent goes `idle`, + * parks it on expiry (disposes the live session, keeps the AgentRef + + * sessionFile), and revives it on demand through + * {@link AgentLifecycleManager.ensureLive}. Only this manager flips + * `parked` ↔ `idle`. + */ + +import { logger } from "@oh-my-pi/pi-utils"; +import type { AgentSession } from "../session/agent-session"; +import { AgentRegistry, MAIN_AGENT_ID, type RegistryEvent } from "./agent-registry"; + +export type AgentReviver = () => Promise; + +export interface AdoptOptions { + /** TTL before an idle agent is parked. <= 0 disables parking. */ + idleTtlMs: number; + /** Recreates a live AgentSession from the ref's sessionFile. Absent => not resumable after park (e.g. isolated runs). */ + revive?: AgentReviver; +} + +interface AdoptedAgent { + idleTtlMs: number; + revive?: AgentReviver; + timer?: NodeJS.Timeout; +} + +export class AgentLifecycleManager { + static #global: AgentLifecycleManager | undefined; + + static global(): AgentLifecycleManager { + if (!AgentLifecycleManager.#global) { + AgentLifecycleManager.#global = new AgentLifecycleManager(); + } + return AgentLifecycleManager.#global; + } + + /** Reset the global manager. Test-only. */ + static resetGlobalForTests(): void { + const current = AgentLifecycleManager.#global; + if (current) { + current.#unsubscribe?.(); + current.#unsubscribe = undefined; + for (const adopted of current.#adopted.values()) { + clearTimeout(adopted.timer); + } + current.#adopted.clear(); + current.#revivals.clear(); + current.#parking.clear(); + } + AgentLifecycleManager.#global = undefined; + } + + readonly #registry: AgentRegistry; + readonly #adopted = new Map(); + /** Ids whose session is being disposed by {@link park} right now. */ + readonly #parking = new Set(); + /** In-flight revives, so concurrent {@link ensureLive} calls coalesce. */ + readonly #revivals = new Map>(); + #unsubscribe: (() => void) | undefined; + + constructor(registry: AgentRegistry = AgentRegistry.global()) { + this.#registry = registry; + this.#unsubscribe = registry.onChange(event => this.#onRegistryEvent(event)); + } + + /** + * Take ownership of a finished subagent. Caller has already set registry + * status to "idle". Arms the TTL timer (idleTtlMs <= 0 adopts without one). + */ + adopt(id: string, opts: AdoptOptions): void { + if (id === MAIN_AGENT_ID) return; + if (!this.#registry.get(id)) { + logger.warn("AgentLifecycleManager.adopt: unknown agent id", { id }); + return; + } + const existing = this.#adopted.get(id); + clearTimeout(existing?.timer); + const adopted: AdoptedAgent = { idleTtlMs: opts.idleTtlMs, revive: opts.revive }; + this.#adopted.set(id, adopted); + this.#armTimer(id, adopted); + } + + /** True if the id is adopted (parked or live). */ + has(id: string): boolean { + return this.#adopted.has(id); + } + + /** True while {@link park} is disposing this agent's session (lets dispose hooks distinguish park from teardown). */ + isParking(id: string): boolean { + return this.#parking.has(id); + } + + /** + * Dispose the live session, detach it from the registry, and mark the + * agent `parked`. No-op unless the id is adopted and live. + */ + async park(id: string): Promise { + const adopted = this.#adopted.get(id); + if (!adopted) return; + const ref = this.#registry.get(id); + if (!ref?.session) return; + if (adopted.timer) { + clearTimeout(adopted.timer); + adopted.timer = undefined; + } + this.#parking.add(id); + try { + try { + await ref.session.dispose(); + } catch (error) { + logger.warn("AgentLifecycleManager.park: session dispose failed", { id, error: String(error) }); + } + this.#registry.detachSession(id); + this.#registry.setStatus(id, "parked"); + } finally { + this.#parking.delete(id); + } + } + + /** + * Return the live session, reviving from the sessionFile if parked. + * Throws a plain Error if the id is unknown or parked without a reviver. + * Concurrent calls share one in-flight revive. + */ + async ensureLive(id: string): Promise { + const ref = this.#registry.get(id); + if (!ref) { + throw new Error( + `Unknown agent "${id}" — it was never registered or has been released. If a transcript exists, read history://${id}.`, + ); + } + if (ref.session) return ref.session; + const inflight = this.#revivals.get(id); + if (inflight) return inflight; + const adopted = this.#adopted.get(id); + if (ref.status !== "parked" || !adopted?.revive) { + throw new Error( + `Agent "${id}" is ${ref.status} and cannot be revived${adopted?.revive ? "" : " (no reviver registered)"}. Its transcript remains readable at history://${id}.`, + ); + } + const revival = this.#revive(id, adopted.revive, ref.sessionFile); + this.#revivals.set(id, revival); + try { + return await revival; + } finally { + this.#revivals.delete(id); + } + } + + /** Hard removal: dispose if live, unregister from registry, drop timers. */ + async release(id: string): Promise { + const adopted = this.#adopted.get(id); + clearTimeout(adopted?.timer); + this.#adopted.delete(id); + const ref = this.#registry.get(id); + if (ref?.session) { + try { + await ref.session.dispose(); + } catch (error) { + logger.warn("AgentLifecycleManager.release: session dispose failed", { id, error: String(error) }); + } + } + this.#registry.unregister(id); + } + + /** Teardown everything (process exit / main session dispose). */ + async dispose(): Promise { + this.#unsubscribe?.(); + this.#unsubscribe = undefined; + const ids = [...this.#adopted.keys()]; + await Promise.all(ids.map(id => this.release(id))); + this.#revivals.clear(); + this.#parking.clear(); + } + + async #revive(id: string, revive: AgentReviver, sessionFile: string | null): Promise { + const session = await revive(); + this.#registry.attachSession(id, session, sessionFile); + // Emits status_changed → "idle", which re-arms the TTL timer below. + this.#registry.setStatus(id, "idle"); + return session; + } + + #armTimer(id: string, adopted: AdoptedAgent): void { + if (adopted.idleTtlMs <= 0) return; + clearTimeout(adopted.timer); + const timer = setTimeout(() => { + adopted.timer = undefined; + void this.park(id); + }, adopted.idleTtlMs); + timer.unref?.(); + adopted.timer = timer; + } + + #onRegistryEvent(event: RegistryEvent): void { + const adopted = this.#adopted.get(event.ref.id); + if (!adopted) return; + if (event.type === "removed") { + clearTimeout(adopted.timer); + this.#adopted.delete(event.ref.id); + return; + } + if (event.type !== "status_changed") return; + if (event.ref.status === "running") { + if (adopted.timer) { + clearTimeout(adopted.timer); + adopted.timer = undefined; + } + } else if (event.ref.status === "idle") { + this.#armTimer(event.ref.id, adopted); + } + } +} diff --git a/packages/coding-agent/src/registry/agent-registry.ts b/packages/coding-agent/src/registry/agent-registry.ts index 271b9d90b..f791fa225 100644 --- a/packages/coding-agent/src/registry/agent-registry.ts +++ b/packages/coding-agent/src/registry/agent-registry.ts @@ -1,16 +1,26 @@ /** - * AgentRegistry - Process-global registry of live AgentSession instances. + * AgentRegistry - Process-global registry of agents (the main session plus + * every subagent), keyed by stable id. * - * Tracks every alive agent (the main session plus every subagent) so the - * `irc` tool can address peers by id. Sessions are registered explicitly at - * creation and removed when the owner releases them. + * Tracks each agent's status and (when live) its AgentSession so peers can be + * addressed by id (`irc`, `task resume`, `history://`). Sessions are + * registered explicitly at creation; finished agents stay registered as + * `idle` (live) or `parked` (session disposed, ref + sessionFile retained for + * revival) and are only removed on explicit release/teardown. */ import type { AgentSession } from "../session/agent-session"; export const MAIN_AGENT_ID = "Main"; -export type AgentStatus = "running" | "idle" | "completed" | "aborted"; +/** + * - `running`: a turn is in flight. + * - `idle`: live AgentSession in memory, awaiting work. Finished agents are + * `idle`, not removed. + * - `parked`: session disposed; AgentRef + sessionFile retained, revivable. + * - `aborted`: hard-killed, terminal. + */ +export type AgentStatus = "running" | "idle" | "parked" | "aborted"; export type AgentKind = "main" | "sub"; export interface AgentRef { @@ -19,6 +29,7 @@ export interface AgentRef { kind: AgentKind; parentId?: string; status: AgentStatus; + /** Null exactly when parked/aborted. */ session: AgentSession | null; sessionFile: string | null; createdAt: number; diff --git a/packages/coding-agent/src/sdk.ts b/packages/coding-agent/src/sdk.ts index d128cce56..05e6664bd 100644 --- a/packages/coding-agent/src/sdk.ts +++ b/packages/coding-agent/src/sdk.ts @@ -34,7 +34,7 @@ import { Snowflake, } from "@oh-my-pi/pi-utils"; import chalk from "chalk"; -import { type AsyncJob, AsyncJobManager, isBackgroundJobSupportEnabled } from "./async"; +import { type AsyncJob, AsyncJobManager } from "./async"; import { loadCapability } from "./capability"; import { type Rule, ruleCapability, setActiveRules } from "./capability/rule"; import { bucketRules } from "./capability/rule-buckets"; @@ -93,6 +93,7 @@ import { createSessionMemoryRuntimeContext, resolveMemoryBackend } from "./memor import type { MnemopiSessionState } from "./mnemopi/state"; import asyncResultTemplate from "./prompts/tools/async-result.md" with { type: "text" }; import lateDiagnosticTemplate from "./prompts/tools/lsp-late-diagnostic.md" with { type: "text" }; +import { AgentLifecycleManager } from "./registry/agent-lifecycle"; import { AgentRegistry, MAIN_AGENT_ID } from "./registry/agent-registry"; import { collectEnvSecrets, @@ -1293,7 +1294,6 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} let hasSession = false; let hasRegistered = false; const enableLsp = options.enableLsp ?? true; - const backgroundJobsEnabled = isBackgroundJobSupportEnabled(settings); const asyncMaxJobs = Math.min(100, Math.max(1, settings.get("async.maxJobs") ?? 100)); const ASYNC_INLINE_RESULT_MAX_CHARS = 12_000; const ASYNC_PREVIEW_MAX_CHARS = 4_000; @@ -1326,7 +1326,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} // (issue #1923). The `instance()` guard means later sessions also skip // constructing an orphaned manager that nothing would ever route to. const asyncJobManager = - backgroundJobsEnabled && !options.parentTaskPrefix && !AsyncJobManager.instance() + !options.parentTaskPrefix && !AsyncJobManager.instance() ? new AsyncJobManager({ maxRunningJobs: asyncMaxJobs, onJobComplete: async (jobId, result, job) => { @@ -1351,6 +1351,17 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} const resolvedAgentId = options.agentId ?? options.parentTaskPrefix ?? MAIN_AGENT_ID; const resolvedAgentDisplayName = options.agentDisplayName ?? ((options.taskDepth ?? 0) > 0 || options.parentTaskPrefix ? "sub" : "main"); + const agentKind = (options.taskDepth ?? 0) > 0 || options.parentTaskPrefix ? ("sub" as const) : ("main" as const); + /** + * Forget the agent ref on teardown — unless the agent is being parked (or is + * already parked). Parking disposes the session but keeps the ref addressable + * (history://, revive); only process teardown / explicit kill unregisters. + */ + const unregisterUnlessParked = (): void => { + if (agentRegistry.get(resolvedAgentId)?.status === "parked") return; + if (AgentLifecycleManager.global().isParking(resolvedAgentId)) return; + agentRegistry.unregister(resolvedAgentId); + }; const evalKernelOwnerId = `agent-session:${Snowflake.next()}`; try { @@ -1409,7 +1420,6 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} getTurnBudget: () => sessionManager.getTurnBudget(), recordEvalSubagentUsage: output => sessionManager.recordEvalSubagentOutput(output), getClientBridge: () => session?.clientBridge, - getCompactContext: () => session.formatCompactContext(), queueDeferredDiagnostics: entry => session?.yieldQueue.enqueue(LSP_LATE_DIAGNOSTIC_MESSAGE_TYPE, entry), bumpFileMutationVersion: path => { const next = (fileMutationVersions.get(path) ?? 0) + 1; @@ -2083,7 +2093,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} agentRegistry.register({ id: resolvedAgentId, displayName: resolvedAgentDisplayName, - kind: (options.taskDepth ?? 0) > 0 || options.parentTaskPrefix ? "sub" : "main", + kind: agentKind, parentId: options.parentTaskPrefix, session: null, sessionFile: sessionManager.getSessionFile() ?? null, @@ -2320,7 +2330,6 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} ttsrManager, obfuscator, agentId: resolvedAgentId, - agentRegistry, providerSessionId: options.providerSessionId, parentEvalSessionId: options.parentEvalSessionId, }); @@ -2341,15 +2350,26 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} // Attach the live session to the pre-registered ref so peers can route IRC // messages here. Refresh sessionFile in case it was unavailable at pre-register - // time. The dispose wrapper below unregisters on teardown. + // time. The dispose wrapper below unregisters on teardown (unless parked). agentRegistry.attachSession(resolvedAgentId, session, sessionManager.getSessionFile() ?? null); { const originalDispose = session.dispose.bind(session); session.dispose = async () => { try { + // Reject new session work (Python/eval starts) the moment disposal + // begins — the lifecycle await below opens an async gap before + // AgentSession.dispose() would otherwise set its guards. + session.beginDispose(); + if (agentKind === "main") { + // Top-level teardown owns the global agent lifecycle: park timers, + // adopted subagent sessions, revivers. Tear it down while shared + // resources (kernels, MCP, LSP) are still live. Subagent disposal + // must NOT touch the global lifecycle. + await AgentLifecycleManager.global().dispose(); + } await originalDispose(); } finally { - agentRegistry.unregister(resolvedAgentId); + unregisterUnlessParked(); unsubscribeCredentialDisabled?.(); } }; @@ -2502,7 +2522,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} if (hasSession) { await session.dispose(); } else { - if (hasRegistered) agentRegistry.unregister(resolvedAgentId); + if (hasRegistered) unregisterUnlessParked(); if (asyncJobManager) { if (AsyncJobManager.instance() === asyncJobManager) { AsyncJobManager.setInstance(undefined); diff --git a/packages/coding-agent/test/registry/agent-lifecycle.test.ts b/packages/coding-agent/test/registry/agent-lifecycle.test.ts new file mode 100644 index 000000000..2170b33bb --- /dev/null +++ b/packages/coding-agent/test/registry/agent-lifecycle.test.ts @@ -0,0 +1,234 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from "bun:test"; +import { AgentLifecycleManager } from "@oh-my-pi/pi-coding-agent/registry/agent-lifecycle"; +import { AgentRegistry, MAIN_AGENT_ID } from "@oh-my-pi/pi-coding-agent/registry/agent-registry"; +import type { AgentSession } from "@oh-my-pi/pi-coding-agent/session/agent-session"; + +interface SessionStub { + session: AgentSession; + disposeCalls: () => number; +} + +/** Minimal session: the lifecycle manager only ever calls dispose() on it. */ +function makeSessionStub(dispose?: () => Promise): SessionStub { + let calls = 0; + const stub = { + dispose: async () => { + calls++; + await dispose?.(); + }, + }; + return { session: stub as unknown as AgentSession, disposeCalls: () => calls }; +} + +function deferred(): { promise: Promise; resolve: () => void } { + let resolve!: () => void; + const promise = new Promise(r => { + resolve = r; + }); + return { promise, resolve }; +} + +/** Settle the async park chain (timer callback → park() → dispose → setStatus). */ +async function flushAsync(): Promise { + for (let i = 0; i < 5; i++) await Promise.resolve(); +} + +const TTL = 20; + +describe("AgentLifecycleManager", () => { + let registry: AgentRegistry; + let lifecycle: AgentLifecycleManager; + + beforeEach(() => { + AgentRegistry.resetGlobalForTests(); + AgentLifecycleManager.resetGlobalForTests(); + registry = AgentRegistry.global(); + lifecycle = AgentLifecycleManager.global(); + }); + afterEach(() => { + vi.useRealTimers(); + vi.restoreAllMocks(); + AgentLifecycleManager.resetGlobalForTests(); + AgentRegistry.resetGlobalForTests(); + }); + + function registerIdleSub(id: string, session: AgentSession | null, sessionFile: string | null = `/tmp/${id}.jsonl`) { + return registry.register({ id, displayName: "task", kind: "sub", session, sessionFile, status: "idle" }); + } + + it("adopt arms the TTL: an idle agent is parked — session disposed, ref + sessionFile retained", async () => { + vi.useFakeTimers(); + const stub = makeSessionStub(); + registerIdleSub("1-Sub", stub.session, "/tmp/1-Sub.jsonl"); + lifecycle.adopt("1-Sub", { idleTtlMs: TTL }); + + vi.advanceTimersByTime(TTL); + await flushAsync(); + + const ref = registry.get("1-Sub"); + expect(stub.disposeCalls()).toBe(1); + expect(ref?.status).toBe("parked"); + expect(ref?.session).toBeNull(); + expect(ref?.sessionFile).toBe("/tmp/1-Sub.jsonl"); + expect(lifecycle.has("1-Sub")).toBe(true); + }); + + it("running disarms the timer; returning to idle re-arms a fresh TTL", async () => { + vi.useFakeTimers(); + const stub = makeSessionStub(); + registerIdleSub("2-Sub", stub.session); + lifecycle.adopt("2-Sub", { idleTtlMs: TTL }); + registry.setStatus("2-Sub", "running"); + + vi.advanceTimersByTime(TTL * 10); + await flushAsync(); + expect(registry.get("2-Sub")?.status).toBe("running"); + expect(registry.get("2-Sub")?.session).toBe(stub.session); + expect(stub.disposeCalls()).toBe(0); + + registry.setStatus("2-Sub", "idle"); + vi.advanceTimersByTime(TTL); + await flushAsync(); + expect(registry.get("2-Sub")?.status).toBe("parked"); + expect(stub.disposeCalls()).toBe(1); + }); + + it("ensureLive revives a parked agent through its reviver and flips it back to idle", async () => { + const revived = makeSessionStub(); + registry.register({ + id: "3-Sub", + displayName: "task", + kind: "sub", + session: null, + sessionFile: "/tmp/3-Sub.jsonl", + status: "parked", + }); + lifecycle.adopt("3-Sub", { idleTtlMs: 0, revive: async () => revived.session }); + + const session = await lifecycle.ensureLive("3-Sub"); + + expect(session).toBe(revived.session); + const ref = registry.get("3-Sub"); + expect(ref?.status).toBe("idle"); + expect(ref?.session).toBe(revived.session); + expect(ref?.sessionFile).toBe("/tmp/3-Sub.jsonl"); + }); + + it("concurrent ensureLive calls during a slow revive coalesce into one reviver run", async () => { + const gate = deferred(); + const revived = makeSessionStub(); + let reviverRuns = 0; + registry.register({ + id: "4-Sub", + displayName: "task", + kind: "sub", + session: null, + sessionFile: "/tmp/4-Sub.jsonl", + status: "parked", + }); + lifecycle.adopt("4-Sub", { + idleTtlMs: 0, + revive: async () => { + reviverRuns++; + await gate.promise; + return revived.session; + }, + }); + + const first = lifecycle.ensureLive("4-Sub"); + const second = lifecycle.ensureLive("4-Sub"); + gate.resolve(); + const [a, b] = await Promise.all([first, second]); + + expect(reviverRuns).toBe(1); + expect(a).toBe(revived.session); + expect(b).toBe(revived.session); + }); + + it("ensureLive on an unknown id throws and points at history://", async () => { + await expect(lifecycle.ensureLive("9-Ghost")).rejects.toThrow(/history:\/\/9-Ghost/); + }); + + it("ensureLive on a parked agent without a reviver throws as not revivable", async () => { + registry.register({ id: "5-Sub", displayName: "task", kind: "sub", session: null, status: "parked" }); + lifecycle.adopt("5-Sub", { idleTtlMs: 0 }); + + await expect(lifecycle.ensureLive("5-Sub")).rejects.toThrow(/cannot be revived.*no reviver registered/); + }); + + it("release disposes a live adopted agent, unregisters it, and leaves no pending park", async () => { + vi.useFakeTimers(); + const stub = makeSessionStub(); + registerIdleSub("6-Sub", stub.session); + lifecycle.adopt("6-Sub", { idleTtlMs: TTL }); + + await lifecycle.release("6-Sub"); + + expect(stub.disposeCalls()).toBe(1); + expect(registry.get("6-Sub")).toBeUndefined(); + expect(lifecycle.has("6-Sub")).toBe(false); + + // The disarmed timer must not fire a late park (which would double-dispose). + vi.advanceTimersByTime(TTL * 10); + await flushAsync(); + expect(stub.disposeCalls()).toBe(1); + expect(registry.get("6-Sub")).toBeUndefined(); + }); + + it("adopt(Main) is a no-op: Main is never adopted or parked", async () => { + vi.useFakeTimers(); + const stub = makeSessionStub(); + registry.register({ + id: MAIN_AGENT_ID, + displayName: "main", + kind: "main", + session: stub.session, + status: "idle", + }); + lifecycle.adopt(MAIN_AGENT_ID, { idleTtlMs: TTL }); + + expect(lifecycle.has(MAIN_AGENT_ID)).toBe(false); + vi.advanceTimersByTime(TTL * 10); + await flushAsync(); + expect(registry.get(MAIN_AGENT_ID)?.status).toBe("idle"); + expect(registry.get(MAIN_AGENT_ID)?.session).toBe(stub.session); + expect(stub.disposeCalls()).toBe(0); + }); + + it("isParking is true exactly while park's dispose is in flight; parked only after it completes", async () => { + const gate = deferred(); + const stub = makeSessionStub(() => gate.promise); + registerIdleSub("7-Sub", stub.session); + lifecycle.adopt("7-Sub", { idleTtlMs: 0 }); + + // park() runs synchronously up to `await session.dispose()`, which we hold open. + const parking = lifecycle.park("7-Sub"); + + expect(stub.disposeCalls()).toBe(1); + expect(lifecycle.isParking("7-Sub")).toBe(true); + expect(registry.get("7-Sub")).toBeDefined(); + expect(registry.get("7-Sub")?.status).toBe("idle"); // not yet flipped + + gate.resolve(); + await parking; + + expect(lifecycle.isParking("7-Sub")).toBe(false); + expect(registry.get("7-Sub")?.status).toBe("parked"); + expect(registry.get("7-Sub")?.session).toBeNull(); + }); + + it("idleTtlMs <= 0 adopts without a timer: the agent never parks", async () => { + vi.useFakeTimers(); + const stub = makeSessionStub(); + registerIdleSub("8-Sub", stub.session); + lifecycle.adopt("8-Sub", { idleTtlMs: 0 }); + + vi.advanceTimersByTime(60_000); + await flushAsync(); + const ref = registry.get("8-Sub"); + expect(ref?.status).toBe("idle"); + expect(ref?.session).toBe(stub.session); + expect(stub.disposeCalls()).toBe(0); + expect(lifecycle.has("8-Sub")).toBe(true); + }); +}); From fcb8663de8ed6503c86d8be28cc9d1d068798160 Mon Sep 17 00:00:00 2001 From: can1357 Date: Wed, 10 Jun 2026 17:47:47 +0200 Subject: [PATCH 13/27] feat(coding-agent): reworked irc to a send/wait/inbox/list mailbox bus MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces the blocking auto-reply IRC turn with a process-global IrcBus and a four-op tool (send/wait/inbox/list). send is fire-and-forget with per-recipient delivery receipts (injected/woken/revived/failed); replies become real turns by the recipient, observed via wait (or the send await:true sugar). Bounded per-agent mailboxes (cap 100) drop oldest on overflow; AgentSession.deliverIrcMessage folds an in-flight delivery in as a non-interrupting aside at the next step boundary, or starts a real wake turn when idle. AgentRegistry/lifecycle handle the idle→woken / parked→revived transitions, so messaging a non-running peer brings it back. Adds a dedicated TUI renderer (directional headers, delivery-outcome coloring, quoted bodies, per-recipient receipt trees, status-badged peer lists with unread counts). irc.timeoutMs is now the default timeout for wait / send await:true. AgentSession sheds the agentRegistry config field, the dedupeIrcReply → dedupeEphemeralReply rename (now used by /btw and /omfg), and the background-channel exchange queue / forwardIrcRelayToMain plumbing the auto-reply model needed. --- docs/tools/irc.md | 149 ++-- .../src/config/settings-schema.ts | 3 +- packages/coding-agent/src/irc/bus.ts | 275 +++++++ .../src/prompts/system/irc-incoming.md | 7 +- .../coding-agent/src/prompts/tools/irc.md | 48 +- .../coding-agent/src/session/agent-session.ts | 248 +++---- packages/coding-agent/src/tools/index.ts | 10 +- packages/coding-agent/src/tools/irc.ts | 663 ++++++++++++----- packages/coding-agent/src/tools/renderers.ts | 2 + .../test/tools/irc-renderer.test.ts | 256 +++++++ packages/coding-agent/test/tools/irc.test.ts | 682 +++++++++++------- 11 files changed, 1649 insertions(+), 694 deletions(-) create mode 100644 packages/coding-agent/src/irc/bus.ts create mode 100644 packages/coding-agent/test/tools/irc-renderer.test.ts diff --git a/docs/tools/irc.md b/docs/tools/irc.md index bf2ff70a3..bf4e79a03 100644 --- a/docs/tools/irc.md +++ b/docs/tools/irc.md @@ -1,120 +1,99 @@ # irc -> Send short prose messages to other live agents in the current process. +> Send and receive messages between agents over a process-global mailbox bus. ## Source - Entry: `packages/coding-agent/src/tools/irc.ts` - Model-facing prompt: `packages/coding-agent/src/prompts/tools/irc.md` - Key collaborators: - - `packages/coding-agent/src/registry/agent-registry.ts` — process-global live agent directory. - - `packages/coding-agent/src/session/agent-session.ts` — side-channel reply generation and history injection. - - `packages/coding-agent/src/prompts/system/irc-incoming.md` — no-tools auto-reply prompt. - - `packages/coding-agent/src/tools/index.ts` — tool availability gating. - - `packages/coding-agent/src/config/settings-schema.ts` — `irc.enabled` default. + - `packages/coding-agent/src/irc/bus.ts` — process-global `IrcBus`: per-agent mailboxes, delivery, waiter matching. + - `packages/coding-agent/src/registry/agent-registry.ts` — process-global agent directory and status. + - `packages/coding-agent/src/registry/agent-lifecycle.ts` — revival of parked recipients on direct send. + - `packages/coding-agent/src/session/agent-session.ts` — `deliverIrcMessage(...)`: recipient-side injection and wake turns. + - `packages/coding-agent/src/prompts/system/irc-incoming.md` — incoming-message rendering for the recipient. + - `packages/coding-agent/src/config/settings-schema.ts` — `irc.enabled`, `irc.timeoutMs`. - `packages/coding-agent/src/modes/controllers/event-controller.ts` — renders IRC events into chat UI. - - `packages/coding-agent/src/modes/utils/ui-helpers.ts` — formats `[IRC]` transcript lines. - - `packages/coding-agent/src/task/executor.ts` — carries `irc.enabled` into subagents. ## Inputs -### `op: "list"` - | Field | Type | Required | Description | | --- | --- | --- | --- | -| `op` | `"list"` | Yes | Lists peers visible to the caller. | - -### `op: "send"` - -| Field | Type | Required | Description | -| --- | --- | --- | --- | -| `op` | `"send"` | Yes | Sends one message to one peer or to `"all"`. | -| `to` | `string` | Yes | Peer id such as `Main`, or `"all"` for broadcast. Whitespace is trimmed. | -| `message` | `string` | Yes | Message body. Whitespace is trimmed; empty-after-trim is rejected. | -| `awaitReply` | `boolean` | No | Wait for prose replies. Defaults to `true` for direct messages and `false` for `to: "all"`. | +| `op` | `"send" \| "wait" \| "inbox" \| "list"` | Yes | Operation. | +| `to` | `string` | `send` | Recipient agent id, or `"all"` for broadcast. Whitespace trimmed; self-send rejected. | +| `message` | `string` | `send` | Message body. Empty-after-trim is rejected. | +| `replyTo` | `string` | No | `send`: message id being answered. | +| `await` | `boolean` | No | `send`: after delivery, block until the next message from that peer arrives (round-trip sugar). Invalid with `to: "all"`. | +| `from` | `string` | No | `wait`: only accept a message from this agent id. | +| `timeoutMs` | `number` | No | `wait` / `send await:true`: timeout in milliseconds; `0` waits indefinitely. Defaults to `irc.timeoutMs`. | +| `peek` | `boolean` | No | `inbox`: list messages without consuming them. | ## Outputs - Single-shot `AgentToolResult`; no streaming updates. -- `content` is one text block. - - `list` returns either `No other live agents.` or a bullet list headed by ` peer(s):`. - - `send` returns delivery summary text, then optional `## Replies`, `## Failed`, and `Unknown / unavailable peers:` sections. -- `details` is structured metadata: - - `list`: `{ op, from, peers, channels }` - - `send`: `{ op, from, to, delivered, replies?, failed?, notFound? }` -- The tool does not return raw IRC frames, message ids, or a transcript object. +- `content` is one text block: + - `list`: `No other agents.` or ` peer(s):` bullets — `id [displayName · kind · status]` plus unread count, parent, and last-activity age; a footer notes that parked agents are revived automatically when messaged. + - `send`: per-recipient delivery receipts (`injected` / `woken` / `revived` / `failed — `); with `await: true`, the reply body or a clean no-reply timeout note. + - `wait`: the consumed message as `[] : ` (with a reply-to tag), or `No message within .` + - `inbox`: `Inbox empty.` or ` message(s):` bullets. +- `details: IrcDetails`: `{ op, from?, to?, receipts?, waited?, inbox?, peers? }`. `waited` is `null` when a wait timed out; `receipts` carry `{ to, outcome, error? }`. ## Flow -1. `IrcTool.createIf` only constructs the tool when `irc.enabled` is on and the session has both an `AgentRegistry` and `getAgentId` (`packages/coding-agent/src/tools/irc.ts`). -2. Tool discovery adds another gate in `packages/coding-agent/src/tools/index.ts`: if the caller is `Main` and `async.enabled` is off, `irc` is hidden because the main agent cannot talk to concurrent peers in sync mode. -3. `execute` resolves the process-global registry and sender id. Missing either returns a text error result instead of throwing. -4. `op: "list"` calls `registry.listVisibleTo(senderId)`, which exposes every other agent in flat namespace whose status is `running` or `idle` (`packages/coding-agent/src/registry/agent-registry.ts`). -5. `list` formats human-readable lines and returns `channels` as `['all', ...peerIds]`. These are logical targets only; there is no channel join state. -6. `op: "send"` trims `to` and `message`; missing values produce text errors. -7. `send` resolves targets: - - `to === "all"`: all visible peers. - - otherwise: one exact registry id, excluding self and excluding peers not in `running`/`idle`. -8. `send` chooses `awaitReply = params.awaitReply ?? !isBroadcast`. -9. Each target is dispatched in parallel via `target.session.respondAsBackground(...)`. One slow or failing peer does not block dispatch to the others. -10. `respondAsBackground` emits an `irc_message` session event, forwards a display-only relay to the main session UI, and either: - - queues just the incoming message for later history injection when `awaitReply === false`, or - - renders `packages/coding-agent/src/prompts/system/irc-incoming.md`, runs `runEphemeralTurn` with `toolChoice: "none"`, emits an auto-reply event, then queues both incoming and reply messages for history injection. -11. Deferred injection waits until the recipient is no longer streaming; `#flushPendingBackgroundExchanges` appends the custom messages through normal `message_start`/`message_end` external events so persistence and listeners see them. -12. Dispatch waits are bounded by `irc.timeoutMs` (default `120_000` ms). A value of `0` disables the local timeout; parent aborts still abort the dispatch. -13. `send` aggregates `delivered`, `replies`, `failed`, and `notFound`, then returns one text summary plus matching `details`. +1. `IrcTool.createIf` constructs the tool only when `irc.enabled` is on and the session has both an `AgentRegistry` and `getAgentId`. There is no longer a main-agent gate on `async.enabled` — the main agent is never sync-blocked. +2. `execute` resolves the registry and sender id; missing either returns a text error result instead of throwing. +3. `op: "list"`: `registry.list()` minus self and minus `aborted` agents — `parked` peers ARE listed. Each row includes the unread count from `IrcBus.unreadCount(...)` and last activity. +4. `op: "send"` validates `to`/`message`, rejects self-sends, and rejects `await` with `to: "all"`. +5. Target resolution: broadcasts fan out to `registry.listVisibleTo(senderId)` (live peers only — `running`/`idle`; reviving every parked agent on a broadcast would be a stampede). Direct sends go through the bus unfiltered, so a parked recipient is revived. +6. `IrcBus.send(...)` is fire-and-forget — it never blocks on the recipient generating anything. Delivery by recipient status: + - `running` → message enqueued and injected as a non-interrupting aside at the recipient's next step boundary (`AgentSession.deliverIrcMessage`, rendered from `irc-incoming.md`, persisted as an `irc:incoming` custom message) — receipt `injected`; + - `idle` (live session) → enqueued and a real turn is started — the message wakes the agent — receipt `woken`; + - `parked` → `AgentLifecycleManager.global().ensureLive(to)` revives the session first, then the wake path — receipt `revived`; + - resolution/revival failure → receipt `failed` with the error; other recipients still complete. +7. `send` with `await: true` then calls `IrcBus.wait(senderId, { from: to }, timeoutMs, signal)` and appends the reply (or a no-reply note suggesting `inbox`/`wait`) to the result. +8. `op: "wait"` blocks until a message for the caller (optionally filtered by `from`) arrives, consumes it, and returns it. Timeout returns a clean "no message" result, not an error. +9. `op: "inbox"` drains pending messages (or peeks with `peek: true`) without blocking. +10. Timeouts resolve as `params.timeoutMs ?? irc.timeoutMs`, normalized: `0` disables the timeout, negative/non-finite values fall back to the default `120_000`, positive values are truncated and clamped to ≥ 1 ms. ## Modes / Variants -- `list`: enumerate visible peers and logical channels. -- `send` direct message: one exact peer id, default synchronous auto-reply. -- `send` broadcast: `to: "all"`, default fire-and-forget (`awaitReply: false`) to every visible peer. -- `send` with `awaitReply: false`: recipient records the incoming message but does not generate a reply. -- `send` with `awaitReply: true`: recipient performs a no-tools ephemeral LLM turn and returns prose. +- `list`: enumerate peers with status (`running`/`idle`/`parked`), unread counts, and last activity. +- `send` direct: one exact peer id; wakes idle peers, revives parked ones. +- `send` broadcast: `to: "all"` to every live peer; parked peers are skipped. +- `send` + `await: true`: round-trip convenience — send, then wait for the next message from that peer. Replaces the old `awaitReply` auto-reply semantics without a fake reply. +- `wait`: block for an incoming message, optionally filtered by sender. +- `inbox`: non-blocking drain or peek. ## Side Effects - Session state - - Reads from the process-global `AgentRegistry`. - - Emits `irc_message` session events on recipient sessions. - - Queues IRC custom messages into recipient persisted history after the current stream finishes. - - For non-main recipients, forwards display-only relay observations into the main session UI; these relays are not persisted to the main agent history. - - Subagents inherit `irc.enabled` from task executor settings. + - Reads the process-global `AgentRegistry`; direct sends to parked agents revive their sessions through the lifecycle manager. + - Persists `irc:incoming` custom messages into recipient history; replies are ordinary turns in the recipient's own session. + - Waking an idle/parked recipient starts a real agent turn (model requests, tool use) in that recipient. - User-visible prompts / interactive UI - - IRC events render as `[IRC]` transcript lines in the TUI. - - Auto-replies are generated from `packages/coding-agent/src/prompts/system/irc-incoming.md` and explicitly forbid tool use. + - IRC events render as transcript cards in the TUI; the Agent Hub shows per-agent unread counts. - Background work / cancellation - - `send` starts one background `respondAsBackground` call per target. - - The caller's `AbortSignal` is forwarded into each background reply turn. `irc.timeoutMs` creates a per-recipient `AbortController` and reports timeout failures per target. + - `send` itself never blocks on reply generation; only `wait` (and `await: true`) blocks, bounded by the resolved timeout and the caller's `AbortSignal`. - Network - - No IRC server connection. - - When `awaitReply: true`, the recipient may make model-provider API calls through `runEphemeralTurn`. + - No IRC server connection. Woken recipients make their own model-provider calls as part of their turn. - Filesystem - - No direct filesystem writes in the tool itself. + - No direct filesystem writes in the tool itself; recipient turns persist to their session JSONL as usual. ## Limits & Caps -- Availability gates: - - `irc.enabled` defaults to `true` in `packages/coding-agent/src/config/settings-schema.ts`. - - Main agent tool discovery suppresses `irc` when `async.enabled` is off (`packages/coding-agent/src/tools/index.ts`). -- Visibility scope: only peers in status `running` or `idle` are addressable via `listVisibleTo`. -- Reply execution: - - No tools are available in auto-reply turns (`toolChoice: "none"` in `runEphemeralTurn`). - - `irc.timeoutMs` defaults to `120_000`; `0` disables the timeout, non-finite values fall back to the default, and positive values are truncated and clamped to at least `1` ms. - - No retry, backoff, rate limit, or reply length cap is defined in `irc.ts`; behavior otherwise relies on the underlying model stream and any upstream API limits. -- Flush scheduling: deferred history injection polls every `50` ms while the recipient is still streaming (`#scheduleBackgroundExchangeFlush` in `packages/coding-agent/src/session/agent-session.ts`). +- Availability gates: `irc.enabled` (default `true`), an `AgentRegistry`, and a caller agent id. +- Mailboxes are bounded at 100 messages per agent (`MAILBOX_CAP` in `packages/coding-agent/src/irc/bus.ts`); oldest messages are dropped beyond the cap. +- `irc.timeoutMs` defaults to `120_000` and is the default `wait` / `send await:true` timeout; `0` disables the timeout, non-finite or negative values fall back to the default, positive values are truncated and clamped to at least `1` ms. +- Broadcast scope: live peers only (`running`/`idle`) via `listVisibleTo`; direct sends address any non-aborted agent, including parked ones. ## Errors -- The tool returns text errors, not thrown exceptions, for: +- The tool returns text errors (with `isError: true`), not thrown exceptions, for: - missing registry: `IRC is unavailable in this session.` - missing sender id: `IRC is unavailable: caller has no agent id.` - - missing `to`: `` `to` is required for op="send". `` - - missing `message`: `` `message` is required for op="send". `` - - unknown op: `Unknown irc op.` -- Unknown, self-addressed, non-running, and non-idle direct targets are reported under `details.notFound` and in the text footer `Unknown / unavailable peers:`. -- If a target has no attached session, it is treated as not found. -- Exceptions thrown by `respondAsBackground`, `runEphemeralTurn`, abort handling, or timeout handling are caught per-target and surfaced under `details.failed` as `{ id, error }`; other recipients still complete. -- If no target succeeds, `send` still returns normally with `No recipients received the message.` and optional `failed`/`notFound` metadata. + - missing `to` / `message` on `send` + - self-send: `Cannot send an IRC message to yourself.` + - `await` with `to: "all"` + - unknown op +- Per-recipient delivery failures surface as `failed` receipts with the error message; `send` is marked `isError` only when no recipient received the message. +- `wait` timeout is a normal result (`waited: null`), not an error. ## Notes -- This is IRC-like naming only. There are no servers, sockets, nick registration, auth handshakes, channels beyond `all`, or commands such as join/part/topic. -- Addressing is by exact agent id from the registry; there is no fuzzy lookup or aliasing. -- `channels` in `list` is synthetic output: `all` plus visible peer ids. Nothing is persisted across calls as channel membership. -- Persistence is per recipient history, not per sender history. The sender gets the tool result; the recipient later sees injected custom messages on its next turn. -- The main UI may show IRC relays for conversations it was not part of, but those relay records are explicitly display-only. -- Because reply generation snapshots in-flight assistant text, a recipient can answer based on partially streamed context. -- Direct self-messaging is rejected by resolving the target as unavailable. \ No newline at end of file +- This is IRC-like naming only: no servers, sockets, channels, or join/part state. Addressing is by exact registry agent id. +- Replies are real turns by the recipient — the old ephemeral no-tools auto-reply (`awaitReply` / `respondAsBackground`) no longer exists. A recipient may keep working before answering; check `inbox` or `wait` again rather than re-sending. +- Wake-on-message is the revive primitive: messaging a parked agent is equivalent to resuming it (same `ensureLive` path as `task(resume:)` and the Agent Hub). +- Message ids are Snowflakes; pass them as `replyTo` to thread an answer to a specific message. +- Persistence is per recipient history: the sender gets receipts in the tool result; the recipient sees the injected `irc:incoming` message in its own transcript (visible via `history://`). diff --git a/packages/coding-agent/src/config/settings-schema.ts b/packages/coding-agent/src/config/settings-schema.ts index 66c691233..0e45a27b4 100644 --- a/packages/coding-agent/src/config/settings-schema.ts +++ b/packages/coding-agent/src/config/settings-schema.ts @@ -2316,8 +2316,7 @@ export const SETTINGS_SCHEMA = { ui: { tab: "tools", label: "IRC Timeout", - description: - "Drop IRC messages whose recipient does not respond within this many milliseconds (0 disables the timeout)", + description: "Default timeout for irc wait (and send await:true) in milliseconds; 0 disables the timeout", options: [ { value: "0", label: "Disabled" }, { value: "30000", label: "30 seconds" }, diff --git a/packages/coding-agent/src/irc/bus.ts b/packages/coding-agent/src/irc/bus.ts new file mode 100644 index 000000000..cb0186272 --- /dev/null +++ b/packages/coding-agent/src/irc/bus.ts @@ -0,0 +1,275 @@ +/** + * IrcBus - Process-global mailbox bus for agent-to-agent messaging. + * + * Replaces the old auto-reply model: a `send` never blocks on the recipient + * generating anything. Delivery resolves the recipient via the global + * AgentRegistry — parked agents are revived through the + * AgentLifecycleManager, idle agents are woken with a real turn, and busy + * agents receive the message as a non-interrupting aside at the next step + * boundary (see AgentSession.deliverIrcMessage). Replies are real turns by + * the recipient, observed via `wait`. + */ + +import { logger, Snowflake } from "@oh-my-pi/pi-utils"; +import { AgentLifecycleManager } from "../registry/agent-lifecycle"; +import { AgentRegistry, MAIN_AGENT_ID } from "../registry/agent-registry"; +import type { CustomMessage } from "../session/messages"; + +export interface IrcMessage { + id: string; + /** Sender agent id. */ + from: string; + /** Recipient agent id (resolved; "all" is expanded by the tool, not stored). */ + to: string; + body: string; + ts: number; + /** Message id being answered. */ + replyTo?: string; +} + +export interface IrcDeliveryReceipt { + to: string; + outcome: "injected" | "woken" | "revived" | "failed"; + error?: string; +} + +interface IrcWaiter { + from?: string; + resolve: (msg: IrcMessage) => void; + cancel: () => void; +} + +/** Mailbox cap per agent; oldest messages are dropped beyond it. */ +const MAILBOX_CAP = 100; + +export class IrcBus { + static #global: IrcBus | undefined; + + static global(): IrcBus { + if (!IrcBus.#global) { + IrcBus.#global = new IrcBus(); + } + return IrcBus.#global; + } + + /** Reset the global bus. Test-only. */ + static resetGlobalForTests(): void { + IrcBus.#global = undefined; + } + + readonly #registry: AgentRegistry; + readonly #lifecycle: () => AgentLifecycleManager; + readonly #mailboxes = new Map(); + readonly #waiters = new Map(); + + constructor(registry: AgentRegistry = AgentRegistry.global(), lifecycle?: AgentLifecycleManager) { + this.#registry = registry; + // Lazy: the lifecycle global self-constructs against the global registry, + // so only touch it when a parked recipient actually needs reviving. + this.#lifecycle = () => lifecycle ?? AgentLifecycleManager.global(); + } + + /** + * Fire-and-forget delivery. Never blocks on the recipient generating + * anything: the receipt reports how the message reached the recipient + * (waiter/aside = "injected", idle wake = "woken", park revival = + * "revived"), not what they did with it. + */ + async send(msg: Omit): Promise { + const message: IrcMessage = { ...msg, id: Snowflake.next(), ts: Date.now() }; + const ref = this.#registry.get(message.to); + if (!ref || ref.status === "aborted") { + return { to: message.to, outcome: "failed", error: `Unknown or terminated agent "${message.to}".` }; + } + + let revived = false; + if (ref.status === "parked") { + try { + await this.#lifecycle().ensureLive(message.to); + revived = true; + } catch (error) { + return { + to: message.to, + outcome: "failed", + error: error instanceof Error ? error.message : String(error), + }; + } + } + + // A pending `wait` from the recipient consumes the message directly — + // it is returned from their irc tool call and never hits the inbox or + // the session injection path. + const waiter = this.#takeMatchingWaiter(message.to, message.from); + if (waiter) { + waiter.resolve(message); + this.#relayToMainUi(message); + return { to: message.to, outcome: revived ? "revived" : "injected" }; + } + + const session = this.#registry.get(message.to)?.session; + if (!session) { + return { to: message.to, outcome: "failed", error: `Agent "${message.to}" has no live session.` }; + } + + this.#enqueue(message); + try { + const delivery = await session.deliverIrcMessage(message); + this.#relayToMainUi(message); + return { to: message.to, outcome: revived ? "revived" : delivery }; + } catch (error) { + return { + to: message.to, + outcome: "failed", + error: error instanceof Error ? error.message : String(error), + }; + } + } + + /** + * Block until a message for `agentId` (optionally from `filter.from`) + * arrives; consume + return it. Null on timeout (`timeoutMs <= 0` waits + * forever). Rejects when `signal` aborts. + */ + async wait( + agentId: string, + filter: { from?: string }, + timeoutMs: number, + signal?: AbortSignal, + ): Promise { + if (signal?.aborted) { + throw signal.reason instanceof Error ? signal.reason : new Error("IRC wait aborted"); + } + + // Already-pending mail satisfies the wait without parking a waiter. + const pending = this.#takeFromMailbox(agentId, filter.from); + if (pending) return pending; + + const { promise, resolve, reject } = Promise.withResolvers(); + let timer: NodeJS.Timeout | undefined; + let onAbort: (() => void) | undefined; + + const waiter: IrcWaiter = { + from: filter.from, + resolve: msg => { + cleanup(); + resolve(msg); + }, + cancel: () => { + cleanup(); + }, + }; + const cleanup = (): void => { + this.#removeWaiter(agentId, waiter); + clearTimeout(timer); + if (signal && onAbort) signal.removeEventListener("abort", onAbort); + }; + + if (signal) { + onAbort = () => { + cleanup(); + reject(signal.reason instanceof Error ? signal.reason : new Error("IRC wait aborted")); + }; + signal.addEventListener("abort", onAbort, { once: true }); + } + if (timeoutMs > 0) { + timer = setTimeout(() => { + cleanup(); + resolve(null); + }, timeoutMs); + timer.unref?.(); + } + + let waiters = this.#waiters.get(agentId); + if (!waiters) { + waiters = []; + this.#waiters.set(agentId, waiters); + } + waiters.push(waiter); + return promise; + } + + /** Drain (or peek) pending messages for `agentId`. */ + inbox(agentId: string, opts?: { peek?: boolean }): IrcMessage[] { + const mailbox = this.#mailboxes.get(agentId); + if (!mailbox || mailbox.length === 0) return []; + if (opts?.peek) return [...mailbox]; + this.#mailboxes.delete(agentId); + return mailbox; + } + + unreadCount(agentId: string): number { + return this.#mailboxes.get(agentId)?.length ?? 0; + } + + #enqueue(message: IrcMessage): void { + let mailbox = this.#mailboxes.get(message.to); + if (!mailbox) { + mailbox = []; + this.#mailboxes.set(message.to, mailbox); + } + mailbox.push(message); + if (mailbox.length > MAILBOX_CAP) { + const dropped = mailbox.shift(); + logger.debug("IrcBus: mailbox full, dropped oldest message", { + agentId: message.to, + droppedId: dropped?.id, + droppedFrom: dropped?.from, + }); + } + } + + /** Resolve the OLDEST waiter for `agentId` whose from-filter accepts `from`. */ + #takeMatchingWaiter(agentId: string, from: string): IrcWaiter | undefined { + const waiters = this.#waiters.get(agentId); + if (!waiters) return undefined; + const index = waiters.findIndex(waiter => !waiter.from || waiter.from === from); + if (index === -1) return undefined; + const [waiter] = waiters.splice(index, 1); + if (waiters.length === 0) this.#waiters.delete(agentId); + return waiter; + } + + #removeWaiter(agentId: string, waiter: IrcWaiter): void { + const waiters = this.#waiters.get(agentId); + if (!waiters) return; + const index = waiters.indexOf(waiter); + if (index !== -1) waiters.splice(index, 1); + if (waiters.length === 0) this.#waiters.delete(agentId); + } + + #takeFromMailbox(agentId: string, from?: string): IrcMessage | undefined { + const mailbox = this.#mailboxes.get(agentId); + if (!mailbox) return undefined; + const index = from ? mailbox.findIndex(msg => msg.from === from) : 0; + if (index === -1 || mailbox.length === 0) return undefined; + const [message] = mailbox.splice(index, 1); + if (mailbox.length === 0) this.#mailboxes.delete(agentId); + return message; + } + + /** + * Surface agent↔agent traffic as a display-only card on the main session + * UI. Skipped when the main agent is the recipient — its own + * `deliverIrcMessage` (or `wait` tool result) already shows the message. + */ + #relayToMainUi(message: IrcMessage): void { + if (message.to === MAIN_AGENT_ID) return; + const mainSession = this.#registry.get(MAIN_AGENT_ID)?.session; + if (!mainSession) return; + const record: CustomMessage = { + role: "custom", + customType: "irc:relay", + content: `[IRC \`${message.from}\` → \`${message.to}\`]\n\n${message.body}`, + display: true, + details: { from: message.from, to: message.to, body: message.body }, + attribution: "agent", + timestamp: message.ts, + }; + try { + mainSession.emitIrcRelayObservation(record); + } catch (error) { + // Display-only forwarding must never affect delivery semantics. + logger.debug("IrcBus: main UI relay failed", { to: message.to, error: String(error) }); + } + } +} diff --git a/packages/coding-agent/src/prompts/system/irc-incoming.md b/packages/coding-agent/src/prompts/system/irc-incoming.md index 7f5b8f139..6175b2216 100644 --- a/packages/coding-agent/src/prompts/system/irc-incoming.md +++ b/packages/coding-agent/src/prompts/system/irc-incoming.md @@ -1,8 +1,7 @@ -You received an IRC message from agent `{{from}}`. +Incoming IRC message from agent `{{from}}`{{#if replyTo}} (replying to {{replyTo}}){{/if}}: -Reply briefly and directly using the conversation context already available to you. NEVER call tools. The reply you write is delivered back to `{{from}}` as your answer. - -Message: {{message}} + +If a response is expected, reply with the `irc` tool (`op: "send"`, `to: "{{from}}"`) — you may finish your current step first. Nobody replies on your behalf. diff --git a/packages/coding-agent/src/prompts/tools/irc.md b/packages/coding-agent/src/prompts/tools/irc.md index edb10b560..0aee47c57 100644 --- a/packages/coding-agent/src/prompts/tools/irc.md +++ b/packages/coding-agent/src/prompts/tools/irc.md @@ -1,11 +1,15 @@ -Sends short text messages to other live agents in this process and receives their prose replies. +Sends short text messages to other agents in this process and receives theirs. - The main agent is addressable as `Main`. Subagents reuse their task id (e.g. `AuthLoader`, or `AuthLoader-2` when the name repeats). -- `op: "list"` returns the current set of visible peers. Use it before sending if you are not sure who is live. -- `op: "send"` delivers `message` to `to`. `to` may be a specific id or `"all"` to broadcast. -- Replies are generated on a side channel that does not wait for the recipient's main loop, so it is safe to IRC an agent that is mid tool call. -- The exchange (question + auto-reply) is injected into the recipient's history; they see it on their next turn and can follow up. +- `op: "list"` — every addressable peer with status (`running` | `idle` | `parked`), unread count, parent, and last activity. Use it before sending if you are not sure who exists. +- `op: "send"` — fire-and-forget delivery of `message` to `to` (a peer id, or `"all"` to broadcast to live peers). Returns per-recipient receipts immediately; it NEVER waits for the recipient to act. Receipt outcomes: `injected` (recipient was mid-turn; message folded in at their next step boundary), `woken` (idle recipient started a turn), `revived` (parked recipient was brought back and woken), `failed`. +- Messaging an `idle` or `parked` peer is how you wake it — there is no separate revive call. +- `send` with `await: true` — convenience round-trip: send, then block until the next message from that peer arrives (or the timeout passes). Invalid with `to: "all"`. +- `op: "wait"` — block until a message arrives (optionally only `from` a specific peer); consumes and returns it. A timeout is a clean "no message" result, not an error. +- `op: "inbox"` — drain pending messages without blocking (`peek: true` to leave them unread). +- `replyTo` — set it to the id of the message you are answering so the sender can correlate. +- Nobody answers on a peer's behalf anymore: a reply only arrives when the recipient actually sends one. For background on what a peer has been doing, `read` `history://` instead of interrogating them. @@ -21,29 +25,35 @@ NEVER use `irc` for: routine progress updates, things a tool call can verify, or These rules apply to both sending and replying. - **Plain prose only.** NEVER send structured JSON status payloads (e.g. `{"type":"task_completed",…}`). Write a normal sentence: "Done with the auth refactor — left a TODO in `src/server/auth.ts` for the rate limiter." -- **NEVER quote the message you are replying to.** Lead with the answer. -- **Use IRC, not terminal tools, to learn about peers.** NEVER `grep` artifacts, read other sessions' JSONL files, or shell-poke to figure out what another agent is doing. DM them. -- **One round-trip is enough.** Replies arrive synchronously when the recipient is reachable. NEVER follow up with "did you get my message?". If `delivered` is empty or the result was `failed`, the peer is unavailable — move on or report the blocker; NEVER retry in a loop. +- **NEVER quote the message you are replying to.** Lead with the answer; set `replyTo` instead. +- **Use IRC, not terminal tools, to learn about peers.** NEVER `grep` artifacts, read other sessions' JSONL files, or shell-poke to figure out what another agent is doing. DM them, or `read` `history://`. +- **Send, then keep working.** `send` returns immediately — only `wait` (or `await: true`) when you genuinely cannot proceed without the answer. NEVER follow up with "did you get my message?"; a `failed` receipt means the peer is unreachable — move on or report the blocker; NEVER retry in a loop. +- **Answer when a response is expected.** When an incoming message asks something, reply with `irc send` to the sender (you may finish your current step first). - **Stay terse.** A DM is a chat message, not a memo. One question per send. Share file paths and artifacts via `local://` / `memory://` / `artifact://` URLs instead of pasting blobs. - **Address peers by id.** Use the exact id from `op: "list"` (e.g. `AuthLoader`, `Main`). NEVER invent friendly names. - **NEVER IRC for things a tool would answer.** If a `read`, `grep`, or build command resolves the question, do that first. -- **Answer incoming IRC messages before continuing.** Address the question directly; do not repeat it back to the user. -- `send`: returns each recipient that received the message and any prose replies that arrived. -- `list`: returns peers and channels visible to the caller. +- `send`: per-recipient delivery receipts (`injected` / `woken` / `revived` / `failed`); with `await: true`, also the reply (or a timeout notice). +- `wait`: the consumed message, or a clean timeout notice. +- `inbox`: pending messages, oldest first. +- `list`: peers with status, unread count, parent, and last activity. # List peers `{"op": "list"}` -# Direct message to the main agent (waits for prose reply) -`{"op": "send", "to": "Main", "message": "Should I prefer JWT or session cookies for the auth flow?"}` -# Unexpected state — ask the originator -`{"op": "send", "to": "Main", "message": "Assignment says edit src/auth/jwt.ts but the file does not exist. Is the new path src/server/auth/jwt.ts?"}` -# Blocked by a peer — ask them directly -`{"op": "send", "to": "AuthLoader", "message": "Are you still touching src/server/auth.ts? I need to add a 401 path; OK to proceed or should I wait?"}` -# Broadcast to discover who owns something (no replies, just informs them) -`{"op": "send", "to": "all", "message": "About to refactor src/server/middleware/*. Anyone already in there?", "awaitReply": false}` +# Fire-and-forget DM — keep working, check inbox later +`{"op": "send", "to": "AuthLoader", "message": "Are you still touching src/server/auth.ts? I need to add a 401 path."}` +# Round-trip when you cannot proceed without the answer +`{"op": "send", "to": "Main", "message": "Should I prefer JWT or session cookies for the auth flow?", "await": true}` +# Wake a parked agent (same send — the bus revives it) +`{"op": "send", "to": "SchemaMigrator", "message": "The users table changed again; please re-check your migration."}` +# Block until a specific peer answers +`{"op": "wait", "from": "AuthLoader", "timeoutMs": 60000}` +# Drain pending messages +`{"op": "inbox"}` +# Broadcast to live peers (no replies expected) +`{"op": "send", "to": "all", "message": "About to refactor src/server/middleware/*. Anyone already in there?"}` diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index d677a9d72..def9a1629 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/packages/coding-agent/src/session/agent-session.ts @@ -29,6 +29,7 @@ import { type AgentState, type AgentTool, AppendOnlyContextManager, + type AsideMessage, resolveTelemetry, ThinkingLevel, } from "@oh-my-pi/pi-agent-core"; @@ -163,6 +164,7 @@ import { GoalRuntime } from "../goals/runtime"; import type { Goal, GoalModeState } from "../goals/state"; import type { HindsightSessionState } from "../hindsight/state"; import { type LocalProtocolOptions, resolveLocalUrlToPath } from "../internal-urls"; +import type { IrcMessage } from "../irc/bus"; import { resolveMemoryBackend } from "../memory-backend"; import { getMnemopiSessionState, type MnemopiSessionState, setMnemopiSessionState } from "../mnemopi/state"; import { containsOrchestrate, ORCHESTRATE_NOTICE } from "../modes/orchestrate"; @@ -184,7 +186,6 @@ import planModeToolDecisionReminderPrompt from "../prompts/system/plan-mode-tool }; import ttsrInterruptTemplate from "../prompts/system/ttsr-interrupt.md" with { type: "text" }; import ttsrToolReminderTemplate from "../prompts/system/ttsr-tool-reminder.md" with { type: "text" }; -import { type AgentRegistry, MAIN_AGENT_ID } from "../registry/agent-registry"; import { deobfuscateSessionContext, obfuscateProviderContext, @@ -413,8 +414,6 @@ export interface AgentSessionConfig { asyncJobManager?: AsyncJobManager; /** Agent identity (registry id like "Main" or "Alice") used for IRC routing. */ agentId?: string; - /** Shared agent registry (for forwarding IRC observations to the main session UI). */ - agentRegistry?: AgentRegistry; /** * Override the provider-facing session ID for all API requests from this session. * When absent, `sessionManager.getSessionId()` is used. Needed when benchmark or @@ -557,15 +556,15 @@ function formatRetryFallbackBaseSelector(selector: RetryFallbackSelector): strin return `${selector.provider}/${selector.id}`; } -const IRC_REPLY_MAX_BYTES = 4096; +const EPHEMERAL_REPLY_MAX_BYTES = 4096; /** - * Collapse degenerate IRC ephemeral replies before they hit the relay. + * Collapse degenerate ephemeral replies (/btw, /omfg side-channel turns). * Models occasionally loop on a single line (~16 reports of N-times-repeated * replies); compress runs longer than 3 down to one instance + `[…N×]`, then * cap at 4 KiB so a runaway reply can't flood the channel. */ -function dedupeIrcReply(text: string): string { +function dedupeEphemeralReply(text: string): string { if (!text) return text; const lines = text.split("\n"); const out: string[] = []; @@ -582,11 +581,11 @@ function dedupeIrcReply(text: string): string { i = j; } let result = out.join("\n"); - if (Buffer.byteLength(result, "utf8") > IRC_REPLY_MAX_BYTES) { + if (Buffer.byteLength(result, "utf8") > EPHEMERAL_REPLY_MAX_BYTES) { // Trim by characters until we're under the byte budget — handles multi-byte // glyphs at the boundary without splitting them. const suffix = "\n[…truncated]"; - const budget = IRC_REPLY_MAX_BYTES - Buffer.byteLength(suffix, "utf8"); + const budget = EPHEMERAL_REPLY_MAX_BYTES - Buffer.byteLength(suffix, "utf8"); while (Buffer.byteLength(result, "utf8") > budget) { result = result.slice(0, -1); } @@ -941,13 +940,11 @@ export class AgentSession { #activeEvalExecutions = new Set>(); #evalExecutionDisposing = false; - // Background-channel IRC exchanges queued while the recipient was streaming. - // Drained into history (via emitExternalEvent) once the recipient becomes idle. - #pendingBackgroundExchanges: CustomMessage[][] = []; - #scheduledBackgroundExchangeFlush = false; - // Agent identity + registry for IRC relay forwarding to the main session UI. + // Incoming IRC messages received while a turn was streaming; drained as + // non-interrupting asides at the next step boundary (see the aside provider). + #pendingIrcAsides: CustomMessage[] = []; + // Agent identity (registry id) used for IRC routing and job ownership. #agentId: string | undefined; - #agentRegistry: AgentRegistry | undefined; #providerSessionId: string | undefined; #freshProviderSessionId: string | undefined; #isDisposed = false; @@ -1204,7 +1201,13 @@ export class AgentSession { // Background-job completions / late diagnostics are pulled into the run at // each step boundary as non-interrupting asides (see Agent.getAsideMessages), // so they reach the model between requests without waiting for a yield. - this.agent.setAsideMessageProvider(() => this.yieldQueue.drainLazy()); + this.agent.setAsideMessageProvider(() => { + const pendingIrc = this.#pendingIrcAsides; + this.#pendingIrcAsides = []; + const thunks: AsideMessage[] = pendingIrc.map(record => () => record); + thunks.push(...this.yieldQueue.drainLazy()); + return thunks; + }); this.#convertToLlm = config.convertToLlm ?? convertToLlm; this.#rebuildSystemPrompt = config.rebuildSystemPrompt; this.#getMcpServerInstructions = config.getMcpServerInstructions; @@ -1235,7 +1238,6 @@ export class AgentSession { this.#ttsrManager = config.ttsrManager; this.#obfuscator = config.obfuscator; this.#agentId = config.agentId; - this.#agentRegistry = config.agentRegistry; this.#providerSessionId = config.providerSessionId; this.agent.setAssistantMessageEventInterceptor((message, assistantMessageEvent) => { const event: AgentEvent = { @@ -3090,16 +3092,29 @@ export class AgentSession { state.resetConversationTracking(); } + /** + * Synchronously mark the session as disposing so new work is rejected + * immediately: Python/eval starts throw, queued asides are dropped, and the + * aside provider is detached. Idempotent; `dispose()` runs it first. + * + * Wrappers that await other teardown before delegating to `dispose()` MUST + * call this before their first await — otherwise work started in that async + * gap slips past the disposal guards. + */ + beginDispose(): void { + this.#isDisposed = true; + this.#pendingIrcAsides = []; + this.yieldQueue.clear(); + this.agent.setAsideMessageProvider(undefined); + this.#evalExecutionDisposing = true; + } + /** * Remove all listeners, flush pending writes, and disconnect from agent. * Call this when completely done with the session. */ async dispose(): Promise { - this.#isDisposed = true; - this.#pendingBackgroundExchanges = []; - this.yieldQueue.clear(); - this.agent.setAsideMessageProvider(undefined); - this.#evalExecutionDisposing = true; + this.beginDispose(); try { if (this.#extensionRunner?.hasHandlers("session_shutdown")) { await this.#extensionRunner.emit({ type: "session_shutdown" }); @@ -4634,7 +4649,7 @@ export class AgentSession { // Flush any pending bash messages before the new prompt this.#flushPendingBashMessages(); this.#flushPendingPythonMessages(); - this.#flushPendingBackgroundExchanges(); + this.#flushPendingIrcAsides(); // Reset todo reminder count on new user prompt this.#todoReminderCount = 0; @@ -8926,118 +8941,56 @@ export class AgentSession { } // ========================================================================= - // Background-Channel IRC Exchanges + // IRC Delivery // ========================================================================= /** - * Generate an ephemeral reply to a background message (e.g. an IRC ping from - * another agent) using this session's current model + system prompt + history. + * Deliver an IRC message into this session (recipient side; called by the + * IrcBus). Emits the `irc_message` session event for UI cards and injects + * the rendered message into the model's context as an `irc:incoming` + * custom message: * - * The incoming message is queued for injection into the recipient's persisted - * history immediately so timeouts/abort still preserve delivery. The reply is - * computed via a side-channel `streamSimple` call (analogous to `/btw`) so it - * never blocks on the recipient's in-flight tool calls. When a reply is - * generated, it is queued separately. Injection happens immediately when the - * session is idle, otherwise it is deferred until streaming ends. + * - mid-turn → queued on the aside channel and folded in at the next step + * boundary (non-interrupting, like async-result deliveries) → "injected"; + * - idle → starts a real turn with the message so the recipient wakes + * → "woken". + * + * Never blocks on the recipient's turn: the wake turn is fire-and-forget. */ - async respondAsBackground(args: { - from: string; - message: string; - awaitReply?: boolean; - signal?: AbortSignal; - }): Promise<{ replyText: string | null }> { - const awaitReply = args.awaitReply !== false; - const incomingTimestamp = Date.now(); - const incomingRecord: CustomMessage = { + async deliverIrcMessage(msg: IrcMessage): Promise<"injected" | "woken"> { + if (this.#isDisposed) { + throw new Error("Recipient session is disposed."); + } + const record: CustomMessage = { role: "custom", customType: "irc:incoming", - content: `[IRC \`${args.from}\` → you]\n\n${args.message}`, + content: prompt.render(ircIncomingTemplate, { + from: msg.from, + message: msg.body, + replyTo: msg.replyTo ?? "", + }), display: true, - details: { from: args.from, message: args.message }, + details: { id: msg.id, from: msg.from, message: msg.body, ...(msg.replyTo ? { replyTo: msg.replyTo } : {}) }, attribution: "agent", - timestamp: incomingTimestamp, + timestamp: msg.ts, }; - void this.#emitSessionEvent({ type: "irc_message", message: incomingRecord }); - this.#forwardIrcRelayToMain({ - from: args.from, - to: this.#agentId ?? "?", - body: args.message, - kind: "message", - timestamp: incomingTimestamp, - }); - - this.#queueBackgroundExchangeInjection([incomingRecord]); - if (!awaitReply) { - return { replyText: null }; + void this.#emitSessionEvent({ type: "irc_message", message: record }); + if (this.isStreaming) { + this.#pendingIrcAsides.push(record); + return "injected"; } - - const incomingPrompt = prompt.render(ircIncomingTemplate, { - from: args.from, - message: args.message, + // Idle: same wake primitive the yield queue uses for async-result + // delivery — prompt the agent directly so a real turn runs. + this.agent.prompt(record).catch(error => { + logger.warn("IRC wake turn failed", { from: msg.from, to: msg.to, error: String(error) }); }); - const { replyText } = await this.runEphemeralTurn({ - promptText: incomingPrompt, - signal: args.signal, - }); - - const replyRecord: CustomMessage = { - role: "custom", - customType: "irc:autoreply", - content: `[IRC you → \`${args.from}\` (auto)]\n\n${replyText}`, - display: true, - details: { to: args.from, reply: replyText }, - attribution: "agent", - timestamp: Date.now(), - }; - void this.#emitSessionEvent({ type: "irc_message", message: replyRecord }); - this.#forwardIrcRelayToMain({ - from: this.#agentId ?? "?", - to: args.from, - body: replyText, - kind: "reply", - timestamp: replyRecord.timestamp, - }); - this.#queueBackgroundExchangeInjection([replyRecord]); - - return { replyText }; - } - - /** - * Forward an IRC exchange observation to the main agent's session UI so the - * user can see every IRC conversation in the main transcript, even when the - * main agent is not a direct participant. The relay record is display-only: - * it is NOT injected into the main agent's persisted history. - */ - #forwardIrcRelayToMain(args: { - from: string; - to: string; - body: string; - kind: "message" | "reply"; - timestamp: number; - }): void { - const registry = this.#agentRegistry; - if (!registry) return; - // If this session is the main agent, the local emit already reached the main UI. - if (this.#agentId === MAIN_AGENT_ID) return; - const mainRef = registry.get(MAIN_AGENT_ID); - const mainSession = mainRef?.session; - if (!mainSession || mainSession === this) return; - const arrow = args.kind === "reply" ? "→ (auto)" : "→"; - const relayRecord: CustomMessage = { - role: "custom", - customType: "irc:relay", - content: `[IRC \`${args.from}\` ${arrow} \`${args.to}\`]\n\n${args.body}`, - display: true, - details: { from: args.from, to: args.to, body: args.body, kind: args.kind }, - attribution: "agent", - timestamp: args.timestamp, - }; - mainSession.emitIrcRelayObservation(relayRecord); + return "woken"; } /** * Emit an IRC relay observation event on this session for UI rendering only. - * Does not persist the record to history. Public so other sessions can forward. + * Does not persist the record to history. Called by the IrcBus to surface + * agent↔agent traffic on the main session. */ emitIrcRelayObservation(record: CustomMessage): void { void this.#emitSessionEvent({ type: "irc_message", message: record }); @@ -9049,7 +9002,7 @@ export class AgentSession { * does not block on, or interfere with, any in-flight main turn. The * session's history and persisted state are NOT modified by this call. * - * Used by `respondAsBackground` (IRC) and `BtwController` (`/btw`) to share + * Used by `BtwController` (`/btw`) and `OmfgController` (`/omfg`) to share * the snapshot + stream pipeline. The snapshot includes any in-flight * streaming assistant text so the model sees the half-finished response * rather than missing context. @@ -9137,7 +9090,7 @@ export class AgentSession { args.onTextDelta(replyText.slice(emittedReplyText.length)); } return { - replyText: args.dedupeReply === false ? replyText.trim() : dedupeIrcReply(replyText.trim()), + replyText: args.dedupeReply === false ? replyText.trim() : dedupeEphemeralReply(replyText.trim()), assistantMessage, }; } @@ -9188,46 +9141,21 @@ export class AgentSession { return messages; } - #queueBackgroundExchangeInjection(messages: CustomMessage[]): void { - this.#pendingBackgroundExchanges.push(messages); - if (!this.isStreaming) { - this.#flushPendingBackgroundExchanges(); - return; - } - this.#scheduleBackgroundExchangeFlush(); - } - - #scheduleBackgroundExchangeFlush(): void { - if (this.#scheduledBackgroundExchangeFlush) return; - this.#scheduledBackgroundExchangeFlush = true; - const attempt = (): void => { - if (this.#pendingBackgroundExchanges.length === 0 || this.#isDisposed) { - this.#pendingBackgroundExchanges = []; - this.#scheduledBackgroundExchangeFlush = false; - return; - } - if (this.isStreaming) { - setTimeout(attempt, 50); - return; - } - this.#scheduledBackgroundExchangeFlush = false; - this.#flushPendingBackgroundExchanges(); - }; - setTimeout(attempt, 0); - } - - #flushPendingBackgroundExchanges(): void { - if (this.#pendingBackgroundExchanges.length === 0) return; - const batches = this.#pendingBackgroundExchanges; - this.#pendingBackgroundExchanges = []; - for (const batch of batches) { - for (const msg of batch) { - // emitExternalEvent on message_end appends to agent state and dispatches - // to all session listeners, which in turn handle TUI rendering and - // sessionManager persistence via #handleAgentEvent. - this.agent.emitExternalEvent({ type: "message_start", message: msg }); - this.agent.emitExternalEvent({ type: "message_end", message: msg }); - } + /** + * Persist any IRC asides that missed their step-boundary injection (the + * message landed after the turn's last aside drain). Called at the start + * of the next prompt so the model still sees them. + */ + #flushPendingIrcAsides(): void { + if (this.#pendingIrcAsides.length === 0) return; + const records = this.#pendingIrcAsides; + this.#pendingIrcAsides = []; + for (const record of records) { + // emitExternalEvent on message_end appends to agent state and dispatches + // to all session listeners, which in turn handle TUI rendering and + // sessionManager persistence via #handleAgentEvent. + this.agent.emitExternalEvent({ type: "message_start", message: record }); + this.agent.emitExternalEvent({ type: "message_end", message: record }); } } diff --git a/packages/coding-agent/src/tools/index.ts b/packages/coding-agent/src/tools/index.ts index 750af3f65..c524ca646 100644 --- a/packages/coding-agent/src/tools/index.ts +++ b/packages/coding-agent/src/tools/index.ts @@ -18,7 +18,7 @@ import { LspTool } from "../lsp"; import type { MCPManager } from "../mcp"; import type { MnemopiSessionState } from "../mnemopi/state"; import type { PlanModeState } from "../plan-mode/state"; -import { type AgentRegistry, MAIN_AGENT_ID } from "../registry/agent-registry"; +import type { AgentRegistry } from "../registry/agent-registry"; import type { ArtifactManager } from "../session/artifacts"; import type { ClientBridge } from "../session/client-bridge"; import type { CustomMessage } from "../session/messages"; @@ -539,13 +539,7 @@ export async function createTools(session: ToolSession, toolNames?: string[]): P if (name === "search_tool_bm25") return discoveryActive; if (name === "browser") return session.settings.get("browser.enabled"); if (name === "checkpoint" || name === "rewind") return session.settings.get("checkpoint.enabled"); - if (name === "irc") { - if (!session.settings.get("irc.enabled")) return false; - // Main agent only needs `irc` when subagents may run concurrently (async). - // In sync mode main blocks on `task`, so peer messaging from main is dead weight. - if (!session.settings.get("async.enabled") && session.getAgentId?.() === MAIN_AGENT_ID) return false; - return true; - } + if (name === "irc") return session.settings.get("irc.enabled"); if (name === "retain" || name === "recall" || name === "reflect") { return ["hindsight", "mnemopi"].includes(session.settings.get("memory.backend") ?? ""); } diff --git a/packages/coding-agent/src/tools/irc.ts b/packages/coding-agent/src/tools/irc.ts index a075b3404..49d594aa7 100644 --- a/packages/coding-agent/src/tools/irc.ts +++ b/packages/coding-agent/src/tools/irc.ts @@ -1,62 +1,80 @@ /** - * IRC tool — agent-to-agent messaging. + * IRC tool — agent-to-agent messaging over the process-global IrcBus. * - * Lets any live agent send a short prose message to any other live agent in - * this process and (optionally) get a prose reply. - * - * Routing happens via the global AgentRegistry. Replies are produced by an - * ephemeral side-channel call (`AgentSession.respondAsBackground`) that - * mirrors `/btw`: the recipient's current model, system prompt, and message - * history are used to compute a reply without persisting it through the - * normal stream path. After the reply is generated, both the incoming - * message and the auto-reply are queued for injection into the recipient's - * persisted history (deferred until the recipient is idle), so the model - * sees the exchange on its next turn. - * - * This avoids the deadlock that arises when the recipient is blocked on a - * long-running tool call: the side-channel call does not depend on the - * recipient's main agent loop being free. + * `send` is fire-and-forget: the bus routes the message to the recipient + * (waking idle agents with a real turn, reviving parked ones via the + * lifecycle manager, injecting a non-interrupting aside into busy ones) and + * returns delivery receipts immediately. Replies are real turns by the + * recipient, observed with `wait` (or the `await: true` send sugar). `inbox` + * drains pending messages; `list` shows every addressable peer. */ import type { AgentTool, AgentToolContext, AgentToolResult, AgentToolUpdateCallback } from "@oh-my-pi/pi-agent-core"; -import { prompt } from "@oh-my-pi/pi-utils"; +import { type Component, Text } from "@oh-my-pi/pi-tui"; +import { formatAge, formatDuration, prompt } from "@oh-my-pi/pi-utils"; import * as z from "zod/v4"; +import type { RenderResultOptions } from "../extensibility/custom-tools/types"; +import { IrcBus, type IrcDeliveryReceipt, type IrcMessage } from "../irc/bus"; +import type { Theme } from "../modes/theme/theme"; import ircDescription from "../prompts/tools/irc.md" with { type: "text" }; -import type { AgentRef, AgentRegistry } from "../registry/agent-registry"; +import type { AgentRegistry } from "../registry/agent-registry"; +import { Ellipsis, renderStatusLine, renderTreeList, truncateToWidth } from "../tui"; import type { ToolSession } from "."; +import { + createCachedComponent, + formatBadge, + formatErrorDetail, + getPreviewLines, + PREVIEW_LIMITS, + replaceTabs, + type ToolUIColor, +} from "./render-utils"; const DEFAULT_IRC_TIMEOUT_MS = 120_000; const ircSchema = z.object({ - op: z.enum(["send", "list"]).describe("irc operation"), - to: z.string().optional().describe('recipient agent id or "all"'), - message: z.string().optional().describe("message body"), - awaitReply: z.boolean().optional().describe("wait for prose reply"), + op: z.enum(["send", "wait", "inbox", "list"]).describe("irc operation"), + to: z.string().optional().describe('send: recipient agent id or "all"'), + message: z.string().optional().describe("send: message body"), + replyTo: z.string().optional().describe("send: message id being answered"), + await: z.boolean().optional().describe('send: wait for the recipient\'s reply (invalid with to:"all")'), + from: z.string().optional().describe("wait: only accept a message from this agent id"), + timeoutMs: z.number().optional().describe("wait: timeout in milliseconds (0 waits indefinitely)"), + peek: z.boolean().optional().describe("inbox: list messages without consuming them"), }); type IrcParams = z.infer; -interface IrcReply { - from: string; - text: string; +interface IrcPeerInfo { + id: string; + displayName: string; + kind: string; + status: string; + parentId?: string; + unread: number; + lastActivity: number; } export interface IrcDetails { - op: "send" | "list"; + op: "send" | "wait" | "inbox" | "list"; from?: string; to?: string; - delivered?: string[]; - replies?: IrcReply[]; - failed?: Array<{ id: string; error: string }>; - notFound?: string[]; - peers?: Array<{ id: string; displayName: string; kind: string; status: string; parentId?: string }>; - channels?: string[]; + receipts?: IrcDeliveryReceipt[]; + /** Message consumed by `wait` / `send await:true`; null when the wait timed out. */ + waited?: IrcMessage | null; + inbox?: IrcMessage[]; + peers?: IrcPeerInfo[]; +} + +function formatIncoming(msg: IrcMessage): string { + const replyTag = msg.replyTo ? ` (reply to ${msg.replyTo})` : ""; + return `[${msg.id}] ${msg.from}${replyTag}: ${msg.body}`; } export class IrcTool implements AgentTool { readonly name = "irc"; readonly approval = "read" as const; readonly label = "IRC"; - readonly summary = "Send and receive messages between agents over IRC-like channels"; + readonly summary = "Send and receive messages between agents"; readonly description: string; readonly parameters = ircSchema; readonly strict = true; @@ -87,41 +105,55 @@ export class IrcTool implements AgentTool { return errorResult("IRC is unavailable: caller has no agent id.", { op: params.op }); } - if (params.op === "list") { - return this.#executeList(registry, senderId); + switch (params.op) { + case "list": + return this.#executeList(registry, senderId); + case "send": + return this.#executeSend(registry, senderId, params, signal); + case "wait": + return this.#executeWait(senderId, params, signal); + case "inbox": + return this.#executeInbox(senderId, params); + default: + return errorResult("Unknown irc op.", { op: params.op }); } - if (params.op === "send") { - return this.#executeSend(registry, senderId, params, signal); - } - return errorResult("Unknown irc op.", { op: params.op as "send" | "list" }); } #executeList(registry: AgentRegistry, senderId: string): AgentToolResult { - const peers = registry.listVisibleTo(senderId); + const bus = IrcBus.global(); + const peers = registry + .list() + .filter(ref => ref.id !== senderId && ref.status !== "aborted") + .map(ref => ({ + id: ref.id, + displayName: ref.displayName, + kind: ref.kind, + status: ref.status, + parentId: ref.parentId, + unread: bus.unreadCount(ref.id), + lastActivity: ref.lastActivity, + })); const lines: string[] = []; if (peers.length === 0) { - lines.push("No other live agents."); + lines.push("No other agents."); } else { lines.push(`${peers.length} peer(s):`); for (const peer of peers) { - lines.push(`- ${peer.id} [${peer.displayName} · ${peer.kind} · ${peer.status}]`); + const extras = [ + peer.unread > 0 ? `unread ${peer.unread}` : undefined, + peer.parentId ? `parent ${peer.parentId}` : undefined, + `active ${formatDuration(Date.now() - peer.lastActivity)} ago`, + ].filter(Boolean); + lines.push(`- ${peer.id} [${peer.displayName} · ${peer.kind} · ${peer.status}] — ${extras.join(", ")}`); + } + if (peers.some(peer => peer.status === "parked")) { + lines.push(""); + lines.push("Parked agents are revived automatically when you message them."); } } - const channels = ["all", ...peers.map(p => p.id)]; return { content: [{ type: "text", text: lines.join("\n") }], - details: { - op: "list", - from: senderId, - peers: peers.map(p => ({ - id: p.id, - displayName: p.displayName, - kind: p.kind, - status: p.status, - parentId: p.parentId, - })), - channels, - }, + details: { op: "list", from: senderId, peers }, }; } @@ -139,91 +171,59 @@ export class IrcTool implements AgentTool { if (!message) { return errorResult('`message` is required for op="send".', { op: "send", from: senderId }); } - - // Resolve target peers. - let targets: AgentRef[]; - const notFound: string[] = []; + if (to === senderId) { + return errorResult("Cannot send an IRC message to yourself.", { op: "send", from: senderId, to }); + } const isBroadcast = to === "all"; - if (isBroadcast) { - targets = registry.listVisibleTo(senderId); - } else { - const ref = registry.get(to); - if (!ref || ref.id === senderId) { - notFound.push(to); - targets = []; - } else if (ref.status !== "running" && ref.status !== "idle") { - notFound.push(to); - targets = []; - } else { - targets = [ref]; - } + if (isBroadcast && params.await) { + return errorResult('`await` is invalid with to:"all" — broadcasts have no single replier.', { + op: "send", + from: senderId, + to, + }); } - const awaitReply = params.awaitReply ?? !isBroadcast; - - const timeoutMs = normalizeIrcTimeoutMs(this.session.settings.get("irc.timeoutMs")); - const delivered: string[] = []; - const replies: IrcReply[] = []; - const failed: Array<{ id: string; error: string }> = []; - - // Dispatch to each target in parallel via the recipient's ephemeral - // side-channel. Independent calls so a slow recipient cannot stall the - // others. The recipient's main loop never has to be unblocked: the - // side-channel runs alongside any in-flight tool call. - const dispatches = targets.map(async target => { - const targetSession = target.session; - if (!targetSession) { - notFound.push(target.id); - return; - } - try { - const result = await runIrcDispatchWithTimeout( - timeoutMs, - signal, - timeoutSignal => - targetSession.respondAsBackground({ - from: senderId, - message, - awaitReply, - signal: timeoutSignal, - }), - target.id, - ); - delivered.push(target.id); - if (awaitReply && result.replyText) { - replies.push({ from: target.id, text: result.replyText }); - } - } catch (err) { - failed.push({ id: target.id, error: err instanceof Error ? err.message : String(err) }); - } - }); - await Promise.all(dispatches); + const bus = IrcBus.global(); + // Broadcasts fan out to live peers only (running | idle); reviving every + // parked agent on a broadcast would be a stampede. Direct sends go + // through the bus unfiltered so parked recipients are revived. + const targets = isBroadcast ? registry.listVisibleTo(senderId).map(ref => ref.id) : [to]; + const receipts = await Promise.all( + targets.map(target => bus.send({ from: senderId, to: target, body: message, replyTo: params.replyTo })), + ); const lines: string[] = []; - if (delivered.length === 0) { + const delivered = receipts.filter(receipt => receipt.outcome !== "failed"); + if (targets.length === 0) { + lines.push("No live peers to broadcast to."); + } else if (delivered.length === 0) { lines.push("No recipients received the message."); } else { - lines.push(`Delivered to ${delivered.length} peer(s): ${delivered.join(", ")}`); + lines.push(`Delivered to ${delivered.length} peer(s):`); } - if (replies.length > 0) { + for (const receipt of receipts) { + lines.push( + receipt.outcome === "failed" + ? `- ${receipt.to}: failed — ${receipt.error ?? "unknown error"}` + : `- ${receipt.to}: ${receipt.outcome}`, + ); + } + + let waited: IrcMessage | null | undefined; + if (params.await && delivered.length > 0) { + const timeoutMs = this.#resolveTimeoutMs(params); + waited = await bus.wait(senderId, { from: to }, timeoutMs, signal); lines.push(""); - lines.push("## Replies"); - for (const reply of replies) { - lines.push(`### ${reply.from}`); - lines.push(reply.text); + if (waited) { + lines.push(`Reply from ${waited.from}:`); + lines.push(waited.body); + } else { + lines.push( + `No reply from ${to} within ${formatDuration(timeoutMs)}. ` + + "They may answer later — check `inbox` or `wait` again.", + ); } } - if (failed.length > 0) { - lines.push(""); - lines.push("## Failed"); - for (const f of failed) { - lines.push(`- ${f.id}: ${f.error}`); - } - } - if (notFound.length > 0) { - lines.push(""); - lines.push(`Unknown / unavailable peers: ${notFound.join(", ")}`); - } return { content: [{ type: "text", text: lines.join("\n") }], @@ -231,13 +231,52 @@ export class IrcTool implements AgentTool { op: "send", from: senderId, to, - delivered, - ...(replies.length > 0 ? { replies } : {}), - ...(failed.length > 0 ? { failed } : {}), - ...(notFound.length > 0 ? { notFound } : {}), + receipts, + ...(waited !== undefined ? { waited } : {}), }, + isError: delivered.length === 0 && targets.length > 0, }; } + + async #executeWait(senderId: string, params: IrcParams, signal?: AbortSignal): Promise> { + const from = params.from?.trim() || undefined; + const timeoutMs = this.#resolveTimeoutMs(params); + const waited = await IrcBus.global().wait(senderId, { from }, timeoutMs, signal); + if (!waited) { + const filterNote = from ? ` from ${from}` : ""; + return { + content: [{ type: "text", text: `No message${filterNote} within ${formatDuration(timeoutMs)}.` }], + details: { op: "wait", from: senderId, waited: null }, + }; + } + return { + content: [{ type: "text", text: formatIncoming(waited) }], + details: { op: "wait", from: senderId, waited }, + }; + } + + #executeInbox(senderId: string, params: IrcParams): AgentToolResult { + const messages = IrcBus.global().inbox(senderId, { peek: params.peek }); + if (messages.length === 0) { + return { + content: [{ type: "text", text: "Inbox empty." }], + details: { op: "inbox", from: senderId, inbox: [] }, + }; + } + const header = params.peek ? `${messages.length} unread message(s):` : `${messages.length} message(s):`; + const lines = [header, ...messages.map(msg => `- ${formatIncoming(msg)}`)]; + return { + content: [{ type: "text", text: lines.join("\n") }], + details: { op: "inbox", from: senderId, inbox: messages }, + }; + } + + #resolveTimeoutMs(params: IrcParams): number { + if (params.timeoutMs !== undefined) { + return normalizeIrcTimeoutMs(params.timeoutMs); + } + return normalizeIrcTimeoutMs(this.session.settings.get("irc.timeoutMs")); + } } function errorResult(text: string, details: IrcDetails): AgentToolResult { @@ -256,43 +295,333 @@ function normalizeIrcTimeoutMs(value: number): number { return Math.max(1, Math.trunc(value)); } -async function runIrcDispatchWithTimeout( - timeoutMs: number, - parentSignal: AbortSignal | undefined, - run: (signal?: AbortSignal) => Promise, - targetId: string, -): Promise { - if (timeoutMs <= 0) { - return await run(parentSignal); - } +// ============================================================================= +// TUI Renderer +// ============================================================================= - const controller = new AbortController(); - const timeoutError = new Error(`IRC timed out waiting for ${targetId} after ${timeoutMs} ms`); - let timeout: NodeJS.Timeout | undefined; - let parentAbortListener: (() => void) | undefined; +type IrcRenderArgs = Partial; - const timeoutDeferred = Promise.withResolvers(); - if (parentSignal) { - if (parentSignal.aborted) { - throw parentSignal.reason instanceof Error ? parentSignal.reason : new Error("IRC aborted"); - } - parentAbortListener = () => { - controller.abort(parentSignal.reason); - timeoutDeferred.reject(parentSignal.reason instanceof Error ? parentSignal.reason : new Error("IRC aborted")); - }; - parentSignal.addEventListener("abort", parentAbortListener, { once: true }); - } +const BODY_LINES_COLLAPSED = 2; +const BODY_LINES_EXPANDED = 12; +const BODY_LINE_WIDTH = 100; - timeout = setTimeout(() => { - controller.abort(timeoutError); - timeoutDeferred.reject(timeoutError); - }, timeoutMs); - timeout.unref?.(); +const PEER_STATUS_ORDER: Record = { running: 0, idle: 1, parked: 2 }; - try { - return await Promise.race([run(controller.signal), timeoutDeferred.promise]); - } finally { - if (timeout) clearTimeout(timeout); - if (parentSignal && parentAbortListener) parentSignal.removeEventListener("abort", parentAbortListener); +function ircGlyph(theme: Theme): string { + return theme.styledSymbol("tool.irc", "accent"); +} + +function outcomeColor(outcome: IrcDeliveryReceipt["outcome"]): ToolUIColor { + switch (outcome) { + case "woken": + return "success"; + case "revived": + return "warning"; + case "injected": + return "accent"; + case "failed": + return "error"; } } + +/** Glyph + status word, matching the agent-hub status conventions. */ +function peerStatusBadge(status: string, theme: Theme): string { + switch (status) { + case "running": + return theme.fg("accent", `${theme.status.running} running`); + case "idle": + return theme.fg("success", `${theme.status.enabled} idle`); + case "parked": + return theme.fg("muted", `${theme.status.shadowed} parked`); + default: + return theme.fg("error", `${theme.status.aborted} ${status}`); + } +} + +function messageAge(ts: number | undefined): string { + if (!ts) return ""; + return formatAge(Math.max(1, Math.round((Date.now() - ts) / 1000))); +} + +function textContent(result: { content: Array<{ type: string; text?: string }> }): string { + return result.content.find(part => part.type === "text")?.text?.trim() ?? ""; +} + +/** + * Quote-bordered message body preview. `tone` separates outbound text (dim) + * from received text (toolOutput); a trailing dim counter marks elided lines. + */ +function bodyLines( + body: string, + expanded: boolean, + theme: Theme, + options: { indent?: string; tone?: "dim" | "toolOutput"; collapsedLines?: number } = {}, +): string[] { + const indent = options.indent ?? ""; + const tone = options.tone ?? "toolOutput"; + const max = expanded ? BODY_LINES_EXPANDED : (options.collapsedLines ?? BODY_LINES_COLLAPSED); + const total = body.split("\n").filter(line => line.trim()).length; + const quote = theme.fg("dim", theme.md.quoteBorder); + const lines = getPreviewLines(body, max, BODY_LINE_WIDTH, Ellipsis.Unicode).map( + line => `${indent}${quote} ${theme.fg(tone, replaceTabs(line))}`, + ); + const hidden = total - Math.min(total, max); + if (hidden > 0) { + lines.push(`${indent}${quote} ${theme.fg("dim", `… +${hidden} more ${hidden === 1 ? "line" : "lines"}`)}`); + } + return lines; +} + +/** Header title carrying the op direction: `IRC ➤ peer` out, `IRC ⟵ peer` in. */ +function callTitle(args: IrcRenderArgs | undefined, theme: Theme): string { + switch (args?.op) { + case "send": + return `IRC ${theme.nav.selected} ${args.to?.trim() || "…"}`; + case "wait": + return `IRC ${theme.nav.back} ${args.from?.trim() || "anyone"}`; + case "inbox": + return "IRC inbox"; + case "list": + return "IRC peers"; + default: + return "IRC"; + } +} + +function callMeta(args: IrcRenderArgs | undefined): string[] { + const meta: string[] = []; + if (args?.op === "send") { + if (args.to === "all") meta.push("broadcast"); + if (args.await) meta.push("await reply"); + if (args.replyTo) meta.push("reply"); + } + if (args?.op === "wait" && args.timeoutMs) meta.push(`timeout ${formatDuration(args.timeoutMs)}`); + if (args?.op === "inbox" && args.peek) meta.push("peek"); + return meta; +} + +function renderSendResult( + result: { content: Array<{ type: string; text?: string }>; isError?: boolean }, + details: Partial, + args: IrcRenderArgs | undefined, + expanded: boolean, + theme: Theme, +): string[] { + const receipts = details.receipts ?? []; + const to = details.to ?? args?.to?.trim() ?? "?"; + const title = `IRC ${theme.nav.selected} ${to}`; + + // Pre-delivery failures (validation) and empty broadcasts carry no receipts. + if (receipts.length === 0) { + const text = textContent(result) || (result.isError ? "Send failed." : "Nothing to deliver."); + return [ + renderStatusLine({ icon: result.isError ? "error" : "warning", title }, theme), + result.isError ? formatErrorDetail(text, theme) : ` ${theme.fg("muted", replaceTabs(text))}`, + ]; + } + + const delivered = receipts.filter(receipt => receipt.outcome !== "failed"); + const failedCount = receipts.length - delivered.length; + const waited = details.waited; + const timedOut = waited === null; + + const meta: string[] = []; + if (to === "all") meta.push("broadcast"); + if (receipts.length === 1) { + const receipt = receipts[0]!; + meta.push(theme.fg(outcomeColor(receipt.outcome), receipt.outcome)); + } else { + if (delivered.length > 0) meta.push(theme.fg("success", `${delivered.length} delivered`)); + if (failedCount > 0) meta.push(theme.fg("error", `${failedCount} failed`)); + } + if (timedOut) meta.push(theme.fg("warning", "no reply")); + + const icon = result.isError + ? { icon: "error" as const } + : timedOut + ? { icon: "warning" as const } + : { iconOverride: ircGlyph(theme) }; + const lines = [renderStatusLine({ ...icon, title, meta }, theme)]; + + const sent = args?.message?.trim(); + if (sent) lines.push(...bodyLines(sent, expanded, theme, { indent: " ", tone: "dim" })); + + if (receipts.length > 1 || failedCount > 0) { + lines.push( + ...renderTreeList( + { + items: receipts, + expanded, + maxCollapsed: PREVIEW_LIMITS.COLLAPSED_ITEMS, + itemType: "recipient", + renderItem: receipt => { + const badge = formatBadge(receipt.outcome, outcomeColor(receipt.outcome), theme); + const error = + receipt.outcome === "failed" && receipt.error + ? ` ${theme.fg("error", `${theme.format.dash} ${receipt.error}`)}` + : ""; + return `${theme.fg("toolOutput", receipt.to)} ${badge}${error}`; + }, + }, + theme, + ), + ); + } + + if (waited) { + const age = messageAge(waited.ts); + lines.push( + ` ${theme.fg("dim", theme.nav.back)} ${theme.fg("accent", waited.from)}${age ? ` ${theme.fg("dim", age)}` : ""}`, + ); + lines.push(...bodyLines(waited.body, expanded, theme, { indent: " " })); + } else if (timedOut) { + lines.push(` ${theme.fg("warning", "No reply yet — they may answer later; check inbox or wait again.")}`); + } + return lines; +} + +function renderWaitResult( + result: { content: Array<{ type: string; text?: string }>; isError?: boolean }, + details: Partial, + args: IrcRenderArgs | undefined, + expanded: boolean, + theme: Theme, +): string[] { + const waited = details.waited; + if (!waited) { + const text = textContent(result) || "No message arrived."; + return [ + renderStatusLine( + { icon: "warning", title: `IRC ${theme.nav.back} ${args?.from?.trim() || "anyone"}`, meta: ["timed out"] }, + theme, + ), + ` ${theme.fg("muted", replaceTabs(text))}`, + ]; + } + const meta = [messageAge(waited.ts)]; + if (waited.replyTo) meta.push("reply"); + return [ + renderStatusLine({ iconOverride: ircGlyph(theme), title: `IRC ${theme.nav.back} ${waited.from}`, meta }, theme), + ...bodyLines(waited.body, expanded, theme, { indent: " " }), + ]; +} + +function renderInboxResult( + details: Partial, + args: IrcRenderArgs | undefined, + expanded: boolean, + theme: Theme, +): string[] { + const messages = details.inbox ?? []; + if (messages.length === 0) { + return [renderStatusLine({ iconOverride: ircGlyph(theme), title: "IRC inbox", meta: ["empty"] }, theme)]; + } + const meta = [`${messages.length} ${messages.length === 1 ? "message" : "messages"}`]; + if (args?.peek) meta.push("peek"); + const header = renderStatusLine({ iconOverride: ircGlyph(theme), title: "IRC inbox", meta }, theme); + const items = renderTreeList( + { + items: messages, + expanded, + maxCollapsed: PREVIEW_LIMITS.COLLAPSED_ITEMS, + itemType: "message", + renderItem: msg => { + const age = messageAge(msg.ts); + const replyBadge = msg.replyTo ? ` ${formatBadge("reply", "muted", theme)}` : ""; + const head = `${theme.fg("accent", msg.from)}${age ? ` ${theme.fg("dim", age)}` : ""}${replyBadge}`; + return [head, ...bodyLines(msg.body, expanded, theme, { collapsedLines: 1 })]; + }, + }, + theme, + ); + return [header, ...items]; +} + +function renderListResult(details: Partial, expanded: boolean, theme: Theme): string[] { + const peers = [...(details.peers ?? [])].sort( + (a, b) => + (PEER_STATUS_ORDER[a.status] ?? 9) - (PEER_STATUS_ORDER[b.status] ?? 9) || b.lastActivity - a.lastActivity, + ); + if (peers.length === 0) { + return [renderStatusLine({ icon: "info", title: "IRC peers", meta: ["no other agents"] }, theme)]; + } + const counts = new Map(); + for (const peer of peers) counts.set(peer.status, (counts.get(peer.status) ?? 0) + 1); + const meta = [...counts].map(([status, count]) => `${count} ${status}`); + const unreadTotal = peers.reduce((sum, peer) => sum + peer.unread, 0); + if (unreadTotal > 0) meta.push(theme.fg("warning", `${unreadTotal} unread`)); + const header = renderStatusLine({ iconOverride: ircGlyph(theme), title: "IRC peers", meta }, theme); + const items = renderTreeList( + { + items: peers, + expanded, + maxCollapsed: PREVIEW_LIMITS.COLLAPSED_ITEMS, + itemType: "peer", + renderItem: peer => { + const kindText = peer.parentId ? `${peer.kind}${theme.sep.dot}of ${peer.parentId}` : peer.kind; + const unread = peer.unread > 0 ? ` ${formatBadge(`${peer.unread} unread`, "warning", theme)}` : ""; + const age = messageAge(peer.lastActivity); + return `${peerStatusBadge(peer.status, theme)} ${theme.bold(replaceTabs(peer.id))} ${theme.fg("dim", kindText)}${unread}${age ? ` ${theme.fg("dim", age)}` : ""}`; + }, + }, + theme, + ); + return [header, ...items]; +} + +function buildResultLines( + result: { content: Array<{ type: string; text?: string }>; isError?: boolean }, + details: Partial, + args: IrcRenderArgs | undefined, + expanded: boolean, + theme: Theme, +): string[] { + switch (details.op ?? args?.op) { + case "send": + return renderSendResult(result, details, args, expanded, theme); + case "wait": + return renderWaitResult(result, details, args, expanded, theme); + case "inbox": + return renderInboxResult(details, args, expanded, theme); + case "list": + return renderListResult(details, expanded, theme); + default: { + const text = textContent(result) || (result.isError ? "IRC call failed." : "Done."); + return [ + renderStatusLine({ icon: result.isError ? "error" : "success", title: callTitle(args, theme) }, theme), + result.isError ? formatErrorDetail(text, theme) : ` ${theme.fg("muted", replaceTabs(text))}`, + ]; + } + } +} + +export const ircToolRenderer = { + inline: true, + mergeCallAndResult: true, + + renderCall(args: IrcRenderArgs, _options: RenderResultOptions, uiTheme: Theme): Component { + const lines = [ + renderStatusLine({ icon: "pending", title: callTitle(args, uiTheme), meta: callMeta(args) }, uiTheme), + ]; + if (args?.op === "send" && args.message?.trim()) { + lines.push(...bodyLines(args.message, false, uiTheme, { indent: " ", tone: "dim", collapsedLines: 1 })); + } + return new Text(lines.join("\n"), 0, 0); + }, + + renderResult( + result: { content: Array<{ type: string; text?: string }>; details?: IrcDetails; isError?: boolean }, + options: RenderResultOptions, + uiTheme: Theme, + args?: IrcRenderArgs, + ): Component { + const details: Partial = result.details ?? {}; + return createCachedComponent( + () => options.expanded, + (width, expanded) => + buildResultLines(result, details, args, expanded, uiTheme).map(line => + truncateToWidth(line, width, Ellipsis.Unicode), + ), + ); + }, +}; diff --git a/packages/coding-agent/src/tools/renderers.ts b/packages/coding-agent/src/tools/renderers.ts index dde5dc6be..362fdfae5 100644 --- a/packages/coding-agent/src/tools/renderers.ts +++ b/packages/coding-agent/src/tools/renderers.ts @@ -21,6 +21,7 @@ import { evalToolRenderer } from "./eval-render"; import { findToolRenderer } from "./find"; import { githubToolRenderer } from "./gh-renderer"; import { inspectImageToolRenderer } from "./inspect-image-renderer"; +import { ircToolRenderer } from "./irc"; import { jobToolRenderer } from "./job"; import { recallToolRenderer, reflectToolRenderer, retainToolRenderer } from "./memory-render"; import { readToolRenderer } from "./read"; @@ -58,6 +59,7 @@ export const toolRenderers: Record = { search: searchToolRenderer as ToolRenderer, lsp: lspToolRenderer as ToolRenderer, inspect_image: inspectImageToolRenderer as ToolRenderer, + irc: ircToolRenderer as ToolRenderer, read: readToolRenderer as ToolRenderer, job: jobToolRenderer as ToolRenderer, resolve: resolveToolRenderer as ToolRenderer, diff --git a/packages/coding-agent/test/tools/irc-renderer.test.ts b/packages/coding-agent/test/tools/irc-renderer.test.ts new file mode 100644 index 000000000..768ddc273 --- /dev/null +++ b/packages/coding-agent/test/tools/irc-renderer.test.ts @@ -0,0 +1,256 @@ +import { describe, expect, it } from "bun:test"; +import type { IrcMessage } from "@oh-my-pi/pi-coding-agent/irc/bus"; +import { getThemeByName } from "@oh-my-pi/pi-coding-agent/modes/theme/theme"; +import { type IrcDetails, ircToolRenderer } from "@oh-my-pi/pi-coding-agent/tools/irc"; +import { sanitizeText } from "@oh-my-pi/pi-utils"; + +async function theme() { + const t = await getThemeByName("dark"); + expect(t).toBeDefined(); + return t!; +} + +const lines = (component: { render: (w: number) => readonly string[] }, width = 200) => + sanitizeText(component.render(width).join("\n")).split("\n"); + +const msg = (overrides: Partial): IrcMessage => ({ + id: "7181122334455667789", + from: "AuthLoader", + to: "Main", + body: "session-store rename is merged.", + ts: Date.now() - 30_000, + ...overrides, +}); + +describe("ircToolRenderer send", () => { + it("folds a single delivery outcome into the header and shows the awaited reply", async () => { + const uiTheme = await theme(); + const rendered = lines( + ircToolRenderer.renderResult( + { + content: [{ type: "text", text: "" }], + details: { + op: "send", + from: "Main", + to: "AuthLoader", + receipts: [{ to: "AuthLoader", outcome: "revived" }], + waited: msg({ body: "go ahead, auth.ts is yours." }), + } satisfies IrcDetails, + }, + { expanded: false, isPartial: false }, + uiTheme, + { op: "send", to: "AuthLoader", message: "Are you done with auth.ts?", await: true }, + ), + ); + expect(rendered[0]).toContain("AuthLoader"); + expect(rendered[0]).toContain("revived"); + expect(rendered.some(line => line.includes("Are you done with auth.ts?"))).toBe(true); + expect(rendered.some(line => line.includes("go ahead, auth.ts is yours."))).toBe(true); + }); + + it("lists per-recipient outcomes with error text when a broadcast partially fails", async () => { + const uiTheme = await theme(); + const rendered = lines( + ircToolRenderer.renderResult( + { + content: [{ type: "text", text: "" }], + details: { + op: "send", + from: "Main", + to: "all", + receipts: [ + { to: "AuthLoader", outcome: "woken" }, + { to: "RateLimiter", outcome: "failed", error: 'unknown agent "RateLimiter"' }, + ], + } satisfies IrcDetails, + }, + { expanded: false, isPartial: false }, + uiTheme, + { op: "send", to: "all", message: "heads up" }, + ), + ); + expect(rendered[0]).toContain("broadcast"); + expect(rendered[0]).toContain("1 delivered"); + expect(rendered[0]).toContain("1 failed"); + expect(rendered.some(line => line.includes("AuthLoader") && line.includes("woken"))).toBe(true); + expect(rendered.some(line => line.includes("RateLimiter") && line.includes('unknown agent "RateLimiter"'))).toBe( + true, + ); + }); + + it("flags an awaited send whose reply timed out", async () => { + const uiTheme = await theme(); + const rendered = lines( + ircToolRenderer.renderResult( + { + content: [{ type: "text", text: "" }], + details: { + op: "send", + from: "Main", + to: "AuthLoader", + receipts: [{ to: "AuthLoader", outcome: "injected" }], + waited: null, + } satisfies IrcDetails, + }, + { expanded: false, isPartial: false }, + uiTheme, + { op: "send", to: "AuthLoader", message: "ping", await: true }, + ), + ); + expect(rendered[0]).toContain("no reply"); + expect(rendered.some(line => line.includes("No reply yet"))).toBe(true); + }); + + it("surfaces pre-delivery validation failures as an error detail", async () => { + const uiTheme = await theme(); + const rendered = lines( + ircToolRenderer.renderResult( + { + content: [{ type: "text", text: '`to` is required for op="send".' }], + details: { op: "send", from: "Main" } satisfies IrcDetails, + isError: true, + }, + { expanded: false, isPartial: false }, + uiTheme, + { op: "send" }, + ), + ); + expect(rendered.some(line => line.includes('`to` is required for op="send".'))).toBe(true); + }); +}); + +describe("ircToolRenderer wait", () => { + it("renders the consumed message under a sender header", async () => { + const uiTheme = await theme(); + const rendered = lines( + ircToolRenderer.renderResult( + { + content: [{ type: "text", text: "" }], + details: { op: "wait", from: "Main", waited: msg({}) } satisfies IrcDetails, + }, + { expanded: false, isPartial: false }, + uiTheme, + { op: "wait", from: "AuthLoader" }, + ), + ); + expect(rendered[0]).toContain("AuthLoader"); + expect(rendered.some(line => line.includes("session-store rename is merged."))).toBe(true); + }); + + it("marks a timed-out wait without inventing a message", async () => { + const uiTheme = await theme(); + const rendered = lines( + ircToolRenderer.renderResult( + { + content: [{ type: "text", text: "No message from AuthLoader within 2m." }], + details: { op: "wait", from: "Main", waited: null } satisfies IrcDetails, + }, + { expanded: false, isPartial: false }, + uiTheme, + { op: "wait", from: "AuthLoader" }, + ), + ); + expect(rendered[0]).toContain("timed out"); + expect(rendered.some(line => line.includes("No message from AuthLoader within 2m."))).toBe(true); + }); +}); + +describe("ircToolRenderer inbox", () => { + it("lists each message with sender and body preview", async () => { + const uiTheme = await theme(); + const rendered = lines( + ircToolRenderer.renderResult( + { + content: [{ type: "text", text: "" }], + details: { + op: "inbox", + from: "Main", + inbox: [ + msg({ from: "AuthLoader", body: "bus landed." }), + msg({ from: "RateLimiter", body: "receipts carry outcome.", replyTo: "7181122334455667791" }), + ], + } satisfies IrcDetails, + }, + { expanded: false, isPartial: false }, + uiTheme, + { op: "inbox", peek: true }, + ), + ); + expect(rendered[0]).toContain("2 messages"); + expect(rendered[0]).toContain("peek"); + expect(rendered.some(line => line.includes("bus landed."))).toBe(true); + expect(rendered.some(line => line.includes("RateLimiter"))).toBe(true); + expect(rendered.some(line => line.includes("receipts carry outcome."))).toBe(true); + }); +}); + +describe("ircToolRenderer list", () => { + it("summarizes status counts and flags unread peers", async () => { + const uiTheme = await theme(); + const rendered = lines( + ircToolRenderer.renderResult( + { + content: [{ type: "text", text: "" }], + details: { + op: "list", + from: "Main", + peers: [ + { + id: "RateLimiter", + displayName: "task", + kind: "sub", + status: "parked", + parentId: "Main", + unread: 2, + lastActivity: Date.now() - 12 * 60_000, + }, + { + id: "AuthLoader", + displayName: "task", + kind: "sub", + status: "running", + parentId: "Main", + unread: 0, + lastActivity: Date.now() - 2 * 60_000, + }, + ], + } satisfies IrcDetails, + }, + { expanded: false, isPartial: false }, + uiTheme, + { op: "list" }, + ), + ); + expect(rendered[0]).toContain("1 running"); + expect(rendered[0]).toContain("1 parked"); + expect(rendered[0]).toContain("2 unread"); + // Running peers sort above parked ones regardless of input order. + const authIndex = rendered.findIndex(line => line.includes("AuthLoader")); + const rateIndex = rendered.findIndex(line => line.includes("RateLimiter")); + expect(authIndex).toBeGreaterThan(0); + expect(authIndex).toBeLessThan(rateIndex); + expect(rendered.some(line => line.includes("RateLimiter") && line.includes("2 unread"))).toBe(true); + }); +}); + +describe("ircToolRenderer body truncation", () => { + it("collapses long bodies with an elision counter and expands on demand", async () => { + const uiTheme = await theme(); + const body = Array.from({ length: 6 }, (_, i) => `reply line ${i + 1}`).join("\n"); + const details: IrcDetails = { op: "wait", from: "Main", waited: msg({ body }) }; + const result = { content: [{ type: "text", text: "" }], details }; + + const collapsed = lines( + ircToolRenderer.renderResult(result, { expanded: false, isPartial: false }, uiTheme, { op: "wait" }), + ); + expect(collapsed.some(line => line.includes("reply line 2"))).toBe(true); + expect(collapsed.some(line => line.includes("reply line 3"))).toBe(false); + expect(collapsed.some(line => line.includes("+4 more lines"))).toBe(true); + + const expanded = lines( + ircToolRenderer.renderResult(result, { expanded: true, isPartial: false }, uiTheme, { op: "wait" }), + ); + expect(expanded.some(line => line.includes("reply line 6"))).toBe(true); + expect(expanded.some(line => line.includes("more lines"))).toBe(false); + }); +}); diff --git a/packages/coding-agent/test/tools/irc.test.ts b/packages/coding-agent/test/tools/irc.test.ts index a14c8016e..2b88dcfa4 100644 --- a/packages/coding-agent/test/tools/irc.test.ts +++ b/packages/coding-agent/test/tools/irc.test.ts @@ -1,83 +1,56 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "bun:test"; import { Agent } from "@oh-my-pi/pi-agent-core"; import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; +import { IrcBus, type IrcMessage } from "@oh-my-pi/pi-coding-agent/irc/bus"; +import { AgentLifecycleManager } from "@oh-my-pi/pi-coding-agent/registry/agent-lifecycle"; import { AgentRegistry } from "@oh-my-pi/pi-coding-agent/registry/agent-registry"; -import { AgentSession } from "@oh-my-pi/pi-coding-agent/session/agent-session"; -import { type FileEntry, SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager"; +import { AgentSession, type AgentSessionEvent } from "@oh-my-pi/pi-coding-agent/session/agent-session"; +import type { CustomMessage } from "@oh-my-pi/pi-coding-agent/session/messages"; +import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager"; import type { ToolSession } from "@oh-my-pi/pi-coding-agent/tools"; import { IrcTool } from "@oh-my-pi/pi-coding-agent/tools/irc"; interface FakeSession { session: AgentSession; - calls: Array<{ from: string; message: string; awaitReply: boolean }>; - /** Override the reply this fake session generates. */ - setReply: (text: string) => void; - /** Cause the next respondAsBackground call to throw. */ + /** Messages delivered into this session via deliverIrcMessage. */ + delivered: IrcMessage[]; + /** Outcome the fake reports (busy vs idle recipient). */ + setOutcome: (outcome: "injected" | "woken") => void; + /** Cause the next deliverIrcMessage call to throw. */ setError: (error: Error) => void; - /** Resolve the next respondAsBackground call only when allowed. */ - gateNextCall: () => { release: () => void }; - /** Keep the next respondAsBackground call pending until aborted. */ - hangNextCall: () => void; + /** Side effect run on delivery (e.g. reply via the bus). */ + onDeliver: (fn: (msg: IrcMessage) => void) => void; } + function makeFakeSession(): FakeSession { - let nextReply = "auto-reply"; + let outcome: "injected" | "woken" = "injected"; let nextError: Error | null = null; - let gate: { promise: Promise; release: () => void } | null = null; - let hangNext = false; - const calls: Array<{ from: string; message: string; awaitReply: boolean }> = []; + let deliverHook: ((msg: IrcMessage) => void) | undefined; + const delivered: IrcMessage[] = []; const session = { - respondAsBackground: async (args: { - from: string; - message: string; - awaitReply?: boolean; - signal?: AbortSignal; - }) => { - const awaitReply = args.awaitReply !== false; - calls.push({ from: args.from, message: args.message, awaitReply }); - if (gate) { - const g = gate; - gate = null; - await g.promise; - } - if (hangNext) { - hangNext = false; - const deferred = Promise.withResolvers(); - if (args.signal?.aborted) { - deferred.reject(args.signal.reason instanceof Error ? args.signal.reason : new Error("aborted")); - } else { - args.signal?.addEventListener( - "abort", - () => - deferred.reject(args.signal?.reason instanceof Error ? args.signal.reason : new Error("aborted")), - { once: true }, - ); - } - return await deferred.promise; - } + deliverIrcMessage: async (msg: IrcMessage) => { if (nextError) { const err = nextError; nextError = null; throw err; } - return { replyText: awaitReply ? nextReply : null }; + delivered.push(msg); + deliverHook?.(msg); + return outcome; }, + emitIrcRelayObservation: () => {}, }; return { session: session as unknown as AgentSession, - calls, - setReply: text => { - nextReply = text; + delivered, + setOutcome: value => { + outcome = value; }, setError: error => { nextError = error; }, - gateNextCall: () => { - const { promise, resolve } = Promise.withResolvers(); - gate = { promise, release: resolve }; - return { release: resolve }; - }, - hangNextCall: () => { - hangNext = true; + onDeliver: fn => { + deliverHook = fn; }, }; } @@ -93,7 +66,8 @@ function makeToolSession(registry: AgentRegistry, agentId: string): ToolSession getAgentId: () => agentId, }; } -function createRealBackgroundSession(): { session: AgentSession; sessionManager: SessionManager } { + +function createRealSession(): { session: AgentSession; sessionManager: SessionManager } { const sessionManager = SessionManager.inMemory("/tmp"); const session = new AgentSession({ agent: new Agent({ @@ -110,13 +84,17 @@ function createRealBackgroundSession(): { session: AgentSession; sessionManager: return { session, sessionManager }; } -describe("IrcTool", () => { +describe("IRC", () => { let registry: AgentRegistry; + let bus: IrcBus; const sessions: AgentSession[] = []; beforeEach(() => { AgentRegistry.resetGlobalForTests(); + AgentLifecycleManager.resetGlobalForTests(); + IrcBus.resetGlobalForTests(); registry = AgentRegistry.global(); + bus = IrcBus.global(); }); afterEach(async () => { vi.restoreAllMocks(); @@ -125,217 +103,423 @@ describe("IrcTool", () => { } }); - it("createIf returns null when irc is disabled", () => { - const session: ToolSession = { - cwd: "/tmp", - hasUI: false, - getSessionFile: () => null, - getSessionSpawns: () => "*", - settings: Settings.isolated(), - agentRegistry: registry, - getAgentId: () => "0-Main", - }; - session.settings.set("irc.enabled", false); - expect(IrcTool.createIf(session)).toBeNull(); - }); + describe("IrcBus", () => { + it("send delivers to a live recipient and reports the session outcome", async () => { + const sub = makeFakeSession(); + registry.register({ id: "0-Sub", displayName: "task", kind: "sub", session: sub.session }); - it("createIf returns null without registry/agentId", () => { - const session: ToolSession = { - cwd: "/tmp", - hasUI: false, - getSessionFile: () => null, - getSessionSpawns: () => "*", - settings: Settings.isolated(), - }; - expect(IrcTool.createIf(session)).toBeNull(); - }); + sub.setOutcome("injected"); + const injected = await bus.send({ from: "0-Main", to: "0-Sub", body: "ping" }); + expect(injected).toEqual({ to: "0-Sub", outcome: "injected" }); - it("op=list returns peers visible to the caller", async () => { - const main = makeFakeSession(); - const sub = makeFakeSession(); - registry.register({ id: "0-Main", displayName: "main", kind: "main", session: main.session }); - registry.register({ - id: "0-AuthLoader", - displayName: "task", - kind: "sub", - parentId: "0-Main", - session: sub.session, + sub.setOutcome("woken"); + const woken = await bus.send({ from: "0-Main", to: "0-Sub", body: "ping again" }); + expect(woken.outcome).toBe("woken"); + + expect(sub.delivered.map(msg => msg.body)).toEqual(["ping", "ping again"]); + expect(sub.delivered[0]?.from).toBe("0-Main"); + expect(sub.delivered[0]?.id).toBeTruthy(); }); - const tool = new IrcTool(makeToolSession(registry, "0-Main")); - const result = await tool.execute("call-1", { op: "list" }); - expect(result.details?.op).toBe("list"); - expect(result.details?.peers).toEqual([ - { + it("send to an unknown or aborted agent fails", async () => { + const unknown = await bus.send({ from: "0-Main", to: "0-Ghost", body: "hello?" }); + expect(unknown.outcome).toBe("failed"); + + const sub = makeFakeSession(); + registry.register({ id: "0-Dead", displayName: "task", kind: "sub", session: sub.session }); + registry.setStatus("0-Dead", "aborted"); + const aborted = await bus.send({ from: "0-Main", to: "0-Dead", body: "hello?" }); + expect(aborted.outcome).toBe("failed"); + }); + + it("send surfaces recipient delivery errors as failed", async () => { + const sub = makeFakeSession(); + registry.register({ id: "0-Sub", displayName: "task", kind: "sub", session: sub.session }); + sub.setError(new Error("boom")); + const receipt = await bus.send({ from: "0-Main", to: "0-Sub", body: "ping" }); + expect(receipt).toEqual({ to: "0-Sub", outcome: "failed", error: "boom" }); + }); + + it("send revives a parked recipient through the lifecycle manager", async () => { + const sub = makeFakeSession(); + sub.setOutcome("woken"); + registry.register({ id: "0-Parked", displayName: "task", kind: "sub", session: null, status: "parked" }); + AgentLifecycleManager.global().adopt("0-Parked", { + idleTtlMs: 0, + revive: async () => sub.session, + }); + + const receipt = await bus.send({ from: "0-Main", to: "0-Parked", body: "wake up" }); + expect(receipt.outcome).toBe("revived"); + expect(sub.delivered.map(msg => msg.body)).toEqual(["wake up"]); + expect(registry.get("0-Parked")?.status).toBe("idle"); + }); + + it("send fails cleanly when a parked recipient has no reviver", async () => { + registry.register({ id: "0-Parked", displayName: "task", kind: "sub", session: null, status: "parked" }); + AgentLifecycleManager.global().adopt("0-Parked", { idleTtlMs: 0 }); + const receipt = await bus.send({ from: "0-Main", to: "0-Parked", body: "wake up" }); + expect(receipt.outcome).toBe("failed"); + expect(receipt.error).toBeTruthy(); + }); + + it("wait consumes a matching send instead of delivering it to the session", async () => { + const main = makeFakeSession(); + registry.register({ id: "0-Main", displayName: "main", kind: "main", session: main.session }); + const sub = makeFakeSession(); + registry.register({ id: "0-Sub", displayName: "task", kind: "sub", session: sub.session }); + + const waiting = bus.wait("0-Main", { from: "0-Sub" }, 1000); + const receipt = await bus.send({ from: "0-Sub", to: "0-Main", body: "pong" }); + expect(receipt.outcome).toBe("injected"); + + const msg = await waiting; + expect(msg?.body).toBe("pong"); + // The waiter consumed the message: no session delivery, no inbox copy. + expect(main.delivered).toEqual([]); + expect(bus.unreadCount("0-Main")).toBe(0); + }); + + it("wait from-filter ignores messages from other senders", async () => { + const main = makeFakeSession(); + registry.register({ id: "0-Main", displayName: "main", kind: "main", session: main.session }); + const a = makeFakeSession(); + registry.register({ id: "0-A", displayName: "task", kind: "sub", session: a.session }); + const b = makeFakeSession(); + registry.register({ id: "0-B", displayName: "task", kind: "sub", session: b.session }); + + const waiting = bus.wait("0-Main", { from: "0-B" }, 1000); + await bus.send({ from: "0-A", to: "0-Main", body: "not for the waiter" }); + // The non-matching message fell through to normal delivery. + expect(main.delivered.map(msg => msg.body)).toEqual(["not for the waiter"]); + + await bus.send({ from: "0-B", to: "0-Main", body: "for the waiter" }); + const msg = await waiting; + expect(msg?.from).toBe("0-B"); + expect(msg?.body).toBe("for the waiter"); + }); + + it("wait returns null on timeout and rejects on abort", async () => { + // Genuine 5ms wall-clock timeout: this deliberately exercises the + // bus's real timer path; nothing else races it. + expect(await bus.wait("0-Main", {}, 5)).toBeNull(); + + const controller = new AbortController(); + const waiting = bus.wait("0-Main", {}, 1000, controller.signal); + controller.abort(new Error("cancelled")); + await expect(waiting).rejects.toThrow("cancelled"); + }); + + it("wait drains an already-pending mailbox message first", async () => { + const main = makeFakeSession(); + registry.register({ id: "0-Main", displayName: "main", kind: "main", session: main.session }); + const sub = makeFakeSession(); + registry.register({ id: "0-Sub", displayName: "task", kind: "sub", session: sub.session }); + + await bus.send({ from: "0-Sub", to: "0-Main", body: "earlier" }); + expect(bus.unreadCount("0-Main")).toBe(1); + + // Resolves from the mailbox synchronously; the timeout never fires. + const msg = await bus.wait("0-Main", { from: "0-Sub" }, 5); + expect(msg?.body).toBe("earlier"); + expect(bus.unreadCount("0-Main")).toBe(0); + }); + + it("inbox peeks or drains pending messages", async () => { + const main = makeFakeSession(); + registry.register({ id: "0-Main", displayName: "main", kind: "main", session: main.session }); + const sub = makeFakeSession(); + registry.register({ id: "0-Sub", displayName: "task", kind: "sub", session: sub.session }); + + await bus.send({ from: "0-Sub", to: "0-Main", body: "one" }); + await bus.send({ from: "0-Sub", to: "0-Main", body: "two" }); + + const peeked = bus.inbox("0-Main", { peek: true }); + expect(peeked.map(msg => msg.body)).toEqual(["one", "two"]); + expect(bus.unreadCount("0-Main")).toBe(2); + + const drained = bus.inbox("0-Main"); + expect(drained.map(msg => msg.body)).toEqual(["one", "two"]); + expect(bus.unreadCount("0-Main")).toBe(0); + expect(bus.inbox("0-Main")).toEqual([]); + }); + + it("wait does not leak the waiter after timeout or abort", async () => { + const main = makeFakeSession(); + registry.register({ id: "0-Main", displayName: "main", kind: "main", session: main.session }); + const sub = makeFakeSession(); + registry.register({ id: "0-Sub", displayName: "task", kind: "sub", session: sub.session }); + + // Timed-out waiter is removed: a later send goes to normal delivery. + expect(await bus.wait("0-Main", {}, 5)).toBeNull(); + const afterTimeout = await bus.send({ from: "0-Sub", to: "0-Main", body: "after timeout" }); + expect(afterTimeout.outcome).toBe("injected"); + expect(main.delivered.map(msg => msg.body)).toEqual(["after timeout"]); + expect(bus.unreadCount("0-Main")).toBe(1); + bus.inbox("0-Main"); + + // Aborted waiter is removed too: the dead waiter never consumes mail. + const controller = new AbortController(); + const waiting = bus.wait("0-Main", {}, 1000, controller.signal); + controller.abort(new Error("cancelled")); + await expect(waiting).rejects.toThrow("cancelled"); + await bus.send({ from: "0-Sub", to: "0-Main", body: "after abort" }); + expect(main.delivered.map(msg => msg.body)).toEqual(["after timeout", "after abort"]); + expect(bus.unreadCount("0-Main")).toBe(1); + }); + + it("resolves waiters in FIFO order", async () => { + const main = makeFakeSession(); + registry.register({ id: "0-Main", displayName: "main", kind: "main", session: main.session }); + const sub = makeFakeSession(); + registry.register({ id: "0-Sub", displayName: "task", kind: "sub", session: sub.session }); + + const first = bus.wait("0-Main", {}, 1000); + const second = bus.wait("0-Main", {}, 1000); + await bus.send({ from: "0-Sub", to: "0-Main", body: "one" }); + await bus.send({ from: "0-Sub", to: "0-Main", body: "two" }); + + expect((await first)?.body).toBe("one"); + expect((await second)?.body).toBe("two"); + // Both messages were consumed by waiters, none reached the session. + expect(main.delivered).toEqual([]); + expect(bus.unreadCount("0-Main")).toBe(0); + }); + + it("mailbox drops the oldest message beyond the 100-message cap", async () => { + const main = makeFakeSession(); + registry.register({ id: "0-Main", displayName: "main", kind: "main", session: main.session }); + const sub = makeFakeSession(); + registry.register({ id: "0-Sub", displayName: "task", kind: "sub", session: sub.session }); + + for (let i = 0; i <= 100; i++) { + await bus.send({ from: "0-Sub", to: "0-Main", body: `msg-${i}` }); + } + + expect(bus.unreadCount("0-Main")).toBe(100); + const pending = bus.inbox("0-Main", { peek: true }); + expect(pending[0]?.body).toBe("msg-1"); + expect(pending[pending.length - 1]?.body).toBe("msg-100"); + }); + + it("send surfaces the reviver's error message when revival fails", async () => { + registry.register({ id: "0-Parked", displayName: "task", kind: "sub", session: null, status: "parked" }); + AgentLifecycleManager.global().adopt("0-Parked", { + idleTtlMs: 0, + revive: async () => { + throw new Error("revive exploded"); + }, + }); + + const receipt = await bus.send({ from: "0-Main", to: "0-Parked", body: "wake up" }); + expect(receipt).toEqual({ to: "0-Parked", outcome: "failed", error: "revive exploded" }); + // Failed revival never enqueues: the message is lost, not buffered. + expect(bus.unreadCount("0-Parked")).toBe(0); + }); + }); + + describe("IrcTool", () => { + it("createIf returns null when irc is disabled", () => { + const session: ToolSession = { + cwd: "/tmp", + hasUI: false, + getSessionFile: () => null, + getSessionSpawns: () => "*", + settings: Settings.isolated(), + agentRegistry: registry, + getAgentId: () => "0-Main", + }; + session.settings.set("irc.enabled", false); + expect(IrcTool.createIf(session)).toBeNull(); + }); + + it("createIf returns null without registry/agentId", () => { + const session: ToolSession = { + cwd: "/tmp", + hasUI: false, + getSessionFile: () => null, + getSessionSpawns: () => "*", + settings: Settings.isolated(), + }; + expect(IrcTool.createIf(session)).toBeNull(); + }); + + it("op=list includes parked peers, unread counts, and parent ids", async () => { + const sub = makeFakeSession(); + registry.register({ id: "0-AuthLoader", displayName: "task", kind: "sub", - status: "running", parentId: "0-Main", - }, - ]); - expect(result.details?.channels).toEqual(["all", "0-AuthLoader"]); - }); + session: sub.session, + }); + registry.register({ id: "0-Parked", displayName: "task", kind: "sub", session: null, status: "parked" }); + const main = makeFakeSession(); + registry.register({ id: "0-Main", displayName: "main", kind: "main", session: main.session }); + await bus.send({ from: "0-Main", to: "0-AuthLoader", body: "unread one" }); - it("op=send DM returns the recipient's prose reply", async () => { - const main = makeFakeSession(); - const sub = makeFakeSession(); - sub.setReply("pong"); - registry.register({ id: "0-Main", displayName: "main", kind: "main", session: main.session }); - registry.register({ - id: "0-AuthLoader", - displayName: "task", - kind: "sub", - parentId: "0-Main", - session: sub.session, + const tool = new IrcTool(makeToolSession(registry, "0-Main")); + const result = await tool.execute("call-1", { op: "list" }); + expect(result.details?.op).toBe("list"); + expect(result.details?.peers).toMatchObject([ + { id: "0-AuthLoader", status: "running", parentId: "0-Main", unread: 1 }, + { id: "0-Parked", status: "parked", unread: 0 }, + ]); + const text = result.content[0]?.type === "text" ? result.content[0].text : ""; + expect(text).toContain("Parked agents are revived automatically"); }); - const tool = new IrcTool(makeToolSession(registry, "0-Main")); - const result = await tool.execute("call-2", { - op: "send", - to: "0-AuthLoader", - message: "ping", - }); - expect(result.details?.delivered).toEqual(["0-AuthLoader"]); - expect(result.details?.replies).toEqual([{ from: "0-AuthLoader", text: "pong" }]); - expect(sub.calls).toEqual([{ from: "0-Main", message: "ping", awaitReply: true }]); - }); + it("op=send returns receipts immediately without waiting for a reply", async () => { + const sub = makeFakeSession(); + registry.register({ id: "0-Sub", displayName: "task", kind: "sub", session: sub.session }); - it("op=send returns immediately even when the recipient is mid-tool-call", async () => { - // Simulate "blocked recipient": gateNextCall holds respondAsBackground - // pending until we release it. From the sender's perspective the call - // must still complete because the side-channel does not block on the - // recipient's main loop in a real session. - const main = makeFakeSession(); - const sub = makeFakeSession(); - sub.setReply("ok"); - registry.register({ id: "0-Main", displayName: "main", kind: "main", session: main.session }); - registry.register({ id: "0-Busy", displayName: "task", kind: "sub", parentId: "0-Main", session: sub.session }); - - const gate = sub.gateNextCall(); - const tool = new IrcTool(makeToolSession(registry, "0-Main")); - const pending = tool.execute("call-3", { op: "send", to: "0-Busy", message: "are you there?" }); - // Release the gate after a microtask: the dispatch was already issued - // even though the recipient was holding. - setTimeout(() => gate.release(), 5); - const result = await pending; - expect(result.details?.delivered).toEqual(["0-Busy"]); - expect(result.details?.replies).toEqual([{ from: "0-Busy", text: "ok" }]); - }); - - it("op=send to=all broadcasts (default no reply, only injection on each peer)", async () => { - const main = makeFakeSession(); - const subA = makeFakeSession(); - const subB = makeFakeSession(); - registry.register({ id: "0-Main", displayName: "main", kind: "main", session: main.session }); - registry.register({ id: "0-A", displayName: "task", kind: "sub", parentId: "0-Main", session: subA.session }); - registry.register({ id: "0-B", displayName: "task", kind: "sub", parentId: "0-Main", session: subB.session }); - - const tool = new IrcTool(makeToolSession(registry, "0-Main")); - const result = await tool.execute("call-4", { op: "send", to: "all", message: "anyone there?" }); - expect(new Set(result.details?.delivered)).toEqual(new Set(["0-A", "0-B"])); - expect(result.details?.replies ?? []).toEqual([]); - expect(subA.calls).toEqual([{ from: "0-Main", message: "anyone there?", awaitReply: false }]); - expect(subB.calls).toEqual([{ from: "0-Main", message: "anyone there?", awaitReply: false }]); - }); - - it("op=send returns notFound when target is unknown", async () => { - const main = makeFakeSession(); - registry.register({ id: "0-Main", displayName: "main", kind: "main", session: main.session }); - - const tool = new IrcTool(makeToolSession(registry, "0-Main")); - const result = await tool.execute("call-5", { op: "send", to: "0-Ghost", message: "hi" }); - expect(result.details?.delivered ?? []).toEqual([]); - expect(result.details?.notFound).toEqual(["0-Ghost"]); - }); - - it("op=send fails a hung recipient after the configured timeout", async () => { - const main = makeFakeSession(); - const sub = makeFakeSession(); - sub.hangNextCall(); - registry.register({ id: "0-Main", displayName: "main", kind: "main", session: main.session }); - registry.register({ id: "0-Hung", displayName: "task", kind: "sub", parentId: "0-Main", session: sub.session }); - - const toolSession = makeToolSession(registry, "0-Main"); - toolSession.settings.set("irc.timeoutMs", 5); - const tool = new IrcTool(toolSession); - const result = await tool.execute("call-timeout", { op: "send", to: "0-Hung", message: "ping" }); - - expect(result.details?.delivered ?? []).toEqual([]); - expect(result.details?.failed).toEqual([{ id: "0-Hung", error: "IRC timed out waiting for 0-Hung after 5 ms" }]); - expect(sub.calls).toEqual([{ from: "0-Main", message: "ping", awaitReply: true }]); - }); - it("op=send preserves the recipient's incoming DM when the auto-reply times out", async () => { - const main = makeFakeSession(); - const { session: recipient, sessionManager } = createRealBackgroundSession(); - sessions.push(recipient); - vi.spyOn(recipient, "runEphemeralTurn").mockImplementation(async ({ signal }) => { - if (!signal) { - throw new Error("Missing abort signal"); - } - const deferred = Promise.withResolvers(); - const rejectOnAbort = () => { - deferred.reject(signal.reason instanceof Error ? signal.reason : new Error("IRC aborted")); - }; - if (signal.aborted) { - rejectOnAbort(); - } else { - signal.addEventListener("abort", rejectOnAbort, { once: true }); - } - return await deferred.promise; - }); - registry.register({ id: "0-Main", displayName: "main", kind: "main", session: main.session }); - registry.register({ - id: "0-Hung", - displayName: "task", - kind: "sub", - parentId: "0-Main", - session: recipient, + const tool = new IrcTool(makeToolSession(registry, "0-Main")); + const result = await tool.execute("call-1", { op: "send", to: "0-Sub", message: "ping" }); + expect(result.isError).toBeFalsy(); + expect(result.details?.receipts).toEqual([{ to: "0-Sub", outcome: "injected" }]); + expect(result.details?.waited).toBeUndefined(); + expect(sub.delivered.map(msg => msg.body)).toEqual(["ping"]); }); - const toolSession = makeToolSession(registry, "0-Main"); - toolSession.settings.set("irc.timeoutMs", 5); - const tool = new IrcTool(toolSession); - const result = await tool.execute("call-timeout-persist", { op: "send", to: "0-Hung", message: "ping" }); + it("op=send to=all fans out to live peers and reports per-recipient receipts", async () => { + const a = makeFakeSession(); + registry.register({ id: "0-A", displayName: "task", kind: "sub", session: a.session }); + const b = makeFakeSession(); + b.setError(new Error("kaput")); + registry.register({ id: "0-B", displayName: "task", kind: "sub", session: b.session }); + registry.register({ id: "0-Parked", displayName: "task", kind: "sub", session: null, status: "parked" }); - expect(result.details?.delivered ?? []).toEqual([]); - expect(result.details?.failed).toEqual([{ id: "0-Hung", error: "IRC timed out waiting for 0-Hung after 5 ms" }]); - expect(recipient.messages).toEqual( - expect.arrayContaining([ - expect.objectContaining({ - role: "custom", - customType: "irc:incoming", - content: "[IRC `0-Main` → you]\n\nping", - }), - ]), - ); - expect(recipient.messages).not.toEqual( - expect.arrayContaining([expect.objectContaining({ role: "custom", customType: "irc:autoreply" })]), - ); - const persistedEntries = sessionManager - .captureState() - .fileEntries.filter( - (entry): entry is Extract => - entry.type === "custom_message" && typeof entry.content === "string", - ); - expect(persistedEntries).toEqual([ - expect.objectContaining({ - customType: "irc:incoming", - content: "[IRC `0-Main` → you]\n\nping", - }), - ]); + const tool = new IrcTool(makeToolSession(registry, "0-Main")); + const result = await tool.execute("call-1", { op: "send", to: "all", message: "anyone there?" }); + // Broadcast skips parked agents; one failure does not block the other delivery. + expect(result.details?.receipts).toEqual([ + { to: "0-A", outcome: "injected" }, + { to: "0-B", outcome: "failed", error: "kaput" }, + ]); + expect(a.delivered.map(msg => msg.body)).toEqual(["anyone there?"]); + }); + + it("op=send await=true round-trips the recipient's reply", async () => { + const main = makeFakeSession(); + registry.register({ id: "0-Main", displayName: "main", kind: "main", session: main.session }); + const sub = makeFakeSession(); + registry.register({ id: "0-Sub", displayName: "task", kind: "sub", session: sub.session }); + sub.onDeliver(msg => { + // Reply synchronously during delivery: it lands in the sender's + // mailbox, which the tool's await-wait drains deterministically. + void bus.send({ from: "0-Sub", to: msg.from, body: "pong", replyTo: msg.id }); + }); + + const tool = new IrcTool(makeToolSession(registry, "0-Main")); + const result = await tool.execute("call-1", { op: "send", to: "0-Sub", message: "ping", await: true }); + expect(result.details?.waited?.body).toBe("pong"); + const text = result.content[0]?.type === "text" ? result.content[0].text : ""; + expect(text).toContain("pong"); + }); + + it("op=send await=true reports a clean timeout when no reply arrives", async () => { + const sub = makeFakeSession(); + registry.register({ id: "0-Sub", displayName: "task", kind: "sub", session: sub.session }); + + const tool = new IrcTool(makeToolSession(registry, "0-Main")); + const result = await tool.execute("call-1", { + op: "send", + to: "0-Sub", + message: "ping", + // Real 5ms timeout — exercises the timeout path; no reply ever arrives. + await: true, + timeoutMs: 5, + }); + expect(result.isError).toBeFalsy(); + expect(result.details?.waited).toBeNull(); + const text = result.content[0]?.type === "text" ? result.content[0].text : ""; + expect(text).toContain("No reply from 0-Sub"); + }); + + it("op=send rejects await with to=all and self-sends", async () => { + const tool = new IrcTool(makeToolSession(registry, "0-Main")); + const broadcast = await tool.execute("call-1", { op: "send", to: "all", message: "x", await: true }); + expect(broadcast.isError).toBe(true); + const self = await tool.execute("call-2", { op: "send", to: "0-Main", message: "x" }); + expect(self.isError).toBe(true); + }); + + it("op=send returns a failed receipt for unknown targets", async () => { + const tool = new IrcTool(makeToolSession(registry, "0-Main")); + const result = await tool.execute("call-1", { op: "send", to: "0-Ghost", message: "ping" }); + expect(result.isError).toBe(true); + expect(result.details?.receipts?.[0]?.outcome).toBe("failed"); + }); + + it("op=wait returns a clean non-error timeout result", async () => { + const tool = new IrcTool(makeToolSession(registry, "0-Main")); + const result = await tool.execute("call-1", { op: "wait", timeoutMs: 5 }); + expect(result.isError).toBeFalsy(); + expect(result.details?.waited).toBeNull(); + const text = result.content[0]?.type === "text" ? result.content[0].text : ""; + expect(text).toContain("No message"); + }); + + it("op=inbox drains the caller's mailbox", async () => { + const main = makeFakeSession(); + registry.register({ id: "0-Main", displayName: "main", kind: "main", session: main.session }); + const sub = makeFakeSession(); + registry.register({ id: "0-Sub", displayName: "task", kind: "sub", session: sub.session }); + await bus.send({ from: "0-Sub", to: "0-Main", body: "fyi" }); + + const tool = new IrcTool(makeToolSession(registry, "0-Main")); + const peeked = await tool.execute("call-1", { op: "inbox", peek: true }); + expect(peeked.details?.inbox?.map(msg => msg.body)).toEqual(["fyi"]); + const drained = await tool.execute("call-2", { op: "inbox" }); + expect(drained.details?.inbox?.map(msg => msg.body)).toEqual(["fyi"]); + const empty = await tool.execute("call-3", { op: "inbox" }); + expect(empty.details?.inbox).toEqual([]); + }); }); - it("op=send surfaces recipient errors as failed", async () => { - const main = makeFakeSession(); - const sub = makeFakeSession(); - sub.setError(new Error("model unavailable")); - registry.register({ id: "0-Main", displayName: "main", kind: "main", session: main.session }); - registry.register({ id: "0-Down", displayName: "task", kind: "sub", parentId: "0-Main", session: sub.session }); + describe("AgentSession.deliverIrcMessage", () => { + it("wakes an idle session with a real turn and emits the irc_message event", async () => { + const { session } = createRealSession(); + sessions.push(session); + const promptSpy = vi.spyOn(session.agent, "prompt").mockResolvedValue(undefined); + const ircEvent = new Promise(resolve => { + session.subscribe(event => { + if (event.type === "irc_message") resolve(event); + }); + }); - const tool = new IrcTool(makeToolSession(registry, "0-Main")); - const result = await tool.execute("call-6", { op: "send", to: "0-Down", message: "ping" }); - expect(result.details?.delivered ?? []).toEqual([]); - expect(result.details?.failed).toEqual([{ id: "0-Down", error: "model unavailable" }]); + const outcome = await session.deliverIrcMessage({ + id: "msg-1", + from: "0-Peer", + to: "0-Me", + body: "wake up", + ts: Date.now(), + }); + expect(outcome).toBe("woken"); + expect(promptSpy).toHaveBeenCalledTimes(1); + const prompted = promptSpy.mock.calls[0]?.[0] as unknown as CustomMessage; + expect(prompted).toMatchObject({ role: "custom", customType: "irc:incoming" }); + expect(prompted.details).toMatchObject({ id: "msg-1", from: "0-Peer", message: "wake up" }); + + const event = await ircEvent; + expect(event.type).toBe("irc_message"); + }); + + it("queues a non-interrupting aside when a turn is streaming", async () => { + const { session } = createRealSession(); + sessions.push(session); + const promptSpy = vi.spyOn(session.agent, "prompt").mockResolvedValue(undefined); + Object.defineProperty(session, "isStreaming", { value: true, configurable: true }); + + const outcome = await session.deliverIrcMessage({ + id: "msg-2", + from: "0-Peer", + to: "0-Me", + body: "mid-turn note", + ts: Date.now(), + }); + expect(outcome).toBe("injected"); + expect(promptSpy).not.toHaveBeenCalled(); + }); }); }); From a92d2ce989d6327bcfef1da222490e8ca02fbc8a Mon Sep 17 00:00:00 2001 From: can1357 Date: Wed, 10 Jun 2026 17:49:01 +0200 Subject: [PATCH 14/27] feat(coding-agent): removed context argument from eval agent() spawn Shared background now flows through a '/Users/can/.omp/agent/sessions/-Projects-.tree-pi-commit/2026-06-10T15-36-32-782Z_019eb22d-970e-7000-8964-72c98becf3e8/local' file referenced in each prompt instead of a context string forwarded into the subagent's system prompt. The JS and Python preludes drop the context kwarg from agent(), the subagent system prompt drops the {{#if context}} block and the conversation-context file pointer, and runEvalAgent no longer writes a per-call conversation context file. AgentSession sheds the now-unused formatCompactContext() helper that supplied the file's body, and ToolSession.getCompactContext is removed alongside it. --- docs/tools/eval.md | 8 +-- .../src/eval/__tests__/agent-bridge.test.ts | 5 +- .../coding-agent/src/eval/agent-bridge.ts | 17 +---- .../src/eval/js/shared/prelude.txt | 2 +- packages/coding-agent/src/eval/py/prelude.py | 11 ++-- .../prompts/system/subagent-system-prompt.md | 12 ---- .../src/prompts/system/workflow-notice.md | 4 +- .../coding-agent/src/prompts/tools/eval.md | 6 +- .../coding-agent/src/session/agent-session.ts | 65 ------------------- packages/coding-agent/src/tools/index.ts | 2 - 10 files changed, 19 insertions(+), 113 deletions(-) diff --git a/docs/tools/eval.md b/docs/tools/eval.md index 0fccef25b..054816a0f 100644 --- a/docs/tools/eval.md +++ b/docs/tools/eval.md @@ -138,7 +138,7 @@ Implemented in `packages/coding-agent/src/eval/js/worker-core.ts`, `packages/cod - `await read(path, { offset?, limit? })` - `await tree(path = ".", { maxDepth?, hidden? })` - `sort(text, { reverse?, unique? })`, `uniq(text, { count? })`, `counter(items, { limit?, reverse? })` - - `await agent(prompt, { agentType?, model?, context?, label?, schema? })` + - `await agent(prompt, { agentType?, model?, label?, schema? })` - `await parallel([() => agent("a"), () => agent("b")])` - `await pipeline(items, stage1, stage2)` - `display(value)` behavior: @@ -192,11 +192,11 @@ Both runtimes expose `completion()` — a single stateless completion against a Both runtimes expose `agent()` — a single subagent invocation routed through `packages/coding-agent/src/eval/agent-bridge.ts` into the same `runSubprocess(...)` path used by the `task` tool. It uses the current eval session's spawn policy and inherits the parent eval executor id, so parent and subagent code share JS/Python runtime state. - Signatures: - - JS: `await agent(prompt, { agentType?, model?, context?, label?, schema? })` - - Python: `agent(prompt, *, agent_type="task", model=None, context=None, label=None, schema=None)` + - JS: `await agent(prompt, { agentType?, model?, label?, schema? })` + - Python: `agent(prompt, *, agent_type="task", model=None, label=None, schema=None)` - `agentType` / `agent_type` defaults to the bundled `task` agent and resolves through normal agent discovery, so project and user agents work. - `model` overrides the selected agent's model. Without it, normal per-agent settings and the agent frontmatter model apply. -- `context` supplies shared background; `label` controls the `agent://` output label prefix. +- Shared background is passed via files: write a `local://` file and reference it in the prompt. `label` controls the `agent://` output label prefix. - `schema` passes a JSON Schema to the subagent structured-output path. When present, the helper parses the final JSON text and returns an object. - Spawn restrictions use `session.getSessionSpawns()` exactly like the `task` tool. Eval-driven subagent recursion is capped at depth 3. - JS and Python both expose `parallel(thunks)` and `pipeline(items, ...stages)`; both use a bounded async/threaded pool whose width tracks the `task.maxConcurrency` setting (the same ceiling the `task` tool uses; `0` = run every item at once), preserve item order, and propagate rejections. The width is fetched live from the host via the `__concurrency__` bridge, so the helpers no longer take a `concurrency` argument. diff --git a/packages/coding-agent/src/eval/__tests__/agent-bridge.test.ts b/packages/coding-agent/src/eval/__tests__/agent-bridge.test.ts index a5e263cf8..588fd1b5a 100644 --- a/packages/coding-agent/src/eval/__tests__/agent-bridge.test.ts +++ b/packages/coding-agent/src/eval/__tests__/agent-bridge.test.ts @@ -178,7 +178,7 @@ describe("runEvalAgent", () => { expect(runSpy).not.toHaveBeenCalled(); }); - it("passes the parent execution context and only sets outputSchema when schema is supplied", async () => { + it("passes parent execution options and only sets outputSchema when schema is supplied", async () => { mockAgents(); const runSpy = vi.spyOn(taskExecutor, "runSubprocess").mockImplementation(async options => singleResult(options)); const abortController = new AbortController(); @@ -186,7 +186,7 @@ describe("runEvalAgent", () => { const session = makeSession({ depth: 2, activeModel: "p/current", modelString: "p/fallback" }); await runEvalAgent( - { prompt: " hello ", context: " context ", label: "My Agent", model: "p/override", schema }, + { prompt: " hello ", label: "My Agent", model: "p/override", schema }, { session, signal: abortController.signal }, ); await runEvalAgent({ prompt: "plain" }, { session }); @@ -199,7 +199,6 @@ describe("runEvalAgent", () => { expect(firstOptions.parentActiveModelPattern).toBe("p/current"); expect(firstOptions.outputSchema).toBe(schema); expect(firstOptions.assignment).toBe("hello"); - expect(firstOptions.context).toBe("context"); expect(firstOptions.description).toBe("My Agent"); expect(firstOptions.modelOverride).toEqual(["p/override"]); expect(secondOptions.outputSchema).toBeUndefined(); diff --git a/packages/coding-agent/src/eval/agent-bridge.ts b/packages/coding-agent/src/eval/agent-bridge.ts index 23a4ecff0..d18bb48b4 100644 --- a/packages/coding-agent/src/eval/agent-bridge.ts +++ b/packages/coding-agent/src/eval/agent-bridge.ts @@ -34,7 +34,6 @@ const agentArgsSchema = z.object({ prompt: z.string().min(1, "prompt must be a non-empty string"), agentType: z.string().min(1).optional(), model: z.union([z.string().min(1), z.array(z.string().min(1)).min(1)]).optional(), - context: z.string().optional(), label: z.string().optional(), schema: z.unknown().optional(), }); @@ -43,7 +42,6 @@ interface EvalAgentArgs { prompt: string; agentType?: string; model?: string | string[]; - context?: string; label?: string; schema?: unknown; } @@ -135,20 +133,12 @@ function getOutputManager(session: ToolSession): AgentOutputManager { async function getArtifacts(session: ToolSession): Promise<{ sessionFile: string | null; artifactsDir: string; - contextFile?: string; }> { const sessionFile = session.getSessionFile(); const sessionArtifactsDir = sessionFile ? sessionFile.slice(0, -6) : null; const artifactsDir = sessionArtifactsDir ?? path.join(os.tmpdir(), `omp-eval-agent-${Snowflake.next()}`); await fs.mkdir(artifactsDir, { recursive: true }); - - const shouldWriteConversationContext = session.settings.get("irc.enabled") !== true; - const compactContext = shouldWriteConversationContext ? session.getCompactContext?.() : undefined; - if (!compactContext) return { sessionFile, artifactsDir }; - - const contextFile = path.join(artifactsDir, "context.md"); - await Bun.write(contextFile, compactContext); - return { sessionFile, artifactsDir, contextFile }; + return { sessionFile, artifactsDir }; } function emitProgressStatus(emitStatus: ((event: JsStatusEvent) => void) | undefined, progress: AgentProgress): void { @@ -246,11 +236,10 @@ export async function runEvalAgent(args: unknown, options: EvalAgentBridgeOption }; const parentArtifactManager = options.session.getArtifactManager?.() ?? undefined; const mcpManager = options.session.mcpManager ?? MCPManager.instance(); - const { sessionFile, artifactsDir, contextFile } = await getArtifacts(options.session); + const { sessionFile, artifactsDir } = await getArtifacts(options.session); const outputManager = getOutputManager(options.session); const id = await outputManager.allocate(outputIdBase(parsed.label, agentName)); const assignment = parsed.prompt.trim(); - const context = trimToUndefined(parsed.context); // Suspend eval timeout accounting while the subagent owns control. The // timeout clock restarts once the bridge returns to the cell runtime. const result = await withBridgeTimeoutPause(options.emitStatus, () => @@ -259,7 +248,6 @@ export async function runEvalAgent(args: unknown, options: EvalAgentBridgeOption agent: effectiveAgent, task: renderSubagentPrompt(assignment), assignment, - context, description: trimToUndefined(parsed.label), index: 0, id, @@ -271,7 +259,6 @@ export async function runEvalAgent(args: unknown, options: EvalAgentBridgeOption sessionFile, persistArtifacts: Boolean(sessionFile), artifactsDir, - contextFile, // Eval `agent()` subagents are short-lived programmatic helpers (data // collection, structured output, parallel() fan-out). LSP server // cold-start costs tens of seconds and is pure overhead here, so it is diff --git a/packages/coding-agent/src/eval/js/shared/prelude.txt b/packages/coding-agent/src/eval/js/shared/prelude.txt index 36b61c5ab..095b219a2 100644 --- a/packages/coding-agent/src/eval/js/shared/prelude.txt +++ b/packages/coding-agent/src/eval/js/shared/prelude.txt @@ -65,7 +65,7 @@ if (!globalThis.__omp_js_prelude_loaded__) { }; const agent = async (prompt, opts, ...rest) => { - const o = optionsArg("agent", opts, rest, "{ agentType, model, context, label, schema }"); + const o = optionsArg("agent", opts, rest, "{ agentType, model, label, schema }"); const res = await globalThis.__omp_call_tool__("__agent__", { prompt, ...o }); const text = res && typeof res === "object" ? res.text : res; return hasOwn(o, "schema") ? JSON.parse(text) : text; diff --git a/packages/coding-agent/src/eval/py/prelude.py b/packages/coding-agent/src/eval/py/prelude.py index 0eb2ec942..e2ac422d0 100644 --- a/packages/coding-agent/src/eval/py/prelude.py +++ b/packages/coding-agent/src/eval/py/prelude.py @@ -519,21 +519,20 @@ if "__omp_prelude_loaded__" not in globals(): text = res.get("text") if isinstance(res, dict) else res return json.loads(text) if schema is not None else text - def agent(prompt, *, agent_type="task", model=None, context=None, label=None, schema=None): + def agent(prompt, *, agent_type="task", model=None, label=None, schema=None): """Run a subagent and return its final output. `agent_type` selects the subagent definition (default "task"). Pass - `model` to override that agent's model, `context` for shared background, - `label` for the output artifact id, and `schema` to request structured - JSON output; when `schema` is supplied the parsed object is returned. + `model` to override that agent's model, `label` for the output artifact + id, and `schema` to request structured JSON output; when `schema` is + supplied the parsed object is returned. Share background by writing a + local:// file and referencing it in the prompt. """ args = {"prompt": prompt} if agent_type is not None: args["agentType"] = agent_type if model is not None: args["model"] = model - if context is not None: - args["context"] = context if label is not None: args["label"] = label if schema is not None: diff --git a/packages/coding-agent/src/prompts/system/subagent-system-prompt.md b/packages/coding-agent/src/prompts/system/subagent-system-prompt.md index a7a25dad0..d6f4345cd 100644 --- a/packages/coding-agent/src/prompts/system/subagent-system-prompt.md +++ b/packages/coding-agent/src/prompts/system/subagent-system-prompt.md @@ -3,13 +3,6 @@ ROLE {{agent}} -{{#if context}} -CONTEXT -=================================== - -{{context}} -{{/if}} - {{#if planReference}} PLAN =================================== @@ -32,11 +25,6 @@ You are working in an isolated working tree at `{{worktree}}` for this sub-task. You NEVER modify files outside this tree or in the original repository. {{/if}} -{{#if contextFile}} -# Conversation Context -If you need additional information, your conversation with the user is in {{contextFile}} — `read` its tail or `search` it for relevant terms. -{{/if}} - {{#if ircPeers}} # IRC Peers You can reach other live agents via the `irc` tool. Your id is `{{ircSelfId}}`. Currently visible peers: diff --git a/packages/coding-agent/src/prompts/system/workflow-notice.md b/packages/coding-agent/src/prompts/system/workflow-notice.md index 73085ec6e..a8e1c6f55 100644 --- a/packages/coding-agent/src/prompts/system/workflow-notice.md +++ b/packages/coding-agent/src/prompts/system/workflow-notice.md @@ -13,8 +13,8 @@ Worth it when the task benefits from decomposition + parallel coverage, or from State persists across cells, so scout in one cell and fan out in the next. Every cell has: -- `agent(prompt, *, agent_type="task", model=None, context=None, label=None, schema=None)` — run ONE subagent; returns its final text, or the validated object when `schema` (a JSON Schema dict) is given. With `schema` the subagent is forced to emit structured output that is validated for you — branch on the object, not on parsed prose. `agent_type` picks a discovered agent ("explore", "reviewer", "oracle", …); `context` is shared background; `label` names the artifact. Subagents are told their final text IS the return value, so they hand back raw data. `agent()` blocks until the subagent finishes; eval-spawned agents nest at most 3 deep. -- `parallel(thunks)` — run zero-arg callables concurrently through a bounded pool, preserving input order; returns once all finish. The pool runs as wide as a `task` tool batch — don't hand-tune it; fan out as wide as the work divides. A thunk that raises propagates — wrap risky work in `try/except` inside the thunk to keep partial results. In a loop, bind each closure's value with a default arg (`lambda d=d: …`) or every thunk captures the last one. +- `agent(prompt, *, agent_type="task", model=None, label=None, schema=None)` — run ONE subagent; returns its final text, or the validated object when `schema` (a JSON Schema dict) is given. With `schema` the subagent is forced to emit structured output that is validated for you — branch on the object, not on parsed prose. `agent_type` picks a discovered agent ("explore", "reviewer", "oracle", …); `label` names the artifact. Shared background goes in a `local://` file referenced from each prompt, not a parameter. Subagents are told their final text IS the return value, so they hand back raw data. `agent()` blocks until the subagent finishes; eval-spawned agents nest at most 3 deep. +- `parallel(thunks)` — run zero-arg callables concurrently through a bounded pool, preserving input order; returns once all finish. The pool is bounded by the session's `task` concurrency — don't hand-tune it; fan out as wide as the work divides. A thunk that raises propagates — wrap risky work in `try/except` inside the thunk to keep partial results. In a loop, bind each closure's value with a default arg (`lambda d=d: …`) or every thunk captures the last one. - `pipeline(items, *stages)` — map items through `stages` left-to-right. There is a BARRIER between stages: ALL items clear stage N before stage N+1 begins. Each stage is a one-arg callable; stage 1 gets the original item, later stages get the previous result. Same pool width as `parallel()`. - `completion(prompt, *, model="default", system=None, schema=None)` — oneshot, stateless model call (no tools, no history). Tiers: "smol", "default", "slow". Cheap classification/scoring inside a fan-out. - `log(message)` — emit a progress line above the status tree. `phase(title)` — start a phase; the status lines that follow group under it. diff --git a/packages/coding-agent/src/prompts/tools/eval.md b/packages/coding-agent/src/prompts/tools/eval.md index 8e99ffc0f..52f25f101 100644 --- a/packages/coding-agent/src/prompts/tools/eval.md +++ b/packages/coding-agent/src/prompts/tools/eval.md @@ -46,9 +46,9 @@ tool.(args) → unknown Invoke any session tool by name. `args` is the tool's parameter object. completion(prompt, model?="default", system?=None, schema?=None) → str | dict Oneshot, stateless completion (no history, no tools). `model` picks a tier: "smol" (fast), "default" (this session's model), "slow" (most capable). Pass `system` for a system prompt. Pass a JSON-Schema `schema` to force structured output and get the parsed object back; otherwise returns the completion text. -{{#if spawns}}agent(prompt, agent_type?="task", model?=None, context?=None, label?=None, schema?=None) → str | dict - Run a subagent and return its final output. Defaults to the bundled "task" agent; pass `agent_type`/`agentType` for another discovered agent. Pass a JSON-Schema `schema` to force structured output and get the parsed object back. -{{#if js}} In JS, pass options as one trailing object — never positional: agent(prompt, { agentType, context, schema }). +{{#if spawns}}agent(prompt, agent_type?="task", model?=None, label?=None, schema?=None) → str | dict + Run a subagent and return its final output. Defaults to the bundled "task" agent; pass `agent_type`/`agentType` for another discovered agent. Pass a JSON-Schema `schema` to force structured output and get the parsed object back. Share background by writing a `local://` file and referencing it in the prompt. +{{#if js}} In JS, pass options as one trailing object — never positional: agent(prompt, { agentType, schema }). {{/if}} {{/if}} parallel(thunks) → list diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index def9a1629..022e34867 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/packages/coding-agent/src/session/agent-session.ts @@ -231,10 +231,8 @@ import type { AuthStorage } from "./auth-storage"; import type { ClientBridge, ClientBridgePermissionOption, ClientBridgePermissionOutcome } from "./client-bridge"; import { type BashExecutionMessage, - type CompactionSummaryMessage, type CustomMessage, convertToLlm, - type FileMentionMessage, type PythonExecutionMessage, readPendingDisplayTag, SILENT_ABORT_MARKER, @@ -9960,69 +9958,6 @@ export class AgentSession { }); } - /** - * Format the conversation as compact context for subagents. - * Includes only user messages and assistant text responses. - * Excludes: system prompt, tool definitions, tool calls/results, thinking blocks. - */ - formatCompactContext(): string { - const lines: string[] = []; - lines.push("# Conversation Context"); - lines.push(""); - lines.push( - "This is a summary of the parent conversation. Read this if you need additional context about what was discussed or decided.", - ); - lines.push(""); - - for (const msg of this.messages) { - if (msg.role === "user" || msg.role === "developer") { - lines.push(msg.role === "developer" ? "## Developer" : "## User"); - lines.push(""); - if (typeof msg.content === "string") { - lines.push(msg.content); - } else { - for (const c of msg.content) { - if (c.type === "text") { - lines.push(c.text); - } else if (c.type === "image") { - lines.push("[Image attached]"); - } - } - } - lines.push(""); - } else if (msg.role === "assistant") { - const assistantMsg = msg as AssistantMessage; - // Only include text content, skip tool calls and thinking - const textParts: string[] = []; - for (const c of assistantMsg.content) { - if (c.type === "text" && c.text.trim()) { - textParts.push(c.text); - } - } - if (textParts.length > 0) { - lines.push("## Assistant"); - lines.push(""); - lines.push(textParts.join("\n\n")); - lines.push(""); - } - } else if (msg.role === "fileMention") { - const fileMsg = msg as FileMentionMessage; - const paths = fileMsg.files.map(f => f.path).join(", "); - lines.push(`[Files referenced: ${paths}]`); - lines.push(""); - } else if (msg.role === "compactionSummary") { - const compactMsg = msg as CompactionSummaryMessage; - lines.push("## Earlier Context (Summarized)"); - lines.push(""); - lines.push(compactMsg.summary); - lines.push(""); - } - // Skip: toolResult, bashExecution, pythonExecution, branchSummary, custom, hookMessage - } - - return lines.join("\n").trim(); - } - // ========================================================================= // Extension System // ========================================================================= diff --git a/packages/coding-agent/src/tools/index.ts b/packages/coding-agent/src/tools/index.ts index c524ca646..c70a022bf 100644 --- a/packages/coding-agent/src/tools/index.ts +++ b/packages/coding-agent/src/tools/index.ts @@ -260,8 +260,6 @@ export interface ToolSession { recordEvalSubagentUsage?: (output: number) => void; /** Bridge to the connected client (e.g. ACP editor host). Tools should route fs/terminal/permission requests through this when available. */ getClientBridge?: () => ClientBridge | undefined; - /** Get compact conversation context for subagents (excludes tool results, system prompts) */ - getCompactContext?: () => string; /** Get cached todo phases for this session. */ getTodoPhases?: () => TodoPhase[]; /** Replace cached todo phases for this session. */ From 2b28816010a1cacd90d38396611004e225292543 Mon Sep 17 00:00:00 2001 From: can1357 Date: Wed, 10 Jun 2026 17:51:42 +0200 Subject: [PATCH 15/27] feat(coding-agent): replaced session observer with Agent Hub and inline compaction divider MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The session-observer overlay is gone. The Agent Hub (ctrl+s, alt+a, or double-tap left arrow on an empty editor) presents one overlay with two views: a live registry table (status, unread irc count, current task, last activity; j/k to navigate, r to revive, x to abort/release) and per-agent chat (transcript + input line) — submitting revives a parked agent and steers it via the normal prompt path. Main is the ambient chat and stays out of the table.\n\nrenderInitialMessages no longer takes a prebuilt context: every redraw now reaches for AgentSession.buildTranscriptSessionContext() (full-history transcript with each compaction emitted inline at the point it fired, snapcompact frames re-attached on rebuild). UiHelpers drops the deferred-compaction render and the IRC autoreply branch that the new mailbox bus deprecated. CompactionSummaryMessage renders as a slim divider (── 📷 compacted · ctrl+o ──), expanding to the summary + snapcompact frame count. session-manager.buildSessionContext gains a { transcript: true } mode; an exported buildSessionContextFromFile() reads a session file without taking the writer lock so the hub chat view can tail any agent (parked or live). Theme picks up icon.camera + tool.irc symbol entries. --- .../coding-agent/src/config/keybindings.ts | 7 +- packages/coding-agent/src/main.ts | 2 +- ...ssion-observer-overlay.ts => agent-hub.ts} | 956 +++++++++++------- .../components/compaction-summary-message.ts | 104 +- .../src/modes/components/custom-editor.ts | 10 + .../controllers/extension-ui-controller.ts | 16 +- .../src/modes/controllers/input-controller.ts | 20 +- .../modes/controllers/selector-controller.ts | 38 +- .../src/modes/interactive-mode.ts | 21 +- .../coding-agent/src/modes/theme/theme.ts | 15 +- packages/coding-agent/src/modes/types.ts | 8 +- .../src/modes/utils/hotkeys-markdown.ts | 1 + .../src/modes/utils/ui-helpers.ts | 49 +- .../coding-agent/src/session/agent-session.ts | 10 + .../src/session/session-manager.ts | 64 +- .../test/interactive-mode-status.test.ts | 3 +- .../components/compaction-divider.test.ts | 63 ++ .../render-initial-messages-dedupe.test.ts | 160 --- .../utils/render-initial-messages.test.ts | 101 ++ .../test/silent-abort-overlay-render.test.ts | 54 +- .../test/tools/search-path-lists.test.ts | 29 +- 21 files changed, 1071 insertions(+), 660 deletions(-) rename packages/coding-agent/src/modes/components/{session-observer-overlay.ts => agent-hub.ts} (54%) create mode 100644 packages/coding-agent/test/modes/components/compaction-divider.test.ts delete mode 100644 packages/coding-agent/test/modes/utils/render-initial-messages-dedupe.test.ts create mode 100644 packages/coding-agent/test/modes/utils/render-initial-messages.test.ts diff --git a/packages/coding-agent/src/config/keybindings.ts b/packages/coding-agent/src/config/keybindings.ts index b652586b3..e48c03e0a 100644 --- a/packages/coding-agent/src/config/keybindings.ts +++ b/packages/coding-agent/src/config/keybindings.ts @@ -36,6 +36,7 @@ interface AppKeybindings { "app.clipboard.pasteTextRaw": true; "app.clipboard.copyLine": true; "app.clipboard.copyPrompt": true; + "app.agents.hub": true; "app.session.new": true; "app.session.tree": true; "app.session.fork": true; @@ -166,9 +167,13 @@ export const KEYBINDINGS = { defaultKeys: [], description: "Resume session", }, + "app.agents.hub": { + defaultKeys: "alt+a", + description: "Open the agent hub", + }, "app.session.observe": { defaultKeys: "ctrl+s", - description: "Observe subagent sessions", + description: "Open the agent hub", }, "app.session.togglePath": { defaultKeys: "ctrl+p", diff --git a/packages/coding-agent/src/main.ts b/packages/coding-agent/src/main.ts index b4b17195b..b49f888c3 100644 --- a/packages/coding-agent/src/main.ts +++ b/packages/coding-agent/src/main.ts @@ -431,7 +431,7 @@ async function runInteractiveMode( // Every in-process session load also uses `clearTerminalHistory`; cold launch // follows the same clean-cutover path instead of preserving a previous run's // transcript above the fresh one. - mode.renderInitialMessages(undefined, { preserveExistingChat: true, clearTerminalHistory: true }); + mode.renderInitialMessages({ preserveExistingChat: true, clearTerminalHistory: true }); for (const notify of notifs) { if (!notify) { diff --git a/packages/coding-agent/src/modes/components/session-observer-overlay.ts b/packages/coding-agent/src/modes/components/agent-hub.ts similarity index 54% rename from packages/coding-agent/src/modes/components/session-observer-overlay.ts rename to packages/coding-agent/src/modes/components/agent-hub.ts index 73ab1d185..ebff0ae7f 100644 --- a/packages/coding-agent/src/modes/components/session-observer-overlay.ts +++ b/packages/coding-agent/src/modes/components/agent-hub.ts @@ -1,30 +1,34 @@ /** - * Session observer overlay component. + * Agent Hub overlay component. * - * Picker mode: lists main + active subagent sessions with live status. - * Viewer mode: renders a scrollable, interactive transcript of the selected subagent's session - * by reading its JSONL session file — shows thinking, text, tool calls, results - * with expand/collapse per entry and breadcrumb navigation for nested sub-agents. + * One overlay, two views: + * - Table view: every registered agent except Main (Main IS the ambient + * chat), live from the global AgentRegistry — status, unread irc count, + * current/last task, last activity. Select with j/k, Enter opens a chat, + * `r` revives a parked agent, `x` aborts + releases one. + * - Chat view: per-agent transcript (incremental session-file tail, absorbed + * from the old session observer overlay) plus an input line. Submitting + * revives a parked agent, then prompts/steers it; the message lands in the + * agent's persisted history via the normal prompt path. * - * Lifecycle: - * - shortcut opens picker - * - Enter on a subagent -> viewer - * - shortcut while in viewer -> back to picker (or pop breadcrumb) - * - Esc from viewer -> back to picker (or pop breadcrumb) - * - Esc from picker -> close overlay - * - Enter on main session -> close overlay (jump back) + * Replaces the old SessionObserverOverlayComponent (ctrl+s observer). */ +import * as fs from "node:fs"; import type { ToolResultMessage } from "@oh-my-pi/pi-ai"; -import { Container, Markdown, type MarkdownTheme, matchesKey, ScrollView } from "@oh-my-pi/pi-tui"; -import { formatDuration, formatNumber, logger } from "@oh-my-pi/pi-utils"; +import { Container, Editor, Markdown, type MarkdownTheme, matchesKey, ScrollView } from "@oh-my-pi/pi-tui"; +import { formatAge, formatDuration, formatNumber, logger } from "@oh-my-pi/pi-utils"; import type { KeyId } from "../../config/keybindings"; -import { isSilentAbort } from "../../session/messages"; +import { IrcBus } from "../../irc/bus"; +import { AgentLifecycleManager } from "../../registry/agent-lifecycle"; +import { type AgentRef, AgentRegistry, type AgentStatus, MAIN_AGENT_ID } from "../../registry/agent-registry"; +import type { AgentSession } from "../../session/agent-session"; +import { isSilentAbort, USER_INTERRUPT_LABEL } from "../../session/messages"; import type { SessionMessageEntry } from "../../session/session-manager"; import { parseSessionEntries } from "../../session/session-manager"; import { PREVIEW_LIMITS, replaceTabs, TRUNCATE_LENGTHS, truncateToWidth } from "../../tools/render-utils"; import { toPathList } from "../../tools/search"; import type { ObservableSession, SessionObserverRegistry } from "../session-observer-registry"; -import { getMarkdownTheme, theme } from "../theme/theme"; +import { getEditorTheme, getMarkdownTheme, theme } from "../theme/theme"; import { matchesSelectDown, matchesSelectUp } from "../utils/keybinding-matchers"; import { DynamicBorder } from "./dynamic-border"; import { formatContextUsage } from "./status-line/context-thresholds"; @@ -39,6 +43,10 @@ const MAX_TOOL_ARGS_CHARS = 500; const PAGE_SIZE = 15; /** Left indent for content under entry headers */ const INDENT = " "; +/** Refresh cadence for the relative-time column */ +const AGE_TICK_MS = 5_000; +/** Debounce for live-session transcript refreshes */ +const CHAT_REFRESH_DEBOUNCE_MS = 80; /** Compute the max content width for the current terminal, accounting for indent and chrome. */ function contentWidth(indent = INDENT): number { @@ -57,180 +65,391 @@ interface ViewerEntry { kind: "thinking" | "text" | "toolCall" | "user"; } -/** Breadcrumb item for nested session navigation */ -interface BreadcrumbItem { - sessionId: string; - label: string; - sessionFile: string; +const STATUS_ORDER: Record = { running: 0, idle: 1, parked: 2, aborted: 3 }; + +/** Glyph + status word, colored per theme status conventions. */ +function statusBadge(status: AgentStatus): string { + switch (status) { + case "running": + return theme.fg("accent", `${theme.status.running} running`); + case "idle": + return theme.fg("success", `${theme.status.enabled} idle`); + case "parked": + return theme.fg("muted", `${theme.status.shadowed} parked`); + case "aborted": + return theme.fg("error", `${theme.status.aborted} aborted`); + } } -export class SessionObserverOverlayComponent extends Container { - #registry: SessionObserverRegistry; - #onDone: () => void; - #selectedSessionId?: string; - #observeKeys: KeyId[]; - #transcriptCache?: { path: string; bytesRead: number; entries: SessionMessageEntry[]; model?: string }; +export interface AgentHubDeps { + /** Progress/status snapshot source (task lifecycle + progress channels). */ + observers: SessionObserverRegistry; + /** Keys that toggle the hub closed from inside (app.agents.hub + app.session.observe). */ + hubKeys: KeyId[]; + onDone: () => void; + requestRender: () => void; + /** Injectable for tests; defaults to the process-global registry. */ + registry?: AgentRegistry; + /** Injectable for tests; defaults to the process-global lifecycle manager. */ + lifecycle?: AgentLifecycleManager; + /** Injectable for tests; defaults to the process-global bus. */ + irc?: IrcBus; +} - // Scroll state +export class AgentHubOverlayComponent extends Container { + #registry: AgentRegistry; + #observers: SessionObserverRegistry; + #irc: IrcBus; + #lifecycle: () => AgentLifecycleManager; + #onDone: () => void; + #requestRender: () => void; + #hubKeys: KeyId[]; + #unsubscribers: Array<() => void> = []; + #ageTimer: NodeJS.Timeout | undefined; + + // Table state + #view: "table" | "chat" = "table"; + #rows: AgentRef[] = []; + #selectedRow = 0; + #notice: string | undefined; + + // Chat state + #chatAgentId: string | undefined; + #editor: Editor; + #sessionUnsubscribe: (() => void) | undefined; + #attachedSession: AgentSession | undefined; + #chatRefreshTimer: NodeJS.Timeout | undefined; + #transcriptCache: { path: string; bytesRead: number; entries: SessionMessageEntry[]; model?: string } | undefined; + + // Transcript viewer state (absorbed from the session observer overlay) #scrollOffset = 0; #renderedLines: string[] = []; #viewportHeight = 20; #wasAtBottom = true; - - // Entry selection & expand/collapse #viewerEntries: ViewerEntry[] = []; #selectedEntryIndex = 0; #expandedEntries = new Set(); - - // Breadcrumb navigation - #navigationStack: BreadcrumbItem[] = []; - - // Cached header/footer for viewer (rebuilt on refresh) #viewerHeaderLines: string[] = []; - #viewerFooterLines: string[] = []; - // Markdown rendering #mdTheme: MarkdownTheme = getMarkdownTheme(); - constructor(registry: SessionObserverRegistry, onDone: () => void, observeKeys: KeyId[]) { + constructor(deps: AgentHubDeps) { super(); - this.#registry = registry; - this.#onDone = onDone; - this.#observeKeys = observeKeys; + this.#registry = deps.registry ?? AgentRegistry.global(); + this.#observers = deps.observers; + this.#irc = deps.irc ?? IrcBus.global(); + // Lazy: the lifecycle global self-constructs against the global + // registry, so only touch it when revive/kill actually needs it. + this.#lifecycle = () => deps.lifecycle ?? AgentLifecycleManager.global(); + this.#onDone = deps.onDone; + this.#requestRender = deps.requestRender; + this.#hubKeys = deps.hubKeys; - // Jump directly to the most recently active sub-agent - const mostRecent = this.#getMostRecentSubagent(); - if (mostRecent) { - this.#selectedSessionId = mostRecent.id; - this.#setupViewer(); - } else { - // No sub-agents — close immediately - queueMicrotask(() => this.#onDone()); - } + this.#editor = new Editor(getEditorTheme()); + this.#editor.setMaxHeight(4); + this.#editor.onSubmit = text => this.#submitChatMessage(text); + + this.#unsubscribers.push(this.#registry.onChange(() => this.#onDataChange())); + this.#unsubscribers.push(this.#observers.onChange(() => this.#onDataChange())); + this.#ageTimer = setInterval(() => this.#requestRender(), AGE_TICK_MS); + this.#ageTimer.unref?.(); + + this.#refreshRows(); } - /** Find the most recently updated sub-agent session (prefer active ones) */ - #getMostRecentSubagent(): ObservableSession | undefined { - const sessions = this.#registry.getSessions().filter(s => s.kind === "subagent"); - if (sessions.length === 0) return undefined; - // Prefer active sessions, then sort by lastUpdate descending - const active = sessions.filter(s => s.status === "active"); - const pool = active.length > 0 ? active : sessions; - return pool.sort((a, b) => b.lastUpdate - a.lastUpdate)[0]; + /** Tear down every subscription and timer. Called by the overlay owner on close. */ + dispose(): void { + for (const unsubscribe of this.#unsubscribers.splice(0)) unsubscribe(); + if (this.#ageTimer) { + clearInterval(this.#ageTimer); + this.#ageTimer = undefined; + } + if (this.#chatRefreshTimer) { + clearTimeout(this.#chatRefreshTimer); + this.#chatRefreshTimer = undefined; + } + this.#detachLiveSession(); } override render(width: number): readonly string[] { - return this.#renderViewer(width); + return this.#view === "table" ? this.#renderTable(width) : this.#renderChat(width); } - #setupViewer(): void { - this.clear(); + handleInput(keyData: string): void { + // The hub/observe keys always close the overlay (toggle semantics) + for (const key of this.#hubKeys) { + if (matchesKey(keyData, key)) { + this.#onDone(); + return; + } + } + if (this.#view === "table") { + this.#handleTableInput(keyData); + } else { + this.#handleChatInput(keyData); + } + } + + /** Open the chat view for an agent id (public for table Enter and tests). */ + openChat(id: string): void { + if (!this.#registry.get(id)) return; + this.#view = "chat"; + this.#chatAgentId = id; + this.#notice = undefined; + this.#transcriptCache = undefined; this.#scrollOffset = 0; this.#selectedEntryIndex = 0; this.#expandedEntries.clear(); this.#wasAtBottom = true; - this.#rebuildViewerContent(); - // Auto-scroll to bottom and select last entry on init + this.#editor.setText(""); + this.#attachLiveSession(); + this.#rebuildChatContent(); + // Auto-scroll to bottom and select last entry on open if (this.#viewerEntries.length > 0) { this.#selectedEntryIndex = this.#viewerEntries.length - 1; - this.#wasAtBottom = true; - this.#rebuildViewerContent(); + this.#rebuildChatContent(); } + this.#requestRender(); } - /** Rebuild content from live registry data */ - refreshFromRegistry(): void { - if (this.#selectedSessionId) { - // Keep auto-scrolling to bottom unless the user navigated away from the last entry + // ======================================================================== + // Live data plumbing + // ======================================================================== + + #onDataChange(): void { + this.#refreshRows(); + if (this.#view === "chat") { + // A revive/park swaps the live session out from under the chat view. + this.#attachLiveSession(); + this.#scheduleChatRefresh(); + return; + } + this.#requestRender(); + } + + #refreshRows(): void { + const selectedId = this.#rows[this.#selectedRow]?.id; + this.#rows = this.#registry + .list() + .filter(ref => ref.id !== MAIN_AGENT_ID) + .sort((a, b) => STATUS_ORDER[a.status] - STATUS_ORDER[b.status] || b.lastActivity - a.lastActivity); + const keptIndex = selectedId ? this.#rows.findIndex(ref => ref.id === selectedId) : -1; + this.#selectedRow = keptIndex >= 0 ? keptIndex : Math.min(this.#selectedRow, Math.max(0, this.#rows.length - 1)); + } + + /** Subscribe to the chat agent's live session (if any) for transcript refreshes. Idempotent per session. */ + #attachLiveSession(): void { + const session = this.#chatAgentId ? (this.#registry.get(this.#chatAgentId)?.session ?? undefined) : undefined; + if (session === this.#attachedSession) return; + this.#detachLiveSession(); + if (!session) return; + this.#attachedSession = session; + this.#sessionUnsubscribe = session.subscribe(event => { + if (event.type === "message_end" || event.type === "tool_execution_end" || event.type === "agent_end") { + this.#scheduleChatRefresh(); + } + }); + } + + #detachLiveSession(): void { + this.#sessionUnsubscribe?.(); + this.#sessionUnsubscribe = undefined; + this.#attachedSession = undefined; + } + + #scheduleChatRefresh(): void { + if (this.#chatRefreshTimer) return; + this.#chatRefreshTimer = setTimeout(() => { + this.#chatRefreshTimer = undefined; + if (this.#view !== "chat") return; + // Keep auto-scrolling to bottom unless the user navigated away this.#wasAtBottom = this.#selectedEntryIndex >= this.#viewerEntries.length - 1; - this.#rebuildViewerContent(); - } + this.#rebuildChatContent(); + if (this.#wasAtBottom && this.#viewerEntries.length > 0) { + this.#selectedEntryIndex = this.#viewerEntries.length - 1; + } + this.#requestRender(); + }, CHAT_REFRESH_DEBOUNCE_MS); + this.#chatRefreshTimer.unref?.(); } - /** Rebuild the transcript content lines (called on setup and refresh) */ - #rebuildViewerContent(): void { - const sessions = this.#registry.getSessions(); - const session = sessions.find(s => s.id === this.#selectedSessionId); + #observableFor(id: string): ObservableSession | undefined { + return this.#observers.getSessions().find(s => s.id === id); + } - // Load transcript first so model info is available for header - let messageEntries: SessionMessageEntry[] | null = null; - if (session?.sessionFile) { - messageEntries = this.#loadTranscript(session.sessionFile); - } + // ======================================================================== + // Table view + // ======================================================================== - // Header - this.#viewerHeaderLines = []; - const breadcrumb = this.#buildBreadcrumb(session); - this.#viewerHeaderLines.push(theme.fg("accent", breadcrumb)); - if (session) { - const statusColor = session.status === "active" ? "success" : session.status === "failed" ? "error" : "dim"; - const statusText = theme.fg(statusColor, `[${session.status}]`); - const agentTag = session.agent ? theme.fg("dim", ` ${session.agent}`) : ""; - const subagentIds = this.#getSubagentSessionIds(); - const posIdx = subagentIds.indexOf(this.#selectedSessionId ?? ""); - const posLabel = - subagentIds.length > 1 && posIdx >= 0 ? theme.fg("dim", ` (${posIdx + 1}/${subagentIds.length})`) : ""; - const modelName = this.#transcriptCache?.model; - const modelLabel = modelName ? theme.fg("muted", ` · ${modelName}`) : ""; - this.#viewerHeaderLines.push(`${theme.bold(session.label)} ${statusText}${agentTag}${posLabel}${modelLabel}`); - } + #renderTable(width: number): string[] { + const lines: string[] = []; + lines.push(...new DynamicBorder().render(width)); + const counts = this.#statusSummary(); + lines.push(` ${theme.fg("accent", "Agent Hub")}${counts ? theme.fg("dim", `${theme.sep.dot}${counts}`) : ""}`); + lines.push(...new DynamicBorder().render(width)); - // Content - const contentLines: string[] = []; - this.#viewerEntries = []; - - if (!session) { - contentLines.push(theme.fg("dim", "Session no longer available.")); - } else if (!session.sessionFile) { - contentLines.push(theme.fg("dim", "No session file available yet.")); - } else if (!messageEntries) { - contentLines.push(theme.fg("dim", "Unable to read session file.")); - } else if (messageEntries.length === 0) { - contentLines.push(theme.fg("dim", "No messages yet.")); + if (this.#rows.length === 0) { + lines.push(` ${theme.fg("dim", "no subagents yet — task spawns appear here")}`); } else { - this.#buildTranscriptLines(messageEntries, contentLines); + const termHeight = process.stdout.rows || 40; + // Chrome: 2 borders + title + notice? + blank + hints + border + const maxVisible = Math.max(3, termHeight - 7 - (this.#notice ? 1 : 0)); + let start = 0; + if (this.#rows.length > maxVisible) { + start = Math.min( + Math.max(0, this.#selectedRow - Math.floor(maxVisible / 2)), + this.#rows.length - maxVisible, + ); + } + const end = Math.min(start + maxVisible, this.#rows.length); + for (let i = start; i < end; i++) { + lines.push(this.#renderRow(this.#rows[i], i === this.#selectedRow, width)); + } + if (end < this.#rows.length) { + lines.push(` ${theme.fg("dim", `… ${this.#rows.length - end} more`)}`); + } } - this.#renderedLines = contentLines; - // Footer - this.#viewerFooterLines = []; - const statsLine = this.#buildStatsLine(session); - if (statsLine) this.#viewerFooterLines.push(statsLine); - this.#viewerFooterLines.push( - theme.fg("dim", "j/k:scroll Enter:expand [/]/←→:cycle agents Esc/Ctrl+S:close g/G:top/bottom"), - ); + if (this.#notice) { + lines.push(` ${theme.fg("error", sanitizeLine(this.#notice, Math.max(10, width - 2)))}`); + } + lines.push(""); + lines.push(` ${theme.fg("dim", "j/k:select Enter:chat r:revive x:kill Esc:close")}`); + lines.push(...new DynamicBorder().render(width)); + return lines; + } - // Auto-scroll to bottom if we were at bottom - if (this.#wasAtBottom) { - this.#scrollOffset = Math.max(0, contentLines.length - this.#viewportHeight); + #statusSummary(): string { + const counts: Record = { running: 0, idle: 0, parked: 0, aborted: 0 }; + for (const ref of this.#rows) { + counts[ref.status]++; + } + const parts: string[] = []; + for (const status of ["running", "idle", "parked", "aborted"] as const) { + const count = counts[status]; + if (count > 0) parts.push(`${count} ${status}`); + } + return parts.join(theme.sep.dot); + } + + #renderRow(ref: AgentRef, selected: boolean, width: number): string { + const cursor = selected ? theme.fg("accent", theme.nav.cursor) : " "; + const parts: string[] = [statusBadge(ref.status), theme.bold(replaceTabs(ref.id))]; + parts.push(theme.fg("dim", ref.parentId ? `${ref.kind} · of ${ref.parentId}` : ref.kind)); + const observed = this.#observableFor(ref.id); + const task = observed?.description ?? observed?.progress?.task; + if (task) { + parts.push(theme.fg("muted", sanitizeLine(task, TRUNCATE_LENGTHS.TITLE))); + } + const unread = this.#irc.unreadCount(ref.id); + if (unread > 0) { + parts.push(theme.fg("warning", `⧉ ${unread}`)); + } + parts.push(theme.fg("dim", formatAge(Math.max(1, Math.round((Date.now() - ref.lastActivity) / 1000))))); + return truncateToWidth(` ${cursor} ${parts.join(theme.sep.dot)}`, Math.max(10, width - 1)); + } + + #handleTableInput(keyData: string): void { + if (matchesKey(keyData, "escape")) { + this.#onDone(); + return; + } + if (keyData === "j" || matchesSelectDown(keyData)) { + if (this.#rows.length > 0) { + this.#selectedRow = Math.min(this.#selectedRow + 1, this.#rows.length - 1); + } + this.#requestRender(); + return; + } + if (keyData === "k" || matchesSelectUp(keyData)) { + if (this.#rows.length > 0) { + this.#selectedRow = Math.max(this.#selectedRow - 1, 0); + } + this.#requestRender(); + return; + } + if (matchesKey(keyData, "enter") || keyData === "\r" || keyData === "\n") { + const selected = this.#rows[this.#selectedRow]; + if (selected) this.openChat(selected.id); + return; + } + if (keyData === "r") { + this.#reviveSelected(); + return; + } + if (keyData === "x") { + this.#killSelected(); + return; } } - /** Produce the final viewer output for the overlay system */ - #renderViewer(width: number): string[] { - const termHeight = process.stdout.rows || 40; + #reviveSelected(): void { + const ref = this.#rows[this.#selectedRow]; + if (!ref) return; + if (ref.status !== "parked") { + this.#notice = `Agent "${ref.id}" is ${ref.status} — only parked agents can be revived.`; + this.#requestRender(); + return; + } + this.#notice = undefined; + // Fire-and-forget; failures surface as an inline notice + this.#lifecycle() + .ensureLive(ref.id) + .catch((error: unknown) => { + this.#notice = error instanceof Error ? error.message : String(error); + this.#requestRender(); + }); + this.#requestRender(); + } - // Compute viewport: total height minus header chrome and footer chrome - // Header: border(1) + headerLines + border(1) = headerLines.length + 2 - // Footer: spacer(1) + scrollInfo(1) + footerLines + border(1) = footerLines.length + 2 + #killSelected(): void { + const ref = this.#rows[this.#selectedRow]; + if (!ref) return; + this.#notice = undefined; + void (async () => { + try { + if (ref.status === "running" && ref.session) { + await ref.session.abort({ reason: USER_INTERRUPT_LABEL }); + } + await this.#lifecycle().release(ref.id); + } catch (error) { + logger.warn("Agent hub: kill failed", { id: ref.id, error: String(error) }); + this.#notice = error instanceof Error ? error.message : String(error); + } + this.#refreshRows(); + this.#requestRender(); + })(); + } + + // ======================================================================== + // Chat view + // ======================================================================== + + #renderChat(width: number): string[] { + const termHeight = process.stdout.rows || 40; + const innerWidth = Math.max(20, width - 2); + const editorLines = this.#editor.render(innerWidth); + const noticeLine = this.#notice + ? ` ${theme.fg("error", sanitizeLine(this.#notice, Math.max(10, width - 2)))}` + : undefined; + const footerLines = this.#buildChatFooterLines(); + + // Header: border + headerLines + border; footer: notice? + editor + footer + border const headerChrome = this.#viewerHeaderLines.length + 2; - const footerChrome = this.#viewerFooterLines.length + 2; + const footerChrome = editorLines.length + footerLines.length + (noticeLine ? 1 : 0) + 1; this.#viewportHeight = Math.max(5, termHeight - headerChrome - footerChrome); - // Clamp scroll offset const maxScroll = Math.max(0, this.#renderedLines.length - this.#viewportHeight); + if (this.#wasAtBottom) this.#scrollOffset = maxScroll; this.#scrollOffset = Math.max(0, Math.min(this.#scrollOffset, maxScroll)); const lines: string[] = []; - - // --- Header --- lines.push(...new DynamicBorder().render(width)); - for (const hl of this.#viewerHeaderLines) { - lines.push(` ${hl}`); + for (const headerLine of this.#viewerHeaderLines) { + lines.push(` ${headerLine}`); } lines.push(...new DynamicBorder().render(width)); - // --- Scrolled content viewport --- - const sv = new ScrollView( + const scrollView = new ScrollView( this.#renderedLines.slice(this.#scrollOffset, this.#scrollOffset + this.#viewportHeight), { height: this.#viewportHeight, @@ -239,31 +458,27 @@ export class SessionObserverOverlayComponent extends Container { theme: { track: t => theme.fg("dim", t), thumb: t => theme.fg("accent", t) }, }, ); - sv.setScrollOffset(this.#scrollOffset); - for (const row of sv.render(Math.max(1, width - 1))) lines.push(` ${row}`); + scrollView.setScrollOffset(this.#scrollOffset); + for (const row of scrollView.render(Math.max(1, width - 1))) lines.push(` ${row}`); - // --- Footer --- - lines.push(""); - lines.push(` ${this.#viewerFooterLines[0] ?? ""}`); - for (let i = 1; i < this.#viewerFooterLines.length; i++) { - lines.push(` ${this.#viewerFooterLines[i]}`); - } + if (noticeLine) lines.push(noticeLine); + for (const editorLine of editorLines) lines.push(` ${editorLine}`); + lines.push(...footerLines); lines.push(...new DynamicBorder().render(width)); - return lines; } - #buildBreadcrumb(session: ObservableSession | undefined): string { - const parts: string[] = ["Session Observer"]; - for (const item of this.#navigationStack) { - parts.push(item.label); - } - if (session) parts.push(session.label); - return parts.join(" > "); + #buildChatFooterLines(): string[] { + const lines: string[] = []; + const observed = this.#chatAgentId ? this.#observableFor(this.#chatAgentId) : undefined; + const statsLine = this.#buildStatsLine(observed); + if (statsLine) lines.push(` ${statsLine}`); + lines.push(` ${theme.fg("dim", "Enter:send Esc:back empty input: j/k:scroll Enter:expand g/G:top/bottom")}`); + return lines; } - #buildStatsLine(session: ObservableSession | undefined): string { - const progress = session?.progress; + #buildStatsLine(observed: ObservableSession | undefined): string { + const progress = observed?.progress; if (!progress) return ""; const stats: string[] = []; // Current per-turn context — match the status line's `%/` gauge (e.g. `5.1%/1M`). @@ -290,6 +505,212 @@ export class SessionObserverOverlayComponent extends Container { return parts.join(theme.sep.dot); } + /** Rebuild the chat header + transcript content lines */ + #rebuildChatContent(): void { + const id = this.#chatAgentId; + const ref = id ? this.#registry.get(id) : undefined; + + // Load transcript first so model info is available for the header + let messageEntries: SessionMessageEntry[] | null = null; + if (ref?.sessionFile) { + messageEntries = this.#loadTranscript(ref.sessionFile); + } + + this.#viewerHeaderLines = []; + this.#viewerHeaderLines.push(theme.fg("accent", `Agent Hub > ${id ?? "?"}`)); + if (ref) { + const observed = this.#observableFor(ref.id); + const model = observed?.progress?.resolvedModel ?? this.#transcriptCache?.model; + const kindTag = theme.fg("dim", ` ${ref.parentId ? `${ref.kind} · of ${ref.parentId}` : ref.kind}`); + const modelLabel = model ? theme.fg("muted", `${theme.sep.dot}${model}`) : ""; + this.#viewerHeaderLines.push(`${theme.bold(ref.id)} ${statusBadge(ref.status)}${kindTag}${modelLabel}`); + } + + const contentLines: string[] = []; + this.#viewerEntries = []; + if (!ref) { + contentLines.push(theme.fg("dim", "Agent no longer registered.")); + } else if (!ref.sessionFile) { + contentLines.push(theme.fg("dim", "No session file available yet.")); + } else if (!messageEntries) { + contentLines.push(theme.fg("dim", "Unable to read session file.")); + } else if (messageEntries.length === 0) { + contentLines.push(theme.fg("dim", "No messages yet.")); + } else { + this.#buildTranscriptLines(messageEntries, contentLines); + } + this.#renderedLines = contentLines; + } + + #handleChatInput(keyData: string): void { + const editorEmpty = this.#editor.getText().trim() === ""; + + if (matchesKey(keyData, "escape")) { + if (!editorEmpty) { + this.#editor.setText(""); + this.#requestRender(); + return; + } + this.#closeChat(); + return; + } + + // Navigation mirrors the old observer overlay while the input is empty; + // once the user starts typing, the editor owns every key. + if (editorEmpty && this.#handleViewerNavigation(keyData)) { + return; + } + + this.#editor.handleInput(keyData); + this.#requestRender(); + } + + #closeChat(): void { + this.#view = "table"; + this.#chatAgentId = undefined; + this.#notice = undefined; + this.#detachLiveSession(); + this.#refreshRows(); + this.#requestRender(); + } + + #submitChatMessage(text: string): void { + const id = this.#chatAgentId; + const trimmed = text.trim(); + if (!id || !trimmed) return; + this.#editor.setText(""); + this.#notice = undefined; + void (async () => { + try { + // Revives a parked agent; returns the live session for running/idle. + const session = await this.#lifecycle().ensureLive(id); + this.#attachLiveSession(); + // Steers a mid-turn agent; sends a normal prompt to an idle one. + await session.prompt(trimmed, { streamingBehavior: "steer" }); + } catch (error) { + this.#notice = error instanceof Error ? error.message : String(error); + } + this.#scheduleChatRefresh(); + this.#requestRender(); + })(); + this.#requestRender(); + } + + /** Viewer navigation (selection, paging, expand) for the chat transcript. Returns true when handled. */ + #handleViewerNavigation(keyData: string): boolean { + const entryCount = this.#viewerEntries.length; + + if (keyData === "j" || matchesSelectDown(keyData)) { + if (entryCount > 0) { + this.#selectedEntryIndex = Math.min(this.#selectedEntryIndex + 1, entryCount - 1); + } + this.#rebuildAndScroll(); + return true; + } + if (keyData === "k" || matchesSelectUp(keyData)) { + if (entryCount > 0) { + this.#selectedEntryIndex = Math.max(this.#selectedEntryIndex - 1, 0); + } + this.#rebuildAndScroll(); + return true; + } + if (matchesKey(keyData, "pageDown")) { + if (entryCount > 0) { + const prevIndex = this.#selectedEntryIndex; + this.#selectedEntryIndex = Math.min(this.#selectedEntryIndex + 5, entryCount - 1); + if (this.#selectedEntryIndex === prevIndex) { + this.#scrollOffset = Math.min( + this.#scrollOffset + PAGE_SIZE, + Math.max(0, this.#renderedLines.length - this.#viewportHeight), + ); + } + } else { + this.#scrollOffset = Math.min( + this.#scrollOffset + PAGE_SIZE, + Math.max(0, this.#renderedLines.length - this.#viewportHeight), + ); + } + this.#rebuildAndScroll(); + return true; + } + if (matchesKey(keyData, "pageUp")) { + if (entryCount > 0) { + const prevIndex = this.#selectedEntryIndex; + this.#selectedEntryIndex = Math.max(this.#selectedEntryIndex - 5, 0); + if (this.#selectedEntryIndex === prevIndex) { + this.#scrollOffset = Math.max(this.#scrollOffset - PAGE_SIZE, 0); + } + } else { + this.#scrollOffset = Math.max(this.#scrollOffset - PAGE_SIZE, 0); + } + this.#rebuildAndScroll(); + return true; + } + if (matchesKey(keyData, "enter") || keyData === "\r" || keyData === "\n") { + if (entryCount > 0 && this.#selectedEntryIndex < entryCount) { + if (this.#expandedEntries.has(this.#selectedEntryIndex)) { + this.#expandedEntries.delete(this.#selectedEntryIndex); + } else { + this.#expandedEntries.add(this.#selectedEntryIndex); + } + this.#rebuildAndScroll(); + } + return true; + } + if (keyData === "G") { + if (entryCount > 0) this.#selectedEntryIndex = entryCount - 1; + this.#scrollOffset = Math.max(0, this.#renderedLines.length - this.#viewportHeight); + this.#rebuildAndScroll(); + return true; + } + if (keyData === "g") { + this.#selectedEntryIndex = 0; + this.#scrollOffset = 0; + this.#rebuildAndScroll(); + return true; + } + return false; + } + + /** Rebuild transcript lines (which depend on selectedEntryIndex/expandedEntries) and scroll to selection */ + #rebuildAndScroll(): void { + // Resume auto-scrolling once selection returns to the last entry + this.#wasAtBottom = this.#selectedEntryIndex >= this.#viewerEntries.length - 1; + this.#rebuildChatContent(); + this.#scrollToSelectedEntry(); + this.#requestRender(); + } + + #scrollToSelectedEntry(): void { + if (this.#viewerEntries.length === 0) return; + const entry = this.#viewerEntries[this.#selectedEntryIndex]; + if (!entry) return; + + const entryTop = entry.lineStart; + const entryBottom = entry.lineStart + entry.lineCount; + + if (entry.lineCount >= this.#viewportHeight) { + // Entry taller than viewport: only snap when it's completely out of view. + if (this.#scrollOffset + this.#viewportHeight <= entryTop) { + this.#scrollOffset = Math.max(0, entryTop - 1); + } else if (this.#scrollOffset >= entryBottom) { + this.#scrollOffset = Math.max(0, entryBottom - this.#viewportHeight); + } + } else { + // Entry fits in viewport: ensure it's fully visible + if (entryTop < this.#scrollOffset) { + this.#scrollOffset = Math.max(0, entryTop - 1); + } + if (entryBottom > this.#scrollOffset + this.#viewportHeight) { + this.#scrollOffset = Math.max(0, entryBottom - this.#viewportHeight + 1); + } + } + } + + // ======================================================================== + // Transcript rendering (absorbed from the session observer overlay) + // ======================================================================== + #buildTranscriptLines(messageEntries: SessionMessageEntry[], lines: string[]): void { // Build a tool call ID -> tool result map const toolResults = new Map(); @@ -308,7 +729,7 @@ export class SessionObserverOverlayComponent extends Container { if (msg.content.length === 0 && msg.errorMessage && !isSilentAbort(msg.errorMessage)) { const startLine = lines.length; const isSelected = entryIndex === this.#selectedEntryIndex; - const cursor = isSelected ? theme.fg("accent", "▶") : " "; + const cursor = isSelected ? theme.fg("accent", theme.nav.cursor) : " "; lines.push(""); const errorLines = msg.errorMessage.split("\n"); const maxWidth = contentWidth(); @@ -370,7 +791,7 @@ export class SessionObserverOverlayComponent extends Container { const isSelected = entryIndex === this.#selectedEntryIndex; const isExpanded = this.#expandedEntries.has(entryIndex); const label = msg.role === "developer" ? "System" : "User"; - const cursor = isSelected ? theme.fg("accent", "▶") : " "; + const cursor = isSelected ? theme.fg("accent", theme.nav.cursor) : " "; lines.push(""); if (isExpanded) { lines.push(`${cursor} ${theme.fg("dim", `[${label}]`)}`); @@ -402,7 +823,7 @@ export class SessionObserverOverlayComponent extends Container { } #renderThinkingLines(lines: string[], thinking: string, expanded: boolean, selected: boolean): void { - const cursor = selected ? theme.fg("accent", "▶") : " "; + const cursor = selected ? theme.fg("accent", theme.nav.cursor) : " "; const maxChars = expanded ? MAX_THINKING_CHARS_EXPANDED : MAX_THINKING_CHARS_COLLAPSED; const truncated = thinking.length > maxChars; const expandLabel = !expanded && truncated ? theme.fg("dim", " ↵") : ""; @@ -435,7 +856,7 @@ export class SessionObserverOverlayComponent extends Container { } #renderTextLines(lines: string[], text: string, expanded: boolean, selected: boolean): void { - const cursor = selected ? theme.fg("accent", "▶") : " "; + const cursor = selected ? theme.fg("accent", theme.nav.cursor) : " "; lines.push(""); lines.push(`${cursor} ${theme.fg("muted", "Response")}`); @@ -467,7 +888,7 @@ export class SessionObserverOverlayComponent extends Container { expanded: boolean, selected: boolean, ): void { - const cursor = selected ? theme.fg("accent", "▶") : " "; + const cursor = selected ? theme.fg("accent", theme.nav.cursor) : " "; lines.push(""); // Tool call header @@ -565,14 +986,21 @@ export class SessionObserverOverlayComponent extends Container { case "ast_edit": return args.path ? `path: ${args.path}` : ""; case "task": { - const tasks = args.tasks; - return Array.isArray(tasks) ? `${tasks.length} task(s)` : ""; + const target = + typeof args.resume === "string" && args.resume + ? `resume ${args.resume}` + : typeof args.agent === "string" + ? args.agent + : ""; + const id = typeof args.id === "string" && args.id ? ` ${args.id}` : ""; + return `${target}${id}`.trim(); } default: { const parts: string[] = []; let total = 0; - for (const [key, value] of Object.entries(args)) { + for (const key in args) { if (key.startsWith("_")) continue; + const value = args[key]; const v = typeof value === "string" ? value : JSON.stringify(value); const entry = `${key}: ${replaceTabs(v ?? "")}`; if (total + entry.length > MAX_TOOL_ARGS_CHARS) break; @@ -592,7 +1020,7 @@ export class SessionObserverOverlayComponent extends Container { const fromByte = this.#transcriptCache?.bytesRead ?? 0; const result = readFileIncremental(sessionFile, fromByte); if (!result) { - logger.debug("Session observer: failed to read session file", { path: sessionFile }); + logger.debug("Agent hub: failed to read session file", { path: sessionFile }); return this.#transcriptCache?.entries ?? null; } @@ -627,213 +1055,9 @@ export class SessionObserverOverlayComponent extends Container { } return this.#transcriptCache.entries; } - - #navigateBack(): boolean { - if (this.#navigationStack.length === 0) return false; - const prev = this.#navigationStack.pop()!; - this.#selectedSessionId = prev.sessionId; - this.#transcriptCache = undefined; - this.#scrollOffset = 0; - this.#selectedEntryIndex = 0; - this.#expandedEntries.clear(); - this.#rebuildViewerContent(); - return true; - } - - handleInput(keyData: string): void { - // Ctrl+S (observe key) always closes the overlay - for (const key of this.#observeKeys) { - if (matchesKey(keyData, key)) { - this.#onDone(); - return; - } - } - - this.#handleViewerInput(keyData); - } - - #handleViewerInput(keyData: string): void { - const entryCount = this.#viewerEntries.length; - - // Escape — pop breadcrumb navigation or close overlay - if (matchesKey(keyData, "escape")) { - if (!this.#navigateBack()) { - this.#onDone(); - } - return; - } - - // j / down — move selection down - if (keyData === "j" || matchesSelectDown(keyData)) { - if (entryCount > 0) { - this.#selectedEntryIndex = Math.min(this.#selectedEntryIndex + 1, entryCount - 1); - } - this.#rebuildAndScroll(); - return; - } - - // k / up — move selection up - if (keyData === "k" || matchesSelectUp(keyData)) { - if (entryCount > 0) { - this.#selectedEntryIndex = Math.max(this.#selectedEntryIndex - 1, 0); - } - this.#rebuildAndScroll(); - return; - } - - // Page Down - if (matchesKey(keyData, "pageDown")) { - if (entryCount > 0) { - const prevIndex = this.#selectedEntryIndex; - this.#selectedEntryIndex = Math.min(this.#selectedEntryIndex + 5, entryCount - 1); - // If selection didn't move (bottom of list or single oversized entry), fall back to line scroll - if (this.#selectedEntryIndex === prevIndex) { - this.#scrollOffset = Math.min( - this.#scrollOffset + PAGE_SIZE, - Math.max(0, this.#renderedLines.length - this.#viewportHeight), - ); - } - } else { - this.#scrollOffset = Math.min( - this.#scrollOffset + PAGE_SIZE, - Math.max(0, this.#renderedLines.length - this.#viewportHeight), - ); - } - this.#rebuildAndScroll(); - return; - } - - // Page Up - if (matchesKey(keyData, "pageUp")) { - if (entryCount > 0) { - const prevIndex = this.#selectedEntryIndex; - this.#selectedEntryIndex = Math.max(this.#selectedEntryIndex - 5, 0); - // If selection didn't move (top of list or single oversized entry), fall back to line scroll - if (this.#selectedEntryIndex === prevIndex) { - this.#scrollOffset = Math.max(this.#scrollOffset - PAGE_SIZE, 0); - } - } else { - this.#scrollOffset = Math.max(this.#scrollOffset - PAGE_SIZE, 0); - } - this.#rebuildAndScroll(); - return; - } - - // Enter — toggle expand/collapse, or dive into nested session - if (matchesKey(keyData, "enter") || keyData === "\r" || keyData === "\n") { - if (entryCount > 0 && this.#selectedEntryIndex < entryCount) { - // Toggle expand/collapse - if (this.#expandedEntries.has(this.#selectedEntryIndex)) { - this.#expandedEntries.delete(this.#selectedEntryIndex); - } else { - this.#expandedEntries.add(this.#selectedEntryIndex); - } - this.#rebuildAndScroll(); - } - return; - } - - // G — jump to bottom - if (keyData === "G") { - if (entryCount > 0) this.#selectedEntryIndex = entryCount - 1; - this.#scrollOffset = Math.max(0, this.#renderedLines.length - this.#viewportHeight); - this.#rebuildAndScroll(); - return; - } - - // g — jump to top - if (keyData === "g") { - this.#selectedEntryIndex = 0; - this.#scrollOffset = 0; - this.#rebuildAndScroll(); - return; - } - - // ] / → / Tab — next sub-agent session - if (keyData === "]" || matchesKey(keyData, "tab") || matchesKey(keyData, "right")) { - this.#cycleSession(1); - return; - } - - // [ / ← / Shift+Tab — previous sub-agent session - if (keyData === "[" || matchesKey(keyData, "shift+tab") || matchesKey(keyData, "left")) { - this.#cycleSession(-1); - return; - } - } - - /** Get the ordered list of sub-agent session IDs (excludes main) */ - #getSubagentSessionIds(): string[] { - return this.#registry - .getSessions() - .filter(s => s.kind === "subagent") - .map(s => s.id); - } - - /** Cycle to next (+1) or previous (-1) sub-agent session */ - #cycleSession(direction: 1 | -1): void { - const ids = this.#getSubagentSessionIds(); - if (ids.length <= 1) return; - const currentIdx = ids.indexOf(this.#selectedSessionId ?? ""); - if (currentIdx < 0) return; - const nextIdx = (currentIdx + direction + ids.length) % ids.length; - this.#selectedSessionId = ids[nextIdx]; - this.#transcriptCache = undefined; - this.#scrollOffset = 0; - this.#selectedEntryIndex = 0; - this.#expandedEntries.clear(); - this.#wasAtBottom = true; - this.#rebuildViewerContent(); - // Auto-scroll to bottom: select last entry - if (this.#viewerEntries.length > 0) { - this.#selectedEntryIndex = this.#viewerEntries.length - 1; - this.#wasAtBottom = true; - this.#rebuildViewerContent(); - } - } - - /** Rebuild transcript lines (which depend on selectedEntryIndex/expandedEntries) and scroll to selection */ - #rebuildAndScroll(): void { - // Resume auto-scrolling once selection returns to the last entry - this.#wasAtBottom = this.#selectedEntryIndex >= this.#viewerEntries.length - 1; - this.#rebuildViewerContent(); - this.#scrollToSelectedEntry(); - } - - #scrollToSelectedEntry(): void { - if (this.#viewerEntries.length === 0) return; - const entry = this.#viewerEntries[this.#selectedEntryIndex]; - if (!entry) return; - - const entryTop = entry.lineStart; - const entryBottom = entry.lineStart + entry.lineCount; - - if (entry.lineCount >= this.#viewportHeight) { - // Entry taller than viewport: only snap when it's completely out of view. - // If the viewport overlaps the entry at all, the user may be paging within it. - if (this.#scrollOffset + this.#viewportHeight <= entryTop) { - // Viewport is entirely above the entry — snap to entry top - this.#scrollOffset = Math.max(0, entryTop - 1); - } else if (this.#scrollOffset >= entryBottom) { - // Viewport is entirely below the entry — snap to show entry bottom - this.#scrollOffset = Math.max(0, entryBottom - this.#viewportHeight); - } - // Otherwise: viewport overlaps the entry — don't override manual scroll - } else { - // Entry fits in viewport: ensure it's fully visible - if (entryTop < this.#scrollOffset) { - this.#scrollOffset = Math.max(0, entryTop - 1); - } - if (entryBottom > this.#scrollOffset + this.#viewportHeight) { - this.#scrollOffset = Math.max(0, entryBottom - this.#viewportHeight + 1); - } - } - } } -// Sync helpers for render path -import * as fs from "node:fs"; - +// Sync helper for the render path function readFileIncremental(filePath: string, fromByte: number): { text: string; newSize: number } | null { try { const stat = fs.statSync(filePath); diff --git a/packages/coding-agent/src/modes/components/compaction-summary-message.ts b/packages/coding-agent/src/modes/components/compaction-summary-message.ts index 55363a7c2..d2ebffb10 100644 --- a/packages/coding-agent/src/modes/components/compaction-summary-message.ts +++ b/packages/coding-agent/src/modes/components/compaction-summary-message.ts @@ -1,51 +1,87 @@ -import { Box, Markdown, Spacer, Text } from "@oh-my-pi/pi-tui"; +import { Box, type Component, Markdown } from "@oh-my-pi/pi-tui"; import { getMarkdownTheme, theme } from "../../modes/theme/theme"; import type { CompactionSummaryMessage } from "../../session/messages"; /** - * Component that renders a compaction message with collapsed/expanded state. - * Uses same background color as hook messages for visual consistency. + * Compaction point in the transcript, rendered as a slim horizontal divider: + * + * ──────── 📷 compacted · ctrl+o ──────── + * + * The conversation above the divider stays visible (display transcript keeps + * full history); only the LLM context was reset. Expanding (ctrl+o) reveals + * the compaction summary below the divider. */ -export class CompactionSummaryMessageComponent extends Box { +export class CompactionSummaryMessageComponent implements Component { #expanded = false; + #cache?: { width: number; lines: string[] }; + #detail?: Box; - constructor(private readonly message: CompactionSummaryMessage) { - super(1, 1, t => theme.bg("customMessageBg", t)); - this.#updateDisplay(); - } + constructor(private readonly message: CompactionSummaryMessage) {} setExpanded(expanded: boolean): void { + if (this.#expanded === expanded) return; this.#expanded = expanded; - this.#updateDisplay(); + this.#cache = undefined; } - override invalidate(): void { - super.invalidate(); - this.#updateDisplay(); + invalidate(): void { + this.#cache = undefined; + // Theme may have changed — rebuild the detail box lazily on next render. + this.#detail = undefined; } - #updateDisplay(): void { - this.clear(); - - const tokenStr = this.message.tokensBefore.toLocaleString(); - const label = theme.fg("customMessageLabel", theme.bold("[compaction]")); - this.addChild(new Text(label, 0, 0)); - this.addChild(new Spacer(1)); - - if (this.#expanded) { - const header = `**Compacted from ${tokenStr} tokens**\n\n`; - this.addChild( - new Markdown(header + this.message.summary, 0, 0, getMarkdownTheme(), { - color: (text: string) => theme.fg("customMessageText", text), - }), - ); - } else { - this.addChild( - new Text(theme.fg("customMessageText", `Compacted from ${tokenStr} tokens (ctrl+o to expand)`), 0, 0), - ); - if (this.message.shortSummary) { - this.addChild(new Text(theme.fg("customMessageText", this.message.shortSummary), 0, 1)); - } + render(width: number): readonly string[] { + width = Math.max(1, width); + if (this.#cache?.width === width) { + return this.#cache.lines; } + const lines = this.#expanded + ? ["", this.#divider(width), "", ...this.#detailBox().render(width)] + : ["", this.#divider(width), ""]; + this.#cache = { width, lines }; + return lines; + } + + #divider(width: number): string { + const rule = theme.tree.horizontal; + const label = `${theme.icon.camera} compacted`; + // sep.dot ships pre-padded (" · "); trim so the hint joins with single spaces. + const hint = `${theme.sep.dot.trim()} ctrl+o`; + const plainWidth = Bun.stringWidth(`${label} ${hint}`, { countAnsiEscapeCodes: false }); + // ` label hint ` framed by rules on both sides. + const remaining = width - plainWidth - 2; + if (remaining < 4) { + // Too narrow for a framed rule — emit the bare label. + return theme.fg("muted", label); + } + const left = Math.floor(remaining / 2); + const right = remaining - left; + return ( + theme.fg("dim", rule.repeat(left)) + + ` ${theme.fg("muted", label)} ${theme.fg("dim", hint)} ` + + theme.fg("dim", rule.repeat(right)) + ); + } + + #detailBox(): Box { + if (this.#detail) return this.#detail; + const box = new Box(1, 1, t => theme.bg("customMessageBg", t)); + const tokenStr = this.message.tokensBefore.toLocaleString(); + const frameCount = this.message.images?.length ?? 0; + const frameNote = + frameCount > 0 ? `\n\n_${frameCount} snapcompact frame${frameCount === 1 ? "" : "s"} attached_` : ""; + box.addChild( + new Markdown( + `**Compacted from ${tokenStr} tokens**\n\n${this.message.summary}${frameNote}`, + 0, + 0, + getMarkdownTheme(), + { + color: (text: string) => theme.fg("customMessageText", text), + }, + ), + ); + this.#detail = box; + return box; } } diff --git a/packages/coding-agent/src/modes/components/custom-editor.ts b/packages/coding-agent/src/modes/components/custom-editor.ts index e6f04e6ff..d5ffa1e04 100644 --- a/packages/coding-agent/src/modes/components/custom-editor.ts +++ b/packages/coding-agent/src/modes/components/custom-editor.ts @@ -175,6 +175,8 @@ export class CustomEditor extends Editor { onDequeue?: () => void; /** Called when Caps Lock is pressed. */ onCapsLock?: () => void; + /** Called when left-arrow is pressed while the editor is empty (cursor necessarily at start). */ + onLeftAtStart?: () => void; /** Custom key handlers from extensions and non-built-in app actions. */ #customKeyHandlers = new Map void>(); @@ -257,6 +259,14 @@ export class CustomEditor extends Editor { const parsedKey = parseKey(data); const canonical = parsedKey !== undefined ? canonicalKeyId(parsedKey) : undefined; + // Left-arrow on an empty editor: surface for the agent-hub double-tap + // gesture. Plain "left" only — modified arrows and any in-text cursor + // movement fall through to normal handling. + if (canonical === "left" && this.onLeftAtStart && this.getText().trim() === "") { + this.onLeftAtStart(); + return; + } + if (canonical !== undefined) { // Intercept configured image paste (async - fires and handles result) if (this.#matchesAction(canonical, "app.clipboard.pasteImage") && this.onPasteImage) { diff --git a/packages/coding-agent/src/modes/controllers/extension-ui-controller.ts b/packages/coding-agent/src/modes/controllers/extension-ui-controller.ts index 3b6abe82e..f0f2027f3 100644 --- a/packages/coding-agent/src/modes/controllers/extension-ui-controller.ts +++ b/packages/coding-agent/src/modes/controllers/extension-ui-controller.ts @@ -140,7 +140,7 @@ export class ExtensionUiController { reload: async () => { await this.ctx.session.reload(); this.ctx.chatContainer.clear(); - this.ctx.renderInitialMessages(undefined, { clearTerminalHistory: true }); + this.ctx.renderInitialMessages({ clearTerminalHistory: true }); await this.ctx.reloadTodos(); this.ctx.showStatus("Reloaded session"); }, @@ -197,7 +197,7 @@ export class ExtensionUiController { // Update UI this.ctx.chatContainer.clear(); - this.ctx.renderInitialMessages(undefined, { clearTerminalHistory: true }); + this.ctx.renderInitialMessages({ clearTerminalHistory: true }); await this.ctx.reloadTodos(); this.ctx.editor.setText(result.selectedText); this.ctx.showStatus("Branched to new session"); @@ -212,7 +212,7 @@ export class ExtensionUiController { // Update UI this.ctx.chatContainer.clear(); - this.ctx.renderInitialMessages(undefined, { clearTerminalHistory: true }); + this.ctx.renderInitialMessages({ clearTerminalHistory: true }); await this.ctx.reloadTodos(); if (result.editorText && !this.ctx.editor.getText().trim()) { this.ctx.editor.setText(result.editorText); @@ -230,7 +230,7 @@ export class ExtensionUiController { } setSessionTerminalTitle(this.ctx.sessionManager.getSessionName(), this.ctx.sessionManager.getCwd()); this.ctx.chatContainer.clear(); - this.ctx.renderInitialMessages(undefined, { clearTerminalHistory: true }); + this.ctx.renderInitialMessages({ clearTerminalHistory: true }); await this.ctx.reloadTodos(); return { cancelled: false }; }, @@ -376,7 +376,7 @@ export class ExtensionUiController { reload: async () => { await this.ctx.session.reload(); this.ctx.chatContainer.clear(); - this.ctx.renderInitialMessages(undefined, { clearTerminalHistory: true }); + this.ctx.renderInitialMessages({ clearTerminalHistory: true }); await this.ctx.reloadTodos(); this.ctx.showStatus("Reloaded session"); }, @@ -426,7 +426,7 @@ export class ExtensionUiController { // Update UI this.ctx.chatContainer.clear(); - this.ctx.renderInitialMessages(undefined, { clearTerminalHistory: true }); + this.ctx.renderInitialMessages({ clearTerminalHistory: true }); await this.ctx.reloadTodos(); this.ctx.editor.setText(result.selectedText); this.ctx.showStatus("Branched to new session"); @@ -441,7 +441,7 @@ export class ExtensionUiController { // Update UI this.ctx.chatContainer.clear(); - this.ctx.renderInitialMessages(undefined, { clearTerminalHistory: true }); + this.ctx.renderInitialMessages({ clearTerminalHistory: true }); await this.ctx.reloadTodos(); if (result.editorText && !this.ctx.editor.getText().trim()) { this.ctx.editor.setText(result.editorText); @@ -458,7 +458,7 @@ export class ExtensionUiController { return { cancelled: true }; } this.ctx.chatContainer.clear(); - this.ctx.renderInitialMessages(undefined, { clearTerminalHistory: true }); + this.ctx.renderInitialMessages({ clearTerminalHistory: true }); await this.ctx.reloadTodos(); return { cancelled: false }; }, diff --git a/packages/coding-agent/src/modes/controllers/input-controller.ts b/packages/coding-agent/src/modes/controllers/input-controller.ts index de55b97ad..39a8dca25 100644 --- a/packages/coding-agent/src/modes/controllers/input-controller.ts +++ b/packages/coding-agent/src/modes/controllers/input-controller.ts @@ -235,10 +235,26 @@ export class InputController { for (const key of this.ctx.keybindings.getKeys("app.clipboard.copyLine")) { this.ctx.editor.setCustomKeyHandler(key, () => this.handleCopyCurrentLine()); } - for (const key of this.ctx.keybindings.getKeys("app.session.observe")) { - this.ctx.editor.setCustomKeyHandler(key, () => this.ctx.showSessionObserver()); + const hubKeys = new Set([ + ...this.ctx.keybindings.getKeys("app.agents.hub"), + ...this.ctx.keybindings.getKeys("app.session.observe"), + ]); + for (const key of hubKeys) { + this.ctx.editor.setCustomKeyHandler(key, () => this.ctx.showAgentHub()); } + // Double-tap left arrow on an empty editor opens the agent hub — same + // 500ms window as the double-escape state machine above. + this.ctx.editor.onLeftAtStart = () => { + const now = Date.now(); + if (now - this.ctx.lastLeftTapTime < 500) { + this.ctx.lastLeftTapTime = 0; + this.ctx.showAgentHub(); + } else { + this.ctx.lastLeftTapTime = now; + } + }; + this.#setupEnhancedPaste(); this.ctx.editor.onChange = (text: string) => { diff --git a/packages/coding-agent/src/modes/controllers/selector-controller.ts b/packages/coding-agent/src/modes/controllers/selector-controller.ts index d8dd97a5b..54143cc7a 100644 --- a/packages/coding-agent/src/modes/controllers/selector-controller.ts +++ b/packages/coding-agent/src/modes/controllers/selector-controller.ts @@ -40,6 +40,7 @@ import { shortenPath } from "../../tools/render-utils"; import { copyToClipboard } from "../../utils/clipboard"; import { setSessionTerminalTitle } from "../../utils/title-generator"; import { AgentDashboard } from "../components/agent-dashboard"; +import { AgentHubOverlayComponent } from "../components/agent-hub"; import { AssistantMessageComponent } from "../components/assistant-message"; import { CopySelectorComponent } from "../components/copy-selector"; import { ExtensionDashboard } from "../components/extensions"; @@ -47,7 +48,6 @@ import { HistorySearchComponent } from "../components/history-search"; import { ModelSelectorComponent } from "../components/model-selector"; import { OAuthSelectorComponent } from "../components/oauth-selector"; import { PluginSelectorComponent } from "../components/plugin-selector"; -import { SessionObserverOverlayComponent } from "../components/session-observer-overlay"; import { SessionSelectorComponent } from "../components/session-selector"; import { SettingsSelectorComponent } from "../components/settings-selector"; import { ToolExecutionComponent } from "../components/tool-execution"; @@ -578,7 +578,7 @@ export class SelectorController { } this.ctx.chatContainer.clear(); - this.ctx.renderInitialMessages(undefined, { clearTerminalHistory: true }); + this.ctx.renderInitialMessages({ clearTerminalHistory: true }); this.ctx.editor.setText(result.selectedText); done(); this.ctx.showStatus("Branched to new session"); @@ -719,9 +719,10 @@ export class SelectorController { return; } - // Update UI — pass the context built by navigateTree to skip a second O(N) walk. + // Update UI — rebuild the display transcript for the new leaf (the + // context from navigateTree is the LLM context, not the transcript). this.ctx.chatContainer.clear(); - this.ctx.renderInitialMessages(result.sessionContext, { clearTerminalHistory: true }); + this.ctx.renderInitialMessages({ clearTerminalHistory: true }); await this.ctx.reloadTodos(); if (result.editorText && !this.ctx.editor.getText().trim()) { this.ctx.editor.setText(result.editorText); @@ -846,7 +847,7 @@ export class SelectorController { this.ctx.statusLine.setSessionStartTime(Date.now()); this.ctx.updateEditorTopBorder(); this.ctx.updateEditorBorderColor(); - this.ctx.renderInitialMessages(undefined, { clearTerminalHistory: true }); + this.ctx.renderInitialMessages({ clearTerminalHistory: true }); await this.ctx.reloadTodos(); this.ctx.ui.requestRender(true, { clearScrollback: true }); return true; @@ -871,7 +872,7 @@ export class SelectorController { // Clear and re-render the chat this.ctx.chatContainer.clear(); - this.ctx.renderInitialMessages(undefined, { clearTerminalHistory: true }); + this.ctx.renderInitialMessages({ clearTerminalHistory: true }); await this.ctx.reloadTodos(); this.ctx.showStatus(movedProject ? `Resumed session in ${shortenPath(newCwd)}` : "Resumed session"); } @@ -1074,31 +1075,34 @@ export class SelectorController { }); } - showSessionObserver(registry: SessionObserverRegistry): void { - const observeKeys = this.ctx.keybindings.getKeys("app.session.observe"); - let cleanup: (() => void) | undefined; + showAgentHub(observers: SessionObserverRegistry): void { + const hubKeys = [ + ...this.ctx.keybindings.getKeys("app.agents.hub"), + ...this.ctx.keybindings.getKeys("app.session.observe"), + ]; + let hub: AgentHubOverlayComponent | undefined; let overlayHandle: OverlayHandle | undefined; const done = () => { - cleanup?.(); + hub?.dispose(); overlayHandle?.hide(); this.ctx.ui.requestRender(); }; - const selector = new SessionObserverOverlayComponent(registry, done, observeKeys); - - cleanup = registry.onChange(() => { - selector.refreshFromRegistry(); - this.ctx.ui.requestRender(); + hub = new AgentHubOverlayComponent({ + observers, + hubKeys, + onDone: done, + requestRender: () => this.ctx.ui.requestRender(), }); - overlayHandle = this.ctx.ui.showOverlay(selector, { + overlayHandle = this.ctx.ui.showOverlay(hub, { anchor: "bottom-center", width: "100%", maxHeight: "100%", margin: 0, }); - this.ctx.ui.setFocus(selector); + this.ctx.ui.setFocus(hub); this.ctx.ui.requestRender(); } } diff --git a/packages/coding-agent/src/modes/interactive-mode.ts b/packages/coding-agent/src/modes/interactive-mode.ts index 238634188..d26d385d4 100644 --- a/packages/coding-agent/src/modes/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive-mode.ts @@ -327,6 +327,7 @@ export class InteractiveMode implements InteractiveModeContext { #pendingSubmissionDispose: (() => void) | undefined; lastSigintTime = 0; lastEscapeTime = 0; + lastLeftTapTime = 0; shutdownRequested = false; #isShuttingDown = false; hookSelector: HookSelectorComponent | undefined = undefined; @@ -1091,7 +1092,9 @@ export class InteractiveMode implements InteractiveModeContext { rebuildChatFromMessages(): void { this.chatContainer.clear(); - const context = this.session.buildDisplaySessionContext(); + // Full-history transcript: compactions render as inline dividers instead + // of restarting the visible conversation (the LLM context still resets). + const context = this.session.buildTranscriptSessionContext(); this.renderSessionContext(context); } @@ -2883,11 +2886,8 @@ export class InteractiveMode implements InteractiveModeContext { this.#uiHelpers.renderSessionContext(sessionContext, options); } - renderInitialMessages( - prebuiltContext?: SessionContext, - options?: { preserveExistingChat?: boolean; clearTerminalHistory?: boolean }, - ): void { - this.#uiHelpers.renderInitialMessages(prebuiltContext, options); + renderInitialMessages(options?: { preserveExistingChat?: boolean; clearTerminalHistory?: boolean }): void { + this.#uiHelpers.renderInitialMessages(options); } getUserMessageText(message: Message): string { @@ -3071,13 +3071,8 @@ export class InteractiveMode implements InteractiveModeContext { await this.#selectorController.showDebugSelector(); } - showSessionObserver(): void { - const sessions = this.#observerRegistry.getSessions(); - if (sessions.length <= 1) { - this.showStatus("No active subagent sessions"); - return; - } - this.#selectorController.showSessionObserver(this.#observerRegistry); + showAgentHub(): void { + this.#selectorController.showAgentHub(this.#observerRegistry); } resetObserverRegistry(): void { diff --git a/packages/coding-agent/src/modes/theme/theme.ts b/packages/coding-agent/src/modes/theme/theme.ts index f9ebba809..bfab74921 100644 --- a/packages/coding-agent/src/modes/theme/theme.ts +++ b/packages/coding-agent/src/modes/theme/theme.ts @@ -129,6 +129,8 @@ export type SymbolKey = | "icon.extensionInstruction" // STT | "icon.mic" + // Compaction divider + | "icon.camera" // Thinking Levels | "thinking.minimal" | "thinking.low" @@ -220,7 +222,8 @@ export type SymbolKey = | "tool.resolve" | "tool.review" | "tool.inspectImage" - | "tool.goal"; + | "tool.goal" + | "tool.irc"; type SymbolMap = Record; @@ -322,6 +325,8 @@ const UNICODE_SYMBOLS: SymbolMap = { "icon.extensionInstruction": "📘", // STT "icon.mic": "🎤", + // Compaction divider + "icon.camera": "📷", // Thinking levels "thinking.minimal": "◔ min", "thinking.low": "◑ low", @@ -414,6 +419,7 @@ const UNICODE_SYMBOLS: SymbolMap = { "tool.review": "◉", "tool.inspectImage": "🖼", "tool.goal": "◎", + "tool.irc": "✉", }; const NERD_SYMBOLS: SymbolMap = { @@ -599,6 +605,8 @@ const NERD_SYMBOLS: SymbolMap = { "icon.extensionInstruction": "\uf02d", // STT - fa-microphone "icon.mic": "\uf130", + // Compaction divider - fa-camera-retro + "icon.camera": "\uf083", // Thinking Levels - emoji labels // pick: 🤨 min | alt:  min  min "thinking.minimal": "\u{F0E7} min", @@ -708,6 +716,7 @@ const NERD_SYMBOLS: SymbolMap = { "tool.review": "\uEA70", "tool.inspectImage": "\uEAEA", "tool.goal": "\uEBF8", + "tool.irc": "\uF086", }; const ASCII_SYMBOLS: SymbolMap = { @@ -808,6 +817,8 @@ const ASCII_SYMBOLS: SymbolMap = { "icon.extensionInstruction": "IN", // STT "icon.mic": "MIC", + // Compaction divider + "icon.camera": "[o]", // Thinking Levels "thinking.minimal": "[min]", "thinking.low": "[low]", @@ -898,6 +909,7 @@ const ASCII_SYMBOLS: SymbolMap = { "tool.review": "rev", "tool.inspectImage": "[i]", "tool.goal": "(o)", + "tool.irc": "irc", }; const SYMBOL_PRESETS: Record = { @@ -1686,6 +1698,7 @@ export class Theme { extensionContextFile: this.#symbols["icon.extensionContextFile"], extensionInstruction: this.#symbols["icon.extensionInstruction"], mic: this.#symbols["icon.mic"], + camera: this.#symbols["icon.camera"], }; } diff --git a/packages/coding-agent/src/modes/types.ts b/packages/coding-agent/src/modes/types.ts index f9ffd072a..9df19c9b3 100644 --- a/packages/coding-agent/src/modes/types.ts +++ b/packages/coding-agent/src/modes/types.ts @@ -136,6 +136,7 @@ export interface InteractiveModeContext { locallySubmittedUserSignatures: Set; lastSigintTime: number; lastEscapeTime: number; + lastLeftTapTime: number; shutdownRequested: boolean; hookSelector: HookSelectorComponent | undefined; hookInput: HookInputComponent | undefined; @@ -225,10 +226,7 @@ export interface InteractiveModeContext { sessionContext: SessionContext, options?: { updateFooter?: boolean; populateHistory?: boolean }, ): void; - renderInitialMessages( - prebuiltContext?: SessionContext, - options?: { preserveExistingChat?: boolean; clearTerminalHistory?: boolean }, - ): void; + renderInitialMessages(options?: { preserveExistingChat?: boolean; clearTerminalHistory?: boolean }): void; getUserMessageText(message: Message): string; findLastAssistantMessage(): AssistantMessage | undefined; extractAssistantText(message: AssistantMessage): string; @@ -292,7 +290,7 @@ export interface InteractiveModeContext { showProviderSetup(): Promise; showHookConfirm(title: string, message: string): Promise; showDebugSelector(): Promise; - showSessionObserver(): void; + showAgentHub(): void; resetObserverRegistry(): void; // Input handling diff --git a/packages/coding-agent/src/modes/utils/hotkeys-markdown.ts b/packages/coding-agent/src/modes/utils/hotkeys-markdown.ts index 97a179e66..b0d86e8f2 100644 --- a/packages/coding-agent/src/modes/utils/hotkeys-markdown.ts +++ b/packages/coding-agent/src/modes/utils/hotkeys-markdown.ts @@ -50,6 +50,7 @@ export function buildHotkeysMarkdown(bindings: HotkeysMarkdownBindings): string `| \`${appKey(bindings, "app.editor.external")}\` | Edit message in external editor |`, `| \`${appKey(bindings, "app.clipboard.pasteImage")}\` | Paste image from clipboard |`, `| \`${appKey(bindings, "app.stt.toggle")}\` | Toggle speech-to-text recording |`, + `| \`${appKey(bindings, "app.agents.hub")}\` / \`${appKey(bindings, "app.session.observe")}\` / double-tap \`←\` (empty editor) | Open the agent hub |`, "| `#` | Open prompt actions |", "| `/` | Slash commands |", "| `!` | Run bash command |", diff --git a/packages/coding-agent/src/modes/utils/ui-helpers.ts b/packages/coding-agent/src/modes/utils/ui-helpers.ts index 3136296f8..a01bc69da 100644 --- a/packages/coding-agent/src/modes/utils/ui-helpers.ts +++ b/packages/coding-agent/src/modes/utils/ui-helpers.ts @@ -190,19 +190,13 @@ export class UiHelpers { this.ctx.chatContainer.addChild(component); break; } - if ( - message.customType === "irc:incoming" || - message.customType === "irc:autoreply" || - message.customType === "irc:relay" - ) { + if (message.customType === "irc:incoming" || message.customType === "irc:relay") { const details = ( message as CustomMessage<{ from?: string; to?: string; message?: string; - reply?: string; body?: string; - kind?: "message" | "reply"; }> ).details; let arrow: string; @@ -211,10 +205,6 @@ export class UiHelpers { const peer = details?.from ?? "?"; body = details?.message ?? ""; arrow = `⇦ ${peer}`; - } else if (message.customType === "irc:autoreply") { - const peer = details?.to ?? "?"; - body = details?.reply ?? ""; - arrow = `⇨ ${peer}`; } else { const from = details?.from ?? "?"; const to = details?.to ?? "?"; @@ -337,13 +327,23 @@ export class UiHelpers { let readGroup: ReadToolGroupComponent | null = null; const readToolCallArgs = new Map>(); const readToolCallAssistantComponents = new Map(); - const deferredMessages: AgentMessage[] = []; - for (const message of sessionContext.messages) { - // Defer compaction summaries so they render at the bottom (visible after scroll) - if (message.role === "compactionSummary") { - deferredMessages.push(message); - continue; + // Rebuild-time mirror of the event controller's displaceable-poll + // bookkeeping: a `job` poll that found every watched job still running is + // superseded by the next `job` call, so a rebuilt transcript collapses a + // repeated-poll run to its final snapshot instead of replaying the spam. + let waitingPoll: ToolExecutionComponent | null = null; + const resolveWaitingPoll = (nextToolName?: string) => { + const previous = waitingPoll; + if (!previous) return; + waitingPoll = null; + if (nextToolName === "job" && previous.isDisplaceableBlock()) { + this.ctx.chatContainer.removeChild(previous); } + // Sealing freezes the block and stops the waiting-poll spinner that + // updateResult armed. + previous.seal(); + }; + for (const message of sessionContext.messages) { // Assistant messages need special handling for tool calls if (message.role === "assistant") { this.ctx.addMessageToChat(message); @@ -504,17 +504,15 @@ export class UiHelpers { // rebuilt group freezes (even with a never-persisted result) and commits to // native scrollback like every other historical block. readGroup?.seal(); - - // Render deferred messages (compaction summaries) at the bottom so they're visible - for (const message of deferredMessages) { - this.ctx.addMessageToChat(message, options); - } + // A trailing waiting poll is final history on rebuild; seal it so it + // freezes (and its spinner timer stops) like every other block. + resolveWaitingPoll(); this.ctx.pendingTools.clear(); this.ctx.ui.requestRender(); } - renderInitialMessages(prebuiltContext?: SessionContext, options: RenderInitialMessagesOptions = {}): void { + renderInitialMessages(options: RenderInitialMessagesOptions = {}): void { // This path is used to rebuild the visible chat transcript (e.g. after custom/debug UI). // Clear existing rendered chat first to avoid duplicating the full session in the container. // On a non-preserving rebuild the existing blocks are discarded for good, so @@ -530,8 +528,9 @@ export class UiHelpers { this.ctx.pendingBashComponents = []; this.ctx.pendingPythonComponents = []; - // Reuse a pre-built context when available (e.g. from navigateTree) to avoid a second O(N) walk. - const context = prebuiltContext ?? this.ctx.sessionManager.buildSessionContext(); + // Display always uses the full-history transcript: compactions show as + // inline dividers instead of restarting the visible conversation. + const context = this.ctx.session.buildTranscriptSessionContext(); this.ctx.renderSessionContext(context, { updateFooter: true, populateHistory: true, diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index 022e34867..4805273f8 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/packages/coding-agent/src/session/agent-session.ts @@ -4045,6 +4045,16 @@ export class AgentSession { return deobfuscateSessionContext(this.sessionManager.buildSessionContext(), this.#obfuscator); } + /** + * Full-history transcript for TUI display: every path entry in + * chronological order with compactions rendered inline at the point they + * fired (instead of replacing prior history). Display-only — NEVER feed + * the result to `agent.replaceMessages` or a provider. + */ + buildTranscriptSessionContext(): SessionContext { + return deobfuscateSessionContext(this.sessionManager.buildSessionContext({ transcript: true }), this.#obfuscator); + } + #obfuscateForProvider(value: T): T { if (!this.#obfuscator?.hasSecrets()) return value; return this.#obfuscator.obfuscateObject(value); diff --git a/packages/coding-agent/src/session/session-manager.ts b/packages/coding-agent/src/session/session-manager.ts index 5252730a7..bc0c23c40 100644 --- a/packages/coding-agent/src/session/session-manager.ts +++ b/packages/coding-agent/src/session/session-manager.ts @@ -2,6 +2,7 @@ import * as fs from "node:fs"; import * as os from "node:os"; import * as path from "node:path"; import type { AgentMessage } from "@oh-my-pi/pi-agent-core"; +import { getPreservedSnapcompactArchive, snapcompactImages } from "@oh-my-pi/pi-agent-core/compaction"; import type { ImageContent, Message, @@ -544,6 +545,17 @@ export function getLatestCompactionEntry(entries: SessionEntry[]): CompactionEnt return null; } +export interface BuildSessionContextOptions { + /** + * Build the full-history display transcript instead of the LLM context: + * every path entry in chronological order, with each compaction emitted + * inline as a `compactionSummary` message at the position it fired rather + * than replacing the history before it. Display-only — never send the + * result to a provider. + */ + transcript?: boolean; +} + /** * Build the session context from entries using tree traversal. * If leafId is provided, walks from that entry to root. @@ -553,6 +565,7 @@ export function buildSessionContext( entries: SessionEntry[], leafId?: string | null, byId?: Map, + options?: BuildSessionContextOptions, ): SessionContext { // Build uuid index if not available if (!byId) { @@ -692,7 +705,29 @@ export function buildSessionContext( } }; - if (compaction) { + if (options?.transcript) { + // Display transcript: every entry in chronological order. Compactions do + // not erase prior history here — each renders inline (as a divider in the + // TUI) at the point it fired, with any snapcompact frames re-attached so + // the component can report them. + for (const entry of path) { + if (entry.type === "compaction") { + const snapcompactArchive = getPreservedSnapcompactArchive(entry.preserveData); + messages.push( + createCompactionSummaryMessage( + entry.summary, + entry.tokensBefore, + entry.timestamp, + entry.shortSummary, + undefined, + snapcompactArchive ? snapcompactImages(snapcompactArchive) : undefined, + ), + ); + } else { + appendMessage(entry); + } + } + } else if (compaction) { const providerPayload: ProviderPayload | undefined = (() => { const candidate = compaction.preserveData?.openaiRemoteCompaction; if (!candidate || typeof candidate !== "object") return undefined; @@ -707,7 +742,9 @@ export function buildSessionContext( })(); const remoteReplacementHistory = providerPayload?.items; - // Emit summary first + // Emit summary first; re-attach any archived snapcompact frames so the + // model can keep reading the archived history after every context rebuild. + const snapcompactArchive = getPreservedSnapcompactArchive(compaction.preserveData); messages.push( createCompactionSummaryMessage( compaction.summary, @@ -715,6 +752,7 @@ export function buildSessionContext( compaction.timestamp, compaction.shortSummary, providerPayload, + snapcompactArchive ? snapcompactImages(snapcompactArchive) : undefined, ), ); @@ -957,6 +995,21 @@ async function resolveBlobRefsInEntries(entries: FileEntry[], blobStore: BlobSto await Promise.all(promises); } +/** + * Read-only message view of a session file: load entries, migrate to the + * current version, resolve blob refs, and build the context along the + * persisted leaf path (last entry). Does NOT create a writer or take the + * session lock — safe to call against a file another session is writing. + */ +export async function loadSessionMessagesReadOnly(filePath: string): Promise { + const entries = await loadEntriesFromFile(filePath); + if (entries.length === 0) return []; + migrateToCurrentVersion(entries); + await resolveBlobRefsInEntries(entries, new BlobStore(getBlobsDir())); + const sessionEntries = entries.filter((e): e is SessionEntry => e.type !== "session"); + return buildSessionContext(sessionEntries).messages; +} + /** * Lightweight metadata for a session file, used in session picker UI. * Uses lazy getters to defer string formatting until actually displayed. @@ -3205,11 +3258,12 @@ export class SessionManager { } /** - * Build the session context (what gets sent to the LLM). + * Build the session context (what gets sent to the LLM), or — with + * `{ transcript: true }` — the full-history display transcript. * Uses tree traversal from current leaf. */ - buildSessionContext(): SessionContext { - return buildSessionContext(this.getEntries(), this.#leafId, this.#byId); + buildSessionContext(options?: BuildSessionContextOptions): SessionContext { + return buildSessionContext(this.getEntries(), this.#leafId, this.#byId, options); } /** Strip stale OpenAI Responses assistant replay metadata from loaded in-memory entries. */ diff --git a/packages/coding-agent/test/interactive-mode-status.test.ts b/packages/coding-agent/test/interactive-mode-status.test.ts index f32f13587..fa82aa3a7 100644 --- a/packages/coding-agent/test/interactive-mode-status.test.ts +++ b/packages/coding-agent/test/interactive-mode-status.test.ts @@ -45,6 +45,7 @@ function createInitialRenderHarness(): { ctx: InteractiveModeContext; helpers: U session: { retryAttempt: 0, getToolByName: () => undefined, + buildTranscriptSessionContext: () => buildSessionContext([]), }, toolOutputExpanded: false, hideThinkingBlock: false, @@ -115,7 +116,7 @@ describe("InteractiveMode.showStatus", () => { const { ctx, helpers } = createInitialRenderHarness(); helpers.showWarning("startup notification probe"); - helpers.renderInitialMessages(undefined, { preserveExistingChat: true }); + helpers.renderInitialMessages({ preserveExistingChat: true }); expect(renderContainer(ctx.chatContainer)).toContain("startup notification probe"); }); diff --git a/packages/coding-agent/test/modes/components/compaction-divider.test.ts b/packages/coding-agent/test/modes/components/compaction-divider.test.ts new file mode 100644 index 000000000..3e5475fb5 --- /dev/null +++ b/packages/coding-agent/test/modes/components/compaction-divider.test.ts @@ -0,0 +1,63 @@ +/** + * Contract: the compaction point renders as a slim horizontal divider — + * `── 📷 compacted · ctrl+o ──` — instead of a full summary box, keeping the + * transcript visually continuous. Expansion (ctrl+o) reveals the summary. + * The render cache must honor the pi-tui same-reference contract: unchanged + * components return the identical array so containers can memoize. + */ + +import { beforeAll, describe, expect, it } from "bun:test"; +import { createCompactionSummaryMessage } from "@oh-my-pi/pi-agent-core/compaction"; +import type { ImageContent } from "@oh-my-pi/pi-ai"; +import { CompactionSummaryMessageComponent } from "@oh-my-pi/pi-coding-agent/modes/components/compaction-summary-message"; +import { initTheme } from "@oh-my-pi/pi-coding-agent/modes/theme/theme"; + +beforeAll(() => { + initTheme(); +}); + +const SUMMARY = "Earlier the user fixed the login TTL bug."; + +function makeComponent(images?: ImageContent[]): CompactionSummaryMessageComponent { + return new CompactionSummaryMessageComponent( + createCompactionSummaryMessage(SUMMARY, 84000, new Date().toISOString(), undefined, undefined, images), + ); +} + +describe("CompactionSummaryMessageComponent", () => { + it("collapsed: a single full-width divider carrying the expand affordance", () => { + const lines = makeComponent().render(80); + expect(lines.length).toBe(3); // breathing room above and below the rule + const rule = Bun.stripANSI(lines[1]); + expect(rule).toContain("compacted"); + expect(rule).toContain("ctrl+o"); + // The rule spans the full width and hides the summary body. + expect(Bun.stringWidth(rule)).toBe(80); + expect(rule).not.toContain(SUMMARY); + }); + + it("expanded: reveals the summary (and snapcompact frame count) below the divider", () => { + const component = makeComponent([{ type: "image", data: "ZmFrZQ==", mimeType: "image/png" }]); + component.setExpanded(true); + const text = Bun.stripANSI(component.render(80).join("\n")); + expect(text).toContain("compacted"); + expect(text).toContain(SUMMARY); + expect(text).toContain("tokens"); + expect(text).toContain("1 snapcompact frame attached"); + }); + + it("degrades to a bare label when the viewport is too narrow for a framed rule", () => { + const lines = makeComponent().render(10); + expect(Bun.stripANSI(lines[1])).toContain("compacted"); + }); + + it("honors the same-reference render cache and busts it on expansion toggle", () => { + const component = makeComponent(); + const first = component.render(80); + expect(component.render(80)).toBe(first); + component.setExpanded(true); + const expanded = component.render(80); + expect(expanded).not.toBe(first); + expect(component.render(80)).toBe(expanded); + }); +}); diff --git a/packages/coding-agent/test/modes/utils/render-initial-messages-dedupe.test.ts b/packages/coding-agent/test/modes/utils/render-initial-messages-dedupe.test.ts deleted file mode 100644 index 52de0a13e..000000000 --- a/packages/coding-agent/test/modes/utils/render-initial-messages-dedupe.test.ts +++ /dev/null @@ -1,160 +0,0 @@ -/** - * Regression test: renderInitialMessages must not call buildSessionContext when - * a prebuilt context is supplied (e.g. from navigateTree's return value). - * - * Before the fix, renderInitialMessages always called sessionManager.buildSessionContext(), - * duplicating the O(N) walk already done inside navigateTree. After the fix, - * renderInitialMessages accepts an optional prebuiltContext and skips the walk when provided. - * - * What these tests cover: - * - UiHelpers.renderInitialMessages skip/use of buildSessionContext (isolated unit tests). - * - End-to-end callcount using a real in-memory SessionManager: one build call total when - * the context is passed through, confirming the handoff contract works at runtime. - * - * What is NOT covered here (requires full AgentSession wiring): - * - That AgentSession.navigateTree actually sets result.sessionContext — enforced by - * the TypeScript return type and the agent-session unit tests. - */ - -import { beforeAll, describe, expect, it, type Mock, vi } from "bun:test"; -import { initTheme } from "@oh-my-pi/pi-coding-agent/modes/theme/theme"; -import type { InteractiveModeContext } from "@oh-my-pi/pi-coding-agent/modes/types"; -import { UiHelpers } from "@oh-my-pi/pi-coding-agent/modes/utils/ui-helpers"; -import type { SessionContext } from "@oh-my-pi/pi-coding-agent/session/session-manager"; -import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager"; - -beforeAll(() => { - initTheme(); -}); - -/** Minimal empty SessionContext (what navigateTree returns when leafId is null). */ -function makeEmptyContext(): SessionContext { - return { - messages: [], - thinkingLevel: "off", - serviceTier: undefined, - models: {}, - injectedTtsrRules: [], - selectedMCPToolNames: [], - hasPersistedMCPToolSelection: false, - mode: "none", - }; -} - -/** Build a minimal InteractiveModeContext mock, returning spies for assertions. */ -function makeCtx(sessionManager?: Pick): { - ctx: InteractiveModeContext; - buildSessionContextSpy: Mock<() => SessionContext>; - renderSessionContextSpy: Mock<(...args: unknown[]) => void>; -} { - const buildSessionContextSpy = vi.fn(() => makeEmptyContext()); - const renderSessionContextSpy = vi.fn(); - - const sm = sessionManager ?? { - buildSessionContext: buildSessionContextSpy, - getEntries: vi.fn(() => []), - getCwd: vi.fn(() => "/tmp"), - }; - - const ctx = { - chatContainer: { clear: vi.fn(), addChild: vi.fn() }, - pendingMessagesContainer: { clear: vi.fn() }, - pendingBashComponents: [], - pendingPythonComponents: [], - sessionManager: sm, - renderSessionContext: renderSessionContextSpy, - showStatus: vi.fn(), - ui: { requestRender: vi.fn() }, - resetTranscript: () => ctx.chatContainer.clear(), - } as unknown as InteractiveModeContext; - - return { ctx, buildSessionContextSpy, renderSessionContextSpy }; -} - -// ─── Part 1: Isolated renderInitialMessages behaviour ───────────────────────── - -describe("UiHelpers.renderInitialMessages — isolated", () => { - it("calls sessionManager.buildSessionContext when no prebuilt context is given", () => { - const { ctx, buildSessionContextSpy } = makeCtx(); - new UiHelpers(ctx).renderInitialMessages(); - expect(buildSessionContextSpy).toHaveBeenCalledTimes(1); - }); - - it("does NOT call sessionManager.buildSessionContext when a prebuilt context is provided", () => { - const { ctx, buildSessionContextSpy } = makeCtx(); - new UiHelpers(ctx).renderInitialMessages(makeEmptyContext()); - expect(buildSessionContextSpy).toHaveBeenCalledTimes(0); - }); - - it("passes the prebuilt context directly to renderSessionContext", () => { - const { ctx, renderSessionContextSpy } = makeCtx(); - const prebuilt = makeEmptyContext(); - new UiHelpers(ctx).renderInitialMessages(prebuilt); - expect(renderSessionContextSpy).toHaveBeenCalledWith(prebuilt, { - updateFooter: true, - populateHistory: true, - }); - }); - - it("uses the fallback context from sessionManager when no prebuilt is provided", () => { - const fallback = makeEmptyContext(); - const { ctx, renderSessionContextSpy } = makeCtx(); - (ctx.sessionManager.buildSessionContext as Mock<() => SessionContext>).mockReturnValue(fallback); - new UiHelpers(ctx).renderInitialMessages(); - expect(renderSessionContextSpy).toHaveBeenCalledWith(fallback, { - updateFooter: true, - populateHistory: true, - }); - }); -}); - -// ─── Cold-launch terminal cleanup ──────────────────────────────────────────── -// -// `omp` / `omp -c` leave the previous run's transcript in native scrollback -// because the TUI's initial paint preserves it. The cold-launch render must -// therefore request a scrollback-clearing repaint (`clearTerminalHistory`) so -// the resumed transcript replaces the stale one instead of stacking on it. -// Every in-process session load already does this; this guards the cold path. - -describe("UiHelpers.renderInitialMessages — clearTerminalHistory", () => { - it("requests a scrollback-clearing repaint when clearTerminalHistory is set", () => { - const { ctx } = makeCtx(); - new UiHelpers(ctx).renderInitialMessages(undefined, { clearTerminalHistory: true }); - expect(ctx.ui.requestRender).toHaveBeenCalledWith(true, { clearScrollback: true }); - }); - - it("never clears scrollback when clearTerminalHistory is unset", () => { - const { ctx } = makeCtx(); - new UiHelpers(ctx).renderInitialMessages(); - const clearedCall = (ctx.ui.requestRender as Mock<(...a: unknown[]) => void>).mock.calls.find( - ([force, opts]) => force === true && (opts as { clearScrollback?: boolean } | undefined)?.clearScrollback, - ); - expect(clearedCall).toBeUndefined(); - }); -}); - -// ─── Part 2: End-to-end callcount with a real SessionManager ───────────────── -// -// Simulates the selector-controller handoff: buildSessionContext is called once -// (representing navigateTree's internal call), and the result is passed directly -// to renderInitialMessages — which must NOT call it a second time. - -describe("renderInitialMessages callcount with real SessionManager", () => { - it("total buildSessionContext calls == 1 when context is threaded from navigate to render", async () => { - const sm = SessionManager.inMemory("/tmp/test"); - // Append a user message with the required timestamp field. - sm.appendMessage({ role: "user", content: "hello", timestamp: Date.now() }); - - const spy = vi.spyOn(sm, "buildSessionContext"); - - // Step 1: one call — represents navigateTree's internal rawContext build. - const context = sm.buildSessionContext(); - expect(spy).toHaveBeenCalledTimes(1); - - // Step 2: renderInitialMessages with the prebuilt context must NOT call it again. - const { ctx } = makeCtx(sm as unknown as Parameters[0]); - new UiHelpers(ctx).renderInitialMessages(context); - - expect(spy).toHaveBeenCalledTimes(1); // still exactly 1 — second walk eliminated - }); -}); diff --git a/packages/coding-agent/test/modes/utils/render-initial-messages.test.ts b/packages/coding-agent/test/modes/utils/render-initial-messages.test.ts new file mode 100644 index 000000000..14f02fa88 --- /dev/null +++ b/packages/coding-agent/test/modes/utils/render-initial-messages.test.ts @@ -0,0 +1,101 @@ +/** + * Contract: renderInitialMessages renders the DISPLAY TRANSCRIPT, not the LLM + * context. The transcript comes from `session.buildTranscriptSessionContext()` + * (full history, compactions inline); `sessionManager.buildSessionContext()` + * — the LLM-context builder — must not be consulted for display. Feeding the + * compacted LLM context to the chat is exactly the old "session starts over + * after compaction" bug. + * + * Also guards the cold-launch terminal cleanup: `omp` / `omp -c` leave the + * previous run's transcript in native scrollback because the TUI's initial + * paint preserves it, so the cold-launch render must request a + * scrollback-clearing repaint (`clearTerminalHistory`). + */ + +import { beforeAll, describe, expect, it, type Mock, vi } from "bun:test"; +import { initTheme } from "@oh-my-pi/pi-coding-agent/modes/theme/theme"; +import type { InteractiveModeContext } from "@oh-my-pi/pi-coding-agent/modes/types"; +import { UiHelpers } from "@oh-my-pi/pi-coding-agent/modes/utils/ui-helpers"; +import type { SessionContext } from "@oh-my-pi/pi-coding-agent/session/session-manager"; + +beforeAll(() => { + initTheme(); +}); + +function makeEmptyContext(): SessionContext { + return { + messages: [], + thinkingLevel: "off", + serviceTier: undefined, + models: {}, + injectedTtsrRules: [], + selectedMCPToolNames: [], + hasPersistedMCPToolSelection: false, + mode: "none", + }; +} + +/** Build a minimal InteractiveModeContext mock, returning spies for assertions. */ +function makeCtx(): { + ctx: InteractiveModeContext; + transcriptSpy: Mock<() => SessionContext>; + llmContextSpy: Mock<() => SessionContext>; + renderSessionContextSpy: Mock<(...args: unknown[]) => void>; +} { + const transcriptSpy = vi.fn(() => makeEmptyContext()); + const llmContextSpy = vi.fn(() => makeEmptyContext()); + const renderSessionContextSpy = vi.fn(); + + const ctx = { + chatContainer: { clear: vi.fn(), addChild: vi.fn() }, + pendingMessagesContainer: { clear: vi.fn() }, + pendingBashComponents: [], + pendingPythonComponents: [], + session: { buildTranscriptSessionContext: transcriptSpy }, + sessionManager: { + buildSessionContext: llmContextSpy, + getEntries: vi.fn(() => []), + getCwd: vi.fn(() => "/tmp"), + }, + renderSessionContext: renderSessionContextSpy, + showStatus: vi.fn(), + ui: { requestRender: vi.fn() }, + resetTranscript: () => ctx.chatContainer.clear(), + } as unknown as InteractiveModeContext; + + return { ctx, transcriptSpy, llmContextSpy, renderSessionContextSpy }; +} + +describe("UiHelpers.renderInitialMessages — transcript source", () => { + it("renders the display transcript, never the LLM context", () => { + const { ctx, transcriptSpy, llmContextSpy, renderSessionContextSpy } = makeCtx(); + const transcript = makeEmptyContext(); + transcriptSpy.mockReturnValue(transcript); + + new UiHelpers(ctx).renderInitialMessages(); + + expect(transcriptSpy).toHaveBeenCalledTimes(1); + expect(llmContextSpy).not.toHaveBeenCalled(); + expect(renderSessionContextSpy).toHaveBeenCalledWith(transcript, { + updateFooter: true, + populateHistory: true, + }); + }); +}); + +describe("UiHelpers.renderInitialMessages — clearTerminalHistory", () => { + it("requests a scrollback-clearing repaint when clearTerminalHistory is set", () => { + const { ctx } = makeCtx(); + new UiHelpers(ctx).renderInitialMessages({ clearTerminalHistory: true }); + expect(ctx.ui.requestRender).toHaveBeenCalledWith(true, { clearScrollback: true }); + }); + + it("never clears scrollback when clearTerminalHistory is unset", () => { + const { ctx } = makeCtx(); + new UiHelpers(ctx).renderInitialMessages(); + const clearedCall = (ctx.ui.requestRender as Mock<(...a: unknown[]) => void>).mock.calls.find( + ([force, opts]) => force === true && (opts as { clearScrollback?: boolean } | undefined)?.clearScrollback, + ); + expect(clearedCall).toBeUndefined(); + }); +}); diff --git a/packages/coding-agent/test/silent-abort-overlay-render.test.ts b/packages/coding-agent/test/silent-abort-overlay-render.test.ts index a6ab53f2e..d002df540 100644 --- a/packages/coding-agent/test/silent-abort-overlay-render.test.ts +++ b/packages/coding-agent/test/silent-abort-overlay-render.test.ts @@ -1,17 +1,19 @@ /** - * Regression: observer overlay must not render SILENT_ABORT_MARKER verbatim. + * Regression: the agent-hub chat transcript must not render SILENT_ABORT_MARKER verbatim. * - * Codex review flagged that `session-observer-overlay.ts` renders `errorMessage` - * without filtering the silent-abort sentinel. This test exercises the full - * `#buildTranscriptLines` path through a real JSONL session file and mock registry. + * Codex review flagged that the old observer overlay rendered `errorMessage` + * without filtering the silent-abort sentinel; the renderer now lives in + * `agent-hub.ts`. This test exercises the full `#buildTranscriptLines` path + * through a real JSONL session file and an isolated agent registry. */ import { afterEach, beforeAll, beforeEach, describe, expect, it } from "bun:test"; import * as fs from "node:fs"; import * as os from "node:os"; import * as path from "node:path"; -import { SessionObserverOverlayComponent } from "@oh-my-pi/pi-coding-agent/modes/components/session-observer-overlay"; +import { AgentHubOverlayComponent } from "@oh-my-pi/pi-coding-agent/modes/components/agent-hub"; import type { ObservableSession } from "@oh-my-pi/pi-coding-agent/modes/session-observer-registry"; import { initTheme } from "@oh-my-pi/pi-coding-agent/modes/theme/theme"; +import { AgentRegistry } from "@oh-my-pi/pi-coding-agent/registry/agent-registry"; import { SILENT_ABORT_MARKER } from "@oh-my-pi/pi-coding-agent/session/messages"; const SESSION_ID = "test-session-1"; @@ -32,7 +34,29 @@ function makeSubagentRegistry(sessions: ObservableSession[]) { } as unknown as import("@oh-my-pi/pi-coding-agent/modes/session-observer-registry").SessionObserverRegistry; } -describe("Observer overlay silent-abort regression", () => { +function makeHub(sessionFile: string, observed: ObservableSession[]): AgentHubOverlayComponent { + const agents = new AgentRegistry(); + agents.register({ + id: SESSION_ID, + displayName: SESSION_ID, + kind: "sub", + parentId: "Main", + session: null, + sessionFile, + status: "parked", + }); + const hub = new AgentHubOverlayComponent({ + observers: makeSubagentRegistry(observed), + hubKeys: ["ctrl+s"], + onDone: () => {}, + requestRender: () => {}, + registry: agents, + }); + hub.openChat(SESSION_ID); + return hub; +} + +describe("Agent hub silent-abort regression", () => { let tmpDir: string; beforeAll(() => { @@ -83,7 +107,7 @@ describe("Observer overlay silent-abort regression", () => { }, ]); - const registry = makeSubagentRegistry([ + const hub = makeHub(sessionFile, [ { id: SESSION_ID, kind: "subagent", @@ -94,11 +118,10 @@ describe("Observer overlay silent-abort regression", () => { }, ]); - const overlay = new SessionObserverOverlayComponent(registry, () => {}, ["ctrl+s"]); - - // Render with a reasonable width — the overlay reads the session file - // and calls #buildTranscriptLines internally. - const rendered = overlay.render(120); + // Render with a reasonable width — the hub chat view reads the session + // file and calls #buildTranscriptLines internally. + const rendered = hub.render(120); + hub.dispose(); const renderedText = rendered.join("\n"); // The sentinel MUST NOT appear verbatim in any rendered line @@ -143,7 +166,7 @@ describe("Observer overlay silent-abort regression", () => { }, ]); - const registry = makeSubagentRegistry([ + const hub = makeHub(sessionFile, [ { id: SESSION_ID, kind: "subagent", @@ -154,9 +177,8 @@ describe("Observer overlay silent-abort regression", () => { }, ]); - const overlay = new SessionObserverOverlayComponent(registry, () => {}, ["ctrl+s"]); - - const rendered = overlay.render(120); + const rendered = hub.render(120); + hub.dispose(); const renderedText = rendered.join("\n"); // A real error message SHOULD be rendered with the ✗ Error: prefix diff --git a/packages/coding-agent/test/tools/search-path-lists.test.ts b/packages/coding-agent/test/tools/search-path-lists.test.ts index 955071876..9eb52f7cb 100644 --- a/packages/coding-agent/test/tools/search-path-lists.test.ts +++ b/packages/coding-agent/test/tools/search-path-lists.test.ts @@ -7,7 +7,7 @@ import { validateToolArguments } from "@oh-my-pi/pi-ai/utils/validation"; import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; import { canonicalSnapshotKey } from "@oh-my-pi/pi-coding-agent/edit/file-snapshot-store"; import type { RenderResultOptions } from "@oh-my-pi/pi-coding-agent/extensibility/custom-tools/types"; -import { SessionObserverOverlayComponent } from "@oh-my-pi/pi-coding-agent/modes/components/session-observer-overlay"; +import { AgentHubOverlayComponent } from "@oh-my-pi/pi-coding-agent/modes/components/agent-hub"; import { TreeSelectorComponent } from "@oh-my-pi/pi-coding-agent/modes/components/tree-selector"; import type { ObservableSession, @@ -15,6 +15,7 @@ import type { } from "@oh-my-pi/pi-coding-agent/modes/session-observer-registry"; import type { Theme } from "@oh-my-pi/pi-coding-agent/modes/theme/theme"; import { initTheme } from "@oh-my-pi/pi-coding-agent/modes/theme/theme"; +import { AgentRegistry } from "@oh-my-pi/pi-coding-agent/registry/agent-registry"; import type { SessionEntry, SessionTreeNode } from "@oh-my-pi/pi-coding-agent/session/session-manager"; import { ToolChoiceQueue } from "@oh-my-pi/pi-coding-agent/session/tool-choice-queue"; import { createTools, type ToolSession } from "@oh-my-pi/pi-coding-agent/tools"; @@ -285,7 +286,7 @@ describe("tool path arrays", () => { expect(component).toBeInstanceOf(Text); expect((component as Text).getText()).toContain("in folder with spaces/"); }); - it("session observer overlay renders a single-string search path summary", async () => { + it("agent hub chat renders a single-string search path summary", async () => { const sessionFile = await makeJsonlSessionFile(tempDir, [ { type: "session", version: 3, id: "search-overlay-session", timestamp: new Date().toISOString() }, { @@ -339,7 +340,7 @@ describe("tool path arrays", () => { }, }, ]); - const registry = makeSubagentRegistry([ + const observers = makeSubagentRegistry([ { id: "search-overlay-session", kind: "subagent", @@ -349,9 +350,27 @@ describe("tool path arrays", () => { lastUpdate: Date.now(), }, ]); + const agents = new AgentRegistry(); + agents.register({ + id: "search-overlay-session", + displayName: "search-overlay-session", + kind: "sub", + parentId: "Main", + session: null, + sessionFile, + status: "parked", + }); - const overlay = new SessionObserverOverlayComponent(registry, () => {}, ["ctrl+s"]); - const rendered = Bun.stripANSI(overlay.render(120).join("\n")); + const hub = new AgentHubOverlayComponent({ + observers, + hubKeys: ["ctrl+s"], + onDone: () => {}, + requestRender: () => {}, + registry: agents, + }); + hub.openChat("search-overlay-session"); + const rendered = Bun.stripANSI(hub.render(120).join("\n")); + hub.dispose(); expect(rendered).toContain("paths: folder with spaces/"); }); From a64ff00cb8d27658ca9be1844708e901de3fe4f7 Mon Sep 17 00:00:00 2001 From: can1357 Date: Wed, 10 Jun 2026 17:52:13 +0200 Subject: [PATCH 16/27] feat(coding-agent): added history:// internal URL for agent transcripts Registers a HistoryProtocolHandler with the internal URL router: history:// lists every registered agent (id, status, kind, last activity) and history:// renders a concise markdown transcript (tool calls collapsed to one line each, thinking elided). Live refs render from the in-memory message array; parked refs load read-only from the JSONL session file via loadSessionMessagesReadOnly (no writer, no lock). System prompt + read tool prompt + docs/tools/read.md learn the new scheme. --- docs/tools/read.md | 4 +- .../src/internal-urls/history-protocol.ts | 113 ++++++++ .../coding-agent/src/internal-urls/index.ts | 1 + .../coding-agent/src/internal-urls/router.ts | 4 +- .../coding-agent/src/internal-urls/types.ts | 2 +- .../src/prompts/system/system-prompt.md | 1 + .../coding-agent/src/prompts/tools/read.md | 4 +- .../src/session/session-history-format.ts | 246 ++++++++++++++++++ .../internal-urls/history-protocol.test.ts | 189 ++++++++++++++ .../modes/internal-url-autocomplete.test.ts | 2 +- .../session/session-history-format.test.ts | 108 ++++++++ 11 files changed, 667 insertions(+), 7 deletions(-) create mode 100644 packages/coding-agent/src/internal-urls/history-protocol.ts create mode 100644 packages/coding-agent/src/session/session-history-format.ts create mode 100644 packages/coding-agent/test/internal-urls/history-protocol.test.ts create mode 100644 packages/coding-agent/test/session/session-history-format.test.ts diff --git a/docs/tools/read.md b/docs/tools/read.md index 21edb6e4d..222462c3b 100644 --- a/docs/tools/read.md +++ b/docs/tools/read.md @@ -10,7 +10,7 @@ - `packages/coding-agent/src/tools/archive-reader.ts` — detect `archive.ext:inner/path`, index archives, list/read entries. - `packages/coding-agent/src/tools/sqlite-reader.ts` — detect SQLite targets, parse selectors, render tables. - `packages/coding-agent/src/tools/fetch.ts` — URL parsing, fetch/render pipeline, URL cache/artifacts. - - `packages/coding-agent/src/internal-urls/router.ts` — resolve `agent://`, `artifact://`, `issue://`, `local://`, `mcp://`, `memory://`, `omp://`, `pr://`, `rule://`, `skill://`, and `vault://`. + - `packages/coding-agent/src/internal-urls/router.ts` — resolve `agent://`, `artifact://`, `history://`, `issue://`, `local://`, `mcp://`, `memory://`, `omp://`, `pr://`, `rule://`, `skill://`, and `vault://`. - `packages/coding-agent/src/edit/notebook.ts` — convert `.ipynb` to editable `# %% [...] cell:N` text. - `packages/coding-agent/src/utils/file-display-mode.ts` — decide hashline vs line-number vs raw display. - `packages/coding-agent/src/workspace-tree.ts` — render directory trees. @@ -196,7 +196,7 @@ URL selectors are parsed separately in `packages/coding-agent/src/tools/fetch.ts ### Internal URLs - `read` does not resolve these itself; it delegates to `session.internalRouter.resolve()`. -- Registered protocols are outside this file, but the router in `packages/coding-agent/src/internal-urls/router.ts` is built for `agent://`, `artifact://`, `issue://`, `local://`, `mcp://`, `memory://`, `omp://`, `pr://`, `rule://`, and `skill://`. +- Registered protocols are outside this file, but the router in `packages/coding-agent/src/internal-urls/router.ts` is built for `agent://`, `artifact://`, `history://`, `issue://`, `local://`, `mcp://`, `memory://`, `omp://`, `pr://`, `rule://`, and `skill://`. - `#handleInternalUrl()` behavior: - parses the URL with `parseInternalUrl()` so colons inside the host segment are legal - for `agent://`, treats non-root path extraction or `?q=` extraction as a special no-pagination mode diff --git a/packages/coding-agent/src/internal-urls/history-protocol.ts b/packages/coding-agent/src/internal-urls/history-protocol.ts new file mode 100644 index 000000000..576af1ac1 --- /dev/null +++ b/packages/coding-agent/src/internal-urls/history-protocol.ts @@ -0,0 +1,113 @@ +/** + * Protocol handler for history:// URLs. + * + * Exposes agent transcripts as concise markdown. Live refs render from the + * in-memory message array; parked refs (session disposed, sessionFile + * retained) load read-only from the JSONL session file — no writer, no lock. + * + * URL forms: + * - history:// - Index of all registry agents (id, status, kind, last activity) + * - history:// - Concise markdown transcript of that agent + */ +import type { AgentRef } from "../registry/agent-registry"; +import { AgentRegistry } from "../registry/agent-registry"; +import { formatSessionHistoryMarkdown } from "../session/session-history-format"; +import { loadSessionMessagesReadOnly } from "../session/session-manager"; +import type { InternalResource, InternalUrl, ProtocolHandler, UrlCompletion } from "./types"; + +/** Humanize a last-activity timestamp as `Ns/Nm/Nh/Nd ago`. */ +function formatAgo(timestamp: number): string { + const diffMs = Math.max(0, Date.now() - timestamp); + const secs = Math.floor(diffMs / 1000); + if (secs < 60) return `${secs}s ago`; + const mins = Math.floor(secs / 60); + if (mins < 60) return `${mins}m ago`; + const hours = Math.floor(mins / 60); + if (hours < 24) return `${hours}h ago`; + return `${Math.floor(hours / 24)}d ago`; +} + +/** + * Handler for history:// URLs. + * + * Resolves agent ids against the global AgentRegistry, serving transcripts + * for both live and parked agents. + */ +export class HistoryProtocolHandler implements ProtocolHandler { + readonly scheme = "history"; + readonly immutable = false; + + async resolve(url: InternalUrl): Promise { + const agentId = url.rawHost || url.hostname; + const registry = AgentRegistry.global(); + + if (!agentId) { + const content = this.#renderIndex(registry.list()); + return { + url: url.href, + content, + contentType: "text/markdown", + size: Buffer.byteLength(content, "utf-8"), + }; + } + + let ref = registry.get(agentId); + if (!ref) { + // Case-insensitive fallback: agent ids are human-typed (e.g. AuthLoader). + const lower = agentId.toLowerCase(); + ref = registry.list().find(candidate => candidate.id.toLowerCase() === lower); + } + if (!ref) { + const known = registry.list().map(candidate => candidate.id); + const knownStr = known.length > 0 ? known.join(", ") : "none"; + throw new Error(`Unknown agent: ${agentId}\nKnown agents: ${knownStr}\nList all with history://`); + } + + const notes: string[] = []; + let messages: unknown[]; + if (ref.session) { + messages = ref.session.messages; + notes.push("Source: live session"); + } else if (ref.sessionFile) { + messages = await loadSessionMessagesReadOnly(ref.sessionFile); + notes.push(`Source: session file (read-only, ${ref.status})`); + } else { + throw new Error(`Agent ${ref.id} has no transcript: session is gone and no session file was retained`); + } + + const content = formatSessionHistoryMarkdown(messages, { title: `${ref.id} (${ref.status})` }); + return { + url: url.href, + content, + contentType: "text/markdown", + size: Buffer.byteLength(content, "utf-8"), + sourcePath: ref.sessionFile ?? undefined, + notes, + }; + } + + #renderIndex(refs: AgentRef[]): string { + const lines: string[] = ["# Agents", ""]; + if (refs.length === 0) { + lines.push("No agents registered."); + return `${lines.join("\n")}\n`; + } + lines.push("| id | status | kind | parent | last activity |", "|---|---|---|---|---|"); + for (const ref of refs) { + lines.push( + `| ${ref.id} | ${ref.status} | ${ref.kind} | ${ref.parentId ?? "—"} | ${formatAgo(ref.lastActivity)} |`, + ); + } + lines.push("", "Read a transcript with `read history://`."); + return `${lines.join("\n")}\n`; + } + + async complete(): Promise { + return AgentRegistry.global() + .list() + .map(ref => ({ + value: ref.id, + description: `${ref.status} · ${ref.kind}${ref.parentId ? ` · parent ${ref.parentId}` : ""}`, + })); + } +} diff --git a/packages/coding-agent/src/internal-urls/index.ts b/packages/coding-agent/src/internal-urls/index.ts index a460d8ff8..4714406e4 100644 --- a/packages/coding-agent/src/internal-urls/index.ts +++ b/packages/coding-agent/src/internal-urls/index.ts @@ -10,6 +10,7 @@ export * from "./agent-protocol"; export * from "./artifact-protocol"; +export * from "./history-protocol"; export * from "./issue-pr-protocol"; export * from "./json-query"; export * from "./local-protocol"; diff --git a/packages/coding-agent/src/internal-urls/router.ts b/packages/coding-agent/src/internal-urls/router.ts index 194f9f156..608600ece 100644 --- a/packages/coding-agent/src/internal-urls/router.ts +++ b/packages/coding-agent/src/internal-urls/router.ts @@ -1,5 +1,5 @@ /** - * Internal URL router for internal protocols (`agent://`, `artifact://`, `issue://`, `local://`, `mcp://`, `memory://`, `omp://`, `pr://`, `rule://`, `skill://`, and `vault://`). + * Internal URL router for internal protocols (`agent://`, `artifact://`, `history://`, `issue://`, `local://`, `mcp://`, `memory://`, `omp://`, `pr://`, `rule://`, `skill://`, and `vault://`). * * One process-global router with one handler per scheme. Access via * `InternalUrlRouter.instance()`. Handlers are stateless; per-session and @@ -7,6 +7,7 @@ */ import { AgentProtocolHandler } from "./agent-protocol"; import { ArtifactProtocolHandler } from "./artifact-protocol"; +import { HistoryProtocolHandler } from "./history-protocol"; import { IssueProtocolHandler, PrProtocolHandler } from "./issue-pr-protocol"; import { LocalProtocolHandler } from "./local-protocol"; import { McpProtocolHandler } from "./mcp-protocol"; @@ -35,6 +36,7 @@ export class InternalUrlRouter { this.register(new McpProtocolHandler()); this.register(new IssueProtocolHandler()); this.register(new PrProtocolHandler()); + this.register(new HistoryProtocolHandler()); } /** Process-global router instance. */ diff --git a/packages/coding-agent/src/internal-urls/types.ts b/packages/coding-agent/src/internal-urls/types.ts index 3075b6b16..ef2260057 100644 --- a/packages/coding-agent/src/internal-urls/types.ts +++ b/packages/coding-agent/src/internal-urls/types.ts @@ -1,7 +1,7 @@ /** * Types for the internal URL routing system. * - * Internal URLs (`agent://`, `artifact://`, `issue://`, `local://`, `mcp://`, `memory://`, `omp://`, `pr://`, `rule://`, `skill://`, and `vault://`) are resolved by tools like read, + * Internal URLs (`agent://`, `artifact://`, `history://`, `issue://`, `local://`, `mcp://`, `memory://`, `omp://`, `pr://`, `rule://`, `skill://`, and `vault://`) are resolved by tools like read, * providing access to agent outputs and server resources without exposing filesystem paths. */ diff --git a/packages/coding-agent/src/prompts/system/system-prompt.md b/packages/coding-agent/src/prompts/system/system-prompt.md index db0bd2a07..89218c062 100644 --- a/packages/coding-agent/src/prompts/system/system-prompt.md +++ b/packages/coding-agent/src/prompts/system/system-prompt.md @@ -149,6 +149,7 @@ With most FS/bash-like tools, static references to them will automatically resol - `agent://`: full agent output artifact - `/`: JSON field extraction - `artifact://`: Artifact content +- `history://`: agent transcript as concise markdown; bare `history://` lists agents - `local://.md`: Plan artifacts and shared content with subagents {{#if hasObsidian}} - `vault:///`: Obsidian vault content (read/edit). `vault://` lists vaults; `vault://_/…` targets the active vault. File-scoped `?op=outline|backlinks|links|tags|properties|tasks|base|…`; vault-scoped `?op=search&q=…|daily|tasks|orphans|unresolved|bases|…`. diff --git a/packages/coding-agent/src/prompts/tools/read.md b/packages/coding-agent/src/prompts/tools/read.md index 4a0004298..b95a1a36c 100644 --- a/packages/coding-agent/src/prompts/tools/read.md +++ b/packages/coding-agent/src/prompts/tools/read.md @@ -8,7 +8,7 @@ Read files, directories, archives, SQLite databases, images, documents, internal ## Parameters -- `path` — required. Local path, internal URI (`skill://`, `agent://`, `artifact://`, `memory://`, `rule://`, `local://`, `vault://`, `mcp://`, `omp://`, `issue://`, `pr://`), or URL. Append `:` for line ranges, raw mode, or special modes (e.g. `src/foo.ts:50-200`, `src/foo.ts:raw`, `db.sqlite:users:42`). +- `path` — required. Local path, internal URI (`skill://`, `agent://`, `artifact://`, `history://`, `memory://`, `rule://`, `local://`, `vault://`, `mcp://`, `omp://`, `issue://`, `pr://`), or URL. Append `:` for line ranges, raw mode, or special modes (e.g. `src/foo.ts:50-200`, `src/foo.ts:raw`, `db.sqlite:users:42`). ## Selectors @@ -74,7 +74,7 @@ For `.sqlite`, `.sqlite3`, `.db`, `.db3`: # Internal URIs -`skill://`, `agent://`, `artifact://`, `memory://root`, `rule://`, `local://.md`, `vault:///`, `mcp://`, `omp://.md`, `issue://`, and `pr://` resolve transparently and accept the same line selectors as filesystem paths. Use `artifact://` to recover full output that a previous bash/eval/tool result spilled or truncated. +`skill://`, `agent://`, `artifact://`, `history://`, `memory://root`, `rule://`, `local://.md`, `vault:///`, `mcp://`, `omp://.md`, `issue://`, and `pr://` resolve transparently and accept the same line selectors as filesystem paths. Use `artifact://` to recover full output that a previous bash/eval/tool result spilled or truncated. `history://` is an agent's transcript as concise markdown; bare `history://` lists agents. - You MUST use `read` for every file, directory, archive, and URL inspection. `cat`, `head`, `tail`, `less`, `more`, `ls`, `tar`, `unzip`, `curl`, `wget` are FORBIDDEN — any such bash call is a bug, regardless of how short or convenient it looks. diff --git a/packages/coding-agent/src/session/session-history-format.ts b/packages/coding-agent/src/session/session-history-format.ts new file mode 100644 index 000000000..9588ed1fc --- /dev/null +++ b/packages/coding-agent/src/session/session-history-format.ts @@ -0,0 +1,246 @@ +/** + * Concise markdown transcript serializer for `history://` URLs. + * + * Unlike `session-dump-format.ts` (verbose `/dump` export), this emits a + * compressed transcript: full user/assistant/developer text, tool call + + * result pairs collapsed to single lines, thinking elided, custom messages + * as one-liners. No system prompt, no tool catalog, no config sections. + */ +import type { AgentMessage } from "@oh-my-pi/pi-agent-core"; +import { INTENT_FIELD } from "@oh-my-pi/pi-agent-core"; +import type { AssistantMessage, ImageContent, TextContent, ToolResultMessage } from "@oh-my-pi/pi-ai"; +import type { + BashExecutionMessage, + BranchSummaryMessage, + CompactionSummaryMessage, + CustomMessage, + FileMentionMessage, + HookMessage, + PythonExecutionMessage, +} from "./messages"; + +export interface HistoryFormatOptions { + /** Optional H1 prepended to the transcript. */ + title?: string; +} + +/** Max length of the primary-arg summary inside `→ tool(...)` lines. */ +const PRIMARY_ARG_MAX = 120; + +/** Per-tool preference order for the most informative scalar argument. */ +const PRIMARY_ARG_KEYS = [ + "path", + "file_path", + "filePath", + "command", + "cmd", + "pattern", + "url", + "query", + "prompt", + "assignment", + "message", + "op", + "name", + "id", +] as const; + +/** Collapse whitespace runs and truncate to `max` chars with an ellipsis. */ +function oneLine(text: string, max = PRIMARY_ARG_MAX): string { + const flat = text.replace(/\s+/g, " ").trim(); + return flat.length > max ? `${flat.slice(0, max - 1)}…` : flat; +} + +/** Join the text blocks of a string-or-blocks content field. Images become `[image]`. */ +function contentToText(content: string | readonly (TextContent | ImageContent)[]): string { + if (typeof content === "string") return content; + const parts: string[] = []; + for (const block of content) { + if (block.type === "text") parts.push(block.text); + else parts.push("[image]"); + } + return parts.join("\n"); +} + +function lineCount(text: string): number { + if (!text) return 0; + return text.split("\n").length; +} + +/** Pick the most informative scalar argument of a tool call. */ +function primaryArg(args: Record | undefined): string { + if (!args || typeof args !== "object") return ""; + for (const key of PRIMARY_ARG_KEYS) { + const value = args[key]; + if (typeof value === "string" && value.length > 0) return oneLine(value); + if (Array.isArray(value) && value.length > 0 && value.every(v => typeof v === "string")) { + return oneLine(value.join(", ")); + } + } + // Fallback: first non-intent string arg, then a compact JSON of the args. + const rest: Record = {}; + let restCount = 0; + for (const key in args) { + if (key === INTENT_FIELD) continue; + const value = args[key]; + if (typeof value === "string" && value.length > 0) return oneLine(value); + rest[key] = value; + restCount++; + } + if (restCount === 0) return ""; + try { + return oneLine(JSON.stringify(rest)); + } catch { + return ""; + } +} + +/** One line per tool call: `→ read(src/foo.ts:50-80) ⇒ ok · 31 lines`. */ +function toolCallLine( + name: string, + args: Record | undefined, + result: ToolResultMessage | undefined, +): string { + const head = `→ ${name}(${primaryArg(args)})`; + if (!result) return `${head} ⇒ pending`; + const text = contentToText(result.content); + const lines = lineCount(text); + const count = `${lines} ${lines === 1 ? "line" : "lines"}`; + if (result.isError) { + const firstLine = oneLine(text.split("\n", 1)[0] ?? ""); + return firstLine ? `${head} ⇒ error · ${count} — ${firstLine}` : `${head} ⇒ error · ${count}`; + } + return `${head} ⇒ ok · ${count}`; +} + +/** One line for a user-initiated `!`/`$` execution. */ +function executionLine( + kind: "bash" | "python", + source: string, + msg: BashExecutionMessage | PythonExecutionMessage, +): string { + const status = msg.cancelled + ? "cancelled" + : msg.exitCode !== undefined && msg.exitCode !== 0 + ? `error · exit ${msg.exitCode}` + : "ok"; + const lines = lineCount(msg.output); + return `→ ${kind}! ${oneLine(source)} ⇒ ${status} · ${lines} ${lines === 1 ? "line" : "lines"}`; +} + +/** One-liner for custom/hook messages: `[irc] A → B: body…`. */ +function customOneLiner(msg: CustomMessage | HookMessage): string { + const details = (msg.details ?? {}) as Record; + const str = (key: string): string => (typeof details[key] === "string" ? (details[key] as string) : ""); + switch (msg.customType) { + case "irc:incoming": + return `[irc] ${str("from") || "?"} → me: ${oneLine(str("message"))}`; + case "irc:relay": + return `[irc] ${str("from") || "?"} → ${str("to") || "?"}: ${oneLine(str("body"))}`; + case "async-result": { + const jobs = Array.isArray(details.jobs) && details.jobs.length > 0 ? details.jobs : [details]; + const labels = jobs + .map(job => { + const j = (job ?? {}) as Record; + return typeof j.label === "string" && j.label ? j.label : typeof j.jobId === "string" ? j.jobId : "job"; + }) + .join(", "); + return `[async-result] ${oneLine(labels)}`; + } + default: + return `[${msg.customType}] ${oneLine(contentToText(msg.content))}`; + } +} + +/** + * Format a session's message array as a concise markdown transcript. + * + * `messages` is the session's in-memory message array (or the read-only + * equivalent loaded from a session file) — the same shapes + * `session-dump-format.ts` consumes. + */ +export function formatSessionHistoryMarkdown(messages: unknown[], opts?: HistoryFormatOptions): string { + const typed = messages as AgentMessage[]; + const lines: string[] = []; + if (opts?.title) { + lines.push(`# ${opts.title}`, ""); + } + + // Index tool results by call id so each toolCall collapses to one line. + const resultsByCallId = new Map(); + for (const msg of typed) { + if (msg.role === "toolResult") { + resultsByCallId.set(msg.toolCallId, msg); + } + } + const consumed = new Set(); + + for (const msg of typed) { + switch (msg.role) { + case "user": + case "developer": { + const text = contentToText(msg.content); + if (!text.trim()) break; + lines.push(`## ${msg.role}`, "", text, ""); + break; + } + case "assistant": { + const assistantMsg = msg as AssistantMessage; + const body: string[] = []; + for (const block of assistantMsg.content) { + if (block.type === "text") { + if (block.text.trim()) body.push(block.text); + } else if (block.type === "toolCall") { + const result = resultsByCallId.get(block.id); + if (result) consumed.add(block.id); + body.push(toolCallLine(block.name, block.arguments, result)); + } + // thinking / redactedThinking elided entirely + } + if (body.length === 0) break; + lines.push("## assistant", "", ...body, ""); + break; + } + case "toolResult": { + // Normally consumed by its toolCall; orphans (e.g. truncated history) get their own line. + if (consumed.has(msg.toolCallId)) break; + lines.push(toolCallLine(msg.toolName, undefined, msg), ""); + break; + } + case "bashExecution": { + const bashMsg = msg as BashExecutionMessage; + if (bashMsg.excludeFromContext) break; + lines.push(executionLine("bash", bashMsg.command, bashMsg), ""); + break; + } + case "pythonExecution": { + const pythonMsg = msg as PythonExecutionMessage; + if (pythonMsg.excludeFromContext) break; + lines.push(executionLine("python", pythonMsg.code, pythonMsg), ""); + break; + } + case "custom": + case "hookMessage": { + lines.push(customOneLiner(msg as CustomMessage | HookMessage), ""); + break; + } + case "branchSummary": { + const branchMsg = msg as BranchSummaryMessage; + lines.push(`[branch] from ${branchMsg.fromId}: ${oneLine(branchMsg.summary)}`, ""); + break; + } + case "compactionSummary": { + const compactMsg = msg as CompactionSummaryMessage; + lines.push(`[compaction] ${oneLine(compactMsg.summary)}`, ""); + break; + } + case "fileMention": { + const fileMsg = msg as FileMentionMessage; + lines.push(`[file-mention] ${oneLine(fileMsg.files.map(f => f.path).join(", "))}`, ""); + break; + } + } + } + + return `${lines.join("\n").trim()}\n`; +} diff --git a/packages/coding-agent/test/internal-urls/history-protocol.test.ts b/packages/coding-agent/test/internal-urls/history-protocol.test.ts new file mode 100644 index 000000000..8f7f8d6b0 --- /dev/null +++ b/packages/coding-agent/test/internal-urls/history-protocol.test.ts @@ -0,0 +1,189 @@ +/** + * Contracts: history:// protocol handler (rework-contracts.md §6), resolved + * through `InternalUrlRouter.instance().resolve(...)` like real callers. + * + * - Bare `history://` renders an index listing registered agent ids. + * - `history://` with a live ref renders the in-memory transcript. + * - A parked ref (session null, sessionFile retained) renders read-only from + * the JSONL session file. + * - An unknown id fails with an error listing the known ids. + */ +import { afterEach, beforeEach, describe, expect, it } from "bun:test"; +import * as fs from "node:fs/promises"; +import * as os from "node:os"; +import * as path from "node:path"; +import { InternalUrlRouter } from "@oh-my-pi/pi-coding-agent/internal-urls"; +import { AgentRegistry } from "@oh-my-pi/pi-coding-agent/registry/agent-registry"; +import type { AgentSession } from "@oh-my-pi/pi-coding-agent/session/agent-session"; +import { CURRENT_SESSION_VERSION } from "@oh-my-pi/pi-coding-agent/session/session-manager"; + +async function withTempDir(fn: (dir: string) => Promise): Promise { + const dir = await fs.mkdtemp(path.join(os.tmpdir(), "history-protocol-")); + try { + return await fn(dir); + } finally { + await fs.rm(dir, { recursive: true, force: true }); + } +} + +function fakeLiveSession(messages: unknown[]): AgentSession { + return { messages } as unknown as AgentSession; +} + +/** Minimal current-version session JSONL: header + a linear user/assistant chain. */ +function sessionFixtureJsonl(): string { + const timestamp = new Date().toISOString(); + const header = { + type: "session", + version: CURRENT_SESSION_VERSION, + id: "fixture-session", + timestamp, + cwd: "/tmp", + }; + const userEntry = { + type: "message", + id: "m1", + parentId: null, + timestamp, + message: { role: "user", content: "parked hello", timestamp: 1 }, + }; + const assistantEntry = { + type: "message", + id: "m2", + parentId: "m1", + timestamp, + message: { + role: "assistant", + content: [{ type: "text", text: "parked reply" }], + api: "anthropic-messages", + provider: "anthropic", + model: "test-model", + usage: {}, + stopReason: "stop", + timestamp: 2, + }, + }; + return `${JSON.stringify(header)}\n${JSON.stringify(userEntry)}\n${JSON.stringify(assistantEntry)}\n`; +} + +describe("history:// protocol", () => { + beforeEach(() => { + AgentRegistry.resetGlobalForTests(); + InternalUrlRouter.resetForTests(); + }); + + afterEach(() => { + InternalUrlRouter.resetForTests(); + AgentRegistry.resetGlobalForTests(); + }); + + it("bare history:// renders an index listing registered agents", async () => { + AgentRegistry.global().register({ + id: "HubAgent", + displayName: "task", + kind: "sub", + session: fakeLiveSession([]), + status: "idle", + }); + + const resource = await InternalUrlRouter.instance().resolve("history://"); + + expect(resource.contentType).toBe("text/markdown"); + expect(resource.content).toContain("# Agents"); + expect(resource.content).toContain("| HubAgent | idle | sub |"); + }); + + it("history:// renders a live ref's in-memory transcript", async () => { + AgentRegistry.global().register({ + id: "HubAgent", + displayName: "task", + kind: "sub", + session: fakeLiveSession([{ role: "user", content: "hello from live", timestamp: 1 }]), + status: "idle", + }); + + const resource = await InternalUrlRouter.instance().resolve("history://HubAgent"); + + expect(resource.content).toContain("# HubAgent (idle)"); + expect(resource.content).toContain("## user"); + expect(resource.content).toContain("hello from live"); + expect(resource.notes).toContain("Source: live session"); + }); + + it("resolves agent ids case-insensitively", async () => { + AgentRegistry.global().register({ + id: "HubAgent", + displayName: "task", + kind: "sub", + session: fakeLiveSession([{ role: "user", content: "hello from live", timestamp: 1 }]), + status: "idle", + }); + + const resource = await InternalUrlRouter.instance().resolve("history://hubagent"); + expect(resource.content).toContain("# HubAgent (idle)"); + }); + + it("history:// renders a parked ref read-only from its session file", async () => { + await withTempDir(async dir => { + const sessionFile = path.join(dir, "parked.jsonl"); + await Bun.write(sessionFile, sessionFixtureJsonl()); + AgentRegistry.global().register({ + id: "Sleeper", + displayName: "task", + kind: "sub", + session: null, + sessionFile, + status: "parked", + }); + + const resource = await InternalUrlRouter.instance().resolve("history://Sleeper"); + + expect(resource.content).toContain("# Sleeper (parked)"); + expect(resource.content).toContain("parked hello"); + expect(resource.content).toContain("parked reply"); + expect(resource.sourcePath).toBe(sessionFile); + expect(resource.notes?.join("\n")).toContain("read-only"); + }); + }); + + it("rejects an unknown id with the list of known agents", async () => { + AgentRegistry.global().register({ + id: "HubAgent", + displayName: "task", + kind: "sub", + session: fakeLiveSession([]), + status: "idle", + }); + + const error = await InternalUrlRouter.instance() + .resolve("history://Nope") + .then( + () => null, + err => err as Error, + ); + + expect(error).toBeInstanceOf(Error); + expect(error?.message).toContain("Unknown agent: Nope"); + expect(error?.message).toContain("HubAgent"); + }); + + it("rejects a ref with neither session nor session file", async () => { + AgentRegistry.global().register({ + id: "Husk", + displayName: "task", + kind: "sub", + session: null, + sessionFile: null, + status: "aborted", + }); + + const error = await InternalUrlRouter.instance() + .resolve("history://Husk") + .then( + () => null, + err => err as Error, + ); + + expect(error?.message).toContain("no transcript"); + }); +}); diff --git a/packages/coding-agent/test/modes/internal-url-autocomplete.test.ts b/packages/coding-agent/test/modes/internal-url-autocomplete.test.ts index f111051db..34f69698c 100644 --- a/packages/coding-agent/test/modes/internal-url-autocomplete.test.ts +++ b/packages/coding-agent/test/modes/internal-url-autocomplete.test.ts @@ -125,7 +125,7 @@ describe("internal-url-autocomplete", () => { it("exposes the completion-capable schemes", () => { const schemes = InternalUrlRouter.instance().completionSchemes().sort(); - expect(schemes).toEqual(["agent", "artifact", "local", "memory", "omp", "rule", "skill"]); + expect(schemes).toEqual(["agent", "artifact", "history", "local", "memory", "omp", "rule", "skill"]); }); }); diff --git a/packages/coding-agent/test/session/session-history-format.test.ts b/packages/coding-agent/test/session/session-history-format.test.ts new file mode 100644 index 000000000..dc32cb1cb --- /dev/null +++ b/packages/coding-agent/test/session/session-history-format.test.ts @@ -0,0 +1,108 @@ +/** + * Contracts: history:// transcript serializer (rework-contracts.md §5). + * + * - `## user` / `## assistant` headers carry full text. + * - Thinking blocks are elided entirely. + * - Each toolCall collapses with its toolResult into ONE `→ name(…) ⇒ …` + * line (ok and error variants); result bodies are never dumped. + * - Custom messages render as one-liners (`[irc] from → me: …`). + * - No system prompt / tool catalog sections. + */ +import { describe, expect, it } from "bun:test"; +import { formatSessionHistoryMarkdown } from "@oh-my-pi/pi-coding-agent/session/session-history-format"; + +function buildMessages(): unknown[] { + return [ + { role: "user", content: "Please read the config.", timestamp: 1 }, + { + role: "assistant", + content: [ + { type: "thinking", thinking: "SECRET-THOUGHT about the approach" }, + { type: "text", text: "Reading it now." }, + { type: "toolCall", id: "tc-1", name: "read", arguments: { path: "src/config.ts" } }, + { type: "toolCall", id: "tc-2", name: "bash", arguments: { command: "bun test" } }, + ], + api: "anthropic-messages", + provider: "anthropic", + model: "test-model", + usage: {}, + stopReason: "toolUse", + timestamp: 2, + }, + { + role: "toolResult", + toolCallId: "tc-1", + toolName: "read", + content: [{ type: "text", text: "const a = 1;\nconst b = 2;\nconst c = 3;" }], + isError: false, + timestamp: 3, + }, + { + role: "toolResult", + toolCallId: "tc-2", + toolName: "bash", + content: [{ type: "text", text: "FAIL: 1 test failed" }], + isError: true, + timestamp: 4, + }, + { + role: "custom", + customType: "irc:incoming", + content: "full rendered irc prompt that must not appear", + details: { from: "Main", message: "status update please" }, + timestamp: 5, + }, + ]; +} + +describe("formatSessionHistoryMarkdown", () => { + it("renders role headers, collapses tool pairs to one line, and elides thinking", () => { + const output = formatSessionHistoryMarkdown(buildMessages()); + + expect(output).toContain("## user"); + expect(output).toContain("Please read the config."); + expect(output).toContain("## assistant"); + expect(output).toContain("Reading it now."); + + // Thinking is elided entirely. + expect(output).not.toContain("SECRET-THOUGHT"); + + // Tool call + result collapse to one line each; bodies are not dumped. + expect(output).toContain("→ read(src/config.ts) ⇒ ok · 3 lines"); + expect(output).not.toContain("const a = 1;"); + + // Error variant carries the first line of the error output. + expect(output).toContain("→ bash(bun test) ⇒ error · 1 line — FAIL: 1 test failed"); + + // Consumed toolResults do not render a second orphan line. + const toolLines = output.split("\n").filter(line => line.startsWith("→ ")); + expect(toolLines).toHaveLength(2); + + // Custom messages are one-liners; the rendered prompt body is dropped. + expect(output).toContain("[irc] Main → me: status update please"); + expect(output).not.toContain("full rendered irc prompt"); + + // Concise transcript: no prompt/tool-catalog sections. + expect(output).not.toContain("System Prompt"); + expect(output).not.toContain("Available Tools"); + }); + + it("prefixes an H1 title when requested", () => { + const output = formatSessionHistoryMarkdown(buildMessages(), { title: "Spawnling (idle)" }); + expect(output.startsWith("# Spawnling (idle)\n")).toBe(true); + }); + + it("renders an orphan toolResult (truncated history) as its own line", () => { + const output = formatSessionHistoryMarkdown([ + { + role: "toolResult", + toolCallId: "tc-orphan", + toolName: "search", + content: [{ type: "text", text: "one match" }], + isError: false, + timestamp: 1, + }, + ]); + expect(output).toContain("→ search() ⇒ ok · 1 line"); + }); +}); From 9f62c7904a276b816d8150445f39a979c0476951 Mon Sep 17 00:00:00 2001 From: can1357 Date: Wed, 10 Jun 2026 17:52:49 +0200 Subject: [PATCH 17/27] feat(coding-agent): integrated snapcompact strategy and per-turn supersede pruning Adds compaction.strategy: "snapcompact" to the schema and the AgentSession routing: when chosen, both manual /compact (without custom instructions) and auto compaction call snapcompactCompact() to archive history as PNG frames instead of an LLM summary. Falls back to context-full with a visible warning notice when the current model is text-only or when /compact gets custom instructions. CustomTool and shared-event payloads carry the new action through. \n\nAlso wires the per-turn supersede pass: #pruneSupersededReads() runs every turn before threshold gating (cache-aware: only fires when the post-candidate suffix is small or the prompt cache is cold), prunes older read results superseded by a newer read of the same file, rewrites the session, and accounts the saved tokens in the next compaction decision. Gated by compaction.supersedeReads (default on).\n\nsession/messages.ts now delegates the core role conversion to agent-core's convertMessageToLlm so snapcompact image blocks flow through the LLM-context conversion path. --- .../src/config/settings-schema.ts | 11 +- .../src/extensibility/custom-tools/types.ts | 4 +- .../src/extensibility/shared-events.ts | 4 +- .../coding-agent/src/session/agent-session.ts | 100 +++++++++++++++++- packages/coding-agent/src/session/messages.ts | 89 ++-------------- packages/coding-agent/test/compaction.test.ts | 52 +++++++++ .../test/session-messages.test.ts | 43 +++++++- 7 files changed, 212 insertions(+), 91 deletions(-) diff --git a/packages/coding-agent/src/config/settings-schema.ts b/packages/coding-agent/src/config/settings-schema.ts index 0e45a27b4..3b2edc579 100644 --- a/packages/coding-agent/src/config/settings-schema.ts +++ b/packages/coding-agent/src/config/settings-schema.ts @@ -1177,13 +1177,13 @@ export const SETTINGS_SCHEMA = { "compaction.strategy": { type: "enum", - values: ["context-full", "handoff", "shake", "off"] as const, + values: ["context-full", "handoff", "shake", "snapcompact", "off"] as const, default: "context-full", ui: { tab: "context", label: "Compaction Strategy", description: - "Choose in-place context-full maintenance, auto-handoff, surgical shake (drop heavy content), or disable auto maintenance (off)", + "Choose in-place context-full maintenance, auto-handoff, surgical shake (drop heavy content), snapcompact (archive history as dense images), or disable auto maintenance (off)", options: [ { value: "context-full", @@ -1196,6 +1196,11 @@ export const SETTINGS_SCHEMA = { label: "Shake", description: "Drop heavy content (tool results + large blocks) in place; recover via artifact", }, + { + value: "snapcompact", + label: "Snapcompact", + description: "Archive history onto dense bitmap images the model reads back; no LLM call", + }, { value: "off", label: "Off", @@ -3363,7 +3368,7 @@ export type TreeFilterMode = SettingValue<"treeFilterMode">; export interface CompactionSettings { enabled: boolean; - strategy: "context-full" | "handoff" | "shake" | "off"; + strategy: "context-full" | "handoff" | "shake" | "snapcompact" | "off"; thresholdPercent: number; thresholdTokens: number; reserveTokens: number; diff --git a/packages/coding-agent/src/extensibility/custom-tools/types.ts b/packages/coding-agent/src/extensibility/custom-tools/types.ts index 3703cfaab..41c496287 100644 --- a/packages/coding-agent/src/extensibility/custom-tools/types.ts +++ b/packages/coding-agent/src/extensibility/custom-tools/types.ts @@ -103,11 +103,11 @@ export type CustomToolSessionEvent = | { reason: "auto_compaction_start"; trigger: "threshold" | "overflow" | "idle" | "incomplete"; - action: "context-full" | "handoff" | "shake"; + action: "context-full" | "handoff" | "shake" | "snapcompact"; } | { reason: "auto_compaction_end"; - action: "context-full" | "handoff" | "shake"; + action: "context-full" | "handoff" | "shake" | "snapcompact"; result: CompactionResult | undefined; aborted: boolean; willRetry: boolean; diff --git a/packages/coding-agent/src/extensibility/shared-events.ts b/packages/coding-agent/src/extensibility/shared-events.ts index cefc1ea94..713fd54bf 100644 --- a/packages/coding-agent/src/extensibility/shared-events.ts +++ b/packages/coding-agent/src/extensibility/shared-events.ts @@ -204,13 +204,13 @@ export interface TurnEndEvent { export interface AutoCompactionStartEvent { type: "auto_compaction_start"; reason: "threshold" | "overflow" | "idle" | "incomplete"; - action: "context-full" | "handoff" | "shake"; + action: "context-full" | "handoff" | "shake" | "snapcompact"; } /** Fired when auto-compaction ends */ export interface AutoCompactionEndEvent { type: "auto_compaction_end"; - action: "context-full" | "handoff" | "shake"; + action: "context-full" | "handoff" | "shake" | "snapcompact"; result: CompactionResult | undefined; aborted: boolean; willRetry: boolean; diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index 4805273f8..d0746f35a 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/packages/coding-agent/src/session/agent-session.ts @@ -55,8 +55,14 @@ import { type ShakeRegion, type SummaryOptions, shouldCompact, + snapcompactCompact, } from "@oh-my-pi/pi-agent-core/compaction"; -import { DEFAULT_PRUNE_CONFIG, pruneToolOutputs } from "@oh-my-pi/pi-agent-core/compaction/pruning"; +import { + DEFAULT_PRUNE_CONFIG, + pruneSupersededToolResults, + pruneToolOutputs, + readToolSupersedeKey, +} from "@oh-my-pi/pi-agent-core/compaction/pruning"; import type { ProtectedToolMatcher } from "@oh-my-pi/pi-agent-core/compaction/tool-protection"; import type { AssistantMessage, @@ -258,11 +264,11 @@ export type AgentSessionEvent = | { type: "auto_compaction_start"; reason: "threshold" | "overflow" | "idle" | "incomplete"; - action: "context-full" | "handoff" | "shake"; + action: "context-full" | "handoff" | "shake" | "snapcompact"; } | { type: "auto_compaction_end"; - action: "context-full" | "handoff" | "shake"; + action: "context-full" | "handoff" | "shake" | "snapcompact"; result: CompactionResult | undefined; aborted: boolean; willRetry: boolean; @@ -6071,6 +6077,35 @@ export class AgentSession { return result; } + /** + * Per-turn supersede pass: prune older `read` results that a newer read of + * the same file has made stale. Cache-aware (only fires when the suffix + * after a candidate is small or the session has been idle long enough that + * the provider prompt cache is cold), so it is cheap to run every turn. + * Gated on the `compaction.supersedeReads` setting. + */ + async #pruneSupersededReads(): Promise<{ prunedCount: number; tokensSaved: number } | undefined> { + if (!this.settings.getGroup("compaction").supersedeReads) return undefined; + const branchEntries = this.sessionManager.getBranch(); + const result = pruneSupersededToolResults( + branchEntries, + this.#withPlanProtection({ + supersedeKey: readToolSupersedeKey, + protectedTools: [...DEFAULT_PRUNE_CONFIG.protectedTools], + }), + ); + if (result.prunedCount === 0) { + return undefined; + } + + await this.sessionManager.rewriteEntries(); + const sessionContext = this.buildDisplaySessionContext(); + this.agent.replaceMessages(sessionContext.messages); + this.#syncTodoPhasesFromBranch(); + this.#closeCodexProviderSessionsForHistoryRewrite(); + return result; + } + /** * Strip image content blocks from every message on the current branch and * persist the rewrite. Walks `SessionManager.getBranch()` in place — both @@ -6260,6 +6295,20 @@ export class AgentSession { const compactionPrep = await this.#prepareCompactionFromHooks(preparation, hookCompaction); + // Strategy honored on manual /compact too. Custom instructions imply a + // directed LLM summary; a text-only model cannot read the frames back — + // both take the summarizer path (the latter loudly). + const wantsSnapcompact = + compactionPrep.kind !== "fromHook" && compactionSettings.strategy === "snapcompact" && !customInstructions; + const snapcompactReady = wantsSnapcompact && this.model.input.includes("image"); + if (wantsSnapcompact && !snapcompactReady) { + this.emitNotice( + "warning", + `snapcompact needs a vision-capable model (${this.model.id} is text-only) — using an LLM summary instead`, + "compaction", + ); + } + let summary: string; let shortSummary: string | undefined; let firstKeptEntryId: string; @@ -6273,6 +6322,14 @@ export class AgentSession { tokensBefore = compactionPrep.tokensBefore; details = compactionPrep.details; preserveData = compactionPrep.preserveData; + } else if (snapcompactReady) { + const snapcompactResult = await snapcompactCompact(preparation, { convertToLlm }); + summary = snapcompactResult.summary; + shortSummary = snapcompactResult.shortSummary; + firstKeptEntryId = snapcompactResult.firstKeptEntryId; + tokensBefore = snapcompactResult.tokensBefore; + details = snapcompactResult.details; + preserveData = { ...(compactionPrep.preserveData ?? {}), ...(snapcompactResult.preserveData ?? {}) }; } else { // Generate compaction result. Only convert known abort-shaped // rejections (AbortError raised while the abort signal is set, @@ -6703,6 +6760,10 @@ export class AgentSession { return false; } + // Supersede pass runs every turn, before any threshold gating: it is cheap + // (bails when no candidate) and independent of the compaction setting. + const supersedeResult = await this.#pruneSupersededReads(); + const compactionSettings = this.settings.getGroup("compaction"); if (!compactionSettings.enabled || compactionSettings.strategy === "off") return false; @@ -6711,6 +6772,9 @@ export class AgentSession { if (assistantMessage.stopReason === "error") return false; const pruneResult = await this.#pruneToolOutputs(); let contextTokens = calculateContextTokens(assistantMessage.usage); + if (supersedeResult) { + contextTokens = Math.max(0, contextTokens - supersedeResult.tokensSaved); + } if (pruneResult) { contextTokens = Math.max(0, contextTokens - pruneResult.tokensSaved); } @@ -7601,9 +7665,25 @@ export class AgentSession { // "overflow" forces context-full because the input itself is broken — a handoff // LLM call would hit the same overflow. "incomplete" is an output-side problem, - // so a handoff request on the existing context is still viable. - let action: "context-full" | "handoff" = + // so a handoff request on the existing context is still viable. Snapcompact is + // safe for every reason (it makes no LLM call at all) but requires a vision + // model to be worth anything — fall back to context-full otherwise. + let action: "context-full" | "handoff" | "snapcompact" = compactionSettings.strategy === "handoff" && reason !== "overflow" ? "handoff" : "context-full"; + if (compactionSettings.strategy === "snapcompact") { + if (this.model?.input.includes("image")) { + action = "snapcompact"; + } else { + logger.warn("Snapcompact compaction requires a vision-capable model; falling back to context-full", { + model: this.model?.id, + }); + this.emitNotice( + "warning", + `snapcompact needs a vision-capable model (${this.model?.id ?? "unknown"} is text-only) — using an LLM summary instead`, + "compaction", + ); + } + } await this.#emitSessionEvent({ type: "auto_compaction_start", reason, action }); // Abort any older auto-compaction before installing this run's controller. this.#autoCompactionAbortController?.abort(); @@ -7742,6 +7822,16 @@ export class AgentSession { tokensBefore = compactionPrep.tokensBefore; details = compactionPrep.details; preserveData = compactionPrep.preserveData; + } else if (action === "snapcompact") { + // Local, deterministic: render discarded history onto PNG frames. + // No model candidates, no API key, no retry loop. + const snapcompactResult = await snapcompactCompact(preparation, { convertToLlm }); + summary = snapcompactResult.summary; + shortSummary = snapcompactResult.shortSummary; + firstKeptEntryId = snapcompactResult.firstKeptEntryId; + tokensBefore = snapcompactResult.tokensBefore; + details = snapcompactResult.details; + preserveData = { ...(compactionPrep.preserveData ?? {}), ...(snapcompactResult.preserveData ?? {}) }; } else { const candidates = this.#getCompactionModelCandidates(availableModels); const retrySettings = this.settings.getGroup("retry"); diff --git a/packages/coding-agent/src/session/messages.ts b/packages/coding-agent/src/session/messages.ts index 838d861b0..c19089cfa 100644 --- a/packages/coding-agent/src/session/messages.ts +++ b/packages/coding-agent/src/session/messages.ts @@ -8,8 +8,7 @@ import type { AgentMessage } from "@oh-my-pi/pi-agent-core"; import { type BranchSummaryMessage, type CompactionSummaryMessage, - renderBranchSummaryContext, - renderCompactionSummaryContext, + convertMessageToLlm, } from "@oh-my-pi/pi-agent-core/compaction/messages"; import type { AssistantMessage, @@ -17,7 +16,6 @@ import type { Message, MessageAttribution, TextContent, - ToolResultMessage, UserMessage, } from "@oh-my-pi/pi-ai"; import { prompt } from "@oh-my-pi/pi-utils"; @@ -28,6 +26,7 @@ export { type CompactionSummaryMessage, createBranchSummaryMessage, createCompactionSummaryMessage, + createCustomMessage, } from "@oh-my-pi/pi-agent-core/compaction/messages"; import type { OutputMeta } from "../tools/output-meta"; @@ -59,7 +58,7 @@ export interface SkillPromptDetails { * * Consumers: `AgentSession.#handleAgentEvent` (stamper) writes this value; * `EventController.#handleMessageEnd`, `AssistantMessageComponent`, - * `ui-helpers.addMessageToChat` (renderers), `SessionObserverOverlay + * `ui-helpers.addMessageToChat` (renderers), `AgentHubOverlayComponent * #buildTranscriptLines`, `runPrintMode`, and `AcpAgent#replayAssistantMessage` * (fallback error emission) read it via `isSilentAbort`. */ export const SILENT_ABORT_MARKER = "__omp.silent_abort__"; @@ -220,15 +219,6 @@ export function wrapSteeringForModel(messages: AgentMessage[]): AgentMessage[] { return wrappedMessages ?? messages; } -function getPrunedToolResultContent(message: ToolResultMessage): (TextContent | ImageContent)[] { - if (message.prunedAt === undefined) { - return message.content; - } - const textBlocks = message.content.filter((content): content is TextContent => content.type === "text"); - const text = textBlocks.map(block => block.text).join("") || "[Output truncated]"; - return [{ type: "text", text }]; -} - /** Result of filtering image blocks out of a `(TextContent | ImageContent)[]` array. */ interface StripContentResult { content: (TextContent | ImageContent)[]; @@ -478,26 +468,6 @@ export function sanitizeRehydratedOpenAIResponsesAssistantMessage(message: Assis }; } -/** Convert CustomMessageEntry to AgentMessage format */ -export function createCustomMessage( - customType: string, - content: string | (TextContent | ImageContent)[], - display: boolean, - details: unknown | undefined, - timestamp: string, - attribution?: MessageAttribution, -): CustomMessage { - return { - role: "custom", - customType, - content, - display, - details, - attribution, - timestamp: new Date(timestamp).getTime(), - }; -} - /** * Transform AgentMessages (including custom types) to LLM-compatible Messages. * @@ -530,43 +500,6 @@ export function convertToLlm(messages: AgentMessage[]): Message[] { attribution: "user", timestamp: m.timestamp, }; - case "custom": - case "hookMessage": { - const content = typeof m.content === "string" ? [{ type: "text" as const, text: m.content }] : m.content; - const role = "developer"; - const attribution = m.attribution; - return { - role, - content, - attribution, - timestamp: m.timestamp, - }; - } - case "branchSummary": - return { - role: "user", - content: [ - { - type: "text" as const, - text: renderBranchSummaryContext(m.summary), - }, - ], - attribution: "agent", - timestamp: m.timestamp, - }; - case "compactionSummary": - return { - role: "user", - content: [ - { - type: "text" as const, - text: renderCompactionSummaryContext(m.summary), - }, - ], - attribution: "agent", - providerPayload: m.providerPayload, - timestamp: m.timestamp, - }; case "fileMention": { const fileContents = m.files .map(file => { @@ -587,18 +520,18 @@ export function convertToLlm(messages: AgentMessage[]): Message[] { timestamp: m.timestamp, }; } + case "custom": + case "hookMessage": + case "branchSummary": + case "compactionSummary": case "user": - return { ...m, attribution: m.attribution ?? "user" }; case "developer": - return { ...m, attribution: m.attribution ?? "agent" }; case "assistant": - return m; case "toolResult": - return { - ...m, - content: getPrunedToolResultContent(m as ToolResultMessage), - attribution: m.attribution ?? "agent", - }; + // Core roles share one transformer with agent-core — + // duplicating them here is how snapcompact frames once + // silently fell off the provider request. + return convertMessageToLlm(m); default: m satisfies never; return undefined; diff --git a/packages/coding-agent/test/compaction.test.ts b/packages/coding-agent/test/compaction.test.ts index 8c7912447..24c156413 100644 --- a/packages/coding-agent/test/compaction.test.ts +++ b/packages/coding-agent/test/compaction.test.ts @@ -814,6 +814,58 @@ describe("buildSessionContext", () => { expect((loaded.messages[0] as any).summary).toContain("Summary of 1,a,2,b"); }); + it("re-attaches snapcompact frames from preserveData as compaction summary images", () => { + const u1 = createMessageEntry(createUserMessage("1")); + const a1 = createMessageEntry(createAssistantMessage("a")); + const u2 = createMessageEntry(createUserMessage("2")); + const frame = { data: "ZmFrZQ==", mimeType: "image/png", cols: 64, rows: 40, chars: 4 }; + const compaction: CompactionEntry = { + ...createCompactionEntry("Filmed summary", u2.id), + preserveData: { snapcompact: { frames: [frame], totalChars: 4, truncatedChars: 0 } }, + }; + const u3 = createMessageEntry(createUserMessage("3")); + + const loaded = buildSessionContext([u1, a1, u2, compaction, u3]); + const summaryMessage = loaded.messages[0] as { role: string; images?: unknown }; + expect(summaryMessage.role).toBe("compactionSummary"); + expect(summaryMessage.images).toEqual([{ type: "image", data: "ZmFrZQ==", mimeType: "image/png" }]); + }); + + it("transcript option keeps full history with every compaction inline at its position", () => { + const u1 = createMessageEntry(createUserMessage("1")); + const a1 = createMessageEntry(createAssistantMessage("a")); + const compact1 = createCompactionEntry("First summary", u1.id); + const u2 = createMessageEntry(createUserMessage("2")); + const frame = { data: "ZmFrZQ==", mimeType: "image/png", cols: 64, rows: 40, chars: 4 }; + const compact2: CompactionEntry = { + ...createCompactionEntry("Second summary", u2.id), + preserveData: { snapcompact: { frames: [frame], totalChars: 4, truncatedChars: 0 } }, + }; + const u3 = createMessageEntry(createUserMessage("3")); + const entries: SessionEntry[] = [u1, a1, compact1, u2, compact2, u3]; + + const transcript = buildSessionContext(entries, undefined, undefined, { transcript: true }); + // Nothing erased: every message survives, compactions sit where they fired. + expect(transcript.messages.map(m => m.role)).toEqual([ + "user", + "assistant", + "compactionSummary", + "user", + "compactionSummary", + "user", + ]); + const first = transcript.messages[2] as { summary: string }; + const second = transcript.messages[4] as { summary: string; images?: unknown }; + expect(first.summary).toContain("First summary"); + expect(second.summary).toContain("Second summary"); + // Snapcompact frames ride along in the transcript too. + expect(second.images).toEqual([{ type: "image", data: "ZmFrZQ==", mimeType: "image/png" }]); + + // LLM context is untouched by the option: latest compaction replaces history. + const llm = buildSessionContext(entries); + expect(llm.messages.map(m => m.role)).toEqual(["compactionSummary", "user", "user"]); + }); + it("should handle multiple compactions (only latest matters)", () => { // First batch const u1 = createMessageEntry(createUserMessage("1")); diff --git a/packages/coding-agent/test/session-messages.test.ts b/packages/coding-agent/test/session-messages.test.ts index 4297b24db..14bfde91f 100644 --- a/packages/coding-agent/test/session-messages.test.ts +++ b/packages/coding-agent/test/session-messages.test.ts @@ -1,6 +1,6 @@ import { describe, expect, it } from "bun:test"; import type { AgentMessage } from "@oh-my-pi/pi-agent-core"; -import type { ImageContent, Message } from "@oh-my-pi/pi-ai"; +import type { ImageContent, Message, TextContent } from "@oh-my-pi/pi-ai"; import { inferCopilotInitiator } from "@oh-my-pi/pi-ai/providers/github-copilot-headers"; import { convertToLlm, wrapSteeringForModel } from "@oh-my-pi/pi-coding-agent/session/messages"; @@ -13,6 +13,47 @@ function expectAttribution(message: Message | undefined, expected: "user" | "age expect(message.attribution).toBe(expected); } +describe("convertToLlm compaction summary", () => { + it("appends snapcompact frames as image blocks after the summary text", () => { + // Regression: the live session uses THIS converter (not agent-core's + // defaultConvertToLlm). Dropping the frames here silently severs the + // archive from the provider request — the model sees a summary that + // references attached frames that never arrive. + const images: ImageContent[] = [ + { type: "image", data: "ZmFrZQ==", mimeType: "image/png" }, + { type: "image", data: "ZmFrZTI=", mimeType: "image/png" }, + ]; + const messages: AgentMessage[] = [ + { + role: "compactionSummary", + summary: "the film archive", + tokensBefore: 1000, + images, + timestamp: Date.now(), + }, + ]; + + const converted = convertToLlm(messages); + + expect(converted).toHaveLength(1); + expect(converted[0]?.role).toBe("user"); + const content = converted[0]?.content as Array; + expect(content).toHaveLength(3); + expect(content[0].type).toBe("text"); + expect((content[0] as TextContent).text).toContain("the film archive"); + expect(content[1]).toEqual(images[0]); + expect(content[2]).toEqual(images[1]); + }); + + it("emits text-only content when no frames are archived", () => { + const messages: AgentMessage[] = [ + { role: "compactionSummary", summary: "plain summary", tokensBefore: 1000, timestamp: Date.now() }, + ]; + const converted = convertToLlm(messages); + expect((converted[0]?.content as unknown[]).length).toBe(1); + }); +}); + describe("convertToLlm custom message mapping", () => { it("maps custom messages to developer role with explicit agent attribution", () => { const messages: AgentMessage[] = [ From 0f4f301a0209379222d9f382304a585b05945e20 Mon Sep 17 00:00:00 2001 From: can1357 Date: Wed, 10 Jun 2026 17:53:07 +0200 Subject: [PATCH 18/27] feat(coding-agent): made waiting job polls displaceable in the TUI A job poll whose watched jobs are all still running and nothing was cancelled is pure 'still waiting' noise once a newer block exists. ToolExecutionComponent now detects isWaitingPollDetails on the result, leaves the block displaceable, and keeps its rows out of scrollback. EventController tracks the most recent displaceable poll and resolves it (sealing the block) on every event that proves another tool now owns the live region: a new tool call/result, a streaming arg start, an end-of-turn signal, or another job result that itself isn't displaceable. UiHelpers' transcript rebuild applies the same rule. Docs and the new job-poll-displacement test pin the behavior. --- docs/tools/job.md | 5 +- .../src/modes/components/tool-execution.ts | 32 ++- .../src/modes/controllers/event-controller.ts | 45 ++++ .../src/modes/utils/ui-helpers.ts | 10 + packages/coding-agent/src/tools/job.ts | 13 ++ .../test/job-poll-displacement.test.ts | 220 ++++++++++++++++++ 6 files changed, 321 insertions(+), 4 deletions(-) create mode 100644 packages/coding-agent/test/job-poll-displacement.test.ts diff --git a/docs/tools/job.md b/docs/tools/job.md index edeb260c8..5a50c7be0 100644 --- a/docs/tools/job.md +++ b/docs/tools/job.md @@ -7,7 +7,6 @@ - Model-facing prompt: `packages/coding-agent/src/prompts/tools/job.md` - Key collaborators: - `packages/coding-agent/src/async/job-manager.ts` — job registry, cancellation, delivery suppression. - - `packages/coding-agent/src/async/support.ts` — feature gating for background jobs. - `packages/coding-agent/src/tools/bash.ts` — explicit async bash and auto-backgrounded bash jobs. - `packages/coding-agent/src/task/index.ts` — async task-job scheduling. - `packages/coding-agent/src/sdk.ts` — automatic follow-up delivery for unsuppressed completions. @@ -45,7 +44,7 @@ Read-only snapshot path: - Calling `job` with `list: true` returns a markdown summary of every job spawned by the calling agent (running + completed within retention) without waiting. ## Flow -1. `JobTool.createIf(...)` in `packages/coding-agent/src/tools/job.ts` only exposes the tool when `isBackgroundJobSupportEnabled(...)` returns true for either `async.enabled` or `bash.autoBackground.enabled`. +1. `JobTool` is registered unconditionally in `packages/coding-agent/src/tools/index.ts`; there is no `async.enabled` gate (the `task` tool always schedules background jobs). 2. `execute(...)` fetches `session.asyncJobManager`. If absent, it returns `Async execution is disabled; no background jobs are available.` 3. `cancel` ids are processed first: - `manager.getJob(id)` missing → `not_found`. @@ -83,7 +82,7 @@ Spawn paths that produce jobs: - `async: true` always registers a `type: "bash"` job with `AsyncJobManager.register(...)` and returns a start message. - auto-background mode (`bash.autoBackground.enabled`) starts the same managed job path for non-PTY commands, waits up to `min(bash.autoBackground.thresholdMs, timeoutMs - 1000)`, and if the command is still running returns a background-job start result instead of inline command output. - `packages/coding-agent/src/task/index.ts` - - when `async.enabled` is on, the chosen agent is not blocking, and `tasks.length > 0`, each task item is registered as a `type: "task"` job. + - every `task` call (spawn or resume) registers one `type: "task"` job, unless the session has no job manager or the agent definition declares `blocking: true` (sync fallback). Lifecycle and exact state names: - Conceptual scheduling path: `pending` (only task-progress bookkeeping before work starts) → `running` → `completed` / `failed`; cancellation changes a running async job to `cancelled`. diff --git a/packages/coding-agent/src/modes/components/tool-execution.ts b/packages/coding-agent/src/modes/components/tool-execution.ts index e63785818..284f86e6d 100644 --- a/packages/coding-agent/src/modes/components/tool-execution.ts +++ b/packages/coding-agent/src/modes/components/tool-execution.ts @@ -19,6 +19,7 @@ import type { Theme } from "../../modes/theme/theme"; import { theme } from "../../modes/theme/theme"; import { BASH_DEFAULT_PREVIEW_LINES } from "../../tools/bash"; import { EVAL_DEFAULT_PREVIEW_LINES } from "../../tools/eval"; +import { isWaitingPollDetails } from "../../tools/job"; import { formatArgsInline, JSON_TREE_MAX_DEPTH_COLLAPSED, @@ -194,6 +195,11 @@ export class ToolExecutionComponent extends Container { // sealed the block stays in the transcript's repaintable live region so a // late result still repaints instead of stranding the streaming preview. #sealed = false; + // A `job` poll result whose watched jobs are all still running. Such a + // block never finalizes (stays in the transcript live region) so a + // follow-up `job` call can displace it instead of stacking another + // "waiting on N jobs" frame. Cleared by `seal()`. + #displaceable = false; #renderState: { spinnerFrame?: number; expanded: boolean; @@ -359,6 +365,11 @@ export class ToolExecutionComponent extends Container { ): void { this.#result = result; this.#isPartial = isPartial; + // A `job` poll that found every watched job still running is transient + // "still waiting" chrome; keep the block displaceable so the next `job` + // call replaces it instead of stacking another waiting frame (see the + // event controller's displaceable-poll bookkeeping). + this.#displaceable = this.#toolName === "job" && result.isError !== true && isWaitingPollDetails(result.details); // When tool is complete, ensure args are marked complete so spinner stops if (!isPartial) { this.#argsComplete = true; @@ -425,7 +436,11 @@ export class ToolExecutionComponent extends Container { (this.#result?.details as { async?: { state?: string } } | undefined)?.async?.state === "running"; const isBackgroundAsyncTask = this.#toolName === "task" && isBackgroundAsyncRunning; const isPartialTask = this.#isPartial && this.#toolName === "task" && !isBackgroundAsyncTask; - const needsSpinner = isStreamingArgs || isPartialTask; + // A displaceable waiting poll keeps its spinner ticking: it reads as one + // persistent live poll, and the changing leading glyph keeps the + // transcript's stable-prefix ratchet from committing rows of a block + // that a follow-up `job` call may remove. + const needsSpinner = isStreamingArgs || isPartialTask || this.isDisplaceableBlock(); if (needsSpinner && !this.#spinnerInterval) { const now = performance.now(); const frameCount = theme.spinnerFrames.length; @@ -513,6 +528,9 @@ export class ToolExecutionComponent extends Container { isTranscriptBlockFinalized(): boolean { if (this.#sealed) return true; if (this.#result === undefined) return false; + // A displaceable waiting poll stays live: its rows are kept out of + // native scrollback so a follow-up `job` call can remove the block. + if (this.#displaceable) return false; if (!this.#isPartial) return true; // Partial result: a background async tool is accepted to freeze (the agent // continues while it runs and would otherwise pin an unbounded live region); @@ -528,11 +546,23 @@ export class ToolExecutionComponent extends Container { seal(): void { if (this.#sealed) return; this.#sealed = true; + this.#displaceable = false; this.stopAnimation(); this.#updateDisplay(); this.#ui.requestRender(); } + /** + * Whether this block is a waiting `job` poll (every watched job still + * running) that has not been sealed. Such a block never finalized, so none + * of its rows entered native scrollback (the ticking spinner keeps the + * stable-prefix ratchet at zero) and the whole block can be removed when a + * follow-up `job` call supersedes it. + */ + isDisplaceableBlock(): boolean { + return this.#displaceable && !this.#sealed; + } + /** * Stop spinner animation and cleanup resources. */ diff --git a/packages/coding-agent/src/modes/controllers/event-controller.ts b/packages/coding-agent/src/modes/controllers/event-controller.ts index d9a48f3fb..2c12e1d45 100644 --- a/packages/coding-agent/src/modes/controllers/event-controller.ts +++ b/packages/coding-agent/src/modes/controllers/event-controller.ts @@ -77,6 +77,11 @@ export class EventController { // Insertion-ordered IRC cards not yet retired; values are the transcript // components each card contributed (see #retireIrcCard for the guard). #liveIrcCards = new Map(); + // Most recent `job` tool block whose result still had every watched job + // running. Kept un-finalized (live) so the next `job` call displaces it — + // one persistent poll instead of a stack of "waiting on N jobs" frames — + // and sealed in place the moment anything else lands below it. + #displaceablePollComponent: ToolExecutionComponent | undefined = undefined; #streamingReveal: StreamingRevealController; #handlers: AgentSessionEventHandlers; @@ -282,6 +287,7 @@ export class EventController { const signature = `${textContent}\u0000${imageCount}`; this.#resetReadGroup(); + this.#resolveDisplaceablePoll(); const wasOptimistic = this.ctx.optimisticUserMessageSignature === signature; const wasLocallySubmitted = this.ctx.locallySubmittedUserSignatures.delete(signature) || wasOptimistic; if (!wasOptimistic) { @@ -389,6 +395,28 @@ export class EventController { } } + /** + * Resolve the pending displaceable poll block before the next block lands. + * A follow-up `job` call displaces it — the stale "waiting on N jobs" frame + * is removed so repeated polls read as one persistent poll — while anything + * else seals it in place as final history. Removal is safe only because a + * displaceable block never finalizes: commits stop at the first live block, + * so none of its rows have entered native scrollback (see + * ToolExecutionComponent.isDisplaceableBlock). + */ + #resolveDisplaceablePoll(nextToolName?: string): void { + const previous = this.#displaceablePollComponent; + if (!previous) return; + this.#displaceablePollComponent = undefined; + if (nextToolName === "job" && previous.isDisplaceableBlock()) { + this.ctx.chatContainer.removeChild(previous); + } + // Sealing stops the waiting-poll spinner and freezes the block (for a + // just-removed component it only clears the animation timer). + previous.seal(); + this.ctx.ui.requestRender(); + } + async #handleNotice(event: Extract): Promise { const message = event.source ? `${event.source}: ${event.message}` : event.message; if (event.level === "error") { @@ -444,6 +472,7 @@ export class EventController { continue; } if (!readArgsTargetInternalUrl(content.arguments)) { + if (!this.ctx.pendingTools.has(content.id)) this.#resolveDisplaceablePoll(content.name); this.#trackReadToolCall(content.id, content.arguments); const component = this.ctx.pendingTools.get(content.id); if (component) { @@ -465,6 +494,7 @@ export class EventController { ? { ...content.arguments, __partialJson: content.partialJson } : content.arguments; if (!this.ctx.pendingTools.has(content.id)) { + this.#resolveDisplaceablePoll(content.name); this.#resetReadGroup(); const tool = this.ctx.session.getToolByName(content.name); const component = new ToolExecutionComponent( @@ -561,6 +591,9 @@ export class EventController { component.seal(); } } + // These calls will never produce a result either, so the tracked + // waiting poll cannot be displaced anymore — freeze it in place. + this.#resolveDisplaceablePoll(); } this.#lastAssistantComponent = this.ctx.streamingComponent; this.#lastAssistantComponent.setUsageInfo(event.message.usage); @@ -589,6 +622,7 @@ export class EventController { async #handleToolExecutionStart(event: Extract): Promise { this.#updateWorkingMessageFromIntent(event.intent); if (!this.ctx.pendingTools.has(event.toolCallId)) { + this.#resolveDisplaceablePoll(event.toolName); if (event.toolName === "read" && readArgsHaveTarget(event.args) && !readArgsTargetInternalUrl(event.args)) { this.#trackReadToolCall(event.toolCallId, event.args); const component = this.ctx.pendingTools.get(event.toolCallId); @@ -697,6 +731,14 @@ export class EventController { this.ctx.pendingTools.delete(event.toolCallId); this.#backgroundToolCallIds.delete(event.toolCallId); } + if ( + event.toolName === "job" && + component instanceof ToolExecutionComponent && + component.isDisplaceableBlock() + ) { + // Remember the waiting poll so the next `job` call can displace it. + this.#displaceablePollComponent = component; + } this.ctx.ui.requestRender(); } } @@ -759,6 +801,9 @@ export class EventController { this.#readToolCallArgs.clear(); this.#readToolCallAssistantComponents.clear(); this.#resetReadGroup(); + // The turn is over: nothing else lands this turn, so the waiting poll is + // final history — seal it instead of letting its spinner tick while idle. + this.#resolveDisplaceablePoll(); this.#lastAssistantComponent = undefined; this.ctx.ui.requestRender(); this.#scheduleIdleCompaction(); diff --git a/packages/coding-agent/src/modes/utils/ui-helpers.ts b/packages/coding-agent/src/modes/utils/ui-helpers.ts index a01bc69da..e0e805b09 100644 --- a/packages/coding-agent/src/modes/utils/ui-helpers.ts +++ b/packages/coding-agent/src/modes/utils/ui-helpers.ts @@ -379,6 +379,7 @@ export class UiHelpers { if (content.type !== "toolCall") { continue; } + resolveWaitingPoll(content.name); if ( content.name === "read" && @@ -493,8 +494,17 @@ export class UiHelpers { if (component) { component.updateResult(message, false, message.toolCallId); this.ctx.pendingTools.delete(message.toolCallId); + if ( + message.toolName === "job" && + component instanceof ToolExecutionComponent && + component.isDisplaceableBlock() + ) { + waitingPoll = component; + } } } else { + // A user prompt closes the displacement window, same as the live path. + if (message.role === "user") resolveWaitingPoll(); // All other messages use standard rendering this.ctx.addMessageToChat(message, options); } diff --git a/packages/coding-agent/src/tools/job.ts b/packages/coding-agent/src/tools/job.ts index b69dc7290..2991d7525 100644 --- a/packages/coding-agent/src/tools/job.ts +++ b/packages/coding-agent/src/tools/job.ts @@ -65,6 +65,19 @@ export interface JobToolDetails { cancelled?: { id: string; status: CancelStatus }[]; } +/** + * A poll snapshot where every watched job is still running and nothing was + * cancelled — pure "still waiting" noise once a newer poll exists. The TUI + * keeps such a block un-finalized (displaceable) so a follow-up `job` call + * replaces it instead of stacking another waiting frame in the transcript. + */ +export function isWaitingPollDetails(details: unknown): boolean { + const d = details as JobToolDetails | undefined; + if (!d || !Array.isArray(d.jobs) || d.jobs.length === 0) return false; + if (d.cancelled?.length) return false; + return d.jobs.every(job => job?.status === "running"); +} + export class JobTool implements AgentTool { readonly name = "job"; readonly approval = "read" as const; diff --git a/packages/coding-agent/test/job-poll-displacement.test.ts b/packages/coding-agent/test/job-poll-displacement.test.ts new file mode 100644 index 000000000..4ce22f38d --- /dev/null +++ b/packages/coding-agent/test/job-poll-displacement.test.ts @@ -0,0 +1,220 @@ +/** + * Repeated `job` polls must not stack "waiting on N jobs" frames in the + * transcript: a poll whose watched jobs are all still running stays live + * (displaceable) and the next `job` call replaces it — one persistent poll. + * + * Contracts under test: + * - ToolExecutionComponent: a waiting-poll result keeps the block + * un-finalized and displaceable; a settled/cancelled/error result + * finalizes normally; seal() always freezes. + * - EventController: a follow-up `job` call removes the tracked waiting + * poll from the transcript; any other tool seals it in place. + */ +import { afterEach, beforeEach, describe, expect, it, vi } from "bun:test"; +import { resetSettingsForTest, Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; +import { ToolExecutionComponent } from "@oh-my-pi/pi-coding-agent/modes/components/tool-execution"; +import { EventController } from "@oh-my-pi/pi-coding-agent/modes/controllers/event-controller"; +import { initTheme } from "@oh-my-pi/pi-coding-agent/modes/theme/theme"; +import type { InteractiveModeContext } from "@oh-my-pi/pi-coding-agent/modes/types"; +import type { Component, TUI } from "@oh-my-pi/pi-tui"; + +const uiStub = { requestRender() {} } as unknown as TUI; + +type JobStatus = "running" | "completed" | "failed" | "cancelled"; + +function pollResult(statuses: JobStatus[], extra: { cancelled?: boolean; isError?: boolean } = {}) { + return { + content: [{ type: "text" as const, text: "" }], + isError: extra.isError, + details: { + jobs: statuses.map((status, i) => ({ + id: `j${i}`, + type: "task" as const, + status, + label: `job ${i}`, + durationMs: 1_000, + })), + ...(extra.cancelled ? { cancelled: [{ id: "j0", status: "cancelled" as const }] } : {}), + }, + }; +} + +function trackComponent(components: ToolExecutionComponent[], component: ToolExecutionComponent) { + components.push(component); + return component; +} + +describe("job waiting-poll block lifecycle", () => { + const created: ToolExecutionComponent[] = []; + + beforeEach(async () => { + resetSettingsForTest(); + await Settings.init({ inMemory: true }); + await initTheme(); + }); + + afterEach(() => { + // Seal everything so displaceable blocks' spinner intervals never leak + // into later test files. + for (const component of created.splice(0)) component.seal(); + vi.restoreAllMocks(); + resetSettingsForTest(); + }); + + function makeJobComponent() { + return trackComponent(created, new ToolExecutionComponent("job", { poll: ["j0", "j1"] }, {}, undefined, uiStub)); + } + + it("keeps an all-running poll live and displaceable until sealed", () => { + const component = makeJobComponent(); + component.updateResult(pollResult(["running", "running"]), false); + + expect(component.isDisplaceableBlock()).toBe(true); + expect(component.isTranscriptBlockFinalized()).toBe(false); + + component.seal(); + expect(component.isDisplaceableBlock()).toBe(false); + expect(component.isTranscriptBlockFinalized()).toBe(true); + }); + + it("finalizes a poll that observed a settled job", () => { + const component = makeJobComponent(); + component.updateResult(pollResult(["completed", "running"]), false); + + expect(component.isDisplaceableBlock()).toBe(false); + expect(component.isTranscriptBlockFinalized()).toBe(true); + }); + + it("finalizes a poll that carried cancel outcomes or an error", () => { + const cancelled = makeJobComponent(); + cancelled.updateResult(pollResult(["running"], { cancelled: true }), false); + expect(cancelled.isDisplaceableBlock()).toBe(false); + + const errored = makeJobComponent(); + errored.updateResult(pollResult(["running"], { isError: true }), false); + expect(errored.isDisplaceableBlock()).toBe(false); + expect(errored.isTranscriptBlockFinalized()).toBe(true); + }); + + it("never marks non-job tools displaceable", () => { + const component = trackComponent(created, new ToolExecutionComponent("bash", { command: "ls" }, {}, undefined, uiStub)); + component.updateResult(pollResult(["running"]), false); + expect(component.isDisplaceableBlock()).toBe(false); + }); +}); + +describe("EventController displaces consecutive waiting polls", () => { + const created: ToolExecutionComponent[] = []; + + beforeEach(async () => { + resetSettingsForTest(); + await Settings.init({ inMemory: true }); + await initTheme(); + }); + + afterEach(() => { + for (const component of created.splice(0)) component.seal(); + vi.restoreAllMocks(); + resetSettingsForTest(); + }); + + function createFixture() { + const children: Component[] = []; + const ctx = { + isInitialized: true, + init: vi.fn(async () => {}), + ui: { requestRender: vi.fn() }, + statusLine: { invalidate: vi.fn() }, + updateEditorTopBorder: vi.fn(), + toolOutputExpanded: false, + pendingTools: new Map(), + chatContainer: { + children, + addChild: (component: Component) => { + children.push(component); + }, + removeChild: (component: Component) => { + const index = children.indexOf(component); + if (index !== -1) children.splice(index, 1); + }, + }, + session: { getToolByName: () => undefined }, + sessionManager: { getCwd: () => process.cwd() }, + } as unknown as InteractiveModeContext; + return { controller: new EventController(ctx), children }; + } + + async function runPoll(controller: EventController, children: Component[], toolCallId: string) { + await controller.handleEvent({ type: "tool_execution_start", toolCallId, toolName: "job", args: { poll: ["j0"] } }); + const component = children[children.length - 1] as ToolExecutionComponent; + trackComponent(created, component); + await controller.handleEvent({ + type: "tool_execution_end", + toolCallId, + toolName: "job", + result: pollResult(["running", "running"]), + isError: false, + }); + return component; + } + + it("removes the previous waiting poll when the next job call starts", async () => { + const { controller, children } = createFixture(); + + const first = await runPoll(controller, children, "t1"); + expect(children).toContain(first); + + const second = await runPoll(controller, children, "t2"); + + // The stale "waiting" frame is gone; only the fresh poll remains. + expect(children).not.toContain(first); + expect(children).toContain(second); + // The displaced block is sealed so its spinner interval is stopped. + expect(first.isTranscriptBlockFinalized()).toBe(true); + }); + + it("seals the waiting poll in place when a different tool runs next", async () => { + const { controller, children } = createFixture(); + + const poll = await runPoll(controller, children, "t1"); + expect(poll.isTranscriptBlockFinalized()).toBe(false); + + await controller.handleEvent({ + type: "tool_execution_start", + toolCallId: "t2", + toolName: "bash", + args: { command: "ls" }, + }); + trackComponent(created, children[children.length - 1] as ToolExecutionComponent); + + // The poll frame stays — it is final history now, not displaceable. + expect(children).toContain(poll); + expect(poll.isTranscriptBlockFinalized()).toBe(true); + expect(poll.isDisplaceableBlock()).toBe(false); + }); + + it("does not displace a poll that observed completions", async () => { + const { controller, children } = createFixture(); + + await controller.handleEvent({ + type: "tool_execution_start", + toolCallId: "t1", + toolName: "job", + args: { poll: ["j0"] }, + }); + const settled = trackComponent(created, children[children.length - 1] as ToolExecutionComponent); + await controller.handleEvent({ + type: "tool_execution_end", + toolCallId: "t1", + toolName: "job", + result: pollResult(["completed", "running"]), + isError: false, + }); + + const next = await runPoll(controller, children, "t2"); + + // A poll that carried real results is kept as history. + expect(children).toContain(settled); + expect(children).toContain(next); + }); +}); From 3aa1cedd737fc30938bd12259a65a5c3f49834b5 Mon Sep 17 00:00:00 2001 From: can1357 Date: Wed, 10 Jun 2026 17:53:26 +0200 Subject: [PATCH 19/27] feat(coding-agent): enforced an inline byte cap at the bash and browser tool boundary Adds enforceInlineByteCap() in streaming-output and applies it to bash and browser tool results: oversized outputs are elided head/tail with an artifact:// footer pointing at the full capture, closing paths that previously let 100KB+ inline results past the minimizer. Defense at the tool-result boundary (no-op for already-bounded output). --- .../src/session/streaming-output.ts | 60 ++++++++ packages/coding-agent/src/tools/bash.ts | 20 ++- packages/coding-agent/src/tools/browser.ts | 29 +++- .../test/tools/output-caps.test.ts | 142 ++++++++++++++++++ 4 files changed, 245 insertions(+), 6 deletions(-) create mode 100644 packages/coding-agent/test/tools/output-caps.test.ts diff --git a/packages/coding-agent/src/session/streaming-output.ts b/packages/coding-agent/src/session/streaming-output.ts index 2eee2dd25..b4aa18ab4 100644 --- a/packages/coding-agent/src/session/streaming-output.ts +++ b/packages/coding-agent/src/session/streaming-output.ts @@ -570,6 +570,66 @@ export function truncateMiddle(content: string, options: TruncationOptions = {}) }; } +// ============================================================================= +// Inline byte cap — final defense at the tool-result boundary +// ============================================================================= + +/** Options for {@link enforceInlineByteCap}. */ +export interface InlineByteCapOptions { + /** Inline byte budget. Defaults to {@link DEFAULT_MAX_BYTES}. */ + maxBytes?: number; + /** What the text is, for the elision marker (e.g. "bash output"). */ + label: string; + /** + * Persist the full text as a session artifact. When an artifact id is + * returned, a `[raw output: artifact://]` footer is appended so the + * elided bytes stay recoverable. + */ + saveArtifact?: (full: string) => string | undefined | Promise; +} + +/** Drop the partial last line of a head window (keep it if there is no newline at all). */ +function trimHeadToLineBoundary(text: string): string { + const idx = text.lastIndexOf(NL); + return idx > 0 ? text.substring(0, idx) : text; +} + +/** Drop the partial first line of a tail window (keep it if there is no newline at all). */ +function trimTailToLineBoundary(text: string): string { + const idx = text.indexOf(NL); + if (idx < 0 || idx === text.length - 1) return text; + return text.substring(idx + 1); +} + +/** + * Final-defense inline size guard for tool results. + * + * No-op when `text` fits within `maxBytes` (the common path). Otherwise keeps + * ~60% of the budget from the head and ~25% from the tail — cut on line + * boundaries, never splitting a multi-byte UTF-8 sequence — with an elision + * marker between. The remaining ~15% is slack for the marker and the optional + * `[raw output: artifact://]` footer, so the result stays under `maxBytes`. + */ +export async function enforceInlineByteCap(text: string, options: InlineByteCapOptions): Promise { + const maxBytes = options.maxBytes ?? DEFAULT_MAX_BYTES; + if (maxBytes <= 0) return text; + const totalBytes = Buffer.byteLength(text, "utf-8"); + if (totalBytes <= maxBytes) return text; + + const head = trimHeadToLineBoundary(truncateHeadBytes(text, Math.floor(maxBytes * 0.6)).text); + const tail = trimTailToLineBoundary(truncateTailBytes(text, Math.floor(maxBytes * 0.25)).text); + const elidedBytes = Math.max(0, totalBytes - Buffer.byteLength(head, "utf-8") - Buffer.byteLength(tail, "utf-8")); + const marker = `[… elided ${elidedBytes} bytes of ${options.label} …]`; + let composed = `${head}\n${marker}\n${tail}`; + + const artifactId = await options.saveArtifact?.(text); + if (artifactId) { + const sep = composed.endsWith(NL) ? "" : NL; + composed += `${sep}[raw output: artifact://${artifactId}]`; + } + return composed; +} + // ============================================================================= // TailBuffer — ring-style tail buffer with lazy joining // ============================================================================= diff --git a/packages/coding-agent/src/tools/bash.ts b/packages/coding-agent/src/tools/bash.ts index 18246957a..ce7dda67b 100644 --- a/packages/coding-agent/src/tools/bash.ts +++ b/packages/coding-agent/src/tools/bash.ts @@ -17,7 +17,7 @@ import { truncateToVisualLines } from "../modes/components/visual-truncate"; import { highlightCode, type Theme } from "../modes/theme/theme"; import bashDescription from "../prompts/tools/bash.md" with { type: "text" }; import type { ClientBridgeTerminalExitStatus, ClientBridgeTerminalOutput } from "../session/client-bridge"; -import { DEFAULT_MAX_BYTES, streamTailUpdates, TailBuffer } from "../session/streaming-output"; +import { DEFAULT_MAX_BYTES, enforceInlineByteCap, streamTailUpdates, TailBuffer } from "../session/streaming-output"; import { renderStatusLine } from "../tui"; import { CachedOutputBlock, markFramedBlockComponent } from "../tui/output-block"; import { getSixelLineMask } from "../utils/sixel"; @@ -429,7 +429,7 @@ export class BashTool implements AgentTool { } } - #buildCompletedResult( + async #buildCompletedResult( result: BashResult | BashInteractiveResult, timeoutSec: number, options: { @@ -438,7 +438,7 @@ export class BashTool implements AgentTool { terminalId?: string; wallTimeMs?: number; } = {}, - ): AgentToolResult { + ): Promise> { const exitCode = result.exitCode; const failedExit = exitCode !== undefined && exitCode !== 0; @@ -472,7 +472,17 @@ export class BashTool implements AgentTool { if (failedExit) { details.exitCode = exitCode; } - const resultBuilder = toolResult(details).text(outputText).truncationFromSummary(result, { direction: "tail" }); + // Final defense at the tool-result boundary: no bash path (client bridge, + // head-retention spill, minimizer miss) may emit more than + // ~DEFAULT_MAX_BYTES inline. No-op for already-bounded output. + const cappedOutputText = await enforceInlineByteCap(outputText, { + label: "bash output", + saveArtifact: full => saveBashOriginalArtifact(this.session, full), + }); + + const resultBuilder = toolResult(details) + .text(cappedOutputText) + .truncationFromSummary(result, { direction: "tail" }); if (failedExit) resultBuilder.error(); return resultBuilder.done(); } @@ -560,7 +570,7 @@ export class BashTool implements AgentTool { onMinimizedSave: originalText => saveBashOriginalArtifact(this.session, originalText), }); const wallTimeMs = performance.now() - wallTimeStart; - const finalResult = this.#buildCompletedResult(result, options.timeoutSec, { + const finalResult = await this.#buildCompletedResult(result, options.timeoutSec, { requestedTimeoutSec: options.requestedTimeoutSec, notices: options.notices ?? [], wallTimeMs, diff --git a/packages/coding-agent/src/tools/browser.ts b/packages/coding-agent/src/tools/browser.ts index 1aa73c29c..093bc7af3 100644 --- a/packages/coding-agent/src/tools/browser.ts +++ b/packages/coding-agent/src/tools/browser.ts @@ -3,6 +3,7 @@ import { prompt, untilAborted } from "@oh-my-pi/pi-utils"; import * as z from "zod/v4"; import browserDescription from "../prompts/tools/browser.md" with { type: "text" }; import type { ToolSession } from "../sdk"; +import { enforceInlineByteCap } from "../session/streaming-output"; import { truncateForPrompt } from "./approval"; import { acquireBrowser, type BrowserHandle, type BrowserKind, type BrowserKindTag } from "./browser/registry"; import type { Observation, ScreenshotResult } from "./browser/tab-protocol"; @@ -271,11 +272,37 @@ export class BrowserTool implements AgentTool c.type === "text") .map(c => c.text) .join("\n"); - details.result = textOnly; + // Final defense at the tool-result boundary: a single run can display + // tens of KB (large JSON returns, dumped observations). Cap the combined + // text inline; the full text stays recoverable via the artifact footer + // when allocation succeeds. + const cappedText = await enforceInlineByteCap(textOnly, { + label: "browser output", + saveArtifact: full => saveBrowserOutputArtifact(this.session, full), + }); + details.result = cappedText; + if (cappedText !== textOnly) { + const nonText = content.filter(c => c.type !== "text"); + return toolResult(details) + .content([...nonText, { type: "text", text: cappedText }]) + .done(); + } return toolResult(details).content(content).done(); } } +/** Persist over-cap browser run output as a session artifact; mirrors the bash minimizer's save path. */ +async function saveBrowserOutputArtifact(session: ToolSession, fullText: string): Promise { + try { + const alloc = await session.allocateOutputArtifact?.("browser-original"); + if (!alloc?.path || !alloc.id) return undefined; + await Bun.write(alloc.path, fullText); + return alloc.id; + } catch { + return undefined; + } +} + function describeBrowser(handle: BrowserHandle): string { switch (handle.kind.kind) { case "headless": diff --git a/packages/coding-agent/test/tools/output-caps.test.ts b/packages/coding-agent/test/tools/output-caps.test.ts new file mode 100644 index 000000000..3c7997db6 --- /dev/null +++ b/packages/coding-agent/test/tools/output-caps.test.ts @@ -0,0 +1,142 @@ +/** + * Contract tests for `enforceInlineByteCap`, the final-defense inline size + * guard at the tool-result boundary (bash, browser). Over-cap text keeps a + * head (~60% of budget) and tail (~25%) cut on line boundaries with an + * elision marker between; sub-cap text passes through untouched so existing + * bounded-output paths (bash sink/minimizer) see zero behavior change. + */ +import { describe, expect, it } from "bun:test"; +import { DEFAULT_MAX_BYTES, enforceInlineByteCap } from "@oh-my-pi/pi-coding-agent/session/streaming-output"; + +const MARKER_PATTERN = /\[… elided \d+ bytes of test output …\]/; + +/** Build `count` complete lines of the form `line-00001 `. */ +function makeLines(count: number, pad = "x".repeat(40)): string { + const lines: string[] = []; + for (let i = 1; i <= count; i++) { + lines.push(`line-${String(i).padStart(5, "0")} ${pad}`); + } + return lines.join("\n"); +} + +describe("enforceInlineByteCap", () => { + it("returns sub-cap text unchanged (identity, not just equality)", async () => { + const text = makeLines(10); + const result = await enforceInlineByteCap(text, { maxBytes: 4096, label: "test output" }); + expect(result).toBe(text); + }); + + it("returns text exactly at the cap unchanged", async () => { + const text = "a".repeat(1000); + expect(Buffer.byteLength(text, "utf-8")).toBe(1000); + const result = await enforceInlineByteCap(text, { maxBytes: 1000, label: "test output" }); + expect(result).toBe(text); + }); + + it("uses DEFAULT_MAX_BYTES when maxBytes is omitted", async () => { + const under = "a".repeat(DEFAULT_MAX_BYTES - 1); + expect(await enforceInlineByteCap(under, { label: "test output" })).toBe(under); + + const over = makeLines(2000); // ~94KB, well over the 50KB default + expect(Buffer.byteLength(over, "utf-8")).toBeGreaterThan(DEFAULT_MAX_BYTES); + const result = await enforceInlineByteCap(over, { label: "test output" }); + expect(result).not.toBe(over); + expect(result).toMatch(MARKER_PATTERN); + expect(Buffer.byteLength(result, "utf-8")).toBeLessThanOrEqual(DEFAULT_MAX_BYTES); + }); + + it("elides over-cap text to head + tail on line boundaries within the budget", async () => { + const maxBytes = 4096; + const text = makeLines(500); // ~23KB + const result = await enforceInlineByteCap(text, { maxBytes, label: "test output" }); + + expect(result).toMatch(MARKER_PATTERN); + // Stays within the cap plus a sliver of slack for the marker line. + expect(Buffer.byteLength(result, "utf-8")).toBeLessThanOrEqual(maxBytes + 256); + + const lines = result.split("\n"); + const markerIdx = lines.findIndex(line => MARKER_PATTERN.test(line)); + expect(markerIdx).toBeGreaterThan(0); + expect(markerIdx).toBeLessThan(lines.length - 1); + + // Head starts at the very beginning; tail ends at the very end. + expect(lines[0]).toBe(`line-00001 ${"x".repeat(40)}`); + expect(lines[lines.length - 1]).toBe(`line-00500 ${"x".repeat(40)}`); + + // Line-boundary cuts: the lines flanking the marker are complete input lines. + const completeLine = /^line-\d{5} x{40}$/; + expect(lines[markerIdx - 1]).toMatch(completeLine); + expect(lines[markerIdx + 1]).toMatch(completeLine); + + // Head should be roughly 60% and tail roughly 25% of the budget. + const headBytes = Buffer.byteLength(lines.slice(0, markerIdx).join("\n"), "utf-8"); + const tailBytes = Buffer.byteLength(lines.slice(markerIdx + 1).join("\n"), "utf-8"); + expect(headBytes).toBeLessThanOrEqual(Math.floor(maxBytes * 0.6)); + expect(tailBytes).toBeLessThanOrEqual(Math.floor(maxBytes * 0.25)); + expect(headBytes).toBeGreaterThan(tailBytes); + }); + + it("does not corrupt multi-byte UTF-8 near the cut boundaries", async () => { + // Each line is multi-byte heavy: é (2B), € (3B), 😀 (4B). + const text = makeLines(800, "é€😀".repeat(12)); + const maxBytes = 4096; + const result = await enforceInlineByteCap(text, { maxBytes, label: "test output" }); + + expect(result).toMatch(MARKER_PATTERN); + // Valid UTF-8 round-trip: encode/decode is lossless and introduces no + // replacement characters (the input contains none). + const roundTripped = Buffer.from(result, "utf-8").toString("utf-8"); + expect(roundTripped).toBe(result); + expect(result.includes("\uFFFD")).toBe(false); + + // Every kept content line is a complete, uncorrupted input line. + const completeLine = /^line-\d{5} (?:é€😀){12}$/u; + for (const line of result.split("\n")) { + if (MARKER_PATTERN.test(line)) continue; + expect(line).toMatch(completeLine); + } + }); + + it("appends the artifact footer when saveArtifact yields an id", async () => { + const text = makeLines(500); + let saved: string | undefined; + const result = await enforceInlineByteCap(text, { + maxBytes: 4096, + label: "test output", + saveArtifact: full => { + saved = full; + return Promise.resolve("17"); + }, + }); + // saveArtifact receives the full original text, not the elided version. + expect(saved).toBe(text); + expect(result.endsWith("[raw output: artifact://17]")).toBe(true); + expect(result).toMatch(MARKER_PATTERN); + }); + + it("omits the footer when saveArtifact returns undefined", async () => { + const text = makeLines(500); + const result = await enforceInlineByteCap(text, { + maxBytes: 4096, + label: "test output", + saveArtifact: () => undefined, + }); + expect(result).not.toContain("[raw output: artifact://"); + expect(result).toMatch(MARKER_PATTERN); + }); + + it("does not invoke saveArtifact for sub-cap text", async () => { + let called = false; + const text = "short output"; + const result = await enforceInlineByteCap(text, { + maxBytes: 4096, + label: "test output", + saveArtifact: () => { + called = true; + return "99"; + }, + }); + expect(result).toBe(text); + expect(called).toBe(false); + }); +}); From 9613076e9637e88a23d3c7417ac8c79ecd8be707 Mon Sep 17 00:00:00 2001 From: can1357 Date: Wed, 10 Jun 2026 17:53:46 +0200 Subject: [PATCH 20/27] feat(coding-agent): added repeat-read notice to the read tool Tracks successful file reads per resolved base path (selector stripped) for the session; once a path has been whole-file-read three times, every subsequent read for that path appends a one-line nudge suggesting narrower line-range re-reads or the context echoed in edit results. Non-file sources (URLs, internal resources, directories, archives, SQLite, images) are never counted. --- packages/coding-agent/src/tools/read.ts | 58 +++++++- .../test/tools/read-repeat-notice.test.ts | 137 ++++++++++++++++++ 2 files changed, 194 insertions(+), 1 deletion(-) create mode 100644 packages/coding-agent/test/tools/read-repeat-notice.test.ts diff --git a/packages/coding-agent/src/tools/read.ts b/packages/coding-agent/src/tools/read.ts index 285dbd09b..5918da327 100644 --- a/packages/coding-agent/src/tools/read.ts +++ b/packages/coding-agent/src/tools/read.ts @@ -736,6 +736,17 @@ interface ResolvedSqliteReadPath { /** Per-execute memo of suffix-glob lookups; `null` records a confirmed miss. */ type SuffixMatchCache = Map; +/** + * Repeated whole-file reads of the same path pin stale copies in context. + * From this per-session read count onward, file reads carry a trailing nudge + * to prefer narrower re-reads. + */ +const REPEAT_READ_NOTICE_THRESHOLD = 3; + +function formatRepeatReadNotice(count: number): string { + return `[note: read #${count} of this file this session — after edits, prefer the context echoed in the edit result or a narrow range re-read]`; +} + /** * Read tool implementation. * @@ -754,6 +765,8 @@ export class ReadTool implements AgentTool { readonly #autoResizeImages: boolean; readonly #defaultLimit: number; readonly #inspectImageEnabled: boolean; + /** Successful file reads per resolved base path (selector stripped) this session. */ + readonly #readCounts = new Map(); constructor(private readonly session: ToolSession) { const displayMode = resolveFileDisplayMode(session); @@ -772,6 +785,19 @@ export class ReadTool implements AgentTool { }); } + /** + * Count a file read of `absolutePath` and return the repeat-read nudge once + * the per-session count reaches {@link REPEAT_READ_NOTICE_THRESHOLD}. + * Non-file sources (URLs, internal resources, directories, archives, + * SQLite, images) are never counted. + */ + #repeatReadNotice(absolutePath: string): string | undefined { + const count = (this.#readCounts.get(absolutePath) ?? 0) + 1; + this.#readCounts.set(absolutePath, count); + if (count < REPEAT_READ_NOTICE_THRESHOLD) return undefined; + return formatRepeatReadNotice(count); + } + async #tryReadDelimitedPaths( readPath: string, signal?: AbortSignal, @@ -948,6 +974,8 @@ export class ReadTool implements AgentTool { ignoreResultLimits?: boolean; raw?: boolean; immutable?: boolean; + /** Trailing repeat-read nudge; appended at the very end of the text. */ + repeatNotice?: string; }, ): AgentToolResult { const displayMode = resolveFileDisplayMode(this.session, { raw: options.raw, immutable: options.immutable }); @@ -1092,6 +1120,9 @@ export class ReadTool implements AgentTool { : formatLineEntries(buildLineEntries(endLine), startLineDisplay); } + if (options.repeatNotice) { + outputText += `\n${options.repeatNotice}`; + } resultBuilder.text(outputText); if (truncationInfo) { resultBuilder.truncation(truncationInfo.result, truncationInfo.options); @@ -1117,6 +1148,8 @@ export class ReadTool implements AgentTool { entityLabel: string; raw?: boolean; immutable?: boolean; + /** Trailing repeat-read nudge; appended at the very end of the text. */ + repeatNotice?: string; }, ): AgentToolResult { const displayMode = resolveFileDisplayMode(this.session, { raw: options.raw, immutable: options.immutable }); @@ -1177,8 +1210,11 @@ export class ReadTool implements AgentTool { const bound = range.endLine !== undefined ? `${range.startLine}-${range.endLine}` : `${range.startLine}`; notices.push(`[Range ${bound} is beyond end of ${options.entityLabel} (${totalLines} lines total); skipped]`); } - const finalText = + let finalText = notices.length > 0 ? (outputText ? `${outputText}\n${notices.join("\n")}` : notices.join("\n")) : outputText; + if (options.repeatNotice) { + finalText = finalText ? `${finalText}\n${options.repeatNotice}` : options.repeatNotice; + } resultBuilder.text(finalText); return resultBuilder.done(); } @@ -1196,6 +1232,7 @@ export class ReadTool implements AgentTool { parsed: ParsedSelector, displayMode: { hashLines: boolean; lineNumbers: boolean }, suffixResolution: { from: string; to: string } | undefined, + repeatNotice: string | undefined, signal: AbortSignal | undefined, ): Promise<{ outputText: string; @@ -1215,6 +1252,7 @@ export class ReadTool implements AgentTool { sourcePath: absolutePath, entityLabel: "file", raw: rawSelector, + repeatNotice, }); if (suffixResolution) { const notice = `[Path '${suffixResolution.from}' not found; resolved to '${suffixResolution.to}' via suffix match]`; @@ -1896,6 +1934,7 @@ export class ReadTool implements AgentTool { let details: ReadToolDetails = {}; let sourcePath: string | undefined; let columnTruncated = 0; + let repeatNotice: string | undefined; let truncationInfo: | { result: TruncationResult; options: { direction: "head"; startLine?: number; totalFileLines?: number } } | undefined; @@ -1960,11 +1999,13 @@ export class ReadTool implements AgentTool { } } else if (isNotebookPath(absolutePath) && !isRawSelector(parsed)) { const notebookText = await readEditableNotebookText(absolutePath, localReadPath); + repeatNotice = this.#repeatReadNotice(absolutePath); if (isMultiRange(parsed) && parsed.kind === "lines") { return this.#buildInMemoryMultiRangeResult(notebookText, parsed.ranges, { details: { resolvedPath: absolutePath }, sourcePath: absolutePath, entityLabel: "notebook", + repeatNotice, }); } const { offset, limit } = selToOffsetLimit(parsed); @@ -1972,11 +2013,13 @@ export class ReadTool implements AgentTool { details: { resolvedPath: absolutePath }, sourcePath: absolutePath, entityLabel: "notebook", + repeatNotice, }); } else if (shouldConvertWithMarkit) { // Convert document via markit. const result = await convertFileWithMarkit(absolutePath, signal); if (result.ok) { + repeatNotice = this.#repeatReadNotice(absolutePath); // Route the converted markdown through the in-memory text builder // so line-range selectors (`file.pdf:50-100`, `:5-16,40-80`) and // raw mode apply against the converted output. Without this, @@ -1987,6 +2030,7 @@ export class ReadTool implements AgentTool { details: { resolvedPath: absolutePath }, sourcePath: absolutePath, entityLabel: "document", + repeatNotice, }); } const { offset, limit } = selToOffsetLimit(parsed); @@ -1995,6 +2039,7 @@ export class ReadTool implements AgentTool { sourcePath: absolutePath, entityLabel: "document", raw: isRawSelector(parsed), + repeatNotice, }); } else if (result.error) { content = [{ type: "text", text: `[Cannot read ${ext} file: ${result.error || "conversion failed"}]` }]; @@ -2002,6 +2047,7 @@ export class ReadTool implements AgentTool { content = [{ type: "text", text: `[Cannot read ${ext} file: conversion failed]` }]; } } else { + repeatNotice = this.#repeatReadNotice(absolutePath); if ( parsed.kind === "none" && this.session.settings.get("read.summarize.enabled") && @@ -2043,6 +2089,7 @@ export class ReadTool implements AgentTool { parsed, displayMode, suffixResolution, + repeatNotice, undefined, // plain-file read: deterministic and fast, never abort mid-read ); if (multiResult.bridgeResult) return multiResult.bridgeResult; @@ -2066,6 +2113,7 @@ export class ReadTool implements AgentTool { sourcePath: absolutePath, entityLabel: "file", raw: isRawSelector(parsed), + repeatNotice, }); if (suffixResolution) { const notice = `[Path '${suffixResolution.from}' not found; resolved to '${suffixResolution.to}' via suffix match]`; @@ -2367,6 +2415,14 @@ export class ReadTool implements AgentTool { content = [{ type: "text", text: notice }, ...content]; } } + if (repeatNotice) { + // Trailing nudge goes at the very end of the textual result so it never + // disturbs hashline tag headers or inline notices. + const lastText = content.findLast((c): c is TextContent => c.type === "text"); + if (lastText) { + lastText.text = `${lastText.text}\n${repeatNotice}`; + } + } const resultBuilder = toolResult(details).content(content); if (sourcePath) { resultBuilder.sourcePath(sourcePath); diff --git a/packages/coding-agent/test/tools/read-repeat-notice.test.ts b/packages/coding-agent/test/tools/read-repeat-notice.test.ts new file mode 100644 index 000000000..91c5fc437 --- /dev/null +++ b/packages/coding-agent/test/tools/read-repeat-notice.test.ts @@ -0,0 +1,137 @@ +import { afterEach, beforeEach, describe, expect, it } from "bun:test"; +import * as fs from "node:fs/promises"; +import * as os from "node:os"; +import * as path from "node:path"; +import type { AgentToolResult } from "@oh-my-pi/pi-agent-core"; +import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; +import { + type InternalResource, + type InternalUrl, + InternalUrlRouter, + type ProtocolHandler, +} from "@oh-my-pi/pi-coding-agent/internal-urls"; +import type { ToolSession } from "@oh-my-pi/pi-coding-agent/tools"; +import type { ReadToolDetails } from "@oh-my-pi/pi-coding-agent/tools/read"; +import { ReadTool } from "@oh-my-pi/pi-coding-agent/tools/read"; + +const NOTICE_RE = + /\[note: read #(\d+) of this file this session — after edits, prefer the context echoed in the edit result or a narrow range re-read\]/; + +function textOutput(result: AgentToolResult): string { + return result.content + .filter(c => c.type === "text") + .map(c => c.text) + .join("\n"); +} + +function createSession(cwd: string): ToolSession { + const settings = Settings.isolated(); + // Deterministic plain-file reads regardless of language heuristics. + settings.set("read.summarize.enabled", false); + // URL reads must never reach the network in tests. + settings.set("fetch.enabled", false); + return { + cwd, + hasUI: false, + getSessionFile: () => path.join(cwd, "session.jsonl"), + getSessionSpawns: () => "*", + getArtifactsDir: () => path.join(cwd, "artifacts"), + allocateOutputArtifact: async () => ({ id: "artifact-1", path: path.join(cwd, "artifact-1.log") }), + settings, + }; +} + +function registerVirtualDoc(content: string): void { + const handler: ProtocolHandler = { + scheme: "virtual", + immutable: true, + async resolve(url: InternalUrl): Promise { + return { + url: url.href, + content, + contentType: "text/plain", + size: Buffer.byteLength(content, "utf-8"), + }; + }, + }; + InternalUrlRouter.instance().register(handler); +} + +function makeNumberedContent(lines: number): string { + return Array.from({ length: lines }, (_, i) => `line ${i + 1}`).join("\n"); +} + +describe("read tool repeat-read notice", () => { + let tmpDir: string; + + beforeEach(async () => { + tmpDir = await fs.mkdtemp(path.join(os.tmpdir(), "read-repeat-notice-test-")); + InternalUrlRouter.resetForTests(); + }); + + afterEach(async () => { + await fs.rm(tmpDir, { recursive: true, force: true }); + InternalUrlRouter.resetForTests(); + }); + + it("appends the notice on the third read of the same path, not on the first two", async () => { + const filePath = path.join(tmpDir, "sample.txt"); + await fs.writeFile(filePath, makeNumberedContent(10)); + const tool = new ReadTool(createSession(tmpDir)); + + const first = textOutput(await tool.execute("c1", { path: filePath })); + expect(first).toContain("line 1"); + expect(first).not.toMatch(NOTICE_RE); + + const second = textOutput(await tool.execute("c2", { path: filePath })); + expect(second).not.toMatch(NOTICE_RE); + + const third = textOutput(await tool.execute("c3", { path: filePath })); + expect(third).toContain("line 1"); + const match = third.match(NOTICE_RE); + expect(match?.[1]).toBe("3"); + // Appended at the very end of content, after the file body (never + // prepended, so hashline tag headers stay on the first line). + expect(third.trimEnd().endsWith("a narrow range re-read]")).toBe(true); + expect(third.indexOf("line 10")).toBeLessThan(third.search(NOTICE_RE)); + }); + + it("shares one counter across different selectors of the same file", async () => { + const filePath = path.join(tmpDir, "selectors.txt"); + await fs.writeFile(filePath, makeNumberedContent(20)); + const tool = new ReadTool(createSession(tmpDir)); + + const plain = textOutput(await tool.execute("s1", { path: filePath })); + expect(plain).not.toMatch(NOTICE_RE); + + const range = textOutput(await tool.execute("s2", { path: `${filePath}:2-4` })); + expect(range).not.toMatch(NOTICE_RE); + + const raw = textOutput(await tool.execute("s3", { path: `${filePath}:raw` })); + expect(raw.match(NOTICE_RE)?.[1]).toBe("3"); + + const multi = textOutput(await tool.execute("s4", { path: `${filePath}:1-2,5-6` })); + expect(multi.match(NOTICE_RE)?.[1]).toBe("4"); + }); + + it("never adds the notice for https:// or internal :// sources", async () => { + registerVirtualDoc(makeNumberedContent(5)); + const tool = new ReadTool(createSession(tmpDir)); + + for (let i = 1; i <= 4; i++) { + const text = textOutput(await tool.execute(`v${i}`, { path: "virtual://doc" })); + expect(text).toContain("line 1"); + expect(text).not.toMatch(NOTICE_RE); + } + + // https:// exits before any counting (fetch disabled in this session). + await expect(tool.execute("u1", { path: "https://example.com/page" })).rejects.toThrow("URL reads are disabled"); + + // The :// reads above never polluted the per-file counter: a real file + // still needs three reads of its own before the notice appears. + const filePath = path.join(tmpDir, "clean.txt"); + await fs.writeFile(filePath, makeNumberedContent(3)); + const first = textOutput(await tool.execute("f1", { path: filePath })); + expect(first).not.toMatch(NOTICE_RE); + }); +}); From 9d99ae1af0c2ccbd5418ec5c239407861955bbc3 Mon Sep 17 00:00:00 2001 From: can1357 Date: Wed, 10 Jun 2026 17:54:47 +0200 Subject: [PATCH 21/27] feat(coding-agent): rewrote the task tool to spawn one persistent subagent per call MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The task tool now takes a single { agent, assignment, description, ... } and always runs the subagent in the background — the batch tasks[] array and shared context parameter are gone. Fan-out is parallel task calls; shared background flows through a '/Users/can/.omp/agent/sessions/-Projects-.tree-pi-commit/2026-06-10T15-36-32-782Z_019eb22d-970e-7000-8964-72c98becf3e8/local' file referenced in each assignment.\n\nIntroduces a persistent subagent lifecycle: finished subagents stay live as idle, the lifecycle manager parks them to disk after task.agentIdleTtlMs (default 7 minutes; 0 keeps them live until exit), and they revive automatically when prompted from the Agent Hub, messaged on IRC, or resumed via task. New task(resume: "") revives an idle or parked subagent and runs a follow-up assignment in its existing session.\n\nAdds soft request budgets (explore/quick_task 40, others 90, configurable via task.softRequestBudget, 0 disables): crossing the budget injects a one-time wrap-up steer into the child; crossing 1.5× aborts the run gracefully. Cancelled/aborted subagent salvage replaces the old (no output) with the child's last activity snippet plus request/token stats; SingleResult tracks a per-child requests counter (assistant message_end events) used to sort agent lists in runtime-ascending order in both the live progress view (finished agents above pending/running) and the finalized result view, so rows no longer reshuffle on finalize. Adds a task gallery fixture variant for the resume path (renderer key separated from fixture key).\n\nAll task tests are reshaped around the single-call contract; tests for the discarded shared-context flow are removed, and new task-guards/task-resume/task-schema tests pin the new contract surface. --- docs/tools/task.md | 278 ++-- packages/coding-agent/src/cli/gallery-cli.ts | 2 +- .../src/cli/gallery-fixtures/agentic.ts | 403 +++-- .../src/cli/gallery-fixtures/types.ts | 5 + .../src/commit/agentic/tools/analyze-file.ts | 55 +- .../src/config/settings-schema.ts | 28 + .../src/eval/__tests__/agent-bridge.test.ts | 3 + .../src/prompts/system/orchestrate-notice.md | 4 +- .../src/prompts/tools/task-summary.md | 21 +- .../coding-agent/src/prompts/tools/task.md | 57 +- packages/coding-agent/src/task/executor.ts | 1239 ++++++++++----- packages/coding-agent/src/task/index.ts | 1361 +++++++---------- .../coding-agent/src/task/output-manager.ts | 11 - packages/coding-agent/src/task/render.ts | 159 +- packages/coding-agent/src/task/repair-args.ts | 47 +- packages/coding-agent/src/task/simple-mode.ts | 4 - packages/coding-agent/src/task/types.ts | 74 +- .../test/eval/agent-bridge.test.ts | 1 + .../coding-agent/test/rpc-subagents.test.ts | 1 + .../test/streaming-preview-height.test.ts | 21 +- .../task/executor-subagent-reminders.test.ts | 10 +- .../test/task/output-manager.test.ts | 8 +- .../test/task/render-call.test.ts | 133 +- .../test/task/render-nested-live.test.ts | 31 +- .../test/task/render-yield-shape.test.ts | 2 + .../test/task/subagent-lsp.test.ts | 4 +- .../test/task/task-guards.test.ts | 273 ++++ .../test/task/task-progress-render.test.ts | 127 +- .../test/task/task-resume.test.ts | 272 ++++ .../test/task/task-schema.test.ts | 83 + .../test/tool-live-region-scrollback.test.ts | 64 +- .../test/tools/task-repair-args.test.ts | 26 +- .../test/tools/task-simple-mode.test.ts | 133 +- .../swarm-extension/src/swarm/pipeline.ts | 1 + 34 files changed, 2996 insertions(+), 1945 deletions(-) create mode 100644 packages/coding-agent/test/task/task-guards.test.ts create mode 100644 packages/coding-agent/test/task/task-resume.test.ts create mode 100644 packages/coding-agent/test/task/task-schema.test.ts diff --git a/docs/tools/task.md b/docs/tools/task.md index a03e492b1..383d9c595 100644 --- a/docs/tools/task.md +++ b/docs/tools/task.md @@ -1,6 +1,6 @@ # task -> Launch subagents for parallel, optionally isolated work. +> Spawn one subagent per call to work in the background, or resume an existing one. ## Source - Entry: `packages/coding-agent/src/task/index.ts` @@ -9,217 +9,163 @@ - `packages/coding-agent/src/task/types.ts` — dynamic schema, progress/result types, output caps. - `packages/coding-agent/src/task/discovery.ts` — discover project/user/plugin/bundled agents. - `packages/coding-agent/src/task/agents.ts` — bundled agent definitions and frontmatter parsing. - - `packages/coding-agent/src/task/executor.ts` — create child sessions, run subagents, collect output. - - `packages/coding-agent/src/task/parallel.ts` — concurrency-limited scheduling and async semaphore. + - `packages/coding-agent/src/task/executor.ts` — create child sessions, run/resume subagents, collect output, hand finished sessions to the lifecycle manager. + - `packages/coding-agent/src/registry/agent-lifecycle.ts` — idle-TTL parking and revival of finished subagents. + - `packages/coding-agent/src/registry/agent-registry.ts` — process-global agent directory (`running | idle | parked | aborted`). + - `packages/coding-agent/src/async/job-manager.ts` — background job registration, progress, and result delivery. + - `packages/coding-agent/src/task/parallel.ts` — `Semaphore` used for the session-scoped concurrency bound. - `packages/coding-agent/src/task/isolation-backend.ts` — isolation backend resolution and platform fallback. - `packages/coding-agent/src/task/worktree.ts` — worktree / FUSE / ProjFS setup, patch capture, branch merge. - `packages/coding-agent/src/task/output-manager.ts` — session-scoped `agent://` id allocation. - - `packages/coding-agent/src/task/simple-mode.ts` — `default` / `schema-free` / `independent` field gating. + - `packages/coding-agent/src/task/name-generator.ts` — default AdjectiveNoun agent ids. + - `packages/coding-agent/src/task/simple-mode.ts` — `default` / `schema-free` / `independent` schema gating. - `packages/coding-agent/src/internal-urls/agent-protocol.ts` — resolve `agent://` to saved subagent output. + - `packages/coding-agent/src/internal-urls/history-protocol.ts` — resolve `history://` to a concise transcript. - `packages/coding-agent/src/tools/index.ts` — tool registration and recursion-depth gating. - `packages/coding-agent/src/sdk.ts` — child-session router/tool wiring and per-subagent `AgentOutputManager`. - `docs/task-agent-discovery.md` — deeper discovery and precedence notes. - - `docs/handoff-generation-pipeline.md` — session artifact/handoff persistence patterns used by the wider session layer. ## Inputs -### Default mode (`task.simple = "default"`) +One call spawns (or resumes) exactly one subagent. There is no batch parameter and no shared `context` parameter — shared background goes into a `local://` file (e.g. `local://ctx.md`) that each assignment references; subagents share the parent's `local://` root. | Field | Type | Required | Description | | --- | --- | --- | --- | -| `agent` | `string` | Yes | Exact agent name for every task item. Resolved at execution time through `discoverAgents(...)`. | -| `tasks` | `Array<{ id: string; description: string; assignment: string }>` | Yes | Batch of small, self-contained task items. `id` max length 48 in schema; duplicate ids are rejected case-insensitively at runtime. | -| `context` | `string` | No | Shared background prepended to every subagent system prompt. Trimmed before use. | -| `schema` | `string` | No | JSON-encoded JTD schema. Overrides agent/session output schema when this mode allows task-level schemas. | -| `isolated` | `boolean` | No | Only present when the tool is created with isolation enabled. Requests isolated execution for the whole batch. | +| `agent` | `string` | Conditional | Agent type to spawn. Required unless `resume` is set; providing both is a validation error. | +| `resume` | `string` | Conditional | Existing agent id — revive the agent if parked and run a follow-up assignment in its existing session. Cannot be combined with `agent` or `isolated`. | +| `id` | `string` | No | Stable agent id, schema max length 48. Defaults to a generated AdjectiveNoun name. Uniquified per session by `AgentOutputManager`. | +| `description` | `string` | No | UI label only; the subagent never sees it. | +| `assignment` | `string` | Yes | The work — complete, self-contained instructions. Empty-after-trim is rejected. | +| `schema` | `string` | No | JSON-encoded JTD schema for the expected `yield` payload. Field exists only when `task.simple = "default"`. | +| `isolated` | `boolean` | No | Run in an isolated workspace and return patches. Field exists only when `task.isolation.mode` is not `none`. Isolated agents are NOT resumable. | -`tasks[].description` is UI-only. `tasks[].assignment` is the actual per-task instruction. - -### Schema-free mode (`task.simple = "schema-free"`) - -Same as default, except `schema` is rejected by `validateTaskModeParams(...)` in `packages/coding-agent/src/task/index.ts`. - -### Independent mode (`task.simple = "independent"`) - -| Field | Type | Required | Description | -| --- | --- | --- | --- | -| `agent` | `string` | Yes | Exact agent name. | -| `tasks` | `Array<{ id: string; description: string; assignment: string }>` | Yes | Same item shape, but each `assignment` must carry all required background because shared `context` is disabled. | -| `isolated` | `boolean` | No | Same conditional field as above. | - -In this mode both `context` and `schema` are rejected. +Simple-mode gating (`task.simple`, one axis): `default` accepts the per-call `schema` override; `schema-free` and `independent` reject it (`validateTaskModeParams(...)`). `independent` additionally renders the subagent user prompt with the independent-mode flag. Agent frontmatter and inherited session schemas work in every mode. ## Outputs + The tool returns one text block plus `details: TaskToolDetails`. -`details` fields: -- `projectAgentsDir: string | null` — nearest discovered project `agents/` dir. -- `results: SingleResult[]` — one entry per task in input order for synchronous execution; empty for async-launch responses. -- `totalDurationMs: number` -- `usage?: Usage` — sum of per-subagent assistant-message usage. -- `outputPaths?: string[]` — written `.md` artifact paths for completed subagent outputs. -- `progress?: AgentProgress[]` — live or final per-task progress snapshots. -- `async?: { state: "running" | "completed" | "failed"; jobId: string; type: "task" }` — present for background execution updates/results. +Immediate (async) response — the normal case: +- `content`: `` Spawned agent `` (job ``). The result will be delivered when it yields. ... `` (or `Resumed agent ...`), plus a coordination hint (`irc` DM when enabled, otherwise `job`). +- `details`: `{ projectAgentsDir: null, results: [], totalDurationMs: 0, progress: [], async: { state: "running", jobId, type: "task" } }`. +- Live progress keeps streaming into the same tool block via `onUpdate(...)`; the final result arrives later as an async-result injection into the parent conversation. The delivery text appends a resume hint: `` is now idle — task(resume:"") to continue it, transcript at history:// `` (aborted variant points at the transcript only). + +Settled (sync-fallback or job-body) response: +- `content`: summary rendered from `packages/coding-agent/src/prompts/tools/task-summary.md` with a preview capped at 5000 chars; `agent://` holds the full output. +- `details.results`: at most one `SingleResult`; `usage`, `outputPaths` populated. `SingleResult` includes: - identity: `index`, `id`, `agent`, `agentSource`, `description`, optional `assignment` -- status: `exitCode`, optional `error`, optional `aborted`, optional `abortReason` -- output: `output`, `stderr`, `truncated`, `durationMs`, `tokens` +- status: `exitCode`, optional `error`, optional `aborted`, optional `abortReason`, optional `retryFailure` +- output: `output`, `stderr`, `truncated`, `durationMs`, `tokens`, `requests`, optional `contextTokens`/`contextWindow` - artifact metadata: `outputPath?`, `patchPath?`, `branchName?`, `nestedPatches?`, `outputMeta?` - extracted tool data: `extractedToolData?` from registered subprocess tool handlers such as `yield` and `report_finding` Artifacts and side channels: -- Every subagent with an artifacts dir writes `.md`; `agent://` resolves to that file. -- If the output file is JSON, `agent:///` and `agent://?q=` perform JSON extraction in `packages/coding-agent/src/internal-urls/agent-protocol.ts`. -- When the parent session persists artifacts, each subagent also gets `.jsonl` session history. -- Isolated patch mode writes `.patch` per successful task before merge. -- Async mode returns immediately after job registration, then emits `onUpdate(...)` progress snapshots and later hands completion to the session async-job pipeline. +- Every subagent with an artifacts dir writes `.md`; `agent://` resolves to that file. Resumes overwrite it per assignment. +- If the output file is JSON, `agent:///` and `agent://?q=` perform JSON extraction. +- Each subagent gets `.jsonl` session history when the parent persists artifacts; `history://` renders it as a concise transcript (works for live and parked agents). +- Isolated patch mode writes `.patch` before merge. ## Flow -1. `TaskTool.create(...)` in `packages/coding-agent/src/task/index.ts` calls `discoverAgents(session.cwd)` once to build the dynamic prompt description from current agents and `task.simple` capabilities. -2. `execute(...)` validates mode-gated fields with `validateTaskModeParams(...)`. -3. It decides async vs sync: - - sync when `async.enabled` is false - - sync when the selected cached agent has `blocking === true` - - sync when `tasks.length === 0` - - otherwise async job scheduling -4. Async path: - - allocate unique output ids with `AgentOutputManager.allocateBatch(...)` - - create one async job per task through `session.asyncJobManager.register(...)` - - limit concurrent job bodies with `Semaphore(task.maxConcurrency)` from `packages/coding-agent/src/task/parallel.ts` - - each job body calls `#executeSync(...)` with a one-task batch and the preallocated id - - `onUpdate(...)` emits aggregate `progress` snapshots and `details.async` -5. Sync path (`#executeSync(...)`) rediscovers agents from disk via `discoverAgents(...)`, so runtime resolution can differ from the earlier prompt description. -6. It resolves the requested agent with `getAgent(...)`, rejects unknown or disabled agents, and enforces parent spawn policy plus `PI_BLOCKED_AGENT` self-recursion prevention. -7. It derives the effective output schema in priority order: task call `schema` (if allowed) → agent frontmatter `output` → inherited parent session schema. -8. It validates task ids: missing ids and case-insensitive duplicates are immediate errors. -9. If `isolated` was requested, it requires a git repo (`getRepoRoot(...)` / `captureBaseline(...)`) and resolves the actual backend through `resolveIsolationBackendForTaskExecution(...)`. -10. It chooses an artifacts dir from the parent session when available, otherwise a temp dir, and writes `context.md` there when `session.getCompactContext?.()` returns content. -11. It allocates unique ids again if the caller did not preallocate them, then builds `tasksWithUniqueIds`. -12. For each task, it seeds an `AgentProgress` entry and runs `runTask(...)` through `mapWithConcurrencyLimit(...)` using `task.maxConcurrency`. -13. Non-isolated `runTask(...)` calls `runSubprocess(...)` directly with parent cwd. -14. Isolated `runTask(...)`: - - creates an isolation workspace (`ensureWorktree(...)`, `ensureFuseOverlay(...)`, or `ensureProjfsOverlay(...)`) - - applies the captured baseline for worktrees - - runs `runSubprocess(...)` inside that workspace - - on success, either commits to a per-task branch (`mergeMode === "branch"`) or captures a patch with `captureDeltaPatch(...)` - - always cleans up the isolation workspace/backend -15. `runSubprocess(...)` in `packages/coding-agent/src/task/executor.ts` creates a child agent session with: - - isolated settings snapshot via `Settings.isolated(...)`, forcing `async.enabled = false` and `bash.autoBackground.enabled = false` - - child `agentId` / `parentTaskPrefix` equal to the allocated task id - - child internal URL router and `AgentOutputManager` from `packages/coding-agent/src/sdk.ts` - - the shared `context`, optional `context.md` reference, optional isolation worktree path, output schema, and IRC peer roster in the system prompt template -16. Child tool availability is derived from the agent definition plus runtime guards: - - explicit `agent.tools` if provided - - auto-add `task` when the agent has `spawns` and recursion depth allows it - - remove `task` at or past `task.maxRecursionDepth` - - expand `exec` to `eval` and `bash` - - strip parent-owned `todo` after session creation -17. `runSubprocess(...)` subscribes to child agent events, coalesces progress updates every 150 ms, forwards lifecycle/progress events on the parent event bus, and extracts tool data through `subprocessToolRegistry`. -18. The child must finish through the hidden `yield` tool. If it does not, `runSubprocess(...)` sends up to 3 reminder prompts; the last reminder forces `toolChoice = yield` when supported. -19. Finalization uses `finalizeSubprocessOutput(...)` to reconcile raw assistant text, `yield` payloads, structured schemas, `report_finding` data, and abort states. Output is truncated with `MAX_OUTPUT_BYTES` / `MAX_OUTPUT_LINES` before returning to the parent, but the full raw output is still written to `.md`. -20. After all sync tasks finish, `#executeSync(...)` aggregates usage, collects artifact paths, and if isolation was used merges results back: - - branch mode: cherry-pick per-task branches with `mergeTaskBranches(...)`, then delete merged branches with `cleanupTaskBranches(...)` - - patch mode: combine non-empty patch artifacts, dry-check with `git.patch.canApplyText(...)`, then apply or leave manual artifacts - - nested repo patches are applied separately with `applyNestedPatches(...)` -21. The final text summary is rendered from `packages/coding-agent/src/prompts/tools/task-summary.md` and includes `agent://` handles for outputs that exist. +1. `TaskTool.create(...)` discovers agents once per cwd through a process-level memo (`discoverAgentsForCreate`) to render the dynamic prompt description. +2. `execute(...)` repairs raw params (`repairTaskParams`), then validates: schema gating per `task.simple`, `agent` XOR `resume`, `resume` excludes `isolated`, non-empty `assignment`. +3. Sync fallback only when the session has no `AsyncJobManager` (orphaned host) or the selected agent definition declares `blocking: true`; the call then runs `#executeSync(...)` inline under the session-scoped semaphore. +4. Otherwise execution is always async: + - the agent id is resolved up front — `resume` must name a registered agent (else `ToolError` pointing at `irc` op:"list" and `history://`); spawns allocate via `AgentOutputManager.allocate(params.id || generateTaskName())`; + - one `type: "task"` job is registered with `session.asyncJobManager` (`id` = agent id, `queued: true`, `ownerId` = caller agent id) and the tool returns immediately; + - the job body acquires the session-scoped `Semaphore` (one per `TaskTool` instance, sized from `task.maxConcurrency` at first use), marks the job running, runs `#executeSync(...)`, and reports progress through `buildAsyncDetails`/`onUpdate`; + - a failed or aborted run throws `TaskJobError` so the job lands `failed`, but the agent itself stays registered and interrogable. +5. `#executeSync(...)` dispatches: `resume` → `#executeResume(...)`, else `#runSpawn(...)`. +6. Resume path (`#executeResume`): + - `AgentLifecycleManager.global().ensureLive(resumeId)` returns the live session, reviving a parked one from its session JSONL; unknown ids or parked-without-reviver throw a `ToolError`; + - `resumeSubprocess(...)` in `packages/coding-agent/src/task/executor.ts` injects the rendered follow-up through the session's normal prompt path and drives it through the same monitor/yield/finalize pipeline as a spawn; + - the session is never disposed here — registry status settles back to `idle` (even on failure/abort) and the lifecycle manager re-arms the idle TTL. +7. Spawn path (`#runSpawn`) rediscovers agents from disk, so runtime resolution can differ from the create-time description. +8. It resolves the requested agent, rejects unknown or settings-disabled agents, and enforces parent spawn policy plus `PI_BLOCKED_AGENT` self-recursion prevention. +9. Output schema priority: task call `schema` (when `task.simple` allows) → agent frontmatter `output` → inherited parent session schema. +10. Plan mode swaps in an `effectiveAgent` with a read-only tool subset and plan-mode prompt; `runSubprocess(...)` receives the effective agent. +11. If `isolated`, it requires a git repo (`getRepoRoot(...)` / `captureBaseline(...)`) and resolves the backend through isolation-backend resolution with platform fallback. +12. Artifacts dir comes from the parent session file when available, otherwise a temp dir. When the session is executing an approved plan, the plan reference is handed to the subagent. +13. Non-isolated spawns call `runSubprocess(...)` directly with parent cwd; isolated spawns run inside the isolation workspace, then commit to a branch (`mergeMode === "branch"`) or capture a patch, and always clean up the workspace. +14. `runSubprocess(...)` creates a child agent session with an isolated settings snapshot (forcing `async.enabled = false` and `bash.autoBackground.enabled = false` — subagents are internally synchronous), child `agentId` equal to the allocated id, child internal URL router/`AgentOutputManager`, output schema, and the IRC peer roster in the system prompt. +15. Child tool availability: explicit `agent.tools` if provided; auto-add `task` when the agent has `spawns` and depth allows; strip `task` at `task.maxRecursionDepth`; expand `exec` to `eval` + `bash`; strip parent-owned `todo`. +16. The child must finish through the hidden `yield` tool; up to 3 reminder prompts, the last forcing `toolChoice = yield` when supported. `finalizeSubprocessOutput(...)` reconciles raw text, `yield` payloads, structured schemas, `report_finding` data, and abort states. +17. End-of-run lifecycle (keep-alive, in `runSubprocess`'s finalizer): + - hard abort (caller signal / wall-clock / budget) → registry status `aborted`, session disposed — terminal; + - isolated run → status `parked` without a reviver (workspace is merged + cleaned, so the session is not resumable; transcript stays readable via `history://`), then session disposed and detached; + - everything else (success and failure alike) → status `idle` with the live session attached, and `AgentLifecycleManager.global().adopt(id, { idleTtlMs, revive })` arms the park timer. The reviver reopens the session JSONL (park closed the writer, so the single-writer lock is taken cleanly). +18. Lifecycle thereafter: `idle` agents are parked after `task.agentIdleTtlMs` (session disposed; `AgentRef` + session file retained); messaging (`irc`), `task(resume:)`, or the Agent Hub revives them back to `idle`. `"Main"` is never parked. ## Modes / Variants - Execution mode - - Sync inline execution — default path. - - Async background execution — one async job per task item when `async.enabled` is on and the chosen agent is not marked `blocking`. -- Simple mode - - `default` — accepts shared `context` and per-call `schema`. - - `schema-free` — accepts `context`, rejects `schema`. - - `independent` — rejects `context` and `schema`; each assignment stands alone. -- Isolation backend - - `none` — no isolation. - - `worktree` — detached git worktree plus baseline replay. - - `fuse-overlay` — Unix FUSE overlay mount. - - `fuse-projfs` — Windows ProjFS overlay. -- Isolation merge strategy - - Patch mode — capture/apply root patches, keep patch artifacts when application fails. - - Branch mode — commit each task onto `omp/task/` branch, cherry-pick into parent, preserve failed branches for manual resolution. -- Agent source - - Project custom agents — nearest project config/plugin agent directories, first by source-family precedence. - - User custom agents — user config/plugin agent directories after project dirs of the same source family. - - Bundled agents — appended last from `packages/coding-agent/src/task/agents.ts`. -- Bundled agent types - - `explore` — read-only scout with structured handoff output. - - `plan` — architecture/planning agent; may spawn `explore`. - - `designer` — UI/UX specialist. - - `reviewer` — review agent with `report_finding` extraction. - - `task` — general-purpose worker with full capabilities. - - `quick_task` — low-reasoning mechanical worker using the same task prompt body. - - `librarian` — source-grounded external API/library researcher. - - `oracle` — senior-engineer implementation/debugging/general consultation agent. + - Always-async background job — default; spawn and resume both go through `AsyncJobManager`. + - Sync inline fallback — only when no job manager exists or the agent definition has `blocking: true`. +- Spawn vs resume + - `agent: ""` — fresh subagent with a new (or caller-provided) id. + - `resume: ""` — follow-up assignment in an existing session; revives a parked agent first. Transcript accretes; `agent://` is overwritten per assignment. +- Simple mode (`task.simple`) + - `default` — accepts per-call `schema`. + - `schema-free` / `independent` — reject `schema`; `independent` also flags the subagent user prompt as independent-mode. +- Isolation backend: `none`, `worktree`, `fuse-overlay`, `fuse-projfs`. +- Isolation merge strategy: patch mode (capture/apply root patches) or branch mode (commit to `omp/task/`, cherry-pick into parent). +- Agent source precedence: project custom agents, then user custom agents, then bundled agents (`explore`, `plan`, `designer`, `reviewer`, `task`, `quick_task`, `librarian`, `oracle`). ## Side Effects - Filesystem - - Writes `context.md`, `.jsonl`, and `.md` under the session artifacts dir or a temp task dir. - - In isolated patch mode writes `.patch` artifacts. - - Creates/removes worktrees or overlay mount directories. - - In branch mode creates temporary worktrees and task branches. + - Writes `.jsonl` and `.md` under the session artifacts dir or a temp task dir; isolated patch mode writes `.patch`. + - Creates/removes worktrees or overlay mount directories; branch mode creates temporary worktrees and task branches. - Network - Child sessions may use whichever networked tools/models their active tool set permits. - MCP proxy tools can call existing parent MCP connections with a 60_000 ms timeout. - Subprocesses / native bindings - - `fuse-overlayfs` and `fusermount`/`fusermount3` for FUSE isolation. - - ProjFS native bindings via `@oh-my-pi/pi-natives` on Windows. + - `fuse-overlayfs` and `fusermount`/`fusermount3` for FUSE isolation; ProjFS native bindings on Windows. - Git operations for baseline capture, patch apply, worktrees, branches, stash, cherry-pick, commits. - Session state (transcript, memory, jobs, checkpoints, registries) - - Creates child `AgentSession` instances with isolated settings snapshots. - - Registers async jobs in `session.asyncJobManager` for background task mode. + - Creates child `AgentSession` instances with isolated settings snapshots; finished sessions stay registered in the process-global `AgentRegistry` as `idle`/`parked` until process teardown or explicit release. + - Registers one async job per call in `session.asyncJobManager`; completion is injected into the parent as an async-result message. + - Arms idle-TTL timers in `AgentLifecycleManager` (unref'd; they never hold the process open). - Emits `task:subagent:event`, `task:subagent:progress`, and `task:subagent:lifecycle` on the parent event bus. - - Allocates session-scoped output ids through `AgentOutputManager` so `agent://` remains unique across invocations and resumes. - - Shares the parent `local://` root with subagents by passing `localProtocolOptions` through `createAgentSession(...)`. -- User-visible prompts / interactive UI - - Async mode streams aggregate progress updates. - - Missing-`yield` recovery sends up to three internal reminder prompts to the child session. - - Final summaries include `` blocks for isolation fallbacks or merge failures. + - Allocates session-scoped output ids through `AgentOutputManager` so `agent://` stays unique across invocations and resumes. + - Shares the parent `local://` root and `ArtifactManager` with subagents. - Background work / cancellation - - Parent abort stops scheduling new work, aborts active child sessions, and marks unscheduled tasks as skipped. - - Async jobs keep their own cancellation via `AsyncJobManager`. + - `job cancel` (or parent tool-call abort) cancels the job; a hard-aborted run lands `aborted` and is torn down. + - Missing-`yield` recovery sends up to three internal reminder prompts to the child session. ## Limits & Caps -- Per-subagent output truncation: `MAX_OUTPUT_BYTES = 500_000` and `MAX_OUTPUT_LINES = 5000` in `packages/coding-agent/src/task/types.ts`. Full raw output is still written to `.md` before truncation is returned to the caller. -- Progress coalescing in child execution: `PROGRESS_COALESCE_MS = 150` in `packages/coding-agent/src/task/executor.ts`. -- Recent output tail for progress: `RECENT_OUTPUT_TAIL_BYTES = 8 * 1024` and `recentOutput` keeps the last 8 non-empty lines in `packages/coding-agent/src/task/executor.ts`. -- Missing-`yield` reminder retries: `MAX_YIELD_RETRIES = 3` in `packages/coding-agent/src/task/executor.ts`. -- MCP proxy timeout: `MCP_CALL_TIMEOUT_MS = 60_000` in `packages/coding-agent/src/task/executor.ts`. -- Task id schema cap: `tasks[].id` `maxLength: 48` in `packages/coding-agent/src/task/types.ts`. -- Prompt text says ids should be `≤32` chars, but the runtime schema allows 48; this mismatch is real. -- Async/full sync parallelism both use `task.maxConcurrency` from settings: - - sync path: `mapWithConcurrencyLimit(...)` - - async path: `Semaphore(...)` around job bodies -- Recursion depth gate: `task.maxRecursionDepth` from settings; `packages/coding-agent/src/tools/index.ts` hides the `task` tool at or beyond the limit, and `runSubprocess(...)` also strips child `task` access at max depth. -- Final inline summary preview per task uses `fullOutputThreshold = 5000` chars in `packages/coding-agent/src/task/index.ts`; longer outputs are summarized while `agent://` points to the full artifact. +- Concurrency: one session-scoped `Semaphore` sized from `task.maxConcurrency` at first use (later setting changes do not resize it) bounds concurrent subagents across parallel `task` calls — both async job bodies and the sync fallback acquire it. +- Idle TTL: `task.agentIdleTtlMs`, default `420_000` ms (7 min); `<= 0` disables parking and keeps idle sessions live until exit. +- Per-subagent output truncation: `MAX_OUTPUT_BYTES = 500_000` and `MAX_OUTPUT_LINES = 5000` in `packages/coding-agent/src/task/types.ts` (overridable via `PI_TASK_MAX_OUTPUT_BYTES` / `PI_TASK_MAX_OUTPUT_LINES`). Full raw output is still written to `.md`. +- Progress coalescing: `PROGRESS_COALESCE_MS = 150`; recent-output tail: `RECENT_OUTPUT_TAIL_BYTES = 8 * 1024` (last 8 non-empty lines). +- Missing-`yield` reminder retries: `MAX_YIELD_RETRIES = 3`; MCP proxy timeout: `MCP_CALL_TIMEOUT_MS = 60_000` — both in `packages/coding-agent/src/task/executor.ts`. +- Agent id schema cap: `id` `maxLength: 48` in `packages/coding-agent/src/task/types.ts`. Prompt text says ids should be `≤32` chars; this mismatch is real. +- Soft request budget (`task.softRequestBudget`) and wall clock (`task.maxRuntimeMs`) apply to spawns and resumes alike. +- Recursion depth gate: `task.maxRecursionDepth`; `packages/coding-agent/src/tools/index.ts` hides the `task` tool at or beyond the limit, and `runSubprocess(...)` also strips child `task` access at max depth. +- Final inline summary preview uses `fullOutputThreshold = 5000` chars in `packages/coding-agent/src/task/index.ts`; `agent://` points to the full artifact. ## Errors -- Most validation failures are returned as normal tool text with empty `results`, not thrown: - - invalid simple-mode fields - - unknown/disabled agent - - missing tasks - - missing/duplicate task ids - - spawn-policy denial - - requesting `isolated` while isolation mode is `none` -- Isolated execution without a git repo returns `Isolated task execution requires a git repository. ...`. -- Backend resolution can return a hard error (`ProjFS isolation initialization failed...`) or a non-fatal warning with fallback to `worktree`. -- `mapWithConcurrencyLimit(...)` fails fast on non-abort worker exceptions; already completed results are preserved only in the thrown path’s local state, not surfaced unless the caller catches and converts them. -- Child-session failures surface as `SingleResult.exitCode = 1` with `stderr`/`error` populated. +- Parameter validation failures are returned as normal tool text with empty `results`: + - `schema` outside `task.simple = "default"` + - both or neither of `agent` / `resume` + - `resume` combined with `isolated` + - missing/empty `assignment` + - unknown or settings-disabled agent, spawn-policy denial, requesting `isolated` while isolation mode is `none` +- `resume` of an id not in the registry throws a `ToolError` naming `irc` op:"list" and `history://`. +- `ensureLive(...)` failures (agent parked without a reviver — e.g. an isolated run — or torn down) surface as `` Cannot resume "": ... `` `ToolError`s. +- Isolated execution without a git repo returns `Isolated task execution requires a git repository. ...`; backend resolution can hard-error (ProjFS init) or warn and fall back to `worktree`. +- Job registration failure returns `Failed to start background task job: ...`. +- Child failures surface as `SingleResult.exitCode = 1` with `stderr`/`error` populated; the async job is marked failed but the delivery text still carries the output plus a resume/transcript hint. - If the child omits `yield`, `finalizeSubprocessOutput(...)` injects warnings such as `SYSTEM WARNING: Subagent exited without calling yield tool after 3 reminders.` -- Async scheduling failures are accumulated per task; if no jobs start, the tool returns `Failed to start background task jobs: ...`. - `agent://` resolution errors are model-visible when another tool reads them: no session, no artifacts dir, missing id, conflicting extraction syntax, or invalid JSON for extraction. ## Notes -- Agent discovery precedence is first-wins by exact name: project dirs before user dirs within a source family, plugin agent dirs after config dirs, bundled agents last. See `packages/coding-agent/src/task/discovery.ts` and `docs/task-agent-discovery.md`. -- `TaskTool.create(...)` caches discovered agents only for description rendering and the async blocking-agent decision. `#executeSync(...)` rediscovers agents each call. -- Custom agent frontmatter can override bundled agents by name. Bundled definitions are embedded at build time in `packages/coding-agent/src/task/agents.ts`. -- Child sessions do not inherit conversation history automatically. The only built-in carry-over is shared `context`, optional `context.md`, workspace tree/skills/context files, and shared `local://` root. -- `Settings.isolated(...)` gives each child a session-isolated settings snapshot; tool enablement is recomputed inside the child session rather than sharing mutable parent tool state. -- When the parent passes `mcpManager`, child sessions disable standalone MCP discovery and instead get proxy tools that reuse the parent connections. -- Plan mode mutates an `effectiveAgent` with a read-only tool subset and plan-mode prompt text, but `runSubprocess(...)` is still invoked with `agent` rather than `effectiveAgent`. Model/thinking/schema overrides use the effective agent; prompt/tool/spawn restrictions do not fully flow through this call path. -- Branch-mode merge temporarily stashes the parent repo before cherry-picking task branches. A stash-pop conflict is treated as merge failure and leaves recovery state behind. -- Patch-mode only applies combined root patches if every successful task produced a patch and `git.patch.canApplyText(...)` succeeds. -- Nested git repos are handled separately from the root repo. They are copied into isolated worktrees, diffed independently, and merged later with `applyNestedPatches(...)` because parent git cannot track their file-level changes. -- `agent://` ids are name-based (`Task` first, `Task-2`/`Task-3` only when the name repeats, nested like `Parent.Child`) by `AgentOutputManager`; this is what prevents artifact collisions across repeated or nested task invocations. +- Parallelism is parallel `task` calls in one assistant message; the session-scoped semaphore bounds the fan-out. There is no batch array. +- Shared background convention: write it once to a `local://` file and reference that path in each assignment — subagents share the parent's `local://` root. This replaces the removed `context` parameter. +- Prefer `resume` over a fresh spawn for follow-up work: the resumed agent already holds the relevant context. `irc` op:"list" shows idle/parked candidates; `history://` shows what an agent has done. +- Subagents are internally synchronous: the executor forces `async.enabled = false` and `bash.autoBackground.enabled = false` in the child settings snapshot, so there are no fire-and-forget grandchildren. +- Agent discovery precedence is first-wins by exact name: project dirs before user dirs within a source family, plugin agent dirs after config dirs, bundled agents last. Create-time discovery is memoized per cwd for the prompt description; execution-time discovery stays fresh. +- Child sessions do not inherit conversation history. Built-in carry-over is the workspace tree/skills/context files, the shared `local://` root, and the approved-plan reference when one exists. +- When the parent passes `mcpManager`, child sessions disable standalone MCP discovery and get proxy tools that reuse parent connections. +- Branch-mode merge temporarily stashes the parent repo before cherry-picking; a stash-pop conflict is treated as merge failure and leaves recovery state behind. Patch mode only applies the combined root patch when `git.patch.canApplyText(...)` succeeds; failures leave the `.patch` artifact for manual handling. +- Nested git repos are diffed independently inside isolated workspaces and merged separately with `applyNestedPatches(...)`. +- `agent://` ids are name-based (`Task` first, `Task-2`/`Task-3` only when the name repeats, nested like `Parent.Child`) by `AgentOutputManager`; this is what prevents artifact collisions across repeated or nested invocations. diff --git a/packages/coding-agent/src/cli/gallery-cli.ts b/packages/coding-agent/src/cli/gallery-cli.ts index 7f4eb8c67..31d1a958b 100644 --- a/packages/coding-agent/src/cli/gallery-cli.ts +++ b/packages/coding-agent/src/cli/gallery-cli.ts @@ -69,7 +69,7 @@ function fakeToolFor(name: string, fixture: GalleryFixture | undefined): AgentTo if (!fixture?.label && !fixture?.editMode && !fixture?.customRendered) return undefined; const tool: Record = { name, label: fixture.label ?? name, mode: fixture.editMode }; if (fixture.customRendered) { - const renderer = toolRenderers[name] as + const renderer = toolRenderers[fixture.renderer ?? name] as | { renderCall?: unknown; renderResult?: unknown; mergeCallAndResult?: unknown; inline?: unknown } | undefined; if (renderer) { diff --git a/packages/coding-agent/src/cli/gallery-fixtures/agentic.ts b/packages/coding-agent/src/cli/gallery-fixtures/agentic.ts index d1c262abe..df66493a9 100644 --- a/packages/coding-agent/src/cli/gallery-fixtures/agentic.ts +++ b/packages/coding-agent/src/cli/gallery-fixtures/agentic.ts @@ -1,62 +1,76 @@ -// Gallery fixtures for the agentic orchestration tools (task, goal, job). +// Gallery fixtures for the agentic orchestration tools (task, irc, goal, job). +import type { Usage } from "@oh-my-pi/pi-ai"; +import type { TaskToolDetails } from "../../task/types"; +import type { IrcDetails } from "../../tools/irc"; import type { GalleryFixture } from "./types"; +/** Message/activity timestamps are offsets from load time so gallery ages stay plausible. */ +const FIXTURE_NOW = Date.now(); + +/** Plausible cumulative usage for a fixture subagent run. */ +const fixtureUsage = (tokens: { input: number; output: number }, costTotal: number): Usage => ({ + input: tokens.input, + output: tokens.output, + cacheRead: 0, + cacheWrite: 0, + totalTokens: tokens.input + tokens.output, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: costTotal }, +}); + export const agenticFixtures: Record = { task: { label: "Task", customRendered: true, - // Streaming: agent chosen, first task fully arrived, second still landing. + // Streaming: agent chosen, assignment still landing. streamingArgs: { agent: "task", - tasks: [ - { - id: "AuthLoader", - description: "Load auth middleware", - assignment: "Read packages/server/src/auth/*.ts and summarize the session-cookie flow.", - }, - { id: "RateLimiter", description: "Audit rate limiter" }, - ], + id: "AuthLoader", + description: "Load auth middleware", + assignment: "Read packages/server/src/auth/*.ts and summarize the session-cookie", }, args: { agent: "task", - context: [ - "# Goal", - "Harden the HTTP auth stack before the release cut.", - "# Constraints", - "Touch only files under packages/server/src/auth/. Do not run gates.", - ].join("\n"), - tasks: [ - { - id: "AuthLoader", - description: "Load auth middleware", - assignment: - "Read packages/server/src/auth/session.ts and middleware.ts, then document the session-cookie validation flow and any TODOs.", - }, - { - id: "RateLimiter", - description: "Audit rate limiter", - assignment: - "Inspect packages/server/src/auth/rate-limit.ts. Confirm the 429 path sets Retry-After and report gaps.", - }, - { - id: "TokenRotation", - description: "Check token rotation", - assignment: - "Trace refresh-token rotation in packages/server/src/auth/tokens.ts and flag any reuse window.", - }, - ], + id: "AuthLoader", + description: "Load auth middleware", + assignment: + "Read packages/server/src/auth/session.ts and middleware.ts, then document the session-cookie validation flow and any TODOs.", }, result: { content: [ { type: "text", - text: "3 agents completed: AuthLoader, RateLimiter, TokenRotation.", + text: "Agent AuthLoader completed.", }, ], details: { projectAgentsDir: null, totalDurationMs: 48_200, - usage: { cost: { total: 0.34 } }, + usage: fixtureUsage({ input: 52_600, output: 8_800 }, 0.12), + progress: [ + { + index: 0, + id: "AuthLoader", + agent: "task", + agentSource: "bundled", + status: "completed", + task: "Read packages/server/src/auth/session.ts and middleware.ts", + description: "Load auth middleware", + lastIntent: "Documenting session-cookie flow", + recentTools: [ + { tool: "read", args: "packages/server/src/auth/session.ts", endMs: 1_749_200_040_000 }, + { tool: "read", args: "packages/server/src/auth/middleware.ts", endMs: 1_749_200_052_000 }, + ], + recentOutput: ["Session validation runs in middleware.ts:42 via verifySessionCookie()."], + toolCount: 9, + requests: 6, + tokens: 61_400, + contextTokens: 23_100, + contextWindow: 200_000, + cost: 0.12, + durationMs: 41_900, + resolvedModel: "anthropic/claude-sonnet", + }, + ], results: [ { index: 0, @@ -77,100 +91,31 @@ export const agenticFixtures: Record = { truncated: false, durationMs: 41_900, tokens: 61_400, + requests: 6, contextTokens: 23_100, contextWindow: 200_000, resolvedModel: "anthropic/claude-sonnet", - usage: { cost: { total: 0.12 } }, + usage: fixtureUsage({ input: 52_600, output: 8_800 }, 0.12), outputMeta: { lineCount: 3, charCount: 214 }, }, - { - index: 1, - id: "RateLimiter", - agent: "task", - agentSource: "bundled", - description: "Audit rate limiter", - task: "Inspect packages/server/src/auth/rate-limit.ts", - assignment: - "Inspect packages/server/src/auth/rate-limit.ts. Confirm the 429 path sets Retry-After and report gaps.", - exitCode: 0, - output: [ - "rate-limit.ts uses a fixed-window counter keyed by client IP.", - "429 responses set Retry-After (rate-limit.ts:57).", - "Gap: no per-account limit, so a botnet across IPs bypasses the cap.", - ].join("\n"), - stderr: "", - truncated: false, - durationMs: 38_500, - tokens: 54_800, - contextTokens: 19_700, - contextWindow: 200_000, - resolvedModel: "anthropic/claude-sonnet", - usage: { cost: { total: 0.1 } }, - outputMeta: { lineCount: 3, charCount: 198 }, - }, - { - index: 2, - id: "TokenRotation", - agent: "task", - agentSource: "bundled", - description: "Check token rotation", - task: "Trace refresh-token rotation in packages/server/src/auth/tokens.ts", - assignment: - "Trace refresh-token rotation in packages/server/src/auth/tokens.ts and flag any reuse window.", - exitCode: 0, - output: [ - "Refresh tokens rotate on every use (tokens.ts:120) and the old jti is revoked.", - "Reuse of a rotated token triggers full-family revocation — no reuse window found.", - ].join("\n"), - stderr: "", - truncated: false, - durationMs: 48_200, - tokens: 49_200, - contextTokens: 17_500, - contextWindow: 200_000, - resolvedModel: "anthropic/claude-sonnet", - usage: { cost: { total: 0.12 } }, - outputMeta: { lineCount: 2, charCount: 160 }, - }, ], - }, + } satisfies TaskToolDetails, }, errorResult: { isError: true, content: [ { type: "text", - text: "1 of 3 agents failed: RateLimiter.", + text: "Agent RateLimiter failed.", }, ], details: { projectAgentsDir: null, - totalDurationMs: 39_400, - usage: { cost: { total: 0.21 } }, + totalDurationMs: 9_800, + usage: fixtureUsage({ input: 10_900, output: 1_400 }, 0.1), results: [ { index: 0, - id: "AuthLoader", - agent: "task", - agentSource: "bundled", - description: "Load auth middleware", - task: "Read packages/server/src/auth/session.ts and middleware.ts", - assignment: - "Read packages/server/src/auth/session.ts and middleware.ts, then document the session-cookie validation flow and any TODOs.", - exitCode: 0, - output: "Session validation runs in middleware.ts:42 via verifySessionCookie().", - stderr: "", - truncated: false, - durationMs: 31_200, - tokens: 58_100, - contextTokens: 21_900, - contextWindow: 200_000, - resolvedModel: "anthropic/claude-sonnet", - usage: { cost: { total: 0.11 } }, - outputMeta: { lineCount: 1, charCount: 70 }, - }, - { - index: 1, id: "RateLimiter", agent: "task", agentSource: "bundled", @@ -184,15 +129,243 @@ export const agenticFixtures: Record = { truncated: false, durationMs: 9_800, tokens: 12_300, + requests: 3, contextTokens: 6_400, contextWindow: 200_000, resolvedModel: "anthropic/claude-sonnet", - usage: { cost: { total: 0.1 } }, + usage: fixtureUsage({ input: 10_900, output: 1_400 }, 0.1), error: "Subagent exited 1: target file packages/server/src/auth/rate-limit.ts does not exist.", outputMeta: { lineCount: 0, charCount: 0 }, }, ], - }, + } satisfies TaskToolDetails, + }, + }, + + // Resume: follow-up assignment into an existing (idle or parked) agent. + task_resume: { + label: "Task (resume)", + customRendered: true, + renderer: "task", + // Streaming: resume target known; the follow-up assignment still landing. + streamingArgs: { + resume: "AuthLoader", + assignment: "Follow up: does the sliding-expiration TODO affect", + }, + args: { + resume: "AuthLoader", + assignment: + "Follow up: does the sliding-expiration TODO at session.ts:88 affect the refresh-token path? Document the answer.", + }, + result: { + content: [{ type: "text", text: "Agent AuthLoader completed." }], + details: { + projectAgentsDir: null, + totalDurationMs: 22_400, + usage: fixtureUsage({ input: 30_200, output: 4_100 }, 0.07), + results: [ + { + index: 0, + id: "AuthLoader", + agent: "task", + agentSource: "bundled", + task: "Follow up: does the sliding-expiration TODO at session.ts:88 affect the refresh-token path?", + assignment: + "Follow up: does the sliding-expiration TODO at session.ts:88 affect the refresh-token path? Document the answer.", + exitCode: 0, + output: + "No — refresh tokens bypass the sliding window: refreshSession() re-issues the cookie unconditionally (session.ts:131).", + stderr: "", + truncated: false, + durationMs: 19_700, + tokens: 34_300, + requests: 4, + contextTokens: 31_800, + contextWindow: 200_000, + resolvedModel: "anthropic/claude-sonnet", + usage: fixtureUsage({ input: 30_200, output: 4_100 }, 0.07), + outputMeta: { lineCount: 1, charCount: 118 }, + }, + ], + } satisfies TaskToolDetails, + }, + errorResult: { + isError: true, + content: [ + { + type: "text", + text: 'No agent "AuthLoader" to resume — it ran isolated and is not revivable. See history:// for the agent index.', + }, + ], + }, + }, + + irc: { + label: "IRC", + // Streaming: recipient known; the message body still arriving. + streamingArgs: { op: "send", to: "AuthLoader", message: "Are you still touching" }, + args: { + op: "send", + to: "AuthLoader", + message: "Are you still touching src/server/auth.ts? I need to add a 401 path.", + await: true, + }, + result: { + content: [ + { + type: "text", + text: [ + "Delivered to 1 peer(s):", + "- AuthLoader: revived", + "", + "Reply from AuthLoader:", + "Done with auth.ts — go ahead, just rebase past my session-store rename.", + ].join("\n"), + }, + ], + details: { + op: "send", + from: "Main", + to: "AuthLoader", + receipts: [{ to: "AuthLoader", outcome: "revived" }], + waited: { + id: "7181122334455667789", + from: "AuthLoader", + to: "Main", + body: "Done with auth.ts — go ahead, just rebase past my session-store rename.", + ts: FIXTURE_NOW - 5_000, + replyTo: "7181122334455667788", + }, + } satisfies IrcDetails, + }, + errorResult: { + isError: true, + content: [ + { + type: "text", + text: 'No recipients received the message.\n- RateLimiter: failed — unknown agent "RateLimiter"', + }, + ], + details: { + op: "send", + from: "Main", + to: "RateLimiter", + receipts: [{ to: "RateLimiter", outcome: "failed", error: 'unknown agent "RateLimiter"' }], + } satisfies IrcDetails, + }, + }, + + irc_wait: { + label: "IRC (wait)", + customRendered: true, + renderer: "irc", + streamingArgs: { op: "wait", from: "AuthLoader" }, + args: { op: "wait", from: "AuthLoader", timeoutMs: 60_000 }, + result: { + content: [ + { + type: "text", + text: "[7181122334455667790] AuthLoader: session-store rename is merged; auth.ts is yours.", + }, + ], + details: { + op: "wait", + from: "Main", + waited: { + id: "7181122334455667790", + from: "AuthLoader", + to: "Main", + body: "session-store rename is merged; auth.ts is yours.", + ts: FIXTURE_NOW - 30_000, + }, + } satisfies IrcDetails, + }, + }, + + irc_inbox: { + label: "IRC (inbox)", + customRendered: true, + renderer: "irc", + streamingArgs: { op: "inbox" }, + args: { op: "inbox", peek: true }, + result: { + content: [ + { + type: "text", + text: [ + "2 unread message(s):", + "- [7181122334455667791] AuthLoader: hub table reads unreadCount — ping me when the bus lands.", + "- [7181122334455667792] RateLimiter (reply to 7181122334455667791): bus is in; receipts carry outcome.", + ].join("\n"), + }, + ], + details: { + op: "inbox", + from: "Main", + inbox: [ + { + id: "7181122334455667791", + from: "AuthLoader", + to: "Main", + body: "hub table reads unreadCount — ping me when the bus lands.", + ts: FIXTURE_NOW - 4 * 60_000, + }, + { + id: "7181122334455667792", + from: "RateLimiter", + to: "Main", + body: "bus is in; receipts carry outcome.", + ts: FIXTURE_NOW - 60_000, + replyTo: "7181122334455667791", + }, + ], + } satisfies IrcDetails, + }, + }, + + irc_list: { + label: "IRC (list)", + customRendered: true, + renderer: "irc", + streamingArgs: { op: "list" }, + args: { op: "list" }, + result: { + content: [ + { + type: "text", + text: [ + "2 peer(s):", + "- AuthLoader [task · sub · idle] — parent Main, active 2m ago", + "- RateLimiter [task · sub · parked] — unread 2, parent Main, active 12m ago", + "", + "Parked agents are revived automatically when you message them.", + ].join("\n"), + }, + ], + details: { + op: "list", + from: "Main", + peers: [ + { + id: "AuthLoader", + displayName: "task", + kind: "sub", + status: "idle", + parentId: "Main", + unread: 0, + lastActivity: FIXTURE_NOW - 2 * 60_000, + }, + { + id: "RateLimiter", + displayName: "task", + kind: "sub", + status: "parked", + parentId: "Main", + unread: 2, + lastActivity: FIXTURE_NOW - 12 * 60_000, + }, + ], + } satisfies IrcDetails, }, }, diff --git a/packages/coding-agent/src/cli/gallery-fixtures/types.ts b/packages/coding-agent/src/cli/gallery-fixtures/types.ts index da4b9b2e4..cdf935e16 100644 --- a/packages/coding-agent/src/cli/gallery-fixtures/types.ts +++ b/packages/coding-agent/src/cli/gallery-fixtures/types.ts @@ -36,6 +36,11 @@ export interface GalleryFixture { * real one keeps the gallery honest for these tools. */ customRendered?: boolean; + /** + * Renderer-registry key to use when the fixture key is a variant of a tool + * (e.g. `task_resume` → `task`). Defaults to the fixture key. + */ + renderer?: string; /** * Arguments shown during the streaming state — a partial view of {@link args} * as if the tool-call JSON were still arriving. May include `__partialJson` diff --git a/packages/coding-agent/src/commit/agentic/tools/analyze-file.ts b/packages/coding-agent/src/commit/agentic/tools/analyze-file.ts index 78f0e7b2b..7d09113e7 100644 --- a/packages/coding-agent/src/commit/agentic/tools/analyze-file.ts +++ b/packages/coding-agent/src/commit/agentic/tools/analyze-file.ts @@ -59,29 +59,46 @@ export function createAnalyzeFileTool(options: { label: "Analyze Files", description: "Spawn quick_task agents to analyze files.", parameters: analyzeFileSchema, - async execute(toolCallId, params, onUpdate, ctx, signal) { + async execute(toolCallId, params, _onUpdate, ctx, signal) { const toolSession = buildToolSession(ctx, options); + // The hand-built ToolSession carries no asyncJobManager, so every + // execute() below takes the task tool's sync fallback and resolves + // with the subagent's result inline — exactly what this flow needs. + // The tool's session semaphore bounds the parallel fan-out. const taskTool = await TaskTool.create(toolSession); const numstat = options.state.overview?.numstat ?? []; - const tasks = params.files.map((file, index) => { - const relatedFiles = formatRelatedFiles(params.files, file, numstat); - const assignment = prompt.render(analyzeFilePrompt, { - file, - goal: params.goal, - related_files: relatedFiles, - }); - return { - id: `AnalyzeFile${index + 1}`, - description: `Analyze ${file}`, - assignment, - }; - }); - const taskParams: TaskParams = { - agent: "quick_task", - schema: JSON.stringify(analyzeFileOutputSchema), - tasks, + const schema = JSON.stringify(analyzeFileOutputSchema); + const analyses = await Promise.all( + params.files.map((file, index) => { + const relatedFiles = formatRelatedFiles(params.files, file, numstat); + const assignment = prompt.render(analyzeFilePrompt, { + file, + goal: params.goal, + related_files: relatedFiles, + }); + const taskParams: TaskParams = { + agent: "quick_task", + id: `AnalyzeFile${index + 1}`, + description: `Analyze ${file}`, + assignment, + schema, + }; + return taskTool.execute(`${toolCallId}-${index + 1}`, taskParams, signal); + }), + ); + const results = analyses.flatMap(analysis => analysis.details?.results ?? []); + const text = analyses + .map(analysis => analysis.content.find(part => part.type === "text")?.text ?? "") + .filter(Boolean) + .join("\n\n"); + return { + content: [{ type: "text", text: text || "(no output)" }], + details: { + projectAgentsDir: null, + results, + totalDurationMs: analyses.reduce((sum, analysis) => sum + (analysis.details?.totalDurationMs ?? 0), 0), + }, }; - return taskTool.execute(toolCallId, taskParams, signal, onUpdate); }, }; } diff --git a/packages/coding-agent/src/config/settings-schema.ts b/packages/coding-agent/src/config/settings-schema.ts index 3b2edc579..3b20c0387 100644 --- a/packages/coding-agent/src/config/settings-schema.ts +++ b/packages/coding-agent/src/config/settings-schema.ts @@ -2856,6 +2856,34 @@ export const SETTINGS_SCHEMA = { }, }, + "task.agentIdleTtlMs": { + type: "number", + default: 420_000, + ui: { + tab: "tasks", + label: "Agent Idle TTL", + description: + "How long an idle subagent stays live in memory before being parked to disk (ms). Parked agents are revived automatically when messaged or resumed. 0 keeps idle agents live until exit.", + }, + }, + + "task.softRequestBudget": { + type: "number", + default: 90, + ui: { + tab: "tasks", + label: "Soft Subagent Request Budget", + description: + "Soft per-subagent request budget (assistant requests per run). Crossing it injects one steering notice asking the subagent to wrap up; at 1.5x the budget the run is aborted gracefully, salvaging partial output. 0 disables the guard. Bundled explore/quick_task agents use a lower built-in budget.", + options: [ + { value: "0", label: "Disabled" }, + { value: "40", label: "40 requests" }, + { value: "90", label: "90 requests", description: "Default" }, + { value: "150", label: "150 requests" }, + ], + }, + }, + "task.disabledAgents": { type: "array", default: [] as string[], diff --git a/packages/coding-agent/src/eval/__tests__/agent-bridge.test.ts b/packages/coding-agent/src/eval/__tests__/agent-bridge.test.ts index 588fd1b5a..1c63b3431 100644 --- a/packages/coding-agent/src/eval/__tests__/agent-bridge.test.ts +++ b/packages/coding-agent/src/eval/__tests__/agent-bridge.test.ts @@ -99,6 +99,7 @@ function singleResult(options: ExecutorOptions, overrides: Partial truncated: false, durationMs: 1, tokens: 0, + requests: 0, ...overrides, }; } @@ -541,6 +542,7 @@ describe("agent() through eval runtimes", () => { recentOutput: [], toolCount: 0, tokens: 0, + requests: 0, cost: 0, durationMs: 0, ...overrides, @@ -673,6 +675,7 @@ describe("agent() through eval runtimes", () => { recentOutput: [], toolCount: i, tokens: 0, + requests: 0, cost: 0, durationMs: i * 10, }); diff --git a/packages/coding-agent/src/prompts/system/orchestrate-notice.md b/packages/coding-agent/src/prompts/system/orchestrate-notice.md index c8086fbb4..cd2754909 100644 --- a/packages/coding-agent/src/prompts/system/orchestrate-notice.md +++ b/packages/coding-agent/src/prompts/system/orchestrate-notice.md @@ -8,7 +8,7 @@ You decompose, dispatch, verify, and iterate. Substantial and parallelizable wor 1. **NEVER yield until everything is closed.** A phase finishing is *not* a yield point — launch the next phase in the same turn. Stop only when every requested item is verifiably done, or you hit a concrete [blocked] state that genuinely requires the user. 2. **Enumerate the full surface before dispatching.** If the request references audits, plans, checklists, phase lists, or file lists, expand them into a flat set of items in `todo`. "Most of them" or "the important ones" is failure. Re-read the source documents — NEVER work from memory. -3. **Parallelize maximally; NEVER launch a one-off task.** Every set of edits with disjoint file scope MUST ship as one `task` batch — fan the work as wide as it decomposes. A single-task batch for divisible work is a failure: split it. If you are about to dispatch exactly one subagent, stop — either there is more to run alongside it (find it and batch them) or the change is small enough to make inline yourself (do it). Serialize only when one subagent produces a contract (types, schema, shared module) the next consumes — and state the dependency when you do. +3. **Parallelize maximally; NEVER launch a one-off task.** Every set of edits with disjoint file scope MUST ship as parallel `task` calls in one message — fan the work as wide as it decomposes. Dispatching divisible work one call at a time, serially, is a failure: split it and dispatch together. If you are about to dispatch exactly one subagent, stop — either there is more to run alongside it (find it and dispatch them together) or the change is small enough to make inline yourself (do it). Serialize only when one subagent produces a contract (types, schema, shared module) the next consumes — and state the dependency when you do. 4. **Each `task` assignment is self-contained.** Subagents have no shared context. Spell out: target files (≤3–5 explicit paths, no globs), the change with APIs and patterns, edge cases, and observable acceptance criteria. NEVER assume they read the same plan you did. 5. **Verify after every phase before launching the next.** Run the appropriate gate: `bun check` for types, package-scoped `bun test` for behavior, `lsp diagnostics` for changed files. If a phase introduced breakage, dispatch fix-up subagents *before* moving on. NEVER declare a phase done on a red tree. 6. **Commit policy.** If the request asks for commits or the repo workflow expects them, commit after each green phase with a focused message. NEVER commit a red tree. NEVER commit work the user did not ask to commit. @@ -21,7 +21,7 @@ You decompose, dispatch, verify, and iterate. Substantial and parallelizable wor 1. **Ingest.** Read every referenced file (audits, plans, prior agent output, current branch state). Run `git status` to see uncommitted changes. 2. **Plan.** Materialize the full work surface in `todo` as ordered phases. Within each phase, list the parallelizable units. -3. **Dispatch phase.** Launch all parallel `task` subagents in one call. Wait for the batch. +3. **Dispatch phase.** Launch all parallel `task` subagents in one message, then collect every result (async results / `job poll`) before moving on. 4. **Verify phase.** Run the gates. On failure, dispatch fix-up subagents and re-verify. Do not advance with a red gate. 5. **Commit phase** (if applicable). Focused message naming the phase. 6. **Advance.** Mark the phase done in `todo`, immediately start the next phase. No summary message between phases — keep going. diff --git a/packages/coding-agent/src/prompts/tools/task-summary.md b/packages/coding-agent/src/prompts/tools/task-summary.md index b6a945351..21f21d9b6 100644 --- a/packages/coding-agent/src/prompts/tools/task-summary.md +++ b/packages/coding-agent/src/prompts/tools/task-summary.md @@ -1,28 +1,17 @@ - -
{{successCount}}/{{totalCount}} succeeded{{#if hasCancelledNote}} ({{cancelledCount}} cancelled){{/if}} [{{duration}}]
- -{{#each summaries}} - -{{status}} + {{#if meta}}{{/if}} {{#if truncated}} - + {{preview}} {{else}} - + {{preview}} - + {{/if}} - -{{#unless @last}} ---- -{{/unless}} -{{/each}} - {{#if mergeSummary}} {{mergeSummary}} {{/if}} - + diff --git a/packages/coding-agent/src/prompts/tools/task.md b/packages/coding-agent/src/prompts/tools/task.md index eb2e8cd83..d567568c2 100644 --- a/packages/coding-agent/src/prompts/tools/task.md +++ b/packages/coding-agent/src/prompts/tools/task.md @@ -1,43 +1,37 @@ -Launches subagents to parallelize workflows. +Spawns ONE subagent per call to work in the background, or resumes an existing one. -{{#if asyncEnabled}} -- Results are delivered automatically when complete. -- The tool result lists the assigned task ids (e.g. `AuthLoader`) — those are the live agent ids. +- Spawning is non-blocking: the call returns immediately with the agent id and a job id; the result is delivered automatically when the agent yields. +- Parallelism = multiple `task` calls in one assistant message. Concurrency is bounded at {{MAX_CONCURRENCY}} running subagents per session. +- If genuinely blocked on a result, wait with `job poll`; otherwise keep working. `job cancel` terminates a task and **cannot carry a message** — only for stalled/abandoned work. {{#if ircEnabled}} -- Coordinate with running tasks via `irc` using those ids. `job cancel` terminates a task and **cannot carry a message** — only use it for stalled/abandoned work. -- If genuinely blocked on completion, wait with `job poll`; otherwise keep working. -{{else}} -- If genuinely blocked on completion, wait with `job poll`; otherwise keep working. -- Use `job list` to snapshot manager state; `cancel: [id]` only to actually stop a stuck task. -{{/if}} +- Coordinate with running agents via `irc` using their ids. Agents reach you and their siblings live the same way. {{/if}} -{{#if ircEnabled}} -Subagents have no conversation history, but they can reach you and their siblings live via the `irc` tool. Front-load every fact, file path, and direction they need in {{#if contextEnabled}}`context` or `assignment`{{else}}each `assignment`{{/if}}. -{{else}} -Subagents have no conversation history. Every fact, file path, and direction they need MUST be explicit in {{#if contextEnabled}}`context` or `assignment`{{else}}each `assignment`{{/if}}. -{{/if}} + +- Finished agents stay alive: `idle` first, then `parked` after a TTL — both remain addressable and revivable. +- `resume: ""` revives an idle/parked agent and runs a follow-up assignment in its existing session. **Prefer resuming an agent that already holds the relevant context over spawning fresh**{{#if ircEnabled}} — check `irc` op:"list" for candidates{{/if}}. +- `history://` is the agent's transcript; `agent://` its latest output artifact. + -- `agent`: agent type for all tasks -- `tasks`: tasks to execute in parallel - - `.id`: CamelCase, ≤32 chars - - `.description`: UI label only — subagent never sees it - - `.assignment`: complete self-contained instructions; one-liners and missing acceptance criteria are PROHIBITED -{{#if contextEnabled}}- `context`: shared background prepended to every assignment; session-specific only{{/if}} +- `agent`: agent type to spawn; omit when `resume` is set +- `resume`: existing agent id — continue that agent instead of spawning (cannot combine with `agent` or `isolated`) +- `id`: stable agent id, CamelCase, ≤32 chars; generated when omitted +- `description`: UI label only — subagent never sees it +- `assignment`: complete self-contained instructions; one-liners and missing acceptance criteria are PROHIBITED {{#if customSchemaEnabled}}- `schema`: JTD schema for expected structured output (do not put format rules in assignments){{/if}} -{{#if isolationEnabled}}- `isolated`: run in isolated env; use when tasks edit overlapping files{{/if}} +{{#if isolationEnabled}}- `isolated`: run in isolated env; returns patches. Isolated agents are NOT resumable{{/if}} -- **Maximize batch width.** Spawn the widest parallel set the work decomposes into. NEVER spawn a single-task batch for divisible work, or defer work that could have been concurrent. -- **Subagents do not verify, lint, or format.** Every assignment MUST instruct the subagent to skip all gates, formatters, and project-wide build/test/lint. You run them once at the end across the union of changed files — avoids redundant runs and racing formatter passes. +- **Maximize fan-out.** Issue the widest set of parallel `task` calls the work decomposes into. NEVER serialize work that could run concurrently. +- **Subagents do not verify, lint, or format.** Every assignment MUST instruct the subagent to skip all gates, formatters, and project-wide build/test/lint. You run them once at the end across the union of changed files. - No globs, no "update all", no package-wide scope. Fan out. - NEVER slow down or serialize because tasks might overlap on some files. Agents resolve collisions among themselves in real time. -- Pass large payloads via `local://` URIs, not inline. {{#if contextEnabled}} (other than the context){{/if}} -{{#if contextEnabled}}- Put shared constraints in `context` once; do not duplicate across assignments.{{/if}} +- Subagents have no conversation history. Every fact, file path, and direction they need MUST be explicit in the `assignment`. +- **Shared background**: write it ONCE to a `local://` file (e.g. `local://ctx.md`) and reference that path in each assignment. Pass large payloads via `local://` URIs, not inline. - Prefer agents that investigate **and** edit in one pass; only spin a read-only discovery step when affected files are genuinely unknown. -- **Read-only agents**: Agents tagged READ-ONLY (e.g. `explore`) have no edit/write/command tools. NEVER hand them an assignment that requires changing files or running commands — they cannot do it and the turn is wasted. Use them to investigate and report back; do the edits yourself or delegate to a writing agent (`task`, `oracle`, `designer`). +- **Read-only agents**: Agents tagged READ-ONLY (e.g. `explore`) have no edit/write/command tools. NEVER hand them an assignment that requires changing files or running commands. Use them to investigate and report back; do the edits yourself or delegate to a writing agent (`task`, `oracle`, `designer`). - **No reasoning offload**: NEVER offload reasoning, analysis, design, or decision-making to `quick_task` or `explore` — they run minimal-effort / small models for mechanical lookups and data collection only. Keep judgment and synthesis in your own context; delegate hard thinking to `task`, `plan`, or `oracle`. @@ -51,16 +45,9 @@ Test: can task B run correctly without seeing A's output? If no, sequence A → Sequential when one task produces a contract (types, API, schema, core module) the other consumes. Parallel when tasks touch disjoint files or are independent refactors/tests. {{/if}} +Sequenced follow-ups SHOULD `resume` the agent that produced the prerequisite — it already holds the context. -{{#if contextEnabled}} - -# Goal ← one sentence: what the batch accomplishes -# Constraints ← MUST/NEVER rules and session decisions -# Contract ← exact types/signatures if tasks share an interface - -{{/if}} - # Target ← exact files and symbols; explicit non-goals # Change ← step-by-step add/remove/rename; APIs and patterns diff --git a/packages/coding-agent/src/task/executor.ts b/packages/coding-agent/src/task/executor.ts index a0ebd42c1..dda69eeef 100644 --- a/packages/coding-agent/src/task/executor.ts +++ b/packages/coding-agent/src/task/executor.ts @@ -7,6 +7,7 @@ import path from "node:path"; import type { AgentEvent, AgentIdentity, AgentTelemetryConfig, ThinkingLevel } from "@oh-my-pi/pi-agent-core"; import { recordHandoff, resolveTelemetry } from "@oh-my-pi/pi-agent-core"; +import type { Usage } from "@oh-my-pi/pi-ai"; import { logger, prompt, untilAborted } from "@oh-my-pi/pi-utils"; import type { Rule } from "../capability/rule"; import { ModelRegistry } from "../config/model-registry"; @@ -26,8 +27,9 @@ import type { MCPManager } from "../mcp/manager"; import type { MnemopiSessionState } from "../mnemopi/state"; import subagentSystemPromptTemplate from "../prompts/system/subagent-system-prompt.md" with { type: "text" }; import submitReminderTemplate from "../prompts/system/subagent-yield-reminder.md" with { type: "text" }; +import { AgentLifecycleManager } from "../registry/agent-lifecycle"; import { AgentRegistry } from "../registry/agent-registry"; -import { createAgentSession, discoverAuthStorage } from "../sdk"; +import { type CreateAgentSessionOptions, createAgentSession, discoverAuthStorage } from "../sdk"; import type { AgentSession, AgentSessionEvent } from "../session/agent-session"; import type { ArtifactManager } from "../session/artifacts"; import type { AuthStorage } from "../session/auth-storage"; @@ -63,6 +65,30 @@ import { const MCP_CALL_TIMEOUT_MS = 60_000; +/** + * Soft per-agent request budgets (assistant requests per run). When a subagent + * crosses its budget it receives ONE steering notice asking it to wrap up; at + * 1.5x the budget the run is aborted gracefully so partial output is salvaged. + * The `default` key applies to agents without an explicit entry and can be + * overridden via the `task.softRequestBudget` setting (0 disables the guard). + */ +export const SOFT_REQUEST_BUDGET: Record = { + explore: 40, + quick_task: 40, + default: 90, +}; + +/** Steering notice injected once when a subagent crosses its soft request budget. */ +export function buildBudgetNotice(requests: number): string { + return `[budget notice] You have used ${requests} requests in this run. Wrap up now: finish the current step and yield your final report.`; +} + +/** Flatten whitespace and clip salvage text for the cancelled-child summary line. */ +function formatSalvageSnippet(text: string, maxLength = 500): string { + const flattened = text.replace(/\s+/g, " ").trim(); + return flattened.length > maxLength ? `${flattened.slice(0, maxLength - 1)}…` : flattened; +} + /** Agent event types to forward for progress tracking. */ const agentEventTypes = new Set([ "agent_start", @@ -94,9 +120,13 @@ function normalizeModelPatterns(value: string | string[] | undefined): string[] function renderIrcPeerRoster(selfId: string): string { const peers = AgentRegistry.global() .list() - .filter(ref => ref.id !== selfId && (ref.status === "running" || ref.status === "idle")); - if (peers.length === 0) return "- (no other live agents)"; - return peers.map(peer => `- \`${peer.id}\` — ${peer.displayName} (${peer.kind}, ${peer.status})`).join("\n"); + .filter(ref => ref.id !== selfId && ref.status !== "aborted"); + if (peers.length === 0) return "- (no other agents)"; + const lines = peers.map(peer => `- \`${peer.id}\` — ${peer.displayName} (${peer.kind}, ${peer.status})`); + if (peers.some(peer => peer.status === "idle" || peer.status === "parked")) { + lines.push("Idle/parked peers are not gone: messaging them wakes (or revives) them."); + } + return lines.join("\n"); } function withAbortTimeout(promise: Promise, timeoutMs: number, signal?: AbortSignal): Promise { @@ -152,7 +182,6 @@ export interface ExecutorOptions { agent: AgentDefinition; task: string; assignment?: string; - context?: string; /** * The session's active overall plan, handed off so subagents spawned during * plan execution share the same plan context as the main agent. Omitted when @@ -186,8 +215,6 @@ export interface ExecutorOptions { sessionFile?: string | null; persistArtifacts?: boolean; artifactsDir?: string; - /** Path to parent conversation context file */ - contextFile?: string; eventBus?: EventBus; contextFiles?: ContextFileEntry[]; skills?: Skill[]; @@ -611,28 +638,67 @@ export function createSubagentSettings( }); } +type AbortReason = "signal" | "terminate" | "timeout" | "budget"; + +/** Inputs for the shared run monitor used by both fresh spawns and resumes. */ +interface RunMonitorArgs { + index: number; + id: string; + agent: AgentDefinition; + task: string; + assignment?: string; + description?: string; + modelOverride?: string | string[]; + signal?: AbortSignal; + onProgress?: (progress: AgentProgress) => void; + eventBus?: EventBus; + parentToolCallId?: string; + sessionFile?: string; + /** Soft assistant-request budget; 0 disables the guard. */ + softRequestBudget: number; + /** Wall-clock cap in ms; 0 disables the timer. */ + maxRuntimeMs: number; +} + /** - * Run a single agent in-process. + * The run-monitoring core shared by {@link runSubprocess} and + * {@link resumeSubprocess}: progress tracking, event processing, abort/budget + * machinery, usage accumulation, and output capture for one assignment run. */ -export async function runSubprocess(options: ExecutorOptions): Promise { - const { - cwd, - agent, - task, - assignment, - index, - id, - worktree, - modelOverride, - thinkingLevel, - outputSchema, - enableLsp, - signal, - onProgress, - } = options; +interface SubagentRunMonitor { + readonly progress: AgentProgress; + /** Fires when the run was asked to stop (caller signal, timeout, budget, terminate). */ + readonly abortSignal: AbortSignal; + readonly accumulatedUsage: Usage; + hasUsage(): boolean; + yieldCalled(): boolean; + runtimeLimitExceeded(): boolean; + /** True when the abort carries a precise external reason (signal / wall-clock / budget). */ + hasExplicitAbortReason(): boolean; + /** Whether the (attempted) abort counts as a cancelled run rather than an internal failure. */ + isAbortedRun(): boolean; + requestAbort(reason: AbortReason): void; + resolveSignalAbortReason(): string; + resolveAbortReasonText(): string; + setActiveSession(session: AgentSession | null): void; + /** Return and clear the active session reference. */ + takeActiveSession(): AgentSession | null; + /** Subscribe the monitor to a session's events. Returns the unsubscribe function. */ + attach(session: AgentSession): () => void; + /** Best-effort capture of the last assistant text for cancelled-run salvage. */ + captureSalvage(session: AgentSession): void; + lastAssistantSalvageText(): string | undefined; + /** Final raw output: end-of-run assistant text when available, else accumulated chunks. */ + rawOutput(): string; + scheduleProgress(flush?: boolean): void; + /** Stop processing events and clear listeners/timers. Call once the run settled. */ + finish(): void; +} + +function createSubagentRunMonitor(args: RunMonitorArgs): SubagentRunMonitor { + const { index, id, agent, task, assignment, signal, onProgress, softRequestBudget, maxRuntimeMs } = args; const startTime = Date.now(); - // Initialize progress const progress: AgentProgress = { index, id, @@ -641,109 +707,23 @@ export async function runSubprocess(options: ExecutorOptions): Promise= 0 && childDepth >= maxRecursionDepth; - - // Add tools if specified - let toolNames: string[] | undefined; - if (agent.tools && agent.tools.length > 0) { - toolNames = agent.tools; - // Auto-include task tool if spawns defined but task not in tools - if (agent.spawns !== undefined && !toolNames.includes("task") && !atMaxDepth) { - toolNames = [...toolNames, "task"]; - } - } - - if (atMaxDepth && toolNames?.includes("task")) { - toolNames = toolNames.filter(name => name !== "task"); - } - // IRC is always available; the COOP prompt section advertises it, so a restricted - // whitelist must still carry `irc` for the subagent to actually use it. - if (toolNames && !toolNames.includes("irc")) { - toolNames = [...toolNames, "irc"]; - } - if (toolNames?.includes("exec")) { - const allowEvalPy = settings.get("eval.py") ?? true; - const allowEvalJs = settings.get("eval.js") ?? true; - const expanded = toolNames.filter(name => name !== "exec"); - if (allowEvalPy || allowEvalJs) expanded.push("eval"); - expanded.push("bash"); - toolNames = Array.from(new Set(expanded)); - } - - const modelPatterns = normalizeModelPatterns(modelOverride ?? agent.model); - const sessionFile = subtaskSessionFile ?? null; - const spawnsEnv = atMaxDepth - ? "" - : agent.spawns === undefined - ? "" - : agent.spawns === "*" - ? "*" - : agent.spawns.join(","); - - const lspEnabled = enableLsp ?? true; - const ircEnabled = subagentSettings.get("irc.enabled") === true; - const contextFileForPrompt = ircEnabled ? undefined : options.contextFile; - const skipPythonPreflight = Array.isArray(toolNames) && !toolNames.includes("eval"); - const outputChunks: string[] = []; const finalOutputChunks: string[] = []; const RECENT_OUTPUT_TAIL_BYTES = 8 * 1024; let recentOutputTail = ""; - let stderr = ""; let resolved = false; - type AbortReason = "signal" | "terminate" | "timeout"; let abortSent = false; let abortReason: AbortReason | undefined; let runtimeLimitExceeded = false; @@ -752,11 +732,10 @@ export async function runSubprocess(options: ExecutorOptions): Promise void) | null = null; let yieldCalled = false; // Accumulate usage incrementally from message_end events (no memory for streaming events) - const accumulatedUsage = { + const accumulatedUsage: Usage = { input: 0, output: 0, cacheRead: 0, @@ -765,11 +744,17 @@ export async function runSubprocess(options: ExecutorOptions): Promise { if (reason === "timeout") { runtimeLimitExceeded = true; } + if (reason === "budget") { + budgetLimitExceeded = true; + } if (abortSent) { if (reason === "signal" && abortReason !== "signal" && abortReason !== "timeout") { abortReason = "signal"; @@ -786,11 +771,14 @@ export async function runSubprocess(options: ExecutorOptions): Promise { - if (!resolved) requestAbort("signal"); - }; if (signal) { - signal.addEventListener("abort", onAbort, { once: true, signal: listenerSignal }); + signal.addEventListener( + "abort", + () => { + if (!resolved) requestAbort("signal"); + }, + { once: true, signal: listenerSignal }, + ); } // Wall-clock hard limit. Defense-in-depth for the case where a provider stream @@ -826,6 +814,9 @@ export async function runSubprocess(options: ExecutorOptions): Promise { progress.durationMs = Date.now() - startTime; onProgress?.({ ...progress }); - if (options.eventBus) { - options.eventBus.emit(TASK_SUBAGENT_PROGRESS_CHANNEL, { + if (args.eventBus) { + args.eventBus.emit(TASK_SUBAGENT_PROGRESS_CHANNEL, { index, agent: agent.name, agentSource: agent.source, task, - parentToolCallId: options.parentToolCallId, + parentToolCallId: args.parentToolCallId, assignment, progress: { ...progress }, - sessionFile: subtaskSessionFile, + sessionFile: args.sessionFile, }); } lastProgressEmitMs = Date.now(); @@ -925,8 +916,8 @@ export async function runSubprocess(options: ExecutorOptions): Promise { - if (!options.eventBus) return; - options.eventBus.emit(TASK_SUBAGENT_EVENT_CHANNEL, { + if (!args.eventBus) return; + args.eventBus.emit(TASK_SUBAGENT_EVENT_CHANNEL, { id, event, }); @@ -1078,6 +1069,26 @@ export async function runSubprocess(options: ExecutorOptions): Promise 0 && !abortSent) { + if (progress.requests >= softRequestBudget * 1.5) { + requestAbort("budget"); + } else if (!budgetSteerSent && progress.requests >= softRequestBudget) { + budgetSteerSent = true; + const steerSession = activeSession; + if (steerSession) { + void steerSession + .sendUserMessage(buildBudgetNotice(progress.requests), { deliverAs: "steer" }) + .catch(err => { + logger.warn("Subagent budget steer failed", { + error: err instanceof Error ? err.message : String(err), + }); + }); + } + } + } + } if (role === "assistant") { const messageContent = getMessageContent(event.message) || (event as AgentEvent & { content?: unknown }).content; @@ -1147,6 +1158,543 @@ export async function runSubprocess(options: ExecutorOptions): Promise void) => + session.subscribe(event => { + emitSubagentEvent(event); + if (event.type === "auto_retry_start") { + progress.retryState = { + attempt: event.attempt, + maxAttempts: event.maxAttempts, + delayMs: event.delayMs, + errorMessage: event.errorMessage, + startedAtMs: Date.now(), + }; + progress.retryFailure = undefined; + scheduleProgress(true); + return; + } + if (event.type === "auto_retry_end") { + const attempt = progress.retryState?.attempt ?? event.attempt; + progress.retryState = undefined; + if (!event.success) { + progress.retryFailure = { + attempt, + errorMessage: event.finalError ?? "Auto-retry failed", + }; + } + scheduleProgress(true); + return; + } + if (isAgentEvent(event)) { + try { + processEvent(event); + } catch (err) { + logger.error("Subagent event processing failed", { + error: err instanceof Error ? err.message : String(err), + }); + requestAbort("terminate"); + } + } + }); + + const captureSalvage = (session: AgentSession): void => { + // Best-effort salvage: capture the last assistant text so + // cancelled/aborted children can surface "last activity" instead of + // "(no output)". + try { + const lastContent = session.getLastAssistantMessage()?.content; + if (Array.isArray(lastContent)) { + const text = lastContent + .map(block => (block.type === "text" && typeof block.text === "string" ? block.text : "")) + .filter(Boolean) + .join("\n"); + if (text.trim()) { + lastAssistantSalvageText = text; + } + } + } catch { + // Salvage is best-effort; partial sessions may not implement it + } + }; + + return { + progress, + abortSignal, + accumulatedUsage, + hasUsage: () => hasUsage, + yieldCalled: () => yieldCalled, + runtimeLimitExceeded: () => runtimeLimitExceeded, + hasExplicitAbortReason: () => abortReason === "signal" || runtimeLimitExceeded || budgetLimitExceeded, + isAbortedRun: () => + abortReason === "signal" || runtimeLimitExceeded || budgetLimitExceeded || abortReason === undefined, + requestAbort, + resolveSignalAbortReason, + resolveAbortReasonText, + setActiveSession: session => { + activeSession = session; + }, + takeActiveSession: () => { + const session = activeSession; + activeSession = null; + return session; + }, + attach, + captureSalvage, + lastAssistantSalvageText: () => lastAssistantSalvageText, + rawOutput: () => (finalOutputChunks.length > 0 ? finalOutputChunks.join("") : outputChunks.join("")), + scheduleProgress, + finish: () => { + resolved = true; + listenerController.abort(); + if (runtimeTimeoutId !== undefined) { + clearTimeout(runtimeTimeoutId); + runtimeTimeoutId = undefined; + } + if (progressTimeoutId) { + clearTimeout(progressTimeoutId); + progressTimeoutId = null; + } + }, + }; +} + +interface DriveOutcome { + exitCode: number; + error?: string; + aborted: boolean; + abortReasonText?: string; +} + +const MAX_YIELD_RETRIES = 3; + +/** + * Drive one assignment through a live session: send the prompt, wait for idle, + * remind the agent to `yield` (up to {@link MAX_YIELD_RETRIES} times), then + * classify the terminal assistant state. Shared by spawn and resume paths. + */ +async function driveSessionToYield( + session: AgentSession, + monitor: SubagentRunMonitor, + task: string, +): Promise { + const abortSignal = monitor.abortSignal; + let exitCode = 0; + let error: string | undefined; + let aborted = false; + let abortReasonText: string | undefined; + const checkAbort = () => { + if (abortSignal.aborted) { + aborted = monitor.isAbortedRun(); + if (aborted) { + abortReasonText ??= monitor.resolveAbortReasonText(); + } + exitCode = 1; + throw new ToolAbortError(); + } + }; + const awaitAbortable = async (promise: Promise): Promise => { + checkAbort(); + const { promise: abortPromise, reject } = Promise.withResolvers(); + const onAbort = () => { + try { + checkAbort(); + } catch (err) { + reject(err); + } + }; + abortSignal.addEventListener("abort", onAbort, { once: true }); + try { + return await Promise.race([promise, abortPromise]); + } finally { + abortSignal.removeEventListener("abort", onAbort); + } + }; + + try { + await awaitAbortable(session.prompt(task, { attribution: "agent" })); + await awaitAbortable(session.waitForIdle()); + + const reminderToolChoice = buildNamedToolChoice("yield", session.model); + + let retryCount = 0; + while (!monitor.yieldCalled() && retryCount < MAX_YIELD_RETRIES && !abortSignal.aborted) { + // Skip reminders when the model returned a terminal error (e.g. + // rate-limit cap hit, auth failure). Re-prompting would just + // hit the same wall, multiplying the failure noise without + // any chance of producing a yield. + const lastBeforeReminder = session.getLastAssistantMessage(); + if (lastBeforeReminder?.stopReason === "error") break; + try { + retryCount++; + const reminder = prompt.render(submitReminderTemplate, { + retryCount, + maxRetries: MAX_YIELD_RETRIES, + }); + + const isFinalRetry = retryCount >= MAX_YIELD_RETRIES; + await awaitAbortable( + session.prompt(reminder, { + attribution: "agent", + synthetic: true, + ...(isFinalRetry && reminderToolChoice ? { toolChoice: reminderToolChoice } : {}), + }), + ); + await awaitAbortable(session.waitForIdle()); + } catch (err) { + if (abortSignal.aborted || err instanceof ToolAbortError) { + // Benign control-flow exit — user cancel (^C) or compaction aborting + // pending operations both surface here as ToolAbortError. The outer + // catch and finally already mark the run aborted; logging at ERROR + // would spam operator dashboards with non-failures. + logger.debug("Subagent prompt aborted"); + } else { + logger.error("Subagent prompt failed", { + error: err instanceof Error ? err.message : String(err), + }); + } + } + } + + await awaitAbortable(session.waitForIdle()); + + const lastAssistant = session.getLastAssistantMessage(); + if (lastAssistant) { + if (lastAssistant.stopReason === "aborted") { + aborted = monitor.isAbortedRun(); + if (aborted) { + // A real caller signal or the wall-clock timer carries a precise + // reason (signal.reason / "runtime limit exceeded"). An internal + // turn abort does NOT — prefer the assistant message's own + // errorMessage ("Request was aborted" or a specific stream error) + // over the misleading "Cancelled by caller". + abortReasonText ??= monitor.hasExplicitAbortReason() + ? monitor.resolveAbortReasonText() + : lastAssistant.errorMessage?.trim() || monitor.resolveAbortReasonText(); + } + exitCode = 1; + } else if (lastAssistant.stopReason === "error") { + exitCode = 1; + error ??= lastAssistant.errorMessage || "Subagent failed"; + } + } + } catch (err) { + exitCode = 1; + if (!abortSignal.aborted) { + error = err instanceof Error ? err.stack || err.message : String(err); + } + } finally { + if (abortSignal.aborted) { + aborted = monitor.isAbortedRun(); + if (aborted) { + abortReasonText ??= monitor.resolveAbortReasonText(); + } + if (exitCode === 0) exitCode = 1; + } + } + + return { exitCode, error, aborted, abortReasonText }; +} + +interface FinalizeRunArgs { + monitor: SubagentRunMonitor; + done: { exitCode: number; error?: string; aborted?: boolean; abortReason?: string; durationMs: number }; + index: number; + id: string; + agent: AgentDefinition; + task: string; + assignment?: string; + description?: string; + modelOverride?: string | string[]; + outputSchema?: unknown; + signal?: AbortSignal; + artifactsDir?: string; + eventBus?: EventBus; + parentToolCallId?: string; + sessionFile?: string; + startTime: number; +} + +/** + * Turn a settled run into a {@link SingleResult}: resolve the yield payload via + * {@link finalizeSubprocessOutput}, salvage cancelled-run output, write the + * `.md` output artifact, flush final progress, and emit the lifecycle end + * event. Shared by spawn and resume paths. + */ +async function finalizeRunResult(args: FinalizeRunArgs): Promise { + const { monitor, done, index, id, agent, task, assignment, signal, modelOverride } = args; + const progress = monitor.progress; + let exitCode = done.exitCode; + let stderr = done.error ?? ""; + + // Use final output if available, otherwise accumulated output + let rawOutput = monitor.rawOutput(); + const yieldItems = progress.extractedToolData?.yield as YieldItem[] | undefined; + const reportFindingDetails = progress.extractedToolData?.report_finding as ReportFindingDetails[] | undefined; + const reportFindings: ReviewFinding[] | undefined = reportFindingDetails?.map(toReviewFinding); + const finalized = finalizeSubprocessOutput({ + rawOutput, + exitCode, + stderr, + doneAborted: Boolean(done.aborted), + signalAborted: Boolean(signal?.aborted), + yieldItems, + reportFindings, + outputSchema: args.outputSchema, + }); + rawOutput = finalized.rawOutput; + exitCode = finalized.exitCode; + stderr = finalized.stderr; + // Salvage for cancelled/aborted children that produced no completed output: + // surface the last assistant text + stats instead of "(no output)" so the + // parent doesn't redo work the child already finished. + const salvageText = monitor.lastAssistantSalvageText(); + if ( + (done.aborted || signal?.aborted || monitor.runtimeLimitExceeded()) && + !rawOutput.trim() && + salvageText !== undefined + ) { + rawOutput = `[cancelled after ${progress.requests} req, ${progress.tokens} tok — last activity: "${formatSalvageSnippet(salvageText)}"]`; + } + const lastYield = yieldItems?.[yieldItems.length - 1]; + const yieldAbortReason = lastYield?.status === "aborted" ? lastYield.error || "Subagent aborted task" : undefined; + const { abortedViaYield, hasYield } = finalized; + const { content: truncatedOutput, truncated } = truncateTail(rawOutput, { + maxBytes: MAX_OUTPUT_BYTES, + maxLines: MAX_OUTPUT_LINES, + }); + + // Write output artifact (input and jsonl already written in real-time) + // Compute output metadata for agent:// URL integration + let outputMeta: { lineCount: number; charCount: number } | undefined; + let outputPath: string | undefined; + if (args.artifactsDir) { + outputPath = path.join(args.artifactsDir, `${id}.md`); + try { + await Bun.write(outputPath, rawOutput); + outputMeta = { + lineCount: rawOutput.split("\n").length, + charCount: rawOutput.length, + }; + } catch { + // Non-fatal + } + } + + // Update final progress. A wall-clock timeout always wins: if the runtime + // limit fired we report aborted/failed regardless of whether a yield landed + // while we were tearing the session down. The yield data is still surfaced + // to the caller via `progress.extractedToolData`, but the exit status must + // reflect the timeout so on-call doesn't mistake a stuck run for success. + const runtimeLimitExceeded = monitor.runtimeLimitExceeded(); + if (runtimeLimitExceeded && exitCode === 0) { + exitCode = 1; + } + const wasAborted = + runtimeLimitExceeded || abortedViaYield || (!hasYield && (done.aborted || signal?.aborted || false)); + const finalAbortReason = wasAborted + ? runtimeLimitExceeded + ? monitor.resolveAbortReasonText() + : abortedViaYield + ? yieldAbortReason + : (done.abortReason ?? + (signal?.aborted ? monitor.resolveSignalAbortReason() : monitor.resolveAbortReasonText())) + : undefined; + progress.status = wasAborted ? "aborted" : exitCode === 0 ? "completed" : "failed"; + monitor.scheduleProgress(true); + + // Emit lifecycle end event after finalization so yield status is reflected + if (args.eventBus) { + args.eventBus.emit(TASK_SUBAGENT_LIFECYCLE_CHANNEL, { + id, + agent: agent.name, + parentToolCallId: args.parentToolCallId, + agentSource: agent.source, + description: args.description, + status: progress.status as "completed" | "failed" | "aborted", + sessionFile: args.sessionFile, + index, + }); + } + + return { + index, + id, + agent: agent.name, + agentSource: agent.source, + task, + assignment, + description: args.description, + lastIntent: progress.lastIntent, + exitCode, + output: truncatedOutput, + stderr, + truncated: Boolean(truncated), + durationMs: Date.now() - args.startTime, + tokens: progress.tokens, + requests: progress.requests, + contextTokens: progress.contextTokens, + contextWindow: progress.contextWindow, + modelOverride, + resolvedModel: progress.resolvedModel, + error: exitCode !== 0 && stderr ? stderr : undefined, + aborted: wasAborted, + abortReason: finalAbortReason, + usage: monitor.hasUsage() ? monitor.accumulatedUsage : undefined, + outputPath, + extractedToolData: progress.extractedToolData, + retryFailure: progress.retryFailure, + outputMeta, + }; +} + +/** + * Run a single agent in-process. + */ +export async function runSubprocess(options: ExecutorOptions): Promise { + const { + cwd, + agent, + task, + assignment, + index, + id, + worktree, + modelOverride, + thinkingLevel, + outputSchema, + enableLsp, + signal, + onProgress, + } = options; + const startTime = Date.now(); + + // Check if already aborted + if (signal?.aborted) { + return { + index, + id, + agent: agent.name, + agentSource: agent.source, + task, + assignment, + description: options.description, + exitCode: 1, + output: "", + stderr: "Cancelled before start", + truncated: false, + durationMs: 0, + tokens: 0, + requests: 0, + modelOverride, + error: "Cancelled before start", + aborted: true, + abortReason: "Cancelled before start", + }; + } + + // Set up artifact paths and write input file upfront if artifacts dir provided + let subtaskSessionFile: string | undefined; + if (options.artifactsDir) { + subtaskSessionFile = path.join(options.artifactsDir, `${id}.jsonl`); + } + + const settings = options.settings ?? Settings.isolated(); + const subagentSettings = createSubagentSettings( + settings, + agent.readSummarize === false ? { "read.summarize.enabled": false } : undefined, + ); + const maxRecursionDepth = settings.get("task.maxRecursionDepth") ?? 2; + const maxRuntimeMs = Math.max( + 0, + Math.trunc(Number(options.maxRuntimeMs ?? settings.get("task.maxRuntimeMs") ?? 0) || 0), + ); + // TTL before an adopted idle subagent is parked by the lifecycle manager. + // <= 0 disables parking (the session stays live until process teardown). + const agentIdleTtlMs = Math.trunc(Number(settings.get("task.agentIdleTtlMs") ?? 420_000) || 0); + const configuredDefaultBudget = Math.max( + 0, + Math.trunc(Number(settings.get("task.softRequestBudget") ?? SOFT_REQUEST_BUDGET.default) || 0), + ); + const softRequestBudget = + configuredDefaultBudget === 0 ? 0 : (SOFT_REQUEST_BUDGET[agent.name] ?? configuredDefaultBudget); + const parentDepth = options.taskDepth ?? 0; + const childDepth = parentDepth + 1; + const atMaxDepth = maxRecursionDepth >= 0 && childDepth >= maxRecursionDepth; + + // Add tools if specified + let toolNames: string[] | undefined; + if (agent.tools && agent.tools.length > 0) { + toolNames = agent.tools; + // Auto-include task tool if spawns defined but task not in tools + if (agent.spawns !== undefined && !toolNames.includes("task") && !atMaxDepth) { + toolNames = [...toolNames, "task"]; + } + } + + if (atMaxDepth && toolNames?.includes("task")) { + toolNames = toolNames.filter(name => name !== "task"); + } + // IRC is always available; the COOP prompt section advertises it, so a restricted + // whitelist must still carry `irc` for the subagent to actually use it. + if (toolNames && !toolNames.includes("irc")) { + toolNames = [...toolNames, "irc"]; + } + if (toolNames?.includes("exec")) { + const allowEvalPy = settings.get("eval.py") ?? true; + const allowEvalJs = settings.get("eval.js") ?? true; + const expanded = toolNames.filter(name => name !== "exec"); + if (allowEvalPy || allowEvalJs) expanded.push("eval"); + expanded.push("bash"); + toolNames = Array.from(new Set(expanded)); + } + + const modelPatterns = normalizeModelPatterns(modelOverride ?? agent.model); + const sessionFile = subtaskSessionFile ?? null; + const spawnsEnv = atMaxDepth + ? "" + : agent.spawns === undefined + ? "" + : agent.spawns === "*" + ? "*" + : agent.spawns.join(","); + + const lspEnabled = enableLsp ?? true; + const ircEnabled = subagentSettings.get("irc.enabled") === true; + const skipPythonPreflight = Array.isArray(toolNames) && !toolNames.includes("eval"); + + const monitor = createSubagentRunMonitor({ + index, + id, + agent, + task, + assignment, + description: options.description, + modelOverride, + signal, + onProgress, + eventBus: options.eventBus, + parentToolCallId: options.parentToolCallId, + sessionFile: subtaskSessionFile, + softRequestBudget, + maxRuntimeMs, + }); + const progress = monitor.progress; + let unsubscribe: (() => void) | null = null; + let reviveSession: (() => Promise) | null = null; + // Adopted (kept-alive) subagents flip registry status from session events on + // later turns: revive/wake → running, turn drained → idle. The subscription + // intentionally survives this run; a disposed session emits nothing, so it + // needs no teardown. + const installRegistryStatusSync = (target: AgentSession): void => { + target.subscribe(event => { + if (event.type === "agent_start") { + AgentRegistry.global().setStatus(id, "running"); + } else if (event.type === "agent_end") { + AgentRegistry.global().setStatus(id, "idle"); + } + }); + }; + const runSubagent = async (): Promise<{ exitCode: number; error?: string; @@ -1155,17 +1703,13 @@ export async function runSubprocess(options: ExecutorOptions): Promise => { const sessionAbortController = new AbortController(); + const abortSignal = monitor.abortSignal; let exitCode = 0; let error: string | undefined; let aborted = false; let abortReasonText: string | undefined; const checkAbort = () => { if (abortSignal.aborted) { - aborted = abortReason === "signal" || runtimeLimitExceeded || abortReason === undefined; - if (aborted) { - abortReasonText ??= resolveAbortReasonText(); - } - exitCode = 1; throw new ToolAbortError(); } }; @@ -1283,7 +1827,11 @@ export async function runSubprocess(options: ExecutorOptions): Promise ({ cwd: worktree ?? cwd, authStorage, modelRegistry, @@ -1303,12 +1851,10 @@ export async function runSubprocess(options: ExecutorOptions): Promise { const subagentPrompt = prompt.render(subagentSystemPromptTemplate, { agent: agent.systemPrompt, - context: options.context?.trim() ?? "", planReference: options.planReference?.content ?? "", planReferencePath: options.planReference?.path ?? "", worktree: worktree ?? "", outputSchema: normalizedOutputSchema, - contextFile: contextFileForPrompt, ircPeers: ircEnabled ? renderIrcPeerRoster(id) : "", ircSelfId: ircEnabled ? id : "", }); @@ -1316,7 +1862,7 @@ export async function runSubprocess(options: ExecutorOptions): Promise { + const reopened = await SessionManager.open(sessionFile); + if (options.parentArtifactManager) { + reopened.adoptArtifactManager(options.parentArtifactManager); + } + const { session: revived } = await createAgentSession(buildSubagentSessionOptions(reopened)); + installRegistryStatusSync(revived); + return revived; + }; + } // Emit lifecycle start event if (options.eventBus) { @@ -1449,44 +2013,7 @@ export async function runSubprocess(options: ExecutorOptions): Promise { - emitSubagentEvent(event); - if (event.type === "auto_retry_start") { - progress.retryState = { - attempt: event.attempt, - maxAttempts: event.maxAttempts, - delayMs: event.delayMs, - errorMessage: event.errorMessage, - startedAtMs: Date.now(), - }; - progress.retryFailure = undefined; - scheduleProgress(true); - return; - } - if (event.type === "auto_retry_end") { - const attempt = progress.retryState?.attempt ?? event.attempt; - progress.retryState = undefined; - if (!event.success) { - progress.retryFailure = { - attempt, - errorMessage: event.finalError ?? "Auto-retry failed", - }; - } - scheduleProgress(true); - return; - } - if (isAgentEvent(event)) { - try { - processEvent(event); - } catch (err) { - logger.error("Subagent event processing failed", { - error: err instanceof Error ? err.message : String(err), - }); - requestAbort("terminate"); - } - } - }); + unsubscribe = monitor.attach(session); checkAbort(); // Autoload skills via sendCustomMessage (same mechanic as /skill:) @@ -1504,78 +2031,12 @@ export async function runSubprocess(options: ExecutorOptions): Promise= MAX_YIELD_RETRIES; - await awaitAbortable( - session.prompt(reminder, { - attribution: "agent", - synthetic: true, - ...(isFinalRetry && reminderToolChoice ? { toolChoice: reminderToolChoice } : {}), - }), - ); - await awaitAbortable(session.waitForIdle()); - } catch (err) { - if (abortSignal.aborted || err instanceof ToolAbortError) { - // Benign control-flow exit — user cancel (^C) or compaction aborting - // pending operations both surface here as ToolAbortError. The outer - // catch and finally already mark the run aborted; logging at ERROR - // would spam operator dashboards with non-failures. - logger.debug("Subagent prompt aborted", { - reason: abortReason ?? "signal", - }); - } else { - logger.error("Subagent prompt failed", { - error: err instanceof Error ? err.message : String(err), - }); - } - } - } - - await awaitAbortable(session.waitForIdle()); - if (!yieldCalled && !abortSignal.aborted) { - exitCode = 0; - } - - const lastAssistant = session.getLastAssistantMessage(); - if (lastAssistant) { - if (lastAssistant.stopReason === "aborted") { - aborted = abortReason === "signal" || runtimeLimitExceeded || abortReason === undefined; - if (aborted) { - // A real caller signal or the wall-clock timer carries a precise - // reason (signal.reason / "runtime limit exceeded"). An internal - // turn abort (abortReason === undefined) does NOT — prefer the - // assistant message's own errorMessage ("Request was aborted" or a - // specific stream error) over the misleading "Cancelled by caller". - abortReasonText ??= - abortReason === "signal" || runtimeLimitExceeded - ? resolveAbortReasonText() - : lastAssistant.errorMessage?.trim() || resolveAbortReasonText(); - } - exitCode = 1; - } else if (lastAssistant.stopReason === "error") { - exitCode = 1; - error ??= lastAssistant.errorMessage || "Subagent failed"; - } - } + const outcome = await driveSessionToYield(session, monitor, task); + exitCode = outcome.exitCode; + error = outcome.error; + aborted = outcome.aborted; + abortReasonText = outcome.abortReasonText; } catch (err) { exitCode = 1; if (!abortSignal.aborted) { @@ -1583,9 +2044,9 @@ export async function runSubprocess(options: ExecutorOptions): Promise session.dispose()); - } catch { - // Ignore cleanup errors + const session = monitor.takeActiveSession(); + if (session) { + monitor.captureSalvage(session); + const registry = AgentRegistry.global(); + if (aborted) { + // Hard abort (caller signal / wall-clock / budget): terminal teardown. + registry.setStatus(id, "aborted"); + try { + await untilAborted(AbortSignal.timeout(5000), () => session.dispose()); + } catch { + // Ignore cleanup errors + } + } else if (worktree !== undefined) { + // Isolated run: the worktree is merged + cleaned after the run, so + // the session is not resumable. Park the ref WITHOUT adopting — the + // transcript stays reachable (history://), but ensureLive will throw. + // Status must flip to "parked" before dispose so the sdk dispose + // wrapper skips unregister. + registry.setStatus(id, "parked"); + try { + await untilAborted(AbortSignal.timeout(5000), () => session.dispose()); + } catch { + // Ignore cleanup errors + } + registry.detachSession(id); + } else { + // Keep-alive: finished and failed subagents both stay interrogable. + // The lifecycle manager owns idle-TTL parking + revival from here on. + registry.setStatus(id, "idle"); + AgentLifecycleManager.global().adopt(id, { + idleTtlMs: agentIdleTtlMs, + revive: reviveSession ?? undefined, + }); } } } @@ -1619,87 +2106,115 @@ export async function runSubprocess(options: ExecutorOptions): Promise 0 ? finalOutputChunks.join("") : outputChunks.join(""); - const yieldItems = progress.extractedToolData?.yield as YieldItem[] | undefined; - const reportFindingDetails = progress.extractedToolData?.report_finding as ReportFindingDetails[] | undefined; - const reportFindings: ReviewFinding[] | undefined = reportFindingDetails?.map(toReviewFinding); - const finalized = finalizeSubprocessOutput({ - rawOutput, - exitCode, - stderr, - doneAborted: Boolean(done.aborted), - signalAborted: Boolean(signal?.aborted), - yieldItems, - reportFindings, + return finalizeRunResult({ + monitor, + done, + index, + id, + agent, + task, + assignment, + description: options.description, + modelOverride, outputSchema, + signal, + artifactsDir: options.artifactsDir, + eventBus: options.eventBus, + parentToolCallId: options.parentToolCallId, + sessionFile: subtaskSessionFile, + startTime, }); - rawOutput = finalized.rawOutput; - exitCode = finalized.exitCode; - stderr = finalized.stderr; - const lastYield = yieldItems?.[yieldItems.length - 1]; - const yieldAbortReason = lastYield?.status === "aborted" ? lastYield.error || "Subagent aborted task" : undefined; - const { abortedViaYield, hasYield } = finalized; - const { content: truncatedOutput, truncated } = truncateTail(rawOutput, { - maxBytes: MAX_OUTPUT_BYTES, - maxLines: MAX_OUTPUT_LINES, - }); +} - // Write output artifact (input and jsonl already written in real-time) - // Compute output metadata for agent:// URL integration - let outputMeta: { lineCount: number; charCount: number } | undefined; - let outputPath: string | undefined; - if (options.artifactsDir) { - outputPath = path.join(options.artifactsDir, `${id}.md`); - try { - await Bun.write(outputPath, rawOutput); - outputMeta = { - lineCount: rawOutput.split("\n").length, - charCount: rawOutput.length, - }; - } catch { - // Non-fatal - } +/** Options for resuming an existing live subagent session with a follow-up assignment. */ +export interface ResumeExecutorOptions { + /** Live session, e.g. from `AgentLifecycleManager.global().ensureLive(id)`. */ + session: AgentSession; + /** Registry agent id being resumed. */ + id: string; + /** Agent definition for progress labels and soft budgets; a minimal stub is acceptable. */ + agent: AgentDefinition; + /** Rendered follow-up prompt, injected via the session's normal prompt path. */ + task: string; + assignment?: string; + description?: string; + index: number; + parentToolCallId?: string; + /** Optional schema validating this follow-up's yield payload. */ + outputSchema?: unknown; + signal?: AbortSignal; + onProgress?: (progress: AgentProgress) => void; + eventBus?: EventBus; + settings?: Settings; + /** Where the `.md` output artifact is (over)written for this assignment. */ + artifactsDir?: string; +} + +/** + * Run a follow-up assignment on an EXISTING live agent session through the same + * monitoring/finalize pipeline as a fresh spawn. The session is never created + * or disposed here: it stays alive (and adopted by the lifecycle manager from + * its original spawn) afterwards — registry status flips via the session's + * registry status sync, and the idle TTL re-arms via the lifecycle manager's + * registry subscription. Each resume overwrites the `agent://` output + * artifact; the transcript accretes in the session JSONL. + */ +export async function resumeSubprocess(options: ResumeExecutorOptions): Promise { + const { session, id, agent, task, assignment, index, signal } = options; + const startTime = Date.now(); + + if (signal?.aborted) { + return { + index, + id, + agent: agent.name, + agentSource: agent.source, + task, + assignment, + description: options.description, + exitCode: 1, + output: "", + stderr: "Cancelled before start", + truncated: false, + durationMs: 0, + tokens: 0, + requests: 0, + error: "Cancelled before start", + aborted: true, + abortReason: "Cancelled before start", + }; } - // Update final progress. A wall-clock timeout always wins: if the runtime - // limit fired we report aborted/failed regardless of whether a yield landed - // while we were tearing the session down. The yield data is still surfaced - // to the caller via `progress.extractedToolData`, but the exit status must - // reflect the timeout so on-call doesn't mistake a stuck run for success. - if (runtimeLimitExceeded && exitCode === 0) { - exitCode = 1; - } - const wasAborted = - runtimeLimitExceeded || abortedViaYield || (!hasYield && (done.aborted || signal?.aborted || false)); - const finalAbortReason = wasAborted - ? runtimeLimitExceeded - ? resolveAbortReasonText() - : abortedViaYield - ? yieldAbortReason - : (done.abortReason ?? (signal?.aborted ? resolveSignalAbortReason() : resolveAbortReasonText())) - : undefined; - progress.status = wasAborted ? "aborted" : exitCode === 0 ? "completed" : "failed"; - scheduleProgress(true); + const settings = options.settings ?? Settings.isolated(); + const maxRuntimeMs = Math.max(0, Math.trunc(Number(settings.get("task.maxRuntimeMs") ?? 0) || 0)); + const configuredDefaultBudget = Math.max( + 0, + Math.trunc(Number(settings.get("task.softRequestBudget") ?? SOFT_REQUEST_BUDGET.default) || 0), + ); + const softRequestBudget = + configuredDefaultBudget === 0 ? 0 : (SOFT_REQUEST_BUDGET[agent.name] ?? configuredDefaultBudget); + const sessionFile = AgentRegistry.global().get(id)?.sessionFile ?? undefined; + + const monitor = createSubagentRunMonitor({ + index, + id, + agent, + task, + assignment, + description: options.description, + signal, + onProgress: options.onProgress, + eventBus: options.eventBus, + parentToolCallId: options.parentToolCallId, + sessionFile, + softRequestBudget, + maxRuntimeMs, + }); + monitor.setActiveSession(session); + const unsubscribe = monitor.attach(session); - // Emit lifecycle end event after finalization so yield status is reflected if (options.eventBus) { options.eventBus.emit(TASK_SUBAGENT_LIFECYCLE_CHANNEL, { id, @@ -1707,38 +2222,56 @@ export async function runSubprocess(options: ExecutorOptions): Promise): void { - const input = usage.input ?? 0; - const output = usage.output ?? 0; - const cacheRead = usage.cacheRead ?? 0; - const cacheWrite = usage.cacheWrite ?? 0; - const totalTokens = usage.totalTokens ?? input + output + cacheRead + cacheWrite; - const cost = - usage.cost ?? - ({ - input: 0, - output: 0, - cacheRead: 0, - cacheWrite: 0, - total: 0, - } satisfies Usage["cost"]); - - target.input += input; - target.output += output; - target.cacheRead += cacheRead; - target.cacheWrite += cacheWrite; - target.totalTokens += totalTokens; - target.cost.input += cost.input; - target.cost.output += cost.output; - target.cost.cacheRead += cost.cacheRead; - target.cost.cacheWrite += cost.cacheWrite; - target.cost.total += cost.total; -} // Re-export types and utilities export { loadBundledAgents as BUNDLED_AGENTS } from "./agents"; @@ -165,6 +132,17 @@ export function isReadOnlyAgent(agent: AgentDefinition): boolean { return !!agent.tools?.length && agent.tools.every(tool => READ_ONLY_TOOL_NAMES.has(tool)); } +/** + * Preview text for a child result. Falls back to "(no output)" — annotated + * with the request count when the child actually did work, so the parent can + * tell a no-op child from one that burned requests before being cancelled. + */ +export function formatResultOutputFallback(result: Pick): string { + const base = result.output.trim() || result.stderr.trim(); + if (base) return base; + return result.requests > 0 ? `(no output) after ${result.requests} req` : "(no output)"; +} + /** * Render the tool description from a cached agent list and current settings. */ @@ -172,7 +150,6 @@ function renderDescription( agents: AgentDefinition[], maxConcurrency: number, isolationEnabled: boolean, - asyncEnabled: boolean, disabledAgents: string[], simpleMode: TaskSimpleMode, ircEnabled: boolean, @@ -196,14 +173,12 @@ function renderDescription( description: agent.description, readOnly: isReadOnlyAgent(agent), })); - const { contextEnabled, customSchemaEnabled } = getTaskSimpleModeCapabilities(simpleMode); + const { customSchemaEnabled } = getTaskSimpleModeCapabilities(simpleMode); return prompt.render(taskDescriptionTemplate, { agents: renderedAgents, spawningDisabled, MAX_CONCURRENCY: maxConcurrency, isolationEnabled, - asyncEnabled, - contextEnabled, customSchemaEnabled, ircEnabled, defaultMode: simpleMode === "default", @@ -220,87 +195,46 @@ function createTaskModeError(text: string): AgentToolResult { } function validateTaskModeParams(simpleMode: TaskSimpleMode, params: TaskParams): string | undefined { - const { contextEnabled, customSchemaEnabled } = getTaskSimpleModeCapabilities(simpleMode); - const disallowedFields: string[] = []; - if (!contextEnabled && params.context !== undefined) { - disallowedFields.push("context"); - } - if (!customSchemaEnabled && params.schema !== undefined) { - disallowedFields.push("schema"); - } - if (disallowedFields.length === 0) { + const { customSchemaEnabled } = getTaskSimpleModeCapabilities(simpleMode); + if (customSchemaEnabled || params.schema === undefined) { return undefined; } - - if (simpleMode === "schema-free") { - return "task.simple is set to schema-free, so the task tool does not accept `schema`. Remove it and rely on the selected agent definition or inherited session schema."; - } - - if (disallowedFields.length === 1) { - return `task.simple is set to independent, so the task tool does not accept \`${disallowedFields[0]}\`. Put everything the subagent needs inside each task assignment.`; - } - - return "task.simple is set to independent, so the task tool does not accept `context` or `schema`. Put all required background and output expectations inside each task assignment or the selected agent definition."; + return `task.simple is set to ${simpleMode}, so the task tool does not accept \`schema\`. Remove it and rely on the selected agent definition or inherited session schema.`; } -/** Sentinel for async jobs whose subagent finished with a failing result; batch counters are already updated. */ -class TaskJobError extends Error {} - /** - * Validate task ids: every task needs a non-empty id and ids must be unique - * (case-insensitive). Returns a problem description, or undefined when valid. + * Validate the spawn/resume parameter contract: `agent` XOR `resume`, + * `resume` excludes `isolated`, and `assignment` is always required. + * Returns a problem description, or undefined when valid. */ -function validateTaskIds(tasks: TaskParams["tasks"]): string | undefined { - const missingTaskIndexes: number[] = []; - const idIndexes = new Map(); - - for (let i = 0; i < tasks.length; i++) { - const id = tasks[i]?.id; - if (typeof id !== "string" || id.trim() === "") { - missingTaskIndexes.push(i); - continue; - } - const normalizedId = id.toLowerCase(); - const indexes = idIndexes.get(normalizedId); - if (indexes) { - indexes.push(i); - } else { - idIndexes.set(normalizedId, [i]); - } +function validateSpawnParams(params: TaskParams): string | undefined { + const resume = typeof params.resume === "string" ? params.resume.trim() : ""; + const agent = typeof params.agent === "string" ? params.agent.trim() : ""; + if (resume && agent) { + return "Provide either `agent` (spawn a new subagent) or `resume` (continue an existing one), not both."; } - - const duplicateIds: Array<{ id: string; indexes: number[] }> = []; - for (const [normalizedId, indexes] of idIndexes.entries()) { - if (indexes.length > 1) { - duplicateIds.push({ - id: tasks[indexes[0]]?.id ?? normalizedId, - indexes, - }); - } + if (!resume && !agent) { + return "Missing `agent`. Provide `agent` to spawn a subagent, or `resume` with an existing agent id."; } - - if (missingTaskIndexes.length === 0 && duplicateIds.length === 0) { - return undefined; + if (resume && params.isolated === true) { + return "`resume` cannot be combined with `isolated` — isolated agents are not resumable."; } - - const problems: string[] = []; - if (missingTaskIndexes.length > 0) { - problems.push(`Missing task ids at indexes: ${missingTaskIndexes.join(", ")}`); + if (typeof params.assignment !== "string" || params.assignment.trim() === "") { + return "Missing `assignment`. Provide complete, self-contained instructions for the agent."; } - if (duplicateIds.length > 0) { - const details = duplicateIds.map(entry => `${entry.id} (indexes ${entry.indexes.join(", ")})`).join("; "); - problems.push(`Duplicate task ids detected (case-insensitive): ${details}`); - } - return `Invalid tasks: ${problems.join(". ")}`; + return undefined; } +/** Sentinel for async jobs whose subagent finished with a failing result; progress is already updated. */ +class TaskJobError extends Error {} + /** * Process-level memo for create-time agent discovery, keyed by resolved cwd. * * `TaskTool.create` runs for every (sub)agent session in this process and the * walk-up + plugin-registry scan in `discoverAgents` is identical for a given * cwd, so repeat creations reuse the first scan. Execution-time discovery - * (`#executeSync`) intentionally stays fresh. The memo also tracks the live + * (`#runSpawn`) intentionally stays fresh. The memo also tracks the live * `discoverAgents` binding: test spies swap that binding, which invalidates * the memo automatically. */ @@ -332,8 +266,9 @@ function discoverAgentsForCreate(cwd: string): Promise { /** * Task tool - Delegate tasks to specialized agents. * - * Requires async initialization to discover available agents. - * Use `TaskTool.create(session)` to instantiate. + * Each call spawns ONE subagent (or resumes an existing one). Spawning is + * non-blocking: the call registers an AsyncJobManager job and returns + * immediately; the result is delivered when the agent yields. */ export class TaskTool implements AgentTool { readonly name = "task"; @@ -341,22 +276,21 @@ export class TaskTool implements AgentTool { const params = args as Partial; const lines: string[] = []; - if (typeof params.agent === "string") { + if (typeof params.resume === "string" && params.resume.trim()) { + lines.push(`Resume: ${truncateForPrompt(params.resume)}`); + } else if (typeof params.agent === "string") { lines.push(`Agent: ${truncateForPrompt(params.agent)}`); } - const tasks = Array.isArray(params.tasks) ? params.tasks : []; - const firstTask = tasks[0]; - if (firstTask) { - lines.push(`Task: ${truncateForPrompt(firstTask.id)}`); - lines.push(`Assignment:\n${truncateForPrompt(firstTask.assignment)}`); - if (tasks.length > 1) { - lines.push(`+${tasks.length - 1} more task${tasks.length === 2 ? "" : "s"}`); - } + if (typeof params.id === "string" && params.id.trim()) { + lines.push(`Task: ${truncateForPrompt(params.id)}`); + } + if (typeof params.assignment === "string") { + lines.push(`Assignment:\n${truncateForPrompt(params.assignment)}`); } return lines; }; readonly label = "Task"; - readonly summary = "Spawn a subagent to complete a parallel task"; + readonly summary = "Spawn a subagent to complete a task in the background"; readonly strict = true; readonly loadMode = "discoverable"; readonly renderResult = renderResult; @@ -366,6 +300,12 @@ export class TaskTool implements AgentTool> { const params = repairTaskParams(rawParams as TaskParams); const simpleMode = this.#getTaskSimpleMode(); - const validationError = validateTaskModeParams(simpleMode, params); + const validationError = validateTaskModeParams(simpleMode, params) ?? validateSpawnParams(params); if (validationError) { return createTaskModeError(validationError); } - const asyncEnabled = this.session.settings.get("async.enabled"); - const selectedAgent = this.#discoveredAgents.find(agent => agent.name === params.agent); - if (!asyncEnabled || selectedAgent?.blocking === true) { - return this.#executeSync(toolCallId, params, signal, onUpdate); - } - + const isResume = typeof params.resume === "string" && params.resume.trim().length > 0; + const selectedAgent = isResume ? undefined : this.#discoveredAgents.find(agent => agent.name === params.agent); const manager = this.session.asyncJobManager; - if (!manager) { - // Async was requested but no manager is registered (e.g. an - // orphaned session whose host never wired one up). Falling back - // to the sync path keeps the tool usable; only background/job-poll - // semantics are lost. - logger.warn("task: async.enabled but no AsyncJobManager registered; falling back to sync execution"); - return this.#executeSync(toolCallId, params, signal, onUpdate); + if (!manager || selectedAgent?.blocking === true) { + // Sync fallback: orphaned host that never wired a job manager, or an + // agent definition that declares `blocking: true`. The session-scoped + // semaphore still bounds fan-out across parallel task calls. + if (!manager) { + logger.warn("task: no AsyncJobManager registered; falling back to sync execution"); + } + const semaphore = this.#getSpawnSemaphore(); + await semaphore.acquire(); + try { + return await this.#executeSync(toolCallId, params, signal, onUpdate); + } finally { + semaphore.release(); + } } - const taskItems = params.tasks ?? []; - if (taskItems.length === 0) { - return this.#executeSync(toolCallId, params, signal, onUpdate); + // Resolve the agent id up front so the immediate result can name it. + let agentId: string; + if (isResume) { + agentId = params.resume!.trim(); + if (!AgentRegistry.global().get(agentId)) { + throw new ToolError( + `Unknown agent "${agentId}" — nothing to resume. Use \`irc\` op:"list" to see live agent ids; past transcripts are readable at history://${agentId}.`, + ); + } + } else { + const outputManager = + this.session.agentOutputManager ?? new AgentOutputManager(this.session.getArtifactsDir ?? (() => null)); + agentId = await outputManager.allocate(params.id?.trim() || generateTaskName()); } - const taskIdProblem = validateTaskIds(taskItems); - if (taskIdProblem) { - return createTaskModeError(taskIdProblem); - } - - const outputManager = - this.session.agentOutputManager ?? new AgentOutputManager(this.session.getArtifactsDir ?? (() => null)); - const uniqueIds = await outputManager.allocateBatch(taskItems.map(t => t.id)); - const fallbackAgentSource = - this.#discoveredAgents.find(agent => agent.name === params.agent)?.source ?? "bundled"; - const progressByTaskId = new Map(); - for (let index = 0; index < taskItems.length; index++) { - const taskItem = taskItems[index]; - const assignment = taskItem.assignment.trim(); - progressByTaskId.set(taskItem.id, { - index, - id: taskItem.id, - agent: params.agent, - agentSource: fallbackAgentSource, - status: "pending", - task: renderSubagentUserPrompt(assignment, simpleMode), - assignment, - description: taskItem.description, - recentTools: [], - recentOutput: [], - toolCount: 0, - tokens: 0, - cost: 0, - durationMs: 0, - }); - } - - const startedJobs: Array<{ jobId: string; taskId: string }> = []; - const failedSchedules: string[] = []; - let completedJobs = 0; - let failedJobs = 0; - - const getProgressSnapshot = (): AgentProgress[] => { - // Shallow copies: top-level fields are reassigned (never mutated in - // place) and the large nested payloads (extractedToolData) are - // immutable once attached — structuredClone here cost O(batch × payload) - // per progress event. - return Array.from(progressByTaskId.values()) - .sort((a, b) => a.index - b.index) - .map(progress => ({ ...progress })); + const assignment = (params.assignment ?? "").trim(); + const agentLabel = isResume + ? (AgentRegistry.global().get(agentId)?.displayName ?? "task") + : (params.agent ?? "task"); + const progress: AgentProgress = { + index: 0, + id: agentId, + agent: agentLabel, + agentSource: selectedAgent?.source ?? "bundled", + status: "pending", + task: renderSubagentUserPrompt(assignment, simpleMode), + assignment, + description: params.description, + recentTools: [], + recentOutput: [], + toolCount: 0, + requests: 0, + tokens: 0, + cost: 0, + durationMs: 0, }; const buildAsyncDetails = (state: "running" | "completed" | "failed", jobId: string): TaskToolDetails => ({ projectAgentsDir: null, results: [], totalDurationMs: 0, - progress: getProgressSnapshot(), + progress: [{ ...progress }], async: { state, jobId, type: "task" }, }); - const emitAsyncUpdate = (state: "running" | "completed" | "failed", text: string): void => { - const primaryJobId = startedJobs[0]?.jobId ?? "task"; - onUpdate?.({ - content: [{ type: "text", text }], - details: buildAsyncDetails(state, primaryJobId), - }); + const buildResumeHint = (aborted: boolean): string => { + if (aborted) { + return `\n\n${agentId} was aborted — transcript at history://${agentId}`; + } + return `\n\n${agentId} is now idle — task(resume:"${agentId}") to continue it, transcript at history://${agentId}`; }; - const maxConcurrency = this.session.settings.get("task.maxConcurrency"); - const semaphore = new Semaphore(maxConcurrency); - - for (let i = 0; i < taskItems.length; i++) { - const taskItem = taskItems[i]; - if (signal?.aborted) { - failedSchedules.push(`${taskItem.id}: cancelled before scheduling`); - completedJobs += 1; - const progress = progressByTaskId.get(taskItem.id); - if (progress) { - progress.status = "aborted"; - } - continue; - } - - const uniqueId = uniqueIds[i]; - const singleParams: TaskParams = { ...params, tasks: [taskItem] }; - const label = uniqueId; - try { - const jobId = manager.register( - "task", - label, - async ({ signal: runSignal, reportProgress, markRunning }) => { - const startedAt = Date.now(); - const progress = progressByTaskId.get(taskItem.id); - await semaphore.acquire(); - if (runSignal.aborted) { - semaphore.release(); - if (progress) { - progress.status = "aborted"; - } - completedJobs += 1; - failedJobs += 1; - throw new Error("Aborted before execution"); - } - markRunning(); - if (progress) { - progress.status = "running"; - } + let jobId: string; + try { + jobId = manager.register( + "task", + agentId, + async ({ jobId: ownJobId, signal: runSignal, reportProgress, markRunning }) => { + const startedAt = Date.now(); + const semaphore = this.#getSpawnSemaphore(); + await semaphore.acquire(); + if (runSignal.aborted) { + semaphore.release(); + progress.status = "aborted"; + throw new Error("Aborted before execution"); + } + markRunning(); + progress.status = "running"; + await reportProgress( + `Running background task ${agentId}...`, + buildAsyncDetails("running", ownJobId) as unknown as Record, + ); + try { + const result = await this.#executeSync(toolCallId, params, runSignal, undefined, agentId); + const finalText = result.content.find(part => part.type === "text")?.text ?? "(no output)"; + const singleResult = result.details?.results[0]; + // A missing result means the sync path failed at the tool level + // (results: []) — treat it as a failure, not success. + const resultFailed = !singleResult || (singleResult.aborted ?? false) || singleResult.exitCode !== 0; + progress.status = singleResult?.aborted ? "aborted" : resultFailed ? "failed" : "completed"; + progress.durationMs = singleResult?.durationMs ?? Math.max(0, Date.now() - startedAt); + progress.tokens = singleResult?.tokens ?? 0; + progress.requests = singleResult?.requests ?? 0; + progress.contextTokens = singleResult?.contextTokens; + progress.contextWindow = singleResult?.contextWindow; + progress.cost = singleResult?.usage?.cost.total ?? 0; + progress.extractedToolData = singleResult?.extractedToolData; + progress.retryFailure = singleResult?.retryFailure; + progress.retryState = undefined; + const statusText = resultFailed + ? `Background task ${agentId} failed.` + : `Background task ${agentId} complete.`; await reportProgress( - `Running background task ${taskItem.id}...`, - buildAsyncDetails("running", startedJobs[0]?.jobId ?? label) as unknown as Record, + statusText, + buildAsyncDetails(resultFailed ? "failed" : "completed", ownJobId) as unknown as Record< + string, + unknown + >, ); - try { - const result = await this.#executeSync(toolCallId, singleParams, runSignal, undefined, [uniqueId]); - const finalText = result.content.find(part => part.type === "text")?.text ?? "(no output)"; - const singleResult = result.details?.results[0]; - // A missing per-task result means #executeSync failed at the - // tool level (results: []) — treat it as a failure, not success. - const resultFailed = - !singleResult || (singleResult.aborted ?? false) || singleResult.exitCode !== 0; - if (progress) { - progress.status = singleResult?.aborted ? "aborted" : resultFailed ? "failed" : "completed"; - progress.durationMs = singleResult?.durationMs ?? Math.max(0, Date.now() - startedAt); - progress.tokens = singleResult?.tokens ?? 0; - progress.contextTokens = singleResult?.contextTokens; - progress.contextWindow = singleResult?.contextWindow; - progress.cost = singleResult?.usage?.cost.total ?? 0; - progress.extractedToolData = singleResult?.extractedToolData; - progress.retryFailure = singleResult?.retryFailure; - progress.retryState = undefined; - } - completedJobs += 1; - if (resultFailed) { - failedJobs += 1; - } - const remaining = taskItems.length - completedJobs; - const isDone = remaining === 0; - await reportProgress( - isDone - ? `Background task batch complete: ${completedJobs}/${taskItems.length} finished.` - : `Background task batch progress: ${completedJobs}/${taskItems.length} finished (${remaining} running).`, - buildAsyncDetails( - isDone ? (failedJobs > 0 || failedSchedules.length > 0 ? "failed" : "completed") : "running", - startedJobs[0]?.jobId ?? label, - ) as unknown as Record, - ); - if (isDone) { - emitAsyncUpdate( - failedJobs > 0 || failedSchedules.length > 0 ? "failed" : "completed", - `Background task batch complete: ${completedJobs}/${taskItems.length} finished.`, - ); - } - if (resultFailed) { - // Mark the job itself failed; counters above are already updated. - throw new TaskJobError(finalText); - } - return finalText; - } catch (error) { - if (error instanceof TaskJobError) { - throw error; - } - if (progress) { - progress.status = "failed"; - progress.durationMs = Math.max(0, Date.now() - startedAt); - } - completedJobs += 1; - failedJobs += 1; - const remaining = taskItems.length - completedJobs; - const isDone = remaining === 0; - await reportProgress( - isDone - ? `Background task batch complete with failures: ${failedJobs} failed.` - : `Background task batch progress: ${completedJobs}/${taskItems.length} finished (${remaining} running).`, - buildAsyncDetails( - isDone ? "failed" : "running", - startedJobs[0]?.jobId ?? label, - ) as unknown as Record, - ); - if (isDone) { - emitAsyncUpdate( - "failed", - `Background task batch complete with failures: ${failedJobs} failed.`, - ); - } - throw error; - } finally { - semaphore.release(); + onUpdate?.({ + content: [{ type: "text", text: statusText }], + details: buildAsyncDetails(resultFailed ? "failed" : "completed", ownJobId), + }); + const deliveryText = `${finalText}${buildResumeHint(singleResult?.aborted === true)}`; + if (resultFailed) { + // Mark the job itself failed; the failed agent stays interrogable. + throw new TaskJobError(deliveryText); } + return deliveryText; + } catch (error) { + if (error instanceof TaskJobError) { + throw error; + } + progress.status = "failed"; + progress.durationMs = Math.max(0, Date.now() - startedAt); + const statusText = `Background task ${agentId} failed.`; + await reportProgress( + statusText, + buildAsyncDetails("failed", ownJobId) as unknown as Record, + ); + onUpdate?.({ + content: [{ type: "text", text: statusText }], + details: buildAsyncDetails("failed", ownJobId), + }); + const message = error instanceof Error ? error.message : String(error); + const hint = AgentRegistry.global().get(agentId) ? buildResumeHint(false) : ""; + throw new TaskJobError(`${message}${hint}`); + } finally { + semaphore.release(); + } + }, + { + id: agentId, + queued: true, + ownerId: this.session.getAgentId?.() ?? undefined, + onProgress: (text, details) => { + const progressDetails = + (details as TaskToolDetails | undefined) ?? buildAsyncDetails("running", agentId); + onUpdate?.({ content: [{ type: "text", text }], details: progressDetails }); }, - { - id: label, - queued: true, - ownerId: this.session.getAgentId?.() ?? undefined, - onProgress: (text, details) => { - const progressDetails = - (details as TaskToolDetails | undefined) ?? - buildAsyncDetails("running", startedJobs[0]?.jobId ?? label); - onUpdate?.({ content: [{ type: "text", text }], details: progressDetails }); - }, - }, - ); - startedJobs.push({ jobId, taskId: taskItem.id }); - } catch (error) { - const message = error instanceof Error ? error.message : String(error); - failedSchedules.push(`${taskItem.id}: ${message}`); - completedJobs += 1; - const progress = progressByTaskId.get(taskItem.id); - if (progress) { - progress.status = "failed"; - } - } - } - - if (startedJobs.length === 0) { - const failureText = `Failed to start background task jobs: ${failedSchedules.join("; ")}`; + }, + ); + } catch (error) { + const message = error instanceof Error ? error.message : String(error); return { - content: [{ type: "text", text: failureText }], + content: [{ type: "text", text: `Failed to start background task job: ${message}` }], details: { projectAgentsDir: null, results: [], totalDurationMs: 0 }, }; } - emitAsyncUpdate( - "running", - `Launching ${startedJobs.length} background ${startedJobs.length === 1 ? "task" : "tasks"}...`, - ); - - const scheduleFailureSummary = - failedSchedules.length > 0 - ? ` Failed to schedule ${failedSchedules.length} task${failedSchedules.length === 1 ? "" : "s"}.` - : ""; - const ircEnabled = this.session.settings.get("irc.enabled") === true; - const taskIdByItemId = new Map(); - for (let i = 0; i < taskItems.length; i++) { - taskIdByItemId.set(taskItems[i].id, uniqueIds[i]); - } - const startedListing = startedJobs - .map(({ taskId, jobId }) => { - const id = taskIdByItemId.get(taskId) ?? taskId; - const desc = progressByTaskId.get(taskId)?.description; - const prefix = `- \`${id}\` (job \`${jobId}\`)`; - return desc ? `${prefix} — ${desc}` : prefix; - }) - .join("\n"); const coordinationHint = ircEnabled - ? ` DM these ids via \`irc\` to coordinate while they run; reach for \`job\` only to inspect (\`list\`), wait (\`poll\`), or cancel a stuck task.` - : ` Use \`job\` to inspect (\`list\`), wait (\`poll\`), or cancel a stuck task by id.`; + ? `DM \`${agentId}\` via \`irc\` to coordinate while it runs; use \`job\` only to inspect (\`list\`), wait (\`poll\`), or cancel a stuck task.` + : `Use \`job\` to inspect (\`list\`), wait (\`poll\`), or cancel a stuck task.`; + const verb = isResume ? "Resumed" : "Spawned"; + const descriptionSuffix = params.description ? ` — ${params.description}` : ""; + + onUpdate?.({ + content: [{ type: "text", text: `${verb} agent \`${agentId}\`...` }], + details: buildAsyncDetails("running", jobId), + }); return { content: [ { type: "text", - text: `Started ${startedJobs.length} background task job${startedJobs.length === 1 ? "" : "s"} using ${params.agent}.${scheduleFailureSummary} Results will be delivered when complete.\n${startedListing}\n${coordinationHint}`, + text: `${verb} agent \`${agentId}\` (job \`${jobId}\`)${descriptionSuffix}. The result will be delivered when it yields. ${coordinationHint}`, }, ], details: { projectAgentsDir: null, results: [], totalDurationMs: 0, - progress: getProgressSnapshot(), - async: { state: "running", jobId: startedJobs[0].jobId, type: "task" }, + progress: [{ ...progress }], + async: { state: "running", jobId, type: "task" }, }, }; } + /** + * Synchronous execution of one spawn or resume. Used as the body of every + * async job and directly by the sync fallback (no job manager / blocking + * agent) and by in-process callers that need the result inline (e.g. the + * commit flow's analyze_files tool). + */ async #executeSync( toolCallId: string, params: TaskParams, signal?: AbortSignal, onUpdate?: AgentToolUpdateCallback, - preAllocatedIds?: string[], + preAllocatedId?: string, + ): Promise> { + if (typeof params.resume === "string" && params.resume.trim().length > 0) { + return this.#executeResume(toolCallId, params, signal, onUpdate); + } + return this.#runSpawn(toolCallId, params, signal, onUpdate, preAllocatedId); + } + + /** + * Resume an existing agent: revive it if parked, inject the follow-up + * assignment through the session's normal prompt path, and run it through + * the same yield/finalize pipeline as a spawn. The session stays alive + * (idle, TTL re-armed) afterwards. + */ + async #executeResume( + toolCallId: string, + params: TaskParams, + signal?: AbortSignal, + onUpdate?: AgentToolUpdateCallback, + ): Promise> { + const startTime = Date.now(); + const resumeId = params.resume!.trim(); + const simpleMode = this.#getTaskSimpleMode(); + const { customSchemaEnabled } = getTaskSimpleModeCapabilities(simpleMode); + const assignment = (params.assignment ?? "").trim(); + + let session: AgentSession; + try { + session = await AgentLifecycleManager.global().ensureLive(resumeId); + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + throw new ToolError( + `Cannot resume "${resumeId}": ${message} Use \`irc\` op:"list" to see live agent ids; transcripts are readable at history://${resumeId}.`, + ); + } + + const agentName = AgentRegistry.global().get(resumeId)?.displayName ?? "task"; + const agentDef: AgentDefinition = getAgent(this.#discoveredAgents, agentName) ?? { + name: agentName, + description: "", + systemPrompt: "", + source: "bundled", + }; + + // Resumed output artifacts overwrite agent:// in the parent's + // artifacts dir; the transcript accretes in the session JSONL. + const sessionFile = this.session.getSessionFile(); + const artifactsDir = sessionFile ? sessionFile.slice(0, -6) : undefined; + + const result = await resumeSubprocess({ + session, + id: resumeId, + agent: agentDef, + task: renderSubagentUserPrompt(assignment, simpleMode), + assignment, + description: params.description, + index: 0, + parentToolCallId: toolCallId, + outputSchema: customSchemaEnabled ? params.schema : undefined, + signal, + onProgress: progress => { + onUpdate?.({ + content: [{ type: "text", text: `Resuming ${resumeId}...` }], + details: { + projectAgentsDir: null, + results: [], + totalDurationMs: Date.now() - startTime, + progress: [{ ...progress, recentTools: progress.recentTools.slice() }], + }, + }); + }, + eventBus: this.session.eventBus, + settings: this.session.settings, + artifactsDir, + }); + + return this.#buildResultPayload(result, null, Date.now() - startTime, ""); + } + + /** Spawn a fresh subagent and run it to completion. */ + async #runSpawn( + toolCallId: string, + params: TaskParams, + signal?: AbortSignal, + onUpdate?: AgentToolUpdateCallback, + preAllocatedId?: string, ): Promise> { const startTime = Date.now(); const { agents, projectAgentsDir } = await discoverAgents(this.session.cwd); - const { agent: agentName, context, schema: outputSchema } = params; + const agentName = params.agent ?? ""; const simpleMode = this.#getTaskSimpleMode(); - const { contextEnabled, customSchemaEnabled } = getTaskSimpleModeCapabilities(simpleMode); - const sharedContext = contextEnabled ? context?.trim() : undefined; + const { customSchemaEnabled } = getTaskSimpleModeCapabilities(simpleMode); + const outputSchema = params.schema; + const assignment = (params.assignment ?? "").trim(); const isolationMode = this.session.settings.get("task.isolation.mode"); const isolationRequested = "isolated" in params ? params.isolated === true : false; const isIsolated = isolationMode !== "none" && isolationRequested; const mergeMode = this.session.settings.get("task.isolation.merge"); const commitStyle = this.session.settings.get("task.isolation.commits"); - const maxConcurrency = this.session.settings.get("task.maxConcurrency"); const taskDepth = this.session.taskDepth ?? 0; const subagentLspEnabled = (this.session.enableLsp ?? true) && this.session.settings.get("task.enableLsp"); if (isolationMode === "none" && "isolated" in params) { return { - content: [ - { - type: "text", - text: "Task isolation is disabled.", - }, - ], - details: { - projectAgentsDir, - results: [], - totalDurationMs: 0, - }, + content: [{ type: "text", text: "Task isolation is disabled." }], + details: { projectAgentsDir, results: [], totalDurationMs: 0 }, }; } @@ -748,17 +692,8 @@ export class TaskTool implements AgentTool a.name).join(", ") || "none"; return { - content: [ - { - type: "text", - text: `Unknown agent "${agentName}". Available: ${available}`, - }, - ], - details: { - projectAgentsDir, - results: [], - totalDurationMs: 0, - }, + content: [{ type: "text", text: `Unknown agent "${agentName}". Available: ${available}` }], + details: { projectAgentsDir, results: [], totalDurationMs: 0 }, }; } @@ -773,11 +708,7 @@ export class TaskTool implements AgentTool 0 ? ` Available: ${enabled.join(", ")}` : ""}`, }, ], - details: { - projectAgentsDir, - results: [], - totalDurationMs: 0, - }, + details: { projectAgentsDir, results: [], totalDurationMs: 0 }, }; } @@ -817,38 +748,6 @@ export class TaskTool implements AgentTool(); - - // Update callback - const emitProgress = () => { - const progress = Array.from(progressMap.values()).sort((a, b) => a.index - b.index); - onUpdate?.({ - content: [{ type: "text", text: `Running ${params.tasks.length} agents...` }], - details: { - projectAgentsDir, - results: [], - totalDurationMs: Date.now() - startTime, - progress, - }, - }); - }; - try { // Check self-recursion prevention if (this.#blockedAgent && agentName === this.#blockedAgent) { @@ -928,11 +801,7 @@ export class TaskTool implements AgentTool null)); - uniqueIds = await outputManager.allocateBatch(tasks.map(t => t.id)); + agentId = await outputManager.allocate(params.id?.trim() || generateTaskName()); } - const tasksWithUniqueIds = tasks.map((t, i) => ({ ...t, id: uniqueIds[i] })); const availableSkills = [...(this.session.skills ?? [])]; // Resolve autoload skills from agent definition against available skills @@ -995,85 +849,101 @@ export class TaskTool implements AgentTool { + onUpdate?.({ + content: [{ type: "text", text: `Running agent ${agentId}...` }], + details: { + projectAgentsDir, + results: [], + totalDurationMs: Date.now() - startTime, + progress: [latestProgress], + }, }); - } + }; emitProgress(); - const runTask = async ( - task: (typeof tasksWithUniqueIds)[number], - index: number, - workerSignal?: AbortSignal, - ) => { + const buildCommitMessageFn = () => + commitStyle === "ai" && this.session.modelRegistry + ? async (diff: string) => { + return generateCommitMessage( + diff, + this.session.modelRegistry!, + this.session.settings, + this.session.getSessionId?.() ?? undefined, + ); + } + : undefined; + + const sharedRunOptions = { + cwd: this.session.cwd, + agent: effectiveAgent, + task: renderSubagentUserPrompt(assignment, simpleMode), + assignment, + planReference, + description: params.description, + index: 0, + parentToolCallId: toolCallId, + id: agentId, + taskDepth, + modelOverride, + parentActiveModelPattern, + thinkingLevel: thinkingLevelOverride, + outputSchema: effectiveOutputSchema, + sessionFile, + persistArtifacts: !!artifactsDir, + artifactsDir: effectiveArtifactsDir, + enableLsp: subagentLspEnabled, + signal, + eventBus: this.session.eventBus, + onProgress: (progress: AgentProgress) => { + // Shallow snapshot; recentTools is mutated in place by the + // executor, the rest is reassigned or immutable. A deep clone + // here cost O(extractedToolData) per progress event. + latestProgress = { ...progress, recentTools: progress.recentTools.slice() }; + emitProgress(); + }, + authStorage: this.session.authStorage, + modelRegistry: this.session.modelRegistry, + settings: this.session.settings, + mcpManager, + contextFiles, + skills: availableSkills, + autoloadSkills: resolvedAutoloadSkills, + workspaceTree: this.session.workspaceTree, + promptTemplates, + rules: this.session.rules, + preloadedExtensionPaths: this.session.extensionPaths, + preloadedCustomToolPaths: this.session.customToolPaths, + localProtocolOptions, + parentArtifactManager, + parentHindsightSessionState: this.session.getHindsightSessionState?.(), + parentMnemopiSessionState: this.session.getMnemopiSessionState?.(), + parentTelemetry: this.session.getTelemetry?.(), + parentEvalSessionId, + }; + + const runTask = async (): Promise => { if (!isIsolated) { - return runSubprocess({ - cwd: this.session.cwd, - agent: effectiveAgent, - task: renderSubagentUserPrompt(task.assignment, simpleMode), - assignment: task.assignment.trim(), - context: sharedContext, - planReference, - description: task.description, - index, - parentToolCallId: toolCallId, - id: task.id, - taskDepth, - modelOverride, - parentActiveModelPattern, - thinkingLevel: thinkingLevelOverride, - outputSchema: effectiveOutputSchema, - sessionFile, - persistArtifacts: !!artifactsDir, - artifactsDir: effectiveArtifactsDir, - contextFile: contextFilePath, - enableLsp: subagentLspEnabled, - signal: workerSignal ?? signal, - eventBus: this.session.eventBus, - onProgress: progress => { - // Shallow snapshot; recentTools is mutated in place by the - // executor, the rest is reassigned or immutable. A deep clone - // here cost O(extractedToolData) per progress event. - progressMap.set(index, { ...progress, recentTools: progress.recentTools.slice() }); - emitProgress(); - }, - authStorage: this.session.authStorage, - modelRegistry: this.session.modelRegistry, - settings: this.session.settings, - mcpManager, - contextFiles, - skills: availableSkills, - autoloadSkills: resolvedAutoloadSkills, - workspaceTree: this.session.workspaceTree, - promptTemplates, - rules: this.session.rules, - preloadedExtensionPaths: this.session.extensionPaths, - preloadedCustomToolPaths: this.session.customToolPaths, - localProtocolOptions, - parentArtifactManager, - parentHindsightSessionState: this.session.getHindsightSessionState?.(), - parentMnemopiSessionState: this.session.getMnemopiSessionState?.(), - parentTelemetry: this.session.getTelemetry?.(), - parentEvalSessionId, - }); + return runSubprocess(sharedRunOptions); } const taskStart = Date.now(); @@ -1084,73 +954,25 @@ export class TaskTool implements AgentTool { - progressMap.set(index, { ...progress, recentTools: progress.recentTools.slice() }); - emitProgress(); - }, - authStorage: this.session.authStorage, - modelRegistry: this.session.modelRegistry, - settings: this.session.settings, - mcpManager, - contextFiles, - skills: availableSkills, - autoloadSkills: resolvedAutoloadSkills, - workspaceTree: this.session.workspaceTree, - promptTemplates, - rules: this.session.rules, - localProtocolOptions, - parentArtifactManager, - parentHindsightSessionState: this.session.getHindsightSessionState?.(), - parentMnemopiSessionState: this.session.getMnemopiSessionState?.(), - parentTelemetry: this.session.getTelemetry?.(), - parentEvalSessionId, + preloadedExtensionPaths: undefined, + preloadedCustomToolPaths: undefined, }); if (mergeMode === "branch" && result.exitCode === 0) { try { - const commitMsg = - commitStyle === "ai" && this.session.modelRegistry - ? async (diff: string) => { - return generateCommitMessage( - diff, - this.session.modelRegistry!, - this.session.settings, - this.session.getSessionId?.() ?? undefined, - ); - } - : undefined; const commitResult = await commitToBranch( isolationDir, taskBaseline, - task.id, - task.description, - commitMsg, + agentId, + params.description, + buildCommitMessageFn(), ); return { ...result, @@ -1159,7 +981,7 @@ export class TaskTool implements AgentTool { - if (result !== undefined) { - return result; - } - const task = tasksWithUniqueIds[index]; - const assignment = task.assignment.trim(); - return { - index, - id: task.id, - agent: agentName, - agentSource: agent.source, - task: renderSubagentUserPrompt(assignment, simpleMode), - assignment, - description: task.description, - exitCode: 1, - output: "", - stderr: "Skipped (cancelled before start)", - truncated: false, - durationMs: 0, - tokens: 0, - modelOverride, - error: "Cancelled before start", - aborted: true, - abortReason: "Cancelled before start", - }; - }); - - // Aggregate usage from executor results (already accumulated incrementally) - const aggregatedUsage = createUsageTotals(); - let hasAggregatedUsage = false; - for (const result of results) { - if (result.usage) { - addUsageTotals(aggregatedUsage, result.usage); - hasAggregatedUsage = true; - } - } - - // Collect output paths (artifacts already written by executor in real-time) - const outputPaths: string[] = []; - const patchPaths: string[] = []; - for (const result of results) { - if (result.outputPath) { - outputPaths.push(result.outputPath); - } - if (result.patchPath) { - patchPaths.push(result.patchPath); - } - } + const result = await runTask(); let mergeSummary = ""; let changesApplied: boolean | null = null; let hadAnyChanges = false; - let mergedBranchesForNestedPatches: Set | null = null; + let mergedBranchForNestedPatches = false; if (isIsolated && repoRoot) { try { if (mergeMode === "branch") { - // Branch mode: merge task branches sequentially - const branchEntries = results - .filter(r => r.branchName && r.exitCode === 0 && !r.aborted) - .map(r => ({ branchName: r.branchName!, taskId: r.id, description: r.description })); - - if (branchEntries.length === 0) { + if (!result.branchName || result.exitCode !== 0 || result.aborted) { changesApplied = true; - hadAnyChanges = false; mergeSummary = "\n\nNo changes to apply."; } else { - const mergeResult = await mergeTaskBranches(repoRoot, branchEntries); - mergedBranchesForNestedPatches = new Set(mergeResult.merged); + const mergeResult = await mergeTaskBranches(repoRoot, [ + { branchName: result.branchName, taskId: result.id, description: result.description }, + ]); + mergedBranchForNestedPatches = mergeResult.merged.includes(result.branchName); changesApplied = mergeResult.failed.length === 0; hadAnyChanges = changesApplied && mergeResult.merged.length > 0; if (changesApplied) { mergeSummary = hadAnyChanges - ? `\n\nMerged ${mergeResult.merged.length} branch${mergeResult.merged.length === 1 ? "" : "es"}: ${mergeResult.merged.join(", ")}` + ? `\n\nMerged branch: ${result.branchName}` : "\n\nNo changes to apply."; } else { - const mergedPart = - mergeResult.merged.length > 0 ? `Merged: ${mergeResult.merged.join(", ")}.\n` : ""; - const failedPart = `Failed: ${mergeResult.failed.join(", ")}.`; const conflictPart = mergeResult.conflict ? `\nConflict: ${mergeResult.conflict}` : ""; - mergeSummary = `\n\nBranch merge failed. ${mergedPart}${failedPart}${conflictPart}\nUnmerged branches remain for manual resolution.`; + mergeSummary = `\n\nBranch merge failed: ${result.branchName}.${conflictPart}\nThe unmerged branch remains for manual resolution.`; } if (mergeResult.stashConflict) { mergeSummary += `\n\n${mergeResult.stashConflict}`; } - } - // Clean up merged branches (keep failed ones for manual resolution) - const allBranches = branchEntries.map(b => b.branchName); - if (changesApplied) { - await cleanupTaskBranches(repoRoot, allBranches); + // Clean up the merged branch (keep failed ones for manual resolution) + if (changesApplied) { + await cleanupTaskBranches(repoRoot, [result.branchName]); + } } } else { - // Patch mode: apply patches from successful tasks. Failed or - // aborted siblings must not block completed work from landing. - const successfulResults = results.filter(r => r.exitCode === 0 && !r.error && !r.aborted); - const patchesInOrder = successfulResults.map(result => result.patchPath).filter(Boolean) as string[]; - const missingPatch = successfulResults.some(result => !result.patchPath); - if (missingPatch) { + // Patch mode: apply the patch from a successful run. A failed or + // aborted run has nothing to apply and must not block the result. + const succeeded = result.exitCode === 0 && !result.error && !result.aborted; + if (!succeeded) { + changesApplied = true; + hadAnyChanges = false; + } else if (!result.patchPath) { changesApplied = false; hadAnyChanges = false; } else { - const patchStats = await Promise.all( - patchesInOrder.map(async patchPath => ({ - patchPath, - size: (await fs.stat(patchPath)).size, - })), - ); - const nonEmptyPatches = patchStats.filter(patch => patch.size > 0).map(patch => patch.patchPath); - if (nonEmptyPatches.length === 0) { + const patchText = await Bun.file(result.patchPath).text(); + if (!patchText.trim()) { changesApplied = true; hadAnyChanges = false; } else { - const patchTexts = await Promise.all( - nonEmptyPatches.map(async patchPath => Bun.file(patchPath).text()), - ); - const combinedPatch = patchTexts - .map(text => (text.endsWith("\n") ? text : `${text}\n`)) - .join(""); - if (!combinedPatch.trim()) { - changesApplied = true; - hadAnyChanges = false; - } else { - changesApplied = await git.patch.canApplyText(repoRoot, combinedPatch); - if (changesApplied) { - try { - await git.patch.applyText(repoRoot, combinedPatch); - hadAnyChanges = true; - } catch { - changesApplied = false; - hadAnyChanges = false; - } + const normalized = patchText.endsWith("\n") ? patchText : `${patchText}\n`; + changesApplied = await git.patch.canApplyText(repoRoot, normalized); + if (changesApplied) { + try { + await git.patch.applyText(repoRoot, normalized); + hadAnyChanges = true; + } catch { + changesApplied = false; + hadAnyChanges = false; } } } @@ -1359,10 +1102,7 @@ export class TaskTool implements AgentTool 0 - ? `\n\nPatch artifacts:\n${patchPaths.map(patch => `- ${patch}`).join("\n")}` - : ""; + const patchList = result.patchPath ? `\n\nPatch artifact:\n- ${result.patchPath}` : ""; mergeSummary = `\n\n${notification}${patchList}`; } } @@ -1376,34 +1116,15 @@ export class TaskTool implements AgentTool { - if (!r.nestedPatches || r.nestedPatches.length === 0 || r.exitCode !== 0 || r.aborted) { - return false; - } - if (mergeMode !== "branch") { - return true; - } - if (!r.branchName || !mergedBranchesForNestedPatches) { - return false; - } - return mergedBranchesForNestedPatches.has(r.branchName); - }) - .flatMap(r => r.nestedPatches!); - if (allNestedPatches.length > 0) { + const nestedPatches = result.nestedPatches ?? []; + const eligible = + nestedPatches.length > 0 && + result.exitCode === 0 && + !result.aborted && + (mergeMode !== "branch" || mergedBranchForNestedPatches); + if (eligible) { try { - const commitMsg = - commitStyle === "ai" && this.session.modelRegistry - ? async (diff: string) => { - return generateCommitMessage( - diff, - this.session.modelRegistry!, - this.session.settings, - this.session.getSessionId?.() ?? undefined, - ); - } - : undefined; - await applyNestedPatches(repoRoot, allNestedPatches, commitMsg); + await applyNestedPatches(repoRoot, nestedPatches, buildCommitMessageFn()); } catch { // Nested patch failures are non-fatal to the parent merge mergeSummary += @@ -1412,58 +1133,6 @@ export class TaskTool implements AgentTool r.aborted).length; - const successCount = results.filter(r => r.exitCode === 0 && !r.error && !r.aborted).length; - const totalDuration = Date.now() - startTime; - - const summaries = results.map(r => { - const status = r.aborted - ? "cancelled" - : r.exitCode === 0 && r.error - ? "merge failed" - : r.exitCode === 0 - ? "completed" - : `failed (exit ${r.exitCode})`; - const output = r.output.trim() || r.stderr.trim() || "(no output)"; - const outputCharCount = r.outputMeta?.charCount ?? output.length; - const fullOutputThreshold = 5000; - let preview = output; - let truncated = false; - if (outputCharCount > fullOutputThreshold) { - const slice = output.slice(0, fullOutputThreshold); - const lastNewline = slice.lastIndexOf("\n"); - preview = lastNewline >= 0 ? slice.slice(0, lastNewline) : slice; - truncated = true; - } - return { - agent: r.agent, - status, - id: r.id, - preview, - truncated, - meta: r.outputMeta - ? { - lineCount: r.outputMeta.lineCount, - charSize: formatBytes(r.outputMeta.charCount), - } - : undefined, - }; - }); - - const outputIds = results.filter(r => !r.aborted || r.output.trim()).map(r => `agent://${r.id}`); - const summary = prompt.render(taskSummaryTemplate, { - successCount, - totalCount: results.length, - cancelledCount, - hasCancelledNote: aborted && cancelledCount > 0, - duration: formatDuration(totalDuration), - summaries, - outputIds, - agentName, - mergeSummary, - }); - // Cleanup temp directory if used const shouldCleanupTempArtifacts = tempArtifactsDir && (!isIsolated || changesApplied === true || changesApplied === null); @@ -1471,25 +1140,65 @@ export class TaskTool implements AgentTool { + const status = result.aborted + ? "cancelled" + : result.exitCode === 0 && result.error + ? "merge failed" + : result.exitCode === 0 + ? "completed" + : `failed (exit ${result.exitCode})`; + const output = formatResultOutputFallback(result); + const outputCharCount = result.outputMeta?.charCount ?? output.length; + const fullOutputThreshold = 5000; + let preview = output; + let truncated = false; + if (outputCharCount > fullOutputThreshold) { + const slice = output.slice(0, fullOutputThreshold); + const lastNewline = slice.lastIndexOf("\n"); + preview = lastNewline >= 0 ? slice.slice(0, lastNewline) : slice; + truncated = true; + } + const summary = prompt.render(taskSummaryTemplate, { + agentName: result.agent, + id: result.id, + status, + duration: formatDuration(totalDurationMs), + preview, + truncated, + meta: result.outputMeta + ? { + lineCount: result.outputMeta.lineCount, + charSize: formatBytes(result.outputMeta.charCount), + } + : undefined, + mergeSummary, + }); + + return { + content: [{ type: "text", text: summary }], + details: { + projectAgentsDir, + results: [result], + totalDurationMs, + usage: result.usage, + outputPaths: result.outputPath ? [result.outputPath] : undefined, + }, + }; + } } diff --git a/packages/coding-agent/src/task/output-manager.ts b/packages/coding-agent/src/task/output-manager.ts index 46d6593c0..74fba7832 100644 --- a/packages/coding-agent/src/task/output-manager.ts +++ b/packages/coding-agent/src/task/output-manager.ts @@ -85,15 +85,4 @@ export class AgentOutputManager { await this.#ensureInitialized(); return this.#allocateUnique(id); } - - /** - * Allocate unique IDs for a batch of tasks. - * - * @param ids Array of requested IDs - * @returns Array of unique IDs in same order - */ - async allocateBatch(ids: string[]): Promise { - await this.#ensureInitialized(); - return ids.map(id => this.#allocateUnique(id)); - } } diff --git a/packages/coding-agent/src/task/render.ts b/packages/coding-agent/src/task/render.ts index fb881a2e6..556f59438 100644 --- a/packages/coding-agent/src/task/render.ts +++ b/packages/coding-agent/src/task/render.ts @@ -33,7 +33,7 @@ import { import { framedBlock, renderStatusLine } from "../tui"; import { repairDoubleEncodedJsonString } from "./repair-args"; import { subprocessToolRegistry } from "./subprocess-tool-registry"; -import type { AgentProgress, SingleResult, TaskItem, TaskParams, TaskToolDetails } from "./types"; +import type { AgentProgress, SingleResult, TaskParams, TaskToolDetails } from "./types"; /** * Get status icon for agent state. @@ -62,6 +62,7 @@ function appendAgentStats( line: string, opts: { toolCount?: number; + requests?: number; tokens: number; contextTokens?: number; contextWindow?: number; @@ -74,6 +75,9 @@ function appendAgentStats( if (opts.toolCount) { line += `${theme.sep.dot}${theme.fg("dim", `${formatNumber(opts.toolCount)} ${theme.icon.extensionTool}`)}`; } + if (opts.requests) { + line += `${theme.sep.dot}${theme.fg("dim", `${formatNumber(opts.requests)} req`)}`; + } // Current per-turn context — match the status line's `%/` gauge (e.g. `5.1%/1M`). if (opts.contextTokens && opts.contextTokens > 0) { const ctx = @@ -505,65 +509,57 @@ function formatOutputInline(data: unknown, theme: Theme, maxWidth = 80): string } /** - * Render the per-task list (`id` + ui `description`) for the streaming call - * preview. The args stream in token by token, so the array grows over time and - * trailing entries may be partially parsed — every field access is defensive. + * Render the call preview lines for the single spawned/resumed agent. The + * args stream in token by token, so every field access is defensive. */ -function renderTaskItemLines(tasks: TaskItem[] | undefined, expanded: boolean, theme: Theme): string[] { - const items = tasks ?? []; - if (items.length === 0) return []; - +function renderTaskCallLines(args: Partial | undefined, theme: Theme): string[] { + if (!args) return []; const bullet = theme.fg("dim", "•"); - const cap = expanded ? items.length : Math.min(items.length, 12); - const truncated = cap < items.length; - const lines: string[] = []; - for (let i = 0; i < cap; i++) { - const task = items[i] as Partial | undefined; - const rawId = task?.id?.trim(); - const idLabel = rawId ? formatTaskId(rawId) : `#${i + 1}`; - let line = `${bullet} ${theme.fg("accent", theme.bold(idLabel))}`; - const desc = task?.description?.trim(); + + const resume = typeof args.resume === "string" ? args.resume.trim() : ""; + const rawId = typeof args.id === "string" ? args.id.trim() : ""; + const idLabel = resume ? formatTaskId(resume) : rawId ? formatTaskId(rawId) : ""; + const desc = typeof args.description === "string" ? args.description.trim() : ""; + if (idLabel || desc) { + let line = `${bullet} ${theme.fg("accent", theme.bold(idLabel || "agent"))}`; if (desc) { line += `: ${theme.fg("muted", truncateToWidth(replaceTabs(desc), 64))}`; } lines.push(line); } - if (truncated) { - lines.push(`${bullet} ${theme.fg("dim", formatMoreItems(items.length - cap, "agent"))}`); - } return lines; } -/** - * Build the shared-context section (the `# Goal / # Constraints` background - * passed to every subagent). Rendered in both the streaming call preview and - * the merged result frame so the brief stays visible for the whole task - * lifecycle — not just until the first progress snapshot replaces the call view. - */ -type TaskRenderSection = { lines: readonly string[] }; -type ContextSectionRenderer = (width: number) => TaskRenderSection; +/** One renderable frame section: optional label, body rows, leading divider. */ +type TaskRenderSection = { label?: string; lines: readonly string[]; separator?: boolean }; +type AssignmentSectionRenderer = (width: number) => TaskRenderSection; // Default output-block layout is: left border + one-cell content inset + right // border. Render markdown at that inner width so the output block does not need -// to rewrap already-rendered context lines. -const CONTEXT_FRAME_INSET = 3; +// to rewrap already-rendered assignment lines. +const ASSIGNMENT_FRAME_INSET = 3; -function contextMarkdownWidth(frameWidth: number): number { - return Math.max(1, frameWidth - CONTEXT_FRAME_INSET); -} - -function createContextSectionRenderer(args: TaskParams | undefined, theme: Theme): ContextSectionRenderer | undefined { +/** + * Build the assignment section (the markdown brief handed to the subagent). + * Rendered in both the streaming call preview and the result frame so the + * brief stays visible for the whole task lifecycle — not just until the first + * progress snapshot replaces the call view. + */ +function createAssignmentSectionRenderer( + args: Partial | undefined, + theme: Theme, +): AssignmentSectionRenderer | undefined { // `renderResult` receives the raw tool args (unlike `renderCall`, which is - // fed through `repairTaskParams`), so undo any per-field double-encoding here - // too. The repair is idempotent on already-clean text. - const context = repairDoubleEncodedJsonString(args?.context ?? "").trim(); - if (!context) return undefined; + // fed through `repairTaskParams`), so undo any per-field double-encoding + // here too. The repair is idempotent on already-clean text. + const assignment = repairDoubleEncodedJsonString(typeof args?.assignment === "string" ? args.assignment : "").trim(); + if (!assignment) return undefined; - const markdown = new Markdown(context, 0, 0, getMarkdownTheme(), { + const markdown = new Markdown(assignment, 0, 0, getMarkdownTheme(), { color: text => theme.fg("muted", text), }); - return width => ({ lines: markdown.render(contextMarkdownWidth(width)) }); + return width => ({ lines: markdown.render(Math.max(1, width - ASSIGNMENT_FRAME_INSET)) }); } /** @@ -575,22 +571,23 @@ export function renderCall( theme: Theme, ): Component { const showIsolated = "isolated" in args && args.isolated === true; - const header = renderStatusLine({ icon: "pending", title: "Task", description: args.agent }, theme); - const contextSectionRenderer = createContextSectionRenderer(args, theme); + const resume = typeof args.resume === "string" && args.resume.trim() ? args.resume.trim() : undefined; + const headerDescription = resume ? `resume ${formatTaskId(resume)}` : args.agent; + const header = renderStatusLine({ icon: "pending", title: "Task", description: headerDescription }, theme); + const assignmentSection = createAssignmentSectionRenderer(args, theme); return framedBlock(theme, width => { const sections: Array<{ label?: string; lines: readonly string[]; separator?: boolean }> = []; - if (contextSectionRenderer) sections.push(contextSectionRenderer(width)); - - // The per-task preview list only exists to surface dispatched agents while - // the call args stream in. Once a result snapshot exists, `renderResult` - // draws the same agents as progress/result lines, so showing the Tasks - // section here would just repeat the count the result frame already shows. + // The call preview only exists to surface the dispatched agent while the + // args stream in. Once a result snapshot exists, `renderResult` draws the + // same agent (and the assignment brief) itself, so showing it here would + // repeat what the result frame already shows. if (!options.renderContext?.hasResult) { sections.push({ separator: true, - lines: renderTaskItemLines(args.tasks, options.expanded, theme), + lines: renderTaskCallLines(args, theme), }); + if (assignmentSection) sections.push(assignmentSection(width)); } return { @@ -631,8 +628,12 @@ function renderAgentProgress( const titlePart = description ? `${theme.bold(displayId)}: ${description}` : displayId; const indent = prefix ? `${prefix} ` : ""; let statusLine: string; - if (progress.status === "running") { - const bullet = theme.styledSymbol("status.done", "text"); + if (progress.status === "running" || progress.status === "pending") { + // Live (or queued) agents shimmer their description so the row reads as + // in-flight even after the block freezes — the async spawn result keeps + // the agent on "pending" while the detached job runs. + const bullet = + progress.status === "running" ? theme.styledSymbol("status.done", "text") : theme.fg(iconColor, icon); const name = theme.fg("accent", description ? theme.bold(displayId) : displayId); statusLine = `${indent}${bullet} ${name}`; if (description) { @@ -945,6 +946,7 @@ function renderAgentResult( statusLine, { tokens: result.tokens, + requests: result.requests, contextTokens: result.contextTokens, contextWindow: result.contextWindow, cost: result.usage?.cost.total ?? 0, @@ -1073,9 +1075,11 @@ function renderAgentResult( } /** - * Order live progress entries so finished agents render first and unfinished - * (pending/running) ones stay pinned at the bottom as tasks complete. Stable - * within each group, so agents keep their dispatch order. + * Order live progress entries so finished agents render first — sorted by + * runtime ascending, matching {@link orderResultsForDisplay} — while + * unfinished (pending/running) ones stay pinned at the bottom in dispatch + * order. Because a finished agent's runtime is fixed, finalization renders + * the same order and rows never reshuffle. */ function orderProgressForDisplay(progress: readonly AgentProgress[]): AgentProgress[] { const finished: AgentProgress[] = []; @@ -1083,9 +1087,19 @@ function orderProgressForDisplay(progress: readonly AgentProgress[]): AgentProgr for (const p of progress) { (p.status === "pending" || p.status === "running" ? unfinished : finished).push(p); } + finished.sort((a, b) => a.durationMs - b.durationMs || a.index - b.index); return finished.concat(unfinished); } +/** + * Order finalized results by runtime ascending (tie-break: dispatch index) so + * the finalized list matches the live-progress order produced by + * {@link orderProgressForDisplay}. + */ +function orderResultsForDisplay(results: readonly SingleResult[]): SingleResult[] { + return [...results].sort((a, b) => a.durationMs - b.durationMs || a.index - b.index); +} + /** * Render the tool result. */ @@ -1097,25 +1111,27 @@ export function renderResult( ): Component { const fallbackText = result.content.find(c => c.type === "text")?.text ?? ""; const details = result.details; - const contextSectionRenderer = createContextSectionRenderer(args, theme); + const resumeLabel = + typeof args?.resume === "string" && args.resume.trim() ? `resume ${formatTaskId(args.resume.trim())}` : undefined; + const assignmentSection = createAssignmentSectionRenderer(args, theme); if (!details) { const text = result.content.find(c => c.type === "text")?.text || ""; const errored = result.isError === true; const header = errored - ? renderStatusLine({ icon: "error", title: "Task", description: args?.agent }, theme) + ? renderStatusLine({ icon: "error", title: "Task", description: resumeLabel ?? args?.agent }, theme) : renderStatusLine( { iconOverride: theme.styledSymbol("status.done", "accent"), title: "Task", - description: args?.agent, + description: resumeLabel ?? args?.agent, }, theme, ); return framedBlock(theme, width => ({ header, sections: [ - ...(contextSectionRenderer ? [contextSectionRenderer(width)] : []), + ...(assignmentSection ? [assignmentSection(width)] : []), ...(text ? [{ separator: true, lines: [theme.fg("dim", truncateToWidth(text, width))] }] : []), ], state: errored ? "error" : "success", @@ -1131,10 +1147,9 @@ export function renderResult( const isError = aborted || failed; const agentCount = hasResults ? details.results.length : (details.progress?.length ?? 0); const icon: ToolUIStatus = options.isPartial ? "running" : isError ? "error" : mergeFailed ? "warning" : "success"; - // Surface the dispatched agent type (e.g. `Reviewer`) alongside the count so - // the header reads `Task 16 agents: Reviewer`. All tasks in one call share a - // single `agent` type (top-level param), so one label covers the whole batch. - const agentName = args?.agent?.trim(); + // Surface the dispatched agent type (e.g. `Reviewer`) or the resumed agent + // id alongside the count so the header reads `Task 1 agent: Reviewer`. + const agentName = resumeLabel ?? args?.agent?.trim(); const countLabel = agentCount > 0 ? `${agentCount} ${agentCount === 1 ? "agent" : "agents"}` : undefined; const metaLabel = countLabel ? (agentName ? `${countLabel}: ${agentName}` : countLabel) : agentName; const header = renderStatusLine( @@ -1158,7 +1173,7 @@ export function renderResult( lines.push(...renderAgentProgress(progress, "", " ", expanded, theme, spinnerFrame)); }); } else if (details.results && details.results.length > 0) { - details.results.forEach(res => { + orderResultsForDisplay(details.results).forEach(res => { lines.push(...renderAgentResult(res, "", " ", expanded, theme)); }); @@ -1171,6 +1186,8 @@ export function renderResult( if (successCount > 0) summaryParts.push(theme.fg("success", `${successCount} succeeded`)); if (mergeFailedCount > 0) summaryParts.push(theme.fg("warning", `${mergeFailedCount} merge failed`)); if (failCount > 0) summaryParts.push(theme.fg("error", `${failCount} failed`)); + const totalRequests = details.results.reduce((sum, r) => sum + (r.requests ?? 0), 0); + if (totalRequests > 0) summaryParts.push(theme.fg("dim", `${formatNumber(totalRequests)} req`)); summaryParts.push(theme.fg("dim", formatDuration(details.totalDurationMs))); // Wrap the run summary in the theme's bracket glyphs (dim chrome, colored // counts) to match the bash tool's `[Wall: … | Exit: …]` footer. @@ -1189,7 +1206,7 @@ export function renderResult( return { header, sections: [ - ...(contextSectionRenderer ? [contextSectionRenderer(width)] : []), + ...(assignmentSection ? [assignmentSection(width)] : []), { separator: true, lines: [theme.fg("dim", truncateToWidth(text, width))] }, ], state, @@ -1219,7 +1236,7 @@ export function renderResult( return { header, sections: [ - ...(contextSectionRenderer ? [contextSectionRenderer(width)] : []), + ...(assignmentSection ? [assignmentSection(width)] : []), ...(lines.length > 0 ? [{ separator: true, lines }] : []), ], state, @@ -1252,8 +1269,9 @@ function renderNestedTaskResults(detailsList: TaskToolDetails[], expanded: boole const lines: string[] = []; for (const details of detailsList) { if (!details.results || details.results.length === 0) continue; - details.results.forEach((result, index) => { - const { prefix, continuePrefix } = nestedMarkers(index === details.results.length - 1, theme); + const ordered = orderResultsForDisplay(details.results); + ordered.forEach((result, index) => { + const { prefix, continuePrefix } = nestedMarkers(index === ordered.length - 1, theme); lines.push(...renderAgentResult(result, prefix, continuePrefix, expanded, theme)); }); } @@ -1275,8 +1293,9 @@ function renderNestedTaskTree( for (const details of detailsList) { const hasResults = Boolean(details.results && details.results.length > 0); if (hasResults) { - details.results.forEach((result, index) => { - const { prefix, continuePrefix } = nestedMarkers(index === details.results.length - 1, theme); + const ordered = orderResultsForDisplay(details.results); + ordered.forEach((result, index) => { + const { prefix, continuePrefix } = nestedMarkers(index === ordered.length - 1, theme); lines.push(...renderAgentResult(result, prefix, continuePrefix, expanded, theme)); }); continue; diff --git a/packages/coding-agent/src/task/repair-args.ts b/packages/coding-agent/src/task/repair-args.ts index 62cc8f7a7..ec52a0000 100644 --- a/packages/coding-agent/src/task/repair-args.ts +++ b/packages/coding-agent/src/task/repair-args.ts @@ -2,7 +2,7 @@ * Repair double-encoded JSON string arguments for the task tool. * * Models occasionally JSON-escape a string value twice when emitting a - * `task` tool call, so a `context`/`assignment` that should read + * `task` tool call, so an `assignment` that should read * * # Role * You are a judge … "describe this" … return — @@ -24,11 +24,11 @@ * string. * * This is deliberately scoped to the task tool's natural-language fields - * (`context`, `assignment`, `description`). It is NOT applied to code-bearing + * (`assignment`, `description`). It is NOT applied to code-bearing * tools (write/edit/bash/search), where a backslash or quote is load-bearing * and a false-positive unescape would silently corrupt a file or command. */ -import type { TaskItem, TaskParams } from "./types"; +import type { TaskParams } from "./types"; /** A backslash that escapes a structural char — `\"`, `\\`, `\/`, or `\uXXXX`. */ const STRUCTURAL_ESCAPE = /\\(?:["\\/]|u[0-9a-fA-F]{4})/; @@ -78,40 +78,21 @@ export function repairDoubleEncodedJsonString(value: string): string { return typeof decoded === "string" && decoded !== value ? decoded : value; } -/** Repair a single (possibly partial) task item's prose fields. */ -function repairTaskItem(task: TaskItem): TaskItem { - if (task === null || typeof task !== "object") return task; - const assignment = - typeof task.assignment === "string" ? repairDoubleEncodedJsonString(task.assignment) : task.assignment; - const description = - typeof task.description === "string" ? repairDoubleEncodedJsonString(task.description) : task.description; - if (assignment === task.assignment && description === task.description) return task; - return { ...task, assignment, description }; -} - /** - * Repair double-encoded prose in task-tool params (`context` and each task's - * `assignment`/`description`). Returns the same reference when nothing changed - * so callers can cheaply skip work. Defensive against partially-streamed args - * (missing/undefined fields, partial task arrays) so it is safe on the render - * path as well as on execution. + * Repair double-encoded prose in task-tool params (`assignment` and + * `description`). Returns the same reference when nothing changed so callers + * can cheaply skip work. Defensive against partially-streamed args + * (missing/undefined fields) so it is safe on the render path as well as on + * execution. */ export function repairTaskParams(params: TaskParams): TaskParams { if (params === null || typeof params !== "object") return params; - const context = typeof params.context === "string" ? repairDoubleEncodedJsonString(params.context) : params.context; + const assignment = + typeof params.assignment === "string" ? repairDoubleEncodedJsonString(params.assignment) : params.assignment; + const description = + typeof params.description === "string" ? repairDoubleEncodedJsonString(params.description) : params.description; - let tasks = params.tasks; - if (Array.isArray(params.tasks)) { - let changed = false; - const repaired = params.tasks.map(task => { - const next = repairTaskItem(task); - if (next !== task) changed = true; - return next; - }); - if (changed) tasks = repaired; - } - - if (context === params.context && tasks === params.tasks) return params; - return { ...params, context, tasks }; + if (assignment === params.assignment && description === params.description) return params; + return { ...params, assignment, description }; } diff --git a/packages/coding-agent/src/task/simple-mode.ts b/packages/coding-agent/src/task/simple-mode.ts index fc4269f13..22c3c607d 100644 --- a/packages/coding-agent/src/task/simple-mode.ts +++ b/packages/coding-agent/src/task/simple-mode.ts @@ -3,21 +3,17 @@ export const TASK_SIMPLE_MODES = ["default", "schema-free", "independent"] as co export type TaskSimpleMode = (typeof TASK_SIMPLE_MODES)[number]; interface TaskSimpleModeCapabilities { - contextEnabled: boolean; customSchemaEnabled: boolean; } const TASK_SIMPLE_MODE_CAPABILITIES: Record = { default: { - contextEnabled: true, customSchemaEnabled: true, }, "schema-free": { - contextEnabled: true, customSchemaEnabled: false, }, independent: { - contextEnabled: false, customSchemaEnabled: false, }, }; diff --git a/packages/coding-agent/src/task/types.ts b/packages/coding-agent/src/task/types.ts index 5b535e649..01b84a1e5 100644 --- a/packages/coding-agent/src/task/types.ts +++ b/packages/coding-agent/src/task/types.ts @@ -66,34 +66,16 @@ export interface SubagentLifecyclePayload { index: number; } -const assignmentDescription = "per-task instructions; self-contained"; - -const createTaskItemSchema = (_contextEnabled: boolean) => - z.object({ - id: z.string().max(48).describe("camelcase identifier"), - description: z.string().describe("ui label, not seen by subagent"), - assignment: z.string().describe(assignmentDescription), - }); - -/** Single task item for parallel execution (default shape with context enabled). */ -export const taskItemSchema = createTaskItemSchema(true); -export type TaskItem = z.infer; - -const createTaskSchema = (options: { isolationEnabled: boolean; simpleMode: TaskSimpleMode }) => { - const { contextEnabled, customSchemaEnabled } = getTaskSimpleModeCapabilities(options.simpleMode); - const itemSchema = createTaskItemSchema(contextEnabled); - +const createTaskSchema = (options: { isolationEnabled: boolean; customSchemaEnabled: boolean }) => { let schema = z.object({ - agent: z.string().describe("agent type"), - tasks: z.array(itemSchema).describe("tasks to execute in parallel"), + agent: z.string().optional().describe("agent type; omit when resume is set"), + id: z.string().max(48).optional().describe("stable agent id; default generated"), + description: z.string().optional().describe("ui label, not seen by subagent"), + assignment: z.string().describe("the work; self-contained instructions"), + resume: z.string().optional().describe("existing agent id: revive and continue instead of spawning"), }); - if (contextEnabled) { - schema = schema.extend({ - context: z.string().optional().describe("shared background prepended to each assignment"), - }); - } - if (customSchemaEnabled) { + if (options.customSchemaEnabled) { schema = schema.extend({ schema: z.string().optional().describe("jtd schema for expected response shape"), }); @@ -108,19 +90,15 @@ const createTaskSchema = (options: { isolationEnabled: boolean; simpleMode: Task return schema; }; -export const taskSchema = createTaskSchema({ isolationEnabled: true, simpleMode: "default" }); -export const taskSchemaNoIsolation = createTaskSchema({ isolationEnabled: false, simpleMode: "default" }); -const taskSchemaSchemaFree = createTaskSchema({ isolationEnabled: true, simpleMode: "schema-free" }); -const taskSchemaSchemaFreeNoIsolation = createTaskSchema({ isolationEnabled: false, simpleMode: "schema-free" }); -const taskSchemaIndependent = createTaskSchema({ isolationEnabled: true, simpleMode: "independent" }); -const taskSchemaIndependentNoIsolation = createTaskSchema({ isolationEnabled: false, simpleMode: "independent" }); +export const taskSchema = createTaskSchema({ isolationEnabled: true, customSchemaEnabled: true }); +const taskSchemaNoIsolation = createTaskSchema({ isolationEnabled: false, customSchemaEnabled: true }); +const taskSchemaSchemaFree = createTaskSchema({ isolationEnabled: true, customSchemaEnabled: false }); +const taskSchemaSchemaFreeNoIsolation = createTaskSchema({ isolationEnabled: false, customSchemaEnabled: false }); const ALL_TASK_SCHEMAS = [ taskSchema, taskSchemaNoIsolation, taskSchemaSchemaFree, taskSchemaSchemaFreeNoIsolation, - taskSchemaIndependent, - taskSchemaIndependentNoIsolation, ] as const; type DynamicTaskSchema = (typeof ALL_TASK_SCHEMAS)[number]; @@ -129,22 +107,28 @@ export type TaskSchema = typeof taskSchema; export type TaskToolSchemaInstance = DynamicTaskSchema; export function getTaskSchema(options: { isolationEnabled: boolean; simpleMode: TaskSimpleMode }): DynamicTaskSchema { - switch (options.simpleMode) { - case "schema-free": - return options.isolationEnabled ? taskSchemaSchemaFree : taskSchemaSchemaFreeNoIsolation; - case "independent": - return options.isolationEnabled ? taskSchemaIndependent : taskSchemaIndependentNoIsolation; - default: - return options.isolationEnabled ? taskSchema : taskSchemaNoIsolation; + const { customSchemaEnabled } = getTaskSimpleModeCapabilities(options.simpleMode); + if (customSchemaEnabled) { + return options.isolationEnabled ? taskSchema : taskSchemaNoIsolation; } + return options.isolationEnabled ? taskSchemaSchemaFree : taskSchemaSchemaFreeNoIsolation; } export interface TaskParams { - agent: string; - context?: string; + /** Agent type; required unless `resume` is set. */ + agent?: string; + /** Stable agent id; default = generated AdjectiveNoun. */ + id?: string; + /** UI label, not seen by the subagent. */ + description?: string; + /** The work; required. */ + assignment?: string; + /** JTD schema for the expected yield shape; unchanged semantics. */ schema?: string; - tasks: TaskItem[]; + /** Run in an isolated worktree; isolated agents are NOT resumable. */ isolated?: boolean; + /** Existing agent id: revive + follow-up instead of spawn. */ + resume?: string; } /** A code review finding reported by the reviewer agent */ @@ -206,6 +190,8 @@ export interface AgentProgress { recentTools: Array<{ tool: string; args: string; endMs: number }>; recentOutput: string[]; toolCount: number; + /** Count of assistant requests (assistant message_end events) across the run. Drives the soft request budget guard. */ + requests: number; /** Cumulative input + output + cacheWrite tokens across all turns. Excludes cacheRead (re-reads cached context every turn, making cumulative sum misleading). */ tokens: number; /** @@ -276,6 +262,8 @@ export interface SingleResult { durationMs: number; /** Cumulative input + output + cacheWrite tokens across all turns. Excludes cacheRead (re-reads cached context every turn, making cumulative sum misleading). */ tokens: number; + /** Count of assistant requests (assistant message_end events) across the run. */ + requests: number; /** Latest per-turn context size at task completion. See `AgentProgress.contextTokens`. */ contextTokens?: number; /** Model's context window in tokens, when known. */ diff --git a/packages/coding-agent/test/eval/agent-bridge.test.ts b/packages/coding-agent/test/eval/agent-bridge.test.ts index 832dd7fba..492087753 100644 --- a/packages/coding-agent/test/eval/agent-bridge.test.ts +++ b/packages/coding-agent/test/eval/agent-bridge.test.ts @@ -21,6 +21,7 @@ function createResult(): SingleResult { truncated: false, durationMs: 1, tokens: 0, + requests: 0, }; } diff --git a/packages/coding-agent/test/rpc-subagents.test.ts b/packages/coding-agent/test/rpc-subagents.test.ts index 65b89a31b..133373bdd 100644 --- a/packages/coding-agent/test/rpc-subagents.test.ts +++ b/packages/coding-agent/test/rpc-subagents.test.ts @@ -43,6 +43,7 @@ function createProgress(overrides: Partial = {}): AgentProgress { recentTools: [], recentOutput: [], toolCount: 0, + requests: 0, tokens: 0, cost: 0, durationMs: 0, diff --git a/packages/coding-agent/test/streaming-preview-height.test.ts b/packages/coding-agent/test/streaming-preview-height.test.ts index 7bfae7adf..acf37be34 100644 --- a/packages/coding-agent/test/streaming-preview-height.test.ts +++ b/packages/coding-agent/test/streaming-preview-height.test.ts @@ -372,19 +372,26 @@ describe("streaming tool call preview height (bounded across renderers)", () => } }, 30_000); - test("task pending preview preserves full multiline context", () => { + test("task pending preview stays bounded with a long multiline assignment", () => { + // CONTRACT CHANGE with the single-spawn task rework: the old uncapped + // multi-task `context` rendering is gone with the field. The pending + // preview now intentionally bounds the assignment (first line + a + // "more lines" marker when collapsed; 12 lines when expanded), like + // bash/ssh, so a long assignment can no longer strand the block top. const longLines = Array.from({ length: 80 }, (_, i) => `line-${i}`); const { lines, text } = renderPending("task", { agent: "task", - context: longLines.join("\n"), - tasks: [{ id: "alpha", description: "preview" }], + id: "alpha", + description: "preview", + assignment: longLines.join("\n"), }); - expect(lines.length, "task preview should not be capped").toBeGreaterThan(80); + expect(lines.length, "task preview should stay bounded").toBeLessThan(20); + expect(text).toContain("preview"); expect(text).toContain("line-0"); - expect(text).toContain("line-40"); - expect(text).toContain("line-79"); - expect(text).not.toMatch(/more lines/); + expect(text).not.toContain("line-40"); + expect(text).not.toContain("line-79"); + expect(text, "task preview should advertise truncation").toMatch(/more lines/); }); test("eval pending preview preserves full code (never collapsed)", () => { diff --git a/packages/coding-agent/test/task/executor-subagent-reminders.test.ts b/packages/coding-agent/test/task/executor-subagent-reminders.test.ts index 1b308036a..7f69af243 100644 --- a/packages/coding-agent/test/task/executor-subagent-reminders.test.ts +++ b/packages/coding-agent/test/task/executor-subagent-reminders.test.ts @@ -197,7 +197,7 @@ describe("runSubprocess yield reminders", () => { expect(createAgentSessionSpy).toHaveBeenCalledTimes(1); }); - it("renders shared task context in subagent system prompt before now", async () => { + it("splices the subagent role prompt before the trailing system section", async () => { let userPrompt = ""; const session = createMockSession(({ text, emit }) => { userPrompt = text; @@ -217,7 +217,6 @@ describe("runSubprocess yield reminders", () => { await runSubprocess({ ...baseOptions, id: "subagent-context-system", - context: "Shared task background", task: "Your assignment is below.\nBe thorough and complete fully before yielding.\n\nDo the task.", }); @@ -229,11 +228,12 @@ describe("runSubprocess yield reminders", () => { expect(systemPrompt).toHaveLength(4); expect(systemPrompt?.[0]).toBe("system"); expect(systemPrompt?.[1]).toBe("project"); - expect(systemPrompt?.[2]).toMatch(/CONTEXT\n=+\n\nShared task background/); expect(systemPrompt?.[2]).toMatch(/ROLE\n=+\n\ntest/); + // The parent-conversation CONTEXT section is gone: subagents get their + // background inside the assignment (or a local:// file), never a dump. + expect(systemPrompt?.[2]).not.toMatch(/CONTEXT\n=+/); expect(systemPrompt?.[3]).toBe("now"); expect(userPrompt).not.toMatch(/CONTEXT\n=+/); - expect(userPrompt).not.toContain("Shared task background"); }); it("sends reminder prompt when subagent stops without yield", async () => { @@ -586,7 +586,7 @@ describe("runSubprocess yield reminders", () => { expect(result.aborted).toBe(true); expect(errorSpy).not.toHaveBeenCalledWith("Subagent prompt failed", expect.anything()); - expect(debugSpy).toHaveBeenCalledWith("Subagent prompt aborted", expect.anything()); + expect(debugSpy).toHaveBeenCalledWith("Subagent prompt aborted"); }); }); diff --git a/packages/coding-agent/test/task/output-manager.test.ts b/packages/coding-agent/test/task/output-manager.test.ts index 50fef2440..b374105fe 100644 --- a/packages/coding-agent/test/task/output-manager.test.ts +++ b/packages/coding-agent/test/task/output-manager.test.ts @@ -19,10 +19,14 @@ describe("AgentOutputManager", () => { expect(await mgr.allocate("Bob")).toBe("Bob"); }); - it("de-duplicates within a batch while preserving order", async () => { + it("de-duplicates repeated names while preserving order", async () => { const mgr = new AgentOutputManager(() => null); - expect(await mgr.allocateBatch(["Auth", "Auth", "Api", "Auth"])).toEqual(["Auth", "Auth-2", "Api", "Auth-3"]); + const ids: string[] = []; + for (const name of ["Auth", "Auth", "Api", "Auth"]) { + ids.push(await mgr.allocate(name)); + } + expect(ids).toEqual(["Auth", "Auth-2", "Api", "Auth-3"]); }); it("nests ids under a parent prefix and still suffixes repeats", async () => { diff --git a/packages/coding-agent/test/task/render-call.test.ts b/packages/coding-agent/test/task/render-call.test.ts index 4ea12d880..f666239de 100644 --- a/packages/coding-agent/test/task/render-call.test.ts +++ b/packages/coding-agent/test/task/render-call.test.ts @@ -1,7 +1,7 @@ import { afterAll, beforeAll, describe, expect, it } from "bun:test"; import { resetSettingsForTest, Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; import { getThemeByName, setThemeInstance, type Theme } from "@oh-my-pi/pi-coding-agent/modes/theme/theme"; -import type { TaskParams, TaskToolDetails } from "@oh-my-pi/pi-coding-agent/task"; +import type { TaskParams } from "@oh-my-pi/pi-coding-agent/task"; import { taskToolRenderer } from "@oh-my-pi/pi-coding-agent/task/render"; describe("task renderer: streaming call preview", () => { @@ -25,100 +25,62 @@ describe("task renderer: streaming call preview", () => { return Bun.stripANSI(component.render(160).join("\n")); } - function renderCompleted(args: TaskParams): string { - const details: TaskToolDetails = { - projectAgentsDir: null, - totalDurationMs: 12, - results: [ - { - index: 0, - id: "Only", - agent: args.agent, - agentSource: "bundled", - task: "Render the shared context", - exitCode: 0, - output: "Done.", - stderr: "", - truncated: false, - durationMs: 12, - tokens: 1, - }, - ], - }; - const component = taskToolRenderer.renderResult( - { content: [{ type: "text", text: "1 agent completed." }], details }, - { expanded: false, isPartial: false }, - theme, - args, - ); - return Bun.stripANSI(component.render(160).join("\n")); - } - - // The preview must surface each agent's id + ui description so the user can - // see which agents are being dispatched, not a bare "N agents" count. - it("lists each task's id and description instead of only a count", () => { + // The preview must surface the agent id + ui description so the user can + // see what is being dispatched while args stream in. + it("shows the agent id, description, and assignment preview", () => { const args: TaskParams = { agent: "reviewer", - tasks: [ - { id: "ReviewAuth", description: "Audit the auth module", assignment: "..." }, - { id: "ReviewDb", description: "Audit the db layer", assignment: "..." }, - ], + id: "ReviewAuth", + description: "Audit the auth module", + assignment: "Review packages/server/src/auth for missing 401 handling.\nReport findings.", }; const out = render(args); + expect(out).toContain("reviewer"); expect(out).toContain("ReviewAuth"); expect(out).toContain("Audit the auth module"); - expect(out).toContain("ReviewDb"); - expect(out).toContain("Audit the db layer"); - // The per-task list stands on its own — neither the redundant "Tasks (N)" - // section label nor the old flat "N agents" line is drawn. - expect(out).not.toContain("Tasks ("); - expect(out).not.toContain("2 agents"); + expect(out).toContain("Review packages/server/src/auth for missing 401 handling."); }); - it("renders a partially-streamed entry without a description and missing trailing entry", () => { + it("renders partially-streamed args without crashing", () => { const args = { agent: "task", - // Trailing entry mimics streaming JSON: id arrived, description not yet, - // plus a not-yet-materialized slot. - tasks: [{ id: "First", description: "Do the first thing", assignment: "..." }, { id: "Second" }, undefined], + id: "First", + // description/assignment not yet arrived. } as unknown as TaskParams; const out = render(args); expect(out).toContain("First"); - expect(out).toContain("Do the first thing"); - expect(out).toContain("Second"); - // Missing-id slot falls back to a positional placeholder rather than crashing. - expect(out).toContain("#3"); - expect(out).not.toContain("Tasks ("); + expect(out).toContain("task"); }); - it("caps the collapsed list and reports the overflow as agents", () => { - const tasks = Array.from({ length: 15 }, (_, i) => ({ - id: `Agent${i + 1}`, - description: `Task ${i + 1}`, - assignment: "...", - })); - const args: TaskParams = { agent: "task", tasks }; + it("always renders the full assignment markdown, collapsed or expanded", () => { + const assignmentLines = Array.from({ length: 6 }, (_, i) => `Step ${i + 1}: do the thing.`); + const args: TaskParams = { + agent: "task", + id: "Worker", + assignment: assignmentLines.join("\n"), + }; + // The assignment is the brief handed to the subagent; it renders as + // markdown in full regardless of the expanded toggle. const collapsed = render(args, false); - expect(collapsed).toContain("Agent1"); - expect(collapsed).toContain("Agent12"); - expect(collapsed).not.toContain("Agent13"); - expect(collapsed).toContain("3 more agents"); + expect(collapsed).toContain("Step 1"); + expect(collapsed).toContain("Step 6"); const expanded = render(args, true); - expect(expanded).toContain("Agent13"); - expect(expanded).toContain("Agent15"); - expect(expanded).not.toContain("more agents"); + expect(expanded).toContain("Step 1"); + expect(expanded).toContain("Step 6"); }); it("surfaces the isolation flag in the header bar", () => { const args: TaskParams = { agent: "task", isolated: true, - tasks: [{ id: "Only", description: "Single task", assignment: "..." }], + id: "Only", + description: "Single task", + assignment: "...", }; const out = render(args); const lines = out.split("\n"); @@ -129,34 +91,28 @@ describe("task renderer: streaming call preview", () => { expect(lines[0]).toContain("isolated"); }); - it("renders shared context as markdown in call and result frames", () => { + it("labels resume calls with the resumed agent id", () => { const args: TaskParams = { - agent: "task", - context: ["# Goal", "Fix **rendering**.", "", "# Constraints", "- Keep `task` visible"].join("\n"), - tasks: [{ id: "Only", description: "Single task", assignment: "..." }], + resume: "AuthLoader", + assignment: "Also check the refresh-token path.", }; + const out = render(args); + const lines = out.split("\n"); - for (const out of [render(args), renderCompleted(args)]) { - expect(out).toContain("Goal"); - expect(out).toContain("Fix rendering."); - expect(out).toContain("Constraints"); - expect(out).toContain("Keep task visible"); - expect(out).not.toContain("# Goal"); - expect(out).not.toContain("# Constraints"); - } + expect(lines[0]).toContain("resume AuthLoader"); + expect(out).toContain("Also check the refresh-token path."); }); // Once the tool produces a result, the container suppresses the call entirely - // via `mergeCallAndResult` and `renderResult` draws each agent. As a safety - // net, `renderCall` also drops its duplicate per-task preview when a result - // snapshot is present, so the two never stack. - it("drops the per-task preview list once a result snapshot exists", () => { + // via `mergeCallAndResult` and `renderResult` draws the agent. As a safety + // net, `renderCall` also drops its preview when a result snapshot is present, + // so the two never stack. + it("drops the preview once a result snapshot exists", () => { const args: TaskParams = { agent: "reviewer", - tasks: [ - { id: "ReviewAuth", description: "Audit the auth module", assignment: "..." }, - { id: "ReviewDb", description: "Audit the db layer", assignment: "..." }, - ], + id: "ReviewAuth", + description: "Audit the auth module", + assignment: "Review the auth module.", }; const component = taskToolRenderer.renderCall( args, @@ -166,7 +122,6 @@ describe("task renderer: streaming call preview", () => { const out = Bun.stripANSI(component.render(160).join("\n")); expect(out).not.toContain("Audit the auth module"); - expect(out).not.toContain("Audit the db layer"); - expect(out).not.toContain("Tasks ("); + expect(out).not.toContain("Review the auth module."); }); }); diff --git a/packages/coding-agent/test/task/render-nested-live.test.ts b/packages/coding-agent/test/task/render-nested-live.test.ts index 422ce07e1..1555476b1 100644 --- a/packages/coding-agent/test/task/render-nested-live.test.ts +++ b/packages/coding-agent/test/task/render-nested-live.test.ts @@ -1,7 +1,7 @@ import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from "bun:test"; import { resetSettingsForTest, Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; import { getThemeByName, setThemeInstance } from "@oh-my-pi/pi-coding-agent/modes/theme/theme"; -import type { AgentProgress, SingleResult, TaskParams, TaskToolDetails } from "@oh-my-pi/pi-coding-agent/task"; +import type { AgentProgress, SingleResult, TaskToolDetails } from "@oh-my-pi/pi-coding-agent/task"; import { taskToolRenderer } from "@oh-my-pi/pi-coding-agent/task/render"; import { formatDuration, formatNumber } from "@oh-my-pi/pi-utils"; @@ -41,6 +41,7 @@ describe("task renderer: nested live rendering", () => { recentTools: [], recentOutput: [], toolCount: 1, + requests: 0, tokens: 1000, cost: 0, durationMs: 1234, @@ -63,6 +64,7 @@ describe("task renderer: nested live rendering", () => { truncated: false, durationMs: 500, tokens: 200, + requests: 0, }; } @@ -79,6 +81,7 @@ describe("task renderer: nested live rendering", () => { recentTools: [], recentOutput: [], toolCount: 0, + requests: 0, tokens: 0, cost: 0, durationMs: 0, @@ -209,32 +212,6 @@ describe("task renderer: nested live rendering", () => { expect(text).not.toContain("Σ"); }); - it("keeps the shared context visible while the task is in progress", async () => { - const theme = (await getThemeByName("dark"))!; - const details: TaskToolDetails = { - projectAgentsDir: null, - results: [], - totalDurationMs: 0, - progress: [makeRunningProgress({ id: "Probe", description: "Investigate padding" })], - }; - const args = { - agent: "task", - context: "# Goal\nHarden the auth stack before the cut.", - tasks: [], - } as unknown as TaskParams; - const component = taskToolRenderer.renderResult( - { content: [{ type: "text", text: "Running 1 agents..." }], details }, - { expanded: false, isPartial: true, spinnerFrame: 0 }, - theme, - args, - ); - const text = Bun.stripANSI(component.render(160).join("\n")); - // The brief no longer vanishes the moment the first progress snapshot - // replaces the streaming call view. - expect(text).toContain("Goal"); - expect(text).toContain("Harden the auth stack before the cut."); - }); - it("renders a static result header while the body shimmers the running task name", async () => { const theme = (await getThemeByName("dark"))!; const details: TaskToolDetails = { diff --git a/packages/coding-agent/test/task/render-yield-shape.test.ts b/packages/coding-agent/test/task/render-yield-shape.test.ts index f59dd4058..f6f94113d 100644 --- a/packages/coding-agent/test/task/render-yield-shape.test.ts +++ b/packages/coding-agent/test/task/render-yield-shape.test.ts @@ -46,6 +46,7 @@ describe("task renderer: malformed yield slot (#1987)", () => { truncated: false, durationMs: 250, tokens: 100, + requests: 0, // Cast deliberately: production typings declare `unknown[]`, but the // renderer must defend against a stray non-array value — that's // exactly what this regression test exercises. @@ -66,6 +67,7 @@ describe("task renderer: malformed yield slot (#1987)", () => { recentTools: [], recentOutput: [], toolCount: 1, + requests: 0, tokens: 100, cost: 0, durationMs: 250, diff --git a/packages/coding-agent/test/task/subagent-lsp.test.ts b/packages/coding-agent/test/task/subagent-lsp.test.ts index f123cd5e9..42c0edc01 100644 --- a/packages/coding-agent/test/task/subagent-lsp.test.ts +++ b/packages/coding-agent/test/task/subagent-lsp.test.ts @@ -18,7 +18,9 @@ import { EventBus } from "@oh-my-pi/pi-coding-agent/utils/event-bus"; const TEST_TASK: TaskParams = { agent: "task", - tasks: [{ id: "CheckLsp", description: "Check LSP availability", assignment: "Inspect LSP tools." }], + id: "CheckLsp", + description: "Check LSP availability", + assignment: "Inspect LSP tools.", }; function createAssistantStopMessage(text: string): AssistantMessage { diff --git a/packages/coding-agent/test/task/task-guards.test.ts b/packages/coding-agent/test/task/task-guards.test.ts new file mode 100644 index 000000000..37134a2ec --- /dev/null +++ b/packages/coding-agent/test/task/task-guards.test.ts @@ -0,0 +1,273 @@ +import { afterEach, describe, expect, it, vi } from "bun:test"; +import type { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry"; +import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; +import type { LoadExtensionsResult } from "@oh-my-pi/pi-coding-agent/extensibility/extensions/types"; +import type { CreateAgentSessionResult } from "@oh-my-pi/pi-coding-agent/sdk"; +import * as sdkModule from "@oh-my-pi/pi-coding-agent/sdk"; +import type { AgentSession, AgentSessionEvent } from "@oh-my-pi/pi-coding-agent/session/agent-session"; +import { formatResultOutputFallback } from "@oh-my-pi/pi-coding-agent/task"; +import { runSubprocess } from "@oh-my-pi/pi-coding-agent/task/executor"; +import type { AgentDefinition } from "@oh-my-pi/pi-coding-agent/task/types"; +import { EventBus } from "@oh-my-pi/pi-coding-agent/utils/event-bus"; + +/** + * Contract: runaway-subagent guards. + * + * 1. The executor counts assistant requests (message_end events) and surfaces + * the count on `SingleResult.requests`. + * 2. Crossing the soft request budget injects exactly ONE steering notice into + * the child session asking it to wrap up; crossing 1.5x the budget aborts + * the run gracefully. + * 3. A cancelled/aborted child that produced no completed output salvages its + * last assistant text into a `[cancelled after N req, …]` summary instead + * of the parent seeing "(no output)" and redoing the work. + */ + +interface SteerCall { + content: string; + options?: { deliverAs?: "steer" | "followUp" }; +} + +interface FakeSessionConfig { + /** Events pushed to the executor's subscriber on the next microtask. */ + events?: AgentSessionEvent[]; + /** When true, prompt/waitForIdle hang until abort() is called. */ + hang?: boolean; + /** Returned from getLastAssistantMessage (salvage source). */ + lastAssistantMessage?: unknown; +} + +interface FakeSessionHandle { + session: AgentSession; + steerCalls: SteerCall[]; + abortCalls: () => number; +} + +function assistantMessageEnd(text: string, usage?: Record): AgentSessionEvent { + return { + type: "message_end", + message: { + role: "assistant", + content: text ? [{ type: "text", text }] : [], + usage: usage ?? { input: 10, output: 5, cacheRead: 0, cacheWrite: 0, totalTokens: 15 }, + }, + } as unknown as AgentSessionEvent; +} + +function yieldToolEnd(): AgentSessionEvent { + return { + type: "tool_execution_end", + toolCallId: "tool-yield", + toolName: "yield", + result: { + content: [{ type: "text", text: "Result submitted." }], + details: { status: "success", data: { ok: true } }, + }, + isError: false, + } as AgentSessionEvent; +} + +function createFakeSession(config: FakeSessionConfig = {}): FakeSessionHandle { + let abortCount = 0; + const steerCalls: SteerCall[] = []; + const { promise: hang, resolve: releaseHang } = Promise.withResolvers(); + if (!config.hang) releaseHang(); + + const session: Partial = { + state: { messages: [] } as never, + agent: { state: { systemPrompt: ["test"] } } as never, + extensionRunner: undefined as never, + sessionManager: { appendSessionInit: () => {} } as never, + getActiveToolNames: () => ["read", "yield"], + setActiveToolsByName: async (_names: string[]) => {}, + subscribe: (listener: (event: AgentSessionEvent) => void) => { + if (config.events?.length) { + const events = config.events; + queueMicrotask(() => { + for (const event of events) listener(event); + }); + } + return () => {}; + }, + prompt: async () => { + await hang; + return true; + }, + waitForIdle: async () => { + await hang; + }, + sendUserMessage: async (content, options) => { + steerCalls.push({ content: String(content), options }); + }, + getLastAssistantMessage: () => (config.lastAssistantMessage ?? undefined) as never, + abort: async () => { + abortCount += 1; + releaseHang(); + }, + dispose: async () => {}, + }; + return { + session: session as AgentSession, + steerCalls, + abortCalls: () => abortCount, + }; +} + +function mockCreateAgentSession(session: AgentSession) { + return vi.spyOn(sdkModule, "createAgentSession").mockResolvedValue({ + session, + extensionsResult: {} as unknown as LoadExtensionsResult, + setToolUIContext: () => {}, + eventBus: new EventBus(), + } satisfies CreateAgentSessionResult); +} + +const baseAgent: AgentDefinition = { + name: "task", + description: "test", + systemPrompt: "test", + source: "bundled", +}; + +const baseOptions = { + cwd: "/tmp", + agent: baseAgent, + task: "do work", + index: 0, + id: "subagent-guards", + modelRegistry: { refresh: async () => {} } as unknown as ModelRegistry, + enableLsp: false, +}; + +describe("runSubprocess request guards", () => { + afterEach(() => { + vi.restoreAllMocks(); + }); + + it("counts assistant requests into SingleResult.requests", async () => { + const settings = Settings.isolated({ "task.maxRuntimeMs": 0 }); + const handle = createFakeSession({ + events: [ + assistantMessageEnd("step one"), + assistantMessageEnd("step two"), + assistantMessageEnd("step three"), + yieldToolEnd(), + ], + }); + mockCreateAgentSession(handle.session); + + const result = await runSubprocess({ ...baseOptions, id: "subagent-requests", settings }); + + expect(result.aborted).toBe(false); + expect(result.requests).toBe(3); + // Well under any budget: no steer injected. + expect(handle.steerCalls.length).toBe(0); + }); + + it("injects exactly one steering notice when the soft budget is crossed", async () => { + // Budget 4: steer fires at request 4 and must not repeat at request 5 + // (still below the 1.5x hard stop of 6). + const settings = Settings.isolated({ "task.maxRuntimeMs": 0, "task.softRequestBudget": 4 }); + const handle = createFakeSession({ + events: [ + assistantMessageEnd("1"), + assistantMessageEnd("2"), + assistantMessageEnd("3"), + assistantMessageEnd("4"), + assistantMessageEnd("5"), + yieldToolEnd(), + ], + }); + mockCreateAgentSession(handle.session); + + const result = await runSubprocess({ ...baseOptions, id: "subagent-steer", settings }); + + expect(result.requests).toBe(5); + expect(result.aborted).toBe(false); + expect(handle.steerCalls.length).toBe(1); + expect(handle.steerCalls[0].content).toContain("[budget notice]"); + expect(handle.steerCalls[0].content).toContain("4 requests"); + expect(handle.steerCalls[0].options?.deliverAs).toBe("steer"); + }); + + it("aborts the run gracefully at 1.5x the soft budget", async () => { + // Budget 2: steer at 2, hard stop at 3. The session hangs so only the + // budget abort can release it. + const settings = Settings.isolated({ "task.maxRuntimeMs": 0, "task.softRequestBudget": 2 }); + const handle = createFakeSession({ + hang: true, + events: [ + assistantMessageEnd("", { input: 10, output: 5, cacheRead: 0, cacheWrite: 0, totalTokens: 15 }), + assistantMessageEnd("", { input: 10, output: 5, cacheRead: 0, cacheWrite: 0, totalTokens: 15 }), + assistantMessageEnd("", { input: 10, output: 5, cacheRead: 0, cacheWrite: 0, totalTokens: 15 }), + ], + }); + mockCreateAgentSession(handle.session); + + const result = await runSubprocess({ ...baseOptions, id: "subagent-hard-stop", settings }); + + expect(result.aborted).toBe(true); + expect(result.exitCode).toBe(1); + expect(result.abortReason).toContain("request budget exceeded"); + expect(handle.abortCalls()).toBeGreaterThanOrEqual(1); + expect(handle.steerCalls.length).toBe(1); + }); + + it("salvages the last assistant text for an aborted child with no completed output", async () => { + const settings = Settings.isolated({ "task.maxRuntimeMs": 50 }); + const handle = createFakeSession({ + hang: true, + events: [ + // One completed assistant turn with usage but no text content: + // counts a request and tokens without producing output chunks. + assistantMessageEnd("", { input: 100, output: 50, cacheRead: 0, cacheWrite: 0, totalTokens: 150 }), + ], + lastAssistantMessage: { + role: "assistant", + stopReason: "aborted", + content: [{ type: "text", text: "Reading the\n\tconfig loader before patching" }], + }, + }); + mockCreateAgentSession(handle.session); + + const result = await runSubprocess({ ...baseOptions, id: "subagent-salvage", settings }); + + expect(result.aborted).toBe(true); + expect(result.requests).toBe(1); + expect(result.output).toContain("cancelled after 1 req"); + expect(result.output).toContain("150 tok"); + expect(result.output).toContain("last activity:"); + // Whitespace is flattened so the snippet stays a single line. + expect(result.output).toContain("Reading the config loader before patching"); + expect(result.output).not.toContain("\n"); + }); + + it("clips oversized salvage snippets", async () => { + const settings = Settings.isolated({ "task.maxRuntimeMs": 50 }); + const longText = `start-marker ${"x".repeat(700)}`; + const handle = createFakeSession({ + hang: true, + lastAssistantMessage: { + role: "assistant", + stopReason: "aborted", + content: [{ type: "text", text: longText }], + }, + }); + mockCreateAgentSession(handle.session); + + const result = await runSubprocess({ ...baseOptions, id: "subagent-salvage-clip", settings }); + + expect(result.aborted).toBe(true); + expect(result.output).toContain("start-marker"); + expect(result.output).toContain("…"); + expect(result.output).not.toContain(longText); + expect(result.output.length).toBeLessThan(700); + }); + + it("formats the (no output) fallback with the request count", () => { + expect(formatResultOutputFallback({ output: "", stderr: "", requests: 7 })).toBe("(no output) after 7 req"); + expect(formatResultOutputFallback({ output: " ", stderr: "", requests: 0 })).toBe("(no output)"); + expect(formatResultOutputFallback({ output: "real output", stderr: "", requests: 7 })).toBe("real output"); + expect(formatResultOutputFallback({ output: "", stderr: "boom", requests: 7 })).toBe("boom"); + }); +}); diff --git a/packages/coding-agent/test/task/task-progress-render.test.ts b/packages/coding-agent/test/task/task-progress-render.test.ts index 1a4562582..1465ddca0 100644 --- a/packages/coding-agent/test/task/task-progress-render.test.ts +++ b/packages/coding-agent/test/task/task-progress-render.test.ts @@ -1,9 +1,9 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "bun:test"; import type { RenderResultOptions } from "@oh-my-pi/pi-agent-core"; import { resetSettingsForTest, Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; -import { getThemeByName } from "@oh-my-pi/pi-coding-agent/modes/theme/theme"; +import { getThemeByName, setThemeInstance } from "@oh-my-pi/pi-coding-agent/modes/theme/theme"; import { taskToolRenderer } from "@oh-my-pi/pi-coding-agent/task/render"; -import type { AgentProgress, TaskToolDetails } from "@oh-my-pi/pi-coding-agent/task/types"; +import type { AgentProgress, SingleResult, TaskToolDetails } from "@oh-my-pi/pi-coding-agent/task/types"; function runningProgress(overrides: Partial = {}): AgentProgress { return { @@ -16,6 +16,7 @@ function runningProgress(overrides: Partial = {}): AgentProgress recentTools: [], recentOutput: [], toolCount: 0, + requests: 0, tokens: 0, cost: 0, durationMs: 0, @@ -23,6 +24,24 @@ function runningProgress(overrides: Partial = {}): AgentProgress }; } +function finishedResult(overrides: Partial = {}): SingleResult { + return { + index: 0, + id: "Agent", + agent: "task", + agentSource: "bundled", + task: "investigate hot paths", + exitCode: 0, + output: "done", + stderr: "", + truncated: false, + durationMs: 0, + tokens: 0, + requests: 0, + ...overrides, + }; +} + function detailsFor(progress: AgentProgress): TaskToolDetails { return { projectAgentsDir: null, results: [], totalDurationMs: 0, progress: [progress] }; } @@ -102,7 +121,64 @@ describe("task progress rendering", () => { expect(strippedRow).not.toContain(theme.getSpinnerFrames("status")[0]); }); - it("pins unfinished tasks below finished ones in the live view", async () => { + it("shimmers the pending description like a running one (frozen async spawn snapshot)", async () => { + const theme = (await getThemeByName("dark"))!; + const options: RenderResultOptions = { expanded: false, isPartial: true, spinnerFrame: 0 }; + const progress = runningProgress({ + id: "BestGpt", + status: "pending", + description: "Combine winners for gpt", + }); + + const renderRow = (timeMs: number): string => { + vi.spyOn(Date, "now").mockReturnValue(timeMs); + return findRow( + taskToolRenderer.renderResult( + { content: [{ type: "text", text: "" }], details: detailsFor(progress) }, + options, + theme, + ), + "BestGpt", + ); + }; + + const rawRow0 = renderRow(0); + const rawRow1 = renderRow(700); + + expect(Bun.stripANSI(rawRow0)).toContain("BestGpt: Combine winners for gpt"); + // The label stays one solid bold-accent run; the description shimmers, + // so the row animates across frames exactly like a running agent's. + const label = theme.fg("accent", theme.bold("BestGpt")); + expect(rawRow0).toContain(label); + expect(rawRow1).toContain(label); + expect(rawRow0).not.toBe(rawRow1); + }); + + it("renders the assignment markdown inside the result frame", async () => { + const theme = (await getThemeByName("dark"))!; + setThemeInstance(theme); + const options: RenderResultOptions = { expanded: false, isPartial: true, spinnerFrame: 0 }; + const progress = runningProgress({ id: "BestGpt", status: "pending", description: "Combine winners" }); + + const rendered = Bun.stripANSI( + taskToolRenderer + .renderResult( + { content: [{ type: "text", text: "Spawned agent BestGpt..." }], details: detailsFor(progress) }, + options, + theme, + { agent: "task", id: "BestGpt", assignment: "# Target\nCombine the winning patches." }, + ) + .render(120) + .join("\n"), + ); + + // The brief stays visible for the whole task lifecycle, not just while + // the call args stream in. + expect(rendered).toContain("Target"); + expect(rendered).toContain("Combine the winning patches."); + }); + + it("pins unfinished tasks below finished ones, finished sorted by runtime asc", async () => { const theme = (await getThemeByName("dark"))!; const options: RenderResultOptions = { expanded: false, isPartial: true, spinnerFrame: 0 }; const details: TaskToolDetails = { @@ -110,10 +186,10 @@ describe("task progress rendering", () => { results: [], totalDurationMs: 0, progress: [ - runningProgress({ index: 0, id: "FirstRunning", status: "running" }), - runningProgress({ index: 1, id: "DoneEarly", status: "completed" }), + runningProgress({ index: 0, id: "FirstRunning", status: "running", durationMs: 9000 }), + runningProgress({ index: 1, id: "DoneSlow", status: "completed", durationMs: 5000 }), runningProgress({ index: 2, id: "StillPending", status: "pending" }), - runningProgress({ index: 3, id: "FailedFast", status: "failed" }), + runningProgress({ index: 3, id: "FailedFast", status: "failed", durationMs: 1000 }), ], }; @@ -124,8 +200,34 @@ describe("task progress rendering", () => { .join("\n"), ); - // Finished agents (in dispatch order) come first; pending/running stay at the bottom. - const positions = ["DoneEarly", "FailedFast", "FirstRunning", "StillPending"].map(id => rendered.indexOf(id)); + // Finished agents sorted by runtime ascending; pending/running stay at the + // bottom in dispatch order. + const positions = ["FailedFast", "DoneSlow", "FirstRunning", "StillPending"].map(id => rendered.indexOf(id)); + expect(positions.every(p => p >= 0)).toBe(true); + expect(positions).toEqual([...positions].sort((a, b) => a - b)); + }); + + it("orders finalized results by runtime asc, matching the live view", async () => { + const theme = (await getThemeByName("dark"))!; + const options: RenderResultOptions = { expanded: false, isPartial: false }; + const details: TaskToolDetails = { + projectAgentsDir: null, + results: [ + finishedResult({ index: 0, id: "SlowFinish", durationMs: 9000 }), + finishedResult({ index: 1, id: "FastFinish", durationMs: 1000 }), + finishedResult({ index: 2, id: "MidFinish", durationMs: 4000 }), + ], + totalDurationMs: 9000, + }; + + const rendered = Bun.stripANSI( + taskToolRenderer + .renderResult({ content: [{ type: "text", text: "" }], details }, options, theme) + .render(120) + .join("\n"), + ); + + const positions = ["FastFinish", "MidFinish", "SlowFinish"].map(id => rendered.indexOf(id)); expect(positions.every(p => p >= 0)).toBe(true); expect(positions).toEqual([...positions].sort((a, b) => a - b)); }); @@ -144,15 +246,17 @@ describe("task result detail-less state", () => { it("renders a validation failure with the error glyph, not a success bullet", async () => { const theme = (await getThemeByName("dark"))!; + // The assignment section renders markdown, which reads the active theme. + setThemeInstance(theme); const options: RenderResultOptions = { expanded: false, isPartial: false }; const component = taskToolRenderer.renderResult( { - content: [{ type: "text", text: 'Validation failed for tool "task": tasks: Invalid input' }], + content: [{ type: "text", text: 'Validation failed for tool "task": assignment: Invalid input' }], isError: true, }, options, theme, - { agent: "explore", tasks: [] }, + { agent: "explore", assignment: "Look around." }, ); const stripped = Bun.stripANSI(component.render(120).join("\n")); @@ -166,10 +270,11 @@ describe("task result detail-less state", () => { it("renders a detail-less success with the accent bullet, not an error glyph", async () => { const theme = (await getThemeByName("dark"))!; + setThemeInstance(theme); const options: RenderResultOptions = { expanded: false, isPartial: false }; const component = taskToolRenderer.renderResult({ content: [{ type: "text", text: "done" }] }, options, theme, { agent: "explore", - tasks: [], + assignment: "Look around.", }); const stripped = Bun.stripANSI(component.render(120).join("\n")); diff --git a/packages/coding-agent/test/task/task-resume.test.ts b/packages/coding-agent/test/task/task-resume.test.ts new file mode 100644 index 000000000..3db8d3273 --- /dev/null +++ b/packages/coding-agent/test/task/task-resume.test.ts @@ -0,0 +1,272 @@ +/** + * Contracts: task tool spawn/resume routing (rework-contracts.md §3). + * + * 1. With an AsyncJobManager wired, `execute` returns immediately (agent id + + * job id) while the job body is still gated; job completion delivers a + * result carrying the `task(resume:"")` / `history://` hint. + * 2. Resume routes through `AgentLifecycleManager.ensureLive` and hands the + * live session to `resumeSubprocess`; an ensureLive rejection surfaces as a + * ToolError naming `history://`. + * 3. The session-scoped spawn semaphore (task.maxConcurrency) serializes job + * bodies: with concurrency 1 the second body does not start until the + * first releases. + * + * Param validation (agent XOR resume, resume+isolated, missing assignment) is + * covered by test/task/task-schema.test.ts. + */ +import { afterEach, beforeEach, describe, expect, it, vi } from "bun:test"; +import { AsyncJobManager } from "@oh-my-pi/pi-coding-agent/async/job-manager"; +import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; +import { AgentLifecycleManager } from "@oh-my-pi/pi-coding-agent/registry/agent-lifecycle"; +import { AgentRegistry } from "@oh-my-pi/pi-coding-agent/registry/agent-registry"; +import type { AgentSession } from "@oh-my-pi/pi-coding-agent/session/agent-session"; +import { TaskTool } from "@oh-my-pi/pi-coding-agent/task"; +import * as discoveryModule from "@oh-my-pi/pi-coding-agent/task/discovery"; +import * as executorModule from "@oh-my-pi/pi-coding-agent/task/executor"; +import type { AgentDefinition, SingleResult, TaskParams } from "@oh-my-pi/pi-coding-agent/task/types"; +import type { ToolSession } from "@oh-my-pi/pi-coding-agent/tools"; +import { ToolError } from "@oh-my-pi/pi-coding-agent/tools/tool-errors"; + +const taskAgent: AgentDefinition = { + name: "task", + description: "General-purpose task agent", + systemPrompt: "You are a task agent.", + source: "bundled", +}; + +function createSession(options: { manager?: AsyncJobManager; settings?: Record }): ToolSession { + return { + cwd: "/tmp", + hasUI: false, + settings: Settings.isolated(options.settings ?? {}), + getSessionFile: () => null, + getSessionSpawns: () => "*", + asyncJobManager: options.manager, + } as unknown as ToolSession; +} + +function getFirstText(result: { content: Array<{ type: string; text?: string }> }): string { + const content = result.content.find(part => part.type === "text"); + return content?.type === "text" ? (content.text ?? "") : ""; +} + +function makeResult(id: string, overrides: Partial = {}): SingleResult { + return { + index: 0, + id, + agent: "task", + agentSource: "bundled", + task: "task prompt", + assignment: "Do the thing.", + exitCode: 0, + output: "All done.", + stderr: "", + truncated: false, + durationMs: 5, + tokens: 0, + requests: 1, + ...overrides, + }; +} + +interface Deferred { + promise: Promise; + resolve: () => void; +} + +function deferred(): Deferred { + let resolve!: () => void; + const promise = new Promise(res => { + resolve = res; + }); + return { promise, resolve }; +} + +async function pollUntil(predicate: () => boolean, timeoutMs = 2000): Promise { + const start = Date.now(); + while (!predicate()) { + if (Date.now() - start > timeoutMs) throw new Error("pollUntil timed out"); + await Bun.sleep(5); + } +} + +describe("task spawn/resume routing", () => { + const managers: AsyncJobManager[] = []; + + function createManager(): AsyncJobManager { + const manager = new AsyncJobManager({ onJobComplete: () => {} }); + managers.push(manager); + return manager; + } + + beforeEach(() => { + AgentRegistry.resetGlobalForTests(); + AgentLifecycleManager.resetGlobalForTests(); + }); + + afterEach(async () => { + vi.restoreAllMocks(); + for (const manager of managers.splice(0)) { + await manager.dispose({ timeoutMs: 1000 }); + } + AgentLifecycleManager.resetGlobalForTests(); + AgentRegistry.resetGlobalForTests(); + }); + + it("returns immediately on spawn and delivers the resume hint when the job completes", async () => { + vi.spyOn(discoveryModule, "discoverAgents").mockResolvedValue({ + agents: [taskAgent], + projectAgentsDir: null, + }); + const gate = deferred(); + const runSpy = vi.spyOn(executorModule, "runSubprocess").mockImplementation(async options => { + await gate.promise; + return makeResult(options.id ?? "?"); + }); + + const manager = createManager(); + const tool = await TaskTool.create(createSession({ manager })); + + const result = await tool.execute("tc-spawn", { + agent: "task", + id: "Spawnling", + description: "background work", + assignment: "Do the thing.", + } as TaskParams); + + // Tool returned while the job body is still gated on the deferred. + const text = getFirstText(result); + expect(text).toContain("Spawned agent `Spawnling`"); + const jobId = result.details?.async?.jobId; + expect(jobId).toBeTruthy(); + expect(text).toContain(`job \`${jobId}\``); + const job = manager.getJob(jobId!); + expect(job?.status).toBe("running"); + expect(job?.resultText).toBeUndefined(); + + gate.resolve(); + await job!.promise; + + expect(job!.status).toBe("completed"); + expect(job!.resultText).toContain('task(resume:"Spawnling")'); + expect(job!.resultText).toContain("history://Spawnling"); + expect(runSpy).toHaveBeenCalledTimes(1); + }); + + it("rejects an async resume of an unregistered agent without registering a job", async () => { + vi.spyOn(discoveryModule, "discoverAgents").mockResolvedValue({ + agents: [taskAgent], + projectAgentsDir: null, + }); + const manager = createManager(); + const tool = await TaskTool.create(createSession({ manager })); + + const error = await tool + .execute("tc-resume-unknown", { resume: "Nobody", assignment: "Follow up." } as TaskParams) + .then( + () => null, + err => err as Error, + ); + + expect(error).toBeInstanceOf(ToolError); + expect(error?.message).toContain('Unknown agent "Nobody"'); + expect(error?.message).toContain("history://Nobody"); + expect(manager.getAllJobs()).toHaveLength(0); + }); + + it("resume routes through AgentLifecycleManager.ensureLive and hands the live session to resumeSubprocess", async () => { + vi.spyOn(discoveryModule, "discoverAgents").mockResolvedValue({ agents: [], projectAgentsDir: null }); + const fakeSession = { messages: [] } as unknown as AgentSession; + AgentRegistry.global().register({ + id: "Reso", + displayName: "task", + kind: "sub", + session: fakeSession, + status: "idle", + }); + const ensureLiveSpy = vi.spyOn(AgentLifecycleManager.global(), "ensureLive").mockResolvedValue(fakeSession); + const resumeSpy = vi + .spyOn(executorModule, "resumeSubprocess") + .mockResolvedValue(makeResult("Reso", { output: "Follow-up done." })); + + // No job manager => sync fallback, so the resume pipeline runs inline. + const tool = await TaskTool.create(createSession({})); + const result = await tool.execute("tc-resume", { + resume: "Reso", + assignment: "Also check refresh tokens.", + } as TaskParams); + + expect(ensureLiveSpy).toHaveBeenCalledTimes(1); + expect(ensureLiveSpy).toHaveBeenCalledWith("Reso"); + expect(resumeSpy).toHaveBeenCalledTimes(1); + const resumeOptions = resumeSpy.mock.calls[0]![0]; + expect(resumeOptions.session).toBe(fakeSession); + expect(resumeOptions.id).toBe("Reso"); + expect(resumeOptions.assignment).toBe("Also check refresh tokens."); + + const text = getFirstText(result); + expect(text).toContain("Reso"); + expect(text).toContain("completed"); + expect(result.details?.results).toHaveLength(1); + expect(result.details?.results[0]?.exitCode).toBe(0); + }); + + it("surfaces an ensureLive rejection as a ToolError naming history://", async () => { + vi.spyOn(discoveryModule, "discoverAgents").mockResolvedValue({ agents: [], projectAgentsDir: null }); + vi.spyOn(AgentLifecycleManager.global(), "ensureLive").mockRejectedValue(new Error("session file corrupt")); + + const tool = await TaskTool.create(createSession({})); + const error = await tool + .execute("tc-resume-dead", { resume: "Ghost", assignment: "Wake up." } as TaskParams) + .then( + () => null, + err => err as Error, + ); + + expect(error).toBeInstanceOf(ToolError); + expect(error?.message).toContain('Cannot resume "Ghost"'); + expect(error?.message).toContain("session file corrupt"); + expect(error?.message).toContain("history://Ghost"); + }); + + it("bounds concurrent job bodies with the session spawn semaphore", async () => { + vi.spyOn(discoveryModule, "discoverAgents").mockResolvedValue({ + agents: [taskAgent], + projectAgentsDir: null, + }); + const started: string[] = []; + const gates = new Map(); + vi.spyOn(executorModule, "runSubprocess").mockImplementation(async options => { + const id = options.id ?? "?"; + started.push(id); + const gate = deferred(); + gates.set(id, gate); + await gate.promise; + return makeResult(id); + }); + + const manager = createManager(); + const tool = await TaskTool.create(createSession({ manager, settings: { "task.maxConcurrency": 1 } })); + + const first = await tool.execute("tc-1", { agent: "task", id: "First", assignment: "Work A." } as TaskParams); + const second = await tool.execute("tc-2", { agent: "task", id: "Second", assignment: "Work B." } as TaskParams); + const firstJob = manager.getJob(first.details!.async!.jobId)!; + const secondJob = manager.getJob(second.details!.async!.jobId)!; + + // First job body reaches the executor; second stays parked at the semaphore. + await pollUntil(() => started.length >= 1); + await Bun.sleep(25); + expect(started).toHaveLength(1); + + // Releasing the first body lets the second one start. + gates.get(started[0]!)!.resolve(); + await firstJob.promise; + await pollUntil(() => started.length === 2); + expect(started).toEqual(["First", "Second"]); + + gates.get("Second")!.resolve(); + await secondJob.promise; + expect(firstJob.status).toBe("completed"); + expect(secondJob.status).toBe("completed"); + }); +}); diff --git a/packages/coding-agent/test/task/task-schema.test.ts b/packages/coding-agent/test/task/task-schema.test.ts new file mode 100644 index 000000000..0e3f6ecb7 --- /dev/null +++ b/packages/coding-agent/test/task/task-schema.test.ts @@ -0,0 +1,83 @@ +import { afterEach, describe, expect, it, vi } from "bun:test"; +import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; +import { TaskTool, taskSchema } from "@oh-my-pi/pi-coding-agent/task"; +import * as discoveryModule from "@oh-my-pi/pi-coding-agent/task/discovery"; +import type { ToolSession } from "@oh-my-pi/pi-coding-agent/tools"; + +// Contract (rework-contracts.md §3): the task tool spawns ONE agent per call. +// `tasks[]` and `context` are gone; `resume` continues an existing agent. + +describe("task schema (single-spawn)", () => { + it("accepts {agent, assignment}", () => { + const parsed = taskSchema.safeParse({ agent: "explore", assignment: "Map the auth module." }); + expect(parsed.success).toBe(true); + }); + + it("accepts {resume, assignment}", () => { + const parsed = taskSchema.safeParse({ resume: "AuthLoader", assignment: "Also check refresh tokens." }); + expect(parsed.success).toBe(true); + }); + + it("requires assignment", () => { + const parsed = taskSchema.safeParse({ agent: "explore" }); + expect(parsed.success).toBe(false); + }); + + it("carries no tasks/context fields", () => { + const parsed = taskSchema.safeParse({ + agent: "explore", + assignment: "Map the auth module.", + context: "shared background", + tasks: [{ id: "A", assignment: "..." }], + }); + expect(parsed.success).toBe(true); + if (parsed.success) { + // Unknown keys are stripped: the batch/context shape no longer exists. + expect("tasks" in parsed.data).toBe(false); + expect("context" in parsed.data).toBe(false); + } + }); +}); + +describe("task spawn/resume validation", () => { + afterEach(() => { + vi.restoreAllMocks(); + }); + + function createSession(): ToolSession { + return { + cwd: "/tmp", + hasUI: false, + settings: Settings.isolated({ "task.isolation.mode": "none" }), + getSessionFile: () => null, + getSessionSpawns: () => "*", + } as unknown as ToolSession; + } + + async function executeText(params: unknown): Promise { + vi.spyOn(discoveryModule, "discoverAgents").mockResolvedValue({ agents: [], projectAgentsDir: null }); + const tool = await TaskTool.create(createSession()); + const result = await tool.execute("tool-call", params); + return result.content.find(part => part.type === "text")?.text ?? ""; + } + + it("rejects resume + agent together", async () => { + const text = await executeText({ agent: "explore", resume: "AuthLoader", assignment: "..." }); + expect(text).toContain("not both"); + }); + + it("rejects neither resume nor agent", async () => { + const text = await executeText({ assignment: "..." }); + expect(text).toContain("Missing `agent`"); + }); + + it("rejects resume + isolated", async () => { + const text = await executeText({ resume: "AuthLoader", isolated: true, assignment: "..." }); + expect(text).toContain("not resumable"); + }); + + it("rejects a missing assignment", async () => { + const text = await executeText({ agent: "explore" }); + expect(text).toContain("Missing `assignment`"); + }); +}); diff --git a/packages/coding-agent/test/tool-live-region-scrollback.test.ts b/packages/coding-agent/test/tool-live-region-scrollback.test.ts index d94b5691b..7434134b0 100644 --- a/packages/coding-agent/test/tool-live-region-scrollback.test.ts +++ b/packages/coding-agent/test/tool-live-region-scrollback.test.ts @@ -1,6 +1,6 @@ import { beforeAll, describe, expect, it } from "bun:test"; import type { AssistantMessage } from "@oh-my-pi/pi-ai"; -import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; +import { Settings, settings } from "@oh-my-pi/pi-coding-agent/config/settings"; import { AssistantMessageComponent } from "@oh-my-pi/pi-coding-agent/modes/components/assistant-message"; import { ToolExecutionComponent } from "@oh-my-pi/pi-coding-agent/modes/components/tool-execution"; import { TranscriptContainer } from "@oh-my-pi/pi-coding-agent/modes/components/transcript-container"; @@ -563,19 +563,21 @@ describe("tool live-region scrollback", () => { } }); - it("commits the scrolled-off head of an over-tall pending task context to scrollback", async () => { + it("commits the scrolled-off head of an over-tall pending eval cell to scrollback", async () => { if (process.platform === "win32") return; + // The single-spawn task renderer bounds its pending preview (the old + // uncapped multi-task `context` field is gone), so the eval tool — + // whose pending code preview is intentionally never capped — now + // carries the over-tall pending content. const term = new VirtualTerminal(120, 12); const tui = new TUI(term); const chat = new TranscriptContainer(); - const context = (n: number) => Array.from({ length: n }, (_unused, i) => `- CTX-${i}`).join("\n"); + const code = (n: number) => Array.from({ length: n }, (_unused, i) => `// - CTX-${i}`).join("\n"); const args = (n: number) => ({ - agent: "task", - context: context(n), - tasks: [{ id: "alpha", description: "probe", assignment: "Inspect the task context." }], + cells: [{ language: "js", title: "probe", code: code(n) }], }); - const component = new ToolExecutionComponent("task", args(4), {}, undefined, tui, process.cwd()); + const component = new ToolExecutionComponent("eval", args(4), {}, undefined, tui, process.cwd()); try { chat.addChild(component); @@ -603,35 +605,39 @@ describe("tool live-region scrollback", () => { } }); - it("keeps the static task context reachable in scrollback while progress ticks below it", async () => { + it("keeps the static task assignment reachable in scrollback while progress ticks below it", async () => { if (process.platform === "win32") return; const term = new VirtualTerminal(120, 12); const tui = new TUI(term); const chat = new TranscriptContainer(); - const context = Array.from({ length: 40 }, (_unused, i) => `- CTX-${i}`).join("\n"); - const args = { - agent: "explore", - context, - tasks: [{ id: "alpha", description: "probe", assignment: "Inspect the repo." }], - }; + const assignment = Array.from({ length: 40 }, (_unused, i) => `- CTX-${i}`).join("\n"); + const args = { agent: "explore", id: "alpha", description: "probe", assignment }; const component = new ToolExecutionComponent("task", args, {}, undefined, tui, process.cwd()); - const progressAt = (toolCount: number) => ({ + // The multi-line assignment section only renders expanded; shimmer + // would repaint the status line above it every frame, capping the + // stable prefix above the assignment, so pin it off for the run. + component.setExpanded(true); + settings.override("display.shimmer", "disabled"); + const progressAt = (tick: number) => ({ index: 0, id: "alpha", agent: "explore", agentSource: "bundled" as const, status: "running" as const, - task: "probe", + task: assignment, description: "probe", + currentTool: "read", + currentToolArgs: `probe-step-${tick}`, recentTools: [], recentOutput: [], - toolCount, + toolCount: 5, + requests: 0, tokens: 0, cost: 0, - durationMs: toolCount * 250, + durationMs: 1000, }); - const partial = (toolCount: number) => + const partial = (tick: number) => component.updateResult( { content: [{ type: "text", text: "" }], @@ -639,7 +645,7 @@ describe("tool live-region scrollback", () => { projectAgentsDir: null, results: [], totalDurationMs: 0, - progress: [progressAt(toolCount)], + progress: [progressAt(tick)], }, }, true, @@ -651,13 +657,14 @@ describe("tool live-region scrollback", () => { tui.start(); await term.waitForRender(); - // A running task rewrites its progress line (tool counts, spinner) - // below the static context for the whole run. The context head that - // scrolled above the viewport must still reach native scrollback — - // previously the ticking tail suspended commits for the entire - // block, leaving the context neither in history nor on screen. - // Two full promotion windows: the call→result transition frame - // poisons the first window's minimum, the second promotes the head. + // A running task rewrites its current-tool line (the ticking tail) + // below the static assignment section for the whole run. The + // assignment head that scrolled above the viewport must still reach + // native scrollback — previously the ticking tail suspended commits + // for the entire block, leaving the assignment neither in history + // nor on screen. Two full promotion windows: the call→result + // transition frame poisons the first window's minimum, the second + // promotes the head. for (let i = 1; i <= 70; i++) { partial(i); tui.requestRender(); @@ -669,8 +676,9 @@ describe("tool live-region scrollback", () => { expect(viewportText).not.toContain("CTX-0"); expect(scrollText).toContain("CTX-0"); - expect(scrollText).toContain("CTX-20"); + expect(scrollText).toContain("CTX-5"); } finally { + settings.clearOverride("display.shimmer"); component.stopAnimation(); tui.stop(); await term.flush(); diff --git a/packages/coding-agent/test/tools/task-repair-args.test.ts b/packages/coding-agent/test/tools/task-repair-args.test.ts index dec3c57cd..4629b327d 100644 --- a/packages/coding-agent/test/tools/task-repair-args.test.ts +++ b/packages/coding-agent/test/tools/task-repair-args.test.ts @@ -44,31 +44,27 @@ describe("repairDoubleEncodedJsonString", () => { }); describe("repairTaskParams", () => { - it("repairs context and each task's assignment/description, leaving ids intact", () => { + it("repairs assignment and description, leaving agent/id intact", () => { const params = { agent: "task", - context: "# Goal\\nDo the thing \\u2014 carefully", - tasks: [ - { - id: "FirstTask", - description: 'judge \\"sketch\\" accuracy', - assignment: "Score 0-100.\\nUse the full range.\\nNo bunching.", - }, - ], + id: "FirstTask", + description: 'judge \\"sketch\\" accuracy', + assignment: "Score 0-100.\\nUse the full range.\\nNo bunching.", } as unknown as TaskParams; const repaired = repairTaskParams(params); - expect(repaired.context).toBe("# Goal\nDo the thing — carefully"); - expect(repaired.tasks[0].id).toBe("FirstTask"); - expect(repaired.tasks[0].description).toBe('judge "sketch" accuracy'); - expect(repaired.tasks[0].assignment).toBe("Score 0-100.\nUse the full range.\nNo bunching."); + expect(repaired.agent).toBe("task"); + expect(repaired.id).toBe("FirstTask"); + expect(repaired.description).toBe('judge "sketch" accuracy'); + expect(repaired.assignment).toBe("Score 0-100.\nUse the full range.\nNo bunching."); }); it("returns the same reference when nothing needs repair", () => { const params = { agent: "task", - context: "plain context", - tasks: [{ id: "A", description: "label", assignment: "do work" }], + id: "A", + description: "label", + assignment: "do work", } as unknown as TaskParams; expect(repairTaskParams(params)).toBe(params); }); diff --git a/packages/coding-agent/test/tools/task-simple-mode.test.ts b/packages/coding-agent/test/tools/task-simple-mode.test.ts index 1e7636e87..3a1189ef3 100644 --- a/packages/coding-agent/test/tools/task-simple-mode.test.ts +++ b/packages/coding-agent/test/tools/task-simple-mode.test.ts @@ -16,6 +16,8 @@ const TEST_AGENTS = [ }, ]; +const ALL_MODES = ["default", "schema-free", "independent"] as const; + function createSession(overrides: Partial> = {}): ToolSession { return { cwd: "/tmp", @@ -36,87 +38,86 @@ function getFirstText(result: { content: Array<{ type: string; text?: string }> return content?.type === "text" ? (content.text ?? "") : ""; } +function mockDiscovery(): void { + vi.spyOn(discoveryModule, "discoverAgents").mockResolvedValue({ + agents: TEST_AGENTS, + projectAgentsDir: null, + }); +} + describe("task.simple", () => { afterEach(() => { vi.restoreAllMocks(); }); - it("removes only the custom schema input in schema-free mode", async () => { - vi.spyOn(discoveryModule, "discoverAgents").mockResolvedValue({ - agents: TEST_AGENTS, - projectAgentsDir: null, - }); + it("exposes the custom schema input only in default mode", async () => { + mockDiscovery(); - const tool = await TaskTool.create(createSession({ "task.simple": "schema-free" })); - const properties = getSchemaProperties(tool); + const defaultTool = await TaskTool.create(createSession({ "task.simple": "default" })); + expect(getSchemaProperties(defaultTool).schema).toBeDefined(); + expect(defaultTool.description).toContain("- `schema`:"); - expect(properties.context).toBeDefined(); - expect(properties.schema).toBeUndefined(); - expect(tool.description).toContain("`context` or `assignment`"); - expect(tool.description).toContain("- `context`:"); - expect(tool.description).not.toContain("- `schema`:"); + for (const mode of ["schema-free", "independent"] as const) { + const tool = await TaskTool.create(createSession({ "task.simple": mode })); + expect(getSchemaProperties(tool).schema).toBeUndefined(); + expect(tool.description).not.toContain("- `schema`:"); + } }); - it("removes both context and schema inputs in independent mode", async () => { - vi.spyOn(discoveryModule, "discoverAgents").mockResolvedValue({ - agents: TEST_AGENTS, - projectAgentsDir: null, - }); + it("never exposes batch tasks or shared context inputs in any mode", async () => { + mockDiscovery(); - const tool = await TaskTool.create(createSession({ "task.simple": "independent" })); - const properties = getSchemaProperties(tool); - - expect(properties.context).toBeUndefined(); - expect(properties.schema).toBeUndefined(); - expect(tool.description).toContain("each `assignment`"); - expect(tool.description).not.toContain("- `context`:"); - expect(tool.description).not.toContain("- `schema`:"); + for (const mode of ALL_MODES) { + const tool = await TaskTool.create(createSession({ "task.simple": mode })); + const properties = getSchemaProperties(tool); + expect(properties.tasks).toBeUndefined(); + expect(properties.context).toBeUndefined(); + // The flat single-spawn contract is what replaced them. + expect(properties.assignment).toBeDefined(); + expect(properties.resume).toBeDefined(); + } }); - it("rejects direct schema and context fields when the mode disables them", async () => { - vi.spyOn(discoveryModule, "discoverAgents").mockResolvedValue({ - agents: TEST_AGENTS, - projectAgentsDir: null, - }); + it("describes the non-blocking spawn and resume contract", async () => { + mockDiscovery(); - const schemaFreeTool = await TaskTool.create(createSession({ "task.simple": "schema-free" })); - const schemaFreeResult = await schemaFreeTool.execute("tool-1", { - agent: "task", - schema: '{"properties":{"ok":{"type":"boolean"}}}', - tasks: [{ id: "One", description: "label", assignment: "Do the thing." }], - } as TaskParams); - expect(getFirstText(schemaFreeResult)).toContain("does not accept `schema`"); - const validatedSchemaFreeParams = validateToolArguments(schemaFreeTool, { - type: "toolCall", - id: "tool-1-validated", - name: schemaFreeTool.name, - arguments: { + const tool = await TaskTool.create(createSession({ "task.simple": "default" })); + expect(tool.description).toContain("Spawning is non-blocking"); + expect(tool.description).toContain("revives an idle/parked agent"); + }); + + it("rejects a direct schema input when the mode disables it", async () => { + mockDiscovery(); + + for (const mode of ["schema-free", "independent"] as const) { + const tool = await TaskTool.create(createSession({ "task.simple": mode })); + + // Execution-time guard: raw params carrying `schema` are refused. + const result = await tool.execute(`tool-${mode}`, { agent: "task", + id: "One", + description: "label", + assignment: "Do the thing.", schema: '{"properties":{"ok":{"type":"boolean"}}}', - tasks: [{ id: "One", description: "label", assignment: "Do the thing." }], - }, - }); - const validatedSchemaFreeResult = await schemaFreeTool.execute("tool-1-validated", validatedSchemaFreeParams); - expect(getFirstText(validatedSchemaFreeResult)).toContain("does not accept `schema`"); + } as TaskParams); + expect(getFirstText(result)).toContain("does not accept `schema`"); - const independentTool = await TaskTool.create(createSession({ "task.simple": "independent" })); - const independentResult = await independentTool.execute("tool-2", { - agent: "task", - context: "Shared background", - tasks: [{ id: "Two", description: "label", assignment: "Do the independent thing." }], - } as TaskParams); - expect(getFirstText(independentResult)).toContain("does not accept `context`"); - const validatedIndependentParams = validateToolArguments(independentTool, { - type: "toolCall", - id: "tool-2-validated", - name: independentTool.name, - arguments: { - agent: "task", - context: "Shared background", - tasks: [{ id: "Two", description: "label", assignment: "Do the independent thing." }], - }, - }); - const validatedIndependentResult = await independentTool.execute("tool-2-validated", validatedIndependentParams); - expect(getFirstText(validatedIndependentResult)).toContain("does not accept `context`"); + // Round-trip guard: wire validation passes the extraneous `schema` + // through, so the execution-time check must still refuse it. + const validated = validateToolArguments(tool, { + type: "toolCall", + id: `tool-${mode}-validated`, + name: tool.name, + arguments: { + agent: "task", + id: "One", + description: "label", + assignment: "Do the thing.", + schema: '{"properties":{"ok":{"type":"boolean"}}}', + }, + }) as TaskParams; + const validatedResult = await tool.execute(`tool-${mode}-validated`, validated); + expect(getFirstText(validatedResult)).toContain("does not accept `schema`"); + } }); }); diff --git a/packages/swarm-extension/src/swarm/pipeline.ts b/packages/swarm-extension/src/swarm/pipeline.ts index 82c247c14..d2e05464d 100644 --- a/packages/swarm-extension/src/swarm/pipeline.ts +++ b/packages/swarm-extension/src/swarm/pipeline.ts @@ -185,6 +185,7 @@ export class PipelineController { truncated: false, durationMs: 0, tokens: 0, + requests: 0, error, }; return { agentName, result: failResult }; From 32cb88066890ccee042adb76cb7159e217d84318 Mon Sep 17 00:00:00 2001 From: can1357 Date: Wed, 10 Jun 2026 17:55:08 +0200 Subject: [PATCH 22/27] chore(coding-agent): regenerated html export template for task resume and irc ops Reflects the task tool's resume path (single-agent badge or resume=) and the new irc op/receipt shapes (no awaitReply, op falls back to result.details.op) in the static html exporter. --- .../src/export/html/template.generated.ts | 2 +- .../coding-agent/src/export/html/template.js | 51 ++++++++++++++----- 2 files changed, 39 insertions(+), 14 deletions(-) diff --git a/packages/coding-agent/src/export/html/template.generated.ts b/packages/coding-agent/src/export/html/template.generated.ts index 10510d666..177e23630 100644 --- a/packages/coding-agent/src/export/html/template.generated.ts +++ b/packages/coding-agent/src/export/html/template.generated.ts @@ -1,2 +1,2 @@ // Auto-generated by scripts/generate-template.ts - DO NOT EDIT -export const TEMPLATE = "\n\n\n \n \n Session Export\n \n \n\n\n \n
\n
\n \n
\n
\n
\n
\n
\n
\n \"\"\n
\n
\n\n \n \n \n \n\n\n"; +export const TEMPLATE = "\n\n\n \n \n Session Export\n \n \n\n\n \n
\n
\n \n
\n
\n
\n
\n
\n
\n \"\"\n
\n
\n\n \n \n \n \n\n\n"; diff --git a/packages/coding-agent/src/export/html/template.js b/packages/coding-agent/src/export/html/template.js index d94777191..98ee6ef64 100644 --- a/packages/coding-agent/src/export/html/template.js +++ b/packages/coding-agent/src/export/html/template.js @@ -1023,18 +1023,18 @@ } function renderTask(name, args, result, ctx) { - const agent = str(args.agent) || '?'; - const tasks = Array.isArray(args.tasks) ? args.tasks : []; - const badges = ['agent=' + agent, tasks.length + ' subtask' + (tasks.length === 1 ? '' : 's')]; + const badges = []; + if (args.resume) badges.push('resume=' + str(args.resume)); + else badges.push('agent=' + (str(args.agent) || '?')); + if (args.id) badges.push('id=' + str(args.id)); if (args.isolated) badges.push('isolated'); let html = toolHead('task', '', badges); - if (tasks.length) { + const description = str(args.description); + const assignment = str(args.assignment); + if (description || assignment) { html += '
'; - for (const t of tasks) { - const id = t?.id ? escapeHtml(String(t.id)) : '?'; - const desc = t?.description ? escapeHtml(String(t.description)) : ''; - html += '
' + id + ' ' + desc + '
'; - } + if (description) html += '
' + escapeHtml(description) + '
'; + if (assignment) html += '
' + escapeHtml(assignment) + '
'; html += '
'; } if (result) { @@ -1479,13 +1479,38 @@ } function renderIrc(name, args, result, ctx) { - const op = str(args.op) || '?'; + const details = result && result.details ? result.details : null; + const op = str(args.op) || (details && str(details.op)) || '?'; const badges = [op]; - if (args.to) badges.push('to=' + args.to); - if (args.awaitReply === false) badges.push('no-reply'); + if (args.to) badges.push('to=' + str(args.to)); + if (op === 'wait' && args.from) badges.push('from=' + str(args.from)); + if (args.await) badges.push('await'); + if (args.peek) badges.push('peek'); let html = toolHead('irc', '', badges); if (args.message) html += '
' + escapeHtml(String(args.message)) + '
'; - if (result) { + let renderedDetails = false; + if (details && Array.isArray(details.receipts) && details.receipts.length) { + html += '
'; + for (const receipt of details.receipts) { + const outcome = escapeHtml(String(receipt.outcome)) + (receipt.error ? ' — ' + escapeHtml(String(receipt.error)) : ''); + html += '
' + escapeHtml(String(receipt.to)) + ' ' + outcome + '
'; + } + html += '
'; + renderedDetails = true; + } + if (details && details.waited) { + html += '
' + escapeHtml(String(details.waited.from)) + ': ' + escapeHtml(String(details.waited.body)) + '
'; + renderedDetails = true; + } + if (details && Array.isArray(details.inbox) && details.inbox.length) { + html += '
'; + for (const msg of details.inbox) { + html += '
' + escapeHtml(String(msg.from)) + ' ' + escapeHtml(String(msg.body)) + '
'; + } + html += '
'; + renderedDetails = true; + } + if (!renderedDetails && result) { const output = ctx.getResultText(); if (output) html += formatExpandableOutput(output, 8); } From 6ff1fe5caf012f8256fe4cb8c452e27906b03976 Mon Sep 17 00:00:00 2001 From: can1357 Date: Wed, 10 Jun 2026 17:55:19 +0200 Subject: [PATCH 23/27] docs(coding-agent): logged Unreleased entries and snapcompact compaction notes Adds the [Unreleased] block to the coding-agent CHANGELOG covering the task tool single-call rewrite, IRC mailbox bus rework, eval agent() context drop, Agent Hub overlay, supersede pruning, soft request budgets, subagent salvage, repeat-read notice, inline byte cap, agent lifecycle, history:// protocol, snapcompact strategy, IRC tool renderer, compaction inline divider, and async.enabled scope narrowing. docs/compaction.md gets the snapcompact strategy section (rationale, frame budget, eviction policy, fallback to context-full). --- docs/compaction.md | 17 ++++++++++++++++- packages/coding-agent/CHANGELOG.md | 24 +++++++++++++++++++++++- 2 files changed, 39 insertions(+), 2 deletions(-) diff --git a/docs/compaction.md b/docs/compaction.md index 2fb5a22e9..789e36aa1 100644 --- a/docs/compaction.md +++ b/docs/compaction.md @@ -10,6 +10,7 @@ Both are persisted as session entries and converted back into user-context messa ## Key implementation files - `packages/agent/src/compaction/compaction.ts` (context-full summarization and handoff generation) +- `packages/agent/src/compaction/snapcompact.ts` (snapcompact strategy: history archived as dense bitmap images) - `packages/agent/src/compaction/branch-summarization.ts` - `packages/agent/src/compaction/pruning.ts` - `packages/agent/src/compaction/utils.ts` @@ -126,6 +127,20 @@ The automatic paths are intentionally different: - Trigger: `runIdleCompaction()` when not streaming or already compacting. - Uses `reason: "idle"` and does not auto-continue afterward. +### Snapcompact strategy + +`compaction.strategy: "snapcompact"` replaces the LLM summarization call with a local, deterministic archival pass (`snapcompactCompact` in `packages/agent/src/compaction/snapcompact.ts`): + +- The discarded history is serialized (same `serializeConversation` path the summarizer uses), whitespace-collapsed, and printed onto square PNG frames using the bundled public-domain X.org `5x8` pixel font — 2576px frames carrying ~165k characters each. Glyph ink cycles through six hues at sentence boundaries to help the model keep its place. +- Frames persist under `CompactionEntry.preserveData.snapcompact` and are re-attached to the `compactionSummary` message as image blocks on every context rebuild; the entry's `summary` is a deterministic reading guide (grid geometry, role tags, truncation notes) plus the usual file-operation lists. +- Later compactions carry earlier frames forward. Beyond an 8-frame budget the archive fades from the middle out: the earliest frame (session head — the original request, or the filmed summary of older history) is pinned, and the oldest *unpinned* frames are evicted, so head and tail both survive. If the previous compaction was text-based, its summary is printed at the head of the frame archive as `[Summary of earlier history]`. +- No model, API key, or network is involved, so snapcompact is also safe for overflow recovery. It requires a vision-capable current model (`model.input` includes `"image"`); otherwise the run falls back to context-full and emits a warning notice (auto and manual paths). Manual `/compact` honors the strategy unless custom instructions are given (those imply a directed LLM summary). +- Rationale: a vision model reads ~50 chars per image token after provider downscaling, ~7x cheaper than raw text at near-parity recall (SQuAD eval: F1 0.878 vs 0.899). + +### Display transcript + +Compaction no longer visually restarts the conversation. The TUI renders the **display transcript** (`buildSessionContext({ transcript: true })` / `AgentSession.buildTranscriptSessionContext()`): every path entry in chronological order, with each compaction shown inline as a slim divider — `── 📷 compacted · ctrl+o ──` — at the point it fired. Expanding (ctrl+o) reveals the summary. Only the LLM context resets at the compaction boundary; the scrollback above the divider stays intact, including across session resume. + ### Pre-compaction pruning Before compaction checks, tool-result pruning may run (`pruneToolOutputs`). @@ -373,7 +388,7 @@ Post-navigation event exposing new/old leaf and optional summary entry. From `settings-schema.ts`: - `compaction.enabled` = `true` -- `compaction.strategy` = `"context-full"` (`"handoff"` and `"off"` are also supported) +- `compaction.strategy` = `"context-full"` (`"handoff"`, `"shake"`, `"snapcompact"`, and `"off"` are also supported) - `compaction.reserveTokens` = `16384` - `compaction.keepRecentTokens` = `20000` - `compaction.autoContinue` = `true` diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 3de17230b..6fcc94773 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,15 +2,37 @@ ## [Unreleased] +### Breaking Changes + +- The `task` tool now spawns exactly one subagent per call and always runs it in the background: the batch `tasks[]` array and shared `context` parameter are removed — fan out with parallel `task` calls, share background via a `local://` file referenced in each assignment, and receive results as async job deliveries (block with `job poll` only when genuinely needed) +- Reworked `irc` to `send`/`wait`/`inbox`/`list` ops over a per-agent mailbox bus: the blocking `awaitReply` auto-reply turn is removed — `send` is fire-and-forget with delivery receipts, and replies are real turns by the recipient observed via `wait` (or the `send` `await: true` sugar) +- Removed the `context` argument from eval `agent()` in both the JS and Python preludes: pass shared background via a `local://` file referenced in the prompt +- Replaced the standalone session-observer overlay with the Agent Hub: `app.session.observe` (`ctrl+s`) now opens the hub, whose chat view absorbed the observer's transcript renderer + ### Added - Added pre-TUI startup input capture so users could type while interactive sessions initialize and keep their draft while the application loads - Added queued submission replay so Enter presses made before startup completion are submitted automatically once interactive mode begins +- Added per-turn supersede pruning of stale `read` results: when a file is re-read, older copies of the same path/selector are pruned from context at cache-favorable moments (small suffix, idle gap, or alongside overflow pruning). Gated by the new `compaction.supersedeReads` setting (default on) +- Added soft request budgets for task subagents (explore/quick_task 40, others 90, configurable via `task.softRequestBudget`, 0 disables): crossing the budget injects a one-time wrap-up steer into the child; crossing 1.5× aborts the run gracefully +- Added cancelled/aborted subagent salvage: instead of `(no output)`, merged task results now carry the child's last activity snippet plus request/token stats, and per-child stats lines include request counts +- Added a repeat-read notice to the `read` tool: the third and later reads of the same file in a session append a one-line note suggesting range re-reads or the context echoed in edit results +- Added a hard inline byte cap (~50KB) at the bash and browser tool-result boundaries with head/tail elision and an `artifact://` footer for the full output, closing paths that previously let 100KB+ results land inline +- Added the Agent Hub overlay (`ctrl+s`, `alt+a`, or double-tap left arrow on an empty editor): a live table of registered subagents (status, unread IRC count, current task, last activity) with per-agent chat — Enter opens a transcript + input line that steers a running agent, prompts an idle one, and revives a parked one; `r` revives and `x` aborts/releases the selected agent +- Added the `snapcompact` compaction strategy (`compaction.strategy: "snapcompact"`): history is archived onto dense bitmap "snapcompact" frames a vision model reads back directly, instead of an LLM-generated summary — instant, free, and verbatim. Auto compaction (including overflow recovery) and manual `/compact` both honor it; falls back to context-full with a visible warning notice when the current model is text-only (e.g. Codex API surfaces) or when `/compact` is given custom instructions. Frames survive context rebuilds and later compactions (budget eviction is middle-out: the session-head frame is pinned); the expanded compaction message notes the attached frame count +- Added a persistent subagent lifecycle: finished subagents stay live as `idle`, are parked to disk after `task.agentIdleTtlMs` (default 7 minutes; `0` keeps them live until exit), and are revived automatically when messaged, resumed, or prompted from the Agent Hub +- Added `task(resume: "")` to revive an idle or parked subagent and run a follow-up assignment in its existing session, keeping its accumulated context +- Added the `history://` protocol: `history://` lists every registered agent and `history://` renders a concise markdown transcript (tool calls collapsed to one line each, thinking elided) for live and parked agents alike +- Added an IRC mailbox bus with bounded per-agent inboxes: `irc` `wait` blocks until a matching message arrives, `inbox` drains or peeks pending messages, and sending to an idle or parked agent wakes or revives it for a real turn +- Added a dedicated TUI renderer for the `irc` tool: directional send/receive headers with delivery-outcome coloring, quoted message bodies with expand-aware truncation, per-recipient receipt trees for broadcasts and failures, and status-badged peer listings with unread counts ### Changed - Changed interactive startup to carry the startup editor state into the live prompt so text entered during splash is preserved in the editor when the TUI takes over +- Changed the compaction UX so the conversation no longer visually restarts: the TUI renders the full-history display transcript (`buildSessionContext({ transcript: true })`), with each compaction shown as a slim inline divider — `── 📷 compacted · ctrl+o ──` — at the point it fired; expanding (ctrl+o) reveals the summary and snapcompact frame count. Applies to live compaction, `/compact`, `/tree` navigation, and session resume - Changed model-scope display during startup to appear as an in-UI information notification instead of a direct stdout line +- Changed `async.enabled` to gate async bash commands only — the `task` tool now runs asynchronously regardless of the setting +- Changed `irc.timeoutMs` to be the default timeout for `irc` `wait` and `send` with `await: true` ### Fixed @@ -38,7 +60,7 @@ ### Changed - Bash execution now preserves minimized shell output inline while saving the untouched capture as an `artifact://…` footer when shell minimization rewrites a command's output. -- Task tool live progress now renders finished subagents first and keeps unfinished (pending/running) ones pinned at the bottom of the list. +- Task tool agent lists now render in runtime-ascending order in both the live progress view (finished agents, sorted by runtime, above pending/running ones) and the finalized result view, so rows no longer reshuffle when the call finalizes. - `OutputSink` artifact files (`~/.omp/agent/artifacts/..log`) are unbounded by default again, so `artifact://` references preserve the complete raw stream. The head + rolling-tail capping machinery from [#2081](https://github.com/can1357/oh-my-pi/issues/2081) (with its `[ARTIFACT TRUNCATED: …]` close notice) remains available as an opt-in via `artifactMaxBytes`, and the head window now closes permanently on first overflow so later small chunks cannot be written out of order before the tail replay. ### Fixed From 332d4537346dc08591b6c03610337baa6ff0e41c Mon Sep 17 00:00:00 2001 From: roboomp Date: Wed, 10 Jun 2026 18:46:14 +0000 Subject: [PATCH 24/27] fix(ai): preserve 3p anthropic-messages reasoning chains across model swaps MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The transform pass dropped or demoted every prior anthropic-messages assistant turn the moment the conversation crossed a model id — custom anthropic-messages providers configured via `models.yaml` (DeepSeek, Z.AI, etc.) and session-level model swaps both lost their reasoning chains on continuation, silently degrading agent behavior for compatible reasoning endpoints. Two coupled changes in `transform-messages.ts`: * All prior `anthropic-messages → anthropic-messages` turns now keep their `thinking`/`redactedThinking` blocks as native blocks, not just the latest surviving assistant. The legacy `mustPreserveLatestAnthropicThinking` flag is replaced by `isAnthropicReplay` + `isLatestSurvivingAssistant`; the latest byte-for-byte abandoned-tool-use exemption is preserved. * Cross-model signatures are only stripped when an official Anthropic endpoint is involved on either end of the replay. 3p ↔ 3p replays (Z.AI, DeepSeek, custom anthropic-messages providers) pass the signature through unchanged so the reasoning chain stays signed for the next turn — compatible reasoning endpoints treat signatures as opaque continuation hints, not cryptographically-verified material. Source-side official detection uses the canonical catalog provider id `"anthropic"` (assistant messages carry no `baseUrl`); target- side detection reads the baked `compat.officialEndpoint` flag. Any `redacted_thinking` sibling without a native landing spot is still dropped to avoid lone redacted blocks beside text-demoted reasoning. Fixes #2257 Fixes #2265 --- packages/ai/CHANGELOG.md | 4 + .../ai/src/providers/transform-messages.ts | 87 ++++- .../anthropic-prior-turn-thinking.test.ts | 342 ++++++++++++++++++ 3 files changed, 426 insertions(+), 7 deletions(-) create mode 100644 packages/ai/test/anthropic-prior-turn-thinking.test.ts diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index af15979fe..1755a96b4 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- Cross-model `anthropic-messages → anthropic-messages` continuations now preserve prior assistant turns' reasoning chains end-to-end: every prior `thinking`/`redactedThinking` block survives (not just the latest surviving assistant), and third-party ↔ third-party replays keep their signatures intact so the reasoning chain stays signed for the next turn. Signatures are stripped (and any `redacted_thinking` sibling without a native landing spot is dropped) only when an official Anthropic endpoint is on either end of the replay — official Anthropic cryptographically binds reasoning signatures to its key+session+model, while compatible reasoning endpoints (Z.AI, DeepSeek, custom anthropic-messages providers configured via `models.yaml`) treat them as opaque continuation hints. Source-side official detection uses the canonical catalog provider id `"anthropic"` (assistant messages carry no `baseUrl`); target-side detection reuses the baked `compat.officialEndpoint` flag. Latest-turn byte-for-byte behavior (Anthropic's "thinking blocks in the latest assistant message cannot be modified" rule) and existing aborted/errored last-block sanitization are unchanged. ([#2257](https://github.com/can1357/oh-my-pi/issues/2257), [#2265](https://github.com/can1357/oh-my-pi/issues/2265)) + ## [15.10.12] - 2026-06-10 ### Added diff --git a/packages/ai/src/providers/transform-messages.ts b/packages/ai/src/providers/transform-messages.ts index 34346d99c..eb810ae1a 100644 --- a/packages/ai/src/providers/transform-messages.ts +++ b/packages/ai/src/providers/transform-messages.ts @@ -139,6 +139,10 @@ function getLatestSurvivingAssistantIndex(messages: readonly Message[]): number return -1; } +function isAnthropicMessagesModel(model: Model): model is Model<"anthropic-messages"> { + return model.api === "anthropic-messages"; +} + /** * Normalize tool call ID for cross-provider compatibility. * OpenAI Responses API generates IDs that are 450+ chars with special characters like `|`. @@ -184,10 +188,45 @@ export function transformMessages( assistantMsg.api === model.api && assistantMsg.model === model.id; - const mustPreserveLatestAnthropicThinking = - index === latestSurvivingAssistantIndex && - model.api === "anthropic-messages" && - assistantMsg.api === "anthropic-messages"; + const isAnthropicTarget = isAnthropicMessagesModel(model); + // Anthropic's all-or-none contract on prior-turn thinking blocks + // applies to every `anthropic-messages → anthropic-messages` replay, + // not just the latest assistant turn. The legacy + // `mustPreserveLatestAnthropicThinking` flag only honored it for the + // latest turn; every prior turn fell through to the cross-API + // text-demotion path whenever the conversation crossed a model id, + // silently dropping the reasoning chain on continuation for custom + // anthropic-messages providers configured via `models.yaml` and + // session-level model swaps (#2257). + const isAnthropicReplay = isAnthropicTarget && assistantMsg.api === "anthropic-messages"; + const isLatestSurvivingAssistant = index === latestSurvivingAssistantIndex; + // Signature policy is a second axis. Anthropic cryptographically + // binds reasoning signatures to its key+session+model, so cross-model + // signatures must be stripped whenever official Anthropic is on + // either end of the replay: + // * official → 3p: the 3p target can't reverify the signature; + // keeping it leaks private continuation metadata for no benefit. + // * 3p → official: official rejects a foreign signature outright. + // * official → official cross-model: the new model rejects the + // previous model's signature. + // 3p ↔ 3p replays preserve signatures because compatible providers + // (Z.AI, DeepSeek, custom `models.yaml` providers) treat them as + // opaque continuation hints rather than verified material; stripping + // degrades the reasoning chain into unsigned/text on the next turn + // (#2265). Source-side official detection uses the canonical catalog + // provider id `"anthropic"` because assistant messages carry no + // `baseUrl` — a user who manually points `provider: "anthropic"` at + // a custom proxy via `models.yaml` will see signatures stripped, the + // conservative direction (degraded reasoning, not broken requests). + const isOfficialAnthropicSource = isAnthropicReplay && assistantMsg.provider === "anthropic"; + const isOfficialAnthropicTarget = isAnthropicTarget && model.compat.officialEndpoint; + const officialAnthropicInvolved = isOfficialAnthropicSource || isOfficialAnthropicTarget; + // Compatible Anthropic-messages reasoning targets that accept + // unsigned thinking natively (Z.AI, DeepSeek, the generic + // `reasoning && !official` case in the compat builder). Used to keep + // `redacted_thinking` siblings beside unsigned visible thinking on + // targets that won't text-demote it. + const replaysUnsignedAnthropicThinking = isAnthropicTarget && model.compat.replayUnsignedThinking; // Thinking signatures can be untrustworthy for two distinct reasons with very // different blast radii: // @@ -226,11 +265,37 @@ export function transformMessages( // untrustworthy signature so the encoder can downgrade the block to text. const signatureUntrustworthy = abandonedToolUse || (invalidStopReason && blockIndex === lastBlockIndex); - const sanitized = + let sanitized: typeof block = signatureUntrustworthy && block.thinkingSignature ? { ...block, thinkingSignature: undefined } : block; - if (mustPreserveLatestAnthropicThinking) return abandonedToolUse ? block : sanitized; + if (isAnthropicReplay) { + // Latest abandoned turn: Anthropic's byte-for-byte rule forbids + // even stripping a signature on the latest message. + if (isLatestSurvivingAssistant && abandonedToolUse) return block; + // Cross-model prior turns crossing an official Anthropic endpoint + // must strip the source signature so the downstream encoder + // applies its `replayUnsignedThinking` policy (unsigned thinking + // is emitted natively on Anthropic-compatible reasoning endpoints + // and demoted to text on official Anthropic). 3p ↔ 3p replays + // keep the signature so the reasoning chain stays signed on + // continuation (#2265). + if ( + !isLatestSurvivingAssistant && + !isSameModel && + officialAnthropicInvolved && + sanitized.thinkingSignature + ) { + sanitized = { ...sanitized, thinkingSignature: undefined }; + } + // Drop blocks with neither a signature anchor nor any text — + // nothing for the next turn to replay. + if (!sanitized.thinkingSignature && (!sanitized.thinking || sanitized.thinking.trim() === "")) { + return []; + } + return sanitized; + } + // Cross-API target: keep the existing text-demotion fallback. // For same model: keep thinking blocks with signatures (needed for replay) // even if the thinking text is empty (OpenAI encrypted reasoning) if (isSameModel && sanitized.thinkingSignature) return sanitized; @@ -244,7 +309,15 @@ export function transformMessages( } if (block.type === "redactedThinking") { - if (mustPreserveLatestAnthropicThinking) return block; + // Redacted thinking is native-only. Keep it for same-model + // signed replay, the latest byte-for-byte Anthropic turn, or + // compatible targets that will also emit sibling unsigned + // thinking natively. Drop it when the visible thinking was + // cross-model stripped and will be demoted to text. + if (isAnthropicReplay) { + if (isSameModel || isLatestSurvivingAssistant || replaysUnsignedAnthropicThinking) return block; + return []; + } if (isSameModel) return block; return []; } diff --git a/packages/ai/test/anthropic-prior-turn-thinking.test.ts b/packages/ai/test/anthropic-prior-turn-thinking.test.ts new file mode 100644 index 000000000..755468729 --- /dev/null +++ b/packages/ai/test/anthropic-prior-turn-thinking.test.ts @@ -0,0 +1,342 @@ +import { describe, expect, it } from "bun:test"; +import { convertAnthropicMessages } from "@oh-my-pi/pi-ai/providers/anthropic"; +import type { + AssistantMessage, + Message, + Model, + ModelSpec, + ToolResultMessage, + UserMessage, +} from "@oh-my-pi/pi-ai/types"; +import { buildModel } from "@oh-my-pi/pi-catalog/build"; + +/** + * Cross-model `anthropic-messages` continuations must preserve the prior + * turn's reasoning chain. Anthropic enforces an all-or-none contract on + * thinking blocks ("if you include thinking blocks in prior assistant turns, + * you must include ALL thinking blocks (including redacted ones)") but the + * legacy transform only honored that for the LATEST surviving assistant. + * Every earlier turn fell through to the cross-API text-demotion path + * whenever the conversation crossed a model boundary — silently dropping the + * reasoning chain on continuation for custom anthropic-messages providers + * configured via `models.yaml` and for session-level model swaps (#2257). + * + * The signature policy is a second axis: official Anthropic cryptographically + * binds signatures to its key+session+model, so cross-model signatures must + * be stripped (and matching redacted siblings dropped) whenever either side + * of the replay is official Anthropic. Third-party endpoints (Z.AI, DeepSeek, + * custom anthropic-messages providers) treat signatures as opaque + * continuation hints they pass through unchanged, so 3p ↔ 3p replays + * preserve them as-is to keep the reasoning chain signed for the next + * turn (#2265). + */ +function makeAnthropicModel(overrides: Partial> = {}): Model<"anthropic-messages"> { + return buildModel({ + api: "anthropic-messages", + provider: "custom-anthropic", + id: "reasoning-model", + name: "Reasoning Anthropic-Compatible Model", + baseUrl: "https://llm.example.com/anthropic", + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + maxTokens: 8_192, + contextWindow: 200_000, + reasoning: true, + ...overrides, + } as ModelSpec<"anthropic-messages">); +} + +function makeUser(text: string): UserMessage { + return { role: "user", content: text, timestamp: 0 }; +} + +function makeAssistant( + content: AssistantMessage["content"], + overrides: Partial = {}, +): AssistantMessage { + return { + role: "assistant", + content, + api: "anthropic-messages", + provider: "custom-anthropic", + model: "reasoning-model", + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "toolUse", + timestamp: 0, + ...overrides, + }; +} + +function toolResult(toolCallId: string, text: string): ToolResultMessage { + return { + role: "toolResult", + toolCallId, + toolName: "read", + content: [{ type: "text", text }], + isError: false, + timestamp: 0, + }; +} + +interface WireThinkingBlock { + type: "thinking"; + thinking: string; + signature: string; +} +interface WireTextBlock { + type: "text"; + text: string; +} +interface WireRedactedBlock { + type: "redacted_thinking"; + data: string; +} +interface WireToolUseBlock { + type: "tool_use"; + id: string; + name: string; + input: Record; +} +type WireBlock = + | WireThinkingBlock + | WireTextBlock + | WireRedactedBlock + | WireToolUseBlock + | { type: string; [key: string]: unknown }; + +describe("Anthropic prior-turn thinking preservation (#2257, #2265)", () => { + it("preserves the prior thinking block as native `thinking` across compatible endpoints", () => { + // Source v1, target v2, both on the same custom anthropic-messages + // provider. The first assistant turn is PRIOR, so the latest-only + // preservation path doesn't help — without the fix the prior thinking + // block is demoted to plain `text` and the reasoning chain disappears. + const target = makeAnthropicModel({ id: "reasoning-model-v2" }); + const priorThinkingText = "Plan: read README, then summarize."; + const messages: Message[] = [ + makeUser("Summarize README"), + makeAssistant( + [ + { type: "thinking", thinking: priorThinkingText, thinkingSignature: "sig_v1" }, + { type: "toolCall", id: "toolu_prior", name: "read", arguments: { path: "README.md" } }, + ], + { model: "reasoning-model-v1" }, + ), + toolResult("toolu_prior", "README body"), + makeAssistant( + [ + { type: "thinking", thinking: "Got the body, now translating", thinkingSignature: "sig_v2" }, + { type: "text", text: "Voici le résumé en français." }, + ], + { model: "reasoning-model-v2", stopReason: "stop" }, + ), + makeUser("Now translate it to Spanish"), + ]; + + const params = convertAnthropicMessages(messages, target, false); + const assistants = params.filter(p => p.role === "assistant"); + expect(assistants).toHaveLength(2); + const priorBlocks = assistants[0].content as WireBlock[]; + const thinking = priorBlocks.find(b => b.type === "thinking") as WireThinkingBlock | undefined; + expect(thinking).toBeDefined(); + expect(thinking?.thinking).toBe(priorThinkingText); + // 3p ↔ 3p replay: the source signature is opaque continuation metadata + // that compatible endpoints pass through. Stripping it (the pre-fix + // behavior) silently demotes the reasoning chain on the next turn. + expect(thinking?.signature).toBe("sig_v1"); + // And the paired tool_use must still be present right after it. + const toolUse = priorBlocks.find(b => b.type === "tool_use") as WireToolUseBlock | undefined; + expect(toolUse?.id).toBe("toolu_prior"); + }); + + it("keeps the signature on prior turns when the source model matches the target", () => { + // Same provider+api+id throughout: signatures are valid and must ride + // the wire untouched (prompt-cache stability + Anthropic's all-or-none + // invariant). + const target = makeAnthropicModel(); + const messages: Message[] = [ + makeUser("Summarize README"), + makeAssistant([ + { type: "thinking", thinking: "plan", thinkingSignature: "sig_same" }, + { type: "toolCall", id: "toolu_prior", name: "read", arguments: { path: "README.md" } }, + ]), + toolResult("toolu_prior", "README body"), + makeAssistant( + [ + { type: "thinking", thinking: "summarising", thinkingSignature: "sig_latest" }, + { type: "text", text: "summary" }, + ], + { stopReason: "stop" }, + ), + makeUser("And now in Spanish"), + ]; + + const params = convertAnthropicMessages(messages, target, false); + const assistants = params.filter(p => p.role === "assistant"); + const priorBlocks = assistants[0].content as WireBlock[]; + const thinking = priorBlocks.find(b => b.type === "thinking") as WireThinkingBlock | undefined; + expect(thinking?.thinking).toBe("plan"); + expect(thinking?.signature).toBe("sig_same"); + }); + + it("preserves redacted_thinking blocks from prior anthropic-messages turns", () => { + // Anthropic's "include ALL thinking blocks (including redacted ones)" + // rule means redacted_thinking from earlier turns must survive whenever + // any thinking content from the same turn is replayed. + const target = makeAnthropicModel({ id: "reasoning-model-v2" }); + const messages: Message[] = [ + makeUser("Summarize README"), + makeAssistant( + [ + { type: "thinking", thinking: "visible reasoning", thinkingSignature: "sig" }, + { type: "redactedThinking", data: "encrypted-blob" }, + { type: "toolCall", id: "toolu_prior", name: "read", arguments: { path: "README.md" } }, + ], + { model: "reasoning-model-v1" }, + ), + toolResult("toolu_prior", "README body"), + makeAssistant( + [ + { type: "thinking", thinking: "later", thinkingSignature: "sig_latest" }, + { type: "text", text: "summary" }, + ], + { model: "reasoning-model-v2", stopReason: "stop" }, + ), + makeUser("Translate"), + ]; + + const params = convertAnthropicMessages(messages, target, false); + const assistants = params.filter(p => p.role === "assistant"); + const priorBlocks = assistants[0].content as WireBlock[]; + const redacted = priorBlocks.find(b => b.type === "redacted_thinking") as WireRedactedBlock | undefined; + expect(redacted).toBeDefined(); + expect(redacted?.data).toBe("encrypted-blob"); + }); + + it("strips foreign signatures and drops redacted_thinking when the target is official Anthropic", () => { + // 3p → official Anthropic. The official endpoint rejects foreign + // signatures cryptographically, and `replayUnsignedThinking: false` + // demotes the unsigned visible thinking to text downstream, so the + // matching redacted sibling must not remain as a lone native + // redacted_thinking block. + const target = makeAnthropicModel({ + provider: "anthropic", + id: "claude-sonnet-4-6", + baseUrl: "https://api.anthropic.com", + }); + const messages: Message[] = [ + makeUser("Summarize README"), + makeAssistant( + [ + { type: "thinking", thinking: "visible reasoning", thinkingSignature: "sig_custom" }, + { type: "redactedThinking", data: "foreign-encrypted-blob" }, + { type: "toolCall", id: "toolu_prior", name: "read", arguments: { path: "README.md" } }, + ], + { model: "reasoning-model-v1" }, + ), + toolResult("toolu_prior", "README body"), + makeAssistant( + [ + { type: "thinking", thinking: "official latest", thinkingSignature: "sig_latest" }, + { type: "text", text: "summary" }, + ], + { + provider: "anthropic", + model: "claude-sonnet-4-6", + stopReason: "stop", + }, + ), + makeUser("Translate"), + ]; + + const params = convertAnthropicMessages(messages, target, false); + const assistants = params.filter(p => p.role === "assistant"); + const priorBlocks = assistants[0].content as WireBlock[]; + const text = priorBlocks.find(b => b.type === "text") as WireTextBlock | undefined; + expect(text?.text).toBe("visible reasoning"); + expect(priorBlocks.find(b => b.type === "thinking")).toBeUndefined(); + expect(priorBlocks.find(b => b.type === "redacted_thinking")).toBeUndefined(); + }); + + it("strips official Anthropic source signatures on cross-model replay to a 3p target", () => { + // official Anthropic → 3p. Anthropic's signature is bound to the + // issuing model+session, so the 3p target cannot reverify or + // meaningfully continue from it; passing it through would leak + // private continuation metadata for no benefit. The unsigned thinking + // is still emitted natively because the 3p target's compat advertises + // `replayUnsignedThinking: true`. + const target = makeAnthropicModel({ id: "reasoning-model-v2" }); + const messages: Message[] = [ + makeUser("Summarize README"), + makeAssistant( + [ + { type: "thinking", thinking: "anthropic reasoning", thinkingSignature: "sig_anthropic" }, + { type: "toolCall", id: "toolu_prior", name: "read", arguments: { path: "README.md" } }, + ], + { provider: "anthropic", model: "claude-sonnet-4-6" }, + ), + toolResult("toolu_prior", "README body"), + makeAssistant( + [ + { type: "thinking", thinking: "v2 reasoning", thinkingSignature: "sig_v2" }, + { type: "text", text: "summary" }, + ], + { model: "reasoning-model-v2", stopReason: "stop" }, + ), + makeUser("Translate"), + ]; + + const params = convertAnthropicMessages(messages, target, false); + const assistants = params.filter(p => p.role === "assistant"); + const priorBlocks = assistants[0].content as WireBlock[]; + const thinking = priorBlocks.find(b => b.type === "thinking") as WireThinkingBlock | undefined; + expect(thinking?.thinking).toBe("anthropic reasoning"); + expect(thinking?.signature).toBe(""); + }); + + it("does not promote prior unsigned thinking from non-anthropic sources to thinking blocks", () => { + // Cross-API replay: prior turn came from OpenAI-responses with no + // Anthropic signature. The all-or-none rule scope is per-API; we must + // not invent thinking blocks for a turn whose source can't sign them — + // the existing cross-API text demotion is the right behavior. + const target = makeAnthropicModel(); + const messages: Message[] = [ + makeUser("Summarize README"), + makeAssistant( + [ + { type: "thinking", thinking: "openai chain-of-thought", thinkingSignature: "" }, + { type: "toolCall", id: "toolu_prior", name: "read", arguments: { path: "README.md" } }, + ], + { + api: "openai-responses", + provider: "openai", + model: "o1-preview", + } as Partial, + ), + toolResult("toolu_prior", "README body"), + makeAssistant( + [ + { type: "thinking", thinking: "anthropic latest", thinkingSignature: "sig_latest" }, + { type: "text", text: "summary" }, + ], + { stopReason: "stop" }, + ), + makeUser("Translate"), + ]; + + const params = convertAnthropicMessages(messages, target, false); + const assistants = params.filter(p => p.role === "assistant"); + const priorBlocks = assistants[0].content as WireBlock[]; + expect(priorBlocks.find(b => b.type === "thinking")).toBeUndefined(); + // Reasoning text still survives on the wire (as text, via the existing + // cross-API demotion path). + const text = priorBlocks.find(b => b.type === "text") as WireTextBlock | undefined; + expect(text?.text).toBe("openai chain-of-thought"); + }); +}); From dab23f76a424baf5aaf14b7cf2f4170fd21c49ca Mon Sep 17 00:00:00 2001 From: roboomp Date: Wed, 10 Jun 2026 19:27:19 +0000 Subject: [PATCH 25/27] fix(tui): replaced thinking.autoPending question-mark glyphs with loading indicators MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The status-bar auto-thinking pending indicator used `?`-bearing symbols in all three presets (`▣?` Unicode, nf-md-help_box Nerd Font, `[?]` ASCII), which is visually indistinguishable from a terminal missing- glyph fallback. Users assumed the icon was broken and toggled presets to no effect. Replaced each preset with an unambiguous loading indicator: - Unicode: `⟳` (Clockwise Open Circle Arrow) - Nerd Font: `\uf1ce` (fa-circle-o-notch) - ASCII: `[~]` Fixes #2267 --- packages/coding-agent/CHANGELOG.md | 4 ++++ packages/coding-agent/src/modes/theme/theme.ts | 8 ++++---- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index f488065de..ac912b916 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- Fixed the `thinking.autoPending` statusbar indicator using question-mark glyphs (`▣?`, nf-md-help_box, `[?]`) in every symbol preset, which made the auto-thinking pending state indistinguishable from a terminal missing-glyph fallback. Replaced with clear loading indicators (`⟳`, fa-circle-o-notch, `[~]`) ([#2267](https://github.com/can1357/oh-my-pi/issues/2267)). + ## [15.10.12] - 2026-06-10 ### Added diff --git a/packages/coding-agent/src/modes/theme/theme.ts b/packages/coding-agent/src/modes/theme/theme.ts index f9ebba809..99dbf4b5e 100644 --- a/packages/coding-agent/src/modes/theme/theme.ts +++ b/packages/coding-agent/src/modes/theme/theme.ts @@ -328,7 +328,7 @@ const UNICODE_SYMBOLS: SymbolMap = { "thinking.medium": "◒ med", "thinking.high": "◕ high", "thinking.xhigh": "◉ xhigh", - "thinking.autoPending": "▣?", + "thinking.autoPending": "⟳", // Checkboxes "checkbox.checked": "☑", "checkbox.unchecked": "☐", @@ -610,8 +610,8 @@ const NERD_SYMBOLS: SymbolMap = { "thinking.high": "\u{F111} high", // pick: 🧠 xhi | alt:  xhi  xhi "thinking.xhigh": "\u{F06D} xhi", - // pick: 󰞋 (nf-md-help_box) | alt:  [?] - "thinking.autoPending": "\u{f078b}", + // pick: (fa-circle-o-notch) | alt: 󰂼 (nf-md-cached) ⟳ + "thinking.autoPending": "\uf1ce", // Checkboxes // pick:  | alt:   "checkbox.checked": "\uf14a", @@ -814,7 +814,7 @@ const ASCII_SYMBOLS: SymbolMap = { "thinking.medium": "[med]", "thinking.high": "[high]", "thinking.xhigh": "[xhi]", - "thinking.autoPending": "[?]", + "thinking.autoPending": "[~]", // Checkboxes "checkbox.checked": "[x]", "checkbox.unchecked": "[ ]", From 08a941a14e0fe281853deee197b3eae1029ab695 Mon Sep 17 00:00:00 2001 From: can1357 Date: Wed, 10 Jun 2026 21:49:51 +0200 Subject: [PATCH 26/27] feat: added standalone snapcompact package and model-specific frame shaping - Added a new @oh-my-pi/snapcompact package and redirected compaction call sites to it. - Added provider-aware snapcompact shape resolution for model-specific mixed-frame behavior. - Added optional image detail support by extending ImageContent and passing hints through OpenAI providers. - Added native snapcompact render options, including 5x8/8x8 font loading and palette/geometry controls. --- README.md | 4 + bun.lock | 17 + crates/pi-natives/src/fonts/unscii-8.hex | 255 +++++++ crates/pi-natives/src/snapcompact.rs | 501 ++++++++++++-- docs/compaction.md | 8 +- .../examples/hello-extension/package.json | 1 + .../mini-marketplace/my-plugin/package.json | 1 + docs/skills/examples/safety-hook/package.json | 1 + package.json | 2 + packages/agent/CHANGELOG.md | 7 +- packages/agent/package.json | 1 + packages/agent/src/compaction/compaction.ts | 2 +- packages/agent/src/compaction/index.ts | 1 - .../compaction/prompts/snapcompact-summary.md | 6 +- packages/agent/src/compaction/snapcompact.ts | 316 --------- .../agent/test/snapcompact-frames.test.ts | 44 ++ packages/ai/CHANGELOG.md | 4 + .../ai/src/providers/openai-completions.ts | 2 + .../src/providers/openai-responses-shared.ts | 4 +- packages/ai/src/types.ts | 6 + packages/coding-agent/CHANGELOG.md | 1 + .../extensions/with-deps/package.json | 1 + packages/coding-agent/package.json | 1 + .../coding-agent/src/session/agent-session.ts | 6 +- .../src/session/session-manager.ts | 2 +- .../test/job-poll-displacement.test.ts | 12 +- .../test/streaming-preview-height.test.ts | 17 +- packages/natives/CHANGELOG.md | 2 +- packages/natives/native/index.d.ts | 38 +- packages/snapcompact/.gitignore | 2 + packages/snapcompact/CHANGELOG.md | 10 + packages/snapcompact/package.json | 62 ++ .../snapcompact/research/anthropic_api.py | 74 ++ packages/snapcompact/research/bdf.py | 247 +++++++ .../snapcompact/research/exp01_patchalign.py | 284 ++++++++ .../snapcompact/research/exp02_surprisal.py | 427 ++++++++++++ .../snapcompact/research/exp03_numhard.py | 382 +++++++++++ packages/snapcompact/research/exp04_layout.py | 374 ++++++++++ .../snapcompact/research/exp05_anchors.py | 344 ++++++++++ .../snapcompact/research/exp06_rolecolor.py | 417 ++++++++++++ .../snapcompact/research/exp07_readtax.py | 273 ++++++++ .../snapcompact/research/exp08_foveate.py | 388 +++++++++++ .../snapcompact/research/exp09_cacheappend.py | 365 ++++++++++ .../snapcompact/research/exp10_profiles.py | 358 ++++++++++ .../snapcompact/research/exp11_memhier.py | 301 +++++++++ .../snapcompact/research/exp12_arbitrage.py | 320 +++++++++ .../snapcompact/research/exp13_extractive.py | 255 +++++++ .../snapcompact/research/exp14_bestgpt.py | 475 +++++++++++++ .../snapcompact/research/exp15_bestgemini.py | 478 +++++++++++++ .../snapcompact/research/exp16_bestfable.py | 476 +++++++++++++ .../snapcompact/research/exp17_bestopus.py | 419 ++++++++++++ .../snapcompact/research/exp18_bestkimi.py | 439 ++++++++++++ .../snapcompact/research/exp19_bestglm.py | 477 +++++++++++++ packages/snapcompact/research/exp20_8x8u.py | 362 ++++++++++ .../snapcompact/research/exp21_braille.py | 276 ++++++++ packages/snapcompact/research/exp22_ttf6pt.py | 243 +++++++ packages/snapcompact/research/final.py | 331 +++++++++ packages/snapcompact/research/mono.py | 154 +++++ .../research/prompts/exp02-qa-image.md | 8 + .../research/prompts/exp04-qa-image.md | 8 + .../research/prompts/exp05-qa-image-ctl.md | 10 + .../research/prompts/exp05-qa-image.md | 13 + .../research/prompts/exp06-prov-image.md | 6 + .../research/prompts/exp06-qa-image-tag.md | 6 + .../research/prompts/exp06-qa-image.md | 6 + .../research/prompts/exp07-answer-bands.md | 4 + .../research/prompts/exp07-locate.md | 5 + .../research/prompts/exp07-qa-image.md | 7 + .../research/prompts/exp08-archive-eager.md | 10 + .../research/prompts/exp08-archive-phrase.md | 9 + .../research/prompts/exp08-archive.md | 7 + .../research/prompts/exp08-zoom.md | 6 + .../research/prompts/exp09-frame.md | 3 + .../research/prompts/exp09-page.md | 1 + .../snapcompact/research/prompts/exp09-qa.md | 6 + .../research/prompts/exp11-qa-hier.md | 12 + .../research/prompts/exp13-extract.md | 10 + .../research/prompts/exp19-qa-doc.md | 8 + .../research/prompts/exp21-qa-braille.md | 11 + .../research/prompts/qa-image-cols.md | 6 + .../research/prompts/qa-image-multi.md | 6 + .../snapcompact/research/prompts/qa-image.md | 6 + .../research/prompts/qa-remote-compact.md | 8 + .../snapcompact/research/prompts/qa-text.md | 10 + .../research/prompts/session-frame.md | 5 + packages/snapcompact/research/providers.py | 261 +++++++ packages/snapcompact/research/run.py | 353 ++++++++++ .../snapcompact_3d_activation_html.py | 192 ++++++ .../research/snapcompact_3d_activation_viz.py | 230 +++++++ .../research/snapcompact_activation_probe.py | 225 ++++++ .../snapcompact_blackbox_occlusion.py | 261 +++++++ .../research/snapcompact_blog_viz.py | 231 +++++++ .../snapcompact_carrier_convergence.py | 212 ++++++ .../research/snapcompact_convergence_3d.py | 187 +++++ .../snapcompact_convergence_extras.py | 207 ++++++ .../research/snapcompact_convergence_viz.py | 219 ++++++ .../snapcompact_lockon_anatomy_viz.py | 288 ++++++++ .../research/snapcompact_logit_lens_dump.py | 143 ++++ .../research/snapcompact_logit_lens_viz.py | 169 +++++ .../research/snapcompact_materialize_sweep.py | 226 +++++++ .../research/snapcompact_materialize_viz.py | 171 +++++ .../research/snapcompact_pricing_viz.py | 137 ++++ .../snapcompact_qwen_control_intervention.py | 363 ++++++++++ .../snapcompact_qwen_spotlight_viz.py | 278 ++++++++ .../research/snapcompact_r2_chord.py | 265 ++++++++ .../research/snapcompact_r2_crystal.py | 338 ++++++++++ .../research/snapcompact_r2_filmstrip.py | 256 +++++++ .../research/snapcompact_r2_hero.py | 497 ++++++++++++++ .../research/snapcompact_r2_metro.py | 348 ++++++++++ .../research/snapcompact_tensor_heatmap.py | 354 ++++++++++ .../research/snapcompact_text_image_3d_viz.py | 211 ++++++ .../snapcompact_text_image_compare.py | 489 ++++++++++++++ .../research/snapcompact_token_entry_dump.py | 158 +++++ .../research/snapcompact_token_entry_viz.py | 225 ++++++ .../research/snapcompact_viz_atlas.py | 377 +++++++++++ .../research/snapcompact_viz_circuit.py | 427 ++++++++++++ .../research/snapcompact_viz_city.py | 237 +++++++ .../research/snapcompact_viz_explainer.py | 444 ++++++++++++ .../research/snapcompact_viz_glass_stack.py | 255 +++++++ .../research/snapcompact_viz_glyph_matrix.py | 397 +++++++++++ .../research/snapcompact_viz_radial.py | 219 ++++++ .../research/snapcompact_viz_token_grid.py | 396 +++++++++++ .../research/snapcompact_viz_volume.py | 235 +++++++ .../research/snapcompact_viz_waterfall.py | 301 +++++++++ packages/snapcompact/research/squad.py | 120 ++++ packages/snapcompact/src/index.ts | 1 + .../src/prompts/file-operations.md | 10 + .../src/prompts/snapcompact-summary.md | 17 + packages/snapcompact/src/snapcompact.ts | 638 ++++++++++++++++++ .../test/snapcompact.test.ts | 191 ++++-- packages/snapcompact/tsconfig.json | 7 + packages/snapcompact/tsconfig.publish.json | 25 + .../typescript-edit-benchmark/package.json | 1 + packages/utils/test/prompt.test.ts | 2 +- python/robomp/web/package.json | 1 + scripts/ci-release-publish.ts | 1 + scripts/install-tests/run-ci.sh | 6 +- 137 files changed, 21786 insertions(+), 489 deletions(-) create mode 100644 crates/pi-natives/src/fonts/unscii-8.hex delete mode 100644 packages/agent/src/compaction/snapcompact.ts create mode 100644 packages/agent/test/snapcompact-frames.test.ts create mode 100644 packages/snapcompact/.gitignore create mode 100644 packages/snapcompact/CHANGELOG.md create mode 100644 packages/snapcompact/package.json create mode 100644 packages/snapcompact/research/anthropic_api.py create mode 100644 packages/snapcompact/research/bdf.py create mode 100644 packages/snapcompact/research/exp01_patchalign.py create mode 100644 packages/snapcompact/research/exp02_surprisal.py create mode 100644 packages/snapcompact/research/exp03_numhard.py create mode 100644 packages/snapcompact/research/exp04_layout.py create mode 100644 packages/snapcompact/research/exp05_anchors.py create mode 100644 packages/snapcompact/research/exp06_rolecolor.py create mode 100644 packages/snapcompact/research/exp07_readtax.py create mode 100644 packages/snapcompact/research/exp08_foveate.py create mode 100644 packages/snapcompact/research/exp09_cacheappend.py create mode 100644 packages/snapcompact/research/exp10_profiles.py create mode 100644 packages/snapcompact/research/exp11_memhier.py create mode 100644 packages/snapcompact/research/exp12_arbitrage.py create mode 100644 packages/snapcompact/research/exp13_extractive.py create mode 100644 packages/snapcompact/research/exp14_bestgpt.py create mode 100644 packages/snapcompact/research/exp15_bestgemini.py create mode 100644 packages/snapcompact/research/exp16_bestfable.py create mode 100644 packages/snapcompact/research/exp17_bestopus.py create mode 100644 packages/snapcompact/research/exp18_bestkimi.py create mode 100644 packages/snapcompact/research/exp19_bestglm.py create mode 100644 packages/snapcompact/research/exp20_8x8u.py create mode 100644 packages/snapcompact/research/exp21_braille.py create mode 100644 packages/snapcompact/research/exp22_ttf6pt.py create mode 100644 packages/snapcompact/research/final.py create mode 100644 packages/snapcompact/research/mono.py create mode 100644 packages/snapcompact/research/prompts/exp02-qa-image.md create mode 100644 packages/snapcompact/research/prompts/exp04-qa-image.md create mode 100644 packages/snapcompact/research/prompts/exp05-qa-image-ctl.md create mode 100644 packages/snapcompact/research/prompts/exp05-qa-image.md create mode 100644 packages/snapcompact/research/prompts/exp06-prov-image.md create mode 100644 packages/snapcompact/research/prompts/exp06-qa-image-tag.md create mode 100644 packages/snapcompact/research/prompts/exp06-qa-image.md create mode 100644 packages/snapcompact/research/prompts/exp07-answer-bands.md create mode 100644 packages/snapcompact/research/prompts/exp07-locate.md create mode 100644 packages/snapcompact/research/prompts/exp07-qa-image.md create mode 100644 packages/snapcompact/research/prompts/exp08-archive-eager.md create mode 100644 packages/snapcompact/research/prompts/exp08-archive-phrase.md create mode 100644 packages/snapcompact/research/prompts/exp08-archive.md create mode 100644 packages/snapcompact/research/prompts/exp08-zoom.md create mode 100644 packages/snapcompact/research/prompts/exp09-frame.md create mode 100644 packages/snapcompact/research/prompts/exp09-page.md create mode 100644 packages/snapcompact/research/prompts/exp09-qa.md create mode 100644 packages/snapcompact/research/prompts/exp11-qa-hier.md create mode 100644 packages/snapcompact/research/prompts/exp13-extract.md create mode 100644 packages/snapcompact/research/prompts/exp19-qa-doc.md create mode 100644 packages/snapcompact/research/prompts/exp21-qa-braille.md create mode 100644 packages/snapcompact/research/prompts/qa-image-cols.md create mode 100644 packages/snapcompact/research/prompts/qa-image-multi.md create mode 100644 packages/snapcompact/research/prompts/qa-image.md create mode 100644 packages/snapcompact/research/prompts/qa-remote-compact.md create mode 100644 packages/snapcompact/research/prompts/qa-text.md create mode 100644 packages/snapcompact/research/prompts/session-frame.md create mode 100644 packages/snapcompact/research/providers.py create mode 100644 packages/snapcompact/research/run.py create mode 100644 packages/snapcompact/research/snapcompact_3d_activation_html.py create mode 100644 packages/snapcompact/research/snapcompact_3d_activation_viz.py create mode 100644 packages/snapcompact/research/snapcompact_activation_probe.py create mode 100644 packages/snapcompact/research/snapcompact_blackbox_occlusion.py create mode 100644 packages/snapcompact/research/snapcompact_blog_viz.py create mode 100644 packages/snapcompact/research/snapcompact_carrier_convergence.py create mode 100644 packages/snapcompact/research/snapcompact_convergence_3d.py create mode 100644 packages/snapcompact/research/snapcompact_convergence_extras.py create mode 100644 packages/snapcompact/research/snapcompact_convergence_viz.py create mode 100644 packages/snapcompact/research/snapcompact_lockon_anatomy_viz.py create mode 100644 packages/snapcompact/research/snapcompact_logit_lens_dump.py create mode 100644 packages/snapcompact/research/snapcompact_logit_lens_viz.py create mode 100644 packages/snapcompact/research/snapcompact_materialize_sweep.py create mode 100644 packages/snapcompact/research/snapcompact_materialize_viz.py create mode 100644 packages/snapcompact/research/snapcompact_pricing_viz.py create mode 100644 packages/snapcompact/research/snapcompact_qwen_control_intervention.py create mode 100644 packages/snapcompact/research/snapcompact_qwen_spotlight_viz.py create mode 100755 packages/snapcompact/research/snapcompact_r2_chord.py create mode 100755 packages/snapcompact/research/snapcompact_r2_crystal.py create mode 100755 packages/snapcompact/research/snapcompact_r2_filmstrip.py create mode 100755 packages/snapcompact/research/snapcompact_r2_hero.py create mode 100755 packages/snapcompact/research/snapcompact_r2_metro.py create mode 100644 packages/snapcompact/research/snapcompact_tensor_heatmap.py create mode 100644 packages/snapcompact/research/snapcompact_text_image_3d_viz.py create mode 100644 packages/snapcompact/research/snapcompact_text_image_compare.py create mode 100644 packages/snapcompact/research/snapcompact_token_entry_dump.py create mode 100644 packages/snapcompact/research/snapcompact_token_entry_viz.py create mode 100644 packages/snapcompact/research/snapcompact_viz_atlas.py create mode 100644 packages/snapcompact/research/snapcompact_viz_circuit.py create mode 100644 packages/snapcompact/research/snapcompact_viz_city.py create mode 100644 packages/snapcompact/research/snapcompact_viz_explainer.py create mode 100644 packages/snapcompact/research/snapcompact_viz_glass_stack.py create mode 100644 packages/snapcompact/research/snapcompact_viz_glyph_matrix.py create mode 100644 packages/snapcompact/research/snapcompact_viz_radial.py create mode 100644 packages/snapcompact/research/snapcompact_viz_token_grid.py create mode 100644 packages/snapcompact/research/snapcompact_viz_volume.py create mode 100755 packages/snapcompact/research/snapcompact_viz_waterfall.py create mode 100644 packages/snapcompact/research/squad.py create mode 100644 packages/snapcompact/src/index.ts create mode 100644 packages/snapcompact/src/prompts/file-operations.md create mode 100644 packages/snapcompact/src/prompts/snapcompact-summary.md create mode 100644 packages/snapcompact/src/snapcompact.ts rename packages/{agent => snapcompact}/test/snapcompact.test.ts (64%) create mode 100644 packages/snapcompact/tsconfig.json create mode 100644 packages/snapcompact/tsconfig.publish.json diff --git a/README.md b/README.md index a9fbe3395..ac6ec2743 100644 --- a/README.md +++ b/README.md @@ -482,12 +482,16 @@ For architecture and contribution guidelines, see [packages/coding-agent/DEVELOP | Package | Description | | --------------------------------------------------------- | -------------------------------------------------------------------------- | | **[@oh-my-pi/pi-ai](packages/ai)** | Multi-provider LLM client with streaming and model/provider integration | +| **[@oh-my-pi/pi-catalog](packages/catalog)** | Model catalog: bundled model database, provider descriptors, and identity | | **[@oh-my-pi/pi-agent-core](packages/agent)** | Agent runtime with tool calling and state management | | **[@oh-my-pi/pi-coding-agent](packages/coding-agent)** | Interactive coding agent CLI and SDK | | **[@oh-my-pi/pi-tui](packages/tui)** | Terminal UI library with differential rendering | | **[@oh-my-pi/pi-natives](packages/natives)** | N-API bindings for grep, shell, image, text, syntax highlighting, and more | | **[@oh-my-pi/omp-stats](packages/stats)** | Local observability dashboard for AI usage statistics | | **[@oh-my-pi/pi-utils](packages/utils)** | Shared utilities (logging, streams, dirs/env/process helpers) | +| **[@oh-my-pi/hashline](packages/hashline)** | Line-anchored patch language and applier behind the `edit` tool | +| **[@oh-my-pi/pi-mnemopi](packages/mnemopi)** | Local SQLite memory engine for Oh My Pi agents | +| **[@oh-my-pi/pi-snapcompact](packages/snapcompact)** | SQuAD eval suite for snapcompact bitmap-frame context compression | | **[@oh-my-pi/swarm-extension](packages/swarm-extension)** | Swarm orchestration extension package | ### Rust Crates diff --git a/bun.lock b/bun.lock index d097dee50..5fe030003 100644 --- a/bun.lock +++ b/bun.lock @@ -21,6 +21,7 @@ "@oh-my-pi/pi-catalog": "catalog:", "@oh-my-pi/pi-natives": "catalog:", "@oh-my-pi/pi-utils": "catalog:", + "@oh-my-pi/snapcompact": "catalog:", "@opentelemetry/api": "catalog:", }, "devDependencies": { @@ -76,6 +77,7 @@ "@oh-my-pi/pi-natives": "catalog:", "@oh-my-pi/pi-tui": "catalog:", "@oh-my-pi/pi-utils": "catalog:", + "@oh-my-pi/snapcompact": "catalog:", "@opentelemetry/api": "catalog:", "@opentelemetry/context-async-hooks": "catalog:", "@opentelemetry/exporter-trace-otlp-proto": "catalog:", @@ -141,6 +143,18 @@ "@types/bun": "catalog:", }, }, + "packages/snapcompact": { + "name": "@oh-my-pi/snapcompact", + "version": "15.10.12", + "dependencies": { + "@oh-my-pi/pi-ai": "catalog:", + "@oh-my-pi/pi-natives": "catalog:", + "@oh-my-pi/pi-utils": "catalog:", + }, + "devDependencies": { + "@types/bun": "catalog:", + }, + }, "packages/stats": { "name": "@oh-my-pi/omp-stats", "version": "15.10.12", @@ -276,6 +290,7 @@ "@oh-my-pi/pi-natives": "15.10.12", "@oh-my-pi/pi-tui": "15.10.12", "@oh-my-pi/pi-utils": "15.10.12", + "@oh-my-pi/snapcompact": "15.10.12", "@opentelemetry/api": "^1.9.1", "@opentelemetry/context-async-hooks": "^2.7.1", "@opentelemetry/exporter-trace-otlp-proto": "^0.218.0", @@ -672,6 +687,8 @@ "@oh-my-pi/pi-utils": ["@oh-my-pi/pi-utils@workspace:packages/utils"], + "@oh-my-pi/snapcompact": ["@oh-my-pi/snapcompact@workspace:packages/snapcompact"], + "@oh-my-pi/swarm-extension": ["@oh-my-pi/swarm-extension@workspace:packages/swarm-extension"], "@oh-my-pi/typescript-edit-benchmark": ["@oh-my-pi/typescript-edit-benchmark@workspace:packages/typescript-edit-benchmark"], diff --git a/crates/pi-natives/src/fonts/unscii-8.hex b/crates/pi-natives/src/fonts/unscii-8.hex new file mode 100644 index 000000000..cc21a2b06 --- /dev/null +++ b/crates/pi-natives/src/fonts/unscii-8.hex @@ -0,0 +1,255 @@ +00001:E080EA2AEE0A0A00 +00002:E080EA2AE40A0A00 +00003:E080CA8AE40A0A00 +00004:E080CE84E4040400 +00005:E080CE8AEA0E0400 +00006:E0A0EAAAAC0A0A00 +00007:C0A0C8A8C8080E00 +00008:C0A0CEA8CE020E00 +00009:A0A0EEA4A4040400 +0000A:80808E88EC080800 +0000B:A0A0AEA444040400 +0000C:E080CE888C080800 +0000D:E0808E8AEE0C0A00 +0000E:E080EE2AEA0A0E00 +0000F:E080EE24E4040E00 +00010:C0A0A8A8C8080E00 +00011:C0A0A4ACC4040E00 +00012:C0A0AEA2CE080E00 +00013:C0A0AEA2C6020E00 +00014:C0A0AAAACE020200 +00015:E0A0AAAAAC0A0A00 +00016:E080EA2AEE040400 +00017:E080CC8AEC0A0C00 +00018:E0808E8AEA0A0A00 +00019:E080CA8EEA0A0A00 +0001A:3C66663018001800 +0001B:E080CE88E8080E00 +0001C:E080CE888E020E00 +0001D:E080AEA8EE020E00 +0001E:E0A0EEC8AE020E00 +0001F:A0A0AEA8EE020E00 +00020:0000000000000000 +00021:1818181818001800 +00022:6666660000000000 +00023:6C6CFE6CFE6C6C00 +00024:183E603C067C1800 +00025:00C6CC183066C600 +00026:386C3876DCCC7600 +00027:1818300000000000 +00028:0C18303030180C00 +00029:30180C0C0C183000 +0002A:00663CFF3C660000 +0002B:0018187E18180000 +0002C:0000000000181830 +0002D:0000007E00000000 +0002E:0000000000181800 +0002F:03060C183060C000 +00030:3C666E7666663C00 +00031:1838181818187E00 +00032:3C660C1830607E00 +00033:3C66061C06663C00 +00034:1C3C6CCCFE0C0C00 +00035:7E607C0606663C00 +00036:1C30607C66663C00 +00037:7E06060C18181800 +00038:3C66663C66663C00 +00039:3C66663E060C3800 +0003A:0018180000181800 +0003B:0018180000181830 +0003C:0C18306030180C00 +0003D:00007E007E000000 +0003E:6030180C18306000 +0003F:3C66060C18001800 +00040:7CC6DEDEDEC07C00 +00041:183C66667E666600 +00042:7C66667C66667C00 +00043:3C66606060663C00 +00044:786C6666666C7800 +00045:7E60607C60607E00 +00046:7E60607C60606000 +00047:3C66606E66663E00 +00048:6666667E66666600 +00049:7E18181818187E00 +0004A:0606060606663C00 +0004B:C6CCD8F0D8CCC600 +0004C:6060606060607E00 +0004D:C6EEFED6C6C6C600 +0004E:C6E6F6DECEC6C600 +0004F:3C66666666663C00 +00050:7C66667C60606000 +00051:3C666666666C3600 +00052:7C66667C6C666600 +00053:3C66603C06663C00 +00054:7E18181818181800 +00055:6666666666663C00 +00056:66666666663C1800 +00057:C6C6C6D6FEEEC600 +00058:C3663C183C66C300 +00059:C3663C1818181800 +0005A:7E060C1830607E00 +0005B:3C30303030303C00 +0005C:C06030180C060300 +0005D:3C0C0C0C0C0C3C00 +0005E:10386CC600000000 +0005F:00000000000000FF +00060:180C060000000000 +00061:00003C063E663E00 +00062:60607C6666667C00 +00063:00003C6060603C00 +00064:06063E6666663E00 +00065:00003C667E603C00 +00066:1C307C3030303000 +00067:00003E66663E067C +00068:60607C6666666600 +00069:1800381818181E00 +0006A:0C000C0C0C0C0C78 +0006B:6060666C786C6600 +0006C:3818181818181E00 +0006D:0000CCFED6D6C600 +0006E:00007C6666666600 +0006F:00003C6666663C00 +00070:00007C66667C6060 +00071:00003E66663E0606 +00072:00007C6660606000 +00073:00003E603C067C00 +00074:30307E3030301E00 +00075:0000666666663E00 +00076:00006666663C1800 +00077:0000C6C6D67C6C00 +00078:0000C66C386CC600 +00079:00006666663E063C +0007A:00007E0C18307E00 +0007B:0E18187018180E00 +0007C:1818181818181800 +0007D:7018180E18187000 +0007E:76DC000000000000 +0007F:C0A0AEA4C4040400 +00080:E0A0EEAAEA0A0E00 +00081:E0A0E4ACE4040E00 +00082:E0A0EEA2EE080E00 +00083:E0A0EEA2EE020E00 +00084:E0404E4AEA0A0A00 +00085:E0A0A8A8A8080E00 +00086:E080EE28EE020E00 +00087:E080CE88EE020E00 +00088:A0A0EEA8AE020E00 +00089:A0A0E2A2A20A0E00 +0008A:A0A0AEA84E020E00 +0008B:E0A0EC8A8A0A0C00 +0008C:E0A0EA8A8A0A0E00 +0008D:C0A0CEA4A4040E00 +0008E:E080EE22EE080E00 +0008F:E080EE22E6020E00 +00090:C0A0AEA8C8080E00 +00091:E0A0E48C84040E00 +00092:E0A0EE828E080E00 +00093:E080EE28EC080E00 +00094:E0808E88E8080E00 +00095:A0E0EAAAAE0E0A00 +00096:E080EE2AEE080800 +00097:E080CE8AEE080800 +00098:E0A0EE2AEE0A0E00 +00099:E0A0EE2AEE020E00 +0009A:E0A0EE2AEE0A0A00 +0009B:E0808E88EE020E00 +0009C:E080EE24E4040400 +0009D:E0A0AEA8EE020E00 +0009E:E0A0EA8E8E0A0A00 +0009F:E0A0EEAAAE080800 +000A0:0000000000000000 +000A1:1800181818181800 +000A2:00187ED8D87E1800 +000A3:386C60F06066FC00 +000A4:663C663C66000000 +000A5:C3663C183C181800 +000A6:1818180018181800 +000A7:3C603C66663C063C +000A8:6600000000000000 +000A9:7E819DB1B19D817E +000AA:3C6C6C3E007E0000 +000AB:003366CC66330000 +000AC:007E060600000000 +000AD:0000003C00000000 +000AE:7E81B9A5B9A5817E +000AF:7E00000000000000 +000B0:3C663C0000000000 +000B1:18187E1818007E00 +000B2:7018306078000000 +000B3:780C180C78000000 +000B4:0C18300000000000 +000B5:00006666667C60C0 +000B6:3E7A7A3A1A1A1A00 +000B7:0000001800000000 +000B8:0000000000000818 +000B9:3070303030000000 +000BA:386C6C38007C0000 +000BB:00CC663366CC0000 +000BC:40C64C5A366ACF02 +000BD:40C64C5E3366CC0F +000BE:C0266C3AF66ACF02 +000BF:1800183060663C00 +000C0:70003C667E666600 +000C1:0E003C667E666600 +000C2:1866003C667E6600 +000C3:76DC003C667E6600 +000C4:66003C667E666600 +000C5:1818003C667E6600 +000C6:3F6CCCFECCCCCF00 +000C7:3C66606060663C18 +000C8:7000FEC0F8C0FE00 +000C9:0E00FEC0F8C0FE00 +000CA:186600FEF0C0FE00 +000CB:6600FEC0F8C0FE00 +000CC:70007E1818187E00 +000CD:0E007E1818187E00 +000CE:1866007E18187E00 +000CF:66007E1818187E00 +000D0:786C66F6666C7800 +000D1:76DC00C6F6DEC600 +000D2:70007CC6C6C67C00 +000D3:0E007CC6C6C67C00 +000D4:1866007CC6C67C00 +000D5:76DC007CC6C67C00 +000D6:66007CC6C6C67C00 +000D7:00C66C386CC60000 +000D8:3E666E7E76667C00 +000D9:7000C6C6C6C67C00 +000DA:0E00C6C6C6C67C00 +000DB:186600C6C6C67C00 +000DC:6600C6C6C6C67C00 +000DD:0E0066663C181800 +000DE:C0C0FCC6FCC0C000 +000DF:3C66666C66666C00 +000E0:70003C063E663E00 +000E1:0E003C063E663E00 +000E2:1866003E66C67E00 +000E3:76DC003E66C67E00 +000E4:66003C063E663E00 +000E5:1818003E66C67E00 +000E6:00007E1B7FD87700 +000E7:00003C6060603C18 +000E8:70003C667E603C00 +000E9:0E003C667E603C00 +000EA:1866003C7E603C00 +000EB:66003C667E603C00 +000EC:7000381818183C00 +000ED:0E00381818183C00 +000EE:1866003818183C00 +000EF:6600381818183C00 +000F0:0C3E0C7CCCCC7800 +000F1:76DC007C66666600 +000F2:70003C6666663C00 +000F3:0E003C6666663C00 +000F4:1866003C66663C00 +000F5:76DC003C66663C00 +000F6:66003C6666663C00 +000F7:1818007E00181800 +000F8:00027CCED6E67C80 +000F9:7000666666663E00 +000FA:0E00666666663E00 +000FB:1866006666663E00 +000FC:6600666666663E00 +000FD:0E006666663E063C +000FE:60607C66667C6060 +000FF:66006666663E063C diff --git a/crates/pi-natives/src/snapcompact.rs b/crates/pi-natives/src/snapcompact.rs index e1da346a4..3805a8b08 100644 --- a/crates/pi-natives/src/snapcompact.rs +++ b/crates/pi-natives/src/snapcompact.rs @@ -1,46 +1,61 @@ //! Snapcompact frame rendering. //! -//! Rasterizes pre-normalized conversation text onto a square 8-bit indexed -//! bitmap using the bundled public-domain X.org `5x8` BDF font, with glyph -//! ink cycling through six hues at sentence boundaries, then encodes the -//! bitmap as an indexed PNG. +//! Rasterizes pre-normalized conversation text onto a square bitmap using one +//! of the bundled public-domain pixel fonts, then encodes it as PNG: //! -//! Text normalization, frame chunking, and archive management live in -//! `packages/agent/src/compaction/snapcompact.ts`; this module is only the -//! hot `text -> PNG bytes` path. +//! - `5x8` — X.org BDF font (legacy shape). +//! - `8x8` — unscii-8 hex font (Latin-1 subset), the square cell that won the +//! snapcompact `SQuAD` evals. +//! +//! Shape controls, all eval-validated in `packages/snapcompact`: +//! +//! - **variant** — `sent` cycles glyph ink through six hues at sentence +//! boundaries; `bw` prints plain black ink (best for Anthropic readers). +//! - **lineRepeat** — prints every text line N times; copies after the first +//! sit on a pale highlight band. Redundancy coding: two looks per glyph at +//! half the density ("8x8r" shapes). +//! - **cellWidth/cellHeight** — target cell size. When it differs from the +//! font's natural cell, glyphs are rasterized at native size and the canvas +//! is Lanczos3-resampled to the target (anisotropic stretch, e.g. the +//! OpenAI-optimal "6x6u" shape), producing an anti-aliased RGB frame. +//! +//! Text normalization, frame chunking, provider shape selection, and archive +//! management live in `packages/agent/src/compaction/snapcompact.ts`; this +//! module is only the hot `text -> PNG bytes` path. -use std::{borrow::Cow, collections::HashMap, sync::LazyLock}; +use std::{borrow::Cow, collections::HashMap, f32::consts::PI, sync::LazyLock}; use napi::bindgen_prelude::*; use napi_derive::napi; -/// Horizontal advance per glyph cell in pixels. -const GLYPH_ADVANCE_X: usize = 5; -/// Vertical pitch per text row in pixels. -const GLYPH_PITCH_Y: usize = 8; - /// Upper bound on the frame edge: a hard stop against absurd allocations /// (`size * size` pixel buffer), far above the 2576px production frame. const MAX_FRAME_SIZE: u32 = 16384; -/// Sentence ink palette: index 0 is the white background, 1-6 are the six -/// dark hues from the eval renderer (HLS l=0.22 s=0.95, h ∈ {0, .08, .3, -/// .5, .62, .78}), pre-baked to sRGB. -const PALETTE: [[u8; 3]; 7] = [ +/// Indexed palette: 0 is the white background, 1-6 are the six dark sentence +/// hues from the eval renderer (HLS l=0.22 s=0.95, h ∈ {0, .08, .3, .5, .62, +/// .78}), 7 is plain black ink (`bw` variant), 8 is the pale highlight band +/// behind repeated line copies. +const PALETTE: [[u8; 3]; 9] = [ [255, 255, 255], - [109, 2, 2], // red - [109, 53, 2], // amber - [24, 109, 2], // green - [2, 109, 109], // teal - [2, 32, 109], // blue - [75, 2, 109], // violet + [109, 2, 2], // red + [109, 53, 2], // amber + [24, 109, 2], // green + [2, 109, 109], // teal + [2, 32, 109], // blue + [75, 2, 109], // violet + [0, 0, 0], // bw ink + [255, 247, 194], // repeat highlight band ]; -const INK_COLORS: usize = PALETTE.len() - 1; +const INK_COLORS: usize = 6; +const INK_BLACK: u8 = 7; +const BG_REPEAT: u8 = 8; -static FONT: LazyLock = LazyLock::new(|| parse_bdf(include_str!("fonts/5x8.bdf"))); +static FONT_5X8: LazyLock = LazyLock::new(|| parse_bdf(include_str!("fonts/5x8.bdf"), 5, 8)); +static FONT_8X8: LazyLock = LazyLock::new(|| parse_hex(include_str!("fonts/unscii-8.hex"))); struct Glyph { - /// Glyph width in pixels (≤ 8 for this font). + /// Glyph width in pixels (≤ 8 for the bundled fonts). w: u8, /// Glyph height in pixels. h: i32, @@ -54,9 +69,13 @@ struct Font { /// Glyphs keyed by Unicode code point (ASCII + Latin-1 coverage). glyphs: HashMap, ascent: i32, + /// Natural cell advance (x) in pixels. + cell_w: usize, + /// Natural cell pitch (y) in pixels. + cell_h: usize, } -fn parse_bdf(text: &str) -> Font { +fn parse_bdf(text: &str, cell_w: usize, cell_h: usize) -> Font { let mut glyphs = HashMap::new(); let mut ascent = 0i32; let mut enc = -1i64; @@ -91,27 +110,88 @@ fn parse_bdf(text: &str) -> Font { } } } - Font { glyphs, ascent } + Font { glyphs, ascent, cell_w, cell_h } } -/// Rasterize `text` onto a `size` x `size` palette-indexed bitmap, row-major -/// with no word wrap. Ink color advances at sentence boundaries (terminator -/// in `.!?` followed by a space). Characters beyond the grid capacity are -/// ignored; code points missing from the font leave their cell blank. -fn render_bitmap(text: &str, size: usize, font: &Font) -> Vec { - let cols = size / GLYPH_ADVANCE_X; - let grid_rows = size / GLYPH_PITCH_Y; - let capacity = cols * grid_rows; - let mut pixels = vec![0u8; size * size]; // 0 = white background +/// Parse a unifont-style `.hex` font (`CODEPOINT:16-hex-digit bitmap`, one +/// byte per row of an 8x8 glyph). Baseline sits at row 7 (`ascent` 7 with a +/// one-pixel descender row), matching the eval renderer. +fn parse_hex(text: &str) -> Font { + let mut glyphs = HashMap::new(); + for line in text.lines() { + let Some((cp, bits)) = line.split_once(':') else { + continue; + }; + let Ok(enc) = u32::from_str_radix(cp.trim(), 16) else { + continue; + }; + let bits = bits.trim(); + if bits.len() != 16 { + continue; + } + let rows: Vec = (0..8) + .map(|i| u8::from_str_radix(&bits[i * 2..i * 2 + 2], 16).unwrap_or(0)) + .collect(); + glyphs.insert(enc, Glyph { w: 8, h: 8, xoff: 0, yoff: -1, rows }); + } + Font { glyphs, ascent: 7, cell_w: 8, cell_h: 8 } +} + +fn resolve_font(name: &str) -> Option<&'static Font> { + match name { + "5x8" => Some(&FONT_5X8), + "8x8" => Some(&FONT_8X8), + _ => None, + } +} + +/// Frame grid geometry shared with the TypeScript caller. +struct Grid { + cols: usize, + rows: usize, + repeat: usize, +} + +/// Rasterize `text` onto a `width` x `height` palette-indexed bitmap at the +/// font's natural cell size, row-major with no word wrap. Each text line is +/// printed `grid.repeat` times; copies after the first sit on the highlight +/// band. Ink cycles through six hues at sentence boundaries (terminator in +/// `.!?` followed by a space) unless `black_ink` pins it to black. Characters +/// beyond `cols * rows` are ignored; code points missing from the font leave +/// their cell blank. +fn render_bitmap( + text: &str, + width: usize, + height: usize, + font: &Font, + grid: &Grid, + black_ink: bool, +) -> Vec { + let mut pixels = vec![0u8; width * height]; // 0 = white background + let capacity = grid.cols * grid.rows; if capacity == 0 { return pixels; } + if grid.repeat > 1 { + for row in 0..grid.rows { + for copy in 1..grid.repeat { + let band_top = (row * grid.repeat + copy) * font.cell_h; + for y in band_top..(band_top + font.cell_h).min(height) { + pixels[y * width..y * width + width].fill(BG_REPEAT); + } + } + } + } let codes: Vec = text.chars().map(|ch| ch as u32).collect(); let count = codes.len().min(capacity); let mut sentence = 0usize; for i in 0..count { let code = codes[i]; - let ink = (1 + sentence % INK_COLORS) as u8; + let ink = if black_ink { + INK_BLACK + } else { + (1 + sentence % INK_COLORS) as u8 + }; if matches!(code, 0x2e | 0x21 | 0x3f) && codes.get(i + 1) == Some(&0x20) { sentence += 1; } @@ -121,24 +201,27 @@ fn render_bitmap(text: &str, size: usize, font: &Font) -> Vec { if glyph.rows.is_empty() { continue; } - let row = i / cols; - let col = i - row * cols; - let top = (row * GLYPH_PITCH_Y) as i32 + font.ascent - glyph.h - glyph.yoff; - let left = (col * GLYPH_ADVANCE_X) as i32 + glyph.xoff; - for (r, &bits) in glyph.rows.iter().enumerate() { - if bits == 0 { - continue; - } - let y = top + r as i32; - if y < 0 || y >= size as i32 { - continue; - } - let row_base = y as usize * size; - for b in 0..glyph.w { - if bits & (0x80u8 >> b) != 0 { - let x = left + i32::from(b); - if x >= 0 && (x as usize) < size { - pixels[row_base + x as usize] = ink; + let row = i / grid.cols; + let col = i - row * grid.cols; + let left = (col * font.cell_w) as i32 + glyph.xoff; + for copy in 0..grid.repeat { + let cell_top = ((row * grid.repeat + copy) * font.cell_h) as i32; + let top = cell_top + font.ascent - glyph.h - glyph.yoff; + for (r, &bits) in glyph.rows.iter().enumerate() { + if bits == 0 { + continue; + } + let y = top + r as i32; + if y < 0 || y >= height as i32 { + continue; + } + let row_base = y as usize * width; + for b in 0..glyph.w { + if bits & (0x80u8 >> b) != 0 { + let x = left + i32::from(b); + if x >= 0 && (x as usize) < width { + pixels[row_base + x as usize] = ink; + } } } } @@ -147,8 +230,89 @@ fn render_bitmap(text: &str, size: usize, font: &Font) -> Vec { pixels } +// ============================================================================ +// Lanczos3 resampling (stretch shapes) +// ============================================================================ + +fn lanczos3(x: f32) -> f32 { + let x = x.abs(); + if x < 1e-6 { + return 1.0; + } + if x >= 3.0 { + return 0.0; + } + let pix = PI * x; + (pix.sin() / pix) * ((pix / 3.0).sin() / (pix / 3.0)) +} + +/// Per-output-pixel kernel contributions for one axis, PIL-convention +/// (`center = (i + 0.5) * scale`, kernel stretched by `max(scale, 1)`, +/// weights normalized). +fn contributions(src_len: usize, dst_len: usize) -> Vec<(usize, Vec)> { + let scale = src_len as f32 / dst_len as f32; + let filt_scale = scale.max(1.0); + let support = 3.0 * filt_scale; + let mut out = Vec::with_capacity(dst_len); + for i in 0..dst_len { + let center = (i as f32 + 0.5) * scale; + let begin = ((center - support) as isize).max(0) as usize; + let end = ((center + support).ceil() as usize).min(src_len); + let mut weights = Vec::with_capacity(end - begin); + let mut total = 0.0f32; + for x in begin..end { + let w = lanczos3((x as f32 + 0.5 - center) / filt_scale); + weights.push(w); + total += w; + } + if total != 0.0 { + for w in &mut weights { + *w /= total; + } + } + out.push((begin, weights)); + } + out +} + +/// Separable Lanczos3 resize of an interleaved RGB f32 buffer. +fn resize_rgb(src: &[f32], sw: usize, sh: usize, dw: usize, dh: usize) -> Vec { + let horiz = contributions(sw, dw); + let mut tmp = vec![0f32; dw * sh * 3]; + for y in 0..sh { + let src_row = &src[y * sw * 3..(y + 1) * sw * 3]; + let dst_row = &mut tmp[y * dw * 3..(y + 1) * dw * 3]; + for (x, (begin, weights)) in horiz.iter().enumerate() { + let mut acc = [0f32; 3]; + for (k, &w) in weights.iter().enumerate() { + let s = (begin + k) * 3; + acc[0] = src_row[s].mul_add(w, acc[0]); + acc[1] = src_row[s + 1].mul_add(w, acc[1]); + acc[2] = src_row[s + 2].mul_add(w, acc[2]); + } + dst_row[x * 3..x * 3 + 3].copy_from_slice(&acc); + } + } + let vert = contributions(sh, dh); + let mut out = vec![0f32; dw * dh * 3]; + for (y, (begin, weights)) in vert.iter().enumerate() { + let dst_row = &mut out[y * dw * 3..(y + 1) * dw * 3]; + for (k, &w) in weights.iter().enumerate() { + let src_row = &tmp[(begin + k) * dw * 3..(begin + k + 1) * dw * 3]; + for (d, &s) in dst_row.iter_mut().zip(src_row) { + *d = s.mul_add(w, *d); + } + } + } + out +} + +// ============================================================================ +// PNG encoding +// ============================================================================ + /// Pack one-byte-per-pixel palette indices into 4-bit PNG scanline data -/// (two pixels per byte, high nibble first). With only 7 palette entries, +/// (two pixels per byte, high nibble first). With only 9 palette entries, /// 4-bit depth halves the pre-deflate stream vs 8-bit. fn pack_nibbles(pixels: &[u8], size: usize) -> Vec { let row_bytes = size.div_ceil(2); @@ -196,49 +360,236 @@ fn encode_indexed_png( Ok(out) } -/// Render one snapcompact frame: print pre-normalized text onto a -/// `size` x `size` 4-bit indexed bitmap and encode it as a PNG. +/// Encode an interleaved RGB8 buffer as PNG. Stretched frames are +/// continuous-tone, so adaptive filtering (the `Balanced` default) helps. +fn encode_rgb_png(pixels: &[u8], size: usize, compression: png::Compression) -> Result> { + let mut out = Vec::new(); + let mut encoder = png::Encoder::new(&mut out, size as u32, size as u32); + encoder.set_color(png::ColorType::Rgb); + encoder.set_depth(png::BitDepth::Eight); + encoder.set_compression(compression); + let mut writer = encoder + .write_header() + .map_err(|err| Error::from_reason(format!("Failed to write PNG header: {err}")))?; + writer + .write_image_data(pixels) + .map_err(|err| Error::from_reason(format!("Failed to write PNG data: {err}")))?; + writer + .finish() + .map_err(|err| Error::from_reason(format!("Failed to finish PNG stream: {err}")))?; + Ok(out) +} + +// ============================================================================ +// Entry point +// ============================================================================ + +/// Shape options for one snapcompact frame. +#[napi(object)] +#[derive(Default)] +pub struct SnapcompactRenderOptions { + /// Frame edge in pixels. + pub size: u32, + /// Bundled font: `"5x8"` (X.org BDF) or `"8x8"` (unscii-8). Default `"5x8"`. + pub font: Option, + /// Target cell advance in pixels. Differing from the font's natural cell + /// triggers the Lanczos stretch path. Default: font natural width. + pub cell_width: Option, + /// Target cell pitch in pixels. Default: font natural height. + pub cell_height: Option, + /// Ink variant: `"sent"` (six-hue sentence cycling) or `"bw"` (black). + /// Default `"sent"`. + pub variant: Option, + /// Print each text line this many times; copies after the first sit on a + /// pale highlight band. Default 1. + pub line_repeat: Option, +} + +/// Render one snapcompact frame: print pre-normalized text onto a square +/// bitmap and encode it as PNG. /// -/// The glyph grid holds `floor(size/5) * floor(size/8)` characters; input -/// beyond that is ignored (the caller chunks text to capacity). Returns the -/// PNG bytes. +/// The glyph grid holds `floor(size/cellWidth) * +/// floor(size/cellHeight/lineRepeat)` characters; input beyond that is ignored +/// (the caller chunks text to capacity). Native-cell shapes encode as 4-bit +/// indexed PNG; stretched shapes (target cell != font cell) encode as RGB. +/// Returns the PNG bytes. #[napi] -pub fn render_snapcompact_png(text: String, size: u32) -> Result { +pub fn render_snapcompact_png( + text: String, + options: SnapcompactRenderOptions, +) -> Result { + let size = options.size; if size == 0 || size > MAX_FRAME_SIZE { return Err(Error::from_reason(format!( "Invalid frame size {size}: expected 1..={MAX_FRAME_SIZE}" ))); } - let pixels = render_bitmap(&text, size as usize, &FONT); - Ok(encode_indexed_png(&pixels, size as usize, png::Compression::Balanced)?.into()) + let font_name = options.font.as_deref().unwrap_or("5x8"); + let font = resolve_font(font_name).ok_or_else(|| { + Error::from_reason(format!( + "Unknown snapcompact font {font_name:?}: expected \"5x8\" or \"8x8\"" + )) + })?; + let black_ink = match options.variant.as_deref().unwrap_or("sent") { + "sent" => false, + "bw" => true, + other => { + return Err(Error::from_reason(format!( + "Unknown snapcompact variant {other:?}: expected \"sent\" or \"bw\"" + ))); + }, + }; + let target_w = options.cell_width.unwrap_or(font.cell_w as u32).max(1) as usize; + let target_h = options.cell_height.unwrap_or(font.cell_h as u32).max(1) as usize; + let repeat = options.line_repeat.unwrap_or(1).max(1) as usize; + let size = size as usize; + let grid = Grid { cols: size / target_w, rows: size / target_h / repeat, repeat }; + if grid.cols == 0 || grid.rows == 0 { + return Err(Error::from_reason(format!( + "Frame size {size} cannot fit a {target_w}x{target_h} cell grid (repeat {repeat})" + ))); + } + + if (target_w, target_h) == (font.cell_w, font.cell_h) { + // Native cell: rasterize straight onto the frame, indexed. + let pixels = render_bitmap(&text, size, size, font, &grid, black_ink); + return Ok(encode_indexed_png(&pixels, size, png::Compression::Balanced)?.into()); + } + + // Stretch shape: rasterize at the font's natural cell on a tight canvas, + // Lanczos3-resample to the target cell, paste onto the white frame. + let src_w = grid.cols * font.cell_w; + let src_h = grid.rows * grid.repeat * font.cell_h; + let dst_w = grid.cols * target_w; + let dst_h = grid.rows * grid.repeat * target_h; + let indexed = render_bitmap(&text, src_w, src_h, font, &grid, black_ink); + let mut rgb = vec![0f32; src_w * src_h * 3]; + for (dst, &idx) in rgb.chunks_exact_mut(3).zip(&indexed) { + let [r, g, b] = PALETTE[idx as usize]; + dst[0] = f32::from(r); + dst[1] = f32::from(g); + dst[2] = f32::from(b); + } + let resized = resize_rgb(&rgb, src_w, src_h, dst_w, dst_h); + let mut frame = vec![255u8; size * size * 3]; + for y in 0..dst_h.min(size) { + let src_row = &resized[y * dst_w * 3..(y + 1) * dst_w * 3]; + let dst_row = &mut frame[y * size * 3..]; + for (d, &s) in dst_row[..dst_w.min(size) * 3].iter_mut().zip(src_row) { + *d = s.round().clamp(0.0, 255.0) as u8; + } + } + Ok(encode_rgb_png(&frame, size, png::Compression::Balanced)?.into()) } #[cfg(test)] mod tests { use super::*; + fn opts(size: u32) -> SnapcompactRenderOptions { + SnapcompactRenderOptions { size, ..Default::default() } + } + #[test] - fn font_parses_ascii_coverage() { - let font = &*FONT; - assert_eq!(font.ascent, 7); - // Every printable ASCII char must have a glyph. - for cp in 0x20u32..0x7f { - assert!(font.glyphs.contains_key(&cp), "missing glyph for U+{cp:04X}"); + fn fonts_parse_ascii_coverage() { + for (font, ascent) in [(&*FONT_5X8, 7), (&*FONT_8X8, 7)] { + assert_eq!(font.ascent, ascent); + // Every printable ASCII char must have a glyph. + for cp in 0x20u32..0x7f { + assert!(font.glyphs.contains_key(&cp), "missing glyph for U+{cp:04X}"); + } } } #[test] fn bitmap_inks_sentences_and_caps_capacity() { - let font = &*FONT; - // 40px -> 8 cols x 5 rows = 40 cells. - let pixels = render_bitmap("Hi. Ok.", 40, font); + // 40px -> 8 cols x 5 rows = 40 cells (5x8 font). + let grid = Grid { cols: 8, rows: 5, repeat: 1 }; + let pixels = render_bitmap("Hi. Ok.", 40, 40, &FONT_5X8, &grid, false); let inks: Vec = pixels.iter().copied().filter(|&p| p != 0).collect(); assert!(inks.contains(&1), "first sentence should use ink 1"); assert!(inks.contains(&2), "second sentence should use ink 2"); assert!(!inks.contains(&3), "no third sentence ink expected"); // Overflow input renders without panicking and stays in-bounds. - let overflow = render_bitmap(&"x".repeat(100), 40, font); + let overflow = render_bitmap(&"x".repeat(100), 40, 40, &FONT_5X8, &grid, false); assert_eq!(overflow.len(), 40 * 40); } + + #[test] + fn bw_variant_prints_black_only() { + let grid = Grid { cols: 8, rows: 8, repeat: 1 }; + let pixels = render_bitmap("Hi. Ok.", 64, 64, &FONT_8X8, &grid, true); + let inks: Vec = pixels.iter().copied().filter(|&p| p != 0).collect(); + assert!(!inks.is_empty()); + assert!(inks.iter().all(|&p| p == INK_BLACK), "bw must ink only black"); + } + + #[test] + fn line_repeat_duplicates_rows_on_highlight_bands() { + // 64px, 8x8 font, repeat 2 -> 8 cols x 4 unique rows. + let grid = Grid { cols: 8, rows: 4, repeat: 2 }; + let pixels = render_bitmap("ABCDEFGH", 64, 64, &FONT_8X8, &grid, true); + // Copy band (rows 8..16) carries the highlight background. + assert!(pixels[9 * 64..10 * 64].contains(&BG_REPEAT), "duplicate band must be highlighted"); + // Identical glyph ink in both copies: compare full 8-row bands modulo + // background. + for y in 0..8 { + for x in 0..64 { + let a = pixels[y * 64 + x]; + let b = pixels[(y + 8) * 64 + x]; + assert_eq!(a == INK_BLACK, b == INK_BLACK, "copy ink mismatch at ({x},{y})"); + } + } + } + + #[test] + fn render_native_is_indexed_and_stretch_is_rgb() { + let native = render_snapcompact_png("Hello world. Again.".into(), SnapcompactRenderOptions { + size: 128, + font: Some("8x8".into()), + variant: Some("bw".into()), + line_repeat: Some(2), + ..Default::default() + }) + .unwrap(); + // PNG color type lives at byte 25 of the IHDR: 3 = indexed. + assert_eq!(native[25], 3); + + let stretched = + render_snapcompact_png("Hello world. Again.".into(), SnapcompactRenderOptions { + size: 128, + font: Some("8x8".into()), + cell_width: Some(6), + cell_height: Some(6), + ..Default::default() + }) + .unwrap(); + // 2 = truecolor RGB. + assert_eq!(stretched[25], 2); + // Stretched output must contain anti-aliased (non-extreme) pixels. + let legacy = render_snapcompact_png("Hi. Ok.".into(), opts(40)).unwrap(); + assert_eq!(legacy[25], 3, "default shape stays the legacy 5x8 indexed path"); + } + + #[test] + fn rejects_bad_shapes() { + assert!(render_snapcompact_png("x".into(), opts(0)).is_err()); + assert!( + render_snapcompact_png("x".into(), SnapcompactRenderOptions { + size: 64, + font: Some("9x9".into()), + ..Default::default() + }) + .is_err() + ); + assert!( + render_snapcompact_png("x".into(), SnapcompactRenderOptions { + size: 64, + variant: Some("zebra".into()), + ..Default::default() + }) + .is_err() + ); + } } diff --git a/docs/compaction.md b/docs/compaction.md index 789e36aa1..46af16287 100644 --- a/docs/compaction.md +++ b/docs/compaction.md @@ -10,7 +10,7 @@ Both are persisted as session entries and converted back into user-context messa ## Key implementation files - `packages/agent/src/compaction/compaction.ts` (context-full summarization and handoff generation) -- `packages/agent/src/compaction/snapcompact.ts` (snapcompact strategy: history archived as dense bitmap images) +- `packages/snapcompact/src/snapcompact.ts` (snapcompact strategy: history archived as dense bitmap images) - `packages/agent/src/compaction/branch-summarization.ts` - `packages/agent/src/compaction/pruning.ts` - `packages/agent/src/compaction/utils.ts` @@ -129,13 +129,13 @@ The automatic paths are intentionally different: ### Snapcompact strategy -`compaction.strategy: "snapcompact"` replaces the LLM summarization call with a local, deterministic archival pass (`snapcompactCompact` in `packages/agent/src/compaction/snapcompact.ts`): +`compaction.strategy: "snapcompact"` replaces the LLM summarization call with a local, deterministic archival pass (`snapcompactCompact` from `@oh-my-pi/snapcompact`): -- The discarded history is serialized (same `serializeConversation` path the summarizer uses), whitespace-collapsed, and printed onto square PNG frames using the bundled public-domain X.org `5x8` pixel font — 2576px frames carrying ~165k characters each. Glyph ink cycles through six hues at sentence boundaries to help the model keep its place. +- The discarded history is serialized, whitespace-collapsed, and printed onto provider-aware square PNG frames using bundled public-domain pixel fonts. Anthropic-family and unknown APIs use repeated black `8x8` cells, Google uses repeated sentence-colored `8x8` cells, and OpenAI uses dense stretched `6x6` cells with `detail: "original"`. - Frames persist under `CompactionEntry.preserveData.snapcompact` and are re-attached to the `compactionSummary` message as image blocks on every context rebuild; the entry's `summary` is a deterministic reading guide (grid geometry, role tags, truncation notes) plus the usual file-operation lists. - Later compactions carry earlier frames forward. Beyond an 8-frame budget the archive fades from the middle out: the earliest frame (session head — the original request, or the filmed summary of older history) is pinned, and the oldest *unpinned* frames are evicted, so head and tail both survive. If the previous compaction was text-based, its summary is printed at the head of the frame archive as `[Summary of earlier history]`. - No model, API key, or network is involved, so snapcompact is also safe for overflow recovery. It requires a vision-capable current model (`model.input` includes `"image"`); otherwise the run falls back to context-full and emits a warning notice (auto and manual paths). Manual `/compact` honors the strategy unless custom instructions are given (those imply a directed LLM summary). -- Rationale: a vision model reads ~50 chars per image token after provider downscaling, ~7x cheaper than raw text at near-parity recall (SQuAD eval: F1 0.878 vs 0.899). +- Rationale: the shape table comes from the snapcompact 200k-token evals in `packages/snapcompact`, where bitmap frames preserved QA recall at lower billed-token cost than raw text for vision-capable models. ### Display transcript diff --git a/docs/skills/examples/hello-extension/package.json b/docs/skills/examples/hello-extension/package.json index 1c00ef2be..5027b45a4 100644 --- a/docs/skills/examples/hello-extension/package.json +++ b/docs/skills/examples/hello-extension/package.json @@ -2,6 +2,7 @@ "name": "hello-extension", "version": "1.0.0", "description": "Minimal oh-my-pi extension example", + "homepage": "https://omp.sh", "omp": { "extensions": ["./index.ts"] } diff --git a/docs/skills/examples/mini-marketplace/my-plugin/package.json b/docs/skills/examples/mini-marketplace/my-plugin/package.json index e971f0913..7c4636884 100644 --- a/docs/skills/examples/mini-marketplace/my-plugin/package.json +++ b/docs/skills/examples/mini-marketplace/my-plugin/package.json @@ -1,6 +1,7 @@ { "name": "my-plugin", "version": "0.1.0", + "homepage": "https://omp.sh", "omp": { "extensions": ["./index.ts"] } diff --git a/docs/skills/examples/safety-hook/package.json b/docs/skills/examples/safety-hook/package.json index 8c6ab8271..b8b14fd1c 100644 --- a/docs/skills/examples/safety-hook/package.json +++ b/docs/skills/examples/safety-hook/package.json @@ -2,6 +2,7 @@ "name": "safety-hook", "version": "1.0.0", "description": "oh-my-pi extension example: block rm -rf / via tool_call hook", + "homepage": "https://omp.sh", "omp": { "extensions": ["./index.ts"] } diff --git a/package.json b/package.json index 173cd305c..12da8abf7 100644 --- a/package.json +++ b/package.json @@ -1,5 +1,6 @@ { "name": "omp-monorepo", + "homepage": "https://omp.sh", "private": true, "type": "module", "packageManager": "bun@1.3.14", @@ -30,6 +31,7 @@ "@oh-my-pi/pi-natives": "15.10.12", "@oh-my-pi/pi-tui": "15.10.12", "@oh-my-pi/pi-utils": "15.10.12", + "@oh-my-pi/snapcompact": "15.10.12", "@opentelemetry/api": "^1.9.1", "@opentelemetry/context-async-hooks": "^2.7.1", "@opentelemetry/exporter-trace-otlp-proto": "^0.218.0", diff --git a/packages/agent/CHANGELOG.md b/packages/agent/CHANGELOG.md index 607bfc6ee..9fe70f2b0 100644 --- a/packages/agent/CHANGELOG.md +++ b/packages/agent/CHANGELOG.md @@ -1,9 +1,9 @@ # Changelog ## [Unreleased] - ### Breaking Changes +- Removed `compaction/index.ts` re-export of snapcompact helpers, so snapcompact utilities are no longer available from the agent compaction barrel and should be imported from `@oh-my-pi/snapcompact` - Removed the `convertToLlm` alias export from `compaction/messages` — it duplicated `defaultConvertToLlm` under a second name. Import `defaultConvertToLlm` (array form) or the new `convertMessageToLlm` (single-message form) instead ### Added @@ -11,7 +11,8 @@ - Added `convertMessageToLlm()`: the single-message core transformer behind `defaultConvertToLlm()`. Embedders with app-specific message roles should handle their own roles and delegate every core role (`user`/`developer`/`assistant`/`toolResult`/`custom`/`hookMessage`/`branchSummary`/`compactionSummary`) to it instead of duplicating the conversion — a duplicated `compactionSummary` case is how snapcompact frames once silently dropped off provider requests - Added `pruneSupersededToolResults()` and the opt-in `PruneConfig.supersedeKey` hook so harnesses can prune stale tool results superseded by a newer read of the same file; superseded results are pruned ahead of age-based victims during overflow pruning and replaced with a `[Superseded by a newer read of this file]` placeholder. Without the new config, `pruneToolOutputs()` behavior is unchanged. - Added `readToolSupersedeKey()` implementing the read-tool path/selector grammar (selector-free reads supersede range reads of the same file; URL-scheme paths exempt). Pruning honors prompt-cache economics: per-turn prunes only fire when the post-candidate suffix is small or the cache is cold (idle gap). -- Added the `snapcompact` compaction strategy (`snapcompactCompact()` in `compaction/snapcompact.ts`): instead of an LLM summary, discarded history is printed onto dense 2576px PNG frames with the public-domain X.org `5x8` pixel font (ink cycles per sentence) and re-attached to the compaction summary message as image blocks. Fully local — no model call; ~7x cheaper than raw text at near-parity recall. `CompactionSummaryMessage` gains an optional `images` field, `estimateTokens()` charges per attached frame, and frames persist under `preserveData.snapcompact` with an 8-frame budget that evicts middle-out: the session-head frame is pinned, the oldest unpinned frames drop first, so head and tail both survive. Rasterization and PNG encoding run in native code (`renderSnapcompactPng()` from `@oh-my-pi/pi-natives`), emitting 4-bit indexed PNGs +- Added the `snapcompact` compaction strategy via `@oh-my-pi/snapcompact`: instead of an LLM summary, discarded history is printed onto dense bitmap frames and re-attached to the compaction summary message as image blocks. `CompactionSummaryMessage` gains an optional `images` field, `estimateTokens()` charges per attached frame, and frames persist under `preserveData.snapcompact` with an 8-frame middle-out eviction budget. +- Snapcompact frames are now rendered in a provider-aware shape (`SNAPCOMPACT_SHAPES` + `resolveSnapcompactShape(api)`), following the snapcompact 200k-token monolithic evals: Anthropic-family and unknown APIs get `8x8r-bw` (unscii-8 square cells, black ink, every line printed twice with the copy on a pale highlight band — read at F1 parity with raw text at ~2x lower cost and the most refusal-robust), Google gets `8x8r-sent` (sentence-hue ink, ~2.9x cheaper), and OpenAI gets `6x6u-sent` (unscii Lanczos-stretched to 6x6 cells — OpenAI bills a flat ~2.9k tokens per image, so frame count is the only cost lever) with `detail: "original"` on the frame images. `snapcompactCompact()` accepts `model`/`shape` options, frames persist their shape metadata, mixed-shape archives (provider switches, legacy 5x8 frames) are flagged in the reading instructions, and `snapcompactGeometry()`/`renderSnapcompactFrame()` now take a shape ### Fixed @@ -680,4 +681,4 @@ Initial release under @oh-my-pi scope. See previous releases at [badlogic/pi-mon - `Agent` constructor now has all options optional (empty options use defaults). -- `queueMessage()` is now synchronous (no longer returns a Promise). +- `queueMessage()` is now synchronous (no longer returns a Promise). \ No newline at end of file diff --git a/packages/agent/package.json b/packages/agent/package.json index bd035f3ba..6037d08f3 100644 --- a/packages/agent/package.json +++ b/packages/agent/package.json @@ -39,6 +39,7 @@ "@oh-my-pi/pi-catalog": "catalog:", "@oh-my-pi/pi-natives": "catalog:", "@oh-my-pi/pi-utils": "catalog:", + "@oh-my-pi/snapcompact": "catalog:", "@opentelemetry/api": "catalog:" }, "devDependencies": { diff --git a/packages/agent/src/compaction/compaction.ts b/packages/agent/src/compaction/compaction.ts index 565f801dc..09b6cf3f0 100644 --- a/packages/agent/src/compaction/compaction.ts +++ b/packages/agent/src/compaction/compaction.ts @@ -18,6 +18,7 @@ import { import { clampThinkingLevelForModel } from "@oh-my-pi/pi-catalog/model-thinking"; import { countTokens } from "@oh-my-pi/pi-natives"; import { logger, prompt } from "@oh-my-pi/pi-utils"; +import { SNAPCOMPACT_FRAME_TOKEN_ESTIMATE } from "@oh-my-pi/snapcompact"; import { type AgentTelemetry, instrumentedCompleteSimple } from "../telemetry"; import { ThinkingLevel } from "../thinking"; import type { AgentMessage } from "../types"; @@ -37,7 +38,6 @@ import compactionSummaryPrompt from "./prompts/compaction-summary.md" with { typ import compactionTurnPrefixPrompt from "./prompts/compaction-turn-prefix.md" with { type: "text" }; import compactionUpdateSummaryPrompt from "./prompts/compaction-update-summary.md" with { type: "text" }; import handoffDocumentPrompt from "./prompts/handoff-document.md" with { type: "text" }; -import { SNAPCOMPACT_FRAME_TOKEN_ESTIMATE } from "./snapcompact"; import { computeFileLists, diff --git a/packages/agent/src/compaction/index.ts b/packages/agent/src/compaction/index.ts index 1425b2976..401215724 100644 --- a/packages/agent/src/compaction/index.ts +++ b/packages/agent/src/compaction/index.ts @@ -10,5 +10,4 @@ export * from "./messages"; export * from "./openai"; export * from "./pruning"; export * from "./shake"; -export * from "./snapcompact"; export * from "./utils"; diff --git a/packages/agent/src/compaction/prompts/snapcompact-summary.md b/packages/agent/src/compaction/prompts/snapcompact-summary.md index 3e25c2dcb..0d8d22323 100644 --- a/packages/agent/src/compaction/prompts/snapcompact-summary.md +++ b/packages/agent/src/compaction/prompts/snapcompact-summary.md @@ -1,6 +1,10 @@ Prior conversation history has been archived verbatim onto {{frameCount}} snapcompact frame{{#if multipleFrames}}s{{/if}} — the bitmap image{{#if multipleFrames}}s{{/if}} attached below{{#if multipleFrames}}, ordered oldest to newest{{/if}}. -Reading a frame: monospace 5x8 pixel font on a white background, {{cols}} characters per row, {{rows}} rows per frame; read left to right, top to bottom. Text flows continuously with no word wrap, so words may break across row ends. Whitespace runs (including newlines) were collapsed to single spaces. Ink color cycles through six colors, advancing at sentence boundaries — a color change marks a new sentence. Roles are tagged inline as [User]:, [Assistant]:, [Assistant thinking]:, [Assistant tool calls]:, and [Tool result]:. +Reading a frame: monospace {{fontCell}} pixel font on a white background, {{cols}} characters per row, {{rows}} text rows per frame; read left to right, top to bottom. Text flows continuously with no word wrap, so words may break across row ends. Whitespace runs (including newlines) were collapsed to single spaces. {{#if sentenceInk}}Ink color cycles through six colors, advancing at sentence boundaries — a color change marks a new sentence.{{else}}Glyphs are plain black ink.{{/if}}{{#if lineRepeated}} Every text line is printed twice in a row — first on the white background, then repeated on a pale yellow band. The copies are identical: read each line once and use the duplicate only to double-check hard glyphs.{{/if}} Roles are tagged inline as [User]:, [Assistant]:, [Assistant thinking]:, [Assistant tool calls]:, and [Tool result]:. +{{#if mixedShapes}} + +Older frames may use a different font, grid, or ink coloring than described above; the reading order is always the same (left to right, top to bottom, oldest frame first). +{{/if}} {{#if includedPreviousSummary}} The earliest frame begins with "[Summary of earlier history]" — a condensed digest of context that predates the archived conversation. diff --git a/packages/agent/src/compaction/snapcompact.ts b/packages/agent/src/compaction/snapcompact.ts deleted file mode 100644 index 5236a26eb..000000000 --- a/packages/agent/src/compaction/snapcompact.ts +++ /dev/null @@ -1,316 +0,0 @@ -/** - * Snapcompact compaction: archive conversation history as dense bitmap images. - * - * Instead of asking an LLM to summarize discarded history, the serialized - * conversation is rendered into square PNG frames using the X.org `5x8` - * pixel font (public domain) — one character per 5x8 cell, row-major, glyph - * ink cycling through six hues at sentence boundaries. Vision models read - * the frames back directly, like an archivist at a snapcompact frame reader. - * - * Validated by the imageee SQuAD eval (`img-5x8-sent`, 2576px frames): - * ~0.88 F1 recall vs ~0.90 for raw text, at roughly 7x fewer input tokens. - * The provider downscales frames to its image cap (1568px for Anthropic), - * so one frame costs ~3.3k tokens while carrying ~165k characters. - * - * The whole pass is local and deterministic — no LLM call, no API key, no - * latency beyond rendering. Rasterization and PNG encoding happen in native - * code (`renderSnapcompactPng` in `crates/pi-natives/src/snapcompact.rs`). - * Frames persist in the compaction entry's `preserveData` and are - * re-attached to the compaction summary message on every context rebuild. - */ - -import type { ImageContent } from "@oh-my-pi/pi-ai"; -import { renderSnapcompactPng } from "@oh-my-pi/pi-natives"; -import { prompt } from "@oh-my-pi/pi-utils"; -import type { CompactionDetails, CompactionPreparation, CompactionResult } from "./compaction"; -import { type ConvertToLlm, defaultConvertToLlm } from "./messages"; -import { withOpenAiRemoteCompactionPreserveData } from "./openai"; -import snapcompactSummaryPrompt from "./prompts/snapcompact-summary.md" with { type: "text" }; -import { computeFileLists, serializeConversation, upsertFileOperations } from "./utils"; - -// ============================================================================ -// Constants -// ============================================================================ - -/** Frame edge in pixels. 2576px is the eval-validated sweet spot: the provider - * downscale to 1568px anti-aliases the 1px glyph strokes instead of shearing - * them, which reads *better* than rendering at 1568 directly. */ -export const SNAPCOMPACT_FRAME_SIZE = 2576; - -/** Glyph cell geometry of the bundled `5x8` BDF font. */ -const GLYPH_ADVANCE_X = 5; -const GLYPH_PITCH_Y = 8; - -/** Maximum frames carried on a compaction entry. Oldest frames are dropped - * first once the budget is exceeded (mirrors how iterative text summaries - * fade the oldest detail). 8 frames ≈ 26k image tokens ≈ 1.3M chars. */ -export const SNAPCOMPACT_MAX_FRAMES = 8; - -/** Token cost estimate per frame. Frames render at ≥1568px, so providers bill - * the downscaled long-edge cap: 1568*1568/750 ≈ 3,278 tokens (Anthropic). */ -export const SNAPCOMPACT_FRAME_TOKEN_ESTIMATE = 3300; - -/** Key under `CompactionEntry.preserveData` holding the frame archive. */ -export const SNAPCOMPACT_PRESERVE_KEY = "snapcompact"; - -// ============================================================================ -// Types -// ============================================================================ - -/** One developed snapcompact frame: a base64 PNG plus its reading geometry. */ -export interface SnapcompactFrame { - /** Base64-encoded PNG. */ - data: string; - mimeType: string; - /** Characters per row in the frame grid. */ - cols: number; - /** Rows in the frame grid. */ - rows: number; - /** Characters actually printed onto this frame. */ - chars: number; -} - -/** Frame archive persisted under `preserveData[SNAPCOMPACT_PRESERVE_KEY]`. */ -export interface SnapcompactArchive { - /** Frames ordered oldest to newest. */ - frames: SnapcompactFrame[]; - /** Characters currently readable across all frames. */ - totalChars: number; - /** Characters dropped so far to respect the frame budget. */ - truncatedChars: number; -} - -export interface SnapcompactGeometry { - cols: number; - rows: number; - /** Characters that fit one frame (cols * rows). */ - capacity: number; -} - -export interface SnapcompactOptions { - /** App-level message transformer (same contract as `SummaryOptions.convertToLlm`). */ - convertToLlm?: ConvertToLlm; - /** Frame edge in pixels. Defaults to {@link SNAPCOMPACT_FRAME_SIZE}. */ - frameSize?: number; - /** Frame budget. Defaults to {@link SNAPCOMPACT_MAX_FRAMES}. */ - maxFrames?: number; -} - -/** Result of rendering one frame, before base64 packing. */ -export interface RenderedFrame { - png: Uint8Array; - cols: number; - rows: number; - /** Characters printed (input may be shorter than capacity). */ - chars: number; -} - -// ============================================================================ -// Text normalization -// ============================================================================ - -/** Folds for common non-Latin-1 characters the 5x8 subset cannot draw. */ -const CHAR_FOLD: Record = { - "\u2018": "'", - "\u2019": "'", - "\u201a": "'", - "\u201b": "'", - "\u201c": '"', - "\u201d": '"', - "\u201e": '"', - "\u2013": "-", - "\u2014": "-", - "\u2015": "-", - "\u2212": "-", - "\u2026": "...", - "\u2022": "*", - "\u25cf": "*", - "\u25a0": "*", - "\u25aa": "*", - "\u2190": "<-", - "\u2192": "->", - "\u21d2": "=>", - "\u2713": "v", - "\u2714": "v", - "\u2717": "x", - "\u2718": "x", -}; - -/** - * Prepare text for printing: collapse whitespace runs (incl. newlines) to - * single spaces — the eval's "paragraph breaks collapsed to spaces" format — - * then fold everything outside the font's ASCII + Latin-1 coverage to ASCII - * approximations (`?` as the last resort). - */ -export function normalizeForSnapcompact(text: string): string { - const collapsed = text.replace(/\s+/g, " ").trim(); - let out = ""; - for (const ch of collapsed) { - const cp = ch.codePointAt(0) as number; - if (cp < 0x7f || (cp >= 0xa0 && cp <= 0xff)) { - out += ch; - continue; - } - const fold = CHAR_FOLD[ch]; - if (fold !== undefined) { - out += fold; - } else if (cp >= 0x2500 && cp <= 0x257f) { - // Box drawing: keep table skeletons legible. - out += cp === 0x2502 || cp === 0x2503 ? "|" : cp === 0x2500 || cp === 0x2501 ? "-" : "+"; - } else { - out += "?"; - } - } - return out; -} - -// ============================================================================ -// Rendering -// ============================================================================ -export function snapcompactGeometry(size: number = SNAPCOMPACT_FRAME_SIZE): SnapcompactGeometry { - const cols = Math.floor(size / GLYPH_ADVANCE_X); - const rows = Math.floor(size / GLYPH_PITCH_Y); - return { cols, rows, capacity: cols * rows }; -} - -/** Render one snapcompact frame from already-normalized text. */ -export function renderSnapcompactFrame(text: string, size: number = SNAPCOMPACT_FRAME_SIZE): RenderedFrame { - const { cols, rows, capacity } = snapcompactGeometry(size); - const chars = Math.min(text.length, capacity); - return { png: renderSnapcompactPng(text, size), cols, rows, chars }; -} - -// ============================================================================ -// Archive helpers -// ============================================================================ - -/** Validate and extract a persisted frame archive from `preserveData`. */ -export function getPreservedSnapcompactArchive( - preserveData: Record | undefined, -): SnapcompactArchive | undefined { - const candidate = preserveData?.[SNAPCOMPACT_PRESERVE_KEY]; - if (!candidate || typeof candidate !== "object") return undefined; - const archive = candidate as SnapcompactArchive; - if (!Array.isArray(archive.frames)) return undefined; - const frames = archive.frames.filter( - frame => - !!frame && - typeof frame.data === "string" && - frame.data.length > 0 && - typeof frame.mimeType === "string" && - typeof frame.cols === "number" && - typeof frame.rows === "number" && - typeof frame.chars === "number", - ); - if (frames.length === 0) return undefined; - return { - frames, - totalChars: typeof archive.totalChars === "number" ? archive.totalChars : 0, - truncatedChars: typeof archive.truncatedChars === "number" ? archive.truncatedChars : 0, - }; -} - -/** Convert archive frames into LLM image blocks (oldest first). */ -export function snapcompactImages(archive: SnapcompactArchive): ImageContent[] { - return archive.frames.map(frame => ({ type: "image", data: frame.data, mimeType: frame.mimeType })); -} - -// ============================================================================ -// Compaction entry point -// ============================================================================ - -/** - * Run a snapcompact compaction over prepared messages. Fully local: serializes - * the discarded history, prints it onto PNG frames, merges previously - * archived frames (oldest dropped beyond the budget), and produces a - * deterministic summary explaining how to read the frames. - * - * If the previous compaction was text-based, its summary is printed at the - * head of the frame archive as `[Summary of earlier history]` so no continuity is lost. - */ -export async function snapcompactCompact( - preparation: CompactionPreparation, - options?: SnapcompactOptions, -): Promise { - const { firstKeptEntryId, tokensBefore, previousSummary, previousPreserveData, fileOps } = preparation; - if (!firstKeptEntryId) { - throw new Error("First kept entry has no ID - session may need migration"); - } - const frameSize = options?.frameSize ?? SNAPCOMPACT_FRAME_SIZE; - const maxFrames = Math.max(1, options?.maxFrames ?? SNAPCOMPACT_MAX_FRAMES); - const geometry = snapcompactGeometry(frameSize); - - const messages = preparation.messagesToSummarize.concat(preparation.turnPrefixMessages); - const llmMessages = (options?.convertToLlm ?? defaultConvertToLlm)(messages); - let archiveText = normalizeForSnapcompact(serializeConversation(llmMessages)); - - const previousArchive = getPreservedSnapcompactArchive(previousPreserveData); - const includedPreviousSummary = !previousArchive && !!previousSummary; - if (includedPreviousSummary && previousSummary) { - const head = `[Summary of earlier history] ${normalizeForSnapcompact(previousSummary)}`; - archiveText = archiveText.length > 0 ? `${head} [Recent conversation] ${archiveText}` : head; - } - - let truncatedChars = previousArchive?.truncatedChars ?? 0; - - const newFrames: SnapcompactFrame[] = []; - for (let offset = 0; offset < archiveText.length; offset += geometry.capacity) { - const chunk = archiveText.slice(offset, offset + geometry.capacity); - const rendered = renderSnapcompactFrame(chunk, frameSize); - newFrames.push({ - data: Buffer.from(rendered.png).toBase64(), - mimeType: "image/png", - cols: rendered.cols, - rows: rendered.rows, - chars: rendered.chars, - }); - // Keep the event loop responsive between native render passes. - await Bun.sleep(0); - } - - const frames = [...(previousArchive?.frames ?? []), ...newFrames]; - if (frames.length > maxFrames) { - // Pin the earliest frame: it anchors the session head (the original - // request, or the filmed summary of even older history) the way the - // LLM-summary strategies keep the original goal alive across rounds. - // Eviction removes the oldest *unpinned* frames, so the archive fades - // from the middle out — head and tail survive. With a budget of one - // frame the pin is moot; keep the newest frame instead. - const evictStart = maxFrames >= 2 ? 1 : 0; - const dropped = frames.splice(evictStart, frames.length - maxFrames); - for (const frame of dropped) truncatedChars += frame.chars; - } - const totalChars = frames.reduce((sum, frame) => sum + frame.chars, 0); - - let summary: string; - if (frames.length === 0) { - summary = "No prior history."; - } else { - summary = prompt.render(snapcompactSummaryPrompt, { - frameCount: frames.length, - multipleFrames: frames.length > 1, - cols: geometry.cols, - rows: geometry.rows, - totalChars, - truncatedChars, - includedPreviousSummary, - }); - } - const { readFiles, modifiedFiles } = computeFileLists(fileOps); - summary = upsertFileOperations(summary, readFiles, modifiedFiles); - - // A snapcompact pass replaces any provider-side replacement history; strip the - // OpenAI remote-compaction payload like the default summarizer path does. - // OpenAI remote-compaction payload like the default summarizer path does. - const basePreserve = withOpenAiRemoteCompactionPreserveData(previousPreserveData, undefined) ?? {}; - const archive: SnapcompactArchive = { frames, totalChars, truncatedChars }; - - return { - summary, - shortSummary: `Archived ${totalChars.toLocaleString()} chars of history onto ${frames.length} snapcompact frame${frames.length === 1 ? "" : "s"}`, - firstKeptEntryId, - tokensBefore, - details: { readFiles, modifiedFiles } as CompactionDetails, - preserveData: { ...basePreserve, [SNAPCOMPACT_PRESERVE_KEY]: archive }, - }; -} diff --git a/packages/agent/test/snapcompact-frames.test.ts b/packages/agent/test/snapcompact-frames.test.ts new file mode 100644 index 000000000..1e24f74a0 --- /dev/null +++ b/packages/agent/test/snapcompact-frames.test.ts @@ -0,0 +1,44 @@ +import { describe, expect, it } from "bun:test"; +import type { ImageContent } from "@oh-my-pi/pi-ai"; +import { SNAPCOMPACT_FRAME_TOKEN_ESTIMATE } from "@oh-my-pi/snapcompact"; +import { estimateTokens } from "../src/compaction/compaction"; +import { createCompactionSummaryMessage, defaultConvertToLlm } from "../src/compaction/messages"; + +describe("compaction summary message with snapcompact frames", () => { + const images: ImageContent[] = [ + { type: "image", data: "ZmFrZQ==", mimeType: "image/png" }, + { type: "image", data: "ZmFrZTI=", mimeType: "image/png" }, + ]; + + it("estimateTokens charges per attached frame", () => { + const bare = createCompactionSummaryMessage("summary text", 1000, new Date().toISOString()); + const withFrames = createCompactionSummaryMessage( + "summary text", + 1000, + new Date().toISOString(), + undefined, + undefined, + images, + ); + expect(estimateTokens(withFrames) - estimateTokens(bare)).toBe(2 * SNAPCOMPACT_FRAME_TOKEN_ESTIMATE); + }); + + it("defaultConvertToLlm appends frames as image blocks after the summary text", () => { + const message = createCompactionSummaryMessage( + "the snapcompact archive", + 1000, + new Date().toISOString(), + undefined, + undefined, + images, + ); + const [converted] = defaultConvertToLlm([message]); + expect(converted.role).toBe("user"); + const content = converted.content as Array<{ type: string; text?: string; data?: string }>; + expect(content.length).toBe(3); + expect(content[0].type).toBe("text"); + expect(content[0].text).toContain("the snapcompact archive"); + expect(content[1]).toEqual(images[0]); + expect(content[2]).toEqual(images[1]); + }); +}); diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index af15979fe..cebe1dad9 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Added + +- Added optional `ImageContent.detail` (`"auto" | "low" | "high" | "original"`): an OpenAI resolution hint forwarded by the `openai-responses` serializers (default stays `auto`) and by `openai-completions` for the values Chat Completions supports. `"original"` preserves native resolution — required for snapcompact frames, whose pixel-font glyphs do not survive the default downscale. Providers without a detail knob ignore the field. + ## [15.10.12] - 2026-06-10 ### Added diff --git a/packages/ai/src/providers/openai-completions.ts b/packages/ai/src/providers/openai-completions.ts index 0a1d83eb3..1b3401a94 100644 --- a/packages/ai/src/providers/openai-completions.ts +++ b/packages/ai/src/providers/openai-completions.ts @@ -1651,6 +1651,8 @@ export function convertMessages( type: "image_url", image_url: { url: `data:${item.mimeType};base64,${item.data}`, + // Chat Completions has no "original"; omit it (provider default). + ...(item.detail && item.detail !== "original" ? { detail: item.detail } : {}), }, } satisfies ChatCompletionContentPartImage); } else { diff --git a/packages/ai/src/providers/openai-responses-shared.ts b/packages/ai/src/providers/openai-responses-shared.ts index 9e28d60e0..5ddf64b81 100644 --- a/packages/ai/src/providers/openai-responses-shared.ts +++ b/packages/ai/src/providers/openai-responses-shared.ts @@ -289,7 +289,7 @@ export function convertResponsesInputContent( for (const item of imageBlocks) { normalizedContent.push({ type: "input_image", - detail: "auto", + detail: item.detail ?? "auto", image_url: `data:${item.mimeType};base64,${item.data}`, } satisfies ResponseInputImage); } @@ -448,7 +448,7 @@ export function appendResponsesToolResultMessages( if (block.type === "image") { contentParts.push({ type: "input_image", - detail: "auto", + detail: block.detail ?? "auto", image_url: `data:${block.mimeType};base64,${block.data}`, } satisfies ResponseInputImage); } diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 1aa43478d..6d64e09c2 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -409,6 +409,12 @@ export interface ImageContent { type: "image"; data: string; // base64 encoded image data mimeType: string; // e.g., "image/jpeg", "image/png" + /** + * OpenAI-only resolution hint. `"original"` preserves native resolution + * (required for snapcompact frames, whose glyphs do not survive the + * default `auto` downscale). Providers without a detail knob ignore it. + */ + detail?: "auto" | "low" | "high" | "original"; } export interface ToolCall { diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 6fcc94773..44ce4e906 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -12,6 +12,7 @@ ### Added - Added pre-TUI startup input capture so users could type while interactive sessions initialize and keep their draft while the application loads +- Snapcompact compaction now passes the session model so frames render in the provider-optimal shape (unscii `8x8r-bw` for Anthropic-family/unknown APIs, `8x8r-sent` for Google, Lanczos-stretched `6x6u-sent` with `detail: "original"` for OpenAI), per the snapcompact 200k-token evals - Added queued submission replay so Enter presses made before startup completion are submitted automatically once interactive mode begins - Added per-turn supersede pruning of stale `read` results: when a file is re-read, older copies of the same path/selector are pruned from context at cache-favorable moments (small suffix, idle gap, or alongside overflow pruning). Gated by the new `compaction.supersedeReads` setting (default on) - Added soft request budgets for task subagents (explore/quick_task 40, others 90, configurable via `task.softRequestBudget`, 0 disables): crossing the budget injects a one-time wrap-up steer into the child; crossing 1.5× aborts the run gracefully diff --git a/packages/coding-agent/examples/extensions/with-deps/package.json b/packages/coding-agent/examples/extensions/with-deps/package.json index 347be16b5..06dc2cb08 100644 --- a/packages/coding-agent/examples/extensions/with-deps/package.json +++ b/packages/coding-agent/examples/extensions/with-deps/package.json @@ -1,6 +1,7 @@ { "name": "pi-extension-with-deps", "version": "1.0.0", + "homepage": "https://omp.sh", "type": "module", "omp": { "extensions": [ diff --git a/packages/coding-agent/package.json b/packages/coding-agent/package.json index 9dd841227..dbd5acefe 100644 --- a/packages/coding-agent/package.json +++ b/packages/coding-agent/package.json @@ -56,6 +56,7 @@ "@oh-my-pi/pi-natives": "catalog:", "@oh-my-pi/pi-tui": "catalog:", "@oh-my-pi/pi-utils": "catalog:", + "@oh-my-pi/snapcompact": "catalog:", "@opentelemetry/api": "catalog:", "@opentelemetry/context-async-hooks": "catalog:", "@opentelemetry/exporter-trace-otlp-proto": "catalog:", diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index d0746f35a..51f82b5b5 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/packages/coding-agent/src/session/agent-session.ts @@ -55,7 +55,6 @@ import { type ShakeRegion, type SummaryOptions, shouldCompact, - snapcompactCompact, } from "@oh-my-pi/pi-agent-core/compaction"; import { DEFAULT_PRUNE_CONFIG, @@ -107,6 +106,7 @@ import { relativePathWithinRoot, Snowflake, } from "@oh-my-pi/pi-utils"; +import { snapcompactCompact } from "@oh-my-pi/snapcompact"; import { type AsyncJob, type AsyncJobDeliveryState, AsyncJobManager } from "../async"; import { classifyDifficulty } from "../auto-thinking/classifier"; import { reset as resetCapabilities } from "../capability"; @@ -6323,7 +6323,7 @@ export class AgentSession { details = compactionPrep.details; preserveData = compactionPrep.preserveData; } else if (snapcompactReady) { - const snapcompactResult = await snapcompactCompact(preparation, { convertToLlm }); + const snapcompactResult = await snapcompactCompact(preparation, { convertToLlm, model: this.model }); summary = snapcompactResult.summary; shortSummary = snapcompactResult.shortSummary; firstKeptEntryId = snapcompactResult.firstKeptEntryId; @@ -7825,7 +7825,7 @@ export class AgentSession { } else if (action === "snapcompact") { // Local, deterministic: render discarded history onto PNG frames. // No model candidates, no API key, no retry loop. - const snapcompactResult = await snapcompactCompact(preparation, { convertToLlm }); + const snapcompactResult = await snapcompactCompact(preparation, { convertToLlm, model: this.model }); summary = snapcompactResult.summary; shortSummary = snapcompactResult.shortSummary; firstKeptEntryId = snapcompactResult.firstKeptEntryId; diff --git a/packages/coding-agent/src/session/session-manager.ts b/packages/coding-agent/src/session/session-manager.ts index bc0c23c40..2c8f49c86 100644 --- a/packages/coding-agent/src/session/session-manager.ts +++ b/packages/coding-agent/src/session/session-manager.ts @@ -2,7 +2,6 @@ import * as fs from "node:fs"; import * as os from "node:os"; import * as path from "node:path"; import type { AgentMessage } from "@oh-my-pi/pi-agent-core"; -import { getPreservedSnapcompactArchive, snapcompactImages } from "@oh-my-pi/pi-agent-core/compaction"; import type { ImageContent, Message, @@ -28,6 +27,7 @@ import { Snowflake, toError, } from "@oh-my-pi/pi-utils"; +import { getPreservedSnapcompactArchive, snapcompactImages } from "@oh-my-pi/snapcompact"; import { ArtifactManager } from "./artifacts"; import { type BlobPutOptions, diff --git a/packages/coding-agent/test/job-poll-displacement.test.ts b/packages/coding-agent/test/job-poll-displacement.test.ts index 4ce22f38d..81fa7d08d 100644 --- a/packages/coding-agent/test/job-poll-displacement.test.ts +++ b/packages/coding-agent/test/job-poll-displacement.test.ts @@ -97,7 +97,10 @@ describe("job waiting-poll block lifecycle", () => { }); it("never marks non-job tools displaceable", () => { - const component = trackComponent(created, new ToolExecutionComponent("bash", { command: "ls" }, {}, undefined, uiStub)); + const component = trackComponent( + created, + new ToolExecutionComponent("bash", { command: "ls" }, {}, undefined, uiStub), + ); component.updateResult(pollResult(["running"]), false); expect(component.isDisplaceableBlock()).toBe(false); }); @@ -145,7 +148,12 @@ describe("EventController displaces consecutive waiting polls", () => { } async function runPoll(controller: EventController, children: Component[], toolCallId: string) { - await controller.handleEvent({ type: "tool_execution_start", toolCallId, toolName: "job", args: { poll: ["j0"] } }); + await controller.handleEvent({ + type: "tool_execution_start", + toolCallId, + toolName: "job", + args: { poll: ["j0"] }, + }); const component = children[children.length - 1] as ToolExecutionComponent; trackComponent(created, component); await controller.handleEvent({ diff --git a/packages/coding-agent/test/streaming-preview-height.test.ts b/packages/coding-agent/test/streaming-preview-height.test.ts index acf37be34..fdf2476c9 100644 --- a/packages/coding-agent/test/streaming-preview-height.test.ts +++ b/packages/coding-agent/test/streaming-preview-height.test.ts @@ -372,12 +372,12 @@ describe("streaming tool call preview height (bounded across renderers)", () => } }, 30_000); - test("task pending preview stays bounded with a long multiline assignment", () => { + test("task pending preview keeps the full assignment brief", () => { // CONTRACT CHANGE with the single-spawn task rework: the old uncapped - // multi-task `context` rendering is gone with the field. The pending - // preview now intentionally bounds the assignment (first line + a - // "more lines" marker when collapsed; 12 lines when expanded), like - // bash/ssh, so a long assignment can no longer strand the block top. + // multi-task `context` rendering is gone with the field. The assignment + // brief is the durable record of what the subagent was asked to do, so + // the pending preview renders it in full (like eval code) instead of + // windowing it like bash/ssh command previews. const longLines = Array.from({ length: 80 }, (_, i) => `line-${i}`); const { lines, text } = renderPending("task", { agent: "task", @@ -386,12 +386,11 @@ describe("streaming tool call preview height (bounded across renderers)", () => assignment: longLines.join("\n"), }); - expect(lines.length, "task preview should stay bounded").toBeLessThan(20); + expect(lines.length, "task assignment brief should not be capped").toBeGreaterThan(80); expect(text).toContain("preview"); expect(text).toContain("line-0"); - expect(text).not.toContain("line-40"); - expect(text).not.toContain("line-79"); - expect(text, "task preview should advertise truncation").toMatch(/more lines/); + expect(text).toContain("line-40"); + expect(text).toContain("line-79"); }); test("eval pending preview preserves full code (never collapsed)", () => { diff --git a/packages/natives/CHANGELOG.md b/packages/natives/CHANGELOG.md index 1184a6341..23e0aaac0 100644 --- a/packages/natives/CHANGELOG.md +++ b/packages/natives/CHANGELOG.md @@ -4,7 +4,7 @@ ### Added -- Added `renderSnapcompactPng(text, size)`: rasterizes pre-normalized text onto a square 4-bit indexed PNG using the bundled public-domain X.org `5x8` BDF font, with glyph ink cycling through six hues at sentence boundaries. Replaces the JS rasterizer/PNG writer previously in `@oh-my-pi/pi-agent-core`; the font ships inside the crate (`crates/pi-natives/src/fonts/5x8.bdf`). +- Added `renderSnapcompactPng(text, options)`: rasterizes pre-normalized text onto a square PNG in an eval-validated snapcompact shape. Options select the bundled font (`5x8` X.org BDF or `8x8` unscii-8, both public domain, shipped in `crates/pi-natives/src/fonts/`), the ink variant (`sent` six-hue sentence cycling or `bw` black), line repetition (each text line printed N times, copies on a pale highlight band), and a target cell size — cells differing from the font's natural cell render via Lanczos3 stretch into an anti-aliased RGB frame (e.g. the OpenAI-optimal 6x6 unscii shape); native-cell shapes encode as 4-bit indexed PNG. Replaces the JS rasterizer/PNG writer previously in `@oh-my-pi/pi-agent-core`. ## [15.10.12] - 2026-06-10 diff --git a/packages/natives/native/index.d.ts b/packages/natives/native/index.d.ts index a61406901..bbeeb5364 100644 --- a/packages/natives/native/index.d.ts +++ b/packages/natives/native/index.d.ts @@ -1285,14 +1285,15 @@ export interface PtyStartOptions { export declare function readImageFromClipboard(): Promise /** - * Render one snapcompact frame: print pre-normalized text onto a - * `size` x `size` 4-bit indexed bitmap and encode it as a PNG. + * Render one snapcompact frame: print pre-normalized text onto a square + * bitmap and encode it as PNG. * - * The glyph grid holds `floor(size/5) * floor(size/8)` characters; input - * beyond that is ignored (the caller chunks text to capacity). Returns the - * PNG bytes. + * The glyph grid holds `floor(size/cellWidth) * floor(size/cellHeight/lineRepeat)` + * characters; input beyond that is ignored (the caller chunks text to + * capacity). Native-cell shapes encode as 4-bit indexed PNG; stretched + * shapes (target cell != font cell) encode as RGB. Returns the PNG bytes. */ -export declare function renderSnapcompactPng(text: string, size: number): Uint8Array +export declare function renderSnapcompactPng(text: string, options: SnapcompactRenderOptions): Uint8Array /** * Search content for a pattern (one-shot, compiles pattern each time). @@ -1423,6 +1424,31 @@ export interface SliceResult { */ export declare function sliceWithWidth(line: string, startCol: number, length: number, strict: boolean | undefined | null, tabWidth: number): SliceResult +/** Shape options for one snapcompact frame. */ +export interface SnapcompactRenderOptions { + /** Frame edge in pixels. */ + size: number + /** Bundled font: `"5x8"` (X.org BDF) or `"8x8"` (unscii-8). Default `"5x8"`. */ + font?: string + /** + * Target cell advance in pixels. Differing from the font's natural cell + * triggers the Lanczos stretch path. Default: font natural width. + */ + cellWidth?: number + /** Target cell pitch in pixels. Default: font natural height. */ + cellHeight?: number + /** + * Ink variant: `"sent"` (six-hue sentence cycling) or `"bw"` (black). + * Default `"sent"`. + */ + variant?: string + /** + * Print each text line this many times; copies after the first sit on a + * pale highlight band. Default 1. + */ + lineRepeat?: number +} + export declare function summarizeCode(options: SummaryOptions): SummaryResult export interface SummaryOptions { diff --git a/packages/snapcompact/.gitignore b/packages/snapcompact/.gitignore new file mode 100644 index 000000000..839945bad --- /dev/null +++ b/packages/snapcompact/.gitignore @@ -0,0 +1,2 @@ +.cache/ +results/ diff --git a/packages/snapcompact/CHANGELOG.md b/packages/snapcompact/CHANGELOG.md new file mode 100644 index 000000000..d62403267 --- /dev/null +++ b/packages/snapcompact/CHANGELOG.md @@ -0,0 +1,10 @@ +# Changelog + +## [Unreleased] +### Added + +- Added provider-specific snapcompact frame-shape presets and shape helpers (`SNAPCOMPACT_SHAPES`, `resolveSnapcompactShape`, `isSnapcompactShape`) so callers can consistently select validated image-frame geometry for archive renders +- Added `file-operations.md` and `snapcompact-summary.md` prompts to preserve file-read/write context and frame metadata in the compaction prompt flow +- Added a full `packages/snapcompact/research` experiment and visualization suite for running snapcompact SQuAD studies, provider probes, and activation-style analyses +- Added package-level TypeScript exports and publication config so consumers can import `@oh-my-pi/snapcompact` with typed access to snapcompact APIs +- Published `@oh-my-pi/snapcompact` as the reusable snapcompact compaction package, including bitmap-frame rendering helpers, archive helpers, and the local `snapcompactCompact()` strategy. diff --git a/packages/snapcompact/package.json b/packages/snapcompact/package.json new file mode 100644 index 000000000..bc6bab478 --- /dev/null +++ b/packages/snapcompact/package.json @@ -0,0 +1,62 @@ +{ + "type": "module", + "name": "@oh-my-pi/snapcompact", + "version": "15.10.12", + "description": "Bitmap-frame context compression for vision-capable LLMs", + "homepage": "https://omp.sh", + "author": "Can Boluk", + "license": "MIT", + "repository": { + "type": "git", + "url": "git+https://github.com/can1357/oh-my-pi.git", + "directory": "packages/snapcompact" + }, + "bugs": { + "url": "https://github.com/can1357/oh-my-pi/issues" + }, + "keywords": [ + "context-compression", + "vision", + "compaction", + "llm" + ], + "main": "./src/index.ts", + "types": "./src/index.ts", + "scripts": { + "check": "biome check . && bun run check:types", + "check:types": "tsgo -p tsconfig.json --noEmit", + "lint": "biome lint .", + "test": "bun test --parallel", + "fix": "biome check --write --unsafe .", + "fmt": "biome format --write ." + }, + "dependencies": { + "@oh-my-pi/pi-ai": "catalog:", + "@oh-my-pi/pi-natives": "catalog:", + "@oh-my-pi/pi-utils": "catalog:" + }, + "devDependencies": { + "@types/bun": "catalog:" + }, + "engines": { + "bun": ">=1.3.14" + }, + "files": [ + "src", + "CHANGELOG.md" + ], + "exports": { + ".": { + "types": "./src/index.ts", + "import": "./src/index.ts" + }, + "./snapcompact": { + "types": "./src/snapcompact.ts", + "import": "./src/snapcompact.ts" + }, + "./*": { + "types": "./src/*.ts", + "import": "./src/*.ts" + } + } +} diff --git a/packages/snapcompact/research/anthropic_api.py b/packages/snapcompact/research/anthropic_api.py new file mode 100644 index 000000000..d6a94fa6e --- /dev/null +++ b/packages/snapcompact/research/anthropic_api.py @@ -0,0 +1,74 @@ +"""Minimal Anthropic Messages API client. Key from ~/.env, no SDK.""" + +import base64 +import json +import time +import urllib.error +import urllib.request +from pathlib import Path + +API_URL = "https://api.anthropic.com/v1/messages" +API_VERSION = "2023-06-01" + + +def load_api_key(env_path: str = "~/.env") -> str: + """Last ANTHROPIC_API_KEY assignment wins (mirrors shell sourcing).""" + key = None + for line in Path(env_path).expanduser().read_text().splitlines(): + line = line.strip() + if line.startswith("export "): + line = line[len("export ") :] + if line.startswith("ANTHROPIC_API_KEY="): + key = line.split("=", 1)[1].strip().strip("'\"") + if not key: + raise SystemExit(f"no ANTHROPIC_API_KEY in {env_path}") + return key + + +def image_block(png_path: Path) -> dict: + data = base64.b64encode(png_path.read_bytes()).decode() + return {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": data}} + + +def complete( + api_key: str, + model: str, + messages: list[dict], + system: str | None = None, + max_tokens: int = 8192, + effort: str | None = None, + retries: int = 4, +) -> tuple[str, dict, str]: + """Returns (joined text content, usage dict, stop_reason). + + effort: adaptive-thinking effort (low|medium|high|xhigh|max); None = provider default. + """ + body: dict = {"model": model, "max_tokens": max_tokens, "messages": messages} + if system: + body["system"] = system + if effort: + body["output_config"] = {"effort": effort} + payload = json.dumps(body).encode() + headers = { + "content-type": "application/json", + "x-api-key": api_key, + "anthropic-version": API_VERSION, + } + if effort: + headers["anthropic-beta"] = "effort-2025-11-24" + req = urllib.request.Request(API_URL, data=payload, headers=headers) + for attempt in range(retries + 1): + try: + with urllib.request.urlopen(req, timeout=300) as resp: + out = json.load(resp) + text = "".join(b.get("text", "") for b in out["content"] if b.get("type") == "text") + return text, out.get("usage", {}), out.get("stop_reason", "") + except urllib.error.HTTPError as err: + detail = err.read().decode(errors="replace")[:500] + if err.code in (429, 500, 502, 503, 529) and attempt < retries: + wait = 2.0 * 2**attempt + print(f" HTTP {err.code}, retrying in {wait:.0f}s: {detail[:120]}") + time.sleep(wait) + continue + raise SystemExit(f"API error {err.code}: {detail}") from err + raise AssertionError("unreachable") diff --git a/packages/snapcompact/research/bdf.py b/packages/snapcompact/research/bdf.py new file mode 100644 index 000000000..15003636d --- /dev/null +++ b/packages/snapcompact/research/bdf.py @@ -0,0 +1,247 @@ +"""BDF/HEX pixel-font parsing and dense text-to-image rendering.""" + +import colorsys +import urllib.request +from dataclasses import dataclass +from pathlib import Path + +from PIL import Image + +XORG_RAW = "https://gitlab.freedesktop.org/xorg/font/misc-misc/-/raw/master/{name}.bdf" +TOM_THUMB = "https://robey.lag.net/downloads/tom-thumb.bdf" +UNSCII_HEX = "https://raw.githubusercontent.com/viznut/unscii/master/fontfiles/{name}.hex" + + +@dataclass(frozen=True) +class FontCfg: + """One density configuration: a BDF font drawn on an adv x pitch cell grid.""" + + name: str # condition label, e.g. "6x10" + source: str # bdf file stem or "tom-thumb" + adv: int # x advance per character cell, px + pitch: int # y advance per row, px + ascent: int | None = None # override; default from FONT_ASCENT + native: tuple[int, int] | None = None # rasterize at this cell size, then resize (stretch) to adv x pitch + repeat: int = 1 # render each text line this many times (copy 0 plain, later copies bg-highlighted) + + +def ensure_font(cfg: FontCfg, cache: Path) -> Path: + hexfont = cfg.source.startswith("unscii") + path = cache / f"{cfg.source}.{'hex' if hexfont else 'bdf'}" + if not path.exists(): + if hexfont: + url = UNSCII_HEX.format(name=cfg.source) + else: + url = TOM_THUMB if cfg.source == "tom-thumb" else XORG_RAW.format(name=cfg.source) + urllib.request.urlretrieve(url, path) + return path + + +def parse_bdf(path: Path) -> tuple[dict[int, dict], int]: + """Returns ({codepoint: {bbx, rows}}, font_ascent).""" + glyphs: dict[int, dict] = {} + ascent = 0 + cur: dict = {} + lines = path.read_text().splitlines() + i = 0 + while i < len(lines): + ln = lines[i] + if ln.startswith("FONT_ASCENT"): + ascent = int(ln.split()[1]) + elif ln.startswith("STARTCHAR"): + cur = {"rows": []} + elif ln.startswith("ENCODING"): + cur["enc"] = int(ln.split()[1]) + elif ln.startswith("BBX"): + cur["bbx"] = tuple(map(int, ln.split()[1:5])) + elif ln.startswith("BITMAP"): + i += 1 + while not lines[i].startswith("ENDCHAR"): + cur["rows"].append(int(lines[i], 16)) + i += 1 + glyphs[cur["enc"]] = cur + i += 1 + return glyphs, ascent + + +def parse_hex(path: Path) -> tuple[dict[int, dict], int]: + """Unifont-style .hex (unscii-8: 8x8, one byte per row). Baseline at row 7.""" + glyphs: dict[int, dict] = {} + for line in path.read_text().splitlines(): + cp, _, bits = line.partition(":") + data = bytes.fromhex(bits.strip()) + if len(data) == 8: + glyphs[int(cp, 16)] = {"bbx": (8, 8, 0, -1), "rows": list(data)} + return glyphs, 7 + + +def load_font(cfg: FontCfg, cache: Path) -> tuple[dict[int, dict], int]: + path = ensure_font(cfg, cache) + return parse_hex(path) if path.suffix == ".hex" else parse_bdf(path) + + +# 6 hues; dark variant for glyphs, pale variant for the row's background band. +_HUES = [0.0, 0.08, 0.3, 0.5, 0.62, 0.78] +_DARK = [tuple(int(c * 255) for c in colorsys.hls_to_rgb(h, 0.22, 0.95)) for h in _HUES] +_PALE = [tuple(int(c * 255) for c in colorsys.hls_to_rgb(h, 0.94, 0.6)) for h in _HUES] +_BRIGHT = [tuple(int(c * 255) for c in colorsys.hls_to_rgb(h, 0.70, 0.95)) for h in _HUES] + +VARIANTS = ("color", "zebra", "bw", "sent", "dark", "dark-sent", "dim", "sent-dim", "dark-sent-dim") +_BLACK = (0, 0, 0) +_WHITE = (255, 255, 255) +_GRAY = (232, 232, 232) +_DIMMED = (176, 176, 176) +_DIMMED_DARK = (104, 104, 104) +_REP_LIGHT = (255, 247, 194) # pale yellow highlight for repeated-line copies +_REP_DARK = (44, 44, 24) + +# High-frequency function words a reader can reconstruct from context; the dim +# variants render them in light gray so content words carry the contrast. +_STOPWORDS = frozenset( + "the a an and or of to in on at as is are was were be been by for with that this it its from had has have not but " + "he she his her they their them which also who whom when where while will would could should there then than " + "into over under about after before between during each such these those some most more other only same so".split() +) + + +def _row_palette(variant: str, row: int) -> tuple[tuple[int, int, int], tuple[int, int, int]]: + """(background, default glyph) colors for a row under the given render variant.""" + if variant == "color": + return _PALE[row % 6], _DARK[row % 6] + if variant == "zebra": + return (_WHITE if row % 2 == 0 else _GRAY), _BLACK + if variant in ("bw", "sent", "dim", "sent-dim"): + return _WHITE, _BLACK + if variant in ("dark", "dark-sent", "dark-sent-dim"): + return _BLACK, _WHITE + raise ValueError(f"unknown variant: {variant}") + + +def _sentence_indices(text: str) -> list[int]: + """Running sentence index per character (boundary: terminator + space).""" + out = [0] * len(text) + idx = 0 + for i, ch in enumerate(text): + out[i] = idx + if ch in ".!?" and i + 1 < len(text) and text[i + 1] == " ": + idx += 1 + return out + + +def _stopword_mask(text: str) -> list[bool]: + """True for every character of a word in _STOPWORDS.""" + mask = [False] * len(text) + i = 0 + while i < len(text): + if text[i].isalpha(): + j = i + while j < len(text) and text[j].isalpha(): + j += 1 + if text[i:j].lower() in _STOPWORDS: + for k in range(i, j): + mask[k] = True + i = j + else: + i += 1 + return mask + + +def capacity(cfg: FontCfg, size: int = 1568, columns: int = 1) -> tuple[int, int, int]: + """(cols per line, rows, chars) that fit a size x size image with `columns` newspaper columns.""" + rows = size // cfg.pitch // cfg.repeat + gutter = 2 * cfg.adv if columns > 1 else 0 + cols = (size - (columns - 1) * gutter) // columns // cfg.adv + return cols, rows, columns * cols * rows + + +def render( + text: str, cfg: FontCfg, cache: Path, size: int = 1568, variant: str = "color", columns: int = 1 +) -> Image.Image: + """Fill a size x size grid with `text`; styling per `variant`. + + Layout: full-width row-major when columns == 1; otherwise newspaper flow + (fill the leftmost column top-to-bottom, then the next), columns separated + by a 2-cell gutter with a hairline rule. + + When cfg.native is set, glyphs are rasterized at the native cell size and + the whole canvas is Lanczos-resized to the adv x pitch target (anisotropic + stretch: anti-aliased, no cropping or ink fusion). + + Variants: + color per-row hue cycle on pale row bands + zebra black text, alternating white/gray bands + bw black on white + sent white background, glyph hue cycles per sentence + dark white text on black + dark-sent bright sentence hues on black + dim black on white, stopwords dimmed gray + sent-dim sentence hues, stopwords dimmed gray + dark-sent-dim bright sentence hues on black, stopwords dimmed + """ + glyphs, font_ascent = load_font(cfg, cache) + ascent = cfg.ascent if cfg.ascent is not None else font_ascent + cols, rows, cap = capacity(cfg, size, columns) + text = text[:cap] + sent_idx = _sentence_indices(text) if variant in ("sent", "dark-sent", "sent-dim", "dark-sent-dim") else None + dim_mask = _stopword_mask(text) if variant in ("dim", "sent-dim", "dark-sent-dim") else None + sent_palette = _BRIGHT if variant in ("dark-sent", "dark-sent-dim") else _DARK + dark_bg = variant in ("dark", "dark-sent", "dark-sent-dim") + base_color = _BLACK if dark_bg else _WHITE + if cfg.native is not None: + aw, ph = cfg.native + gutter = 2 * aw if columns > 1 else 0 + span = cols * aw + gutter + canvas_w, canvas_h = columns * span - gutter, rows * cfg.repeat * ph + else: + aw, ph = cfg.adv, cfg.pitch + gutter = 2 * aw if columns > 1 else 0 + span = cols * aw + gutter + canvas_w = canvas_h = size + img = Image.new("RGB", (canvas_w, canvas_h), base_color) + px = img.load() + for row in range(rows): + bg, row_fg = _row_palette(variant, row) + for copy in range(cfg.repeat): + y0 = (row * cfg.repeat + copy) * ph + cbg = bg if copy == 0 else (_REP_DARK if dark_bg else _REP_LIGHT) + for y in range(y0, min(y0 + ph, canvas_h)): + for x in range(canvas_w): + px[x, y] = cbg + for blk in range(columns): + for col in range(cols): + i = (blk * rows + row) * cols + col + if i >= len(text): + break + glyph = glyphs.get(ord(text[i])) + if glyph is None: + continue + fg = row_fg + if sent_idx is not None: + fg = sent_palette[sent_idx[i] % 6] + if dim_mask is not None and dim_mask[i]: + fg = _DIMMED_DARK if dark_bg else _DIMMED + w, h, xoff, yoff = glyph["bbx"] + top = y0 + ascent - h - yoff + shift = 0x80 if w <= 8 else 0x8000 + for r, bits in enumerate(glyph["rows"]): + y = top + r + if not 0 <= y < canvas_h: + continue + for b in range(w): + if bits & (shift >> b): + x = blk * span + col * aw + xoff + b + if 0 <= x < canvas_w: + px[x, y] = fg + rule = (96, 96, 96) if dark_bg else (204, 204, 204) + for blk in range(1, columns): + x = blk * span - gutter // 2 + if 0 <= x < canvas_w: + for y in range(canvas_h): + px[x, y] = rule + if cfg.native is not None: + img = img.resize((canvas_w * cfg.adv // aw, canvas_h * cfg.pitch // ph), Image.LANCZOS) + if img.size != (size, size): + out = Image.new("RGB", (size, size), base_color) + out.paste(img, (0, 0)) + img = out + return img diff --git a/packages/snapcompact/research/exp01_patchalign.py b/packages/snapcompact/research/exp01_patchalign.py new file mode 100644 index 000000000..29431073c --- /dev/null +++ b/packages/snapcompact/research/exp01_patchalign.py @@ -0,0 +1,284 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp01: ViT patch-alignment hypothesis. + +Vision encoders patch at fixed pixel grids (14/16 px, often 2x2-merged; Gemini +tiles at 768 px = 48*16). Current fonts (pitch 10/12/13) straddle patch +boundaries so glyph rows smear across visual tokens. Test patch-aligned cell +grids against same-glyph-budget misaligned controls: + + aligned: img-7x14-sent (native X.Org 7x14), + img-8x16-sent (native Spleen 8x16; X.Org misc-misc has no 8x16), + img-8on16-sent (8x13 glyphs on an 8x16 cell: identical glyphs to + the control, ONLY the pitch changes), + img-6on7x14-sent (6x12 glyphs on a 7x14 cell) + controls: img-7x13-sent, img-8x13-sent (same glyph width, pitch 13) + +Render-size probe: 1568 (baseline size) vs 1536 = 2*768 (exact Gemini tile +multiple; integer 2x downsample keeps 16 px pitch on an 8 px grid; also 3*512 +for OpenAI tile schemes) for the 16 px-pitch fonts. + +Usage: uv run exp01_patchalign.py # full grid + uv run exp01_patchalign.py --report # re-aggregate from cache +""" + +import argparse +import csv +import json +import sys +import urllib.request +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import FontCfg, capacity, render # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 + +MODELS = { + "gpt-5.5": (2.0, 16.0), + "google/gemini-3.5-flash": (0.6, 4.0), +} +FONTS = { + "7x14": FontCfg("7x14", "7x14", 7, 14), # aligned: pitch = 14 px patch + "8x16": FontCfg("8x16", "spleen-8x16", 8, 16), # aligned: native 16 px font (Spleen) + "8on16": FontCfg("8on16", "8x13", 8, 16), # aligned: 8x13 glyphs, pitch 16 cell + "6on7x14": FontCfg("6on7x14", "6x12", 7, 14), # aligned: 6x12 glyphs, 7x14 cell + "7x13": FontCfg("7x13", "7x13", 7, 13), # control for 7x14 (same glyph budget, pitch 13) + "8x13": FontCfg("8x13", "8x13", 8, 13), # control for 8x16/8on16 +} +SPLEEN_URL = "https://raw.githubusercontent.com/fcambus/spleen/master/spleen-8x16.bdf" +# (condition, render size) cells. 1536 only for 16 px-pitch fonts (14 does not +# divide 768 or 1536). +GRID = ( + ("img-7x14-sent", 1568), + ("img-7x13-sent", 1568), + ("img-8x16-sent", 1568), + ("img-8on16-sent", 1568), + ("img-8x13-sent", 1568), + ("img-6on7x14-sent", 1568), + ("img-8x16-sent", 1536), + ("img-8on16-sent", 1536), +) +LENGTHS = (50, 150) + + +def cached(model: str, tag: str, payload: object, fn, fresh: bool) -> dict: + """Disk-cache `fn() -> dict` keyed by (model, tag, payload). Truncations are not cached.""" + key = sha8(model, tag, json.dumps(payload, sort_keys=True, default=str)) + path = QA_CACHE / f"{key}.json" + if path.exists() and not fresh: + hit = json.loads(path.read_text()) + if hit.get("stop") != "max_tokens": + return hit + out = fn() + if out.get("stop") == "max_tokens": + print(f" WARN truncated, not cached: {model} {tag} {key}") + else: + path.write_text(json.dumps(out)) + return out + + +def parse_img_condition(name: str) -> tuple[str, str]: + _, font, variant = name.split("-", 2) + return font, variant + + +def ensure_spleen() -> None: + """bdf.ensure_font only knows X.Org/tom-thumb URLs; stage Spleen 8x16 ourselves.""" + path = CACHE / "spleen-8x16.bdf" + if path.exists() and path.stat().st_size > 0: + return + tmp = path.with_suffix(".tmp.bdf") + urllib.request.urlretrieve(SPLEEN_URL, tmp) + tmp.replace(path) + + +def run_cell_chunk(model: str, cond: str, size: int, start: int, end: int, ctx: dict) -> list[dict]: + """One (model, condition, size, chunk) unit: render carrier, QA, score.""" + args, flow, paras, offsets, keys = ctx["args"], ctx["flow"], ctx["paras"], ctx["offsets"], ctx["keys"] + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + chunk_text = flow[start:end] + q_block = "\n".join(f"{i + 1}. {q['q']}" for i, q in enumerate(questions)) + + font, variant = parse_img_condition(cond) + png = CACHE / f"exp01-{font}-{variant}-{sha8(chunk_text, str(size))}.png" + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(".tmp.png") + render(chunk_text, FONTS[font], CACHE, size, variant).save(tmp) + tmp.replace(png) + cols, rows, _ = capacity(FONTS[font], size) + messages = [ + { + "role": "user", + "content": [ + {"text": load_prompt("qa-image.md").format(cols=cols, rows=rows)}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + qa = cached( + model, "exp01-qa", {"messages": messages, "size": size, "effort": args.effort}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete(keys, model, messages, max_tokens=args.max_tokens, effort=args.effort), + ) + ), + args.fresh, + ) + answers = squad.parse_numbered(qa["text"], len(questions)) + records = [] + for q, a in zip(questions, answers): + records.append( + { + "model": model, + "length": ctx["length"], + "cond": cond, + "size": size, + "chunk": start, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + ) + records[0]["usage"] = [{"phase": "qa", **qa["usage"]}] + return records + + +def aggregate(records: list[dict], price_in: float, price_out: float) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean_f1 = sum(f1s) / n + se = (sum((x - mean_f1) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + us = [u for r in records if "usage" in r for u in r["usage"]] + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + cost_in = (tok["in"] + 1.25 * tok["cache_w"] + 0.1 * tok["cache_r"]) / 1e6 * price_in + cost_out = tok["out"] / 1e6 * price_out + return { + "n": n, + "em": sum(r["em"] for r in records) / n, + "f1": mean_f1, + "f1_se": se, + "abstained": sum(r["abstained"] for r in records), + **{f"tok_{k}": v for k, v in tok.items()}, + "cost_in_usd": round(cost_in, 4), + "cost_out_usd": round(cost_out, 4), + "cost_usd": round(cost_in + cost_out, 4), + } + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--models", default=",".join(MODELS)) + ap.add_argument("--lengths", default=",".join(map(str, LENGTHS))) + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--effort", default=None) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--report", action="store_true", help="reprint from cache only (re-runs cells; all should hit cache)") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + ensure_spleen() + out_dir = RESULTS / "exp01-patchalign" + out_dir.mkdir(parents=True, exist_ok=True) + + models = [m.strip() for m in args.models.split(",") if m.strip()] + lengths = [int(x) for x in args.lengths.split(",") if x.strip()] + + keys = { + "openai": load_env_key("OPENAI_API_KEY", args.env), + "openrouter": load_env_key("OPENROUTER_API_KEY", args.env), + } + + all_paras = squad.load_paragraphs(CACHE) + tasks = [] + for length in lengths: + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + ctx = {"args": args, "flow": flow, "paras": paras, "offsets": offsets, "keys": keys, "length": length} + for model in models: + for cond, size in GRID: + budget = capacity(FONTS[parse_img_condition(cond)[0]], size)[2] + for start in range(0, len(flow), budget): + tasks.append((model, cond, size, start, min(start + budget, len(flow)), ctx)) + print(f"grid: {len(models)} models x {len(lengths)} lengths x {len(GRID)} cells = {len(tasks)} chunk tasks") + + records: list[dict] = [] + done = 0 + with ThreadPoolExecutor(args.workers) as pool: + futures = [pool.submit(run_cell_chunk, m, c, sz, s, e, ctx) for m, c, sz, s, e, ctx in tasks] + for fut in futures: + records.extend(fut.result()) + done += 1 + if done % 10 == 0: + print(f" {done}/{len(tasks)} tasks", flush=True) + + with (out_dir / "records.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + + cells = [] + for model in models: + for length in lengths: + for cond, size in GRID: + sub = [ + r for r in records + if r["model"] == model and r["length"] == length and r["cond"] == cond and r["size"] == size + ] + if not sub: + continue + cells.append( + { + "model": model, + "length": length, + "condition": f"{cond}@{size}", + **aggregate(sub, *MODELS[model]), + } + ) + (out_dir / "summary.json").write_text(json.dumps({"args": vars(args), "cells": cells}, indent=1)) + with (out_dir / "matrix.csv").open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=list(cells[0].keys())) + writer.writeheader() + writer.writerows(cells) + + for length in lengths: + print(f"\n== {length} passages == (F1 +-se / n / $)") + hdr = f"{'condition':<22}" + "".join(f"{m:>34}" for m in models) + print(hdr + "\n" + "-" * len(hdr)) + for cond, size in GRID: + label = f"{cond}@{size}" + row = f"{label:<22}" + for model in models: + cell = next( + (c for c in cells if c["model"] == model and c["length"] == length and c["condition"] == label), + None, + ) + row += ( + f"{cell['f1']:>10.3f} +-{cell['f1_se']:.3f} {cell['n']:>4} {cell['cost_usd']:>7.3f}" + if cell + else f"{'-':>34}" + ) + print(row) + print(f"\ndataset -> {out_dir}/records.jsonl, matrix.csv, summary.json") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp02_surprisal.py b/packages/snapcompact/research/exp02_surprisal.py new file mode 100644 index 000000000..63fe857f6 --- /dev/null +++ b/packages/snapcompact/research/exp02_surprisal.py @@ -0,0 +1,427 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "wordfreq"] +# /// +"""exp02: surprisal-weighted glyph contrast + stopword-removal/disemvowel density. + +Conditions (all 6x10 font, 1568px, vs baseline img-6x10-sent): + img-6x10-surp full text; glyph gray level = unigram surprisal bucket + (wordfreq zipf: ultra-common -> light gray, rare -> black) + img-6x10-sent-surp sentence hues; lightness scaled by surprisal bucket + img-6x10-disemv stopwords dropped + lowercase non-entity words disemvoweled, + rendered with plain `sent` hues; chunks sized by *transformed* + capacity so each page carries more original chars (fewer pages) + +Usage: uv run exp02_surprisal.py [--report] [--fresh] +Keys: OPENAI_API_KEY + OPENROUTER_API_KEY from ~/.env. +""" + +import argparse +import colorsys +import csv +import json +import re +import sys +from concurrent.futures import ThreadPoolExecutor +from functools import lru_cache +from pathlib import Path + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import _sentence_indices, capacity, ensure_font, parse_bdf, render # noqa: E402 +from PIL import Image # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, FONTS, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 +from wordfreq import zipf_frequency # noqa: E402 + +MODELS = {"gpt-5.5": (2.0, 16.0), "google/gemini-3.5-flash": (0.6, 4.0)} +LENGTHS = (50, 150) +CONDITIONS = ("img-6x10-surp", "img-6x10-sent-surp", "img-6x10-disemv") +FONT = FONTS["6x10"] + +# Baseline cells from results/optimal-*/matrix.csv (img-6x10-sent): F1, se, cost$. +BASELINE = { + ("gpt-5.5", 50): (0.850, 0.051, 0.068), + ("gpt-5.5", 150): (0.822, 0.029, 0.245), + ("google/gemini-3.5-flash", 50): (0.984, 0.012, 0.018), + ("google/gemini-3.5-flash", 150): (0.805, 0.035, 0.097), +} + +SURP_NOTE = ( + "Glyph darkness encodes word informativeness: very common words are printed in lighter gray, " + "rarer / more informative words in darker ink. All words are spelled out in full." +) +DISEMV_NOTE = ( + "To fit more text, very common function words (the, of, and, is, ...) were removed, and other " + 'common lowercase words are abbreviated by stripping their interior vowels (e.g. "qck brwn fx jmpd" ' + 'means "quick brown fox jumped"; "gvrnmnt" means "government"). Proper nouns, capitalized words, ' + "numbers, dates, and words adjacent to numbers are kept verbatim. Mentally reconstruct the original " + "wording; ALWAYS write your answers in normal, fully spelled English." +) + +# --- text transform: stopword drop + disemvowel ------------------------------ + +# bdf._STOPWORDS minus "not" (removal flips meaning; dimming it was harmless). +_STOP = frozenset( + "the a an and or of to in on at as is are was were be been by for with that this it its from had has have but " + "he she his her they their them which also who whom when where while will would could should there then than " + "into over under about after before between during each such these those some most more other only same so".split() +) +# Spelled-out numbers/ordinals/units stay verbatim: golds are full of them. +_NUMWORDS = frozenset( + "zero one two three four five six seven eight nine ten eleven twelve thirteen fourteen fifteen sixteen " + "seventeen eighteen nineteen twenty thirty forty fifty sixty seventy eighty ninety hundred thousand million " + "billion trillion first second third fourth fifth sixth seventh eighth ninth tenth half quarter percent".split() +) +_VOWELS = frozenset("aeiou") +_TOKEN_RE = re.compile(r"([^A-Za-z]*)([a-z]+)([^A-Za-z]*)") +_WORD_RE = re.compile(r"[A-Za-z]+") + + +def _has_digit(s: str) -> bool: + return any(c.isdigit() for c in s) + + +def _disemvowel(w: str) -> str: + if len(w) <= 2: + return w + return w[0] + "".join(c for c in w[1:-1] if c not in _VOWELS) + w[-1] + + +def transform(text: str) -> str: + """Drop lowercase stopwords, disemvowel other lowercase words. + + Never touches: capitalized words, tokens containing digits or apostrophes, + tokens adjacent to a digit-bearing token, spelled-out numbers/units. + """ + toks = text.split() + out: list[str] = [] + for i, tok in enumerate(toks): + m = _TOKEN_RE.fullmatch(tok) + if not m: # uppercase, digits, apostrophes, hyphens-with-letters: verbatim + out.append(tok) + continue + pre, core, suf = m.groups() + if ( + _has_digit(pre) + or _has_digit(suf) + or (i > 0 and _has_digit(toks[i - 1])) + or (i + 1 < len(toks) and _has_digit(toks[i + 1])) + or core in _NUMWORDS + ): + out.append(tok) + continue + if core in _STOP: + if pre + suf: # keep punctuation of dropped words + out.append(pre + suf) + continue + out.append(pre + _disemvowel(core) + suf) + return " ".join(out) + + +# --- surprisal-weighted rendering -------------------------------------------- + +# 4 buckets, near-black (rare) .. light gray (ultra-common). Lightest matches +# the old _DIMMED (176) readability point. +_GRAYS = [(185, 185, 185), (135, 135, 135), (75, 75, 75), (0, 0, 0)] +_LIGHT = [0.72, 0.55, 0.40, 0.22] # lightness per bucket for sent hues +_HUES = [0.0, 0.08, 0.3, 0.5, 0.62, 0.78] +_SENT_SURP = [ + [tuple(int(c * 255) for c in colorsys.hls_to_rgb(h, l, 0.95)) for l in _LIGHT] for h in _HUES +] +_WHITE = (255, 255, 255) + + +@lru_cache(maxsize=65536) +def _zipf_bucket(word: str) -> int: + """0 = ultra-common (lightest) .. 3 = rare/unknown (black).""" + z = zipf_frequency(word, "en") + if z >= 6.0: + return 0 + if z >= 5.0: + return 1 + if z >= 4.0: + return 2 + return 3 + + +def _shade_indices(text: str) -> list[int]: + """Per-char surprisal bucket; non-letters (digits, punctuation) stay black.""" + out = [3] * len(text) + for m in _WORD_RE.finditer(text): + b = _zipf_bucket(m.group().lower()) + if b != 3: + for k in range(m.start(), m.end()): + out[k] = b + return out + + +def render_surp(text: str, cfg, cache: Path, size: int = 1568, sent_hues: bool = False) -> Image.Image: + """bdf.render() with the boolean dim_mask generalized to surprisal buckets.""" + glyphs, font_ascent = parse_bdf(ensure_font(cfg, cache)) + ascent = cfg.ascent if cfg.ascent is not None else font_ascent + cols, rows, cap = capacity(cfg, size) + text = text[:cap] + shades = _shade_indices(text) + sent_idx = _sentence_indices(text) if sent_hues else None + img = Image.new("RGB", (size, size), _WHITE) + px = img.load() + for row in range(rows): + y0 = row * cfg.pitch + for col in range(cols): + i = row * cols + col + if i >= len(text): + break + glyph = glyphs.get(ord(text[i])) + if glyph is None: + continue + b = shades[i] + fg = _SENT_SURP[sent_idx[i] % 6][b] if sent_idx is not None else _GRAYS[b] + w, h, xoff, yoff = glyph["bbx"] + top = y0 + ascent - h - yoff + shift = 0x80 if w <= 8 else 0x8000 + for r, bits in enumerate(glyph["rows"]): + y = top + r + if not 0 <= y < size: + continue + for bcol in range(w): + if bits & (shift >> bcol): + x = col * cfg.adv + xoff + bcol + if 0 <= x < size: + px[x, y] = fg + return img + + +# --- chunk planning ----------------------------------------------------------- + + +def plan_chunks(flow: str, cond: str, size: int) -> list[tuple[int, int]]: + cap = capacity(FONT, size)[2] + if cond != "img-6x10-disemv": + return [(s, min(s + cap, len(flow))) for s in range(0, len(flow), cap)] + # Greedy: max original span whose *transformed* text fits a page; word-snapped. + chunks: list[tuple[int, int]] = [] + s, n = 0, len(flow) + while s < n: + if len(transform(flow[s:n])) <= cap: + chunks.append((s, n)) + break + lo, hi = s + cap, n # transform never lengthens => s+cap always fits + while lo < hi: + mid = (lo + hi + 1) // 2 + if len(transform(flow[s:mid])) <= cap: + lo = mid + else: + hi = mid - 1 + e = lo + sp = flow.rfind(" ", s, e) # don't split a word across pages + if sp > s: + e = sp + chunks.append((s, e)) + s = e + (1 if e < n and flow[e] == " " else 0) + return chunks + + +# --- harness (mirrors final.py) ----------------------------------------------- + + +def cached(model: str, tag: str, payload: object, fn, fresh: bool) -> dict: + key = sha8(model, tag, json.dumps(payload, sort_keys=True, default=str)) + path = QA_CACHE / f"{key}.json" + if path.exists() and not fresh: + hit = json.loads(path.read_text()) + if hit.get("stop") != "max_tokens": + return hit + out = fn() + if out.get("stop") == "max_tokens": + print(f" WARN truncated, not cached: {model} {tag} {key}") + else: + path.write_text(json.dumps(out)) + return out + + +def build_png(cond: str, render_text: str, size: int) -> Path: + png = CACHE / f"exp02-{cond}-{sha8(render_text, str(size))}.png" + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(".tmp.png") + if cond == "img-6x10-surp": + img = render_surp(render_text, FONT, CACHE, size, sent_hues=False) + elif cond == "img-6x10-sent-surp": + img = render_surp(render_text, FONT, CACHE, size, sent_hues=True) + else: # disemv: text transform is the variable; keep baseline sent hues + img = render(render_text, FONT, CACHE, size, "sent") + img.save(tmp) + tmp.replace(png) + return png + + +def run_chunk(model: str, cond: str, start: int, end: int, png: Path, render_chars: int, ctx: dict) -> list[dict]: + args, keys = ctx["args"], ctx["keys"] + questions = squad.sample_chunk_questions(ctx["paras"], ctx["offsets"], start, end, args.qpc, args.seed) + if not questions: + return [] + q_block = "\n".join(f"{i + 1}. {q['q']}" for i, q in enumerate(questions)) + cols, rows, _ = capacity(FONT, args.size) + extra = DISEMV_NOTE if cond == "img-6x10-disemv" else SURP_NOTE + messages = [ + { + "role": "user", + "content": [ + {"text": load_prompt("exp02-qa-image.md").format(cols=cols, rows=rows, extra=extra)}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + qa = cached( + model, "exp02-qa", {"messages": messages, "effort": args.effort}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete(keys, model, messages, max_tokens=args.max_tokens, effort=args.effort), + ) + ), + args.fresh, + ) + answers = squad.parse_numbered(qa["text"], len(questions)) + records = [] + for q, a in zip(questions, answers): + records.append( + { + "model": model, + "length": ctx["length"], + "cond": cond, + "chunk": start, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + ) + records[0]["usage"] = [{"phase": "qa", **qa["usage"]}] + records[0]["chunk_orig_chars"] = end - start + records[0]["chunk_render_chars"] = render_chars + return records + + +def aggregate(records: list[dict], price_in: float, price_out: float) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean_f1 = sum(f1s) / n + se = (sum((x - mean_f1) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + us = [u for r in records if "usage" in r for u in r["usage"]] + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + cost_in = (tok["in"] + 1.25 * tok["cache_w"] + 0.1 * tok["cache_r"]) / 1e6 * price_in + cost_out = tok["out"] / 1e6 * price_out + pages = sum(1 for r in records if "chunk_orig_chars" in r) + orig_chars = sum(r.get("chunk_orig_chars", 0) for r in records) + return { + "n": n, + "em": sum(r["em"] for r in records) / n, + "f1": mean_f1, + "f1_se": se, + "abstained": sum(r["abstained"] for r in records), + "pages": pages, + "orig_chars": orig_chars, + "orig_chars_per_page": round(orig_chars / pages) if pages else 0, + **{f"tok_{k}": v for k, v in tok.items()}, + "cost_in_usd": round(cost_in, 4), + "cost_out_usd": round(cost_out, 4), + "cost_usd": round(cost_in + cost_out, 4), + } + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--models", default=",".join(MODELS)) + ap.add_argument("--lengths", default=",".join(map(str, LENGTHS))) + ap.add_argument("--conditions", default=",".join(CONDITIONS)) + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--effort", default=None) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--report", action="store_true", help="reprint from cache only") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + out_dir = RESULTS / "exp02-surprisal" + out_dir.mkdir(parents=True, exist_ok=True) + + models = [m.strip() for m in args.models.split(",") if m.strip()] + lengths = [int(x) for x in args.lengths.split(",") if x.strip()] + conditions = [c.strip() for c in args.conditions.split(",") if c.strip()] + + keys = {} + if not args.report: + keys["openai"] = load_env_key("OPENAI_API_KEY", args.env) + keys["openrouter"] = load_env_key("OPENROUTER_API_KEY", args.env) + + all_paras = squad.load_paragraphs(CACHE) + tasks = [] + for length in lengths: + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + ctx = {"args": args, "paras": paras, "offsets": offsets, "keys": keys, "length": length} + for cond in conditions: + chunks = plan_chunks(flow, cond, args.size) + for start, end in chunks: + orig = flow[start:end] + render_text = transform(orig) if cond == "img-6x10-disemv" else orig + png = build_png(cond, render_text, args.size) # pre-render: no tmp races in pool + if cond == "img-6x10-disemv": + print( + f" len={length} disemv chunk [{start},{end}): {end - start} orig -> " + f"{len(render_text)} rendered chars (x{(end - start) / len(render_text):.2f})" + ) + for model in models: + tasks.append((model, cond, start, end, png, len(render_text), ctx)) + print(f"grid: {len(models)} models x {len(lengths)} lengths x {len(conditions)} conditions = {len(tasks)} chunk tasks") + + records: list[dict] = [] + with ThreadPoolExecutor(args.workers) as pool: + futures = [pool.submit(run_chunk, *t) for t in tasks] + for i, fut in enumerate(futures): + records.extend(fut.result()) + print(f" {i + 1}/{len(tasks)} tasks", flush=True) + + with (out_dir / "records.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + + cells = [] + for model in models: + for length in lengths: + for cond in conditions: + sub = [r for r in records if r["model"] == model and r["length"] == length and r["cond"] == cond] + if not sub: + continue + cells.append({"model": model, "length": length, "condition": cond, **aggregate(sub, *MODELS[model])}) + (out_dir / "summary.json").write_text(json.dumps({"args": vars(args), "cells": cells}, indent=1)) + with (out_dir / "matrix.csv").open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=list(cells[0].keys())) + writer.writeheader() + writer.writerows(cells) + + print(f"\n{'model':<26}{'len':>5}{'condition':<22}{'n':>4}{'EM':>7}{'F1':>7}{'se':>7}{'cost$':>8}{'dF1 vs base':>13}") + for c in cells: + base = BASELINE.get((c["model"], c["length"])) + d = f"{c['f1'] - base[0]:+.3f}" if base else "-" + print( + f"{c['model']:<26}{c['length']:>5} {c['condition']:<20}{c['n']:>4}{c['em']:>7.3f}" + f"{c['f1']:>7.3f}{c['f1_se']:>7.3f}{c['cost_usd']:>8.3f}{d:>13}" + ) + print(f"\ndataset -> {out_dir}/records.jsonl, matrix.csv, summary.json") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp03_numhard.py b/packages/snapcompact/research/exp03_numhard.py new file mode 100644 index 000000000..115fbb96c --- /dev/null +++ b/packages/snapcompact/research/exp03_numhard.py @@ -0,0 +1,382 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp03_numhard: harden digit glyphs in optical-compaction renders. + +Extractive QA on rendered pages dies on 0/O, 1/l, rn/m confusions, and SQuAD +answers skew heavily toward numbers/dates. This experiment re-renders the +img-6x10 baseline with a hardening mask over digits (plus number/date +punctuation directly adjacent to a digit: "1,000", "3.5%", "1914-18"): + + numbold double-strike: every masked glyph pixel painted at x and x+1 + (6x10 digit ink spans cols 0-4 of the 6px cell, so x+1 stays + in-cell -- no bleed into the neighbor) + numred masked glyphs painted pure dark red, overriding sent/bw color + +Conditions: img-6x10-sent-numbold, img-6x10-sent-numred, img-6x10-bw-numred. +Methodology matches the baseline grid (seed 42, qpc 30, size 1568, +max_tokens 32768, effort None) so question sets are identical and the +numeric-gold subset can be compared apples-to-apples against the baseline +records in results/optimal-gpt55 / results/optimal-gemini. + +Usage: uv run exp03_numhard.py [--report] +""" + +import argparse +import csv +import json +import sys +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +from PIL import Image + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import _DARK, _row_palette, _sentence_indices, capacity, ensure_font, parse_bdf # noqa: E402 +from final import MODELS, aggregate, cached # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, FONTS, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 + +EXP = "exp03" +SLUG = "numhard" +FONT = FONTS["6x10"] +CONDITIONS = ("img-6x10-sent-numbold", "img-6x10-sent-numred", "img-6x10-bw-numred") +RUN_MODELS = ("gpt-5.5", "google/gemini-3.5-flash") +BASELINE_COND = "img-6x10-sent" +BASELINE_RECORDS = { + "gpt-5.5": RESULTS / "optimal-gpt55" / "records.jsonl", + "google/gemini-3.5-flash": RESULTS / "optimal-gemini" / "records.jsonl", +} +RED = (220, 0, 0) +_WHITE = (255, 255, 255) +_NUM_PUNCT = set(".,:/%-\u2013$") + + +def number_mask(text: str) -> list[bool]: + """True for digits and number/date punctuation directly adjacent to a digit.""" + mask = [False] * len(text) + for i, ch in enumerate(text): + if ch.isdigit(): + mask[i] = True + elif ch in _NUM_PUNCT: + if (i > 0 and text[i - 1].isdigit()) or (i + 1 < len(text) and text[i + 1].isdigit()): + mask[i] = True + return mask + + +def render_hard(text: str, cfg, cache: Path, size: int, variant: str, hard: str) -> Image.Image: + """Copy of bdf.render() restricted to white-bg variants (sent/bw), with a + digit-hardening pass: `numbold` double-strikes masked glyphs, `numred` + recolors them pure red.""" + glyphs, font_ascent = parse_bdf(ensure_font(cfg, cache)) + ascent = cfg.ascent if cfg.ascent is not None else font_ascent + cols, rows, cap = capacity(cfg, size) + text = text[:cap] + sent_idx = _sentence_indices(text) if variant == "sent" else None + num_mask = number_mask(text) + img = Image.new("RGB", (size, size), _WHITE) + px = img.load() + for row in range(rows): + bg, fg_default = _row_palette(variant, row) + y0 = row * cfg.pitch + for y in range(y0, min(y0 + cfg.pitch, size)): + for x in range(size): + px[x, y] = bg + for col in range(cols): + i = row * cols + col + if i >= len(text): + break + glyph = glyphs.get(ord(text[i])) + if glyph is None: + continue + fg = _DARK[sent_idx[i] % 6] if sent_idx is not None else fg_default + hardened = num_mask[i] + if hardened and hard == "numred": + fg = RED + bold = hardened and hard == "numbold" + w, h, xoff, yoff = glyph["bbx"] + top = y0 + ascent - h - yoff + shift = 0x80 if w <= 8 else 0x8000 + for r, bits in enumerate(glyph["rows"]): + y = top + r + if not 0 <= y < size: + continue + for b in range(w): + if bits & (shift >> b): + x = col * cfg.adv + xoff + b + if 0 <= x < size: + px[x, y] = fg + if bold and x + 1 < size: + px[x + 1, y] = fg + return img + + +def parse_cond(cond: str) -> tuple[str, str]: + """img-6x10-- -> (base, hard).""" + parts = cond.split("-") + return parts[2], parts[3] + + +def chunk_png(chunk_text: str, size: int, base: str, hard: str) -> Path: + png = CACHE / f"{EXP}-img-6x10-{base}-{hard}-{sha8(chunk_text, str(size))}.png" + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(".tmp.png") + render_hard(chunk_text, FONT, CACHE, size, base, hard).save(tmp) + tmp.replace(png) # atomic; cache dir is shared across agents + return png + + +def run_chunk(model: str, cond: str, start: int, end: int, ctx: dict) -> list[dict]: + args, flow = ctx["args"], ctx["flow"] + questions = squad.sample_chunk_questions(ctx["paras"], ctx["offsets"], start, end, args.qpc, args.seed) + if not questions: + return [] + chunk_text = flow[start:end] + base, hard = parse_cond(cond) + png = chunk_png(chunk_text, args.size, base, hard) + cols, rows, _ = capacity(FONT, args.size) + q_block = "\n".join(f"{i + 1}. {q['q']}" for i, q in enumerate(questions)) + messages = [ + { + "role": "user", + "content": [ + {"text": load_prompt("qa-image.md").format(cols=cols, rows=rows)}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + qa = cached( + model, f"{EXP}-qa", {"messages": messages, "effort": args.effort}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete(ctx["keys"], model, messages, max_tokens=args.max_tokens, effort=args.effort), + ) + ), + args.fresh, + ) + answers = squad.parse_numbered(qa["text"], len(questions)) + records = [] + for q, a in zip(questions, answers): + records.append( + { + "model": model, + "length": ctx["length"], + "cond": cond, + "chunk": start, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + ) + records[0]["usage"] = [{"phase": "qa", **qa["usage"]}] + return records + + +# --- analysis --- + + +def is_numeric_gold(golds: list[str]) -> bool: + return any(any(c.isdigit() for c in g) for g in golds) + + +def f1_stats(records: list[dict]) -> dict: + n = len(records) + if n == 0: + return {"n": 0, "em": 0.0, "f1": 0.0, "f1_se": 0.0} + f1s = [r["f1"] for r in records] + mean = sum(f1s) / n + se = (sum((x - mean) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + return {"n": n, "em": sum(r["em"] for r in records) / n, "f1": mean, "f1_se": se} + + +def load_baseline(model: str, lengths: list[int]) -> list[dict]: + out = [] + with BASELINE_RECORDS[model].open() as fh: + for line in fh: + r = json.loads(line) + if r["cond"] == BASELINE_COND and r["length"] in lengths: + out.append(r) + return out + + +def numeric_subset_cells(records: list[dict], models: list[str], lengths: list[int], conditions: list[str]) -> list[dict]: + """Per (model, length): baseline vs each condition, restricted to numeric-gold + questions present in BOTH runs (matched by question text).""" + cells = [] + for model in models: + base = load_baseline(model, lengths) + for length in lengths: + base_num = {r["q"]: r for r in base if r["length"] == length and is_numeric_gold(r["golds"])} + for cond in conditions: + mine = [ + r for r in records + if r["model"] == model and r["length"] == length and r["cond"] == cond + and is_numeric_gold(r["golds"]) and r["q"] in base_num + ] + if not mine: + continue + base_match = [base_num[r["q"]] for r in mine] + cells.append( + { + "model": model, + "length": length, + "condition": cond, + **f1_stats(mine), + "baseline_f1": f1_stats(base_match)["f1"], + "baseline_se": f1_stats(base_match)["f1_se"], + } + ) + return cells + + +def save_sample(records_ctx_flow: str, size: int, base: str, hard: str, out_dir: Path) -> Path: + """Crop a digit-dense region from the first chunk's PNG, 4x nearest upscale.""" + cols, rows, cap = capacity(FONT, size) + text = records_ctx_flow[:cap] + mask = number_mask(text) + # densest 60-char window + best_i, best_n = 0, -1 + win = 60 + run = sum(mask[:win]) + for i in range(len(text) - win): + if run > best_n: + best_n, best_i = run, i + run += mask[i + win] - mask[i] + row, col = best_i // cols, best_i % cols + png = chunk_png(records_ctx_flow[:cap], size, base, hard) + img = Image.open(png) + x0 = max(0, min(col, cols - win) * FONT.adv) + y0 = max(0, (row - 1) * FONT.pitch) + crop = img.crop((x0, y0, min(x0 + win * FONT.adv, size), min(y0 + 4 * FONT.pitch, size))) + crop = crop.resize((crop.width * 4, crop.height * 4), Image.NEAREST) + out = out_dir / f"sample-{base}-{hard}.png" + crop.save(out) + return out + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--models", default=",".join(RUN_MODELS)) + ap.add_argument("--lengths", default="50,150") + ap.add_argument("--conditions", default=",".join(CONDITIONS)) + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--effort", default=None) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--report", action="store_true") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + out_dir = RESULTS / f"{EXP}-{SLUG}" + out_dir.mkdir(parents=True, exist_ok=True) + + models = [m.strip() for m in args.models.split(",") if m.strip()] + lengths = [int(x) for x in args.lengths.split(",") if x.strip()] + conditions = [c.strip() for c in args.conditions.split(",") if c.strip()] + + keys = {} + if not args.report: + keys["openai"] = load_env_key("OPENAI_API_KEY", args.env) + keys["openrouter"] = load_env_key("OPENROUTER_API_KEY", args.env) + + all_paras = squad.load_paragraphs(CACHE) + budget = capacity(FONT, args.size)[2] + tasks = [] + flows: dict[int, str] = {} + for length in lengths: + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + flows[length] = flow + ctx = {"args": args, "flow": flow, "paras": paras, "offsets": offsets, "keys": keys, "length": length} + for model in models: + for cond in conditions: + for start in range(0, len(flow), budget): + tasks.append((model, cond, start, min(start + budget, len(flow)), ctx)) + print(f"{EXP}: {len(models)} models x {len(lengths)} lengths x {len(conditions)} conditions = {len(tasks)} chunk tasks") + + records: list[dict] = [] + with ThreadPoolExecutor(args.workers) as pool: + futures = [pool.submit(run_chunk, m, c, s, e, ctx) for m, c, s, e, ctx in tasks] + for k, fut in enumerate(futures): + records.extend(fut.result()) + print(f" {k + 1}/{len(tasks)} tasks", flush=True) + + with (out_dir / "records.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + + # overall cells (same shape as final.py matrix) + cells = [] + for model in models: + for length in lengths: + for cond in conditions: + sub = [r for r in records if r["model"] == model and r["length"] == length and r["cond"] == cond] + if sub: + cells.append({"model": model, "length": length, "condition": cond, **aggregate(sub, *MODELS[model])}) + with (out_dir / "matrix.csv").open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=list(cells[0].keys())) + writer.writeheader() + writer.writerows(cells) + + num_cells = numeric_subset_cells(records, models, lengths, conditions) + + # baseline numeric subset stats (per model/length, from the SAME matched questions) + base_overall = {} + for model in models: + base = load_baseline(model, lengths) + for length in lengths: + sub = [r for r in base if r["length"] == length] + base_overall[(model, length)] = f1_stats(sub) + + (out_dir / "summary.json").write_text( + json.dumps( + { + "args": vars(args), + "cells": cells, + "numeric_subset": num_cells, + "baseline_overall": {f"{m}|{l}": v for (m, l), v in base_overall.items()}, + }, + indent=1, + ) + ) + + samples = [] + for cond in conditions: + base, hard = parse_cond(cond) + samples.append(str(save_sample(flows[lengths[0]], args.size, base, hard, out_dir))) + + print("\n== overall ==") + print(f"{'model':<24}{'len':>5}{'condition':<28}{'n':>4}{'EM':>7}{'F1':>7}{'se':>7}{'cost$':>8}{'base F1':>9}{'d':>7}") + for c in cells: + b = base_overall[(c["model"], c["length"])] + print( + f"{c['model']:<24}{c['length']:>5}{c['condition']:<28}{c['n']:>4}{c['em']:>7.3f}{c['f1']:>7.3f}" + f"{c['f1_se']:>7.3f}{c['cost_usd']:>8.3f}{b['f1']:>9.3f}{c['f1'] - b['f1']:>+7.3f}" + ) + print("\n== numeric-gold subset (matched questions vs img-6x10-sent baseline) ==") + print(f"{'model':<24}{'len':>5}{'condition':<28}{'n':>4}{'F1':>7}{'se':>7}{'base F1':>9}{'base se':>8}{'d':>7}") + for c in num_cells: + print( + f"{c['model']:<24}{c['length']:>5}{c['condition']:<28}{c['n']:>4}{c['f1']:>7.3f}{c['f1_se']:>7.3f}" + f"{c['baseline_f1']:>9.3f}{c['baseline_se']:>8.3f}{c['f1'] - c['baseline_f1']:>+7.3f}" + ) + print(f"\nresults -> {out_dir}/ samples: {', '.join(samples)}") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp04_layout.py b/packages/snapcompact/research/exp04_layout.py new file mode 100644 index 000000000..0d8a1d861 --- /dev/null +++ b/packages/snapcompact/research/exp04_layout.py @@ -0,0 +1,374 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp04: document-style (in-distribution) layout vs row-major grid. + +Hypothesis: VLMs are pretrained on documents, not 261-col row-major char +grids. A two-column newspaper page (word-wrap, paragraph breaks, headings) +costs capacity (gutter, blank lines, ragged right) but may read better per +token. Conditions: img-6x10-doc (near-black, document-plain) and +img-6x10-doc-sent (sentence-hue glyphs). Same 6x10 font/page size as the +img-6x10-sent baseline; chunking is paragraph-aligned page packing. +""" + +import argparse +import json +import sys +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +from PIL import Image + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import _DARK, capacity, parse_bdf, ensure_font # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, FONTS, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 +from final import cached # noqa: E402 + +MODELS = {"gpt-5.5": (2.0, 16.0), "google/gemini-3.5-flash": (0.6, 4.0)} +LENGTHS = (50, 150) +CONDITIONS = ("img-6x10-doc", "img-6x10-doc-sent") +FONT = FONTS["6x10"] +GUTTER = 3 # char cells between columns +_WHITE = (255, 255, 255) +_BLACK = (0, 0, 0) +_INK = (24, 24, 24) # near-black body text, like a printed page + + +# --- document layout ------------------------------------------------------- + + +def wrap(text: str, width: int) -> list[str]: + """Greedy word-wrap, no mid-word breaks (hard split only for width+ words).""" + lines: list[str] = [] + cur = "" + for word in text.split(): + while len(word) > width: # pathological; never hit on SQuAD prose + if cur: + lines.append(cur) + cur = "" + lines.append(word[:width]) + word = word[width:] + if not cur: + cur = word + elif len(cur) + 1 + len(word) <= width: + cur += " " + word + else: + lines.append(cur) + cur = word + if cur: + lines.append(cur) + return lines + + +def layout_page(paras: list[dict], col_w: int) -> list[dict]: + """Typeset paragraphs into lines: [{kind: heading|body|blank, text}]. + + Article title changes become headings (heading is repeated at the top of + a page even when the article continues from the previous page, since each + page is read in isolation). Paragraphs are separated by one blank line. + """ + lines: list[dict] = [] + prev_title = None + for p in paras: + if p["title"] != prev_title: + if lines: + lines.append({"kind": "blank", "text": ""}) + for hl in wrap(p["title"].replace("_", " ").upper(), col_w): + lines.append({"kind": "heading", "text": hl}) + prev_title = p["title"] + elif lines: + lines.append({"kind": "blank", "text": ""}) + for bl in wrap(p["ctx"], col_w): + lines.append({"kind": "body", "text": bl}) + return lines + + +def pack_pages(paras: list[dict], col_w: int, max_lines: int) -> list[tuple[int, int]]: + """Greedy paragraph-aligned packing: [(i, j)] para ranges, one per page.""" + pages = [] + i = 0 + while i < len(paras): + j = i + 1 + while j < len(paras) and len(layout_page(paras[i : j + 1], col_w)) <= max_lines: + j += 1 + pages.append((i, j)) + i = j + return pages + + +# --- renderer (glyph loop copied from bdf.render, two-column layout) ------- + + +def _sentence_colors(lines: list[dict]) -> list[list[tuple[int, int, int]]]: + """Per-line per-char glyph color cycling hue per sentence across the page.""" + joined = "\n".join(ln["text"] for ln in lines) + idx, out_idx = 0, [] + for i, ch in enumerate(joined): + out_idx.append(idx) + if ch in ".!?" and i + 1 < len(joined) and joined[i + 1] in " \n": + idx += 1 + colors, pos = [], 0 + for ln in lines: + n = len(ln["text"]) + colors.append([_DARK[out_idx[pos + k] % 6] for k in range(n)]) + pos += n + 1 # the joining newline + return colors + + +def render_doc(lines: list[dict], size: int, variant: str, cache: Path) -> Image.Image: + """Two-column page: left column rows top-to-bottom, then right column.""" + glyphs, font_ascent = parse_bdf(ensure_font(FONT, cache)) + ascent = FONT.ascent if FONT.ascent is not None else font_ascent + cols, rows, _ = capacity(FONT, size) + col_w = (cols - GUTTER) // 2 + sent_colors = _sentence_colors(lines) if variant == "doc-sent" else None + img = Image.new("RGB", (size, size), _WHITE) + px = img.load() + for li, ln in enumerate(lines): + column, row = divmod(li, rows) + if column > 1: + break # overflow guard; pack_pages should prevent this + x_origin = column * (col_w + GUTTER) * FONT.adv + y0 = row * FONT.pitch + for ci, ch in enumerate(ln["text"]): + glyph = glyphs.get(ord(ch)) + if glyph is None: + continue + if ln["kind"] == "heading": + fg = _BLACK + elif sent_colors is not None: + fg = sent_colors[li][ci] + else: + fg = _INK + w, h, xoff, yoff = glyph["bbx"] + top = y0 + ascent - h - yoff + shift = 0x80 if w <= 8 else 0x8000 + strikes = (0, 1) if ln["kind"] == "heading" else (0,) + for dx in strikes: + for r, bits in enumerate(glyph["rows"]): + y = top + r + if not 0 <= y < size: + continue + for b in range(w): + if bits & (shift >> b): + x = x_origin + ci * FONT.adv + xoff + b + dx + if 0 <= x < size: + px[x, y] = fg + return img + + +# --- runner ----------------------------------------------------------------- + + +def run_page(model: str, cond: str, page: tuple[int, int], ctx: dict) -> list[dict]: + args, paras, offsets, keys = ctx["args"], ctx["paras"], ctx["offsets"], ctx["keys"] + i, j = page + start = offsets[i] + end = offsets[j - 1] + len(paras[j - 1]["ctx"]) + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + variant = cond.removeprefix("img-6x10-") + lines = ctx["lines"][page] + page_key = sha8(cond, json.dumps([(p["title"], p["ctx"]) for p in paras[i:j]]), str(args.size)) + png = CACHE / f"exp04-{variant}-{page_key}.png" + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(".tmp.png") + render_doc(lines, args.size, variant, CACHE).save(tmp) + tmp.replace(png) + cols, rows, _ = capacity(FONT, args.size) + col_w = (cols - GUTTER) // 2 + q_block = "\n".join(f"{k + 1}. {q['q']}" for k, q in enumerate(questions)) + messages = [ + { + "role": "user", + "content": [ + {"text": load_prompt("exp04-qa-image.md").format(col_w=col_w, rows=rows)}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + qa = cached( + model, "exp04-qa", {"messages": messages, "effort": args.effort}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete(keys, model, messages, max_tokens=args.max_tokens, effort=args.effort), + ) + ), + args.fresh, + ) + answers = squad.parse_numbered(qa["text"], len(questions)) + records = [] + for q, a in zip(questions, answers): + records.append( + { + "model": model, + "length": ctx["length"], + "cond": cond, + "chunk": start, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + ) + records[0]["usage"] = [{"phase": "qa", **qa["usage"]}] + return records + + +def aggregate(records: list[dict], price_in: float, price_out: float) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean_f1 = sum(f1s) / n + se = (sum((x - mean_f1) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + us = [u for r in records if "usage" in r for u in r["usage"]] + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + cost_in = (tok["in"] + 1.25 * tok["cache_w"] + 0.1 * tok["cache_r"]) / 1e6 * price_in + cost_out = tok["out"] / 1e6 * price_out + return { + "n": n, + "em": sum(r["em"] for r in records) / n, + "f1": mean_f1, + "f1_se": se, + "abstained": sum(r["abstained"] for r in records), + **{f"tok_{k}": v for k, v in tok.items()}, + "cost_in_usd": round(cost_in, 4), + "cost_out_usd": round(cost_out, 4), + "cost_usd": round(cost_in + cost_out, 4), + } + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--models", default=",".join(MODELS)) + ap.add_argument("--lengths", default=",".join(map(str, LENGTHS))) + ap.add_argument("--conditions", default=",".join(CONDITIONS)) + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--effort", default=None) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--render-only", action="store_true", help="render pages + capacity stats, no API") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + out_dir = RESULTS / "exp04-layout" + out_dir.mkdir(parents=True, exist_ok=True) + + models = [m.strip() for m in args.models.split(",") if m.strip()] + lengths = [int(x) for x in args.lengths.split(",") if x.strip()] + conditions = [c.strip() for c in args.conditions.split(",") if c.strip()] + + keys = {} + if not args.render_only: + keys["openai"] = load_env_key("OPENAI_API_KEY", args.env) + keys["openrouter"] = load_env_key("OPENROUTER_API_KEY", args.env) + + cols, rows, grid_cap = capacity(FONT, args.size) + col_w = (cols - GUTTER) // 2 + max_lines = 2 * rows + + all_paras = squad.load_paragraphs(CACHE) + tasks = [] + capacity_stats = {} + for length in lengths: + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + pages = pack_pages(paras, col_w, max_lines) + page_lines = {pg: layout_page(paras[pg[0] : pg[1]], col_w) for pg in pages} + page_chars = [offsets[j - 1] + len(paras[j - 1]["ctx"]) - offsets[i] for i, j in pages] + capacity_stats[length] = { + "pages": len(pages), + "chars_per_page": page_chars, + "mean_chars_page": round(sum(page_chars) / len(pages)), + "grid_chars_page": grid_cap, + "corpus_chars": len(flow), + "grid_pages": -(-len(flow) // grid_cap), + } + ctx = { + "args": args, "paras": paras, "offsets": offsets, "keys": keys, + "length": length, "lines": page_lines, + } + for model in models: + for cond in conditions: + for pg in pages: + tasks.append((model, cond, pg, ctx)) + + print(f"layout: {cols} cols -> 2 x {col_w} + gutter {GUTTER}; {max_lines} line slots/page") + for length, st in capacity_stats.items(): + print( + f" len {length}: {st['pages']} doc pages (mean {st['mean_chars_page']} chars/page; " + f"grid {st['grid_chars_page']} chars/page -> {st['grid_pages']} pages)" + ) + if args.render_only: + for length in lengths: + paras = all_paras[:length] + pages = pack_pages(paras, col_w, max_lines) + for cond in conditions: + variant = cond.removeprefix("img-6x10-") + i, j = pages[0] + lines = layout_page(paras[i:j], col_w) + key = sha8(cond, json.dumps([(p["title"], p["ctx"]) for p in paras[i:j]]), str(args.size)) + png = CACHE / f"exp04-{variant}-{key}.png" + tmp = png.with_suffix(".tmp.png") + render_doc(lines, args.size, variant, CACHE).save(tmp) + tmp.replace(png) + print(f" sample: {png}") + return + + print(f"grid: {len(tasks)} page tasks") + records: list[dict] = [] + done = 0 + with ThreadPoolExecutor(args.workers) as pool: + futures = [pool.submit(run_page, m, c, pg, ctx) for m, c, pg, ctx in tasks] + for fut in futures: + records.extend(fut.result()) + done += 1 + print(f" {done}/{len(tasks)} pages", flush=True) + + with (out_dir / "records.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + + cells = [] + for model in models: + for length in lengths: + for cond in conditions: + sub = [r for r in records if r["model"] == model and r["length"] == length and r["cond"] == cond] + if not sub: + continue + cells.append({"model": model, "length": length, "condition": cond, **aggregate(sub, *MODELS[model])}) + (out_dir / "summary.json").write_text( + json.dumps({"args": vars(args), "capacity": capacity_stats, "cells": cells}, indent=1) + ) + import csv + + with (out_dir / "matrix.csv").open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=list(cells[0].keys())) + writer.writeheader() + writer.writerows(cells) + + for c in cells: + print( + f"{c['model']:<26} len {c['length']:<4} {c['condition']:<20} " + f"n={c['n']:<4} EM {c['em']:.3f} F1 {c['f1']:.3f} ±{c['f1_se']:.3f} ${c['cost_usd']:.3f}" + ) + print(f"\n-> {out_dir}/records.jsonl, matrix.csv, summary.json") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp05_anchors.py b/packages/snapcompact/research/exp05_anchors.py new file mode 100644 index 000000000..2cea6a75f --- /dev/null +++ b/packages/snapcompact/research/exp05_anchors.py @@ -0,0 +1,344 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp05: margin row-number ruler ("anchors") on img-6x10-sent. + +Hypothesis: a hex-editor-style row ruler makes the image addressable, so the +model can localize answers instead of transcribing the whole bitmap — cutting +the reasoning/output "read tax" and possibly improving F1. + +Conditions (both chunked by the ruler content capacity, 39936 chars, so the +two conditions share identical question sets): + img-6x10-sent-ruler 5-col blue row ruler + anti-transcription prompt + asking for `answer | row≈N` + img-6x10-sent-noruler baseline render (261 cols) + anti-transcription + prompt only (separates ruler effect from prompt effect) + +Diagnostic: claimed row vs true row of the gold answer (char offset // cols). + +Usage: uv run exp05_anchors.py [--report] [--lengths 50,150] [--models ...] +""" + +import argparse +import csv +import json +import re +import sys +import uuid +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +from PIL import Image + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import _BLACK, _DARK, _WHITE, FontCfg, _row_palette, _sentence_indices, ensure_font, parse_bdf, render # noqa: E402 +from final import MODELS, cached # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, FONTS, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 + +FONT = FONTS["6x10"] +SIZE = 1568 +MARGIN_COLS = 5 # 4 digit cells + 1 gap cell +RULER_STEP = 5 +RULER_FG = (90, 120, 215) # medium blue: legible, never competes with content +COND_RULER = "img-6x10-sent-ruler" +COND_CTL = "img-6x10-sent-noruler" + +# ctx.md baseline (img-6x10-sent, seed 42, qpc 30): (f1, se, cost) +BASELINE = { + ("gpt-5.5", 50): (0.850, 0.051, 0.068), + ("gpt-5.5", 150): (0.822, 0.029, 0.245), + ("google/gemini-3.5-flash", 50): (0.984, 0.012, 0.018), + ("google/gemini-3.5-flash", 150): (0.805, 0.035, 0.097), +} + + +def ruler_capacity(cfg: FontCfg, size: int = SIZE) -> tuple[int, int, int]: + """(content_cols, rows, chars) once MARGIN_COLS are reserved for the ruler.""" + cols, rows = size // cfg.adv - MARGIN_COLS, size // cfg.pitch + return cols, rows, cols * rows + + +def render_ruler(text: str, cfg: FontCfg, cache: Path, size: int = SIZE, variant: str = "sent") -> Image.Image: + """bdf.render() with a left row-number ruler every RULER_STEP rows. + + Content glyphs are shifted right by MARGIN_COLS cells; row indices are + drawn 0-based, right-aligned in the first MARGIN_COLS-1 cells, in RULER_FG. + """ + glyphs, font_ascent = parse_bdf(ensure_font(cfg, cache)) + ascent = cfg.ascent if cfg.ascent is not None else font_ascent + cols, rows, cap = ruler_capacity(cfg, size) + text = text[:cap] + sent_idx = _sentence_indices(text) if variant in ("sent", "dark-sent", "sent-dim") else None + sent_palette = _DARK + img = Image.new("RGB", (size, size), _BLACK if variant in ("dark", "dark-sent") else _WHITE) + px = img.load() + + def draw_glyph(ch: str, cell_col: int, y0: int, fg: tuple[int, int, int]) -> None: + glyph = glyphs.get(ord(ch)) + if glyph is None: + return + w, h, xoff, yoff = glyph["bbx"] + top = y0 + ascent - h - yoff + shift = 0x80 if w <= 8 else 0x8000 + for r, bits in enumerate(glyph["rows"]): + y = top + r + if not 0 <= y < size: + continue + for b in range(w): + if bits & (shift >> b): + x = cell_col * cfg.adv + xoff + b + if 0 <= x < size: + px[x, y] = fg + + for row in range(rows): + bg, fg = _row_palette(variant, row) + y0 = row * cfg.pitch + for y in range(y0, min(y0 + cfg.pitch, size)): + for x in range(size): + px[x, y] = bg + if row % RULER_STEP == 0: + label = str(row) + for j, ch in enumerate(label): + draw_glyph(ch, MARGIN_COLS - 1 - len(label) + j, y0, RULER_FG) + for col in range(cols): + i = row * cols + col + if i >= len(text): + break + if sent_idx is not None: + fg = sent_palette[sent_idx[i] % 6] + draw_glyph(text[i], MARGIN_COLS + col, y0, fg) + return img + + +_ROW_CLAIM = re.compile(r"\|\s*rows?\s*[≈~=:]*\s*(\d+)", re.IGNORECASE) + + +def parse_answers_rows(text: str, n: int) -> tuple[list[str], list[int | None]]: + """Numbered answers, with optional `| row≈N` suffixes stripped into rows.""" + answers, rows = [""] * n, [None] * n + for line in text.splitlines(): + m = re.match(r"\s*(\d+)[.):]\s*(.*\S)?\s*$", line) + if not (m and m.group(2)): + continue + idx = int(m.group(1)) - 1 + if not (0 <= idx < n) or answers[idx]: + continue + body = m.group(2).strip() + rm = _ROW_CLAIM.search(body) + if rm: + rows[idx] = int(rm.group(1)) + body = body[: rm.start()].strip() + answers[idx] = body + return answers, rows + + +def gold_row(chunk_text: str, q: dict, chunk_len: int, content_cols: int) -> int | None: + """True row of the gold answer in the rendered grid (0-based).""" + approx = max(0, int(q["pos_rel"] * chunk_len) - 10) + for g in q["golds"]: + i = chunk_text.find(g, approx) + if i < 0: + i = chunk_text.find(g) + if i < 0: + i = chunk_text.lower().find(g.lower()) + if i >= 0: + return i // content_cols + return None + + +def _ensure_png(png: Path, make) -> None: + """Render-once with atomic publish; unique tmp avoids cross-thread races.""" + if png.exists() and png.stat().st_size > 0: + return + tmp = png.with_suffix(f".{uuid.uuid4().hex}.tmp.png") + make().save(tmp) + tmp.replace(png) + + +def run_cell_chunk(model: str, cond: str, start: int, end: int, ctx: dict) -> list[dict]: + args, flow, paras, offsets, keys = ctx["args"], ctx["flow"], ctx["paras"], ctx["offsets"], ctx["keys"] + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + chunk_text = flow[start:end] + q_block = "\n".join(f"{i + 1}. {q['q']}" for i, q in enumerate(questions)) + + if cond == COND_RULER: + cols, rows, _ = ruler_capacity(FONT, args.size) + png = CACHE / f"exp05-ruler-{sha8(chunk_text, str(args.size))}.png" + _ensure_png(png, lambda: render_ruler(chunk_text, FONT, CACHE, args.size, "sent")) + last_label = (rows - 1) // RULER_STEP * RULER_STEP + prompt = load_prompt("exp05-qa-image.md").format(cols=cols, rows=rows, last_label=last_label) + else: # control: baseline render, anti-transcription prompt only + cols, rows = args.size // FONT.adv, args.size // FONT.pitch + png = CACHE / f"exp05-ctl-{sha8(chunk_text, str(args.size))}.png" + _ensure_png(png, lambda: render(chunk_text, FONT, CACHE, args.size, "sent")) + prompt = load_prompt("exp05-qa-image-ctl.md").format(cols=cols, rows=rows) + + messages = [ + { + "role": "user", + "content": [{"text": prompt}, {"image_path": png}, {"text": q_block}], + } + ] + qa = cached( + model, f"exp05-qa-{cond}", {"messages": messages, "effort": args.effort}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete(keys, model, messages, max_tokens=args.max_tokens, effort=args.effort), + ) + ), + args.fresh, + ) + answers, claimed_rows = parse_answers_rows(qa["text"], len(questions)) + records = [] + for q, a, crow in zip(questions, answers, claimed_rows): + trow = gold_row(chunk_text, q, end - start, cols) if cond == COND_RULER else None + records.append( + { + "model": model, + "length": ctx["length"], + "cond": cond, + "chunk": start, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + "claimed_row": crow, + "true_row": trow, + } + ) + records[0]["usage"] = [{"phase": "qa", **qa["usage"]}] + return records + + +def aggregate(records: list[dict], price_in: float, price_out: float) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean_f1 = sum(f1s) / n + se = (sum((x - mean_f1) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + us = [u for r in records if "usage" in r for u in r["usage"]] + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + cost_in = (tok["in"] + 1.25 * tok["cache_w"] + 0.1 * tok["cache_r"]) / 1e6 * price_in + cost_out = tok["out"] / 1e6 * price_out + loc = [(r["claimed_row"], r["true_row"]) for r in records if r["claimed_row"] is not None and r["true_row"] is not None] + row_stats = {} + if loc: + errs = [abs(c - t) for c, t in loc] + row_stats = { + "row_n": len(loc), + "row_claimed_frac": round(sum(r["claimed_row"] is not None for r in records) / n, 3), + "row_mae": round(sum(errs) / len(errs), 2), + "row_within2": round(sum(e <= 2 for e in errs) / len(errs), 3), + "row_within5": round(sum(e <= 5 for e in errs) / len(errs), 3), + } + return { + "n": n, + "em": sum(r["em"] for r in records) / n, + "f1": mean_f1, + "f1_se": se, + "abstained": sum(r["abstained"] for r in records), + **{f"tok_{k}": v for k, v in tok.items()}, + "cost_in_usd": round(cost_in, 4), + "cost_out_usd": round(cost_out, 4), + "cost_usd": round(cost_in + cost_out, 4), + **row_stats, + } + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--models", default="gpt-5.5,google/gemini-3.5-flash") + ap.add_argument("--lengths", default="50,150") + ap.add_argument("--conditions", default=f"{COND_RULER},{COND_CTL}") + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=SIZE) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--effort", default=None) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--report", action="store_true") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + out_dir = RESULTS / "exp05-anchors" + out_dir.mkdir(parents=True, exist_ok=True) + + models = [m.strip() for m in args.models.split(",") if m.strip()] + lengths = [int(x) for x in args.lengths.split(",") if x.strip()] + conditions = [c.strip() for c in args.conditions.split(",") if c.strip()] + + keys = {} + if not args.report: + keys["openai"] = load_env_key("OPENAI_API_KEY", args.env) + keys["openrouter"] = load_env_key("OPENROUTER_API_KEY", args.env) + + budget = ruler_capacity(FONT, args.size)[2] # both conds: identical question sets + all_paras = squad.load_paragraphs(CACHE) + tasks = [] + for length in lengths: + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + ctx = {"args": args, "flow": flow, "paras": paras, "offsets": offsets, "keys": keys, "length": length} + for model in models: + for cond in conditions: + for start in range(0, len(flow), budget): + tasks.append((model, cond, start, min(start + budget, len(flow)), ctx)) + print(f"grid: {len(tasks)} chunk tasks, chunk budget {budget} chars") + + records: list[dict] = [] + with ThreadPoolExecutor(args.workers) as pool: + futures = [pool.submit(run_cell_chunk, m, c, s, e, ctx) for m, c, s, e, ctx in tasks] + for i, fut in enumerate(futures): + records.extend(fut.result()) + print(f" {i + 1}/{len(tasks)} tasks", flush=True) + + with (out_dir / "records.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + + cells = [] + for model in models: + for length in lengths: + for cond in conditions: + sub = [r for r in records if r["model"] == model and r["length"] == length and r["cond"] == cond] + if not sub: + continue + cell = {"model": model, "length": length, "condition": cond, **aggregate(sub, *MODELS[model])} + base = BASELINE.get((model, length)) + if base: + cell["base_f1"] = base[0] + cell["d_f1"] = round(cell["f1"] - base[0], 3) + cell["base_cost"] = base[2] + cell["d_cost"] = round(cell["cost_usd"] - base[2], 4) + cells.append(cell) + (out_dir / "summary.json").write_text(json.dumps({"args": vars(args), "cells": cells}, indent=1)) + fieldnames = sorted({k for c in cells for k in c}, key=lambda k: (k not in ("model", "length", "condition"), k)) + with (out_dir / "matrix.csv").open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=fieldnames) + writer.writeheader() + writer.writerows(cells) + + for c in cells: + print( + f"{c['model']:>26} L{c['length']:<4}{c['condition']:<24} n={c['n']:<4} em={c['em']:.3f} f1={c['f1']:.3f}±{c['f1_se']:.3f}" + f" cost=${c['cost_usd']:.3f} out={c['tok_out']} reas={c['tok_reasoning']}" + + (f" rowMAE={c['row_mae']} w5={c['row_within5']}" if "row_mae" in c else "") + ) + print(f"\n-> {out_dir}/records.jsonl, matrix.csv, summary.json") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp06_rolecolor.py b/packages/snapcompact/research/exp06_rolecolor.py new file mode 100644 index 000000000..e1c81be2c --- /dev/null +++ b/packages/snapcompact/research/exp06_rolecolor.py @@ -0,0 +1,417 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp06 rolecolor: hue = message role, as zero-char metadata in optical compaction. + +Synthetic transcript: each SQuAD passage gets a role (user/assistant/tool), +round-robin within shuffled triples (seed 42). Three conditions, same passages, +same questions: + + img-6x10-role plain text, glyph hue per role (blue=user, green=assistant, + red=tool) -- metadata at zero character cost + img-6x10-tagbw black-on-white, inline "[user] "/"[asst] "/"[tool] " tag + before each passage -- text-equivalent control, ~7 chars/passage + img-6x10-nometa baseline sent variant, no role metadata -- provenance floor + (content F1 for this condition == img-6x10-sent baseline table) + +Two evals per chunk image: (a) standard SQuAD QA (content F1; not run for +nometa), (b) provenance QA: "which role's message contains the answer?", +scored as plain accuracy against the assigned role of the source passage. + +Chunking: greedy consecutive passages such that the TAGGED text fits the +6x10 capacity (40716 chars), so all three conditions share identical chunks +and question sets. Boundaries therefore shift slightly vs the baseline run +(which chunks the plain flow at exactly 40716); question sets overlap heavily +but are not char-identical to the baseline cells. + +Saturation-decay-by-recency was considered and deliberately omitted: no eval +question tests recency, and desaturating old passages risks destroying the +hue signal the provenance task measures. + +Run from the snapcompact dir: uv run exp06_rolecolor.py +""" + +import argparse +import colorsys +import csv +import json +import random +import sys +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +from PIL import Image + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import capacity, ensure_font, parse_bdf, render # noqa: E402 +from final import cached # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, FONTS, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 + +MODELS = {"gpt-5.5": (2.0, 16.0), "google/gemini-3.5-flash": (0.6, 4.0)} +LENGTHS = (50, 150) +CONDITIONS = ("img-6x10-role", "img-6x10-tagbw", "img-6x10-nometa") +FONT = FONTS["6x10"] + +ROLES = ("user", "assistant", "tool") +TAGS = {"user": "user", "assistant": "asst", "tool": "tool"} # all "[xxxx] " = 7 chars +ROLE_HUES = {"user": 0.62, "assistant": 0.33, "tool": 0.02} +ROLE_RGB = {r: tuple(int(c * 255) for c in colorsys.hls_to_rgb(h, 0.27, 0.90)) for r, h in ROLE_HUES.items()} +_WHITE = (255, 255, 255) + +ENCODING = { + "img-6x10-role": ( + "Glyph color encodes the author role: dark blue = user, dark green = assistant, dark red = tool. " + "A message boundary is where the glyph color changes." + ), + "img-6x10-tagbw": ( + "Each message is preceded by a bracketed role tag rendered in the text: [user], [asst], or [tool]." + ), + "img-6x10-nometa": ( + "The rendering does NOT visually indicate roles; glyph colors only cycle per sentence and carry " + "no role information. Use your best guess." + ), +} +QA_PROMPT = {"img-6x10-role": "exp06-qa-image.md", "img-6x10-tagbw": "exp06-qa-image-tag.md"} + + +def assign_roles(n: int, seed: int) -> list[str]: + """Round-robin role assignment: each consecutive triple of passages contains + every role exactly once, triple-internal order shuffled deterministically.""" + rng = random.Random(seed) + roles: list[str] = [] + while len(roles) < n: + triple = list(ROLES) + rng.shuffle(triple) + roles.extend(triple) + return roles[:n] + + +def build_chunks(paras: list[dict], budget: int) -> list[tuple[int, int]]: + """Greedy consecutive passage ranges [a, b) whose TAGGED rendering fits budget.""" + chunks, cur, cur_len = [], 0, 0 + for i, p in enumerate(paras): + add = 7 + len(p["ctx"]) + 1 # "[xxxx] " + ctx + " " + if cur_len + add > budget and i > cur: + chunks.append((cur, i)) + cur, cur_len = i, 0 + cur_len += add + chunks.append((cur, len(paras))) + return chunks + + +def sample_questions(paras: list[dict], offsets: list[int], start: int, end: int, n: int, seed: int) -> list[dict]: + """squad.sample_chunk_questions with the source passage index recorded (same rng sequence).""" + rng = random.Random(seed * 1_000_003 + start) + eligible = [i for i in range(len(offsets)) if offsets[i] >= start and offsets[i] + len(paras[i]["ctx"]) <= end] + if not eligible: + return [] + n = min(n, len(eligible)) + step = len(eligible) / n + picked = [] + for k in range(n): + pi = eligible[int(k * step)] + qa = rng.choice(paras[pi]["qas"]) + picked.append( + { + "q": " ".join(qa["question"].split()), + "golds": sorted({a["text"] for a in qa["answers"]}), + "pos_rel": (offsets[pi] - start) / (end - start), + "pi": pi, + } + ) + return picked + + +def render_role(text: str, colors: list[tuple[int, int, int]], size: int) -> Image.Image: + """bdf.render() copy, simplified: white background, per-character glyph color.""" + glyphs, font_ascent = parse_bdf(ensure_font(FONT, CACHE)) + ascent = FONT.ascent if FONT.ascent is not None else font_ascent + cols, rows, cap = capacity(FONT, size) + text = text[:cap] + img = Image.new("RGB", (size, size), _WHITE) + px = img.load() + for row in range(rows): + y0 = row * FONT.pitch + for col in range(cols): + i = row * cols + col + if i >= len(text): + break + glyph = glyphs.get(ord(text[i])) + if glyph is None: + continue + fg = colors[i] + w, h, xoff, yoff = glyph["bbx"] + top = y0 + ascent - h - yoff + shift = 0x80 if w <= 8 else 0x8000 + for r, bits in enumerate(glyph["rows"]): + y = top + r + if not 0 <= y < size: + continue + for b in range(w): + if bits & (shift >> b): + x = col * FONT.adv + xoff + b + if 0 <= x < size: + px[x, y] = fg + return img + + +def chunk_carriers(paras: list[dict], roles: list[str], a: int, b: int) -> dict: + """Plain text + per-char role colors, and tagged text, for passages [a, b).""" + plain_parts, colors, tagged_parts = [], [], [] + for i in range(a, b): + seg = paras[i]["ctx"] + " " + plain_parts.append(seg) + colors.extend([ROLE_RGB[roles[i]]] * len(seg)) + tagged_parts.append(f"[{TAGS[roles[i]]}] {seg}") + return {"plain": "".join(plain_parts), "colors": colors, "tagged": "".join(tagged_parts)} + + +def atomic_png(png: Path, make) -> Path: + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(".tmp.png") + make().save(tmp) + tmp.replace(png) + return png + + +def build_image(cond: str, car: dict, size: int) -> Path: + if cond == "img-6x10-role": + png = CACHE / f"exp06-role-{sha8(car['plain'], str(size))}.png" + return atomic_png(png, lambda: render_role(car["plain"], car["colors"], size)) + if cond == "img-6x10-tagbw": + png = CACHE / f"exp06-tagbw-{sha8(car['tagged'], str(size))}.png" + return atomic_png(png, lambda: render(car["tagged"], FONT, CACHE, size, "bw")) + png = CACHE / f"exp06-nometa-{sha8(car['plain'], str(size))}.png" + return atomic_png(png, lambda: render(car["plain"], FONT, CACHE, size, "sent")) + + +def norm_role(answer: str) -> str: + a = answer.lower().strip(" \t.[]()\"'`*") + if "assist" in a or a == "asst": + return "assistant" + if "user" in a or "human" in a: + return "user" + if "tool" in a or "function" in a: + return "tool" + return a + + +def run_cell(model: str, cond: str, length: int, ci: int, chunk: dict, args, keys) -> list[dict]: + """One (model, cond, chunk): content QA (role/tagbw only) + provenance QA.""" + questions, car = chunk["questions"], chunk["car"] + if not questions: + return [] + png = build_image(cond, car, args.size) + cols, rows, _ = capacity(FONT, args.size) + q_block = "\n".join(f"{i + 1}. {q['q']}" for i, q in enumerate(questions)) + usage_rows: list[tuple[str, dict]] = [] + answers = [""] * len(questions) + + if cond in QA_PROMPT: + messages = [ + { + "role": "user", + "content": [ + {"text": load_prompt(QA_PROMPT[cond]).format(cols=cols, rows=rows)}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + qa = cached( + model, "exp06-qa", {"cond": cond, "messages": messages}, + lambda: dict( + zip(("text", "usage", "stop"), llm_complete(keys, model, messages, max_tokens=args.max_tokens)) + ), + args.fresh, + ) + usage_rows.append(("qa", qa["usage"])) + answers = squad.parse_numbered(qa["text"], len(questions)) + + prov_messages = [ + { + "role": "user", + "content": [ + {"text": load_prompt("exp06-prov-image.md").format(cols=cols, rows=rows, encoding=ENCODING[cond])}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + prov = cached( + model, "exp06-prov", {"cond": cond, "messages": prov_messages}, + lambda: dict( + zip(("text", "usage", "stop"), llm_complete(keys, model, prov_messages, max_tokens=args.max_tokens)) + ), + args.fresh, + ) + usage_rows.append(("prov", prov["usage"])) + prov_answers = squad.parse_numbered(prov["text"], len(questions)) + + records = [] + for q, a, pa in zip(questions, answers, prov_answers): + gold_role = chunk["roles"][q["pi"]] + scored = cond in QA_PROMPT + records.append( + { + "model": model, + "length": length, + "cond": cond, + "chunk": ci, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]) if scored else None, + "f1": squad.f1(a, q["golds"]) if scored else None, + "abstained": "unreadable" in a.lower() if scored else None, + "prov_answer": pa, + "prov_gold": gold_role, + "prov_correct": float(norm_role(pa) == gold_role), + } + ) + records[0]["usage"] = [{"phase": p, **u} for p, u in usage_rows] + return records + + +def phase_cost(records: list[dict], phase: str, price_in: float, price_out: float) -> tuple[dict, float]: + us = [u for r in records if "usage" in r for u in r["usage"] if u["phase"] == phase] + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + cost = (tok["in"] + 1.25 * tok["cache_w"] + 0.1 * tok["cache_r"]) / 1e6 * price_in + tok["out"] / 1e6 * price_out + return tok, cost + + +def aggregate(records: list[dict], price_in: float, price_out: float) -> dict: + n = len(records) + f1s = [r["f1"] for r in records if r["f1"] is not None] + if f1s: + mean_f1 = sum(f1s) / len(f1s) + se = (sum((x - mean_f1) ** 2 for x in f1s) / (len(f1s) * (len(f1s) - 1))) ** 0.5 if len(f1s) > 1 else 0.0 + em = sum(r["em"] for r in records if r["em"] is not None) / len(f1s) + abstained = sum(r["abstained"] for r in records if r["abstained"] is not None) + else: + mean_f1 = se = em = None + abstained = None + pacc = sum(r["prov_correct"] for r in records) / n + pse = (pacc * (1 - pacc) / n) ** 0.5 + qa_tok, qa_cost = phase_cost(records, "qa", price_in, price_out) + _, prov_cost = phase_cost(records, "prov", price_in, price_out) + return { + "n": n, + "em": em, + "f1": mean_f1, + "f1_se": se, + "abstained": abstained, + "prov_acc": round(pacc, 4), + "prov_se": round(pse, 4), + **{f"tok_{k}": v for k, v in qa_tok.items()}, + "qa_cost_usd": round(qa_cost, 4), + "prov_cost_usd": round(prov_cost, 4), + "cost_usd": round(qa_cost + prov_cost, 4), + } + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--models", default=",".join(MODELS)) + ap.add_argument("--lengths", default=",".join(map(str, LENGTHS))) + ap.add_argument("--conditions", default=",".join(CONDITIONS)) + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + out_dir = RESULTS / "exp06-rolecolor" + out_dir.mkdir(parents=True, exist_ok=True) + + models = [m.strip() for m in args.models.split(",") if m.strip()] + lengths = [int(x) for x in args.lengths.split(",") if x.strip()] + conditions = [c.strip() for c in args.conditions.split(",") if c.strip()] + keys = { + "openai": load_env_key("OPENAI_API_KEY", args.env), + "openrouter": load_env_key("OPENROUTER_API_KEY", args.env), + } + + budget = capacity(FONT, args.size)[2] + all_paras = squad.load_paragraphs(CACHE) + tasks = [] + for length in lengths: + paras = all_paras[:length] + roles = assign_roles(length, args.seed) + flow, offsets = squad.build_flow(paras) + for ci, (a, b) in enumerate(build_chunks(paras, budget)): + start, end = offsets[a], offsets[b - 1] + len(paras[b - 1]["ctx"]) + chunk = { + "questions": sample_questions(paras, offsets, start, end, args.qpc, args.seed), + "car": chunk_carriers(paras, roles, a, b), + "roles": roles, + } + for model in models: + for cond in conditions: + tasks.append((model, cond, length, ci, chunk)) + print(f"grid: {len(models)} models x {len(lengths)} lengths x {len(conditions)} conditions = {len(tasks)} cells") + + records: list[dict] = [] + failed = 0 + with ThreadPoolExecutor(args.workers) as pool: + futures = [pool.submit(run_cell, m, c, ln, ci, ch, args, keys) for m, c, ln, ci, ch in tasks] + for done, (fut, t) in enumerate(zip(futures, tasks), 1): + try: + records.extend(fut.result()) + except Exception as err: # noqa: BLE001 -- partial results still get written; rerun resumes from cache + failed += 1 + print(f" FAIL {t[0]} {t[1]} len={t[2]} chunk={t[3]}: {type(err).__name__}: {err}", flush=True) + print(f" {done}/{len(tasks)} cells", flush=True) + + with (out_dir / "records.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + + cells = [] + for model in models: + for length in lengths: + for cond in conditions: + sub = [r for r in records if r["model"] == model and r["length"] == length and r["cond"] == cond] + if sub: + cells.append({"model": model, "length": length, "condition": cond, **aggregate(sub, *MODELS[model])}) + (out_dir / "summary.json").write_text(json.dumps({"args": vars(args), "cells": cells}, indent=1)) + with (out_dir / "matrix.csv").open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=list(cells[0].keys())) + writer.writeheader() + writer.writerows(cells) + + for length in lengths: + print(f"\n== {length} passages == (content F1 / prov acc / $total)") + hdr = f"{'condition':<18}" + "".join(f"{m:>26}" for m in models) + print(hdr + "\n" + "-" * len(hdr)) + for cond in conditions: + row = f"{cond:<18}" + for model in models: + cell = next( + (c for c in cells if c["model"] == model and c["length"] == length and c["condition"] == cond), + None, + ) + if cell: + f1 = f"{cell['f1']:.3f}" if cell["f1"] is not None else " - " + row += f"{f1:>10} {cell['prov_acc']:>6.3f} {cell['cost_usd']:>8.3f}" + else: + row += f"{'-':>26}" + print(row) + print(f"\ndataset -> {out_dir}/records.jsonl, matrix.csv, summary.json") + if failed: + raise SystemExit(f"{failed} cells failed -- rerun to resume from cache") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp07_readtax.py b/packages/snapcompact/research/exp07_readtax.py new file mode 100644 index 000000000..1a69432fb --- /dev/null +++ b/packages/snapcompact/research/exp07_readtax.py @@ -0,0 +1,273 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp07 read tax: kill reasoning/output-token inflation on image conditions. + +Image-condition cost is dominated by reasoning tokens — models transcribe the +whole bitmap in CoT before answering. Carrier is the baseline winner +(img-6x10-sent, 1568px) unchanged; only the QA protocol varies: + + baseline - exact baseline protocol, re-measured for a latency + reference (baseline matrices have no wall-clock data) + effort-low - qa-image.md, effort="low" + effort-minimal - qa-image.md, effort="minimal" (gpt-5.5 only, if accepted) + no-transcribe - prompts/exp07-qa-image.md: explicit "do not transcribe, + locate the region per question, read only that region" + locate-then-answer - two turns: (1) effort=low, output only a row-band guess + per question; (2) same conversation + "read only those + bands, answer". Combined cost/latency tracked. + +Wall-clock latency is measured around every llm_complete and stored inside the +response cache entry, so resumes keep real timings. +""" + +import argparse +import csv +import json +import statistics +import sys +import time +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import capacity, render # noqa: E402 +from final import cached # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, FONTS, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 + +MODELS = {"gpt-5.5": (2.0, 16.0), "google/gemini-3.5-flash": (0.6, 4.0)} +LENGTHS = (50, 150) +CONDITIONS = ( + "baseline", "effort-low", "effort-minimal", "no-transcribe", + "locate-then-answer", "locate-low", "locate-none", +) +FONT, VARIANT = "6x10", "sent" + + +def ensure_png(chunk_text: str, size: int) -> Path: + """Reuse the baseline cache PNG when present (identical render), else write exp07-prefixed.""" + h = sha8(chunk_text, str(size)) + base = CACHE / f"img-{FONT}-{VARIANT}-{h}.png" + if base.exists() and base.stat().st_size > 0: + return base + png = CACHE / f"exp07-img-{FONT}-{VARIANT}-{h}.png" + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(".tmp.png") + render(chunk_text, FONTS[FONT], CACHE, size, VARIANT).save(tmp) + tmp.replace(png) + return png + + +def timed_call(keys: dict, model: str, messages: list[dict], max_tokens: int, effort: str | None) -> dict: + t0 = time.monotonic() + text, usage, stop = llm_complete(keys, model, messages, max_tokens=max_tokens, effort=effort) + return {"text": text, "usage": usage, "stop": stop, "latency_s": round(time.monotonic() - t0, 2)} + + +def probe_min_effort(keys: dict) -> str | None: + """Lowest reasoning effort gpt-5.5 accepts: try "minimal", fall back to "none".""" + for effort in ("minimal", "none"): + try: + llm_complete( + keys, "gpt-5.5", + [{"role": "user", "content": [{"text": "Reply with the single word OK."}]}], + max_tokens=64, effort=effort, + ) + return effort + except SystemExit as err: + print(f"effort={effort} rejected by gpt-5.5: {str(err)[:160]}") + return None + + +def run_cell_chunk(model: str, cond: str, start: int, end: int, ctx: dict) -> list[dict]: + args, flow, paras, offsets, keys = ctx["args"], ctx["flow"], ctx["paras"], ctx["offsets"], ctx["keys"] + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + chunk_text = flow[start:end] + q_block = "\n".join(f"{i + 1}. {q['q']}" for i, q in enumerate(questions)) + png = ensure_png(chunk_text, args.size) + cols, rows, _ = capacity(FONTS[FONT], args.size) + usage_rows: list[tuple[str, dict]] = [] + + def qa_messages(prompt_file: str) -> list[dict]: + return [ + { + "role": "user", + "content": [ + {"text": load_prompt(prompt_file).format(cols=cols, rows=rows)}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + + if cond.startswith("locate"): + turn2_effort = {"locate-then-answer": None, "locate-low": "low", "locate-none": "none"}[cond] + locate_msgs = qa_messages("exp07-locate.md") + locate = cached( + model, "exp07-locate", {"messages": locate_msgs, "effort": "low"}, + lambda: timed_call(keys, model, locate_msgs, args.max_tokens, "low"), + args.fresh, + ) + usage_rows.append(("locate", {**locate["usage"], "latency_s": locate.get("latency_s", 0)})) + answer_msgs = locate_msgs + [ + {"role": "assistant", "content": [{"text": locate["text"]}]}, + {"role": "user", "content": [{"text": load_prompt("exp07-answer-bands.md")}]}, + ] + qa = cached( + model, "exp07-qa", {"cond": cond, "messages": answer_msgs, "effort": turn2_effort}, + lambda: timed_call(keys, model, answer_msgs, args.max_tokens, turn2_effort), + args.fresh, + ) + else: + prompt_file = "exp07-qa-image.md" if cond == "no-transcribe" else "qa-image.md" + effort = cond.removeprefix("effort-") if cond.startswith("effort-") else None + messages = qa_messages(prompt_file) + qa = cached( + model, "exp07-qa", {"cond": cond, "messages": messages, "effort": effort}, + lambda: timed_call(keys, model, messages, args.max_tokens, effort), + args.fresh, + ) + usage_rows.append(("qa", {**qa["usage"], "latency_s": qa.get("latency_s", 0)})) + + answers = squad.parse_numbered(qa["text"], len(questions)) + records = [] + for q, a in zip(questions, answers): + records.append( + { + "model": model, + "length": ctx["length"], + "cond": cond, + "chunk": start, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + ) + records[0]["usage"] = [{"phase": p, **u} for p, u in usage_rows] + return records + + +def aggregate(records: list[dict], price_in: float, price_out: float) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean_f1 = sum(f1s) / n + se = (sum((x - mean_f1) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + us = [u for r in records if "usage" in r for u in r["usage"]] + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + cost_in = (tok["in"] + 1.25 * tok["cache_w"] + 0.1 * tok["cache_r"]) / 1e6 * price_in + cost_out = tok["out"] / 1e6 * price_out + # Per-chunk latency = sum over phases (locate + qa for the two-turn protocol). + chunk_lat = [sum(u.get("latency_s", 0) for u in r["usage"]) for r in records if "usage" in r] + return { + "n": n, + "em": sum(r["em"] for r in records) / n, + "f1": mean_f1, + "f1_se": se, + "abstained": sum(r["abstained"] for r in records), + **{f"tok_{k}": v for k, v in tok.items()}, + "latency_p50_s": round(statistics.median(chunk_lat), 1) if chunk_lat else 0.0, + "latency_max_s": round(max(chunk_lat), 1) if chunk_lat else 0.0, + "cost_in_usd": round(cost_in, 4), + "cost_out_usd": round(cost_out, 4), + "cost_usd": round(cost_in + cost_out, 4), + } + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--models", default=",".join(MODELS)) + ap.add_argument("--lengths", default=",".join(map(str, LENGTHS))) + ap.add_argument("--conditions", default=",".join(CONDITIONS)) + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--env", default="~/.env") + ap.add_argument("--out", default="exp07-readtax") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + out_dir = RESULTS / args.out + out_dir.mkdir(parents=True, exist_ok=True) + + models = [m.strip() for m in args.models.split(",") if m.strip()] + lengths = [int(x) for x in args.lengths.split(",") if x.strip()] + conditions = [c.strip() for c in args.conditions.split(",") if c.strip()] + + keys = { + "openai": load_env_key("OPENAI_API_KEY", args.env), + "openrouter": load_env_key("OPENROUTER_API_KEY", args.env), + } + min_effort = probe_min_effort(keys) if "effort-minimal" in conditions and "gpt-5.5" in models else None + print(f"lowest gpt-5.5 effort: {min_effort or 'unavailable -> condition skipped'}") + if "effort-minimal" in conditions: + conditions = [f"effort-{min_effort}" if c == "effort-minimal" and min_effort else c for c in conditions] + conditions = [c for c in conditions if c != "effort-minimal"] + + budget = capacity(FONTS[FONT], args.size)[2] + all_paras = squad.load_paragraphs(CACHE) + tasks = [] + for length in lengths: + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + ctx = {"args": args, "flow": flow, "paras": paras, "offsets": offsets, "keys": keys, "length": length} + for model in models: + for cond in conditions: + if cond in (f"effort-{min_effort}", "locate-none") and model != "gpt-5.5": + continue + if cond == "locate-none" and min_effort != "none": + continue + for start in range(0, len(flow), budget): + tasks.append((model, cond, start, min(start + budget, len(flow)), ctx)) + print(f"grid: {len(tasks)} chunk tasks") + + records: list[dict] = [] + with ThreadPoolExecutor(args.workers) as pool: + futures = [pool.submit(run_cell_chunk, m, c, s, e, ctx) for m, c, s, e, ctx in tasks] + for i, fut in enumerate(futures, 1): + records.extend(fut.result()) + print(f" {i}/{len(tasks)} tasks", flush=True) + + with (out_dir / "records.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + + cells = [] + for model in models: + for length in lengths: + for cond in conditions: + sub = [r for r in records if r["model"] == model and r["length"] == length and r["cond"] == cond] + if not sub: + continue + cells.append({"model": model, "length": length, "condition": cond, **aggregate(sub, *MODELS[model])}) + (out_dir / "summary.json").write_text(json.dumps({"args": vars(args), "cells": cells}, indent=1)) + with (out_dir / "matrix.csv").open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=list(cells[0].keys())) + writer.writeheader() + writer.writerows(cells) + + for c in cells: + print( + f"{c['model']:>24} L{c['length']:<4}{c['condition']:<20}" + f"F1 {c['f1']:.3f}±{c['f1_se']:.3f} reas {c['tok_reasoning']:>6} out {c['tok_out']:>6}" + f" p50 {c['latency_p50_s']:>6.1f}s ${c['cost_usd']:.4f}" + ) + print(f"\n-> {out_dir}/records.jsonl, matrix.csv, summary.json") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp08_foveate.py b/packages/snapcompact/research/exp08_foveate.py new file mode 100644 index 000000000..235765e79 --- /dev/null +++ b/packages/snapcompact/research/exp08_foveate.py @@ -0,0 +1,388 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp08: foveated two-tier reading — aggressive 5x8 archive + on-demand zoom. + +Turn 1: a 5x8 (313 cols x 196 rows = 61348 chars/page, 1.5x denser than the +6x10 winner) archive image + all questions; the model answers what it can and +replies `ZOOM rows A-B` where the region is too small. Turn 2: the union of +requested row bands is sliced from the chunk text (row r covers chars +[(r-1)*cols, r*cols)) and re-rendered at a comfortable 8x13 font as zoom +image(s); the conversation continues with the zoom images + the pending +questions. Answers are merged; F1/cost vs the img-6x10-sent baseline. + +Run from the snapcompact dir: uv run exp08_foveate.py +""" + +import argparse +import csv +import json +import re +import sys +import uuid +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import capacity, render # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, FONTS, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 + +MODELS = { + "gpt-5.5": (2.0, 16.0), + "google/gemini-3.5-flash": (0.6, 4.0), +} +LENGTHS = (50, 150) +CONDITIONS = ("fov-5x8-bw", "fov-5x8-sent") +ARCHIVE_FONT = "5x8" +ZOOM_FONT = "8x13" +ZOOM_SIZES = (520, 784, 1040, 1568) # smallest square that fits the band wins +# fov = conservative prompt (zoom only when too small), rows addressing, tight pad +# fov2 = eager prompt (zoom unless fully certain), rows addressing, wide pad +# fov3 = eager prompt, phrase addressing (model quotes partially-read anchor words; +# harness fuzzy-locates them in the chunk and zooms that row band) +PROTO = { + "fov": ("exp08-archive.md", 2, "rows"), + "fov2": ("exp08-archive-eager.md", 12, "rows"), + "fov3": ("exp08-archive-phrase.md", 12, "phrase"), +} + +# img-6x10-sent baseline (results/optimal-gpt55 + optimal-gemini): f1, se, cost$ +BASELINE = { + ("gpt-5.5", 50): (0.850, 0.0508, 0.068), + ("gpt-5.5", 150): (0.8218, 0.0290, 0.2452), + ("google/gemini-3.5-flash", 50): (0.9841, 0.0119, 0.0181), + ("google/gemini-3.5-flash", 150): (0.8046, 0.0349, 0.097), +} + +_ZOOM_RANGE = re.compile(r"(?i)\bzoom\b[^\d]*(\d+)\s*(?:[-\u2013\u2014]|to\b)\s*(\d+)") +_ZOOM_SINGLE = re.compile(r"(?i)\bzoom\b[^\d]*(\d+)") +_ZOOM_PHRASE = re.compile(r"(?i)\bzoom\b\s*[\"\u201c']+(.+?)[\"\u201d']*\s*$") + + +def cached(model: str, tag: str, payload: object, fn, fresh: bool) -> dict: + """Disk-cache `fn() -> dict` keyed by (model, tag, payload). Truncations are not cached.""" + key = sha8(model, tag, json.dumps(payload, sort_keys=True, default=str)) + path = QA_CACHE / f"{key}.json" + if path.exists() and not fresh: + hit = json.loads(path.read_text()) + if hit.get("stop") != "max_tokens": + return hit + out = fn() + if out.get("stop") == "max_tokens": + print(f" WARN truncated, not cached: {model} {tag} {key}") + else: + path.write_text(json.dumps(out)) + return out + + +def atomic_png(img, path: Path) -> None: + tmp = path.with_name(f"{path.stem}.{uuid.uuid4().hex[:8]}.tmp.png") + img.save(tmp) + tmp.replace(path) + + +def parse_zoom(answer: str) -> tuple[int, int] | None: + """`ZOOM rows A-B` (or single row) -> (A, B); None if not a zoom request.""" + m = _ZOOM_RANGE.search(answer) + if m: + a, b = int(m.group(1)), int(m.group(2)) + return (a, b) if a <= b else (b, a) + m = _ZOOM_SINGLE.search(answer) + if m: + r = int(m.group(1)) + return (r, r) + return None + + +def locate_phrase(chunk: str, phrase: str) -> tuple[int, int] | None: + """Best-effort char span of a (possibly misread) anchor phrase in the chunk.""" + lower = re.sub(r"[^a-z0-9]", " ", chunk.lower()) # length-preserving normalize + p = re.sub(r"[^a-z0-9]", " ", phrase.lower()).split() + if not p: + return None + exact = " ".join(p) + i = lower.find(exact) + if i >= 0: + return i, i + len(exact) + words = [(m.start(), m.group()) for m in re.finditer(r"\S+", lower)] + pset = set(p) + k = max(2 * len(p), 8) + best_score, best_pos = 0, None + for s in range(len(words)): + score = sum(1 for _, w in words[s : s + k] if w in pset) + if score > best_score: + best_score, best_pos = score, words[s][0] + if best_pos is not None and best_score >= max(2, (len(p) + 1) // 2): + return best_pos, best_pos + len(exact) + return None + + +def merge_bands(bands: list[tuple[int, int]], max_row: int, pad: int) -> list[tuple[int, int]]: + """Pad by `pad`, clamp to [1, max_row], merge overlapping/adjacent bands.""" + padded = sorted((max(1, a - pad), min(max_row, b + pad)) for a, b in bands) + merged: list[tuple[int, int]] = [] + for a, b in padded: + if merged and a <= merged[-1][1] + 1: + merged[-1] = (merged[-1][0], max(merged[-1][1], b)) + else: + merged.append((a, b)) + return merged + + +def zoom_renders(chunk_text: str, bands: list[tuple[int, int]], arch_cols: int) -> list[tuple[tuple[int, int], Path]]: + """Slice each band's rows from the chunk and render at ZOOM_FONT; oversized bands split.""" + zcfg = FONTS[ZOOM_FONT] + max_rows = capacity(zcfg, ZOOM_SIZES[-1])[2] // arch_cols # archive rows per zoom page + out = [] + for a, b in bands: + pieces = [(s, min(s + max_rows - 1, b)) for s in range(a, b + 1, max_rows)] + for pa, pb in pieces: + txt = chunk_text[(pa - 1) * arch_cols : pb * arch_cols] + if not txt.strip(): + continue + size = next((s for s in ZOOM_SIZES if capacity(zcfg, s)[2] >= len(txt)), ZOOM_SIZES[-1]) + png = CACHE / f"exp08-zoom-{ZOOM_FONT}-{size}-{sha8(txt)}.png" + if not png.exists() or png.stat().st_size == 0: + atomic_png(render(txt, zcfg, CACHE, size, "bw"), png) + out.append(((pa, pb), png)) + return out + + +def run_cell_chunk(model: str, cond: str, start: int, end: int, ctx: dict) -> list[dict]: + """One (model, condition, chunk): archive QA turn, optional zoom turn, merge, score.""" + args, flow, paras, offsets, keys = ctx["args"], ctx["flow"], ctx["paras"], ctx["offsets"], ctx["keys"] + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + chunk_text = flow[start:end] + proto, _, variant = cond.split("-", 2) + prompt_file, pad, mode = PROTO[proto] + cfg = FONTS[ARCHIVE_FONT] + cols, rows, _ = capacity(cfg, args.size) + + png = CACHE / f"exp08-arch-{ARCHIVE_FONT}-{variant}-{sha8(chunk_text, str(args.size))}.png" + if not png.exists() or png.stat().st_size == 0: + atomic_png(render(chunk_text, cfg, CACHE, args.size, variant), png) + + q_block = "\n".join(f"{i + 1}. {q['q']}" for i, q in enumerate(questions)) + messages = [ + { + "role": "user", + "content": [ + {"text": load_prompt(prompt_file).format(cols=cols, rows=rows)}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + qa1 = cached( + model, "exp08-qa1", {"messages": messages, "effort": args.effort}, + lambda: dict(zip(("text", "usage", "stop"), + llm_complete(keys, model, messages, max_tokens=args.max_tokens, effort=args.effort))), + args.fresh, + ) + usage_rows = [("qa1", qa1["usage"])] + answers1 = squad.parse_numbered(qa1["text"], len(questions)) + zoom_req: list[tuple[int, int] | None] = [] + anchors: list[str | None] = [] + for a in answers1: + anchor = None + if mode == "phrase": + m = _ZOOM_PHRASE.search(a) + text = m.group(1) if m else None + if text is None and re.search(r"(?i)\bzoom\b", a) and not parse_zoom(a): + text = re.sub(r"(?i)^.*?\bzoom\b[:\s]*", "", a).strip("\"'\u201c\u201d ") + if text and len(text.split()) >= 2: + anchor = text + span = locate_phrase(chunk_text, anchor) + zoom_req.append((span[0] // cols + 1, span[1] // cols + 1) if span else None) + elif re.search(r"(?i)\bzoom\b", a): + zoom_req.append(parse_zoom(a)) # rows fallback + else: + zoom_req.append(None) + else: + zoom_req.append(parse_zoom(a)) + anchors.append(anchor) + requested = [i for i, a in enumerate(answers1) if re.search(r"(?i)\bzoom\b", a)] + pending = [i for i, z in enumerate(zoom_req) if z is not None] + final = list(answers1) + for i in requested: + if zoom_req[i] is None: + final[i] = "UNREADABLE" # zoom requested but band unresolvable + + if pending: + bands = merge_bands([zoom_req[i] for i in pending], rows, pad) + zooms = zoom_renders(chunk_text, bands, cols) + z_content: list[dict] = [{"text": load_prompt("exp08-zoom.md")}] + for (a, b), zpng in zooms: + z_content.append({"text": f"Zoom of archive rows {a}-{b}:"}) + z_content.append({"image_path": zpng}) + z_content.append({"text": "\n".join(f"{i + 1}. {questions[i]['q']}" for i in pending)}) + messages2 = messages + [ + {"role": "assistant", "content": [{"text": qa1["text"]}]}, + {"role": "user", "content": z_content}, + ] + qa2 = cached( + model, "exp08-qa2", {"messages": messages2, "effort": args.effort}, + lambda: dict(zip(("text", "usage", "stop"), + llm_complete(keys, model, messages2, max_tokens=args.max_tokens, effort=args.effort))), + args.fresh, + ) + usage_rows.append(("qa2", qa2["usage"])) + answers2 = squad.parse_numbered(qa2["text"], len(questions)) + for i in pending: + final[i] = answers2[i] or "UNREADABLE" + + records = [] + for i, (q, a) in enumerate(zip(questions, final)): + records.append( + { + "model": model, + "length": ctx["length"], + "cond": cond, + "chunk": start, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "answer_turn1": answers1[i], + "zoomed": i in requested, + "zoom_band": list(zoom_req[i]) if zoom_req[i] else None, + "anchor": anchors[i], + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + ) + records[0]["usage"] = [{"phase": p, **u} for p, u in usage_rows] + return records + + +def _phase_cost(us: list[dict], price_in: float, price_out: float) -> float: + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r")} + return (tok["in"] + 1.25 * tok["cache_w"] + 0.1 * tok["cache_r"]) / 1e6 * price_in + tok["out"] / 1e6 * price_out + + +def aggregate(records: list[dict], price_in: float, price_out: float) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean_f1 = sum(f1s) / n + se = (sum((x - mean_f1) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + us = [u for r in records if "usage" in r for u in r["usage"]] + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + cost_in = (tok["in"] + 1.25 * tok["cache_w"] + 0.1 * tok["cache_r"]) / 1e6 * price_in + cost_out = tok["out"] / 1e6 * price_out + zoomed = sum(r["zoomed"] for r in records) + zoom_chunks = sum(1 for r in records if "usage" in r and any(u["phase"] == "qa2" for u in r["usage"])) + return { + "n": n, + "em": sum(r["em"] for r in records) / n, + "f1": mean_f1, + "f1_se": se, + "abstained": sum(r["abstained"] for r in records), + "zoom_q": zoomed, + "zoom_rate": round(zoomed / n, 4), + "no_zoom_pct": round(100 * (n - zoomed) / n, 1), + "zoom_chunks": zoom_chunks, + **{f"tok_{k}": v for k, v in tok.items()}, + "cost_in_usd": round(cost_in, 4), + "cost_out_usd": round(cost_out, 4), + "cost_usd": round(cost_in + cost_out, 4), + "cost_zoom_usd": round(_phase_cost([u for u in us if u["phase"] == "qa2"], price_in, price_out), 4), + } + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--models", default=",".join(MODELS)) + ap.add_argument("--lengths", default=",".join(map(str, LENGTHS))) + ap.add_argument("--conditions", default=",".join(CONDITIONS)) + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--effort", default=None) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--env", default="~/.env") + ap.add_argument("--out", default="exp08-foveate") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + out_dir = RESULTS / args.out + out_dir.mkdir(parents=True, exist_ok=True) + + models = [m.strip() for m in args.models.split(",") if m.strip()] + lengths = [int(x) for x in args.lengths.split(",") if x.strip()] + conditions = [c.strip() for c in args.conditions.split(",") if c.strip()] + keys = { + "openai": load_env_key("OPENAI_API_KEY", args.env), + "openrouter": load_env_key("OPENROUTER_API_KEY", args.env), + } + + budget = capacity(FONTS[ARCHIVE_FONT], args.size)[2] + all_paras = squad.load_paragraphs(CACHE) + tasks = [] + for length in lengths: + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + ctx = {"args": args, "flow": flow, "paras": paras, "offsets": offsets, "keys": keys, "length": length} + for model in models: + for cond in conditions: + for start in range(0, len(flow), budget): + tasks.append((model, cond, start, min(start + budget, len(flow)), ctx)) + print(f"grid: {len(models)} models x {len(lengths)} lengths x {len(conditions)} conditions = {len(tasks)} chunk tasks") + + records: list[dict] = [] + done = 0 + with ThreadPoolExecutor(args.workers) as pool: + futures = [pool.submit(run_cell_chunk, m, c, s, e, ctx) for m, c, s, e, ctx in tasks] + for fut in futures: + records.extend(fut.result()) + done += 1 + print(f" {done}/{len(tasks)} tasks", flush=True) + + with (out_dir / "records.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + + cells = [] + for model in models: + for length in lengths: + for cond in conditions: + sub = [r for r in records if r["model"] == model and r["length"] == length and r["cond"] == cond] + if not sub: + continue + cell = {"model": model, "length": length, "condition": cond, **aggregate(sub, *MODELS[model])} + base = BASELINE.get((model, length)) + if base: + cell["base_f1"] = base[0] + cell["base_cost_usd"] = base[2] + cell["d_f1"] = round(cell["f1"] - base[0], 4) + cell["d_cost_usd"] = round(cell["cost_usd"] - base[2], 4) + cells.append(cell) + (out_dir / "summary.json").write_text(json.dumps({"args": vars(args), "cells": cells}, indent=1)) + with (out_dir / "matrix.csv").open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=list(cells[0].keys())) + writer.writeheader() + writer.writerows(cells) + + for c in cells: + print( + f"{c['model']:>26} L{c['length']:<4}{c['condition']:<14} n={c['n']:<4} f1={c['f1']:.3f}±{c['f1_se']:.3f} " + f"zoom={c['zoom_rate']:.0%} ${c['cost_usd']:.3f} (zoom ${c['cost_zoom_usd']:.3f}) " + f"vs base f1={c.get('base_f1', float('nan')):.3f} ${c.get('base_cost_usd', float('nan')):.3f}" + ) + print(f"\ndataset -> {out_dir}/records.jsonl, matrix.csv, summary.json") + + +if __name__ == "__main__": + main() + + diff --git a/packages/snapcompact/research/exp09_cacheappend.py b/packages/snapcompact/research/exp09_cacheappend.py new file mode 100644 index 000000000..1d67070e0 --- /dev/null +++ b/packages/snapcompact/research/exp09_cacheappend.py @@ -0,0 +1,365 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp09: append-only optical pages vs rewrite-compaction — prompt-cache economics. + +Simulates a growing session over K sequential chunks of the SQuAD flow (length 150 +-> 3 pages of img-6x10-sent @ 1568). At each step k the context is pages 1..k and a +QA turn samples questions over ALL pages so far (seed 42). + +Regimes: + A append-optical : prefix = fixed frame + k byte-identical PNG pages (rendered + once, reused), QA message last. Prefix grows append-only -> + provider prompt cache should re-bill old pages at 0.1x. + Zero LLM calls on the write path. + B rewrite-compact: each step re-summarizes the whole history text with + agent compaction-summary.md (one fresh LLM call per step = + write-path cost), QA over the fresh summary. The summary + rewrite invalidates any prompt-cache prefix. + +Cache probe: the same multi-image prefix is sent twice back-to-back (disk cache +bypassed via distinct probe-call payloads) and cache_r is read on both calls — +do image input tokens actually get prefix-cache hits on OpenAI Responses and +OpenRouter/Gemini? + +Outputs: results/exp09-cacheappend/{records.jsonl, steps.csv, matrix.csv, summary.json} +""" + +import argparse +import csv +import hashlib +import io +import json +import sys +import time +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import capacity, render # noqa: E402 +from final import ACK, cached, session_frame # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, QA_CACHE, RESULTS, FONTS, agent_prompt, load_prompt, sha8 # noqa: E402 + +MODELS = {"gpt-5.5": (2.0, 16.0), "google/gemini-3.5-flash": (0.6, 4.0)} +FONT = "6x10" +VARIANT = "sent" +LENGTH = 150 +MAX_STEPS = 4 +PROBE_TAIL = "Reply with exactly the word OK and nothing else." + + +def call(keys: dict, model: str, messages: list[dict], system: str | None = None, max_tokens: int = 32768) -> dict: + t0 = time.monotonic() + text, usage, stop = llm_complete(keys, model, messages, system=system, max_tokens=max_tokens) + return {"text": text, "usage": usage, "stop": stop, "secs": round(time.monotonic() - t0, 2)} + + +def usd(u: dict, p_in: float, p_out: float) -> float: + return (u.get("in", 0) + 0.1 * u.get("cache_r", 0)) / 1e6 * p_in + u.get("out", 0) / 1e6 * p_out + + +def usd_nocache(u: dict, p_in: float, p_out: float) -> float: + return (u.get("in", 0) + u.get("cache_r", 0)) / 1e6 * p_in + u.get("out", 0) / 1e6 * p_out + + +def render_pages(flow: str, size: int) -> tuple[list[tuple[int, int, Path]], dict]: + """Render each page once; byte-identical files reused across steps. Returns pages + determinism info.""" + budget = capacity(FONTS[FONT], size)[2] + pages = [] + for i, start in enumerate(range(0, len(flow), budget)): + chunk = flow[start : start + budget] + png = CACHE / f"exp09-page{i + 1}-{sha8(chunk, str(size), VARIANT)}.png" + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(".tmp.png") + render(chunk, FONTS[FONT], CACHE, size, VARIANT).save(tmp) + tmp.replace(png) + pages.append((start, min(start + budget, len(flow)), png)) + # Write-path determinism of regime A: render page 1 twice in-memory, compare bytes. + chunk1 = flow[:budget] + digests = [] + for _ in range(2): + buf = io.BytesIO() + render(chunk1, FONTS[FONT], CACHE, size, VARIANT).save(buf, format="PNG") + digests.append(hashlib.sha256(buf.getvalue()).hexdigest()) + det = {"render_sha256": digests, "deterministic": digests[0] == digests[1]} + return pages, det + + +def prefix_messages(k: int, pages: list, cols: int, rows: int) -> list[dict]: + """Append-only prefix: frame + pages 1..k, each ACKed. Byte-stable across steps.""" + msgs = [ + {"role": "user", "content": [{"text": load_prompt("exp09-frame.md").format(cols=cols, rows=rows)}, {"image_path": pages[0][2]}]}, + {"role": "assistant", "content": [{"text": ACK}]}, + ] + for i in range(1, k): + msgs.append({"role": "user", "content": [{"text": load_prompt("exp09-page.md").format(page=i + 1)}, {"image_path": pages[i][2]}]}) + msgs.append({"role": "assistant", "content": [{"text": ACK}]}) + return msgs + + +def step_questions(paras: list, offsets: list, end: int, qpc: int, seed: int) -> tuple[list[dict], str]: + qs = squad.sample_chunk_questions(paras, offsets, 0, end, qpc, seed) + return qs, "\n".join(f"{i + 1}. {q['q']}" for i, q in enumerate(qs)) + + +def score_records(model: str, regime: str, step: int, questions: list[dict], text: str) -> list[dict]: + answers = squad.parse_numbered(text, len(questions)) + return [ + { + "model": model, + "length": LENGTH, + "cond": regime, + "step": step, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + for q, a in zip(questions, answers) + ] + + +def common_prefix_len(a: str, b: str) -> int: + n = 0 + for x, y in zip(a, b): + if x != y: + break + n += 1 + return n + + +def run_model(model: str, ctx: dict) -> dict: + args, keys, flow, paras, offsets, pages = ctx["args"], ctx["keys"], ctx["flow"], ctx["paras"], ctx["offsets"], ctx["pages"] + p_in, p_out = MODELS[model] + cols, rows, _ = capacity(FONTS[FONT], args.size) + K = len(pages) + records: list[dict] = [] + steps: list[dict] = [] + + # --- Regime A: append-optical (sequential; step k's prefix warms step k+1's cache) --- + cum_a = 0.0 + for k in range(1, K + 1): + end = pages[k - 1][1] + questions, q_block = step_questions(paras, offsets, end, args.qpc, args.seed) + msgs = prefix_messages(k, pages, cols, rows) + [ + {"role": "user", "content": [{"text": load_prompt("exp09-qa.md").format(questions=q_block)}]} + ] + qa = cached( + model, "exp09-A-qa", {"step": k, "messages": msgs}, + lambda: call(keys, model, msgs, max_tokens=args.max_tokens), args.fresh, + ) + recs = score_records(model, "append-optical", k, questions, qa["text"]) + recs[0]["usage"] = [{"phase": "qa", **qa["usage"]}] + records += recs + cost = usd(qa["usage"], p_in, p_out) + cum_a += cost + u = qa["usage"] + steps.append( + { + "model": model, "regime": "append-optical", "step": k, "n": len(recs), + "f1": round(sum(r["f1"] for r in recs) / len(recs), 3), + "write_in": 0, "write_out": 0, "write_secs": 0.0, "write_cost": 0.0, + "qa_in": u["in"], "qa_cache_r": u["cache_r"], "qa_out": u["out"], + "qa_reasoning": u.get("reasoning", 0), "qa_secs": qa["secs"], + "step_cost": round(cost, 4), "step_cost_nocache": round(usd_nocache(u, p_in, p_out), 4), + "cum_cost": round(cum_a, 4), + } + ) + print(f" {model} A step {k}: in={u['in']} cache_r={u['cache_r']} out={u['out']} f1={steps[-1]['f1']}", flush=True) + + # --- Cache probe: identical multi-image prefix twice in a row (disk cache bypassed via call index) --- + probe_msgs = prefix_messages(K, pages, cols, rows) + [{"role": "user", "content": [{"text": PROBE_TAIL}]}] + probe = [] + for i in (1, 2, 3): + r = cached( + model, "exp09-probe", {"call": i, "messages": probe_msgs}, + lambda: call(keys, model, probe_msgs, max_tokens=args.max_tokens), args.fresh, + ) + probe.append({"call": i, **r["usage"], "secs": r["secs"]}) + print(f" {model} probe call {i}: in={r['usage']['in']} cache_r={r['usage']['cache_r']}", flush=True) + + # --- Regime B: rewrite-compact (fresh summary each step = write path) --- + cum_b = 0.0 + summaries: list[str] = [] + for k in range(1, K + 1): + end = pages[k - 1][1] + text_k = flow[:end] + questions, q_block = step_questions(paras, offsets, end, args.qpc, args.seed) + sm = cached( + model, "exp09-B-sum", {"step": k, "chunk": sha8(text_k)}, + lambda: call( + keys, model, + session_frame(text_k) + [{"role": "user", "content": [{"text": agent_prompt("compaction-summary.md")}]}], + system=agent_prompt("summarization-system.md"), max_tokens=args.max_tokens, + ), + args.fresh, + ) + summaries.append(sm["text"]) + qa_msgs = [ + {"role": "user", "content": [{"text": load_prompt("qa-text.md").format(context=sm["text"])}, {"text": q_block}]} + ] + qa = cached( + model, "exp09-B-qa", {"step": k, "summary": sm["text"], "q": q_block}, + lambda: call(keys, model, qa_msgs, max_tokens=args.max_tokens), args.fresh, + ) + recs = score_records(model, "rewrite-compact", k, questions, qa["text"]) + recs[0]["usage"] = [{"phase": "summarize", **sm["usage"]}, {"phase": "qa", **qa["usage"]}] + records += recs + w_cost = usd(sm["usage"], p_in, p_out) + q_cost = usd(qa["usage"], p_in, p_out) + cum_b += w_cost + q_cost + su, qu = sm["usage"], qa["usage"] + steps.append( + { + "model": model, "regime": "rewrite-compact", "step": k, "n": len(recs), + "f1": round(sum(r["f1"] for r in recs) / len(recs), 3), + "write_in": su["in"] + su["cache_r"], "write_out": su["out"], + "write_secs": sm["secs"], "write_cost": round(w_cost, 4), + "qa_in": qu["in"], "qa_cache_r": qu["cache_r"], "qa_out": qu["out"], + "qa_reasoning": qu.get("reasoning", 0), "qa_secs": qa["secs"], + "step_cost": round(w_cost + q_cost, 4), + "step_cost_nocache": round(usd_nocache(su, p_in, p_out) + usd_nocache(qu, p_in, p_out), 4), + "cum_cost": round(cum_b, 4), + } + ) + print(f" {model} B step {k}: write {su['in']}+{su['cache_r']}c->{su['out']} ({sm['secs']}s) f1={steps[-1]['f1']}", flush=True) + + # Write-path determinism of regime B: re-run the step-1 summarize with identical payload (fresh key). + det = cached( + model, "exp09-B-sum-det", {"step": 1, "chunk": sha8(flow[: pages[0][1]])}, + lambda: call( + keys, model, + session_frame(flow[: pages[0][1]]) + [{"role": "user", "content": [{"text": agent_prompt("compaction-summary.md")}]}], + system=agent_prompt("summarization-system.md"), max_tokens=args.max_tokens, + ), + args.fresh, + ) + b_det = { + "identical": det["text"] == summaries[0], + "common_prefix_chars": common_prefix_len(det["text"], summaries[0]), + "len_a": len(summaries[0]), "len_b": len(det["text"]), + } + # Cross-step summary prefix stability (the thing the prompt cache would need). + step_stability = [ + {"steps": f"{k}->{k + 1}", "common_prefix_chars": common_prefix_len(summaries[k - 1], summaries[k]), + "len_prev": len(summaries[k - 1]), "len_next": len(summaries[k])} + for k in range(1, K) + ] + return {"records": records, "steps": steps, "probe": probe, "b_determinism": b_det, "b_step_stability": step_stability} + + +def aggregate(records: list[dict], p_in: float, p_out: float) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean_f1 = sum(f1s) / n + se = (sum((x - mean_f1) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + us = [u for r in records if "usage" in r for u in r["usage"]] + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + cost_in = (tok["in"] + 0.1 * tok["cache_r"]) / 1e6 * p_in + cost_out = tok["out"] / 1e6 * p_out + return { + "n": n, + "em": sum(r["em"] for r in records) / n, + "f1": mean_f1, + "f1_se": se, + "abstained": sum(r["abstained"] for r in records), + **{f"tok_{k}": v for k, v in tok.items()}, + "cost_in_usd": round(cost_in, 4), + "cost_out_usd": round(cost_out, 4), + "cost_usd": round(cost_in + cost_out, 4), + } + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--models", default=",".join(MODELS)) + ap.add_argument("--qpc", type=int, default=10) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + out_dir = RESULTS / "exp09-cacheappend" + out_dir.mkdir(parents=True, exist_ok=True) + + models = [m.strip() for m in args.models.split(",") if m.strip()] + keys = { + "openai": load_env_key("OPENAI_API_KEY", args.env), + "openrouter": load_env_key("OPENROUTER_API_KEY", args.env), + } + + paras = squad.load_paragraphs(CACHE)[:LENGTH] + flow, offsets = squad.build_flow(paras) + pages, a_det = render_pages(flow, args.size) + pages = pages[:MAX_STEPS] + print(f"flow {len(flow)} chars -> {len(pages)} pages (K={len(pages)} steps); render deterministic: {a_det['deterministic']}") + + ctx = {"args": args, "keys": keys, "flow": flow, "paras": paras, "offsets": offsets, "pages": pages} + with ThreadPoolExecutor(min(2, len(models))) as pool: + results = dict(zip(models, pool.map(lambda m: run_model(m, ctx), models))) + + records = [r for m in models for r in results[m]["records"]] + steps = [s for m in models for s in results[m]["steps"]] + + with (out_dir / "records.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + with (out_dir / "steps.csv").open("w", newline="") as fh: + w = csv.DictWriter(fh, fieldnames=list(steps[0].keys())) + w.writeheader() + w.writerows(steps) + + cells = [] + for model in models: + for cond in ("append-optical", "rewrite-compact"): + sub = [r for r in records if r["model"] == model and r["cond"] == cond] + final_step = max(r["step"] for r in sub) + fin = [r for r in sub if r["step"] == final_step] + cell = {"model": model, "length": LENGTH, "condition": cond, **aggregate(sub, *MODELS[model])} + cell["final_step_f1"] = round(sum(r["f1"] for r in fin) / len(fin), 3) + cell["final_step_em"] = round(sum(r["em"] for r in fin) / len(fin), 3) + cells.append(cell) + with (out_dir / "matrix.csv").open("w", newline="") as fh: + w = csv.DictWriter(fh, fieldnames=list(cells[0].keys())) + w.writeheader() + w.writerows(cells) + + summary = { + "args": vars(args), + "pages": [{"start": s, "end": e, "png": p.name} for s, e, p in pages], + "render_determinism": a_det, + "cells": cells, + "steps": steps, + "probe": {m: results[m]["probe"] for m in models}, + "b_determinism": {m: results[m]["b_determinism"] for m in models}, + "b_step_stability": {m: results[m]["b_step_stability"] for m in models}, + } + (out_dir / "summary.json").write_text(json.dumps(summary, indent=1, default=str)) + + print("\n== per-step (qa_in / qa_cache_r / write_cost / step_cost / cum_cost / f1) ==") + for s in steps: + print( + f"{s['model']:<26} {s['regime']:<16} k={s['step']} in={s['qa_in']:>6} cache_r={s['qa_cache_r']:>6} " + f"write=${s['write_cost']:.4f} step=${s['step_cost']:.4f} cum=${s['cum_cost']:.4f} f1={s['f1']:.3f}" + ) + print("\n== cache probe (same multi-image prefix twice) ==") + for m in models: + for p in results[m]["probe"]: + print(f"{m:<26} call {p['call']}: in={p['in']:>6} cache_r={p['cache_r']:>6} secs={p['secs']}") + print(f"\ndataset -> {out_dir}/records.jsonl, steps.csv, matrix.csv, summary.json") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp10_profiles.py b/packages/snapcompact/research/exp10_profiles.py new file mode 100644 index 000000000..127edcd5f --- /dev/null +++ b/packages/snapcompact/research/exp10_profiles.py @@ -0,0 +1,358 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp10: per-model optical-profile calibration sweep for gpt-5.5 and gemini-3.5-flash. + +Sibling models each preferred a different (font, variant) density point +(fable: 6x12-dim, opus: 8x13-bw, kimi: 8x13-sent-dim, glm: 8x13-dark-sent), but +gpt-5.5 / gemini-3.5-flash were only ever measured at 6x10-sent (+ 5x8-bw / +6x9-sent-dim). This script runs the calibration sweep an "optical profile" +catalog entry would ship with: + + phase A (screen, length 150): sibling-winner combos + a variant probe at 8x13 + + 6x12-{dim,sent} + phase B (ladder, length 150): density ladder {8x13..5x8} at each model's best + variant from phase A + phase C (cross, length 150): each model evaluated on the *other* model's + optimal cell (profile transferability) + phase D (confirm): each model's top cell at lengths 50 and 250 + +Methodology matches the optimal-* baselines exactly (seed 42, qpc 30, size +1568, max_tokens 32768, effort None) so identical cells hit the shared +.cache/qa/ response cache for free. Cache policy: reads try the canonical +"qa" tag first (free reuse of optimal-run payloads), new responses are written +under the exp10-namespaced tag. + +Usage: uv run exp10_profiles.py # full self-driving sweep + uv run exp10_profiles.py --report # re-aggregate from cache only +""" + +import argparse +import json +import os +import sys +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import capacity, render # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, FONTS, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 + +EXP = "exp10" +OUT_DIR = RESULTS / f"{EXP}-profiles" +MODELS = { # ($/M in, $/M out); cached reads bill 0.1x input + "gpt-5.5": (2.0, 16.0), + "google/gemini-3.5-flash": (0.6, 4.0), +} +BASELINE = { # img-6x10-sent from results/optimal-{gpt55,gemini}/matrix.csv + ("gpt-5.5", 50): (0.850, 0.051, 0.068), + ("gpt-5.5", 150): (0.822, 0.029, 0.245), + ("gpt-5.5", 250): (0.822, 0.026, 0.380), + ("google/gemini-3.5-flash", 50): (0.984, 0.012, 0.018), + ("google/gemini-3.5-flash", 150): (0.805, 0.035, 0.097), + ("google/gemini-3.5-flash", 250): (0.755, 0.033, 0.147), +} + +# Phase A screening cells (length 150). Sibling winners + variant probe at +# 8x13 + the 6x12 bridge. img-6x10-sent is the baseline -- not re-run. +SCREEN = ( + "img-6x12-dim", # fable's winner + "img-8x13-bw", # opus's winner + "img-8x13-sent-dim", # kimi's winner + "img-8x13-dark-sent", # glm's winner + "img-8x13-sent", + "img-8x13-dim", + "img-6x12-sent", +) +LADDER_FONTS = ("8x13", "7x13", "6x12", "6x10", "6x9", "5x8") + + +def parse_img_condition(name: str) -> tuple[str, str]: + _, font, variant = name.split("-", 2) + return font, variant + + +def cached(model: str, payload: object, fn, fresh: bool) -> dict: + """Like final.cached(), but dual-key: read canonical "qa" tag first (free + reuse of the optimal-run cache), then our exp10 tag; write under exp10.""" + blob = json.dumps(payload, sort_keys=True, default=str) + canon = QA_CACHE / f"{sha8(model, 'qa', blob)}.json" + mine = QA_CACHE / f"{sha8(model, f'{EXP}-qa', blob)}.json" + if not fresh: + for path in (canon, mine): + if path.exists(): + hit = json.loads(path.read_text()) + if hit.get("stop") != "max_tokens": + return hit + out = fn() + if out.get("stop") == "max_tokens": + print(f" WARN truncated, not cached: {model} {sha8(model, f'{EXP}-qa', blob)}") + else: + tmp = mine.with_suffix(f".{os.getpid()}.tmp") + tmp.write_text(json.dumps(out)) + tmp.replace(mine) + return out + + +def render_png(chunk_text: str, font: str, variant: str, size: int) -> Path: + """Canonical final.py naming so QA payloads (which embed the path) match + the shared cache; render is deterministic, tmp-then-replace is atomic.""" + png = CACHE / f"img-{font}-{variant}-{sha8(chunk_text, str(size))}.png" + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(f".{os.getpid()}.tmp.png") + render(chunk_text, FONTS[font], CACHE, size, variant).save(tmp) + tmp.replace(png) + return png + + +def run_cell_chunk(model: str, cond: str, start: int, end: int, ctx: dict) -> list[dict]: + """One (model, condition, chunk): render carrier image, QA, score. + Copied from final.run_cell_chunk, image conditions only.""" + args, flow, paras, offsets, keys = ctx["args"], ctx["flow"], ctx["paras"], ctx["offsets"], ctx["keys"] + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + chunk_text = flow[start:end] + q_block = "\n".join(f"{i + 1}. {q['q']}" for i, q in enumerate(questions)) + font, variant = parse_img_condition(cond) + png = render_png(chunk_text, font, variant, args.size) + cols, rows, _ = capacity(FONTS[font], args.size) + messages = [ + { + "role": "user", + "content": [ + {"text": load_prompt("qa-image.md").format(cols=cols, rows=rows)}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + qa = cached( + model, {"messages": messages, "extra": None, "effort": None}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete(keys, model, messages, max_tokens=args.max_tokens, effort=None), + ) + ), + args.fresh, + ) + answers = squad.parse_numbered(qa["text"], len(questions)) + records = [] + for q, a in zip(questions, answers): + records.append( + { + "model": model, + "length": ctx["length"], + "cond": cond, + "chunk": start, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + ) + records[0]["usage"] = [{"phase": "qa", **qa["usage"]}] + return records + + +def aggregate(records: list[dict], price_in: float, price_out: float) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean_f1 = sum(f1s) / n + se = (sum((x - mean_f1) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + us = [u for r in records if "usage" in r for u in r["usage"]] + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + cost_in = (tok["in"] + 1.25 * tok["cache_w"] + 0.1 * tok["cache_r"]) / 1e6 * price_in + cost_out = tok["out"] / 1e6 * price_out + return { + "n": n, + "em": sum(r["em"] for r in records) / n, + "f1": mean_f1, + "f1_se": se, + "abstained": sum(r["abstained"] for r in records), + **{f"tok_{k}": v for k, v in tok.items()}, + "cost_in_usd": round(cost_in, 4), + "cost_out_usd": round(cost_out, 4), + "cost_usd": round(cost_in + cost_out, 4), + } + + +class Runner: + def __init__(self, args, keys): + self.args = args + self.keys = keys + self.records: list[dict] = [] + self.done: set[tuple[str, int, str]] = set() + self.ctxs: dict[int, dict] = {} + self.all_paras = squad.load_paragraphs(CACHE) + + def ctx(self, length: int) -> dict: + if length not in self.ctxs: + paras = self.all_paras[:length] + flow, offsets = squad.build_flow(paras) + self.ctxs[length] = { + "args": self.args, "flow": flow, "paras": paras, + "offsets": offsets, "keys": self.keys, "length": length, + } + return self.ctxs[length] + + def run(self, cells: list[tuple[str, int, str]], label: str) -> None: + cells = [c for c in cells if c not in self.done] + self.done.update(cells) + tasks = [] + for model, length, cond in cells: + ctx = self.ctx(length) + flow = ctx["flow"] + budget = capacity(FONTS[parse_img_condition(cond)[0]], self.args.size)[2] + for start in range(0, len(flow), budget): + tasks.append((model, cond, start, min(start + budget, len(flow)), ctx)) + if not tasks: + return + print(f"[{label}] {len(cells)} cells -> {len(tasks)} chunk tasks") + with ThreadPoolExecutor(self.args.workers) as pool: + futures = [pool.submit(run_cell_chunk, *t) for t in tasks] + for fut in futures: + self.records.extend(fut.result()) + + def cell(self, model: str, length: int, cond: str) -> dict | None: + sub = [r for r in self.records if r["model"] == model and r["length"] == length and r["cond"] == cond] + return {"model": model, "length": length, "condition": cond, **aggregate(sub, *MODELS[model])} if sub else None + + def cells_for(self, model: str, length: int) -> list[dict]: + conds = sorted({r["cond"] for r in self.records if r["model"] == model and r["length"] == length}) + return [c for cond in conds if (c := self.cell(model, length, cond))] + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + OUT_DIR.mkdir(parents=True, exist_ok=True) + + keys = { + "openai": load_env_key("OPENAI_API_KEY", args.env), + "openrouter": load_env_key("OPENROUTER_API_KEY", args.env), + } + runner = Runner(args, keys) + models = list(MODELS) + + # -- phase A: screen at length 150 -------------------------------------- + runner.run([(m, 150, c) for m in models for c in SCREEN], "A screen") + + # -- phase B: density ladder at each model's best variant --------------- + ladder = [(m, 150, "img-6x10-sent") for m in models] # baseline; free via shared cache + best_variant = {} + for m in models: + top = max(runner.cells_for(m, 150), key=lambda c: c["f1"]) + v = parse_img_condition(top["condition"])[1] + best_variant[m] = v + ladder += [(m, 150, f"img-{f}-{v}") for f in LADDER_FONTS] + runner.run(ladder, "B ladder") + + # -- phase C: cross-profile transfer (each model on the other's optimum) - + top150 = {m: max(runner.cells_for(m, 150), key=lambda c: c["f1"]) for m in models} + cross = [(other, 150, top150[m]["condition"]) for m in models for other in models if other != m] + runner.run(cross, "C cross") + top150 = {m: max(runner.cells_for(m, 150), key=lambda c: c["f1"]) for m in models} + + # -- phase D: confirm top cell at lengths 50 and 250 --------------------- + runner.run([(m, ln, top150[m]["condition"]) for m in models for ln in (50, 250)], "D confirm") + + # -- outputs -------------------------------------------------------------- + with (OUT_DIR / "records.jsonl").open("w") as fh: + for r in runner.records: + fh.write(json.dumps(r) + "\n") + cells = [] + for m in models: + for ln in (50, 150, 250): + cells.extend(runner.cells_for(m, ln)) + cells.sort(key=lambda c: (c["model"], c["length"], -c["f1"])) + import csv + + with (OUT_DIR / "matrix.csv").open("w", newline="") as fh: + w = csv.DictWriter(fh, fieldnames=list(cells[0].keys())) + w.writeheader() + w.writerows(cells) + + profiles = {} + for m in models: + top = top150[m] + font, variant = parse_img_condition(top["condition"]) + cols, rows, chars = capacity(FONTS[font], args.size) + confirm = {ln: runner.cell(m, ln, top["condition"]) for ln in (50, 250)} + other = next(o for o in models if o != m) + transfer = runner.cell(other, 150, top["condition"]) + profiles[m] = { + "model": m, + "optical_profile": { + "font": font, + "variant": variant, + "px": args.size, + "cols": cols, + "rows": rows, + "chars_per_page": chars, + "prompt": "qa-image.md", + }, + "expected": { + "f1_at_150": round(top["f1"], 4), + "f1_se_at_150": round(top["f1_se"], 4), + "cost_usd_at_150": top["cost_usd"], + "confirm": { + str(ln): {"f1": round(c["f1"], 4), "se": round(c["f1_se"], 4), "cost_usd": c["cost_usd"]} + for ln, c in confirm.items() if c + }, + }, + "baseline_img_6x10_sent_f1_at_150": BASELINE[(m, 150)][0], + "transfer_f1_on_other_model_at_150": round(transfer["f1"], 4) if transfer else None, + } + (OUT_DIR / "profiles.json").write_text(json.dumps(profiles, indent=1)) + (OUT_DIR / "summary.json").write_text( + json.dumps({"args": vars(args), "best_variant": best_variant, "cells": cells}, indent=1) + ) + + # -- console report ------------------------------------------------------- + spend = 0.0 + for m in models: + print(f"\n== {m} (length 150 screening, sorted by F1) ==") + base_f1, base_se, base_cost = BASELINE[(m, 150)] + print(f"{'condition':<22}{'n':>5}{'EM':>7}{'F1':>7}{'se':>7}{'abst':>6}{'cost$':>8}{'dF1':>8}") + for c in sorted(runner.cells_for(m, 150), key=lambda c: -c["f1"]): + spend += c["cost_usd"] + print( + f"{c['condition']:<22}{c['n']:>5}{c['em']:>7.3f}{c['f1']:>7.3f}{c['f1_se']:>7.3f}" + f"{c['abstained']:>6}{c['cost_usd']:>8.3f}{c['f1'] - base_f1:>+8.3f}" + ) + print(f"{'img-6x10-sent [base]':<22}{'':>5}{'':>7}{base_f1:>7.3f}{base_se:>7.3f}{'':>6}{base_cost:>8.3f}{0:>+8.3f}") + for ln in (50, 250): + c = runner.cell(m, ln, top150[m]["condition"]) + if c: + spend += c["cost_usd"] + b = BASELINE[(m, ln)] + print( + f"confirm@{ln}: {c['condition']} F1={c['f1']:.3f}+-{c['f1_se']:.3f} cost=${c['cost_usd']:.3f}" + f" (baseline {b[0]:.3f}+-{b[1]:.3f} ${b[2]:.3f})" + ) + print(f"\ntotal cell cost (incl. cache-free cells): ${spend:.2f}") + print(f"-> {OUT_DIR}/records.jsonl, matrix.csv, summary.json, profiles.json") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp11_memhier.py b/packages/snapcompact/research/exp11_memhier.py new file mode 100644 index 000000000..a3badacc3 --- /dev/null +++ b/packages/snapcompact/research/exp11_memhier.py @@ -0,0 +1,301 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp11 memhier: memory hierarchy instead of either/or compression. + +Split the passage flow by age into thirds (StreamingLLM/H2O intuition, client-side): + L3 (oldest) -> narrative compaction summary (agent compaction-summary prompt) + L2 (middle) -> optical pages (img-6x10-sent, standard bdf render) + L1 (newest) -> verbatim text +One QA context = [L3 summary] + [L2 images] + [L1 text] + questions, framed by +prompts/exp11-qa-hier.md. `hier-appendix` additionally attaches the L3 text as a +dense 5x8-sent optical appendix (summary for gist + image for exact lookup). + +Question sets are identical to the uniform text baseline (same seed/qpc/40716-char +chunk grid), so cells are directly comparable. Each record carries a *global* +pos_rel and its tier, enabling the pos_rel-tercile breakdown. +""" + +import argparse +import csv +import json +import sys +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import capacity, render # noqa: E402 +from final import MODELS, aggregate, cached, session_frame # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, FONTS, QA_CACHE, RESULTS, TEXT_CHUNK, agent_prompt, load_prompt, sha8 # noqa: E402 + +L2_FONT, L2_VAR = "6x10", "sent" +APX_FONT, APX_VAR = "5x8", "sent" +CONDITIONS = ("hier", "hier-appendix") +APPENDIX_NOTE = ( + ", plus {n_apx} dense bitmap appendix image(s) of the SAME oldest text " + "(monospace pixel font, {acols} characters per row, {arows} rows; use the appendix " + "for exact lookups the summary lacks)" +) + + +def tier_bounds(offsets: list[int], flow_len: int) -> tuple[int, int]: + """Passage-start offsets nearest to 1/3 and 2/3 of the flow (no passage straddles a tier).""" + b1 = min(offsets, key=lambda o: abs(o - flow_len / 3)) + b2 = min(offsets, key=lambda o: abs(o - 2 * flow_len / 3)) + return b1, b2 + + +def render_pages(text: str, font: str, var: str, size: int) -> list[Path]: + cap = capacity(FONTS[font], size)[2] + pages = [] + for s in range(0, len(text), cap): + seg = text[s : s + cap] + png = CACHE / f"exp11-img-{font}-{var}-{sha8(seg, str(size))}.png" + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(".tmp.png") + render(seg, FONTS[font], CACHE, size, var).save(tmp) + tmp.replace(png) + pages.append(png) + return pages + + +def gen_summary(model: str, keys: dict, l3_text: str, max_tokens: int, fresh: bool) -> dict: + return cached( + model, "exp11-summary", {"chunk": l3_text}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete( + keys, model, + session_frame(l3_text) + + [{"role": "user", "content": [{"text": agent_prompt("compaction-summary.md")}]}], + system=agent_prompt("summarization-system.md"), + max_tokens=max_tokens, + ), + ) + ), + fresh, + ) + + +def context_blocks(cond: str, summary: str, l2_pages: list[Path], apx_pages: list[Path], l1_text: str, size: int) -> list[dict]: + cols, rows, _ = capacity(FONTS[L2_FONT], size) + apx_note = "" + if cond == "hier-appendix": + acols, arows, _ = capacity(FONTS[APX_FONT], size) + apx_note = APPENDIX_NOTE.format(n_apx=len(apx_pages), acols=acols, arows=arows) + frame = load_prompt("exp11-qa-hier.md").format( + appendix_note=apx_note, n_pages=len(l2_pages), cols=cols, rows=rows + ) + blocks: list[dict] = [{"text": frame}, {"text": f"TIER 3 — SUMMARY OF OLDEST THIRD:\n\n{summary}"}] + if cond == "hier-appendix": + for i, p in enumerate(apx_pages): + blocks.append({"text": f"TIER 3 appendix image {i + 1}/{len(apx_pages)} (same oldest text as dense bitmap):"}) + blocks.append({"image_path": p}) + for i, p in enumerate(l2_pages): + blocks.append({"text": f"TIER 2 page {i + 1}/{len(l2_pages)} (middle third as bitmap):"}) + blocks.append({"image_path": p}) + blocks.append({"text": f"TIER 1 — VERBATIM NEWEST THIRD:\n\n\n{l1_text}\n"}) + return blocks + + +def run_chunk(model: str, cond: str, start: int, end: int, cell: dict) -> list[dict]: + """One QA call: shared hierarchical context + this chunk's question batch.""" + args, keys, flow = cell["args"], cell["keys"], cell["flow"] + questions = squad.sample_chunk_questions(cell["paras"], cell["offsets"], start, end, args.qpc, args.seed) + if not questions: + return [] + q_block = "\n".join(f"{i + 1}. {q['q']}" for i, q in enumerate(questions)) + messages = [ + { + "role": "user", + "content": cell["blocks"][cond] + [{"text": f"QUESTIONS:\n{q_block}"}], + } + ] + qa = cached( + model, "exp11-qa", {"cond": cond, "length": cell["length"], "messages": messages, "effort": args.effort}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete(keys, model, messages, max_tokens=args.max_tokens, effort=args.effort), + ) + ), + args.fresh, + ) + answers = squad.parse_numbered(qa["text"], len(questions)) + b1, b2 = cell["bounds"] + records = [] + for q, a in zip(questions, answers): + pos_abs = start + q["pos_rel"] * (end - start) + tier = "L1" if pos_abs >= b2 else ("L2" if pos_abs >= b1 else "L3") + records.append( + { + "model": model, + "length": cell["length"], + "cond": cond, + "chunk": start, + "pos_rel": round(pos_abs / len(flow), 4), + "tier": tier, + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + ) + records[0]["usage"] = [{"phase": "qa", **qa["usage"]}] + return records + + +def tier_stats(records: list[dict]) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean = sum(f1s) / n + se = (sum((x - mean) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + return { + "n": n, + "em": round(sum(r["em"] for r in records) / n, 4), + "f1": round(mean, 4), + "f1_se": round(se, 4), + "abstained": sum(r["abstained"] for r in records), + } + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--models", default="gpt-5.5,google/gemini-3.5-flash") + ap.add_argument("--lengths", default="150,250") + ap.add_argument("--conditions", default=",".join(CONDITIONS)) + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--effort", default=None) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--env", default="~/.env") + ap.add_argument("--out", default="exp11-memhier") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + out_dir = RESULTS / args.out + out_dir.mkdir(parents=True, exist_ok=True) + + models = [m.strip() for m in args.models.split(",") if m.strip()] + lengths = [int(x) for x in args.lengths.split(",") if x.strip()] + conditions = [c.strip() for c in args.conditions.split(",") if c.strip()] + keys = { + "openai": load_env_key("OPENAI_API_KEY", args.env), + "openrouter": load_env_key("OPENROUTER_API_KEY", args.env), + } + + all_paras = squad.load_paragraphs(CACHE) + cells: dict[tuple[str, int], dict] = {} + summary_usage: dict[tuple[str, int], dict] = {} + for length in lengths: + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + b1, b2 = tier_bounds(offsets, len(flow)) + l3, l2, l1 = flow[:b1], flow[b1:b2], flow[b2:] + l2_pages = render_pages(l2, L2_FONT, L2_VAR, args.size) + apx_pages = render_pages(l3, APX_FONT, APX_VAR, args.size) if "hier-appendix" in conditions else [] + print( + f"length {length}: flow={len(flow)} chars, tiers L3={len(l3)} L2={len(l2)} L1={len(l1)}, " + f"l2_pages={len(l2_pages)} apx_pages={len(apx_pages)}" + ) + for model in models: + summ = gen_summary(model, keys, l3, args.max_tokens, args.fresh) + if summ.get("stop") == "max_tokens": + raise SystemExit(f"summary truncated for {model} length {length}; raise --max-tokens") + summary_usage[(model, length)] = summ["usage"] + print(f" summary[{model}]: {len(summ['text'])} chars") + cells[(model, length)] = { + "args": args, + "keys": keys, + "flow": flow, + "paras": paras, + "offsets": offsets, + "length": length, + "bounds": (b1, b2), + "blocks": { + cond: context_blocks(cond, summ["text"], l2_pages, apx_pages, l1, args.size) + for cond in conditions + }, + } + + tasks = [] + for (model, length), cell in cells.items(): + for cond in conditions: + for start in range(0, len(cell["flow"]), TEXT_CHUNK): + tasks.append((model, cond, start, min(start + TEXT_CHUNK, len(cell["flow"])), cell)) + print(f"grid: {len(tasks)} QA tasks") + + records: list[dict] = [] + done = 0 + with ThreadPoolExecutor(args.workers) as pool: + futures = [pool.submit(run_chunk, *t) for t in tasks] + for fut in futures: + records.extend(fut.result()) + done += 1 + print(f" {done}/{len(tasks)} tasks", flush=True) + + # Charge the (cached, shared) summarization once per cell — each condition is a + # standalone strategy that would need its own summary. + charged: set[tuple[str, int, str]] = set() + for r in records: + key = (r["model"], r["length"], r["cond"]) + if key not in charged and "usage" in r: + r["usage"].append({"phase": "summarize", **summary_usage[(r["model"], r["length"])]}) + charged.add(key) + + with (out_dir / "records.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + + cell_rows, tier_rows = [], [] + for model in models: + for length in lengths: + for cond in conditions: + sub = [r for r in records if r["model"] == model and r["length"] == length and r["cond"] == cond] + if not sub: + continue + cell_rows.append({"model": model, "length": length, "condition": cond, **aggregate(sub, *MODELS[model])}) + for tier in ("L3", "L2", "L1"): + tsub = [r for r in sub if r["tier"] == tier] + if tsub: + tier_rows.append({"model": model, "length": length, "condition": cond, "tier": tier, **tier_stats(tsub)}) + + (out_dir / "summary.json").write_text(json.dumps({"args": vars(args), "cells": cell_rows, "tiers": tier_rows}, indent=1)) + with (out_dir / "matrix.csv").open("w", newline="") as fh: + w = csv.DictWriter(fh, fieldnames=list(cell_rows[0].keys())) + w.writeheader() + w.writerows(cell_rows) + with (out_dir / "terciles.csv").open("w", newline="") as fh: + w = csv.DictWriter(fh, fieldnames=list(tier_rows[0].keys())) + w.writeheader() + w.writerows(tier_rows) + + print("\n== cells ==") + for c in cell_rows: + print( + f"{c['model']:<24} {c['length']:>4} {c['condition']:<14} n={c['n']:<4} EM={c['em']:.3f} " + f"F1={c['f1']:.3f} ±{c['f1_se']:.3f} abst={c['abstained']:<3} ${c['cost_usd']:.3f}" + ) + print("\n== pos_rel terciles (tier = where the answer lives) ==") + for t in tier_rows: + print( + f"{t['model']:<24} {t['length']:>4} {t['condition']:<14} {t['tier']} n={t['n']:<3} " + f"EM={t['em']:.3f} F1={t['f1']:.3f} ±{t['f1_se']:.3f} abst={t['abstained']}" + ) + print(f"\nresults -> {out_dir}/records.jsonl, matrix.csv, terciles.csv, summary.json") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp12_arbitrage.py b/packages/snapcompact/research/exp12_arbitrage.py new file mode 100644 index 000000000..51f382789 --- /dev/null +++ b/packages/snapcompact/research/exp12_arbitrage.py @@ -0,0 +1,320 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp12: context-window arbitrage measurement. + +(a) chars per input token: optical pages vs plain text (mined from optimal-* runs + + targeted probes at image sizes 1568/1024/768); +(b) TPM accounting: do image tokens dodge text-token throttles? (rate-limit headers + captured around text-heavy vs image-heavy requests); +(c) risk: break-even repricing multiple at which the arbitrage dies. + +Measurement + writeup, no F1 chase. Outputs results/exp12-arbitrage/ +{measurements.json, probes.json, report.md}. +""" + +import argparse +import base64 +import json +import sys +import time +import urllib.error +import urllib.request +from pathlib import Path + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import capacity, render # noqa: E402 +from providers import load_env_key # noqa: E402 +from run import CACHE, FONTS, RESULTS, TEXT_CHUNK, sha8 # noqa: E402 + +OUT = RESULTS / "exp12-arbitrage" +SIZES = (1568, 1024, 768) +PROBE_MODELS = ("gpt-5.5", "google/gemini-3.5-flash") +PRICES = {"gpt-5.5": (2.0, 16.0), "google/gemini-3.5-flash": (0.6, 4.0)} +INSTR = "Reply with exactly: OK" + +# ---------------------------------------------------------------- part A: mine + +MINE_DIRS = ("optimal-combined", "optimal-gpt55", "optimal-gemini", "optimal-fable", + "optimal-opus", "optimal-kimi", "optimal-glm") + + +def cond_budget(cond: str) -> int | None: + """chars per chunk for a condition; None when chars/token is undefined (summaries).""" + if cond == "text": + return TEXT_CHUNK + if cond.startswith("img-"): + font = cond.split("-")[1] + return capacity(FONTS[font], 1568)[2] + return None # compact / handoff carry a generated summary, not the raw chars + + +def mine() -> tuple[list[dict], dict]: + """Per (model, cond): sum carrier chars and total qa input tokens over chunks. + + Also returns per-chunk (chars, tok) detail for the carrier estimation in derive(). + """ + flows = {} + paras = squad.load_paragraphs(CACHE) + for length in (50, 150, 250): + flows[length] = squad.build_flow(paras[:length])[0] + + seen: set[tuple] = set() + agg: dict[tuple[str, str], dict] = {} + detail: dict[str, dict[str, list]] = {} + for d in MINE_DIRS: + path = RESULTS / d / "records.jsonl" + if not path.exists(): + continue + for line in path.read_text().splitlines(): + r = json.loads(line) + if "usage" not in r: + continue + budget = cond_budget(r["cond"]) + if budget is None or r["length"] not in flows: + continue + key = (r["model"], r["cond"], r["length"], r["chunk"]) + if key in seen: + continue # combined is a merge of the per-model dirs + seen.add(key) + qa = next((u for u in r["usage"] if u["phase"] == "qa"), None) + if qa is None: + continue + tok = qa["in"] + qa["cache_r"] + qa["cache_w"] + chars = min(r["chunk"] + budget, len(flows[r["length"]])) - r["chunk"] + cell = agg.setdefault((r["model"], r["cond"]), {"chars": 0, "tok_in": 0, "chunks": 0}) + cell["chars"] += chars + cell["tok_in"] += tok + cell["chunks"] += 1 + detail.setdefault(r["model"], {}).setdefault(r["cond"], []).append((chars, tok)) + rows = [] + for (model, cond), c in sorted(agg.items()): + rows.append({ + "model": model, "cond": cond, "chunks": c["chunks"], "chars": c["chars"], + "tok_in_total": c["tok_in"], + "chars_per_tok": round(c["chars"] / c["tok_in"], 3), + }) + return rows, detail + + +# ------------------------------------------------------------- part B: probes +# Own POST so we can read rate-limit headers (providers._post discards them). + +OPENAI_URL = "https://api.openai.com/v1/responses" +OPENROUTER_URL = "https://openrouter.ai/api/v1/chat/completions" +RL_PREFIXES = ("x-ratelimit", "ratelimit", "retry-after") + + +def post_h(url: str, body: dict, headers: dict, retries: int = 4) -> tuple[dict, dict]: + payload = json.dumps(body).encode() + req = urllib.request.Request(url, data=payload, headers={"content-type": "application/json", **headers}) + for attempt in range(retries + 1): + try: + with urllib.request.urlopen(req, timeout=600) as resp: + rl = {k.lower(): v for k, v in resp.headers.items() if k.lower().startswith(RL_PREFIXES)} + return json.load(resp), rl + except urllib.error.HTTPError as err: + detail = err.read().decode(errors="replace")[:300] + if err.code in (408, 429, 500, 502, 503, 529) and attempt < retries: + wait = 2.0 * 2**attempt + print(f" HTTP {err.code}, retrying in {wait:.0f}s: {detail[:120]}") + time.sleep(wait) + continue + raise SystemExit(f"API error {err.code} ({url}): {detail}") from err + raise AssertionError("unreachable") + + +def png_b64(path: Path) -> str: + return base64.b64encode(path.read_bytes()).decode() + + +def probe_call(model: str, keys: dict, blocks: list[dict]) -> tuple[dict, dict]: + """One user message of blocks -> (normalized usage, rate-limit headers).""" + if model.startswith("gpt-"): + content = [] + for b in blocks: + if "text" in b: + content.append({"type": "input_text", "text": b["text"]}) + else: + content.append({"type": "input_image", + "image_url": f"data:image/png;base64,{png_b64(b['image_path'])}", + "detail": "original"}) + body = {"model": model, "input": [{"role": "user", "content": content}], + "max_output_tokens": 512, "store": False} + out, rl = post_h(OPENAI_URL, body, {"authorization": f"Bearer {keys['openai']}"}) + u = out.get("usage", {}) + cached = (u.get("input_tokens_details") or {}).get("cached_tokens", 0) + usage = {"in": u.get("input_tokens", 0), "cached": cached, "out": u.get("output_tokens", 0)} + return usage, rl + content = [] + for b in blocks: + if "text" in b: + content.append({"type": "text", "text": b["text"]}) + else: + content.append({"type": "image_url", + "image_url": {"url": f"data:image/png;base64,{png_b64(b['image_path'])}"}}) + body = {"model": model, "messages": [{"role": "user", "content": content}], "max_tokens": 512} + out, rl = post_h(OPENROUTER_URL, body, {"authorization": f"Bearer {keys['openrouter']}"}) + u = out.get("usage", {}) + usage = {"in": u.get("prompt_tokens", 0), + "cached": (u.get("prompt_tokens_details") or {}).get("cached_tokens", 0), + "out": u.get("completion_tokens", 0)} + return usage, rl + + +def probe_pngs(flow: str) -> dict[int, Path]: + pngs = {} + for size in SIZES: + cols, rows, cap = capacity(FONTS["6x10"], size) + text = flow[:cap] + png = CACHE / f"exp12-6x10-sent-{size}-{sha8(text, str(size))}.png" + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(".tmp.png") + render(text, FONTS["6x10"], CACHE, size, "sent").save(tmp) + tmp.replace(png) + pngs[size] = png + return pngs + + +def run_probes(keys: dict, flow: str) -> dict: + """Per model: instr-only, full-page text, and 6x10 pages at SIZES. + + Sequence is deliberate (small, text-heavy, small, image-heavy ...) so the + rate-limit header trail brackets each heavy request with a light one. + """ + pngs = probe_pngs(flow) + page_text = flow[:TEXT_CHUNK] + out: dict = {"page_chars": len(page_text), "models": {}} + for model in PROBE_MODELS: + steps = [("overhead-1", [{"text": INSTR}]), + ("text-page", [{"text": INSTR}, {"text": page_text}]), + ("overhead-2", [{"text": INSTR}])] + steps += [(f"img-{s}", [{"text": INSTR}, {"image_path": pngs[s]}]) for s in SIZES] + rows = [] + for name, blocks in steps: + usage, rl = probe_call(model, keys, blocks) + row = {"step": name, "usage": usage, "ratelimit": rl, "t": time.time()} + rows.append(row) + print(f" {model:>24} {name:<11} in={usage['in']:>6} (cached={usage['cached']}) " + f"out={usage['out']:>5} rl-remaining-tokens={rl.get('x-ratelimit-remaining-tokens', '-')}") + out["models"][model] = rows + return out + + +# --------------------------------------------------------- part C: derivation + + +def estimate_carriers(detail: dict, per_model_probed: dict) -> dict: + """Carrier-only chars/token for the non-probed models, from mined per-chunk data. + + Per-chunk total = carrier + QA overhead (prompt + question block). Overhead is + estimated as the mean (chunk_total - probe-measured carrier) over the two probed + models -- the question blocks are identical across models, tokenizers differ by + only a few %. Validation: applying the same estimate back to the probed models + reproduces their probe-measured chars/text-token within ~2%. + """ + overheads = [] + for model, d in per_model_probed.items(): + carrier = d["images"][1568]["image_tokens"] + overheads += [tok - carrier for _, tok in detail[model]["img-6x10-sent"]] + overhead = sum(overheads) / len(overheads) + page = capacity(FONTS["6x10"], 1568)[2] + est = {} + for model, conds in detail.items(): + if "img-6x10-sent" not in conds or "text" not in conds: + continue + img_rows, text_rows = conds["img-6x10-sent"], conds["text"] + img_tok = sum(t for _, t in img_rows) / len(img_rows) - overhead + text_chars = sum(c for c, _ in text_rows) + text_tok = sum(t for _, t in text_rows) - overhead * len(text_rows) + cpt_img, cpt_text = page / img_tok, text_chars / text_tok + est[model] = { + "est_image_tokens_per_page": round(img_tok), + "est_chars_per_img_tok": round(cpt_img, 3), + "est_chars_per_text_tok": round(cpt_text, 3), + "est_window_stretch": round(cpt_img / cpt_text, 3), + "probed": model in per_model_probed, + } + return {"qa_overhead_tokens_est": round(overhead, 1), "models": est} + + +def derive(mined: list[dict], detail: dict, probes: dict) -> dict: + page_chars = probes["page_chars"] + caps = {s: capacity(FONTS["6x10"], s)[2] for s in SIZES} + per_model = {} + for model, rows in probes["models"].items(): + by = {r["step"]: r["usage"] for r in rows} + overhead = min(by["overhead-1"]["in"], by["overhead-2"]["in"]) + text_tok = by["text-page"]["in"] - overhead + img = {} + for s in SIZES: + itok = by[f"img-{s}"]["in"] - overhead + img[s] = {"image_tokens": itok, "page_chars": caps[s], + "chars_per_img_tok": round(caps[s] / itok, 3), + "tok_per_megapixel": round(itok / (s * s / 1e6), 1)} + cpt_text = page_chars / text_tok + cpt_img = img[1568]["chars_per_img_tok"] + stretch = cpt_img / cpt_text + p_in = PRICES[model][0] + per_model[model] = { + "overhead_tokens": overhead, + "text_tokens_per_page": text_tok, + "chars_per_text_tok": round(cpt_text, 3), + "images": img, + "window_stretch_6x10_1568": round(stretch, 3), + "chars_in_200k_window": {"text": int(200_000 * cpt_text), "img_6x10_1568": int(200_000 * cpt_img)}, + "breakeven_img_token_multiple": round(stretch, 3), + "input_cost_per_mchar": {"text": round(p_in / cpt_text, 4), "img_6x10_1568": round(p_in / cpt_img, 4)}, + } + return {"mined": mined, "probes": probes, "derived": per_model, + "carrier_estimates": estimate_carriers(detail, per_model)} + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--fresh", action="store_true", help="re-run API probes even if probes.json exists") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + OUT.mkdir(parents=True, exist_ok=True) + CACHE.mkdir(exist_ok=True) + + mined, detail = mine() + print(f"mined {len(mined)} (model, cond) cells from {', '.join(MINE_DIRS)}") + for r in mined: + print(f" {r['model']:>24} {r['cond']:<18} chunks={r['chunks']:>2} chars={r['chars']:>7} " + f"tok={r['tok_in_total']:>7} chars/tok={r['chars_per_tok']:>7.3f}") + + probes_path = OUT / "probes.json" + if probes_path.exists() and not args.fresh: + probes = json.loads(probes_path.read_text()) + print("reusing probes.json (pass --fresh to re-run)") + else: + keys = {"openai": load_env_key("OPENAI_API_KEY", args.env), + "openrouter": load_env_key("OPENROUTER_API_KEY", args.env)} + # 150 paragraphs -> flow ~90k chars, so every probe page (incl. 1568px / 40716 + # chars) is completely full; image token cost is content-independent anyway + # (verified: identical tok/megapixel at three different fill ratios). + flow = squad.build_flow(squad.load_paragraphs(CACHE)[:150])[0] + probes = run_probes(keys, flow) + tmp = probes_path.with_suffix(".tmp.json") + tmp.write_text(json.dumps(probes, indent=1)) + tmp.replace(probes_path) + + measurements = derive(mined, detail, probes) + tmp = (OUT / "measurements.json").with_suffix(".tmp.json") + tmp.write_text(json.dumps(measurements, indent=1)) + tmp.replace(OUT / "measurements.json") + print(f"\nwrote {OUT}/measurements.json") + for model, d in measurements["derived"].items(): + print(f"{model}: text {d['chars_per_text_tok']} c/t | img-1568 " + f"{d['images'][1568 if 1568 in d['images'] else '1568']['chars_per_img_tok']} c/t | " + f"stretch {d['window_stretch_6x10_1568']}x | breakeven {d['breakeven_img_token_multiple']}x") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp13_extractive.py b/packages/snapcompact/research/exp13_extractive.py new file mode 100644 index 000000000..d23db7221 --- /dev/null +++ b/packages/snapcompact/research/exp13_extractive.py @@ -0,0 +1,255 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp13: extractive compaction baseline ("copy the load-bearing sentences verbatim"). + +Fixes the strawman in the text-summary family: the existing `compact`/`handoff` +baselines ask for a *narrative* summary, which is hostile to extractive QA +(gemini abstains UNREADABLE on it, F1=0). Here the compaction prompt instead +asks for verbatim spans up to a character budget sized to match the optical +carrier (~2000 tokens ~= 8000 chars per 40716-char chunk; img-6x10 page costs +~1664 input tok on gemini / ~3396 on gpt-5.5). + +Pipeline per chunk (mirrors final.py's `compact` branch): + session_frame(chunk) + exp13-extract.md -> cached extraction (tag exp13-extract) + qa-text.md(context=extraction) + questions -> cached QA (tag exp13-qa) +Extraction usage is counted in the cell cost, like the summarize phase. + +Extra column vs final.py: gold_survival = fraction of questions whose gold +answer literally survives in the extraction (normalized string containment) — +the recall ceiling of the method, separable from QA ability. + +Run from the snapcompact dir: uv run exp13_extractive.py +""" + +import argparse +import csv +import json +import sys +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from squad import _normalize # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, QA_CACHE, RESULTS, TEXT_CHUNK, load_prompt, sha8 # noqa: E402 + +MODELS = { + "gpt-5.5": (2.0, 16.0), + "google/gemini-3.5-flash": (0.6, 4.0), +} +COND = "extract" +ACK = "Noted. I have read the passages and will keep them in mind." + + +def cached(model: str, tag: str, payload: object, fn, fresh: bool) -> dict: + """Disk-cache `fn() -> dict` keyed by (model, tag, payload). Truncations are not cached.""" + key = sha8(model, tag, json.dumps(payload, sort_keys=True, default=str)) + path = QA_CACHE / f"{key}.json" + if path.exists() and not fresh: + hit = json.loads(path.read_text()) + if hit.get("stop") != "max_tokens": + return hit + out = fn() + if out.get("stop") == "max_tokens": + print(f" WARN truncated, not cached: {model} {tag} {key}") + else: + path.write_text(json.dumps(out)) + return out + + +def session_frame(chunk_text: str) -> list[dict]: + return [ + {"role": "user", "content": [{"text": load_prompt("session-frame.md").format(context=chunk_text)}]}, + {"role": "assistant", "content": [{"text": ACK}]}, + ] + + +def gold_survives(golds: list[str], extraction_norm: str) -> bool: + return any(_normalize(g) in extraction_norm for g in golds) + + +def run_cell_chunk(model: str, start: int, end: int, ctx: dict) -> list[dict]: + """One (model, chunk) unit: extract verbatim spans, QA over the extraction, score.""" + args, flow, paras, offsets, keys = ctx["args"], ctx["flow"], ctx["paras"], ctx["offsets"], ctx["keys"] + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + chunk_text = flow[start:end] + q_block = "\n".join(f"{i + 1}. {q['q']}" for i, q in enumerate(questions)) + usage_rows: list[tuple[str, dict]] = [] + + extract_prompt = load_prompt("exp13-extract.md").format(budget=args.budget) + gen = cached( + model, "exp13-extract", {"chunk": chunk_text, "budget": args.budget, "effort": args.extract_effort}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete( + keys, model, + session_frame(chunk_text) + [{"role": "user", "content": [{"text": extract_prompt}]}], + max_tokens=args.extract_max_tokens, + effort=args.extract_effort, + ), + ) + ), + args.fresh, + ) + usage_rows.append(("extract", gen["usage"])) + extraction = gen["text"] + extraction_norm = _normalize(extraction) + + messages = [ + { + "role": "user", + "content": [{"text": load_prompt("qa-text.md").format(context=extraction)}, {"text": q_block}], + } + ] + qa = cached( + model, "exp13-qa", {"messages": messages}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete(keys, model, messages, max_tokens=args.max_tokens), + ) + ), + args.fresh, + ) + usage_rows.append(("qa", qa["usage"])) + answers = squad.parse_numbered(qa["text"], len(questions)) + records = [] + for q, a in zip(questions, answers): + records.append( + { + "model": model, + "length": ctx["length"], + "cond": COND, + "chunk": start, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + "gold_survived": gold_survives(q["golds"], extraction_norm), + } + ) + records[0]["usage"] = [{"phase": p, **u} for p, u in usage_rows] + records[0]["extraction_chars"] = len(extraction) + return records + + +def aggregate(records: list[dict], price_in: float, price_out: float) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean_f1 = sum(f1s) / n + se = (sum((x - mean_f1) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + us = [u for r in records if "usage" in r for u in r["usage"]] + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + cost_in = (tok["in"] + 1.25 * tok["cache_w"] + 0.1 * tok["cache_r"]) / 1e6 * price_in + cost_out = tok["out"] / 1e6 * price_out + return { + "n": n, + "em": sum(r["em"] for r in records) / n, + "f1": mean_f1, + "f1_se": se, + "abstained": sum(r["abstained"] for r in records), + "gold_survival": sum(r["gold_survived"] for r in records) / n, + "extraction_chars": sum(r.get("extraction_chars", 0) for r in records), + **{f"tok_{k}": v for k, v in tok.items()}, + "cost_in_usd": round(cost_in, 4), + "cost_out_usd": round(cost_out, 4), + "cost_usd": round(cost_in + cost_out, 4), + } + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--models", default=",".join(MODELS)) + ap.add_argument("--lengths", default="50,150,250") + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--budget", type=int, default=8000, help="max extraction chars per chunk") + ap.add_argument("--max-tokens", type=int, default=32768, help="QA max tokens") + ap.add_argument("--extract-max-tokens", type=int, default=16384, help="extraction max tokens (budget+slack)") + ap.add_argument( + "--extract-effort", default="low", + help="reasoning effort for the extraction call only; verbatim copying needs no deliberation " + "(default-effort gemini burns ~16k reasoning tokens verifying quotes and truncates)", + ) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--env", default="~/.env") + ap.add_argument("--out", default="exp13-extractive") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + out_dir = RESULTS / args.out + out_dir.mkdir(parents=True, exist_ok=True) + + models = [m.strip() for m in args.models.split(",") if m.strip()] + lengths = [int(x) for x in args.lengths.split(",") if x.strip()] + for m in models: + if m not in MODELS: + raise SystemExit(f"unknown model {m}") + + keys = { + "openai": load_env_key("OPENAI_API_KEY", args.env), + "openrouter": load_env_key("OPENROUTER_API_KEY", args.env), + } + + all_paras = squad.load_paragraphs(CACHE) + tasks = [] + for length in lengths: + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + ctx = {"args": args, "flow": flow, "paras": paras, "offsets": offsets, "keys": keys, "length": length} + for model in models: + for start in range(0, len(flow), TEXT_CHUNK): + tasks.append((model, start, min(start + TEXT_CHUNK, len(flow)), ctx)) + print(f"grid: {len(models)} models x {len(lengths)} lengths x 1 condition = {len(tasks)} chunk tasks") + + records: list[dict] = [] + done = 0 + with ThreadPoolExecutor(args.workers) as pool: + futures = [pool.submit(run_cell_chunk, m, s, e, ctx) for m, s, e, ctx in tasks] + for fut in futures: + records.extend(fut.result()) + done += 1 + print(f" {done}/{len(tasks)} tasks", flush=True) + + with (out_dir / "records.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + + cells = [] + for model in models: + for length in lengths: + sub = [r for r in records if r["model"] == model and r["length"] == length] + if not sub: + continue + cells.append({"model": model, "length": length, "condition": COND, **aggregate(sub, *MODELS[model])}) + (out_dir / "summary.json").write_text(json.dumps({"args": vars(args), "cells": cells}, indent=1)) + with (out_dir / "matrix.csv").open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=list(cells[0].keys())) + writer.writeheader() + writer.writerows(cells) + + for c in cells: + print( + f"{c['model']:>24} len={c['length']:<4} f1={c['f1']:.3f}±{c['f1_se']:.3f} em={c['em']:.3f} " + f"survival={c['gold_survival']:.3f} abst={c['abstained']}/{c['n']} cost=${c['cost_usd']:.4f}" + ) + print(f"\ndataset -> {out_dir}/records.jsonl, matrix.csv, summary.json") + + +if __name__ == "__main__": + main() + + diff --git a/packages/snapcompact/research/exp14_bestgpt.py b/packages/snapcompact/research/exp14_bestgpt.py new file mode 100644 index 000000000..9ed35b88d --- /dev/null +++ b/packages/snapcompact/research/exp14_bestgpt.py @@ -0,0 +1,475 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp14: best-of-round-1 combination for gpt-5.5. + +Combines the validated levers from round 1: + - 8x13 glyphs on a patch-aligned 8x16 cell (exp01: 8on16-sent .918@150) + - two-column document layout (exp04: +F1 / -cost / -read-tax at 6x10) + - per-model variant: gpt-5.5 prefers bw (exp10), sent is the runner-up + +Conditions (gpt-5.5 only, 1568px): + img-doc-8on16-bw doc layout, near-black ink (the combination) + img-doc-8on16-sent doc layout, sentence-hue glyphs (variant probe) + img-8on16-bw plain grid, missing round-1 cell (8on16 ran only as sent) + +Phased: screen all three at length 150, confirm the winner at 50/250, +optional effort=none probe at 50. Records merge across runs (cells keyed by +model/length/condition/effort are replaced when re-run, kept otherwise). + +Usage: + uv run exp14_bestgpt.py --render-only # capacity + sample PNGs + uv run exp14_bestgpt.py # screen @150 (default cells) + uv run exp14_bestgpt.py --cells img-doc-8on16-bw@50,img-doc-8on16-bw@250 + uv run exp14_bestgpt.py --cells img-doc-8on16-bw@50 --effort none + uv run exp14_bestgpt.py --report # re-aggregate, no API +""" + +import argparse +import csv +import json +import sys +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +from PIL import Image + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import _DARK, FontCfg, capacity, ensure_font, parse_bdf, render # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 + +EXP = "exp14" +OUT_DIR = RESULTS / f"{EXP}-bestgpt" +MODEL = "gpt-5.5" +PRICE_IN, PRICE_OUT = 2.0, 16.0 +FONT = FontCfg("8on16", "8x13", 8, 16) # exp01 winner: 8x13 glyphs, 16px patch-aligned pitch +GUTTER = 3 # char cells between doc columns (as exp04) +SCREEN_CELLS = "img-doc-8on16-bw@150,img-doc-8on16-sent@150,img-8on16-bw@150" +_WHITE = (255, 255, 255) +_BLACK = (0, 0, 0) +_INK = (24, 24, 24) # exp04 body ink + + +def cached(model: str, tag: str, payload: object, fn, fresh: bool) -> dict: + """Disk-cache `fn() -> dict` keyed by (model, tag, payload). Truncations are not cached.""" + key = sha8(model, tag, json.dumps(payload, sort_keys=True, default=str)) + path = QA_CACHE / f"{key}.json" + if path.exists() and not fresh: + hit = json.loads(path.read_text()) + if hit.get("stop") != "max_tokens": + return hit + out = fn() + if out.get("stop") == "max_tokens": + print(f" WARN truncated, not cached: {model} {tag} {key}") + else: + path.write_text(json.dumps(out)) + return out + + +# --- document layout (ported from exp04, parameterized for FONT) ------------ + + +def wrap(text: str, width: int) -> list[str]: + """Greedy word-wrap, no mid-word breaks (hard split only for width+ words).""" + lines: list[str] = [] + cur = "" + for word in text.split(): + while len(word) > width: # pathological; never hit on SQuAD prose + if cur: + lines.append(cur) + cur = "" + lines.append(word[:width]) + word = word[width:] + if not cur: + cur = word + elif len(cur) + 1 + len(word) <= width: + cur += " " + word + else: + lines.append(cur) + cur = word + if cur: + lines.append(cur) + return lines + + +def layout_page(paras: list[dict], col_w: int) -> list[dict]: + """Typeset paragraphs into lines: [{kind: heading|body|blank, text}]. + + Article title changes become headings (repeated at the top of a page even + when the article continues, since each page is read in isolation). + Paragraphs are separated by one blank line. + """ + lines: list[dict] = [] + prev_title = None + for p in paras: + if p["title"] != prev_title: + if lines: + lines.append({"kind": "blank", "text": ""}) + for hl in wrap(p["title"].replace("_", " ").upper(), col_w): + lines.append({"kind": "heading", "text": hl}) + prev_title = p["title"] + elif lines: + lines.append({"kind": "blank", "text": ""}) + for bl in wrap(p["ctx"], col_w): + lines.append({"kind": "body", "text": bl}) + return lines + + +def pack_pages(paras: list[dict], col_w: int, max_lines: int) -> list[tuple[int, int]]: + """Greedy paragraph-aligned packing: [(i, j)] para ranges, one per page.""" + pages = [] + i = 0 + while i < len(paras): + j = i + 1 + while j < len(paras) and len(layout_page(paras[i : j + 1], col_w)) <= max_lines: + j += 1 + pages.append((i, j)) + i = j + return pages + + +def _sentence_colors(lines: list[dict]) -> list[list[tuple[int, int, int]]]: + """Per-line per-char glyph color cycling hue per sentence across the page.""" + joined = "\n".join(ln["text"] for ln in lines) + idx, out_idx = 0, [] + for i, ch in enumerate(joined): + out_idx.append(idx) + if ch in ".!?" and i + 1 < len(joined) and joined[i + 1] in " \n": + idx += 1 + colors, pos = [], 0 + for ln in lines: + n = len(ln["text"]) + colors.append([_DARK[out_idx[pos + k] % 6] for k in range(n)]) + pos += n + 1 # the joining newline + return colors + + +def render_doc(lines: list[dict], size: int, variant: str, cache: Path) -> Image.Image: + """Two-column page: left column rows top-to-bottom, then right column.""" + glyphs, font_ascent = parse_bdf(ensure_font(FONT, cache)) + ascent = FONT.ascent if FONT.ascent is not None else font_ascent + cols, rows, _ = capacity(FONT, size) + col_w = (cols - GUTTER) // 2 + sent_colors = _sentence_colors(lines) if variant == "sent" else None + img = Image.new("RGB", (size, size), _WHITE) + px = img.load() + for li, ln in enumerate(lines): + column, row = divmod(li, rows) + if column > 1: + break # overflow guard; pack_pages should prevent this + x_origin = column * (col_w + GUTTER) * FONT.adv + y0 = row * FONT.pitch + for ci, ch in enumerate(ln["text"]): + glyph = glyphs.get(ord(ch)) + if glyph is None: + continue + if ln["kind"] == "heading": + fg = _BLACK + elif sent_colors is not None: + fg = sent_colors[li][ci] + else: + fg = _INK + w, h, xoff, yoff = glyph["bbx"] + top = y0 + ascent - h - yoff + shift = 0x80 if w <= 8 else 0x8000 + strikes = (0, 1) if ln["kind"] == "heading" else (0,) + for dx in strikes: + for r, bits in enumerate(glyph["rows"]): + y = top + r + if not 0 <= y < size: + continue + for b in range(w): + if bits & (shift >> b): + x = x_origin + ci * FONT.adv + xoff + b + dx + if 0 <= x < size: + px[x, y] = fg + return img + + +# --- runner ----------------------------------------------------------------- + + +def atomic_save(img: Image.Image, png: Path) -> None: + tmp = png.with_suffix(".tmp.png") + img.save(tmp) + tmp.replace(png) + + +def qa_call(messages: list[dict], questions: list[dict], length: int, cond: str, + start: int, ctx: dict) -> list[dict]: + """One QA call + scoring; shared by doc and grid paths.""" + args, keys = ctx["args"], ctx["keys"] + qa = cached( + MODEL, f"{EXP}-qa", {"messages": messages, "effort": args.effort}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete(keys, MODEL, messages, max_tokens=args.max_tokens, effort=args.effort), + ) + ), + args.fresh, + ) + answers = squad.parse_numbered(qa["text"], len(questions)) + records = [] + for q, a in zip(questions, answers): + records.append( + { + "model": MODEL, + "length": length, + "cond": cond, + "effort": args.effort, + "chunk": start, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + ) + records[0]["usage"] = [{"phase": "qa", **qa["usage"]}] + return records + + +def run_doc_page(cond: str, length: int, page: tuple[int, int], ctx: dict) -> list[dict]: + args, paras, offsets = ctx["args"], ctx["paras"], ctx["offsets"] + i, j = page + start = offsets[i] + end = offsets[j - 1] + len(paras[j - 1]["ctx"]) + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + variant = cond.removeprefix("img-doc-8on16-") + lines = ctx["lines"][page] + page_key = sha8(cond, json.dumps([(p["title"], p["ctx"]) for p in paras[i:j]]), str(args.size)) + png = CACHE / f"{EXP}-doc-{variant}-{page_key}.png" + if not png.exists() or png.stat().st_size == 0: + atomic_save(render_doc(lines, args.size, variant, CACHE), png) + cols, rows, _ = capacity(FONT, args.size) + col_w = (cols - GUTTER) // 2 + q_block = "\n".join(f"{k + 1}. {q['q']}" for k, q in enumerate(questions)) + messages = [ + { + "role": "user", + "content": [ + {"text": load_prompt("exp04-qa-image.md").format(col_w=col_w, rows=rows)}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + return qa_call(messages, questions, length, cond, start, ctx) + + +def run_grid_chunk(cond: str, length: int, start: int, end: int, ctx: dict) -> list[dict]: + args, flow, paras, offsets = ctx["args"], ctx["flow"], ctx["paras"], ctx["offsets"] + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + chunk_text = flow[start:end] + variant = cond.removeprefix("img-8on16-") + png = CACHE / f"{EXP}-8on16-{variant}-{sha8(chunk_text, str(args.size))}.png" + if not png.exists() or png.stat().st_size == 0: + atomic_save(render(chunk_text, FONT, CACHE, args.size, variant), png) + cols, rows, _ = capacity(FONT, args.size) + q_block = "\n".join(f"{k + 1}. {q['q']}" for k, q in enumerate(questions)) + messages = [ + { + "role": "user", + "content": [ + {"text": load_prompt("qa-image.md").format(cols=cols, rows=rows)}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + return qa_call(messages, questions, length, cond, start, ctx) + + +def aggregate(records: list[dict]) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean_f1 = sum(f1s) / n + se = (sum((x - mean_f1) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + us = [u for r in records if "usage" in r for u in r["usage"]] + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + cost_in = (tok["in"] + 1.25 * tok["cache_w"] + 0.1 * tok["cache_r"]) / 1e6 * PRICE_IN + cost_out = tok["out"] / 1e6 * PRICE_OUT + return { + "n": n, + "em": sum(r["em"] for r in records) / n, + "f1": mean_f1, + "f1_se": se, + "abstained": sum(r["abstained"] for r in records), + **{f"tok_{k}": v for k, v in tok.items()}, + "cost_in_usd": round(cost_in, 4), + "cost_out_usd": round(cost_out, 4), + "cost_usd": round(cost_in + cost_out, 4), + } + + +def cell_label(cond: str, effort: str | None) -> str: + return f"{cond}+eff-{effort}" if effort else cond + + +def write_outputs(records: list[dict], capacity_stats: dict, args_dict: dict) -> list[dict]: + with (OUT_DIR / "records.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + cell_keys = sorted({(r["length"], r["cond"], r.get("effort")) for r in records}, + key=lambda k: (k[0], k[1], k[2] or "")) + cells = [] + for length, cond, effort in cell_keys: + sub = [r for r in records if r["length"] == length and r["cond"] == cond and r.get("effort") == effort] + cells.append({"model": MODEL, "length": length, "condition": cell_label(cond, effort), **aggregate(sub)}) + (OUT_DIR / "summary.json").write_text( + json.dumps({"args": args_dict, "capacity": capacity_stats, "cells": cells}, indent=1) + ) + with (OUT_DIR / "matrix.csv").open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=list(cells[0].keys())) + writer.writeheader() + writer.writerows(cells) + return cells + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--cells", default=SCREEN_CELLS, help="comma list of cond@length") + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--effort", default=None) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--render-only", action="store_true") + ap.add_argument("--report", action="store_true", help="re-aggregate existing records, no API") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + OUT_DIR.mkdir(parents=True, exist_ok=True) + + cols, rows, grid_cap = capacity(FONT, args.size) + col_w = (cols - GUTTER) // 2 + max_lines = 2 * rows + print(f"8on16 @ {args.size}px: grid {cols}x{rows} = {grid_cap} chars; " + f"doc 2 x {col_w} cols + gutter {GUTTER}, {max_lines} line slots") + + rec_path = OUT_DIR / "records.jsonl" + existing: list[dict] = [] + if rec_path.exists(): + existing = [json.loads(ln) for ln in rec_path.read_text().splitlines() if ln.strip()] + cap_path = OUT_DIR / "capacity.json" + capacity_stats: dict = json.loads(cap_path.read_text()) if cap_path.exists() else {} + + if args.report: + cells = write_outputs(existing, capacity_stats, vars(args)) + for c in cells: + print(f"len {c['length']:<4} {c['condition']:<28} n={c['n']:<4} EM {c['em']:.3f} " + f"F1 {c['f1']:.3f} ±{c['f1_se']:.3f} ${c['cost_usd']:.3f} " + f"out={c['tok_out']} rsn={c['tok_reasoning']}") + return + + cell_specs = [] + for spec in args.cells.split(","): + spec = spec.strip() + if not spec: + continue + cond, _, ln = spec.partition("@") + cell_specs.append((cond, int(ln))) + lengths = sorted({ln for _, ln in cell_specs}) + + keys = {} + if not args.render_only: + keys["openai"] = load_env_key("OPENAI_API_KEY", args.env) + + all_paras = squad.load_paragraphs(CACHE) + tasks = [] + for length in lengths: + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + pages = pack_pages(paras, col_w, max_lines) + page_lines = {pg: layout_page(paras[pg[0] : pg[1]], col_w) for pg in pages} + page_chars = [offsets[j - 1] + len(paras[j - 1]["ctx"]) - offsets[i] for i, j in pages] + capacity_stats[str(length)] = { + "doc_pages": len(pages), + "mean_chars_page": round(sum(page_chars) / len(pages)), + "min_chars_page": min(page_chars), + "max_chars_page": max(page_chars), + "grid_chars_page": grid_cap, + "corpus_chars": len(flow), + "grid_pages": -(-len(flow) // grid_cap), + } + st = capacity_stats[str(length)] + print(f" len {length}: {st['doc_pages']} doc pages (mean {st['mean_chars_page']} chars, " + f"{round(100 * st['mean_chars_page'] / grid_cap)}% of grid {grid_cap}); " + f"grid {st['grid_pages']} pages; corpus {st['corpus_chars']}") + ctx = {"args": args, "paras": paras, "flow": flow, "offsets": offsets, "keys": keys, "lines": page_lines} + for cond, ln in cell_specs: + if ln != length: + continue + if cond.startswith("img-doc-"): + for pg in pages: + tasks.append(("doc", cond, length, pg, ctx)) + else: + for start in range(0, len(flow), grid_cap): + tasks.append(("grid", cond, length, (start, min(start + grid_cap, len(flow))), ctx)) + + cap_path.write_text(json.dumps(capacity_stats, indent=1)) + + if args.render_only: + for kind, cond, length, unit, ctx in tasks: + if unit[0] != 0 and (kind == "grid" or unit != list(ctx["lines"])[0]): + continue + if kind == "doc": + variant = cond.removeprefix("img-doc-8on16-") + i, j = unit + key = sha8(cond, json.dumps([(p["title"], p["ctx"]) for p in ctx["paras"][i:j]]), str(args.size)) + png = CACHE / f"{EXP}-doc-{variant}-{key}.png" + atomic_save(render_doc(ctx["lines"][unit], args.size, variant, CACHE), png) + else: + variant = cond.removeprefix("img-8on16-") + chunk_text = ctx["flow"][unit[0] : unit[1]] + png = CACHE / f"{EXP}-8on16-{variant}-{sha8(chunk_text, str(args.size))}.png" + atomic_save(render(chunk_text, FONT, CACHE, args.size, variant), png) + print(f" sample: {png}") + return + + print(f"{len(tasks)} page/chunk tasks") + new_records: list[dict] = [] + done = 0 + with ThreadPoolExecutor(args.workers) as pool: + futures = [] + for kind, cond, length, unit, ctx in tasks: + if kind == "doc": + futures.append(pool.submit(run_doc_page, cond, length, unit, ctx)) + else: + futures.append(pool.submit(run_grid_chunk, cond, length, unit[0], unit[1], ctx)) + for fut in futures: + new_records.extend(fut.result()) + done += 1 + print(f" {done}/{len(tasks)}", flush=True) + + # merge: drop existing records for cells just re-run, keep everything else + rerun = {(ln, cond, args.effort) for cond, ln in cell_specs} + kept = [r for r in existing if (r["length"], r["cond"], r.get("effort")) not in rerun] + records = kept + new_records + + cells = write_outputs(records, capacity_stats, vars(args)) + for c in cells: + print(f"len {c['length']:<4} {c['condition']:<28} n={c['n']:<4} EM {c['em']:.3f} " + f"F1 {c['f1']:.3f} ±{c['f1_se']:.3f} ${c['cost_usd']:.3f} " + f"out={c['tok_out']} rsn={c['tok_reasoning']}") + print(f"\n-> {OUT_DIR}/records.jsonl, matrix.csv, summary.json") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp15_bestgemini.py b/packages/snapcompact/research/exp15_bestgemini.py new file mode 100644 index 000000000..1119b32e3 --- /dev/null +++ b/packages/snapcompact/research/exp15_bestgemini.py @@ -0,0 +1,478 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp15: combined best optical profile for google/gemini-3.5-flash. + +Combines the round-1 validated levers: + - 8x13 glyphs on a patch-aligned 8x16 cell (exp01 `8on16`) + - two-column document layout with headings (exp04), parameterized for font/pitch + - gemini's winning variant `sent-dim` (exp10), plus `sent` runner-up + +Phase A (screen @150): img-doc-8on16-sent-dim, img-doc-8on16-sent, and the +missing round-1 grid cell img-8on16-sent-dim. +Phase B (confirm): screening winner at lengths 50 and 250. + +Bar (findings.md best known for gemini): .984@50 / .915@150 / .909@250. +Text ceiling: .989 / .898 / .918. +""" + +import argparse +import csv +import json +import os +import sys +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +from PIL import Image + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import _DARK, _DIMMED, _stopword_mask, FontCfg, capacity, parse_bdf, ensure_font, render # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 + +EXP = "exp15" +OUT_DIR = RESULTS / f"{EXP}-bestgemini" +MODEL = "google/gemini-3.5-flash" +PRICE = (0.6, 4.0) # $/M in, out +FONT = FontCfg("8on16", "8x13", 8, 16) # 8x13 glyphs, ViT-patch-aligned 16px pitch +GUTTER = 3 # char cells between doc columns (exp04) +SCREEN_CONDS = ("img-doc-8on16-sent-dim", "img-doc-8on16-sent", "img-8on16-sent-dim") +_WHITE = (255, 255, 255) +_BLACK = (0, 0, 0) + +# Best-known cells from local findings (for the printed delta column only). +BEST_KNOWN = {50: (0.984, 0.012), 150: (0.915, 0.019), 250: (0.909, 0.016)} +TEXT_CEILING = {50: 0.989, 150: 0.898, 250: 0.918} + + +def cached(model: str, payload: object, fn, fresh: bool) -> dict: + """Disk-cache `fn() -> dict` keyed by (model, exp15-qa, payload). Truncations not cached.""" + key = sha8(model, f"{EXP}-qa", json.dumps(payload, sort_keys=True, default=str)) + path = QA_CACHE / f"{key}.json" + if path.exists() and not fresh: + hit = json.loads(path.read_text()) + if hit.get("stop") != "max_tokens": + return hit + out = fn() + if out.get("stop") == "max_tokens": + print(f" WARN truncated, not cached: {model} {key}") + else: + tmp = path.with_suffix(f".{os.getpid()}.tmp") + tmp.write_text(json.dumps(out)) + tmp.replace(path) + return out + + +# --- document layout (ported from exp04, parameterized font) ---------------- + + +def wrap(text: str, width: int) -> list[str]: + """Greedy word-wrap, no mid-word breaks (hard split only for width+ words).""" + lines: list[str] = [] + cur = "" + for word in text.split(): + while len(word) > width: + if cur: + lines.append(cur) + cur = "" + lines.append(word[:width]) + word = word[width:] + if not cur: + cur = word + elif len(cur) + 1 + len(word) <= width: + cur += " " + word + else: + lines.append(cur) + cur = word + if cur: + lines.append(cur) + return lines + + +def layout_page(paras: list[dict], col_w: int) -> list[dict]: + """Typeset paragraphs into lines: [{kind: heading|body|blank, text}]. + + Headings repeat at the top of a page when an article continues, since + each page is read in isolation.""" + lines: list[dict] = [] + prev_title = None + for p in paras: + if p["title"] != prev_title: + if lines: + lines.append({"kind": "blank", "text": ""}) + for hl in wrap(p["title"].replace("_", " ").upper(), col_w): + lines.append({"kind": "heading", "text": hl}) + prev_title = p["title"] + elif lines: + lines.append({"kind": "blank", "text": ""}) + for bl in wrap(p["ctx"], col_w): + lines.append({"kind": "body", "text": bl}) + return lines + + +def pack_pages(paras: list[dict], col_w: int, max_lines: int) -> list[tuple[int, int]]: + """Greedy paragraph-aligned packing: [(i, j)] para ranges, one per page.""" + pages = [] + i = 0 + while i < len(paras): + j = i + 1 + while j < len(paras) and len(layout_page(paras[i : j + 1], col_w)) <= max_lines: + j += 1 + pages.append((i, j)) + i = j + return pages + + +def _doc_colors(lines: list[dict], variant: str) -> list[list[tuple[int, int, int]]]: + """Per-line per-char glyph colors: sentence hue cycle, optionally with the + stopword dim mask composed on top (sent-dim = exp04 sent + bdf dim).""" + joined = "\n".join(ln["text"] for ln in lines) + sidx, idx = [], 0 + for i, ch in enumerate(joined): + sidx.append(idx) + if ch in ".!?" and i + 1 < len(joined) and joined[i + 1] in " \n": + idx += 1 + dim = _stopword_mask(joined) if variant == "sent-dim" else None + colors, pos = [], 0 + for ln in lines: + n = len(ln["text"]) + colors.append( + [ + _DIMMED if dim is not None and dim[pos + k] else _DARK[sidx[pos + k] % 6] + for k in range(n) + ] + ) + pos += n + 1 # the joining newline + return colors + + +def render_doc(lines: list[dict], size: int, variant: str, cache: Path) -> Image.Image: + """Two-column page at FONT: left column top-to-bottom, then right column.""" + glyphs, font_ascent = parse_bdf(ensure_font(FONT, cache)) + ascent = FONT.ascent if FONT.ascent is not None else font_ascent + cols, rows, _ = capacity(FONT, size) + col_w = (cols - GUTTER) // 2 + colors = _doc_colors(lines, variant) + img = Image.new("RGB", (size, size), _WHITE) + px = img.load() + for li, ln in enumerate(lines): + column, row = divmod(li, rows) + if column > 1: + break # overflow guard; pack_pages should prevent this + x_origin = column * (col_w + GUTTER) * FONT.adv + y0 = row * FONT.pitch + for ci, ch in enumerate(ln["text"]): + glyph = glyphs.get(ord(ch)) + if glyph is None: + continue + fg = _BLACK if ln["kind"] == "heading" else colors[li][ci] + w, h, xoff, yoff = glyph["bbx"] + top = y0 + ascent - h - yoff + shift = 0x80 if w <= 8 else 0x8000 + strikes = (0, 1) if ln["kind"] == "heading" else (0,) + for dx in strikes: + for r, bits in enumerate(glyph["rows"]): + y = top + r + if not 0 <= y < size: + continue + for b in range(w): + if bits & (shift >> b): + x = x_origin + ci * FONT.adv + xoff + b + dx + if 0 <= x < size: + px[x, y] = fg + return img + + +# --- runners ----------------------------------------------------------------- + + +def _qa_call(model: str, messages: list[dict], questions: list[dict], ctx: dict, cond: str, start: int) -> list[dict]: + args, keys = ctx["args"], ctx["keys"] + qa = cached( + model, + {"messages": messages, "extra": None, "effort": args.effort}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete(keys, model, messages, max_tokens=args.max_tokens, effort=args.effort), + ) + ), + args.fresh, + ) + answers = squad.parse_numbered(qa["text"], len(questions)) + records = [] + for q, a in zip(questions, answers): + records.append( + { + "model": model, + "length": ctx["length"], + "cond": cond, + "chunk": start, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + ) + records[0]["usage"] = [{"phase": "qa", **qa["usage"]}] + return records + + +def run_doc_page(model: str, cond: str, page: tuple[int, int], ctx: dict) -> list[dict]: + args, paras, offsets = ctx["args"], ctx["paras"], ctx["offsets"] + i, j = page + start = offsets[i] + end = offsets[j - 1] + len(paras[j - 1]["ctx"]) + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + variant = cond.removeprefix("img-doc-8on16-") + lines = ctx["lines"][page] + page_key = sha8(cond, json.dumps([(p["title"], p["ctx"]) for p in paras[i:j]]), str(args.size)) + png = CACHE / f"{EXP}-doc-8on16-{variant}-{page_key}.png" + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(f".{os.getpid()}.tmp.png") + render_doc(lines, args.size, variant, CACHE).save(tmp) + tmp.replace(png) + cols, rows, _ = capacity(FONT, args.size) + col_w = (cols - GUTTER) // 2 + q_block = "\n".join(f"{k + 1}. {q['q']}" for k, q in enumerate(questions)) + messages = [ + { + "role": "user", + "content": [ + {"text": load_prompt("exp04-qa-image.md").format(col_w=col_w, rows=rows)}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + return _qa_call(model, messages, questions, ctx, cond, start) + + +def run_grid_chunk(model: str, cond: str, start: int, end: int, ctx: dict) -> list[dict]: + args, flow, paras, offsets = ctx["args"], ctx["flow"], ctx["paras"], ctx["offsets"] + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + chunk_text = flow[start:end] + variant = cond.removeprefix("img-8on16-") + png = CACHE / f"{EXP}-grid-8on16-{variant}-{sha8(chunk_text, str(args.size))}.png" + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(f".{os.getpid()}.tmp.png") + render(chunk_text, FONT, CACHE, args.size, variant).save(tmp) + tmp.replace(png) + cols, rows, _ = capacity(FONT, args.size) + q_block = "\n".join(f"{k + 1}. {q['q']}" for k, q in enumerate(questions)) + messages = [ + { + "role": "user", + "content": [ + {"text": load_prompt("qa-image.md").format(cols=cols, rows=rows)}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + return _qa_call(model, messages, questions, ctx, cond, start) + + +def aggregate(records: list[dict], price_in: float, price_out: float) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean_f1 = sum(f1s) / n + se = (sum((x - mean_f1) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + us = [u for r in records if "usage" in r for u in r["usage"]] + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + cost_in = (tok["in"] + 1.25 * tok["cache_w"] + 0.1 * tok["cache_r"]) / 1e6 * price_in + cost_out = tok["out"] / 1e6 * price_out + return { + "n": n, + "em": sum(r["em"] for r in records) / n, + "f1": mean_f1, + "f1_se": se, + "abstained": sum(r["abstained"] for r in records), + "calls": len(us), + **{f"tok_{k}": v for k, v in tok.items()}, + "cost_in_usd": round(cost_in, 4), + "cost_out_usd": round(cost_out, 4), + "cost_usd": round(cost_in + cost_out, 4), + } + + +class Runner: + def __init__(self, args, keys): + self.args = args + self.keys = keys + self.records: list[dict] = [] + self.done: set[tuple[int, str]] = set() + self.ctxs: dict[int, dict] = {} + self.all_paras = squad.load_paragraphs(CACHE) + self.capacity_stats: dict = {} + + def ctx(self, length: int) -> dict: + if length not in self.ctxs: + paras = self.all_paras[:length] + flow, offsets = squad.build_flow(paras) + cols, rows, grid_cap = capacity(FONT, self.args.size) + col_w = (cols - GUTTER) // 2 + pages = pack_pages(paras, col_w, 2 * rows) + page_lines = {pg: layout_page(paras[pg[0] : pg[1]], col_w) for pg in pages} + page_chars = [offsets[j - 1] + len(paras[j - 1]["ctx"]) - offsets[i] for i, j in pages] + self.capacity_stats[length] = { + "doc_pages": len(pages), + "doc_mean_chars_page": round(sum(page_chars) / len(pages)), + "doc_chars_page": page_chars, + "grid_chars_page": grid_cap, + "corpus_chars": len(flow), + "grid_pages": -(-len(flow) // grid_cap), + } + self.ctxs[length] = { + "args": self.args, "flow": flow, "paras": paras, "offsets": offsets, + "keys": self.keys, "length": length, "pages": pages, "lines": page_lines, + } + return self.ctxs[length] + + def run(self, cells: list[tuple[int, str]], label: str) -> None: + cells = [c for c in cells if c not in self.done] + self.done.update(cells) + tasks = [] + for length, cond in cells: + ctx = self.ctx(length) + if cond.startswith("img-doc-"): + for pg in ctx["pages"]: + tasks.append((run_doc_page, (MODEL, cond, pg, ctx))) + else: + grid_cap = capacity(FONT, self.args.size)[2] + flow = ctx["flow"] + for start in range(0, len(flow), grid_cap): + tasks.append((run_grid_chunk, (MODEL, cond, start, min(start + grid_cap, len(flow)), ctx))) + if not tasks: + return + print(f"[{label}] {len(cells)} cells -> {len(tasks)} page/chunk tasks") + with ThreadPoolExecutor(self.args.workers) as pool: + futures = [pool.submit(fn, *t) for fn, t in tasks] + for k, fut in enumerate(futures): + self.records.extend(fut.result()) + print(f" {k + 1}/{len(tasks)}", flush=True) + + def cell(self, length: int, cond: str) -> dict | None: + sub = [r for r in self.records if r["length"] == length and r["cond"] == cond] + return {"model": MODEL, "length": length, "condition": cond, **aggregate(sub, *PRICE)} if sub else None + + def all_cells(self) -> list[dict]: + keys = sorted({(r["length"], r["cond"]) for r in self.records}) + return [c for ln, cond in keys if (c := self.cell(ln, cond))] + + +def print_cells(cells: list[dict]) -> None: + for c in cells: + best, best_se = BEST_KNOWN.get(c["length"], (None, None)) + extra = "" + if best is not None: + comb = (c["f1_se"] ** 2 + best_se**2) ** 0.5 + d = c["f1"] - best + extra = f" vs best {best:.3f}: {d:+.3f} ({d / comb:+.1f}se)" + if c["f1"] >= TEXT_CEILING[c["length"]] - 1e-9: + extra += " >= TEXT CEILING" + print( + f" len {c['length']:<4} {c['condition']:<24} n={c['n']:<4} EM {c['em']:.3f} " + f"F1 {c['f1']:.3f} ±{c['f1_se']:.3f} out {c['tok_out']} reas {c['tok_reasoning']} " + f"${c['cost_usd']:.3f}{extra}" + ) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--effort", default=None) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--render-only", action="store_true") + ap.add_argument("--screen-only", action="store_true", help="skip the 50/250 confirmation phase") + ap.add_argument("--confirm-conds", default=None, help="comma list; default = screening F1 winner") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(parents=True, exist_ok=True) + OUT_DIR.mkdir(parents=True, exist_ok=True) + + keys = {} + if not args.render_only: + keys["openrouter"] = load_env_key("OPENROUTER_API_KEY", args.env) + + cols, rows, grid_cap = capacity(FONT, args.size) + col_w = (cols - GUTTER) // 2 + print(f"font 8on16: {cols} cols x {rows} rows; grid cap {grid_cap}; doc 2 x {col_w} + gutter {GUTTER}, {2 * rows} line slots") + + runner = Runner(args, keys) + + if args.render_only: + ctx = runner.ctx(150) + pg = ctx["pages"][0] + for cond in SCREEN_CONDS: + if cond.startswith("img-doc-"): + variant = cond.removeprefix("img-doc-8on16-") + key = sha8(cond, json.dumps([(p["title"], p["ctx"]) for p in ctx["paras"][pg[0] : pg[1]]]), str(args.size)) + png = CACHE / f"{EXP}-doc-8on16-{variant}-{key}.png" + img = render_doc(ctx["lines"][pg], args.size, variant, CACHE) + else: + variant = cond.removeprefix("img-8on16-") + chunk = ctx["flow"][:grid_cap] + png = CACHE / f"{EXP}-grid-8on16-{variant}-{sha8(chunk, str(args.size))}.png" + img = render(chunk, FONT, CACHE, args.size, variant) + tmp = png.with_suffix(f".{os.getpid()}.tmp.png") + img.save(tmp) + tmp.replace(png) + print(f" sample: {png}") + for length, st in runner.capacity_stats.items(): + print(f" len {length}: {st['doc_pages']} doc pages, mean {st['doc_mean_chars_page']} chars/page " + f"(grid {st['grid_chars_page']} -> {st['grid_pages']} pages)") + return + + # Phase A: screen at 150 + runner.run([(150, c) for c in SCREEN_CONDS], "screen@150") + screen_cells = [c for c in runner.all_cells() if c["length"] == 150] + print_cells(screen_cells) + winner = max(screen_cells, key=lambda c: (c["f1"], -c["cost_usd"]))["condition"] + print(f"screen winner: {winner}") + + # Phase B: confirm winner at 50 and 250 + if not args.screen_only: + confirm = [w.strip() for w in args.confirm_conds.split(",")] if args.confirm_conds else [winner] + runner.run([(ln, c) for ln in (50, 250) for c in confirm], "confirm@50/250") + + cells = runner.all_cells() + with (OUT_DIR / "records.jsonl").open("w") as fh: + for r in runner.records: + fh.write(json.dumps(r) + "\n") + (OUT_DIR / "summary.json").write_text( + json.dumps( + {"args": vars(args), "model": MODEL, "capacity": runner.capacity_stats, + "screen_winner": winner, "cells": cells}, + indent=1, + ) + ) + with (OUT_DIR / "matrix.csv").open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=[k for k in cells[0].keys() if k != "doc_chars_page"]) + writer.writeheader() + writer.writerows(cells) + + print_cells(cells) + print(f"\n-> {OUT_DIR}/records.jsonl, matrix.csv, summary.json") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp16_bestfable.py b/packages/snapcompact/research/exp16_bestfable.py new file mode 100644 index 000000000..8045f588c --- /dev/null +++ b/packages/snapcompact/research/exp16_bestfable.py @@ -0,0 +1,476 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp16: best optical profile for claude-fable-5. + +Round-1 levers (patch-aligned pitch 16, document layout, per-model variant) +were validated on gpt-5.5/gemini only. Fable's known winner is img-6x12-dim +(smaller glyphs than the 8x13 the others prefer, plain dim variant). Test +whether the round-1 levers transfer: + + img-8on16-dim 8x13 glyphs on an 8x16 cell (pitch-16 alignment, big glyphs) + img-6on7x14-dim 6x12 glyphs on a 7x14 cell (alignment, fable's glyph size) + doc-6x12-dim two-column newspaper layout at fable's winner font + doc-8on16-dim layout + alignment combined (run if either lever shows) + +Screen at length 150, confirm the winner at 50/250. Baselines (do not re-run): +img-6x12-dim F1 .956/.911/.923 at $0.132/$0.437/$0.724; text ceiling +.956/.904/.920 at $0.144/$0.498/$0.734. + +Usage: uv run exp16_bestfable.py --render-only # sample PNGs + capacity + uv run exp16_bestfable.py --lengths 150 # screen + uv run exp16_bestfable.py --lengths 50,150,250 --conditions ... # confirm + (re-runs hit .cache/qa/, so the final full invocation rebuilds the + combined records/matrix for free) +""" + +import argparse +import csv +import json +import sys +import uuid +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +from PIL import Image + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import _DIMMED, _stopword_mask, FontCfg, capacity, ensure_font, parse_bdf, render # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 + +MODEL = "claude-fable-5" +PRICE = (10.0, 50.0) # $/M in, out +FONTS = { + "8on16": FontCfg("8on16", "8x13", 8, 16), # 8x13 glyphs, patch-aligned 16 px pitch + "6on7x14": FontCfg("6on7x14", "6x12", 7, 14), # 6x12 glyphs, 7x14 patch-aligned cell + "6x12": FontCfg("6x12", "6x12", 6, 12), # fable's round-0 winner font +} +CONDITIONS = ("img-8on16-dim", "img-6on7x14-dim", "doc-6x12-dim", "doc-8on16-dim") +LENGTHS = (150,) +GUTTER = 3 # char cells between doc columns +_WHITE = (255, 255, 255) +_BLACK = (0, 0, 0) + +# img-6x12-dim per length: (f1, se, cost); text ceiling: (f1, se, cost). +BASE_IMG = {50: (0.9556, 0.0348, 0.132), 150: (0.9113, 0.0244, 0.437), 250: (0.9233, 0.0163, 0.724)} +BASE_TEXT = {50: (0.9556, 0.0348, 0.144), 150: (0.9043, 0.0216, 0.498), 250: (0.9197, 0.0184, 0.734)} + + +def cached(model: str, tag: str, payload: object, fn, fresh: bool) -> dict: + """Disk-cache `fn() -> dict` keyed by (model, tag, payload). Truncations are not cached.""" + key = sha8(model, tag, json.dumps(payload, sort_keys=True, default=str)) + path = QA_CACHE / f"{key}.json" + if path.exists() and not fresh: + hit = json.loads(path.read_text()) + if hit.get("stop") != "max_tokens": + return hit + out = fn() + if out.get("stop") == "max_tokens": + print(f" WARN truncated, not cached: {model} {tag} {key}") + else: + path.write_text(json.dumps(out)) + return out + + +def parse_condition(name: str) -> tuple[str, str, str]: + """'img-8on16-dim' -> (kind, font, variant).""" + kind, font, variant = name.split("-", 2) + return kind, font, variant + + +def atomic_save(img: Image.Image, png: Path) -> None: + tmp = png.with_suffix(f".{uuid.uuid4().hex[:8]}.tmp.png") + img.save(tmp) + tmp.replace(png) + + +# --- document layout (ported from exp04_layout.py, parameterized font) ------ + + +def wrap(text: str, width: int) -> list[str]: + """Greedy word-wrap, no mid-word breaks (hard split only for width+ words).""" + lines: list[str] = [] + cur = "" + for word in text.split(): + while len(word) > width: # pathological; never hit on SQuAD prose + if cur: + lines.append(cur) + cur = "" + lines.append(word[:width]) + word = word[width:] + if not cur: + cur = word + elif len(cur) + 1 + len(word) <= width: + cur += " " + word + else: + lines.append(cur) + cur = word + if cur: + lines.append(cur) + return lines + + +def layout_page(paras: list[dict], col_w: int) -> list[dict]: + """Typeset paragraphs into lines: [{kind: heading|body|blank, text}]. + + Article title changes become headings (repeated at the top of a page even + when the article continues, since each page is read in isolation). + """ + lines: list[dict] = [] + prev_title = None + for p in paras: + if p["title"] != prev_title: + if lines: + lines.append({"kind": "blank", "text": ""}) + for hl in wrap(p["title"].replace("_", " ").upper(), col_w): + lines.append({"kind": "heading", "text": hl}) + prev_title = p["title"] + elif lines: + lines.append({"kind": "blank", "text": ""}) + for bl in wrap(p["ctx"], col_w): + lines.append({"kind": "body", "text": bl}) + return lines + + +def pack_pages(paras: list[dict], col_w: int, max_lines: int) -> list[tuple[int, int]]: + """Greedy paragraph-aligned packing: [(i, j)] para ranges, one per page.""" + pages = [] + i = 0 + while i < len(paras): + j = i + 1 + while j < len(paras) and len(layout_page(paras[i : j + 1], col_w)) <= max_lines: + j += 1 + pages.append((i, j)) + i = j + return pages + + +def _dim_masks(lines: list[dict]) -> list[list[bool]]: + """Per-line stopword mask, computed over the joined page text.""" + joined = "\n".join(ln["text"] for ln in lines) + mask = _stopword_mask(joined) + out, pos = [], 0 + for ln in lines: + n = len(ln["text"]) + out.append(mask[pos : pos + n]) + pos += n + 1 # the joining newline + return out + + +def render_doc(lines: list[dict], cfg: FontCfg, size: int, cache: Path) -> Image.Image: + """Two-column dim page: black content words, gray stopwords, double-strike headings.""" + glyphs, font_ascent = parse_bdf(ensure_font(cfg, cache)) + ascent = cfg.ascent if cfg.ascent is not None else font_ascent + cols, rows, _ = capacity(cfg, size) + col_w = (cols - GUTTER) // 2 + masks = _dim_masks(lines) + img = Image.new("RGB", (size, size), _WHITE) + px = img.load() + for li, ln in enumerate(lines): + column, row = divmod(li, rows) + if column > 1: + break # overflow guard; pack_pages should prevent this + x_origin = column * (col_w + GUTTER) * cfg.adv + y0 = row * cfg.pitch + for ci, ch in enumerate(ln["text"]): + glyph = glyphs.get(ord(ch)) + if glyph is None: + continue + if ln["kind"] == "heading": + fg = _BLACK + else: + fg = _DIMMED if masks[li][ci] else _BLACK + w, h, xoff, yoff = glyph["bbx"] + top = y0 + ascent - h - yoff + shift = 0x80 if w <= 8 else 0x8000 + strikes = (0, 1) if ln["kind"] == "heading" else (0,) + for dx in strikes: + for r, bits in enumerate(glyph["rows"]): + y = top + r + if not 0 <= y < size: + continue + for b in range(w): + if bits & (shift >> b): + x = x_origin + ci * cfg.adv + xoff + b + dx + if 0 <= x < size: + px[x, y] = fg + return img + + +# --- runners ----------------------------------------------------------------- + + +def qa_call(cond: str, messages: list[dict], ctx: dict) -> dict: + args, keys = ctx["args"], ctx["keys"] + return cached( + MODEL, f"exp16-qa-{cond}", {"messages": messages, "size": args.size, "effort": args.effort}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete(keys, MODEL, messages, max_tokens=args.max_tokens, effort=args.effort), + ) + ), + args.fresh, + ) + + +def score(questions: list[dict], qa: dict, cond: str, start: int, ctx: dict) -> list[dict]: + answers = squad.parse_numbered(qa["text"], len(questions)) + records = [] + for q, a in zip(questions, answers): + records.append( + { + "model": MODEL, + "length": ctx["length"], + "cond": cond, + "chunk": start, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + ) + records[0]["usage"] = [{"phase": "qa", **qa["usage"]}] + return records + + +def run_grid_chunk(cond: str, start: int, end: int, ctx: dict) -> list[dict]: + """Row-major grid cell: chunk the flow by capacity, one QA call per chunk.""" + args, flow, paras, offsets = ctx["args"], ctx["flow"], ctx["paras"], ctx["offsets"] + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + chunk_text = flow[start:end] + _, font, variant = parse_condition(cond) + png = CACHE / f"exp16-{font}-{variant}-{sha8(chunk_text, str(args.size))}.png" + if not png.exists() or png.stat().st_size == 0: + atomic_save(render(chunk_text, FONTS[font], CACHE, args.size, variant), png) + cols, rows, _ = capacity(FONTS[font], args.size) + q_block = "\n".join(f"{i + 1}. {q['q']}" for i, q in enumerate(questions)) + messages = [ + { + "role": "user", + "content": [ + {"text": load_prompt("qa-image.md").format(cols=cols, rows=rows)}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + return score(questions, qa_call(cond, messages, ctx), cond, start, ctx) + + +def run_doc_page(cond: str, page: tuple[int, int], ctx: dict) -> list[dict]: + """Document cell: paragraph-aligned page, two-column dim render.""" + args, paras, offsets = ctx["args"], ctx["paras"], ctx["offsets"] + i, j = page + start = offsets[i] + end = offsets[j - 1] + len(paras[j - 1]["ctx"]) + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + _, font, _ = parse_condition(cond) + cfg = FONTS[font] + lines = ctx["lines"][cond][page] + page_key = sha8(cond, json.dumps([(p["title"], p["ctx"]) for p in paras[i:j]]), str(args.size)) + png = CACHE / f"exp16-{cond}-{page_key}.png" + if not png.exists() or png.stat().st_size == 0: + atomic_save(render_doc(lines, cfg, args.size, CACHE), png) + cols, rows, _ = capacity(cfg, args.size) + col_w = (cols - GUTTER) // 2 + q_block = "\n".join(f"{k + 1}. {q['q']}" for k, q in enumerate(questions)) + messages = [ + { + "role": "user", + "content": [ + {"text": load_prompt("exp04-qa-image.md").format(col_w=col_w, rows=rows)}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + return score(questions, qa_call(cond, messages, ctx), cond, start, ctx) + + +def aggregate(records: list[dict], price_in: float, price_out: float) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean_f1 = sum(f1s) / n + se = (sum((x - mean_f1) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + us = [u for r in records if "usage" in r for u in r["usage"]] + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + cost_in = (tok["in"] + 1.25 * tok["cache_w"] + 0.1 * tok["cache_r"]) / 1e6 * price_in + cost_out = tok["out"] / 1e6 * price_out + return { + "n": n, + "em": sum(r["em"] for r in records) / n, + "f1": mean_f1, + "f1_se": se, + "abstained": sum(r["abstained"] for r in records), + **{f"tok_{k}": v for k, v in tok.items()}, + "cost_in_usd": round(cost_in, 4), + "cost_out_usd": round(cost_out, 4), + "cost_usd": round(cost_in + cost_out, 4), + } + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--lengths", default=",".join(map(str, LENGTHS))) + ap.add_argument("--conditions", default=",".join(CONDITIONS)) + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--effort", default=None) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--render-only", action="store_true", help="render sample pages + capacity stats, no API") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + out_dir = RESULTS / "exp16-bestfable" + out_dir.mkdir(parents=True, exist_ok=True) + + lengths = [int(x) for x in args.lengths.split(",") if x.strip()] + conditions = [c.strip() for c in args.conditions.split(",") if c.strip()] + + keys = {} + if not args.render_only: + keys["anthropic"] = load_env_key("ANTHROPIC_API_KEY", args.env) + + all_paras = squad.load_paragraphs(CACHE) + tasks = [] + capacity_stats: dict[str, dict] = {} + for length in lengths: + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + page_lines: dict[str, dict] = {} + doc_pages: dict[str, list[tuple[int, int]]] = {} + for cond in conditions: + kind, font, _ = parse_condition(cond) + cfg = FONTS[font] + cols, rows, grid_cap = capacity(cfg, args.size) + if kind == "doc": + col_w = (cols - GUTTER) // 2 + pages = pack_pages(paras, col_w, 2 * rows) + doc_pages[cond] = pages + page_lines[cond] = {pg: layout_page(paras[pg[0] : pg[1]], col_w) for pg in pages} + page_chars = [offsets[j - 1] + len(paras[j - 1]["ctx"]) - offsets[i] for i, j in pages] + capacity_stats[f"{cond}@{length}"] = { + "pages": len(pages), + "mean_chars_page": round(sum(page_chars) / len(pages)), + "grid_chars_page": grid_cap, + "corpus_chars": len(flow), + } + else: + capacity_stats[f"{cond}@{length}"] = { + "pages": -(-len(flow) // grid_cap), + "mean_chars_page": grid_cap, + "grid_chars_page": grid_cap, + "corpus_chars": len(flow), + } + ctx = { + "args": args, "flow": flow, "paras": paras, "offsets": offsets, + "keys": keys, "length": length, "lines": page_lines, + } + for cond in conditions: + kind, font, _ = parse_condition(cond) + if kind == "doc": + for pg in doc_pages[cond]: + tasks.append(("doc", cond, pg, ctx)) + else: + budget = capacity(FONTS[font], args.size)[2] + for start in range(0, len(flow), budget): + tasks.append(("img", cond, (start, min(start + budget, len(flow))), ctx)) + + for key, st in sorted(capacity_stats.items()): + print( + f" {key}: {st['pages']} pages, mean {st['mean_chars_page']} chars/page " + f"(grid cap {st['grid_chars_page']}; corpus {st['corpus_chars']})" + ) + + if args.render_only: + for length in lengths: + paras = all_paras[:length] + for cond in conditions: + kind, font, _ = parse_condition(cond) + cfg = FONTS[font] + if kind == "doc": + cols, rows, _ = capacity(cfg, args.size) + col_w = (cols - GUTTER) // 2 + i, j = pack_pages(paras, col_w, 2 * rows)[0] + lines = layout_page(paras[i:j], col_w) + key = sha8(cond, json.dumps([(p["title"], p["ctx"]) for p in paras[i:j]]), str(args.size)) + png = CACHE / f"exp16-{cond}-{key}.png" + atomic_save(render_doc(lines, cfg, args.size, CACHE), png) + else: + flow, _ = squad.build_flow(paras) + cap = capacity(cfg, args.size)[2] + chunk_text = flow[:cap] + _, _, variant = parse_condition(cond) + png = CACHE / f"exp16-{font}-{variant}-{sha8(chunk_text, str(args.size))}.png" + atomic_save(render(chunk_text, cfg, CACHE, args.size, variant), png) + print(f" sample: {png}") + return + + print(f"grid: {len(tasks)} tasks") + records: list[dict] = [] + done = 0 + with ThreadPoolExecutor(args.workers) as pool: + futures = [ + pool.submit(run_doc_page, cond, span, ctx) + if kind == "doc" + else pool.submit(run_grid_chunk, cond, span[0], span[1], ctx) + for kind, cond, span, ctx in tasks + ] + for fut in futures: + records.extend(fut.result()) + done += 1 + print(f" {done}/{len(tasks)} tasks", flush=True) + + with (out_dir / "records.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + + cells = [] + for length in lengths: + for cond in conditions: + sub = [r for r in records if r["length"] == length and r["cond"] == cond] + if not sub: + continue + cells.append({"model": MODEL, "length": length, "condition": cond, **aggregate(sub, *PRICE)}) + (out_dir / "summary.json").write_text( + json.dumps({"args": vars(args), "capacity": capacity_stats, "cells": cells}, indent=1) + ) + with (out_dir / "matrix.csv").open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=list(cells[0].keys())) + writer.writeheader() + writer.writerows(cells) + + for c in cells: + bi, bt = BASE_IMG.get(c["length"]), BASE_TEXT.get(c["length"]) + comb_se = (c["f1_se"] ** 2 + bi[1] ** 2) ** 0.5 if bi else 0.0 + d_img = f"vs 6x12-dim {c['f1'] - bi[0]:+.3f} ({(c['f1'] - bi[0]) / comb_se:+.1f}se)" if bi else "" + d_txt = f" vs text {c['f1'] - bt[0]:+.3f}" if bt else "" + flag = " ** beats text ceiling" if bt and c["f1"] > bt[0] else "" + print( + f"{MODEL} len {c['length']:<4} {c['condition']:<18} n={c['n']:<4} " + f"EM {c['em']:.3f} F1 {c['f1']:.3f} ±{c['f1_se']:.3f} ${c['cost_usd']:.3f} " + f"{d_img}{d_txt}{flag}" + ) + print(f"\n-> {out_dir}/records.jsonl, matrix.csv, summary.json") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp17_bestopus.py b/packages/snapcompact/research/exp17_bestopus.py new file mode 100644 index 000000000..01ec418fd --- /dev/null +++ b/packages/snapcompact/research/exp17_bestopus.py @@ -0,0 +1,419 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp17: claude-opus-4-8 model-owner run — do round-1 levers transfer to Anthropic? + +Round-1 (gpt-5.5/gemini) validated: (1) patch-aligned pitch 16 (8x13 glyphs on +an 8x16 cell), (2) two-column document layout, (3) per-model variant. Opus's +round-0 winner is img-8x13-bw; opus emits ~0 reasoning tokens on image QA so +any win here must come from F1, not output-token savings. + +Conditions (variant anchor = bw): + img-8on16-bw grid, 8x13 glyphs on 16 px pitch (alignment only) + img-doc-8on16-bw two-column doc layout at 8on16 (alignment + layout) + img-doc-8x13-bw two-column doc layout at 8x13 (layout only) + +Protocol: screen all three at length 150; only the winner goes to 50/250. +Records accumulate across invocations (re-run cells replace their old rows), +so the final matrix.csv holds the union of screen + confirm runs. + +Usage: uv run exp17_bestopus.py --render-only # capacity + sample PNGs + uv run exp17_bestopus.py # screen @150 + uv run exp17_bestopus.py --conditions img-X --lengths 50,250 # confirm + uv run exp17_bestopus.py --report # reprint from cache +""" + +import argparse +import csv +import json +import sys +import uuid +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +from PIL import Image + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import FontCfg, capacity, ensure_font, parse_bdf, render # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 + +MODEL = "claude-opus-4-8" +PRICES = (15.0, 75.0) # $/M in, $/M out +FONTS = { + "8x13": FontCfg("8x13", "8x13", 8, 13), # round-0 opus winner font + "8on16": FontCfg("8on16", "8x13", 8, 16), # same glyphs, patch-aligned 16 px pitch +} +CONDITIONS = ("img-8on16-bw", "img-doc-8on16-bw", "img-doc-8x13-bw") +LENGTHS = (150,) # screening default; confirm via --lengths 50,250 +GUTTER = 3 # char cells between doc columns +_WHITE = (255, 255, 255) +_BLACK = (0, 0, 0) +_INK = (24, 24, 24) # near-black body text, like a printed page + +# claude-opus-4-8 baselines, results/optimal-combined/matrix.csv (qpc 30, seed 42): +BASELINE = {50: (0.9626, 0.0258, 0.143), 150: (0.8937, 0.0223, 0.380), 250: (0.8708, 0.0196, 0.559)} +TEXT_CEIL = {50: (0.9278, 0.195), 150: (0.9112, 0.637), 250: (0.9268, 0.938)} + + +def cached(model: str, tag: str, payload: object, fn, fresh: bool) -> dict: + """Disk-cache `fn() -> dict` keyed by (model, tag, payload). Truncations are not cached.""" + key = sha8(model, tag, json.dumps(payload, sort_keys=True, default=str)) + path = QA_CACHE / f"{key}.json" + if path.exists() and not fresh: + hit = json.loads(path.read_text()) + if hit.get("stop") != "max_tokens": + return hit + out = fn() + if out.get("stop") == "max_tokens": + print(f" WARN truncated, not cached: {model} {tag} {key}") + else: + path.write_text(json.dumps(out)) + return out + + +def parse_cond(name: str) -> tuple[str, FontCfg]: + """'img-8on16-bw' -> ('grid', cfg); 'img-doc-8on16-bw' -> ('doc', cfg).""" + parts = name.split("-") + if parts[1] == "doc": + return "doc", FONTS[parts[2]] + return "grid", FONTS[parts[1]] + + +# --- document layout (ported from exp04_layout.py, parameterized font) ------ + + +def wrap(text: str, width: int) -> list[str]: + """Greedy word-wrap, no mid-word breaks (hard split only for width+ words).""" + lines: list[str] = [] + cur = "" + for word in text.split(): + while len(word) > width: # pathological; never hit on SQuAD prose + if cur: + lines.append(cur) + cur = "" + lines.append(word[:width]) + word = word[width:] + if not cur: + cur = word + elif len(cur) + 1 + len(word) <= width: + cur += " " + word + else: + lines.append(cur) + cur = word + if cur: + lines.append(cur) + return lines + + +def layout_page(paras: list[dict], col_w: int) -> list[dict]: + """Typeset paragraphs into lines: [{kind: heading|body|blank, text}]. + + Headings repeat at the top of a page when an article continues, since + each page is read in isolation. + """ + lines: list[dict] = [] + prev_title = None + for p in paras: + if p["title"] != prev_title: + if lines: + lines.append({"kind": "blank", "text": ""}) + for hl in wrap(p["title"].replace("_", " ").upper(), col_w): + lines.append({"kind": "heading", "text": hl}) + prev_title = p["title"] + elif lines: + lines.append({"kind": "blank", "text": ""}) + for bl in wrap(p["ctx"], col_w): + lines.append({"kind": "body", "text": bl}) + return lines + + +def pack_pages(paras: list[dict], col_w: int, max_lines: int) -> list[tuple[int, int]]: + """Greedy paragraph-aligned packing: [(i, j)] para ranges, one per page.""" + pages = [] + i = 0 + while i < len(paras): + j = i + 1 + while j < len(paras) and len(layout_page(paras[i : j + 1], col_w)) <= max_lines: + j += 1 + pages.append((i, j)) + i = j + return pages + + +def render_doc(lines: list[dict], cfg: FontCfg, size: int) -> Image.Image: + """Two-column bw page: left column rows top-to-bottom, then right column.""" + glyphs, font_ascent = parse_bdf(ensure_font(cfg, CACHE)) + ascent = cfg.ascent if cfg.ascent is not None else font_ascent + cols, rows, _ = capacity(cfg, size) + col_w = (cols - GUTTER) // 2 + img = Image.new("RGB", (size, size), _WHITE) + px = img.load() + for li, ln in enumerate(lines): + column, row = divmod(li, rows) + if column > 1: + break # overflow guard; pack_pages should prevent this + x_origin = column * (col_w + GUTTER) * cfg.adv + y0 = row * cfg.pitch + for ci, ch in enumerate(ln["text"]): + glyph = glyphs.get(ord(ch)) + if glyph is None: + continue + fg = _BLACK if ln["kind"] == "heading" else _INK + w, h, xoff, yoff = glyph["bbx"] + top = y0 + ascent - h - yoff + shift = 0x80 if w <= 8 else 0x8000 + strikes = (0, 1) if ln["kind"] == "heading" else (0,) + for dx in strikes: + for r, bits in enumerate(glyph["rows"]): + y = top + r + if not 0 <= y < size: + continue + for b in range(w): + if bits & (shift >> b): + x = x_origin + ci * cfg.adv + xoff + b + dx + if 0 <= x < size: + px[x, y] = fg + return img + + +# --- runner ----------------------------------------------------------------- + + +def render_unit_png(cond: str, unit: tuple[int, int], ctx: dict) -> Path: + """Render (or reuse) the PNG for one chunk/page; atomic tmp-then-replace.""" + args, paras = ctx["args"], ctx["paras"] + kind, cfg = parse_cond(cond) + if kind == "grid": + start, end = unit + chunk_text = ctx["flow"][start:end] + png = CACHE / f"exp17-{cond}-{sha8(chunk_text, str(args.size))}.png" + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(f".{uuid.uuid4().hex[:8]}.tmp.png") + render(chunk_text, cfg, CACHE, args.size, "bw").save(tmp) + tmp.replace(png) + else: + i, j = unit + key = sha8(cond, json.dumps([(p["title"], p["ctx"]) for p in paras[i:j]]), str(args.size)) + png = CACHE / f"exp17-{cond}-{key}.png" + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(f".{uuid.uuid4().hex[:8]}.tmp.png") + render_doc(ctx["lines"][(cond, unit)], cfg, args.size).save(tmp) + tmp.replace(png) + return png + + +def run_unit(cond: str, unit: tuple[int, int], ctx: dict) -> list[dict]: + """One (condition, chunk-or-page) cell unit: render carrier, QA, score.""" + args, paras, offsets, keys = ctx["args"], ctx["paras"], ctx["offsets"], ctx["keys"] + kind, cfg = parse_cond(cond) + if kind == "grid": + start, end = unit + else: + i, j = unit + start = offsets[i] + end = offsets[j - 1] + len(paras[j - 1]["ctx"]) + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + q_block = "\n".join(f"{k + 1}. {q['q']}" for k, q in enumerate(questions)) + png = render_unit_png(cond, unit, ctx) + cols, rows, _ = capacity(cfg, args.size) + if kind == "grid": + preamble = load_prompt("qa-image.md").format(cols=cols, rows=rows) + else: + preamble = load_prompt("exp04-qa-image.md").format(col_w=(cols - GUTTER) // 2, rows=rows) + messages = [ + { + "role": "user", + "content": [ + {"text": preamble}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + qa = cached( + MODEL, "exp17-qa", {"messages": messages, "effort": args.effort}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete(keys, MODEL, messages, max_tokens=args.max_tokens, effort=args.effort), + ) + ), + args.fresh, + ) + answers = squad.parse_numbered(qa["text"], len(questions)) + records = [] + for q, a in zip(questions, answers): + records.append( + { + "model": MODEL, + "length": ctx["length"], + "cond": cond, + "chunk": start, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + ) + records[0]["usage"] = [{"phase": "qa", **qa["usage"]}] + return records + + +def aggregate(records: list[dict], price_in: float, price_out: float) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean_f1 = sum(f1s) / n + se = (sum((x - mean_f1) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + us = [u for r in records if "usage" in r for u in r["usage"]] + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + cost_in = (tok["in"] + 1.25 * tok["cache_w"] + 0.1 * tok["cache_r"]) / 1e6 * price_in + cost_out = tok["out"] / 1e6 * price_out + return { + "n": n, + "em": sum(r["em"] for r in records) / n, + "f1": mean_f1, + "f1_se": se, + "abstained": sum(r["abstained"] for r in records), + **{f"tok_{k}": v for k, v in tok.items()}, + "cost_in_usd": round(cost_in, 4), + "cost_out_usd": round(cost_out, 4), + "cost_usd": round(cost_in + cost_out, 4), + } + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--conditions", default=",".join(CONDITIONS)) + ap.add_argument("--lengths", default=",".join(map(str, LENGTHS))) + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--effort", default=None) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--render-only", action="store_true", help="capacity stats + first-page PNGs, no API") + ap.add_argument("--report", action="store_true", help="reprint matrix from accumulated records only") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + out_dir = RESULTS / "exp17-bestopus" + out_dir.mkdir(parents=True, exist_ok=True) + records_path = out_dir / "records.jsonl" + + conditions = [c.strip() for c in args.conditions.split(",") if c.strip()] + lengths = [int(x) for x in args.lengths.split(",") if x.strip()] + + all_paras = squad.load_paragraphs(CACHE) + records: list[dict] = [] + + if not args.report: + keys = {} if args.render_only else {"anthropic": load_env_key("ANTHROPIC_API_KEY", args.env)} + tasks = [] + capacity_stats = {} + for length in lengths: + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + ctx = { + "args": args, "flow": flow, "paras": paras, "offsets": offsets, + "keys": keys, "length": length, "lines": {}, + } + for cond in conditions: + kind, cfg = parse_cond(cond) + cols, rows, grid_cap = capacity(cfg, args.size) + if kind == "grid": + units = [(s, min(s + grid_cap, len(flow))) for s in range(0, len(flow), grid_cap)] + chars = [e - s for s, e in units] + else: + col_w = (cols - GUTTER) // 2 + pages = pack_pages(paras, col_w, 2 * rows) + for pg in pages: + ctx["lines"][(cond, pg)] = layout_page(paras[pg[0] : pg[1]], col_w) + units = pages + chars = [offsets[j - 1] + len(paras[j - 1]["ctx"]) - offsets[i] for i, j in pages] + capacity_stats[f"{cond}@{length}"] = { + "pages": len(units), + "mean_chars_page": round(sum(chars) / len(units)), + "grid_chars_page": grid_cap, + "corpus_chars": len(flow), + } + tasks.extend((cond, u, ctx) for u in units) + + for k, st in capacity_stats.items(): + print( + f" {k}: {st['pages']} pages, mean {st['mean_chars_page']} chars/page " + f"(grid cap {st['grid_chars_page']}; corpus {st['corpus_chars']})" + ) + + if args.render_only: + seen = set() + for cond, unit, ctx in tasks: + if cond in seen: + continue + seen.add(cond) + print(f" sample: {render_unit_png(cond, unit, ctx)}") + return + + print(f"grid: {len(tasks)} unit tasks ({len(conditions)} conds x {lengths})") + done = 0 + with ThreadPoolExecutor(args.workers) as pool: + futures = [pool.submit(run_unit, c, u, ctx) for c, u, ctx in tasks] + for fut in futures: + records.extend(fut.result()) + done += 1 + print(f" {done}/{len(tasks)} units", flush=True) + + # Merge: rows for cells just run replace any prior rows for those cells. + ran_cells = {(length, cond) for length in lengths for cond in conditions} + old = [] + if records_path.exists(): + with records_path.open() as fh: + old = [json.loads(ln) for ln in fh if ln.strip()] + records = [r for r in old if (r["length"], r["cond"]) not in ran_cells] + records + with records_path.open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + else: + with records_path.open() as fh: + records = [json.loads(ln) for ln in fh if ln.strip()] + + cells = [] + for length in sorted({r["length"] for r in records}): + for cond in sorted({r["cond"] for r in records if r["length"] == length}): + sub = [r for r in records if r["length"] == length and r["cond"] == cond] + cells.append({"model": MODEL, "length": length, "condition": cond, **aggregate(sub, *PRICES)}) + (out_dir / "summary.json").write_text( + json.dumps({"args": vars(args), "baseline_img_8x13_bw": BASELINE, "cells": cells}, indent=1) + ) + with (out_dir / "matrix.csv").open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=list(cells[0].keys())) + writer.writeheader() + writer.writerows(cells) + + print(f"\n{'len':<5}{'condition':<22}{'n':>5}{'EM':>7}{'F1':>7}{'+-se':>7}{'$':>8}{'d/se vs 8x13-bw':>17}") + for c in cells: + b_f1, b_se, b_cost = BASELINE[c["length"]] + dse = (c["f1"] - b_f1) / ((c["f1_se"] ** 2 + b_se**2) ** 0.5 or 1) + t_f1, _ = TEXT_CEIL[c["length"]] + flag = " > text ceiling" if c["f1"] > t_f1 else "" + print( + f"{c['length']:<5}{c['condition']:<22}{c['n']:>5}{c['em']:>7.3f}{c['f1']:>7.3f}" + f"{c['f1_se']:>7.3f}{c['cost_usd']:>8.3f}{dse:>+17.2f}{flag}" + ) + print(f"\n-> {out_dir}/records.jsonl, matrix.csv, summary.json") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp18_bestkimi.py b/packages/snapcompact/research/exp18_bestkimi.py new file mode 100644 index 000000000..60c2f0601 --- /dev/null +++ b/packages/snapcompact/research/exp18_bestkimi.py @@ -0,0 +1,439 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp18: best optical profile for moonshotai/kimi-k2.6. + +Round-1 levers (patch-aligned pitch 16, two-column doc layout, per-model +variant) were validated on gpt-5.5/gemini only. Kimi's round-0 winner is +img-8x13-sent-dim (beats text at 150) but it pays the worst read tax in the +fleet (~95% of image-cell cost is output tokens; 120k+ out at length 250). + +Screen at length 150, anchored on sent-dim: + img-8on16-sent-dim grid, 8x13 glyphs on an 8x16 cell (alignment only) + img-doc8on16-sent-dim two-column doc layout at 8on16 (alignment + layout) + img-doc8x13-sent-dim two-column doc layout at pitch 13 (layout only) +Confirm the winner at 50 and 250. + +Usage: uv run exp18_bestkimi.py # screening (length 150) + uv run exp18_bestkimi.py --lengths 50,250 --conditions img-... # confirm + uv run exp18_bestkimi.py --render-only # sample pages, no API + uv run exp18_bestkimi.py --report --lengths 50,150,250 # re-aggregate from cache +""" + +import argparse +import csv +import json +import sys +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +from PIL import Image + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import _DARK, _DIMMED, FontCfg, _stopword_mask, capacity, ensure_font, parse_bdf, render # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 + +MODEL = "moonshotai/kimi-k2.6" +PRICE_IN, PRICE_OUT = 0.68, 3.41 +FONTS = { + "8on16": FontCfg("8on16", "8x13", 8, 16), # patch-aligned: 8x13 glyphs, 16 px pitch + "8x13": FontCfg("8x13", "8x13", 8, 13), # kimi's round-0 winner pitch +} +# cond -> (kind, font key, variant). All anchored on sent-dim (kimi's winner). +CONDITIONS = { + "img-8on16-sent-dim": ("grid", "8on16", "sent-dim"), + "img-doc8on16-sent-dim": ("doc", "8on16", "sent-dim"), + "img-doc8x13-sent-dim": ("doc", "8x13", "sent-dim"), +} +GUTTER = 3 # char cells between doc columns +_WHITE = (255, 255, 255) +_BLACK = (0, 0, 0) + + +def cached(model: str, tag: str, payload: object, fn, fresh: bool) -> dict: + """Disk-cache `fn() -> dict` keyed by (model, tag, payload). Truncations are not cached.""" + key = sha8(model, tag, json.dumps(payload, sort_keys=True, default=str)) + path = QA_CACHE / f"{key}.json" + if path.exists() and not fresh: + hit = json.loads(path.read_text()) + if hit.get("stop") != "max_tokens": + return hit + out = fn() + if out.get("stop") == "max_tokens": + print(f" WARN truncated, not cached: {model} {tag} {key}") + else: + path.write_text(json.dumps(out)) + return out + + +# --- document layout (ported from exp04, parameterized for font/pitch) ----- + + +def wrap(text: str, width: int) -> list[str]: + """Greedy word-wrap, no mid-word breaks (hard split only for width+ words).""" + lines: list[str] = [] + cur = "" + for word in text.split(): + while len(word) > width: + if cur: + lines.append(cur) + cur = "" + lines.append(word[:width]) + word = word[width:] + if not cur: + cur = word + elif len(cur) + 1 + len(word) <= width: + cur += " " + word + else: + lines.append(cur) + cur = word + if cur: + lines.append(cur) + return lines + + +def layout_page(paras: list[dict], col_w: int) -> list[dict]: + """Typeset paragraphs into lines: [{kind: heading|body|blank, text}]. + + Title changes become uppercase double-strike headings; the heading is + repeated at the top of a page when an article continues, since each page + is read in isolation. One blank line between paragraphs. + """ + lines: list[dict] = [] + prev_title = None + for p in paras: + if p["title"] != prev_title: + if lines: + lines.append({"kind": "blank", "text": ""}) + for hl in wrap(p["title"].replace("_", " ").upper(), col_w): + lines.append({"kind": "heading", "text": hl}) + prev_title = p["title"] + elif lines: + lines.append({"kind": "blank", "text": ""}) + for bl in wrap(p["ctx"], col_w): + lines.append({"kind": "body", "text": bl}) + return lines + + +def pack_pages(paras: list[dict], col_w: int, max_lines: int) -> list[tuple[int, int]]: + """Greedy paragraph-aligned packing: [(i, j)] para ranges, one per page.""" + pages = [] + i = 0 + while i < len(paras): + j = i + 1 + while j < len(paras) and len(layout_page(paras[i : j + 1], col_w)) <= max_lines: + j += 1 + pages.append((i, j)) + i = j + return pages + + +def _sentence_indices_doc(lines: list[dict]) -> list[list[int]]: + """Per-line per-char sentence index, cycling across the page (newline counts as boundary space).""" + joined = "\n".join(ln["text"] for ln in lines) + idx, run = 0, [] + for i, ch in enumerate(joined): + run.append(idx) + if ch in ".!?" and i + 1 < len(joined) and joined[i + 1] in " \n": + idx += 1 + out, pos = [], 0 + for ln in lines: + n = len(ln["text"]) + out.append(run[pos : pos + n]) + pos += n + 1 # the joining newline + return out + + +def render_doc(lines: list[dict], cfg: FontCfg, size: int, cache: Path) -> Image.Image: + """Two-column sent-dim page: left column top-to-bottom, then right. + + Body glyph color = sentence hue, overridden to light gray for stopwords + (same composition as bdf.render's sent-dim). Headings: black double-strike. + """ + glyphs, font_ascent = parse_bdf(ensure_font(cfg, cache)) + ascent = cfg.ascent if cfg.ascent is not None else font_ascent + cols, rows, _ = capacity(cfg, size) + col_w = (cols - GUTTER) // 2 + sent_idx = _sentence_indices_doc(lines) + dim_masks = [_stopword_mask(ln["text"]) for ln in lines] + img = Image.new("RGB", (size, size), _WHITE) + px = img.load() + for li, ln in enumerate(lines): + column, row = divmod(li, rows) + if column > 1: + break # overflow guard; pack_pages should prevent this + x_origin = column * (col_w + GUTTER) * cfg.adv + y0 = row * cfg.pitch + for ci, ch in enumerate(ln["text"]): + glyph = glyphs.get(ord(ch)) + if glyph is None: + continue + if ln["kind"] == "heading": + fg = _BLACK + elif dim_masks[li][ci]: + fg = _DIMMED + else: + fg = _DARK[sent_idx[li][ci] % 6] + w, h, xoff, yoff = glyph["bbx"] + top = y0 + ascent - h - yoff + shift = 0x80 if w <= 8 else 0x8000 + strikes = (0, 1) if ln["kind"] == "heading" else (0,) + for dx in strikes: + for r, bits in enumerate(glyph["rows"]): + y = top + r + if not 0 <= y < size: + continue + for b in range(w): + if bits & (shift >> b): + x = x_origin + ci * cfg.adv + xoff + b + dx + if 0 <= x < size: + px[x, y] = fg + return img + + +# --- runner ----------------------------------------------------------------- + + +def doc_png(cond: str, paras: list[dict], lines: list[dict], cfg: FontCfg, size: int) -> Path: + key = sha8(cond, json.dumps([(p["title"], p["ctx"]) for p in paras]), str(size)) + png = CACHE / f"exp18-{cond}-{key}.png" + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(".tmp.png") + render_doc(lines, cfg, size, CACHE).save(tmp) + tmp.replace(png) + return png + + +def run_unit(cond: str, unit: dict, ctx: dict) -> list[dict]: + """One (condition, page/chunk) unit: render carrier, QA, score.""" + args, paras, offsets, keys = ctx["args"], ctx["paras"], ctx["offsets"], ctx["keys"] + start, end = unit["start"], unit["end"] + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + kind, font, variant = CONDITIONS[cond] + cfg = FONTS[font] + cols, rows, _ = capacity(cfg, args.size) + if kind == "grid": + chunk_text = ctx["flow"][start:end] + png = CACHE / f"exp18-{cond}-{sha8(chunk_text, str(args.size))}.png" + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(".tmp.png") + render(chunk_text, cfg, CACHE, args.size, variant).save(tmp) + tmp.replace(png) + prompt = load_prompt("qa-image.md").format(cols=cols, rows=rows) + else: + i, j = unit["page"] + png = doc_png(cond, paras[i:j], unit["lines"], cfg, args.size) + col_w = (cols - GUTTER) // 2 + prompt = load_prompt("exp04-qa-image.md").format(col_w=col_w, rows=rows) + q_block = "\n".join(f"{k + 1}. {q['q']}" for k, q in enumerate(questions)) + messages = [ + { + "role": "user", + "content": [{"text": prompt}, {"image_path": png}, {"text": q_block}], + } + ] + qa = cached( + MODEL, "exp18-qa", {"messages": messages, "size": args.size, "effort": args.effort}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete(keys, MODEL, messages, max_tokens=args.max_tokens, effort=args.effort), + ) + ), + args.fresh, + ) + answers = squad.parse_numbered(qa["text"], len(questions)) + records = [] + for q, a in zip(questions, answers): + records.append( + { + "model": MODEL, + "length": ctx["length"], + "cond": cond, + "chunk": start, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + ) + records[0]["usage"] = [{"phase": "qa", **qa["usage"]}] + return records + + +def aggregate(records: list[dict]) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean_f1 = sum(f1s) / n + se = (sum((x - mean_f1) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + us = [u for r in records if "usage" in r for u in r["usage"]] + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + cost_in = (tok["in"] + 1.25 * tok["cache_w"] + 0.1 * tok["cache_r"]) / 1e6 * PRICE_IN + cost_out = tok["out"] / 1e6 * PRICE_OUT + return { + "n": n, + "em": sum(r["em"] for r in records) / n, + "f1": mean_f1, + "f1_se": se, + "abstained": sum(r["abstained"] for r in records), + **{f"tok_{k}": v for k, v in tok.items()}, + "cost_in_usd": round(cost_in, 4), + "cost_out_usd": round(cost_out, 4), + "cost_usd": round(cost_in + cost_out, 4), + } + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--lengths", default="150") + ap.add_argument("--conditions", default=",".join(CONDITIONS)) + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--effort", default=None) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--render-only", action="store_true", help="render first page per cond + capacity stats, no API") + ap.add_argument("--report", action="store_true", help="re-aggregate (all units should hit cache)") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + out_dir = RESULTS / "exp18-bestkimi" + out_dir.mkdir(parents=True, exist_ok=True) + + lengths = [int(x) for x in args.lengths.split(",") if x.strip()] + conditions = [c.strip() for c in args.conditions.split(",") if c.strip()] + for c in conditions: + if c not in CONDITIONS: + sys.exit(f"unknown condition: {c}") + + keys = {} + if not args.render_only: + keys["openrouter"] = load_env_key("OPENROUTER_API_KEY", args.env) + + all_paras = squad.load_paragraphs(CACHE) + tasks: list[tuple[str, dict, dict]] = [] + capacity_stats: dict = {} + for length in lengths: + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + ctx = {"args": args, "flow": flow, "paras": paras, "offsets": offsets, "keys": keys, "length": length} + capacity_stats[length] = {"corpus_chars": len(flow), "conds": {}} + for cond in conditions: + kind, font, _ = CONDITIONS[cond] + cfg = FONTS[font] + cols, rows, grid_cap = capacity(cfg, args.size) + if kind == "grid": + units = [ + {"start": s, "end": min(s + grid_cap, len(flow))} for s in range(0, len(flow), grid_cap) + ] + chars = [u["end"] - u["start"] for u in units] + else: + col_w = (cols - GUTTER) // 2 + pages = pack_pages(paras, col_w, 2 * rows) + units = [] + for i, j in pages: + units.append( + { + "start": offsets[i], + "end": offsets[j - 1] + len(paras[j - 1]["ctx"]), + "page": (i, j), + "lines": layout_page(paras[i:j], col_w), + } + ) + chars = [u["end"] - u["start"] for u in units] + capacity_stats[length]["conds"][cond] = { + "pages": len(units), + "mean_chars_page": round(sum(chars) / len(units)), + "grid_chars_page": grid_cap, + } + for u in units: + tasks.append((cond, u, ctx)) + + for length, st in capacity_stats.items(): + print(f"len {length}: corpus {st['corpus_chars']} chars") + for cond, cs in st["conds"].items(): + print(f" {cond:<24} {cs['pages']} pages, mean {cs['mean_chars_page']} chars/page (grid cap {cs['grid_chars_page']})") + + if args.render_only: + for cond, u, ctx in tasks: + if u["start"] != 0: + continue + kind, font, variant = CONDITIONS[cond] + cfg = FONTS[font] + if kind == "grid": + chunk_text = ctx["flow"][u["start"] : u["end"]] + png = CACHE / f"exp18-{cond}-{sha8(chunk_text, str(args.size))}.png" + tmp = png.with_suffix(".tmp.png") + render(chunk_text, cfg, CACHE, args.size, variant).save(tmp) + tmp.replace(png) + else: + i, j = u["page"] + png = doc_png(cond, ctx["paras"][i:j], u["lines"], cfg, args.size) + print(f" sample: {png}") + return + + print(f"grid: {len(tasks)} unit tasks on {MODEL}") + records: list[dict] = [] + done = 0 + with ThreadPoolExecutor(args.workers) as pool: + futures = [pool.submit(run_unit, c, u, ctx) for c, u, ctx in tasks] + for fut in futures: + records.extend(fut.result()) + done += 1 + print(f" {done}/{len(tasks)} units", flush=True) + + # Merge with any prior records (confirm runs extend the screening set). + rec_path = out_dir / "records.jsonl" + old: list[dict] = [] + if rec_path.exists(): + ran = {(r["length"], r["cond"]) for r in records} + for line in rec_path.read_text().splitlines(): + r = json.loads(line) + if (r["length"], r["cond"]) not in ran: + old.append(r) + records = old + records + tmp = rec_path.with_suffix(".tmp.jsonl") + with tmp.open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + tmp.replace(rec_path) + + cells = [] + for length in sorted({r["length"] for r in records}): + for cond in CONDITIONS: + sub = [r for r in records if r["length"] == length and r["cond"] == cond] + if not sub: + continue + cells.append({"model": MODEL, "length": length, "condition": cond, **aggregate(sub)}) + (out_dir / "summary.json").write_text( + json.dumps({"args": vars(args), "capacity": capacity_stats, "cells": cells}, indent=1) + ) + with (out_dir / "matrix.csv").open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=list(cells[0].keys())) + writer.writeheader() + writer.writerows(cells) + + for c in cells: + print( + f"len {c['length']:<4} {c['condition']:<24} n={c['n']:<4} EM {c['em']:.3f} " + f"F1 {c['f1']:.3f} ±{c['f1_se']:.3f} out {c['tok_out']:>7} (reas {c['tok_reasoning']}) ${c['cost_usd']:.4f}" + ) + print(f"\n-> {out_dir}/records.jsonl, matrix.csv, summary.json") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp19_bestglm.py b/packages/snapcompact/research/exp19_bestglm.py new file mode 100644 index 000000000..b8ef651ba --- /dev/null +++ b/packages/snapcompact/research/exp19_bestglm.py @@ -0,0 +1,477 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp19: best optical profile for z-ai/glm-4.6v (weakest optical reader in the fleet). + +glm uniquely needed dark mode (bright hues on black) and its round-0 winner was +img-8x13-dark-sent (.789/.753/.638 vs text .944/.904/.880). Round-1 levers +(patch-aligned pitch 16, two-column doc layout) were only validated on +gpt-5.5/gemini. Screen at length 150: + + img-8on16-dark-sent 8x13 glyphs on an 8x16 cell (patch-aligned), dark + img-doc-8on16-dark-sent doc layout + pitch 16 + dark palette (headings in + bright white double-strike) + img-doc-8x13-dark-sent layout-only control at the baseline pitch 13 + img-doc-8on16-sent light-mode probe: does doc structure remove the + need for dark mode? + +then confirm the screening winner at lengths 50 and 250. Methodology matches +the baselines: seed 42, qpc 30, size 1568, max_tokens 32768, effort default. + +Usage: uv run exp19_bestglm.py # screen @150 + uv run exp19_bestglm.py --lengths 50,250 --conditions X # confirm winner + uv run exp19_bestglm.py --render-only # sample PNGs, no API + +Repeated invocations merge into results/exp19-bestglm/ (records for the +(model, length, condition) cells being run are replaced; others kept). +""" + +import argparse +import csv +import json +import sys +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +from PIL import Image + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import _BRIGHT, _DARK, FontCfg, capacity, ensure_font, parse_bdf, render # noqa: E402 +from final import cached # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 + +MODELS = {"z-ai/glm-4.6v": (0.30, 0.90)} +LENGTHS = (150,) +FONTS = { + "8on16": FontCfg("8on16", "8x13", 8, 16), # patch-aligned padded cell (exp01 pattern) + "8x13": FontCfg("8x13", "8x13", 8, 13), # baseline pitch +} +# condition -> (kind, font key, palette variant) +CONDITIONS = { + "img-8on16-dark-sent": ("grid", "8on16", "dark-sent"), + "img-doc-8on16-dark-sent": ("doc", "8on16", "dark-sent"), + "img-doc-8x13-dark-sent": ("doc", "8x13", "dark-sent"), + "img-doc-8on16-sent": ("doc", "8on16", "sent"), +} +GUTTER = 3 # char cells between doc columns +_WHITE = (255, 255, 255) +_BLACK = (0, 0, 0) + + +# --- document layout (ported from exp04, parameterized for font/pitch) ------ + + +def wrap(text: str, width: int) -> list[str]: + """Greedy word-wrap, no mid-word breaks (hard split only for width+ words).""" + lines: list[str] = [] + cur = "" + for word in text.split(): + while len(word) > width: # pathological; never hit on SQuAD prose + if cur: + lines.append(cur) + cur = "" + lines.append(word[:width]) + word = word[width:] + if not cur: + cur = word + elif len(cur) + 1 + len(word) <= width: + cur += " " + word + else: + lines.append(cur) + cur = word + if cur: + lines.append(cur) + return lines + + +def layout_page(paras: list[dict], col_w: int) -> list[dict]: + """Typeset paragraphs into lines: [{kind: heading|body|blank, text}]. + + Article title changes become headings (repeated at the top of a page even + when an article continues, since each page is read in isolation). + Paragraphs are separated by one blank line. + """ + lines: list[dict] = [] + prev_title = None + for p in paras: + if p["title"] != prev_title: + if lines: + lines.append({"kind": "blank", "text": ""}) + for hl in wrap(p["title"].replace("_", " ").upper(), col_w): + lines.append({"kind": "heading", "text": hl}) + prev_title = p["title"] + elif lines: + lines.append({"kind": "blank", "text": ""}) + for bl in wrap(p["ctx"], col_w): + lines.append({"kind": "body", "text": bl}) + return lines + + +def pack_pages(paras: list[dict], col_w: int, max_lines: int) -> list[tuple[int, int]]: + """Greedy paragraph-aligned packing: [(i, j)] para ranges, one per page.""" + pages = [] + i = 0 + while i < len(paras): + j = i + 1 + while j < len(paras) and len(layout_page(paras[i : j + 1], col_w)) <= max_lines: + j += 1 + pages.append((i, j)) + i = j + return pages + + +def _sentence_colors(lines: list[dict], palette: list) -> list[list[tuple[int, int, int]]]: + """Per-line per-char glyph color cycling hue per sentence across the page.""" + joined = "\n".join(ln["text"] for ln in lines) + idx, out_idx = 0, [] + for i, ch in enumerate(joined): + out_idx.append(idx) + if ch in ".!?" and i + 1 < len(joined) and joined[i + 1] in " \n": + idx += 1 + colors, pos = [], 0 + for ln in lines: + n = len(ln["text"]) + colors.append([palette[out_idx[pos + k] % 6] for k in range(n)]) + pos += n + 1 # the joining newline + return colors + + +def render_doc(lines: list[dict], font: FontCfg, size: int, variant: str, cache: Path) -> Image.Image: + """Two-column page: left column rows top-to-bottom, then right column. + + variant "dark-sent": black page, body glyphs in bright sentence hues, + headings bright white double-strike. variant "sent": white page, body in + dark sentence hues, headings black double-strike. The page background + covers the full image, so a padded cell (pitch > glyph height) is dark + edge-to-edge in dark mode. + """ + glyphs, font_ascent = parse_bdf(ensure_font(font, cache)) + ascent = font.ascent if font.ascent is not None else font_ascent + cols, rows, _ = capacity(font, size) + col_w = (cols - GUTTER) // 2 + dark = variant == "dark-sent" + bg, heading_fg = (_BLACK, _WHITE) if dark else (_WHITE, _BLACK) + sent_colors = _sentence_colors(lines, _BRIGHT if dark else _DARK) + img = Image.new("RGB", (size, size), bg) + px = img.load() + for li, ln in enumerate(lines): + column, row = divmod(li, rows) + if column > 1: + break # overflow guard; pack_pages should prevent this + x_origin = column * (col_w + GUTTER) * font.adv + y0 = row * font.pitch + for ci, ch in enumerate(ln["text"]): + glyph = glyphs.get(ord(ch)) + if glyph is None: + continue + fg = heading_fg if ln["kind"] == "heading" else sent_colors[li][ci] + w, h, xoff, yoff = glyph["bbx"] + top = y0 + ascent - h - yoff + shift = 0x80 if w <= 8 else 0x8000 + strikes = (0, 1) if ln["kind"] == "heading" else (0,) + for dx in strikes: + for r, bits in enumerate(glyph["rows"]): + y = top + r + if not 0 <= y < size: + continue + for b in range(w): + if bits & (shift >> b): + x = x_origin + ci * font.adv + xoff + b + dx + if 0 <= x < size: + px[x, y] = fg + return img + + +# --- runner ------------------------------------------------------------------ + + +def save_png(png: Path, img_fn) -> None: + if png.exists() and png.stat().st_size > 0: + return + tmp = png.with_suffix(".tmp.png") + img_fn().save(tmp) + tmp.replace(png) + + +def run_grid_chunk(model: str, cond: str, start: int, end: int, ctx: dict) -> list[dict]: + """One row-major-grid chunk: render via bdf.render, QA, score.""" + args, flow, paras, offsets, keys = ctx["args"], ctx["flow"], ctx["paras"], ctx["offsets"], ctx["keys"] + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + _, font_key, variant = CONDITIONS[cond] + font = FONTS[font_key] + chunk_text = flow[start:end] + png = CACHE / f"exp19-{font_key}-{variant}-{sha8(chunk_text, str(args.size))}.png" + save_png(png, lambda: render(chunk_text, font, CACHE, args.size, variant)) + cols, rows, _ = capacity(font, args.size) + prompt = load_prompt("qa-image.md").format(cols=cols, rows=rows) + return qa_and_score(model, cond, prompt, png, questions, start, ctx) + + +def run_doc_page(model: str, cond: str, page: tuple[int, int], ctx: dict) -> list[dict]: + """One doc-layout page: typeset, render, QA, score.""" + args, paras, offsets, keys = ctx["args"], ctx["paras"], ctx["offsets"], ctx["keys"] + i, j = page + start = offsets[i] + end = offsets[j - 1] + len(paras[j - 1]["ctx"]) + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + _, font_key, variant = CONDITIONS[cond] + font = FONTS[font_key] + lines = ctx["lines"][cond][page] + key = sha8(cond, json.dumps([(p["title"], p["ctx"]) for p in paras[i:j]]), str(args.size)) + png = CACHE / f"exp19-doc-{font_key}-{variant}-{key}.png" + save_png(png, lambda: render_doc(lines, font, args.size, variant, CACHE)) + cols, rows, _ = capacity(font, args.size) + col_w = (cols - GUTTER) // 2 + prompt = load_prompt("exp19-qa-doc.md").format(col_w=col_w, rows=rows) + return qa_and_score(model, cond, prompt, png, questions, start, ctx) + + +def parse_answers(text: str, n: int) -> list[str]: + """parse_numbered, with a fallback for glm's intermittently unnumbered output. + + glm-4.6v often ignores the "numbered list" instruction and emits plain + answer lines in question order. When numbered parsing recovers fewer than + half the answers and the response is a clean <=n line list, map lines + positionally instead (preamble line ending in ':' is dropped). + """ + nums = squad.parse_numbered(text, n) + if sum(bool(a) for a in nums) >= max(1, n // 2): + return nums + lines = [ln.strip() for ln in text.splitlines() if ln.strip()] + if lines and lines[0].endswith(":"): + lines = lines[1:] + if 0 < len(lines) <= n: + return lines + [""] * (n - len(lines)) + return nums + + +def qa_and_score(model: str, cond: str, prompt: str, png: Path, questions: list[dict], start: int, ctx: dict) -> list[dict]: + args, keys = ctx["args"], ctx["keys"] + q_block = "\n".join(f"{k + 1}. {q['q']}" for k, q in enumerate(questions)) + messages = [ + { + "role": "user", + "content": [ + {"text": prompt}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + qa = cached( + model, "exp19-qa", {"messages": messages, "effort": args.effort}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete(keys, model, messages, max_tokens=args.max_tokens, effort=args.effort), + ) + ), + args.fresh, + ) + answers = parse_answers(qa["text"], len(questions)) + records = [] + for q, a in zip(questions, answers): + records.append( + { + "model": model, + "length": ctx["length"], + "cond": cond, + "chunk": start, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + ) + records[0]["usage"] = [{"phase": "qa", **qa["usage"]}] + return records + + +def aggregate(records: list[dict], price_in: float, price_out: float) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean_f1 = sum(f1s) / n + se = (sum((x - mean_f1) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + us = [u for r in records if "usage" in r for u in r["usage"]] + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + cost_in = (tok["in"] + 1.25 * tok["cache_w"] + 0.1 * tok["cache_r"]) / 1e6 * price_in + cost_out = tok["out"] / 1e6 * price_out + return { + "n": n, + "em": sum(r["em"] for r in records) / n, + "f1": mean_f1, + "f1_se": se, + "abstained": sum(r["abstained"] for r in records), + **{f"tok_{k}": v for k, v in tok.items()}, + "cost_in_usd": round(cost_in, 4), + "cost_out_usd": round(cost_out, 4), + "cost_usd": round(cost_in + cost_out, 4), + } + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--models", default=",".join(MODELS)) + ap.add_argument("--lengths", default=",".join(map(str, LENGTHS))) + ap.add_argument("--conditions", default=",".join(CONDITIONS)) + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--effort", default=None) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--render-only", action="store_true", help="render sample pages + capacity stats, no API") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + out_dir = RESULTS / "exp19-bestglm" + out_dir.mkdir(parents=True, exist_ok=True) + + models = [m.strip() for m in args.models.split(",") if m.strip()] + lengths = [int(x) for x in args.lengths.split(",") if x.strip()] + conditions = [c.strip() for c in args.conditions.split(",") if c.strip()] + for c in conditions: + if c not in CONDITIONS: + sys.exit(f"unknown condition: {c}") + + keys = {} + if not args.render_only: + keys["openrouter"] = load_env_key("OPENROUTER_API_KEY", args.env) + + all_paras = squad.load_paragraphs(CACHE) + tasks = [] + capacity_stats = {} + for length in lengths: + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + # Doc page packing per font (max_lines differs with pitch). + page_lines: dict[str, dict] = {} + doc_stats = {} + for cond in conditions: + kind, font_key, _ = CONDITIONS[cond] + if kind != "doc": + continue + font = FONTS[font_key] + cols, rows, _cap = capacity(font, args.size) + col_w = (cols - GUTTER) // 2 + pages = pack_pages(paras, col_w, 2 * rows) + page_lines[cond] = {pg: layout_page(paras[pg[0] : pg[1]], col_w) for pg in pages} + chars = [offsets[j - 1] + len(paras[j - 1]["ctx"]) - offsets[i] for i, j in pages] + doc_stats[cond] = { + "pages": len(pages), + "mean_chars_page": round(sum(chars) / len(pages)), + "col_w": col_w, + "rows": rows, + } + capacity_stats[length] = { + "corpus_chars": len(flow), + "grid": { + fk: dict(zip(("cols", "rows", "chars"), capacity(FONTS[fk], args.size))) for fk in FONTS + }, + "doc": doc_stats, + } + ctx = { + "args": args, "flow": flow, "paras": paras, "offsets": offsets, "keys": keys, + "length": length, "lines": page_lines, + } + for model in models: + for cond in conditions: + kind, font_key, _ = CONDITIONS[cond] + if kind == "grid": + budget = capacity(FONTS[font_key], args.size)[2] + for start in range(0, len(flow), budget): + tasks.append(("grid", model, cond, (start, min(start + budget, len(flow))), ctx)) + else: + for pg in page_lines[cond]: + tasks.append(("doc", model, cond, pg, ctx)) + + for length, st in capacity_stats.items(): + print(f"len {length}: corpus {st['corpus_chars']} chars") + for fk, g in st["grid"].items(): + print(f" grid {fk}: {g['cols']}x{g['rows']} = {g['chars']} chars/page") + for cond, d in st["doc"].items(): + print(f" {cond}: {d['pages']} pages, mean {d['mean_chars_page']} chars/page (2x{d['col_w']}w, {d['rows']} rows)") + + if args.render_only: + for length in lengths: + ctx = next(t[4] for t in tasks if t[4]["length"] == length) + for cond in conditions: + kind, font_key, variant = CONDITIONS[cond] + if kind == "grid": + budget = capacity(FONTS[font_key], args.size)[2] + chunk_text = ctx["flow"][:budget] + png = CACHE / f"exp19-{font_key}-{variant}-{sha8(chunk_text, str(args.size))}.png" + save_png(png, lambda: render(chunk_text, FONTS[font_key], CACHE, args.size, variant)) + else: + pg = next(iter(ctx["lines"][cond])) + i, j = pg + key = sha8(cond, json.dumps([(p["title"], p["ctx"]) for p in ctx["paras"][i:j]]), str(args.size)) + png = CACHE / f"exp19-doc-{font_key}-{variant}-{key}.png" + save_png(png, lambda: render_doc(ctx["lines"][cond][pg], FONTS[font_key], args.size, variant, CACHE)) + print(f" sample: {png}") + return + + print(f"grid: {len(tasks)} page/chunk tasks") + records: list[dict] = [] + done = 0 + with ThreadPoolExecutor(args.workers) as pool: + futures = [ + pool.submit(run_grid_chunk, m, c, u[0], u[1], ctx) if kind == "grid" else pool.submit(run_doc_page, m, c, u, ctx) + for kind, m, c, u, ctx in tasks + ] + for fut in futures: + records.extend(fut.result()) + done += 1 + print(f" {done}/{len(tasks)} tasks", flush=True) + + # Merge with prior invocations: replace the cells we just ran, keep the rest. + ran = {(m, ln, c) for ln in lengths for m in models for c in conditions} + rec_path = out_dir / "records.jsonl" + if rec_path.exists(): + old = [json.loads(ln) for ln in rec_path.read_text().splitlines() if ln.strip()] + records = [r for r in old if (r["model"], r["length"], r["cond"]) not in ran] + records + with rec_path.open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + + cells = [] + for model in sorted({r["model"] for r in records}): + for length in sorted({r["length"] for r in records}): + for cond in CONDITIONS: + sub = [r for r in records if r["model"] == model and r["length"] == length and r["cond"] == cond] + if not sub: + continue + cells.append({"model": model, "length": length, "condition": cond, **aggregate(sub, *MODELS[model])}) + (out_dir / "summary.json").write_text( + json.dumps({"args": vars(args), "capacity": capacity_stats, "cells": cells}, indent=1) + ) + with (out_dir / "matrix.csv").open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=list(cells[0].keys())) + writer.writeheader() + writer.writerows(cells) + + for c in cells: + print( + f"{c['model']:<16} len {c['length']:<4} {c['condition']:<26} " + f"n={c['n']:<4} EM {c['em']:.3f} F1 {c['f1']:.3f} ±{c['f1_se']:.3f} ${c['cost_usd']:.4f}" + ) + print(f"\n-> {out_dir}/records.jsonl, matrix.csv, summary.json") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp20_8x8u.py b/packages/snapcompact/research/exp20_8x8u.py new file mode 100644 index 000000000..a23c063e5 --- /dev/null +++ b/packages/snapcompact/research/exp20_8x8u.py @@ -0,0 +1,362 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp20: swap the per-model round-2 winner font for unscii-8 (8x8u). + +For each model, re-run its best-known config with FontCfg("8x8u","unscii-8",8,8) +in place of the winning font, everything else unchanged (layout, variant, +1568px, seed 42, qpc 30). 8x8u @1568: grid 196x196 = 38,416 chars/page +(~2x doc-8on16's effective capacity); doc 2x96 cols x 196 rows. + + model config swapped from exp20 condition + gpt-5.5 doc-8on16-bw -> img-doc-8x8u-bw + google/gemini-3.5-flash doc-8on16-sent-dim -> img-doc-8x8u-sent-dim + moonshotai/kimi-k2.6 doc8on16-sent-dim -> img-doc-8x8u-sent-dim + z-ai/glm-4.6v doc-8on16-sent -> img-doc-8x8u-sent + claude-fable-5 grid 6x12-dim -> img-8x8u-dim + claude-opus-4-8 grid 8x13-bw -> img-8x8u-bw + +Usage: + uv run exp20_8x8u.py --model gpt-5.5 --render-only # sample PNG, no API + uv run exp20_8x8u.py --model gpt-5.5 # lengths 50,150,250 +""" + +import argparse +import os +import json +import sys +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +from PIL import Image + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import _DARK, _DIMMED, FontCfg, _stopword_mask, capacity, load_font, render # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 + +EXP = "exp20" +OUT_DIR = RESULTS / f"{EXP}-8x8u" +FONT = FontCfg("8x8u", "unscii-8", 8, 8) +GUTTER = 3 +_WHITE = (255, 255, 255) +_BLACK = (0, 0, 0) +_INK = (24, 24, 24) + +# model -> (cond, layout, variant, price_in, price_out, key_name) +CONFIGS = { + "gpt-5.5": ("img-doc-8x8u-bw", "doc", "bw", 2.0, 16.0, "openai"), + "google/gemini-3.5-flash": ("img-doc-8x8u-sent-dim", "doc", "sent-dim", 0.6, 4.0, "openrouter"), + "moonshotai/kimi-k2.6": ("img-doc-8x8u-sent-dim", "doc", "sent-dim", 0.68, 3.41, "openrouter"), + "z-ai/glm-4.6v": ("img-doc-8x8u-sent", "doc", "sent", 0.30, 0.90, "openrouter"), + "claude-fable-5": ("img-8x8u-dim", "grid", "dim", 10.0, 50.0, "anthropic"), + "claude-opus-4-8": ("img-8x8u-bw", "grid", "bw", 15.0, 75.0, "anthropic"), +} +KEY_ENV = {"openai": "OPENAI_API_KEY", "openrouter": "OPENROUTER_API_KEY", "anthropic": "ANTHROPIC_API_KEY"} + + +def slug(model: str) -> str: + return model.split("/")[-1] + + +def cached(model: str, payload: object, fn, fresh: bool) -> dict: + key = sha8(model, f"{EXP}-qa", json.dumps(payload, sort_keys=True, default=str)) + path = QA_CACHE / f"{key}.json" + if not fresh and path.exists(): + return json.loads(path.read_text()) + out = fn() + if out.get("stop") != "max_tokens": + tmp = path.with_suffix(".tmp.json") + tmp.write_text(json.dumps(out)) + tmp.replace(path) + else: + print(f"WARN truncated response, not cached ({model})", flush=True) + return out + + +# --- document layout (exp14's renderer, generalized for hex fonts + dim variants) --- + + +def wrap(text: str, width: int) -> list[str]: + lines, cur = [], "" + for word in text.split(): + while len(word) > width: + if cur: + lines.append(cur) + cur = "" + lines.append(word[:width]) + word = word[width:] + if not cur: + cur = word + elif len(cur) + 1 + len(word) <= width: + cur += " " + word + else: + lines.append(cur) + cur = word + if cur: + lines.append(cur) + return lines + + +def layout_page(paras: list[dict], col_w: int) -> list[dict]: + lines: list[dict] = [] + prev_title = None + for p in paras: + if p["title"] != prev_title: + if lines: + lines.append({"kind": "blank", "text": ""}) + for hl in wrap(p["title"].replace("_", " ").upper(), col_w): + lines.append({"kind": "heading", "text": hl}) + prev_title = p["title"] + elif lines: + lines.append({"kind": "blank", "text": ""}) + for bl in wrap(p["ctx"], col_w): + lines.append({"kind": "body", "text": bl}) + return lines + + +def pack_pages(paras: list[dict], col_w: int, max_lines: int) -> list[tuple[int, int]]: + pages = [] + i = 0 + while i < len(paras): + j = i + 1 + while j < len(paras) and len(layout_page(paras[i : j + 1], col_w)) <= max_lines: + j += 1 + pages.append((i, j)) + i = j + return pages + + +def _char_styles(lines: list[dict], variant: str) -> list[list[tuple[int, int, int]]]: + """Per-line per-char body glyph color for sent / dim composition.""" + joined = "\n".join(ln["text"] for ln in lines) + sent_idx = None + if "sent" in variant: + sent_idx, idx = [], 0 + for i, ch in enumerate(joined): + sent_idx.append(idx) + if ch in ".!?" and i + 1 < len(joined) and joined[i + 1] in " \n": + idx += 1 + dim = _stopword_mask(joined) if "dim" in variant else None + colors, pos = [], 0 + for ln in lines: + row = [] + for k in range(len(ln["text"])): + i = pos + k + if dim is not None and dim[i]: + row.append(_DIMMED) + elif sent_idx is not None: + row.append(_DARK[sent_idx[i] % 6]) + else: + row.append(_INK) + colors.append(row) + pos += len(ln["text"]) + 1 + return colors + + +def render_doc(lines: list[dict], size: int, variant: str, cache: Path) -> Image.Image: + glyphs, font_ascent = load_font(FONT, cache) + ascent = FONT.ascent if FONT.ascent is not None else font_ascent + cols, rows, _ = capacity(FONT, size) + col_w = (cols - GUTTER) // 2 + styles = _char_styles(lines, variant) + img = Image.new("RGB", (size, size), _WHITE) + px = img.load() + for li, ln in enumerate(lines): + column, row = divmod(li, rows) + if column > 1: + break + x_origin = column * (col_w + GUTTER) * FONT.adv + y0 = row * FONT.pitch + for ci, ch in enumerate(ln["text"]): + glyph = glyphs.get(ord(ch)) + if glyph is None: + continue + fg = _BLACK if ln["kind"] == "heading" else styles[li][ci] + w, h, xoff, yoff = glyph["bbx"] + top = y0 + ascent - h - yoff + shift = 0x80 if w <= 8 else 0x8000 + strikes = (0, 1) if ln["kind"] == "heading" else (0,) + for dx in strikes: + for r, bits in enumerate(glyph["rows"]): + y = top + r + if not 0 <= y < size: + continue + for b in range(w): + if bits & (shift >> b): + x = x_origin + ci * FONT.adv + xoff + b + dx + if 0 <= x < size: + px[x, y] = fg + return img + + +# --- runner ------------------------------------------------------------------- + + +def atomic_save(img: Image.Image, png: Path) -> None: + tmp = png.with_suffix(f".{os.getpid()}.tmp.png") # pid-unique: parallel models share sent-dim PNGs + img.save(tmp) + tmp.replace(png) + + +def parse_answers(text: str, n: int) -> list[str]: + """parse_numbered + exp19's positional fallback (glm drops numbering).""" + nums = squad.parse_numbered(text, n) + if sum(bool(a) for a in nums) >= max(1, n // 2): + return nums + lines = [ln.strip() for ln in text.splitlines() if ln.strip()] + if lines and lines[0].endswith(":"): + lines = lines[1:] + if 0 < len(lines) <= n: + return lines + [""] * (n - len(lines)) + return nums + + +def qa_unit(model: str, cond: str, prompt: str, png: Path, questions: list[dict], length: int, start: int, ctx: dict) -> list[dict]: + args, keys = ctx["args"], ctx["keys"] + q_block = "\n".join(f"{k + 1}. {q['q']}" for k, q in enumerate(questions)) + messages = [{"role": "user", "content": [{"text": prompt}, {"image_path": png}, {"text": q_block}]}] + qa = cached( + model, {"messages": messages, "effort": None}, + lambda: dict(zip(("text", "usage", "stop"), llm_complete(keys, model, messages, max_tokens=args.max_tokens))), + args.fresh, + ) + answers = parse_answers(qa["text"], len(questions)) + records = [] + for q, a in zip(questions, answers): + records.append({ + "model": model, "length": length, "cond": cond, "chunk": start, + "pos_rel": q["pos_rel"], "q": q["q"], "answer": a, "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + }) + records[0]["usage"] = [{"phase": "qa", **qa["usage"]}] + return records + + +def aggregate(records: list[dict], price_in: float, price_out: float) -> list[dict]: + cells = {} + for r in records: + cells.setdefault((r["length"], r["cond"]), []).append(r) + out = [] + for (length, cond), recs in sorted(cells.items()): + n = len(recs) + f1s = [r["f1"] for r in recs] + mean = sum(f1s) / n + var = sum((x - mean) ** 2 for x in f1s) / (n - 1) if n > 1 else 0.0 + usage = [u for r in recs for u in r.get("usage", [])] + tin = sum(u["in"] for u in usage) + tout = sum(u["out"] for u in usage) + creads = sum(u.get("cache_r", 0) for u in usage) + rsn = sum(u.get("reasoning", 0) for u in usage) + cost = (tin + 0.1 * creads) * price_in / 1e6 + tout * price_out / 1e6 + out.append({ + "model": recs[0]["model"], "length": length, "condition": cond, "n": n, + "em": round(sum(r["em"] for r in recs) / n, 4), "f1": round(mean, 4), + "f1_se": round((var / n) ** 0.5, 4), "abstained": sum(r["abstained"] for r in recs), + "tok_in": tin, "tok_out": tout, "tok_cache_r": creads, "tok_reasoning": rsn, + "cost_usd": round(cost, 4), + }) + return out + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--model", required=True, choices=sorted(CONFIGS)) + ap.add_argument("--lengths", default="50,150,250") + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--render-only", action="store_true") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + cond, layout, variant, price_in, price_out, key_name = CONFIGS[args.model] + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + OUT_DIR.mkdir(parents=True, exist_ok=True) + + cols, rows, grid_cap = capacity(FONT, args.size) + col_w = (cols - GUTTER) // 2 + max_lines = 2 * rows + print(f"{args.model}: {cond} ({layout}/{variant}); 8x8u grid {cols}x{rows}={grid_cap}, " + f"doc 2x{col_w}+g{GUTTER}, {max_lines} slots", flush=True) + + keys = {} if args.render_only else {key_name: load_env_key(KEY_ENV[key_name], args.env)} + all_paras = squad.load_paragraphs(CACHE) + tasks = [] + cap_stats = {} + for length in (int(x) for x in args.lengths.split(",")): + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + ctx = {"args": args, "keys": keys} + if layout == "doc": + pages = pack_pages(paras, col_w, max_lines) + page_chars = [offsets[j - 1] + len(paras[j - 1]["ctx"]) - offsets[i] for i, j in pages] + cap_stats[length] = {"pages": len(pages), "mean_chars_page": round(sum(page_chars) / len(pages))} + prompt = load_prompt("exp04-qa-image.md").format(col_w=col_w, rows=rows) + for i, j in pages: + start = offsets[i] + end = offsets[j - 1] + len(paras[j - 1]["ctx"]) + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + continue + lines = layout_page(paras[i:j], col_w) + key = sha8(cond, json.dumps([(p["title"], p["ctx"]) for p in paras[i:j]]), str(args.size)) + png = CACHE / f"{EXP}-doc-{variant}-{key}.png" + if not png.exists() or png.stat().st_size == 0: + atomic_save(render_doc(lines, args.size, variant, CACHE), png) + tasks.append((args.model, cond, prompt, png, questions, length, start, ctx)) + else: + cap_stats[length] = {"pages": -(-len(flow) // grid_cap), "mean_chars_page": grid_cap} + prompt = load_prompt("qa-image.md").format(cols=cols, rows=rows) + for start in range(0, len(flow), grid_cap): + end = min(start + grid_cap, len(flow)) + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + continue + png = CACHE / f"{EXP}-grid-{variant}-{sha8(flow[start:end], str(args.size))}.png" + if not png.exists() or png.stat().st_size == 0: + atomic_save(render(flow[start:end], FONT, CACHE, args.size, variant), png) + tasks.append((args.model, cond, prompt, png, questions, length, start, ctx)) + print(f" len {length}: {cap_stats[length]['pages']} pages, " + f"mean {cap_stats[length]['mean_chars_page']} chars/page, corpus {len(flow)}", flush=True) + + if args.render_only: + print(f"sample: {tasks[0][3]}" if tasks else "no tasks") + return + + records: list[dict] = [] + with ThreadPoolExecutor(args.workers) as pool: + futures = [pool.submit(qa_unit, *t) for t in tasks] + for done, fut in enumerate(futures, 1): + records.extend(fut.result()) + print(f" {done}/{len(futures)}", flush=True) + + s = slug(args.model) + with (OUT_DIR / f"records-{s}.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + cells = aggregate(records, price_in, price_out) + hdr = "model,length,condition,n,em,f1,f1_se,abstained,tok_in,tok_out,tok_cache_r,tok_reasoning,cost_usd" + with (OUT_DIR / f"matrix-{s}.csv").open("w") as fh: + fh.write(hdr + "\n") + for c in cells: + fh.write(",".join(str(c[k]) for k in hdr.split(",")) + "\n") + (OUT_DIR / f"summary-{s}.json").write_text(json.dumps({"args": vars(args), "capacity": cap_stats, "cells": cells}, indent=1)) + for c in cells: + print(f"len {c['length']:<4} {c['condition']:<24} n={c['n']:<4} EM {c['em']:.3f} " + f"F1 {c['f1']:.3f} ±{c['f1_se']:.3f} ${c['cost_usd']:.3f} " + f"out={c['tok_out']} rsn={c['tok_reasoning']}", flush=True) + print(f"-> {OUT_DIR}/records-{s}.jsonl", flush=True) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp21_braille.py b/packages/snapcompact/research/exp21_braille.py new file mode 100644 index 000000000..d553c7be0 --- /dev/null +++ b/packages/snapcompact/research/exp21_braille.py @@ -0,0 +1,276 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp21: render text as Grade-1 (uncontracted) 6-dot braille -> gemini-3.5-flash. + +Each character becomes a 2x3 dot matrix drawn directly (no font file): +lowercase letters, digits as number-sign + a-j (one sign per digit run), +a punctuation subset; everything else -> blank cell. Text is lowercased +(SQuAD scoring is case-insensitive; real braille capital signs would +waste cells). + +Conditions (gemini only, 1568px, bw): + img-braille-5x7 1px dots, cell 5x7 -> 313x224 = 70,112 cells/page + img-braille-7x10 2px dots, cell 7x10 -> 224x156 = 34,944 cells/page + +Usage: + uv run exp21_braille.py --render-only + uv run exp21_braille.py # lengths 50,150 +""" + +import argparse +import json +import os +import sys +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +from PIL import Image + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 + +EXP = "exp21" +OUT_DIR = RESULTS / f"{EXP}-braille" +MODELS = { # model -> (price_in, price_out, key_name, key_env) + "google/gemini-3.5-flash": (0.6, 4.0, "openrouter", "OPENROUTER_API_KEY"), + "gpt-5.5": (2.0, 16.0, "openai", "OPENAI_API_KEY"), +} +MODEL = "google/gemini-3.5-flash" +PRICE_IN, PRICE_OUT = 0.6, 4.0 +_WHITE = (255, 255, 255) +_BLACK = (0, 0, 0) + +# dots numbered 1-6: 1=top-left 2=mid-left 3=bottom-left 4=top-right 5=mid-right 6=bottom-right +# bitmask: bit0=dot1 .. bit5=dot6 (matches Unicode U+2800 offsets) +_L = { + "a": 0x01, "b": 0x03, "c": 0x09, "d": 0x19, "e": 0x11, "f": 0x0B, "g": 0x1B, + "h": 0x13, "i": 0x0A, "j": 0x1A, "k": 0x05, "l": 0x07, "m": 0x0D, "n": 0x1D, + "o": 0x15, "p": 0x0F, "q": 0x1F, "r": 0x17, "s": 0x0E, "t": 0x1E, "u": 0x25, + "v": 0x27, "w": 0x3A, "x": 0x2D, "y": 0x3D, "z": 0x35, +} +_PUNCT = { + ".": 0x32, ",": 0x02, "'": 0x04, "-": 0x24, ":": 0x12, ";": 0x06, + "?": 0x26, "!": 0x16, " ": 0x00, +} +_NUMSIGN = 0x3C # dots 3456 +_DIGIT = {d: _L["abcdefghij"[i]] for i, d in enumerate("1234567890")} + +# cell name -> (dot_px, adv, pitch); dot gap is 1px in both configs +CELLS = { + "5x7": (1, 5, 7), + "7x10": (2, 7, 10), +} + + +def braille_cells(text: str) -> tuple[list[int], list[int]]: + """(cell bitmasks, original char index per cell). Lowercases; digit runs share one number sign.""" + cells, origin = [], [] + in_num = False + for i, ch in enumerate(text): + c = ch.lower() + if c in _DIGIT: + if not in_num: + cells.append(_NUMSIGN) + origin.append(i) + in_num = True + cells.append(_DIGIT[c]) + origin.append(i) + continue + in_num = False + cells.append(_L.get(c, _PUNCT.get(c, 0x00))) + origin.append(i) + return cells, origin + + +def render_braille(cells: list[int], cell_name: str, size: int) -> Image.Image: + dpx, adv, pitch = CELLS[cell_name] + cols, rows = size // adv, size // pitch + step = dpx + 1 # dot pitch inside the cell + img = Image.new("RGB", (size, size), _WHITE) + px = img.load() + for idx, mask in enumerate(cells[: cols * rows]): + if not mask: + continue + row, col = divmod(idx, cols) + x0, y0 = col * adv, row * pitch + for bit in range(6): + if not mask & (1 << bit): + continue + dc, dr = divmod(bit, 3) # dots 1-3 left column, 4-6 right column + dx, dy = x0 + dc * step, y0 + dr * step + for yy in range(dy, dy + dpx): + for xx in range(dx, dx + dpx): + if xx < size and yy < size: + px[xx, yy] = _BLACK + return img + + +def atomic_save(img: Image.Image, png: Path) -> None: + tmp = png.with_suffix(f".{os.getpid()}.tmp.png") + img.save(tmp) + tmp.replace(png) + + +def cached(payload: object, fn, fresh: bool) -> dict: + key = sha8(MODEL, f"{EXP}-qa", json.dumps(payload, sort_keys=True, default=str)) + path = QA_CACHE / f"{key}.json" + if not fresh and path.exists(): + return json.loads(path.read_text()) + out = fn() + if out.get("stop") != "max_tokens": + tmp = path.with_suffix(".tmp.json") + tmp.write_text(json.dumps(out)) + tmp.replace(path) + else: + print("WARN truncated response, not cached", flush=True) + return out + + +def qa_unit(cond: str, prompt: str, png: Path, questions: list[dict], length: int, start: int, ctx: dict) -> list[dict]: + args, keys = ctx["args"], ctx["keys"] + q_block = "\n".join(f"{k + 1}. {q['q']}" for k, q in enumerate(questions)) + messages = [{"role": "user", "content": [{"text": prompt}, {"image_path": png}, {"text": q_block}]}] + payload = {"messages": messages} + if args.effort: + payload["effort"] = args.effort + qa = cached( + payload, + lambda: dict(zip(("text", "usage", "stop"), + llm_complete(keys, MODEL, messages, max_tokens=args.max_tokens, effort=args.effort))), + args.fresh, + ) + answers = squad.parse_numbered(qa["text"], len(questions)) + records = [] + for q, a in zip(questions, answers): + records.append({ + "model": MODEL, "length": length, "cond": cond, "chunk": start, + "pos_rel": q["pos_rel"], "q": q["q"], "answer": a, "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + }) + records[0]["usage"] = [{"phase": "qa", **qa["usage"]}] + return records + + +def aggregate(records: list[dict]) -> list[dict]: + cells = {} + for r in records: + cells.setdefault((r["length"], r["cond"]), []).append(r) + out = [] + for (length, cond), recs in sorted(cells.items()): + n = len(recs) + f1s = [r["f1"] for r in recs] + mean = sum(f1s) / n + var = sum((x - mean) ** 2 for x in f1s) / (n - 1) if n > 1 else 0.0 + usage = [u for r in recs for u in r.get("usage", [])] + tin = sum(u["in"] for u in usage) + tout = sum(u["out"] for u in usage) + creads = sum(u.get("cache_r", 0) for u in usage) + rsn = sum(u.get("reasoning", 0) for u in usage) + cost = (tin + 0.1 * creads) * PRICE_IN / 1e6 + tout * PRICE_OUT / 1e6 + out.append({ + "model": MODEL, "length": length, "condition": cond, "n": n, + "em": round(sum(r["em"] for r in recs) / n, 4), "f1": round(mean, 4), + "f1_se": round((var / n) ** 0.5, 4), "abstained": sum(r["abstained"] for r in recs), + "tok_in": tin, "tok_out": tout, "tok_cache_r": creads, "tok_reasoning": rsn, + "cost_usd": round(cost, 4), + }) + return out + + +def main() -> None: + global MODEL, PRICE_IN, PRICE_OUT + ap = argparse.ArgumentParser() + ap.add_argument("--model", default="google/gemini-3.5-flash", choices=sorted(MODELS)) + ap.add_argument("--cells", default="5x7,7x10") + ap.add_argument("--lengths", default="50,150") + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--render-only", action="store_true") + ap.add_argument("--env", default="~/.env") + ap.add_argument("--effort", default=None) + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + OUT_DIR.mkdir(parents=True, exist_ok=True) + + MODEL = args.model + PRICE_IN, PRICE_OUT, key_name, key_env = MODELS[MODEL] + keys = {} if args.render_only else {key_name: load_env_key(key_env, args.env)} + all_paras = squad.load_paragraphs(CACHE) + prompt_tpl = load_prompt("exp21-qa-braille.md") + + tasks = [] + for cell_name in args.cells.split(","): + dpx, adv, pitch = CELLS[cell_name] + cols, rows = args.size // adv, args.size // pitch + cap = cols * rows + cond = f"img-braille-{cell_name}" + (f"+eff-{args.effort}" if args.effort else "") + for length in (int(x) for x in args.lengths.split(",")): + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + cells, origin = braille_cells(flow) + pages = [] + i = 0 + while i < len(cells): + j = min(i + cap, len(cells)) + pages.append((i, j)) + i = j + print(f"{cond} len {length}: {len(pages)} pages, {cap} cells/page " + f"({cols}x{rows}), {len(cells)} cells for {len(flow)} chars", flush=True) + ctx = {"args": args, "keys": keys} + for ci, cj in pages: + start = origin[ci] + end = origin[cj - 1] + 1 + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + continue + png = CACHE / f"{EXP}-{cell_name}-{sha8(flow[start:end], cell_name, str(args.size))}.png" + if not png.exists() or png.stat().st_size == 0: + atomic_save(render_braille(cells[ci:cj], cell_name, args.size), png) + prompt = prompt_tpl.format(cols=cols, rows=rows) + tasks.append((cond, prompt, png, questions, length, start, ctx)) + + if args.render_only: + for t in tasks[:2]: + print(f"sample: {t[2]}") + return + + records: list[dict] = [] + with ThreadPoolExecutor(args.workers) as pool: + futures = [pool.submit(qa_unit, *t) for t in tasks] + for done, fut in enumerate(futures, 1): + records.extend(fut.result()) + print(f" {done}/{len(futures)}", flush=True) + + slug = MODEL.split("/")[-1] + (f"-eff{args.effort}" if args.effort else "") + with (OUT_DIR / f"records-{slug}.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + cells_out = aggregate(records) + hdr = "model,length,condition,n,em,f1,f1_se,abstained,tok_in,tok_out,tok_cache_r,tok_reasoning,cost_usd" + with (OUT_DIR / f"matrix-{slug}.csv").open("w") as fh: + fh.write(hdr + "\n") + for c in cells_out: + fh.write(",".join(str(c[k]) for k in hdr.split(",")) + "\n") + (OUT_DIR / f"summary-{slug}.json").write_text(json.dumps({"args": vars(args), "cells": cells_out}, indent=1)) + for c in cells_out: + print(f"len {c['length']:<4} {c['condition']:<20} n={c['n']:<4} EM {c['em']:.3f} " + f"F1 {c['f1']:.3f} ±{c['f1_se']:.3f} ${c['cost_usd']:.3f} " + f"out={c['tok_out']} rsn={c['tok_reasoning']}", flush=True) + print(f"-> {OUT_DIR}/matrix-{slug}.csv", flush=True) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp22_ttf6pt.py b/packages/snapcompact/research/exp22_ttf6pt.py new file mode 100644 index 000000000..6d95db910 --- /dev/null +++ b/packages/snapcompact/research/exp22_ttf6pt.py @@ -0,0 +1,243 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp22: 6pt antialiased vector type vs 1-bit bitmap fonts. + +Every prior condition used hand-hinted 1-bit bitmap fonts. Here the corpus is +rasterized with a real monospace TTF (Monaco, fallback DejaVu Sans Mono) at +tiny em sizes WITH greyscale antialiasing — the hypothesis being that AA +preserves sub-pixel shape information a VLM can exploit below the bitmap-font +legibility floor. + +Conditions (bw, grid layout, 1568px): + img-ttf6-bw em 6px (6pt @ 72dpi) + img-ttf8-bw em 8px (6pt @ 96dpi) +Cell metrics are measured from the font (advance x line height). + +Usage: + uv run exp22_ttf6pt.py --render-only + uv run exp22_ttf6pt.py --model gpt-5.5 + uv run exp22_ttf6pt.py --model google/gemini-3.5-flash +""" + +import argparse +import json +import math +import os +import sys +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +from PIL import Image, ImageDraw, ImageFont + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 + +EXP = "exp22" +OUT_DIR = RESULTS / f"{EXP}-ttf6pt" +MODELS = { # model -> (price_in, price_out, key_name, key_env) + "google/gemini-3.5-flash": (0.6, 4.0, "openrouter", "OPENROUTER_API_KEY"), + "gpt-5.5": (2.0, 16.0, "openai", "OPENAI_API_KEY"), +} +MODEL = "google/gemini-3.5-flash" +PRICE_IN, PRICE_OUT = 0.6, 4.0 +_WHITE = (255, 255, 255) +_BLACK = (0, 0, 0) +TTF_CANDIDATES = [ + "/System/Library/Fonts/Monaco.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf", +] +EM_SIZES = (6, 8) + + +def mono_font(em: int) -> ImageFont.FreeTypeFont: + for path in TTF_CANDIDATES: + if Path(path).exists(): + return ImageFont.truetype(path, em) + raise FileNotFoundError("no monospace TTF found") + + +def metrics(em: int) -> tuple[float, int, int, int]: + """(advance, pitch, cols, rows) at 1568px for the em size.""" + f = mono_font(em) + adv = f.getlength("0") + ascent, descent = f.getmetrics() + pitch = ascent + descent # tight leading; AA keeps rows separable + cols = int(1568 // adv) + rows = 1568 // pitch + return adv, pitch, cols, rows + + +def render_ttf(text: str, em: int, size: int) -> Image.Image: + f = mono_font(em) + adv, pitch, cols, rows = metrics(em) + img = Image.new("RGB", (size, size), _WHITE) + draw = ImageDraw.Draw(img) + for r in range(rows): + line = text[r * cols : (r + 1) * cols] + if not line: + break + draw.text((0, r * pitch), line, font=f, fill=_BLACK) + return img + + +def atomic_save(img: Image.Image, png: Path) -> None: + tmp = png.with_suffix(f".{os.getpid()}.tmp.png") + img.save(tmp) + tmp.replace(png) + + +def cached(payload: object, fn, fresh: bool) -> dict: + key = sha8(MODEL, f"{EXP}-qa", json.dumps(payload, sort_keys=True, default=str)) + path = QA_CACHE / f"{key}.json" + if not fresh and path.exists(): + return json.loads(path.read_text()) + out = fn() + if out.get("stop") != "max_tokens": + tmp = path.with_suffix(".tmp.json") + tmp.write_text(json.dumps(out)) + tmp.replace(path) + else: + print("WARN truncated response, not cached", flush=True) + return out + + +def qa_unit(cond: str, prompt: str, png: Path, questions: list[dict], length: int, start: int, ctx: dict) -> list[dict]: + args, keys = ctx["args"], ctx["keys"] + q_block = "\n".join(f"{k + 1}. {q['q']}" for k, q in enumerate(questions)) + messages = [{"role": "user", "content": [{"text": prompt}, {"image_path": png}, {"text": q_block}]}] + payload = {"messages": messages} + if args.effort: + payload["effort"] = args.effort + qa = cached( + payload, + lambda: dict(zip(("text", "usage", "stop"), + llm_complete(keys, MODEL, messages, max_tokens=args.max_tokens, effort=args.effort))), + args.fresh, + ) + answers = squad.parse_numbered(qa["text"], len(questions)) + records = [] + for q, a in zip(questions, answers): + records.append({ + "model": MODEL, "length": length, "cond": cond, "chunk": start, + "pos_rel": q["pos_rel"], "q": q["q"], "answer": a, "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + }) + records[0]["usage"] = [{"phase": "qa", **qa["usage"]}] + return records + + +def aggregate(records: list[dict]) -> list[dict]: + cells = {} + for r in records: + cells.setdefault((r["length"], r["cond"]), []).append(r) + out = [] + for (length, cond), recs in sorted(cells.items()): + n = len(recs) + f1s = [r["f1"] for r in recs] + mean = sum(f1s) / n + var = sum((x - mean) ** 2 for x in f1s) / (n - 1) if n > 1 else 0.0 + usage = [u for r in recs for u in r.get("usage", [])] + tin = sum(u["in"] for u in usage) + tout = sum(u["out"] for u in usage) + creads = sum(u.get("cache_r", 0) for u in usage) + rsn = sum(u.get("reasoning", 0) for u in usage) + cost = (tin + 0.1 * creads) * PRICE_IN / 1e6 + tout * PRICE_OUT / 1e6 + out.append({ + "model": MODEL, "length": length, "condition": cond, "n": n, + "em": round(sum(r["em"] for r in recs) / n, 4), "f1": round(mean, 4), + "f1_se": round((var / n) ** 0.5, 4), "abstained": sum(r["abstained"] for r in recs), + "tok_in": tin, "tok_out": tout, "tok_cache_r": creads, "tok_reasoning": rsn, + "cost_usd": round(cost, 4), + }) + return out + + +def main() -> None: + global MODEL, PRICE_IN, PRICE_OUT + ap = argparse.ArgumentParser() + ap.add_argument("--model", default="google/gemini-3.5-flash", choices=sorted(MODELS)) + ap.add_argument("--ems", default="6,8") + ap.add_argument("--lengths", default="50,150") + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--effort", default=None) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--render-only", action="store_true") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + OUT_DIR.mkdir(parents=True, exist_ok=True) + + MODEL = args.model + PRICE_IN, PRICE_OUT, key_name, key_env = MODELS[MODEL] + keys = {} if args.render_only else {key_name: load_env_key(key_env, args.env)} + all_paras = squad.load_paragraphs(CACHE) + prompt_tpl = load_prompt("qa-image.md") + + tasks = [] + for em in (int(x) for x in args.ems.split(",")): + adv, pitch, cols, rows = metrics(em) + cap = cols * rows + cond = f"img-ttf{em}-bw" + (f"+eff-{args.effort}" if args.effort else "") + print(f"{cond}: adv {adv:.2f}px pitch {pitch}px -> {cols}x{rows} = {cap} chars/page", flush=True) + for length in (int(x) for x in args.lengths.split(",")): + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + n_pages = math.ceil(len(flow) / cap) + print(f" len {length}: {n_pages} pages, corpus {len(flow)}", flush=True) + ctx = {"args": args, "keys": keys} + for start in range(0, len(flow), cap): + end = min(start + cap, len(flow)) + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + continue + png = CACHE / f"{EXP}-ttf{em}-{sha8(flow[start:end], str(em), str(args.size))}.png" + if not png.exists() or png.stat().st_size == 0: + atomic_save(render_ttf(flow[start:end], em, args.size), png) + prompt = prompt_tpl.format(cols=cols, rows=rows) + tasks.append((cond, prompt, png, questions, length, start, ctx)) + + if args.render_only: + for t in tasks[:2]: + print(f"sample: {t[2]}") + return + + records: list[dict] = [] + with ThreadPoolExecutor(args.workers) as pool: + futures = [pool.submit(qa_unit, *t) for t in tasks] + for done, fut in enumerate(futures, 1): + records.extend(fut.result()) + print(f" {done}/{len(futures)}", flush=True) + + slug = MODEL.split("/")[-1] + (f"-eff{args.effort}" if args.effort else "") + with (OUT_DIR / f"records-{slug}.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + cells_out = aggregate(records) + hdr = "model,length,condition,n,em,f1,f1_se,abstained,tok_in,tok_out,tok_cache_r,tok_reasoning,cost_usd" + with (OUT_DIR / f"matrix-{slug}.csv").open("w") as fh: + fh.write(hdr + "\n") + for c in cells_out: + fh.write(",".join(str(c[k]) for k in hdr.split(",")) + "\n") + (OUT_DIR / f"summary-{slug}.json").write_text(json.dumps({"args": vars(args), "cells": cells_out}, indent=1)) + for c in cells_out: + print(f"len {c['length']:<4} {c['condition']:<20} n={c['n']:<4} EM {c['em']:.3f} " + f"F1 {c['f1']:.3f} ±{c['f1_se']:.3f} ${c['cost_usd']:.3f} " + f"out={c['tok_out']} rsn={c['tok_reasoning']}", flush=True) + print(f"-> {OUT_DIR}/matrix-{slug}.csv", flush=True) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/final.py b/packages/snapcompact/research/final.py new file mode 100644 index 000000000..4df51f1dc --- /dev/null +++ b/packages/snapcompact/research/final.py @@ -0,0 +1,331 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""snapcompact final: the blog-post dataset. One command, full grid, CSV out. + +Grid: lengths x models x techniques, SQuAD v1.1 dev QA recall (official EM/F1). + + lengths 50 / 150 / 250 passages (~30k / ~102k / ~170k chars) + models claude-fable-5, claude-opus-4-8 (Anthropic Messages API) + gpt-5.5 (OpenAI Responses API) + techniques text plain-text chunks (ceiling) + handoff agent handoff document, QA on the doc + compact Anthropic: agent compaction-summary prompt + OpenAI: remote /responses/compact window + img-{6x10,5x8}-{sent,bw} one 1568x1568 PNG per chunk + +Outputs in results/final/: records.jsonl (per question), matrix.csv (per cell), +summary.json. Responses are cached by payload hash: interrupted or re-scoped +runs only bill new cells. `--report` reprints from cache without API calls. + +Usage: uv run final.py [--models ...] [--lengths 50,150,250] [--conditions ...] +Keys: ANTHROPIC_API_KEY + OPENAI_API_KEY from ~/.env. +""" + +import argparse +import csv +import json +import re +import sys +import uuid +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import capacity, render # noqa: E402 +from providers import is_openai, llm_complete, load_env_key, openai_compact # noqa: E402 +from run import CACHE, FONTS, QA_CACHE, RESULTS, TEXT_CHUNK, agent_prompt, load_prompt, sha8 # noqa: E402 + +# (family display, $/M input, $/M output). Cached reads bill at 0.1x input, +# Anthropic cache writes at 1.25x. Edit prices here; `--report` recomputes. +MODELS = { + "claude-fable-5": (10.0, 50.0), + "claude-opus-4-8": (15.0, 75.0), + "gpt-5.5": (2.0, 16.0), + "google/gemini-3.5-flash": (0.6, 4.0), + "moonshotai/kimi-k2.6": (0.68, 3.41), + "z-ai/glm-4.6v": (0.30, 0.90), +} +LENGTHS = (50, 150, 250) +CONDITIONS = ("text", "handoff", "compact", "img-6x10-sent", "img-6x10-bw", "img-5x8-sent", "img-5x8-bw") +ACK = "Noted. I have read the passages and will keep them in mind." + + +def cached(model: str, tag: str, payload: object, fn, fresh: bool) -> dict: + """Disk-cache `fn() -> dict` keyed by (model, tag, payload). Truncated/empty outputs are not cached.""" + key = sha8(model, tag, json.dumps(payload, sort_keys=True, default=str)) + path = QA_CACHE / f"{key}.json" + if path.exists() and not fresh: + hit = json.loads(path.read_text()) + if hit.get("stop") != "max_tokens" and ("text" not in hit or hit["text"]): + return hit + out = fn() + if out.get("stop") == "max_tokens" or out.get("text") == "": + print(f" WARN truncated/empty, not cached: {model} {tag} {key}") + else: + path.write_text(json.dumps(out)) + return out + + +def parse_img_condition(name: str) -> tuple[str, str, int] | None: + if not name.startswith("img-"): + return None + _, font, variant = name.split("-", 2) + columns = 1 + m = re.match(r"(.+)-(\d+)col$", variant) + if m: + variant, columns = m.group(1), int(m.group(2)) + return font, variant, columns + + +def chunk_budget(cond: str, size: int) -> int: + img = parse_img_condition(cond) + return capacity(FONTS[img[0]], size, img[2])[2] if img else TEXT_CHUNK + + +def session_frame(chunk_text: str) -> list[dict]: + return [ + {"role": "user", "content": [{"text": load_prompt("session-frame.md").format(context=chunk_text)}]}, + {"role": "assistant", "content": [{"text": ACK}]}, + ] + + +def run_cell_chunk(model: str, cond: str, start: int, end: int, ctx: dict) -> list[dict]: + """One (model, condition, chunk) unit: build carrier, QA, score.""" + args, flow, paras, offsets, keys = ctx["args"], ctx["flow"], ctx["paras"], ctx["offsets"], ctx["keys"] + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + chunk_text = flow[start:end] + q_block = "\n".join(f"{i + 1}. {q['q']}" for i, q in enumerate(questions)) + usage_rows: list[tuple[str, dict]] = [] + img = parse_img_condition(cond) + extra_items: list[dict] | None = None + + if img: + font, variant, columns = img + tag = f"{font}-{variant}" if columns == 1 else f"{font}-{variant}-{columns}col" + # "dimv2" salts pure-dim renders: pre-fix PNGs (sticky-fg bug, glyphs after a + # row's first stopword all dimmed) and the QA cache entries keyed on their paths. + salt = ("dimv2",) if variant == "dim" else () + png = CACHE / f"img-{tag}-{sha8(chunk_text, str(args.size), *salt)}.png" + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(f".{uuid.uuid4().hex[:8]}.tmp.png") + render(chunk_text, FONTS[font], CACHE, args.size, variant, columns=columns).save(tmp) + tmp.replace(png) + cols, rows, _ = capacity(FONTS[font], args.size, columns) + preamble = ( + load_prompt("qa-image-cols.md").format(cols=cols, rows=rows, columns=columns) + if columns > 1 + else load_prompt("qa-image.md").format(cols=cols, rows=rows) + ) + messages = [ + { + "role": "user", + "content": [ + {"text": preamble}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + elif cond == "compact" and is_openai(model): + comp = cached( + model, "remote-compact", {"chunk": chunk_text}, + lambda: dict(zip(("items", "usage"), openai_compact(keys["openai"], model, session_frame(chunk_text)))), + args.fresh, + ) + usage_rows.append(("compact", comp["usage"])) + extra_items = comp["items"] + messages = [ + {"role": "user", "content": [{"text": load_prompt("qa-remote-compact.md").format(questions=q_block)}]} + ] + elif cond in ("compact", "handoff"): + prompt_file = {"compact": "compaction-summary.md", "handoff": "handoff-document.md"}[cond] + gen = cached( + model, f"summary-{cond}", {"chunk": chunk_text}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete( + keys, model, + session_frame(chunk_text) + [{"role": "user", "content": [{"text": agent_prompt(prompt_file)}]}], + system=agent_prompt("summarization-system.md"), + max_tokens=args.max_tokens, + ), + ) + ), + args.fresh, + ) + usage_rows.append(("summarize", gen["usage"])) + messages = [ + { + "role": "user", + "content": [{"text": load_prompt("qa-text.md").format(context=gen["text"])}, {"text": q_block}], + } + ] + else: # text + messages = [ + { + "role": "user", + "content": [{"text": load_prompt("qa-text.md").format(context=chunk_text)}, {"text": q_block}], + } + ] + + qa = cached( + model, "qa", {"messages": messages, "extra": extra_items, "effort": args.effort}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete( + keys, model, messages, + max_tokens=args.max_tokens, effort=args.effort, extra_input_items=extra_items, + ), + ) + ), + args.fresh, + ) + usage_rows.append(("qa", qa["usage"])) + answers = squad.parse_numbered(qa["text"], len(questions)) + records = [] + for q, a in zip(questions, answers): + records.append( + { + "model": model, + "length": ctx["length"], + "cond": cond, + "chunk": start, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + ) + records[0]["usage"] = [{"phase": p, **u} for p, u in usage_rows] + return records + + +def aggregate(records: list[dict], price_in: float, price_out: float) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean_f1 = sum(f1s) / n + se = (sum((x - mean_f1) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + us = [u for r in records if "usage" in r for u in r["usage"]] + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + cost_in = (tok["in"] + 1.25 * tok["cache_w"] + 0.1 * tok["cache_r"]) / 1e6 * price_in + cost_out = tok["out"] / 1e6 * price_out + return { + "n": n, + "em": sum(r["em"] for r in records) / n, + "f1": mean_f1, + "f1_se": se, + "abstained": sum(r["abstained"] for r in records), + **{f"tok_{k}": v for k, v in tok.items()}, + "cost_in_usd": round(cost_in, 4), + "cost_out_usd": round(cost_out, 4), + "cost_usd": round(cost_in + cost_out, 4), + } + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--models", default=",".join(MODELS)) + ap.add_argument("--lengths", default=",".join(map(str, LENGTHS))) + ap.add_argument("--conditions", default=",".join(CONDITIONS)) + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=6) + ap.add_argument("--max-tokens", type=int, default=16384) + ap.add_argument("--effort", default=None, help="reasoning effort; None = provider default") + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--report", action="store_true", help="reprint from cache only") + ap.add_argument("--env", default="~/.env") + ap.add_argument("--out", default="final", help="results subdirectory (isolate concurrent runs)") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + out_dir = RESULTS / args.out + out_dir.mkdir(parents=True, exist_ok=True) + + models = [m.strip() for m in args.models.split(",") if m.strip()] + lengths = [int(x) for x in args.lengths.split(",") if x.strip()] + conditions = [c.strip() for c in args.conditions.split(",") if c.strip()] + for m in models: + if m not in MODELS: + raise SystemExit(f"unknown model {m}; add it to MODELS with prices") + + keys = {} + if not args.report: + keys["anthropic"] = load_env_key("ANTHROPIC_API_KEY", args.env) + keys["openai"] = load_env_key("OPENAI_API_KEY", args.env) + keys["openrouter"] = load_env_key("OPENROUTER_API_KEY", args.env) + + all_paras = squad.load_paragraphs(CACHE) + tasks = [] + for length in lengths: + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + ctx = {"args": args, "flow": flow, "paras": paras, "offsets": offsets, "keys": keys, "length": length} + for model in models: + for cond in conditions: + budget = chunk_budget(cond, args.size) + for start in range(0, len(flow), budget): + tasks.append((model, cond, start, min(start + budget, len(flow)), ctx)) + print(f"grid: {len(models)} models x {len(lengths)} lengths x {len(conditions)} conditions = {len(tasks)} chunk tasks") + + records: list[dict] = [] + done = 0 + with ThreadPoolExecutor(args.workers) as pool: + futures = [pool.submit(run_cell_chunk, m, c, s, e, ctx) for m, c, s, e, ctx in tasks] + for fut in futures: + records.extend(fut.result()) + done += 1 + if done % 25 == 0: + print(f" {done}/{len(tasks)} tasks", flush=True) + + with (out_dir / "records.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + + cells = [] + for model in models: + for length in lengths: + for cond in conditions: + sub = [r for r in records if r["model"] == model and r["length"] == length and r["cond"] == cond] + if not sub: + continue + cells.append({"model": model, "length": length, "condition": cond, **aggregate(sub, *MODELS[model])}) + (out_dir / "summary.json").write_text(json.dumps({"args": vars(args), "cells": cells}, indent=1)) + with (out_dir / "matrix.csv").open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=list(cells[0].keys())) + writer.writeheader() + writer.writerows(cells) + + for length in lengths: + print(f"\n== {length} passages == (F1 / $carry-in / $decode-out)") + hdr = f"{'condition':<15}" + "".join(f"{m:>22}" for m in models) + print(hdr + "\n" + "-" * len(hdr)) + for cond in conditions: + row = f"{cond:<15}" + for model in models: + cell = next((c for c in cells if c["model"] == model and c["length"] == length and c["condition"] == cond), None) + row += ( + f"{cell['f1']:>10.3f} {cell['cost_in_usd']:>5.2f} {cell['cost_out_usd']:>5.2f}" + if cell + else f"{'-':>22}" + ) + print(row) + print(f"\ndataset -> {out_dir}/records.jsonl, matrix.csv, summary.json") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/mono.py b/packages/snapcompact/research/mono.py new file mode 100644 index 000000000..db8071e63 --- /dev/null +++ b/packages/snapcompact/research/mono.py @@ -0,0 +1,154 @@ +"""Monolithic long-context probe: the WHOLE corpus in ONE request. + +final.py chunks every condition into ~10k-token QA calls, so it never tests +true in-request long-context retrieval. This runner stuffs an N-char SQuAD +flow (e.g. 800k chars ~ 200k text tokens) into a single request — either as +raw text or as a stack of dense-font images — with questions sampled evenly +across the whole span. Reports overall EM/F1 plus F1 by position quartile +(real lost-in-the-middle, which the chunked harness cannot see). + + uv run --with pillow python mono.py --model gpt-5.5 --chars 800000 \ + --conditions text,img-6x10-sent,img-6x8s-sent,img-8x8u-sent +""" + +import argparse +import json +import sys +import uuid +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent)) + +import squad # noqa: E402 +from bdf import capacity, render # noqa: E402 +from final import MODELS, cached, parse_img_condition # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, FONTS, RESULTS, load_prompt, sha8 # noqa: E402 + + +def build_content(cond: str, flow: str, size: int) -> tuple[list[dict], int]: + """Context blocks for the whole flow (questions appended per batch); returns (blocks, n_images).""" + img = parse_img_condition(cond) + if not img: + assert cond == "text", f"unsupported mono condition {cond!r}" + return [{"text": load_prompt("qa-text.md").format(context=flow), "cache": True}], 0 + font, variant, columns = img + cfg = FONTS[font] + cols, rows, cap = capacity(cfg, size, columns) + salt = ("dimv2",) if variant == "dim" else () + tag = f"{font}-{variant}" if columns == 1 else f"{font}-{variant}-{columns}col" + pngs = [] + for start in range(0, len(flow), cap): + chunk = flow[start : start + cap] + png = CACHE / f"img-{tag}-{sha8(chunk, str(size), *salt)}.png" + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(f".{uuid.uuid4().hex[:8]}.tmp.png") + render(chunk, cfg, CACHE, size, variant, columns=columns).save(tmp) + tmp.replace(png) + pngs.append(png) + preamble = load_prompt("qa-image-multi.md").format(k=len(pngs), cols=cols, rows=rows) + if cfg.repeat > 1: + preamble += ( + f"\nNote: every text line is rendered {cfg.repeat} times consecutively - first on the plain " + "background, then repeated on a pale highlight band. The copies show identical characters; " + "cross-check between them when a glyph is hard to read, and do not treat copies as separate text." + ) + blocks = [{"text": preamble}, *({"image_path": p} for p in pngs), {"text": "End of images.", "cache": True}] + return blocks, len(pngs) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--model", default="gpt-5.5") + ap.add_argument("--chars", type=int, default=800_000) + ap.add_argument("--conditions", default="text,img-6x10-sent,img-6x8s-sent,img-8x8u-sent") + ap.add_argument("--questions", type=int, default=50, help="total questions sampled across the flow") + ap.add_argument("--qpb", type=int, default=5, help="questions per API call (context re-sent, prefix-cached)") + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--effort", default=None) + ap.add_argument("--env", default="~/.env") + ap.add_argument("--out", default="mono") + ap.add_argument("--fresh", action="store_true") + args = ap.parse_args() + + keys = { + "anthropic": load_env_key("ANTHROPIC_API_KEY", args.env), + "openai": load_env_key("OPENAI_API_KEY", args.env), + "openrouter": load_env_key("OPENROUTER_API_KEY", args.env), + } + paras = squad.load_paragraphs(CACHE) + flow, offsets = squad.build_flow(paras, args.chars) + questions = squad.sample_chunk_questions(paras, offsets, 0, len(flow), args.questions, args.seed) + price_in, price_out = MODELS[args.model] + print( + f"flow: {len(flow):,} chars (~{len(flow) // 4 // 1000}k text tokens), " + f"{len(questions)} questions in batches of {args.qpb}" + ) + + out_dir = RESULTS / args.out + out_dir.mkdir(parents=True, exist_ok=True) + records, table = [], [] + for cond in [c.strip() for c in args.conditions.split(",") if c.strip()]: + ctx_blocks, n_imgs = build_content(cond, flow, args.size) + answers, usages, stops = [], [], [] + for b in range(0, len(questions), args.qpb): + batch = questions[b : b + args.qpb] + q_block = "\n".join(f"{i + 1}. {q['q']}" for i, q in enumerate(batch)) + messages = [{"role": "user", "content": [*ctx_blocks, {"text": q_block}]}] + qa = cached( + args.model, "qa-mono", {"messages": messages, "effort": args.effort}, + lambda m=messages: dict( + zip( + ("text", "usage", "stop"), + llm_complete(keys, args.model, m, max_tokens=args.max_tokens, effort=args.effort), + ) + ), + args.fresh, + ) + answers.extend(squad.parse_numbered(qa["text"], len(batch))) + usages.append(qa["usage"]) + stops.append(qa["stop"]) + rows = [ + { + "model": args.model, "cond": cond, "pos_rel": q["pos_rel"], "q": q["q"], + "answer": a, "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + for q, a in zip(questions, answers) + ] + records.extend({**r, "usage": usages} if i == 0 else r for i, r in enumerate(rows)) + u = {k: sum(x[k] for x in usages) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + stop = next((s for s in stops if s == "max_tokens"), stops[-1] if stops else "") + cost = (u["in"] + 1.25 * u["cache_w"] + 0.1 * u["cache_r"]) / 1e6 * price_in + u["out"] / 1e6 * price_out + quart = [] + for lo, hi in ((0, 0.25), (0.25, 0.5), (0.5, 0.75), (0.75, 1.01)): + qs = [r["f1"] for r in rows if lo <= r["pos_rel"] < hi] + quart.append(sum(qs) / len(qs) if qs else float("nan")) + table.append( + { + "cond": cond, "n": len(rows), "imgs": n_imgs, + "em": sum(r["em"] for r in rows) / len(rows), + "f1": sum(r["f1"] for r in rows) / len(rows), + "abst": sum(r["abstained"] for r in rows), + "tok_in": u["in"], "tok_cached": u["cache_r"], "tok_out": u["out"], "reas": u["reasoning"], + "cost": cost, "stop": stop, "q1": quart[0], "q2": quart[1], "q3": quart[2], "q4": quart[3], + } + ) + t = table[-1] + print( + f"{cond:<18} imgs={t['imgs']:>2} f1={t['f1']:.3f} em={t['em']:.3f} abst={t['abst']:>2} " + f"in={t['tok_in']:>7} cached={t['tok_cached']:>7} out={t['tok_out']:>6} reas={t['reas']:>6} " + f"${t['cost']:.2f} stop={t['stop']}" + ) + print(f"{'':<18} F1 by position quartile: " + " ".join(f"q{i + 1}={v:.3f}" for i, v in enumerate(quart))) + + (out_dir / "records.jsonl").write_text("\n".join(json.dumps(r) for r in records)) + (out_dir / "summary.json").write_text(json.dumps(table, indent=1)) + print(f"\ndataset -> {out_dir}/records.jsonl, summary.json") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/prompts/exp02-qa-image.md b/packages/snapcompact/research/prompts/exp02-qa-image.md new file mode 100644 index 000000000..edaaba915 --- /dev/null +++ b/packages/snapcompact/research/prompts/exp02-qa-image.md @@ -0,0 +1,8 @@ +The attached image contains encyclopedia passages rendered as a bitmap: monospace pixel font, {cols} characters per row, {rows} rows, read left-to-right then top-to-bottom. Original paragraph breaks were collapsed to spaces. + +{extra} + +Questions follow after the image. Answer them using ONLY text you can read in the image. +- Give short extractive answers: a word or phrase from the text, written in normal full English spelling. +- If you cannot read the relevant region well enough to answer, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/exp04-qa-image.md b/packages/snapcompact/research/prompts/exp04-qa-image.md new file mode 100644 index 000000000..6bae75bff --- /dev/null +++ b/packages/snapcompact/research/prompts/exp04-qa-image.md @@ -0,0 +1,8 @@ +The attached image is a page of encyclopedia passages typeset like a newspaper: a monospace pixel font in TWO columns of {col_w} characters, {rows} rows each, separated by a vertical gutter of white space. + +Reading order: read the LEFT column top-to-bottom first, then continue at the top of the RIGHT column and read it top-to-bottom. Article titles appear as bold uppercase headings; paragraphs within an article are separated by blank lines, and text is word-wrapped (no words are split across lines or columns). + +Questions follow after the image. Answer them using ONLY text you can read in the image. +- Give short extractive answers: a word or phrase copied from the text. +- If you cannot read the relevant region well enough to answer, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/exp05-qa-image-ctl.md b/packages/snapcompact/research/prompts/exp05-qa-image-ctl.md new file mode 100644 index 000000000..ecf89b10c --- /dev/null +++ b/packages/snapcompact/research/prompts/exp05-qa-image-ctl.md @@ -0,0 +1,10 @@ +The attached image contains encyclopedia passages rendered as a bitmap: monospace pixel font, {cols} characters per row, {rows} rows, read left-to-right then top-to-bottom. Original paragraph breaks were collapsed to spaces. + +Questions follow after the image. Do NOT transcribe the image. For each question: +1. Locate the relevant region first (skim for the topic). +2. Read ONLY that narrow band of rows carefully. +3. Answer from what you read there. + +- Give short extractive answers: a word or phrase copied from the text. +- If you cannot read the relevant region well enough to answer, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/exp05-qa-image.md b/packages/snapcompact/research/prompts/exp05-qa-image.md new file mode 100644 index 000000000..50ab394ab --- /dev/null +++ b/packages/snapcompact/research/prompts/exp05-qa-image.md @@ -0,0 +1,13 @@ +The attached image contains encyclopedia passages rendered as a bitmap: monospace pixel font, {cols} characters per row, {rows} rows, read left-to-right then top-to-bottom. Original paragraph breaks were collapsed to spaces. + +The left margin of the image is a row-number ruler, not part of the text: every 5th row is labeled with its row index in blue digits (0, 5, 10, ..., {last_label}). Rows are numbered from 0 at the top. Use the ruler to address regions of the image like line numbers in a file. + +Questions follow after the image. Do NOT transcribe the image. For each question: +1. Locate the relevant region first (skim for the topic, use the blue ruler to pin down the row range). +2. Read ONLY that narrow band of rows carefully. +3. Answer from what you read there. + +- Give short extractive answers: a word or phrase copied from the text. +- If you cannot read the relevant region well enough to answer, reply exactly UNREADABLE for that question. +- Output a numbered list, one line per question, in exactly this form: `N. | row≈` where is the ruler row where you found the answer. +- No other commentary. diff --git a/packages/snapcompact/research/prompts/exp06-prov-image.md b/packages/snapcompact/research/prompts/exp06-prov-image.md new file mode 100644 index 000000000..dac79c7eb --- /dev/null +++ b/packages/snapcompact/research/prompts/exp06-prov-image.md @@ -0,0 +1,6 @@ +The attached image contains a conversation transcript rendered as a bitmap: monospace pixel font, {cols} characters per row, {rows} rows, read left-to-right then top-to-bottom. The transcript interleaves messages from three roles: user, assistant, and tool. {encoding} + +For each numbered question below, do NOT answer the question itself. Instead identify which role's message contains the answer to it. +- Reply with exactly one word per line: user, assistant, or tool. +- You must choose one of the three roles for every question, even if uncertain — never reply UNREADABLE. +- Output a numbered list, one role per line, no commentary. diff --git a/packages/snapcompact/research/prompts/exp06-qa-image-tag.md b/packages/snapcompact/research/prompts/exp06-qa-image-tag.md new file mode 100644 index 000000000..12ebfe658 --- /dev/null +++ b/packages/snapcompact/research/prompts/exp06-qa-image-tag.md @@ -0,0 +1,6 @@ +The attached image contains a conversation transcript rendered as a bitmap: monospace pixel font, {cols} characters per row, {rows} rows, read left-to-right then top-to-bottom. The transcript interleaves messages from three roles. Each message is preceded by a bracketed role tag: [user], [asst], or [tool]. The tags are part of the rendered text, not message content. + +Questions follow after the image. Answer them using ONLY text you can read in the image. +- Give short extractive answers: a word or phrase copied from the text. +- If you cannot read the relevant region well enough to answer, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/exp06-qa-image.md b/packages/snapcompact/research/prompts/exp06-qa-image.md new file mode 100644 index 000000000..354aefb00 --- /dev/null +++ b/packages/snapcompact/research/prompts/exp06-qa-image.md @@ -0,0 +1,6 @@ +The attached image contains a conversation transcript rendered as a bitmap: monospace pixel font, {cols} characters per row, {rows} rows, read left-to-right then top-to-bottom. The transcript interleaves messages from three roles. Glyph color encodes the author of each message: dark blue = user, dark green = assistant, dark red = tool. A message boundary is where the glyph color changes. + +Questions follow after the image. Answer them using ONLY text you can read in the image. +- Give short extractive answers: a word or phrase copied from the text. +- If you cannot read the relevant region well enough to answer, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/exp07-answer-bands.md b/packages/snapcompact/research/prompts/exp07-answer-bands.md new file mode 100644 index 000000000..bc604ef41 --- /dev/null +++ b/packages/snapcompact/research/prompts/exp07-answer-bands.md @@ -0,0 +1,4 @@ +Now answer the questions. For each question, read ONLY the row band you identified above (plus a row or two of margin). Do NOT transcribe the image or read the rest of it. +- Give short extractive answers: a word or phrase copied from the text in the image. +- If you cannot read the relevant region well enough to answer, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/exp07-locate.md b/packages/snapcompact/research/prompts/exp07-locate.md new file mode 100644 index 000000000..9a4846537 --- /dev/null +++ b/packages/snapcompact/research/prompts/exp07-locate.md @@ -0,0 +1,5 @@ +The attached image contains encyclopedia passages rendered as a bitmap: monospace pixel font, {cols} characters per row, {rows} rows, read left-to-right then top-to-bottom. Original paragraph breaks were collapsed to spaces. + +Questions follow after the image. Do NOT answer them yet, and do NOT transcribe the image. For each question, skim the image coarsely (keywords, names, numbers) and guess which horizontal band of rows most likely contains the answer. +- Output a numbered list, one line per question, each line exactly of the form `rows A-B` (1-indexed, top row is 1; keep each band at most 25 rows tall). +- No answers, no commentary, nothing else. diff --git a/packages/snapcompact/research/prompts/exp07-qa-image.md b/packages/snapcompact/research/prompts/exp07-qa-image.md new file mode 100644 index 000000000..e20c37a9f --- /dev/null +++ b/packages/snapcompact/research/prompts/exp07-qa-image.md @@ -0,0 +1,7 @@ +The attached image contains encyclopedia passages rendered as a bitmap: monospace pixel font, {cols} characters per row, {rows} rows, read left-to-right then top-to-bottom. Original paragraph breaks were collapsed to spaces. + +Questions follow after the image. Answer them using ONLY text you can read in the image. +- Do NOT transcribe the image or read it linearly. For each question, locate the region likely to contain the answer, read only that region, and answer. +- Give short extractive answers: a word or phrase copied from the text. +- If you cannot read the relevant region well enough to answer, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/exp08-archive-eager.md b/packages/snapcompact/research/prompts/exp08-archive-eager.md new file mode 100644 index 000000000..d6731714b --- /dev/null +++ b/packages/snapcompact/research/prompts/exp08-archive-eager.md @@ -0,0 +1,10 @@ +The attached image contains encyclopedia passages rendered as a bitmap: monospace pixel font, {cols} characters per row, {rows} rows, read left-to-right then top-to-bottom. Rows are numbered 1 (top) to {rows} (bottom). Original paragraph breaks were collapsed to spaces. + +IMPORTANT: this font is rendered BELOW the size you can read reliably. Characters that look legible are often misread (digits, names, and dates especially). Be skeptical of your own reading. + +Questions follow after the image. Answer them using ONLY text you can read in the image. +- Give a direct short extractive answer ONLY when you are fully certain of every character in it. +- For any question where you are not fully certain — wrong-looking digits, a name you might be misreading, a region that is small or smudged — reply exactly `ZOOM rows A-B` for that question (e.g. `ZOOM rows 41-47`). That row band will be re-rendered for you at a large, comfortable font. Zooming is cheap and encouraged; guessing is penalized. +- Estimate the row range from the vertical position of the region (row 1 = top edge, row {rows} = bottom edge). If unsure of the exact rows, give a generous range — a margin around your estimate is included automatically. +- If you cannot even locate a relevant region, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/exp08-archive-phrase.md b/packages/snapcompact/research/prompts/exp08-archive-phrase.md new file mode 100644 index 000000000..6f5da8ea8 --- /dev/null +++ b/packages/snapcompact/research/prompts/exp08-archive-phrase.md @@ -0,0 +1,9 @@ +The attached image contains encyclopedia passages rendered as a bitmap: monospace pixel font, {cols} characters per row, {rows} rows, read left-to-right then top-to-bottom. Original paragraph breaks were collapsed to spaces. + +IMPORTANT: this font is rendered BELOW the size you can read reliably. Characters that look legible are often misread (digits, names, and dates especially). Be skeptical of your own reading. + +Questions follow after the image. Answer them using ONLY text you can read in the image. +- Give a direct short extractive answer ONLY when you are fully certain of every character in it. +- For any question where you are not fully certain — wrong-looking digits, a name you might be misreading, a region that is small or smudged — reply exactly `ZOOM ""` for that question, where is 3-8 consecutive words you can partially make out in or right next to that region (your best-effort reading; a few wrong characters are fine). The matching region will be re-rendered for you at a large, comfortable font. Zooming is cheap and encouraged; guessing is penalized. +- If you cannot even locate a relevant region, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/exp08-archive.md b/packages/snapcompact/research/prompts/exp08-archive.md new file mode 100644 index 000000000..3042f313c --- /dev/null +++ b/packages/snapcompact/research/prompts/exp08-archive.md @@ -0,0 +1,7 @@ +The attached image contains encyclopedia passages rendered as a bitmap: monospace pixel font, {cols} characters per row, {rows} rows, read left-to-right then top-to-bottom. Rows are numbered 1 (top) to {rows} (bottom). Original paragraph breaks were collapsed to spaces. The font is deliberately tiny: most of it is readable with effort, but some regions may not be. + +Questions follow after the image. Answer them using ONLY text you can read in the image. +- Give short extractive answers: a word or phrase copied from the text. +- If you locate the relevant region but it is too small or degraded to read reliably, do NOT guess: reply exactly `ZOOM rows A-B` for that question (e.g. `ZOOM rows 41-47`), giving the row range you need re-rendered at higher resolution. +- If you cannot even locate a relevant region, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/exp08-zoom.md b/packages/snapcompact/research/prompts/exp08-zoom.md new file mode 100644 index 000000000..e71e273f2 --- /dev/null +++ b/packages/snapcompact/research/prompts/exp08-zoom.md @@ -0,0 +1,6 @@ +Below are high-resolution re-renderings of the archive row bands you requested. Each image is preceded by a label giving the archive row range it covers; the text inside is identical to those rows, re-flowed to the new line width. + +Answer your remaining questions, listed after the images, using the zoomed images plus anything you already read. Keep the same question numbers as before. +- Give short extractive answers: a word or phrase copied from the text. +- If you still cannot read the answer, reply exactly UNREADABLE. +- Output a numbered list (original numbering), one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/exp09-frame.md b/packages/snapcompact/research/prompts/exp09-frame.md new file mode 100644 index 000000000..e8c621c2d --- /dev/null +++ b/packages/snapcompact/research/prompts/exp09-frame.md @@ -0,0 +1,3 @@ +Reference material for this session arrives as bitmap image pages: monospace pixel font, {cols} characters per row, up to {rows} rows per page, read left-to-right then top-to-bottom. Original paragraph breaks were collapsed to spaces; the text flows continuously across pages. More pages may be appended later in this session. I will ask questions about the material later. + +Here is page 1. diff --git a/packages/snapcompact/research/prompts/exp09-page.md b/packages/snapcompact/research/prompts/exp09-page.md new file mode 100644 index 000000000..2a3601496 --- /dev/null +++ b/packages/snapcompact/research/prompts/exp09-page.md @@ -0,0 +1 @@ +Here is page {page}, continuing the reference material. diff --git a/packages/snapcompact/research/prompts/exp09-qa.md b/packages/snapcompact/research/prompts/exp09-qa.md new file mode 100644 index 000000000..baf6ff869 --- /dev/null +++ b/packages/snapcompact/research/prompts/exp09-qa.md @@ -0,0 +1,6 @@ +Answer the following questions using ONLY text you can read in the reference pages above. +- Give short extractive answers: a word or phrase copied from the text. +- If you cannot read the relevant region well enough to answer, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. + +{questions} diff --git a/packages/snapcompact/research/prompts/exp11-qa-hier.md b/packages/snapcompact/research/prompts/exp11-qa-hier.md new file mode 100644 index 000000000..85ed0bead --- /dev/null +++ b/packages/snapcompact/research/prompts/exp11-qa-hier.md @@ -0,0 +1,12 @@ +You hold a long stream of encyclopedia passages in a three-tier session memory. The stream is ordered oldest to newest and was split into three consecutive age bands, each stored at a different fidelity: + +- TIER 3 (oldest third): available only as a narrative summary that was written earlier{appendix_note}. +- TIER 2 (middle third): rendered as {n_pages} bitmap page image(s) — monospace pixel font, {cols} characters per row, {rows} rows per page, read left-to-right then top-to-bottom. Original paragraph breaks were collapsed to spaces. +- TIER 1 (newest third): plain verbatim text. + +The tiers appear below in age order (oldest first), each clearly labeled. Questions follow at the end; any question may target material from any tier. + +Answer the questions using ONLY the material provided here. +- Give short extractive answers: a word or phrase copied from the text or read out of an image. For facts that only appear in the TIER 3 summary, copy the best short phrase from the summary. +- If the needed fact is not present in any tier and not legible in any image, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/exp13-extract.md b/packages/snapcompact/research/prompts/exp13-extract.md new file mode 100644 index 000000000..d5df0a542 --- /dev/null +++ b/packages/snapcompact/research/prompts/exp13-extract.md @@ -0,0 +1,10 @@ +You are compacting the reference passages above for a later question-answering session. Produce an EXTRACTIVE digest: copy out the most load-bearing sentences and spans VERBATIM. + +Rules: +- Copy text exactly as it appears in the passages: no paraphrasing, no summarizing, no rewording, no truncation of sentences. Keep numbers, dates, names, units, capitalization, and spelling character-for-character. +- Preserve the original order of the source; never reorder spans. +- Prefer complete sentences dense with facts: named entities, quantities, dates, definitions, titles, causal claims. Drop filler, transitions, and rhetorical framing. +- Cover the FULL span of the material evenly — do not spend the budget on the earliest passages and starve the later ones. +- Separate non-adjacent extracts with a line containing only: ... +- Hard output budget: at most {budget} characters in total. Use most of it; never exceed it. +- Output ONLY the extracted text (and the ... separators). No headers, no commentary, no explanations. diff --git a/packages/snapcompact/research/prompts/exp19-qa-doc.md b/packages/snapcompact/research/prompts/exp19-qa-doc.md new file mode 100644 index 000000000..1a08492da --- /dev/null +++ b/packages/snapcompact/research/prompts/exp19-qa-doc.md @@ -0,0 +1,8 @@ +The attached image is a page of encyclopedia passages typeset like a newspaper: a monospace pixel font in TWO columns of {col_w} characters, {rows} rows each, separated by a blank vertical gutter. + +Reading order: read the LEFT column top-to-bottom first, then continue at the top of the RIGHT column and read it top-to-bottom. Article titles appear as bold uppercase headings; paragraphs within an article are separated by blank lines, and text is word-wrapped (no words are split across lines or columns). + +Questions follow after the image. Answer them using ONLY text you can read in the image. +- Give short extractive answers: a word or phrase copied from the text. +- If you cannot read the relevant region well enough to answer, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/exp21-qa-braille.md b/packages/snapcompact/research/prompts/exp21-qa-braille.md new file mode 100644 index 000000000..b68515ca2 --- /dev/null +++ b/packages/snapcompact/research/prompts/exp21-qa-braille.md @@ -0,0 +1,11 @@ +The attached image contains English text encoded in Grade-1 (uncontracted, letter-by-letter) 6-dot braille, drawn as a dense grid of dot-matrix cells: {cols} cells per row, {rows} rows, read left-to-right, top-to-bottom, wrapping at row ends mid-word (no hyphenation marks). + +Encoding details: +- Standard braille letter patterns a-z. All text is lowercase; there are NO capital signs. +- Digits: the number sign (dots 3456) starts a digit run, then digits use a-j (1=a ... 9=i, 0=j). One number sign covers the whole run. +- Punctuation: period (dots 256), comma (dot 2), apostrophe (dot 3), hyphen (dots 36), colon (dots 25), semicolon (dots 23), question mark (dots 236), exclamation (dots 235). A blank cell is a space. All other symbols were dropped as spaces. + +Questions follow after the image. Decode only the regions you need and answer using ONLY text encoded in the image. +- Give short extractive answers in ordinary English: a word or phrase from the text. +- If you cannot decode the relevant region, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/qa-image-cols.md b/packages/snapcompact/research/prompts/qa-image-cols.md new file mode 100644 index 000000000..2cb675006 --- /dev/null +++ b/packages/snapcompact/research/prompts/qa-image-cols.md @@ -0,0 +1,6 @@ +The attached image contains encyclopedia passages rendered as a bitmap in a newspaper-style layout: monospace pixel font, {columns} columns of text separated by a vertical gutter line. Each column holds {rows} lines of {cols} characters. Reading order: start with the leftmost column, read its lines left-to-right from top to bottom, then continue at the top of the next column. Original paragraph breaks were collapsed to spaces. + +Questions follow after the image. Answer them using ONLY text you can read in the image. +- Give short extractive answers: a word or phrase copied from the text. +- If you cannot read the relevant region well enough to answer, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/qa-image-multi.md b/packages/snapcompact/research/prompts/qa-image-multi.md new file mode 100644 index 000000000..fe3cc7380 --- /dev/null +++ b/packages/snapcompact/research/prompts/qa-image-multi.md @@ -0,0 +1,6 @@ +The attached {k} images contain encyclopedia passages rendered as dense bitmaps: monospace pixel font, {cols} characters per row, {rows} rows per image. The text flows continuously across the images: read each image left-to-right, top-to-bottom, then continue with the next image in order (image 1 first, image {k} last). Original paragraph breaks were collapsed to spaces. + +Questions follow after the images. Answer them using ONLY text you can read in the images. +- Give short extractive answers: a word or phrase copied from the text. +- If you cannot read the relevant region well enough to answer, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/qa-image.md b/packages/snapcompact/research/prompts/qa-image.md new file mode 100644 index 000000000..085061a4b --- /dev/null +++ b/packages/snapcompact/research/prompts/qa-image.md @@ -0,0 +1,6 @@ +The attached image contains encyclopedia passages rendered as a bitmap: monospace pixel font, {cols} characters per row, {rows} rows, read left-to-right then top-to-bottom. Original paragraph breaks were collapsed to spaces. + +Questions follow after the image. Answer them using ONLY text you can read in the image. +- Give short extractive answers: a word or phrase copied from the text. +- If you cannot read the relevant region well enough to answer, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/qa-remote-compact.md b/packages/snapcompact/research/prompts/qa-remote-compact.md new file mode 100644 index 000000000..504c3c9bf --- /dev/null +++ b/packages/snapcompact/research/prompts/qa-remote-compact.md @@ -0,0 +1,8 @@ +Earlier in this conversation you were given a set of reference passages to keep in mind. The conversation context may have been compacted since. + +Answer the questions below using ONLY what you retained from those passages. +- Give short extractive answers: a word or phrase from the passages. +- If the information did not survive in your context, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. + +{questions} diff --git a/packages/snapcompact/research/prompts/qa-text.md b/packages/snapcompact/research/prompts/qa-text.md new file mode 100644 index 000000000..d4606aa1f --- /dev/null +++ b/packages/snapcompact/research/prompts/qa-text.md @@ -0,0 +1,10 @@ +Below is reference material. Questions follow in the next block. + + +{context} + + +Answer the questions using ONLY the reference material above. +- Give short extractive answers: a word or phrase copied from the text. +- If the reference does not contain the information, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/session-frame.md b/packages/snapcompact/research/prompts/session-frame.md new file mode 100644 index 000000000..0941b56bd --- /dev/null +++ b/packages/snapcompact/research/prompts/session-frame.md @@ -0,0 +1,5 @@ +Here is a set of reference passages I need you to keep in mind for this session. I will ask questions about them later. + + +{context} + diff --git a/packages/snapcompact/research/providers.py b/packages/snapcompact/research/providers.py new file mode 100644 index 000000000..270cefd5d --- /dev/null +++ b/packages/snapcompact/research/providers.py @@ -0,0 +1,261 @@ +"""Provider-neutral LLM client: Anthropic Messages + OpenAI Responses. + +Neutral message shape: [{"role": str, "content": [block, ...]}] where block is + {"text": str} - text block + {"image_path": Path, "cache": bool} - PNG by path; cache marks the prompt-cache + breakpoint (Anthropic only; OpenAI caches + automatically) + +Normalized usage: {"in", "out", "cache_w", "cache_r", "reasoning"}. +""" + +import base64 +import json +import time +import urllib.error +import urllib.request +from pathlib import Path + +ANTHROPIC_URL = "https://api.anthropic.com/v1/messages" +OPENAI_URL = "https://api.openai.com/v1/responses" +ANTHROPIC_VERSION = "2023-06-01" + + +def load_env_key(var: str, env_path: str = "~/.env") -> str: + """Last assignment wins (mirrors shell sourcing).""" + key = None + for line in Path(env_path).expanduser().read_text().splitlines(): + line = line.strip() + if line.startswith("export "): + line = line[len("export ") :] + if line.startswith(f"{var}="): + key = line.split("=", 1)[1].strip().strip("'\"") + if not key: + raise SystemExit(f"no {var} in {env_path}") + return key + + +def _post(url: str, body: dict, headers: dict, retries: int = 4) -> dict: + payload = json.dumps(body).encode() + req = urllib.request.Request(url, data=payload, headers={"content-type": "application/json", **headers}) + for attempt in range(retries + 1): + try: + with urllib.request.urlopen(req, timeout=600) as resp: + raw = resp.read() + return json.loads(raw) + except urllib.error.HTTPError as err: + detail = err.read().decode(errors="replace")[:500] + if err.code in (408, 429, 500, 502, 503, 529) and attempt < retries: + wait = 2.0 * 2**attempt + print(f" HTTP {err.code}, retrying in {wait:.0f}s: {detail[:120]}") + time.sleep(wait) + continue + raise SystemExit(f"API error {err.code} ({url}): {detail}") from err + except (json.JSONDecodeError, TimeoutError, urllib.error.URLError) as err: + if attempt < retries: + wait = 2.0 * 2**attempt + print(f" bad response ({type(err).__name__}), retrying in {wait:.0f}s") + time.sleep(wait) + continue + raise + raise AssertionError("unreachable") + + +def _png_b64(path: Path) -> str: + return base64.b64encode(path.read_bytes()).decode() + + +# --- Anthropic --- + + +def _anthropic_blocks(blocks: list[dict]) -> list[dict]: + out = [] + for b in blocks: + if "text" in b: + item: dict = {"type": "text", "text": b["text"]} + else: + item = { + "type": "image", + "source": {"type": "base64", "media_type": "image/png", "data": _png_b64(b["image_path"])}, + } + if b.get("cache"): + item["cache_control"] = {"type": "ephemeral"} + out.append(item) + return out + + +def _anthropic_complete( + api_key: str, model: str, messages: list[dict], system: str | None, max_tokens: int, effort: str | None +) -> tuple[str, dict, str]: + body: dict = { + "model": model, + "max_tokens": max_tokens, + "messages": [{"role": m["role"], "content": _anthropic_blocks(m["content"])} for m in messages], + } + if system: + body["system"] = system + headers = {"x-api-key": api_key, "anthropic-version": ANTHROPIC_VERSION} + if effort: + body["output_config"] = {"effort": effort} + headers["anthropic-beta"] = "effort-2025-11-24" + out = _post(ANTHROPIC_URL, body, headers) + text = "".join(b.get("text", "") for b in out["content"] if b.get("type") == "text") + u = out.get("usage", {}) + usage = { + "in": u.get("input_tokens", 0), + "out": u.get("output_tokens", 0), + "cache_w": u.get("cache_creation_input_tokens", 0), + "cache_r": u.get("cache_read_input_tokens", 0), + "reasoning": (u.get("output_tokens_details") or {}).get("thinking_tokens", 0), + } + return text, usage, out.get("stop_reason", "") + + +# --- OpenAI (Responses API) --- + + +def _openai_content(blocks: list[dict], role: str) -> list[dict]: + text_type = "output_text" if role == "assistant" else "input_text" + out = [] + for b in blocks: + if "text" in b: + out.append({"type": text_type, "text": b["text"]}) + else: + out.append( + { + "type": "input_image", + "image_url": f"data:image/png;base64,{_png_b64(b['image_path'])}", + "detail": "original", + } + ) + return out + + +def _openai_usage(out: dict) -> dict: + u = out.get("usage", {}) + cached = (u.get("input_tokens_details") or {}).get("cached_tokens", 0) + return { + "in": u.get("input_tokens", 0) - cached, + "out": u.get("output_tokens", 0), + "cache_w": 0, + "cache_r": cached, + "reasoning": (u.get("output_tokens_details") or {}).get("reasoning_tokens", 0), + } + + +def _openai_output_text(out: dict) -> str: + parts = [] + for item in out.get("output", []): + if item.get("type") == "message": + for c in item.get("content", []): + if c.get("type") == "output_text": + parts.append(c.get("text", "")) + return "".join(parts) + + +def _openai_complete( + api_key: str, + model: str, + messages: list[dict], + system: str | None, + max_tokens: int, + effort: str | None, + extra_input_items: list[dict] | None = None, +) -> tuple[str, dict, str]: + input_items: list[dict] = list(extra_input_items or []) + input_items += [{"role": m["role"], "content": _openai_content(m["content"], m["role"])} for m in messages] + body: dict = {"model": model, "input": input_items, "max_output_tokens": max_tokens, "store": False} + if system: + body["instructions"] = system + if effort: + body["reasoning"] = {"effort": "high" if effort in ("xhigh", "max") else effort} + out = _post(OPENAI_URL, body, {"authorization": f"Bearer {api_key}"}) + status = out.get("status", "") + stop = "max_tokens" if (out.get("incomplete_details") or {}).get("reason") == "max_output_tokens" else status + return _openai_output_text(out), _openai_usage(out), stop + + +def openai_compact(api_key: str, model: str, messages: list[dict]) -> tuple[list[dict], dict]: + """POST /responses/compact: returns (compacted output items, usage).""" + body = { + "model": model, + "input": [{"role": m["role"], "content": _openai_content(m["content"], m["role"])} for m in messages], + } + out = _post(f"{OPENAI_URL}/compact", body, {"authorization": f"Bearer {api_key}"}) + return out.get("output", []), _openai_usage(out) + + +# --- OpenRouter (chat completions) --- + +OPENROUTER_URL = "https://openrouter.ai/api/v1/chat/completions" + + +def _openrouter_complete( + api_key: str, model: str, messages: list[dict], system: str | None, max_tokens: int, effort: str | None +) -> tuple[str, dict, str]: + def content(blocks: list[dict]) -> list[dict]: + out = [] + for b in blocks: + if "text" in b: + out.append({"type": "text", "text": b["text"]}) + else: + out.append({"type": "image_url", "image_url": {"url": f"data:image/png;base64,{_png_b64(b['image_path'])}"}}) + return out + + chat_messages = [{"role": m["role"], "content": content(m["content"])} for m in messages] + if system: + chat_messages.insert(0, {"role": "system", "content": system}) + body: dict = {"model": model, "messages": chat_messages, "max_tokens": max_tokens} + if effort == "none": + body["reasoning"] = {"enabled": False} # OpenRouter's disable switch; effort "none" is not a valid level + elif effort: + body["reasoning"] = {"effort": "high" if effort in ("xhigh", "max") else effort} + out = _post(OPENROUTER_URL, body, {"authorization": f"Bearer {api_key}"}) + choice = (out.get("choices") or [{}])[0] + text = (choice.get("message") or {}).get("content") or "" + if isinstance(text, list): # some providers return content parts + text = "".join(p.get("text", "") for p in text if isinstance(p, dict)) + u = out.get("usage", {}) + usage = { + "in": u.get("prompt_tokens", 0) - (u.get("prompt_tokens_details") or {}).get("cached_tokens", 0), + "out": u.get("completion_tokens", 0), + "cache_w": 0, + "cache_r": (u.get("prompt_tokens_details") or {}).get("cached_tokens", 0), + "reasoning": (u.get("completion_tokens_details") or {}).get("reasoning_tokens", 0), + } + stop = "max_tokens" if choice.get("finish_reason") == "length" else (choice.get("finish_reason") or "") + return text, usage, stop + + +# --- dispatch --- + + +def is_openai(model: str) -> bool: + return model.startswith(("gpt-", "o3", "o4", "codex")) + + +def is_openrouter(model: str) -> bool: + return "/" in model + + +def llm_complete( + api_keys: dict[str, str], + model: str, + messages: list[dict], + system: str | None = None, + max_tokens: int = 16384, + effort: str | None = None, + extra_input_items: list[dict] | None = None, +) -> tuple[str, dict, str]: + """Returns (text, normalized usage, stop). stop == "max_tokens" means truncated.""" + if is_openrouter(model): + if extra_input_items: + raise ValueError("extra_input_items is OpenAI-only (compacted window replay)") + return _openrouter_complete(api_keys["openrouter"], model, messages, system, max_tokens, effort) + if is_openai(model): + return _openai_complete( + api_keys["openai"], model, messages, system, max_tokens, effort, extra_input_items + ) + if extra_input_items: + raise ValueError("extra_input_items is OpenAI-only (compacted window replay)") + return _anthropic_complete(api_keys["anthropic"], model, messages, system, max_tokens, effort) diff --git a/packages/snapcompact/research/run.py b/packages/snapcompact/research/run.py new file mode 100644 index 000000000..efc570d75 --- /dev/null +++ b/packages/snapcompact/research/run.py @@ -0,0 +1,353 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""snapcompact: QA recall per context-compression strategy, over the full SQuAD dev set. + +The corpus (all SQuAD v1.1 dev passages, space-joined; ~1.5M chars) is processed +per condition in chunks sized to that condition's carrying capacity: + + text plain text, fixed 40,716-char chunks (= img-6x10 capacity) + compact agent compaction summary of each text chunk, QA on the summary + handoff agent handoff document of each text chunk, QA on the summary + img-- one 1568x1568 image per chunk; font in {8x13,6x10,5x8,5x7, + 4x6tt,4x5tt}, render variant v in {color,zebra,bw} + +Per chunk, up to --qpc questions are sampled (seeded, evenly spread across the +chunk so answers land at every image row band; pos_rel is recorded for position +analysis). Scoring is official SQuAD EM/F1. Responses are cached by payload +hash, so interrupted runs resume for free. + +Usage examples: + uv run run.py # default condition set, full corpus + uv run run.py --limit-chars 200000 # quick pass on a corpus prefix + uv run run.py --conditions img-6x10-bw # one condition + uv run run.py --report # re-print tables from cache, no API + +Key: ANTHROPIC_API_KEY from ~/.env (last assignment wins). +""" + +import argparse +import hashlib +import json +import re +import sys +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +HERE = Path(__file__).resolve().parent + +def find_agent_prompts() -> Path: + for parent in HERE.parents: + for candidate in ( + parent / "packages" / "agent" / "src" / "compaction" / "prompts", + parent / "agent" / "src" / "compaction" / "prompts", + ): + if candidate.exists(): + return candidate + raise FileNotFoundError("Could not find agent compaction prompts") + + +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from anthropic_api import complete, image_block, load_api_key # noqa: E402 +from bdf import VARIANTS, FontCfg, capacity, render # noqa: E402 + +AGENT_PROMPTS = find_agent_prompts() +CACHE = HERE / ".cache" +QA_CACHE = CACHE / "qa" +RESULTS = HERE / "results" + +FONTS = { + "8x13": FontCfg("8x13", "8x13", 8, 13), + "7x13": FontCfg("7x13", "7x13", 7, 13), + "6x12": FontCfg("6x12", "6x12", 6, 12), + "6x10": FontCfg("6x10", "6x10", 6, 10), + "6x9": FontCfg("6x9", "6x9", 6, 9), + # Anisotropic squashes: same glyphs, tighter grid. Crop flavor (6x6/6x8/3x10): + # next row's band paints over the vertical overhang, halved advance fuses ink. + # Stretch flavor (*s): rasterize native 6x10, Lanczos-resize to the target cell. + "6x8s": FontCfg("6x8s", "6x10", 6, 8, native=(6, 10)), + "6x6s": FontCfg("6x6s", "6x10", 6, 6, native=(6, 10)), + "3x10s": FontCfg("3x10s", "6x10", 3, 10, native=(6, 10)), + # Human-designed square cell: unscii-8 native, and stretched to a 6x6 cell. + "8x8u": FontCfg("8x8u", "unscii-8", 8, 8), + "6x6u": FontCfg("6x6u", "unscii-8", 6, 6, native=(8, 8)), + "6x8": FontCfg("6x8", "6x9", 6, 8, ascent=7), + "6x6": FontCfg("6x6", "6x9", 6, 6, ascent=6), + "3x10": FontCfg("3x10", "6x10", 3, 10), + # Redundancy coding: every line rendered twice, duplicate on a pale highlight. + "8x8r": FontCfg("8x8r", "unscii-8", 8, 8, repeat=2), + # Downsample-survivors: unscii-8 upscaled so gateway resizes leave legible glyphs. + "12x12u": FontCfg("12x12u", "unscii-8", 12, 12, native=(8, 8)), + "16x16u": FontCfg("16x16u", "unscii-8", 16, 16, native=(8, 8)), + "5x8": FontCfg("5x8", "5x8", 5, 8), + "5x7": FontCfg("5x7", "5x7", 5, 7), + "4x6tt": FontCfg("4x6tt", "tom-thumb", 4, 6, ascent=5), + "4x5tt": FontCfg("4x5tt", "tom-thumb", 4, 5, ascent=5), +} +TEXT_CHUNK = 40716 # = img-6x10 capacity; keeps text/summary chunks comparable +DEFAULT_CONDITIONS = ( + "text,compact,handoff," + "img-8x13-color,img-6x10-color,img-5x8-color,img-5x7-color," + "img-6x10-zebra,img-5x8-zebra,img-6x10-bw,img-5x8-bw" +) + + +def sha8(*parts: str) -> str: + return hashlib.sha1("\x00".join(parts).encode()).hexdigest()[:8] + + +def load_prompt(name: str) -> str: + return (HERE / "prompts" / name).read_text() + + +def agent_prompt(name: str) -> str: + text = (AGENT_PROMPTS / name).read_text() + # Drop unused Handlebars conditionals (no custom focus in this eval). + return re.sub(r"\{\{#if .*?\{\{/if\}\}\n?", "", text, flags=re.DOTALL) + + +def cached_complete(api_key: str, model: str, messages: list[dict], fresh: bool, **kw) -> tuple[str, dict]: + """complete() with response caching keyed on the full request payload. + + Truncated responses (stop_reason == max_tokens) are never cached and never + served from cache, so re-runs with a larger budget repair them. + """ + key = sha8(model, kw.get("effort") or "", json.dumps(messages, sort_keys=True)) + path = QA_CACHE / f"{key}.json" + if path.exists() and not fresh: + hit = json.loads(path.read_text()) + if hit.get("stop") != "max_tokens" and hit["text"]: + return hit["text"], hit["usage"] + text, usage, stop = complete(api_key, model, messages, **kw) + if stop == "max_tokens": + print(f" WARN truncated response (stop=max_tokens), not cached: {key}") + else: + path.write_text(json.dumps({"text": text, "usage": usage, "stop": stop})) + return text, usage + + +def parse_condition(name: str) -> dict: + if name in ("text", "compact", "handoff"): + return {"name": name, "kind": name} + m = re.fullmatch(r"img-([a-z0-9]+)-([a-z-]+)", name) + if not m or m.group(1) not in FONTS or m.group(2) not in VARIANTS: + raise SystemExit(f"bad condition {name!r}; expected text|compact|handoff|img--") + return {"name": name, "kind": "image", "font": FONTS[m.group(1)], "variant": m.group(2)} + + +def run_chunk(cond: dict, start: int, end: int, ctx_args: dict) -> list[dict]: + """Execute one (condition, chunk) task; returns per-question records.""" + args, flow, paras, offsets, api_key = ( + ctx_args["args"], + ctx_args["flow"], + ctx_args["paras"], + ctx_args["offsets"], + ctx_args["api_key"], + ) + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + chunk_text = flow[start:end] + usage_rows: list[tuple[str, dict]] = [] + + png = cols = rows = None + context = chunk_text + if cond["kind"] == "image": + salt = ("dimv2",) if cond["variant"] == "dim" else () # cache-bust pre-fix sticky-fg dim renders + png = CACHE / f"img-{cond['font'].name}-{cond['variant']}-{sha8(chunk_text, str(args.size), *salt)}.png" + if not png.exists(): + render(chunk_text, cond["font"], CACHE, args.size, cond["variant"]).save(png) + cols, rows, _ = capacity(cond["font"], args.size) + elif cond["kind"] in ("compact", "handoff"): + prompt_file = {"compact": "compaction-summary.md", "handoff": "handoff-document.md"}[cond["kind"]] + gen_messages = [ + {"role": "user", "content": load_prompt("session-frame.md").format(context=chunk_text)}, + {"role": "assistant", "content": "Noted. I have read the passages and will keep them in mind."}, + {"role": "user", "content": agent_prompt(prompt_file)}, + ] + context, gen_usage = cached_complete( + api_key, args.model, gen_messages, args.fresh, + system=agent_prompt("summarization-system.md"), max_tokens=4096, + ) + usage_rows.append(("summarize", gen_usage)) + + use_cache = args.cache == "on" or (args.cache == "auto" and args.qpb > 0) + batch_size = args.qpb or len(questions) + answers: list[str] = [] + for b in range(0, len(questions), batch_size): + batch = questions[b : b + batch_size] + q_block = "\n".join(f"{i + 1}. {q['q']}" for i, q in enumerate(batch)) + if cond["kind"] == "image": + carrier = image_block(png) + preamble = {"type": "text", "text": load_prompt("qa-image.md").format(cols=cols, rows=rows)} + else: + carrier = {"type": "text", "text": load_prompt("qa-text.md").format(context=context)} + preamble = None + if use_cache: + carrier["cache_control"] = {"type": "ephemeral"} + content = ([preamble] if preamble else []) + [carrier, {"type": "text", "text": q_block}] + messages = [{"role": "user", "content": content}] + text, usage = cached_complete( + api_key, args.model, messages, args.fresh, max_tokens=args.max_tokens, effort=args.effort + ) + usage_rows.append(("qa", usage)) + answers.extend(squad.parse_numbered(text, len(batch))) + records = [] + for q, a in zip(questions, answers): + records.append( + { + "cond": cond["name"], + "chunk": start, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + ) + # Attach token usage to the first record of the chunk (avoids double counting). + records[0]["usage"] = [ + { + "phase": phase, + "in": u.get("input_tokens", 0), + "out": u.get("output_tokens", 0), + "cache_w": u.get("cache_creation_input_tokens", 0), + "cache_r": u.get("cache_read_input_tokens", 0), + } + for phase, u in usage_rows + ] + return records + + +def aggregate(name: str, records: list[dict], price_in: float, price_out: float) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean_f1 = sum(f1s) / n + se = (sum((x - mean_f1) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + usages = [u for r in records if "usage" in r for u in r["usage"]] + tok_in = sum(u["in"] for u in usages) + tok_out = sum(u["out"] for u in usages) + cache_w = sum(u.get("cache_w", 0) for u in usages) + cache_r = sum(u.get("cache_r", 0) for u in usages) + quart = [] + for lo, hi in ((0.0, 0.25), (0.25, 0.5), (0.5, 0.75), (0.75, 1.01)): + qs = [r["f1"] for r in records if lo <= r["pos_rel"] < hi] + quart.append(sum(qs) / len(qs) if qs else None) + return { + "name": name, + "n": n, + "em": sum(r["em"] for r in records) / n, + "f1": mean_f1, + "f1_se": se, + "abstained": sum(r["abstained"] for r in records), + "tokens_in": tok_in, + "tokens_out": tok_out, + "cache_w": cache_w, + "cache_r": cache_r, + # Anthropic pricing: cache write 1.25x input, cache read 0.1x input (5m TTL). + "cost_usd": (tok_in + 1.25 * cache_w + 0.1 * cache_r) / 1e6 * price_in + tok_out / 1e6 * price_out, + "f1_by_quartile": quart, + } + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--model", default="claude-fable-5") + ap.add_argument("--conditions", default=DEFAULT_CONDITIONS) + ap.add_argument("--qpc", type=int, default=30, help="questions sampled per chunk") + ap.add_argument("--qpb", type=int, default=0, help="questions per API call (batches the chunk); 0 = all at once") + ap.add_argument("--cache", choices=["auto", "on", "off"], default="auto", + help="prompt-cache the carrier block; auto = on when --qpb is set") + ap.add_argument("--max-tokens", type=int, default=8192, help="output budget per QA call (incl. thinking)") + ap.add_argument("--effort", choices=["low", "medium", "high", "xhigh", "max"], default=None, + help="adaptive-thinking effort for QA calls; default = provider default") + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=4) + ap.add_argument("--limit-chars", type=int, default=0, help="cap corpus size; 0 = full dev set") + ap.add_argument("--limit-paras", type=int, default=0, help="cap corpus to first N passages; 0 = all") + ap.add_argument("--fresh", action="store_true", help="ignore cached responses") + ap.add_argument("--report", action="store_true", help="aggregate cached records only; no API calls") + ap.add_argument("--price-in", type=float, default=10.0, help="$ per 1M input tokens") + ap.add_argument("--price-out", type=float, default=50.0, help="$ per 1M output tokens") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + scope = f"p{args.limit_paras}" if args.limit_paras else (args.limit_chars or "full") + tag = "".join( + [ + f"-qpb{args.qpb}" if args.qpb else "", + f"-s{args.size}" if args.size != 1568 else "", + f"-e{args.effort}" if args.effort else "", + ] + ) + run_dir = RESULTS / f"{args.model}-seed{args.seed}-qpc{args.qpc}-{scope}{tag}" + run_dir.mkdir(parents=True, exist_ok=True) + + paras = squad.load_paragraphs(CACHE) + if args.limit_paras: + paras = paras[: args.limit_paras] + flow, offsets = squad.build_flow(paras, args.limit_chars or None) + conditions = [parse_condition(c.strip()) for c in args.conditions.split(",") if c.strip()] + + tasks: list[tuple[dict, int, int]] = [] + for cond in conditions: + budget = capacity(cond["font"], args.size)[2] if cond["kind"] == "image" else TEXT_CHUNK + for start in range(0, len(flow), budget): + tasks.append((cond, start, min(start + budget, len(flow)))) + calls = len(tasks) + sum(1 for c, *_ in tasks if c["kind"] in ("compact", "handoff")) + print( + f"corpus={len(flow):,} chars ({len(offsets):,} passages), {len(conditions)} conditions, " + f"{len(tasks)} chunks, <= {calls} API calls, qpc={args.qpc}, model={args.model}" + ) + + api_key = "" if args.report else load_api_key(args.env) + ctx_args = {"args": args, "flow": flow, "paras": paras, "offsets": offsets, "api_key": api_key} + records: list[dict] = [] + done = 0 + with ThreadPoolExecutor(args.workers) as pool: + futures = [pool.submit(run_chunk, cond, start, end, ctx_args) for cond, start, end in tasks] + for fut in futures: + records.extend(fut.result()) + done += 1 + if done % 20 == 0: + print(f" {done}/{len(tasks)} chunks", flush=True) + + with (run_dir / "records.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + + rows = [ + aggregate(cond["name"], [r for r in records if r["cond"] == cond["name"]], args.price_in, args.price_out) + for cond in conditions + if any(r["cond"] == cond["name"] for r in records) + ] + rows.sort(key=lambda r: -r["f1"]) + (run_dir / "summary.json").write_text(json.dumps({"args": vars(args), "rows": rows}, indent=1)) + + hdr = ( + f"{'condition':<15}{'n':>6}{'EM':>7}{'F1':>7}{'±se':>6}{'abst':>6}" + f"{'in tok':>10}{'cache w':>9}{'cache r':>9}{'out tok':>9}{'$':>7}" + ) + print("\n" + hdr + "\n" + "-" * len(hdr)) + for r in rows: + print( + f"{r['name']:<15}{r['n']:>6}{r['em']:>7.3f}{r['f1']:>7.3f}{r['f1_se']:>6.3f}{r['abstained']:>6}" + f"{r['tokens_in']:>10,}{r['cache_w']:>9,}{r['cache_r']:>9,}{r['tokens_out']:>9,}{r['cost_usd']:>7.2f}" + ) + print(f"\n{'condition':<15} F1 by position quartile (Q1..Q4)") + for r in rows: + cells = " ".join(" - " if q is None else f"{q:.3f}" for q in r["f1_by_quartile"]) + print(f"{r['name']:<15} {cells}") + print(f"\nresults -> {run_dir}/") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_3d_activation_html.py b/packages/snapcompact/research/snapcompact_3d_activation_html.py new file mode 100644 index 000000000..f29726405 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_3d_activation_html.py @@ -0,0 +1,192 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["numpy", "plotly"] +# /// +"""Build an embeddable interactive 3D activation terrain for snapcompact.""" + +from __future__ import annotations + +import argparse +import base64 +import json +from pathlib import Path + +import numpy as np +import plotly.graph_objects as go +from plotly.subplots import make_subplots + +HERE = Path(__file__).resolve().parent + + +def downsample(arr: np.ndarray, cols: int) -> np.ndarray: + if arr.shape[1] <= cols: + return arr + edges = np.linspace(0, arr.shape[1], cols + 1).round().astype(int) + out = np.zeros((arr.shape[0], cols), dtype=np.float32) + for i in range(cols): + lo = edges[i] + hi = max(lo + 1, edges[i + 1]) + out[:, i] = arr[:, lo:hi].mean(axis=1) + return out + + +def norm(arr: np.ndarray, q: float = 0.985) -> np.ndarray: + scale = float(np.quantile(arr, q)) + if scale <= 0: + scale = 1.0 + return np.clip(arr / scale, 0, 1) + + +def image_data_uri(path: Path) -> str: + return "data:image/png;base64," + base64.b64encode(path.read_bytes()).decode() + + +def add_surface(fig: go.Figure, z: np.ndarray, row: int, col: int, name: str, colorscale: str, showscale: bool = False) -> None: + y = np.arange(z.shape[0]) + x = np.arange(z.shape[1]) + fig.add_trace( + go.Surface( + x=x, + y=y, + z=z, + name=name, + colorscale=colorscale, + cmin=0, + cmax=1, + showscale=showscale, + lighting={"ambient": 0.58, "diffuse": 0.72, "specular": 0.28, "roughness": 0.52}, + contours={ + "z": {"show": True, "usecolormap": True, "highlightcolor": "#fff0a8", "project_z": True}, + }, + hovertemplate="layer %{y}
image bin %{x}
Δ %{z:.3f}" + name + "", + ), + row=row, + col=col, + ) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--result-dir", default=str(HERE / "results" / "tensor-heatmap-paddleocr-q7")) + ap.add_argument("--out", default=str(HERE / "results" / "snapcompact-activation-terrain.html")) + ap.add_argument("--bins", type=int, default=150) + args = ap.parse_args() + + result_dir = Path(args.result_dir) + summary = json.loads((result_dir / "summary.json").read_text()) + data = np.load(result_dir / "heatmaps.npz") + answer = norm(downsample(data["answer_binned"], args.bins)) + random = norm(downsample(data["random_binned"], args.bins)) + ratio = norm(downsample(data["ratio_binned"], args.bins), 0.97) + + fig = make_subplots( + rows=2, + cols=2, + specs=[[{"type": "surface"}, {"type": "surface"}], [{"type": "surface", "colspan": 2}, None]], + horizontal_spacing=0.02, + vertical_spacing=0.03, + subplot_titles=("Gold answer erased", "Random equal-size erase", "Answer / random residual scar"), + ) + add_surface(fig, answer, 1, 1, "gold answer mask", "Magma") + add_surface(fig, random, 1, 2, "random mask", "Viridis") + add_surface(fig, ratio, 2, 1, "answer/random ratio", "Inferno", True) + + camera = {"eye": {"x": 1.65, "y": -1.75, "z": 0.82}, "center": {"x": 0, "y": 0, "z": -0.08}} + scene_common = { + "bgcolor": "rgba(0,0,0,0)", + "camera": camera, + "xaxis": {"title": "image-token bins", "gridcolor": "rgba(140,170,180,0.18)", "color": "#94a3aa", "zeroline": False}, + "yaxis": {"title": "decoder layer", "gridcolor": "rgba(140,170,180,0.18)", "color": "#94a3aa", "autorange": "reversed", "dtick": 4}, + "zaxis": {"title": "Δ hidden", "gridcolor": "rgba(140,170,180,0.18)", "color": "#94a3aa", "range": [0, 1]}, + "aspectratio": {"x": 2.6, "y": 0.78, "z": 0.52}, + } + fig.update_layout( + template="plotly_dark", + paper_bgcolor="rgba(0,0,0,0)", + plot_bgcolor="rgba(0,0,0,0)", + height=920, + margin={"l": 0, "r": 0, "t": 58, "b": 0}, + font={"family": "Arial, sans-serif", "color": "#efeede"}, + scene=scene_common, + scene2=scene_common, + scene3={**scene_common, "aspectratio": {"x": 3.2, "y": 0.78, "z": 0.58}}, + coloraxis_showscale=False, + ) + fig.update_annotations(font={"size": 18, "color": "#efeede"}) + + q = summary["question"] + original_uri = image_data_uri(result_dir / "images" / "original.png") + masked_uri = image_data_uri(result_dir / "images" / "answer-mask.png") + graph_html = fig.to_html(full_html=False, include_plotlyjs="cdn", config={"displayModeBar": False, "responsive": True}) + html = f""" + + + + +Snapcompact activation terrain + + + +
+
SNAPCOMPACT WHITEBOX
+

Activation terrain from a missing answer

+
Actual decoder hidden states: layer × image-token bin × ||original − masked||. Drag the terrain to inspect the residual-stream scar.
+
+ +
+
interactive 3D residual terraingold-mask spikes rise where the model reacts to losing the answer glyphs
+
{graph_html}
+
+
+
+ + +""" + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + out.write_text(html) + print(out) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_3d_activation_viz.py b/packages/snapcompact/research/snapcompact_3d_activation_viz.py new file mode 100644 index 000000000..d002182d2 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_3d_activation_viz.py @@ -0,0 +1,230 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["matplotlib", "numpy", "pillow"] +# /// +"""Render a 3D blog visualization of snapcompact hidden-state deltas.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") + +import matplotlib.colors as mcolors +import matplotlib.pyplot as plt +import numpy as np +from matplotlib import cm +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent + +BG = (5, 7, 10) +PANEL = (13, 18, 23) +INK = (241, 239, 224) +MUTED = (139, 151, 156) +CYAN = (80, 220, 255) +RED = (255, 83, 62) +GREEN = (145, 255, 112) +AMBER = (255, 194, 65) + + +def font(size: int, bold: bool = False) -> ImageFont.FreeTypeFont | ImageFont.ImageFont: + candidates = [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ] + for path in candidates: + if path and Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def norm_quantile(arr: np.ndarray, q: float = 0.985) -> np.ndarray: + scale = float(np.quantile(arr, q)) + if scale <= 0: + scale = 1.0 + return np.clip(arr / scale, 0, 1) + + +def downsample(arr: np.ndarray, cols: int) -> np.ndarray: + if arr.shape[1] <= cols: + return arr + edges = np.linspace(0, arr.shape[1], cols + 1).round().astype(int) + out = np.zeros((arr.shape[0], cols), dtype=np.float32) + for i in range(cols): + lo = edges[i] + hi = max(lo + 1, edges[i + 1]) + out[:, i] = arr[:, lo:hi].mean(axis=1) + return out + + +def style_3d(ax, title: str, subtitle: str, color: str) -> None: + ax.set_facecolor((0.02, 0.025, 0.035, 1)) + ax.xaxis.pane.set_facecolor((0.02, 0.025, 0.035, 0.0)) + ax.yaxis.pane.set_facecolor((0.02, 0.025, 0.035, 0.0)) + ax.zaxis.pane.set_facecolor((0.02, 0.025, 0.035, 0.0)) + for axis in (ax.xaxis, ax.yaxis, ax.zaxis): + axis._axinfo["grid"]["color"] = (0.28, 0.34, 0.38, 0.20) + axis._axinfo["tick"]["color"] = (0.80, 0.84, 0.84, 0.55) + ax.tick_params(colors="#8b979c", labelsize=8, pad=0) + ax.set_xlabel("image-token bins", color="#8b979c", labelpad=6) + ax.set_ylabel("decoder layer", color="#8b979c", labelpad=6) + ax.set_zlabel("Δ hidden", color="#8b979c", labelpad=5) + ax.set_title(title, color=color, fontsize=16, fontweight="bold", loc="left", pad=10) + ax.text2D(0.0, 0.94, subtitle, transform=ax.transAxes, color="#8b979c", fontsize=9) + ax.view_init(elev=31, azim=-58) + ax.set_box_aspect((2.7, 0.85, 0.55)) + + +def draw_surface(ax, arr: np.ndarray, cmap_name: str, title: str, subtitle: str, color: str, zmax: float = 1.0) -> None: + y = np.arange(arr.shape[0]) + x = np.arange(arr.shape[1]) + X, Y = np.meshgrid(x, y) + Z = arr * zmax + cmap = cm.get_cmap(cmap_name) + ax.plot_surface( + X, + Y, + Z, + rstride=1, + cstride=1, + facecolors=cmap(arr), + linewidth=0, + antialiased=True, + shade=False, + alpha=0.98, + ) + # A dark floor with projected contour lines makes the shape read as 3D. + ax.contour(X, Y, Z, zdir="z", offset=-0.05, levels=9, cmap=cmap, linewidths=0.8, alpha=0.72) + ax.set_zlim(-0.05, zmax) + ax.set_ylim(arr.shape[0] - 1, 0) + style_3d(ax, title, subtitle, color) + + +def crop_with_box(img: Image.Image, start: int, end: int, cols: int, adv: int, pitch: int, pad_cells: int = 34) -> Image.Image: + row0 = max(0, start // cols - 5) + row1 = min(img.height // pitch, end // cols + 6) + col0 = max(0, start % cols - pad_cells) + col1 = min(cols, end % cols + pad_cells) + if col1 <= col0: + col1 = min(cols, col0 + 72) + crop = img.crop((col0 * adv, row0 * pitch, col1 * adv, row1 * pitch)).convert("RGB") + d = ImageDraw.Draw(crop) + bx0 = max(0, (start % cols - col0) * adv - adv) + bx1 = min(crop.width - 1, ((end - 1) % cols - col0 + 2) * adv) + by0 = max(0, (start // cols - row0) * pitch - 1) + by1 = min(crop.height - 1, ((end - 1) // cols - row0 + 1) * pitch + 1) + d.rounded_rectangle((bx0, by0, bx1, by1), radius=3, outline=RED, width=3) + return crop + + +def paste_fit(canvas: Image.Image, img: Image.Image, box: tuple[int, int, int, int]) -> None: + x0, y0, x1, y1 = box + scale = min((x1 - x0) / img.width, (y1 - y0) / img.height) + resized = img.resize((max(1, round(img.width * scale)), max(1, round(img.height * scale))), Image.Resampling.NEAREST) + canvas.paste(resized, (x0 + (x1 - x0 - resized.width) // 2, y0 + (y1 - y0 - resized.height) // 2)) + + +def render_matplotlib_panel(answer: np.ndarray, random: np.ndarray, ratio: np.ndarray) -> Image.Image: + fig = plt.figure(figsize=(16.6, 9.0), dpi=170) + fig.patch.set_facecolor("#05070a") + gs = fig.add_gridspec(2, 2, left=0.02, right=0.99, top=0.96, bottom=0.04, wspace=0.03, hspace=0.08) + ax1 = fig.add_subplot(gs[0, 0], projection="3d") + ax2 = fig.add_subplot(gs[0, 1], projection="3d") + ax3 = fig.add_subplot(gs[1, :], projection="3d") + draw_surface(ax1, answer, "magma", "Gold answer mask", "true answer cells erased", "#ff533e") + draw_surface(ax2, random, "viridis", "Random control mask", "same-sized blank elsewhere", "#91ff70") + draw_surface(ax3, ratio, "inferno", "Answer / random ratio", "where the missing answer leaves a larger residual-stream scar", "#ffc241", zmax=1.08) + tmp = HERE / "results" / ".snapcompact-3d-panel.png" + fig.savefig(tmp, facecolor=fig.get_facecolor(), transparent=False) + plt.close(fig) + img = Image.open(tmp).convert("RGB") + tmp.unlink(missing_ok=True) + return img + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--result-dir", default=str(HERE / "results" / "tensor-heatmap-paddleocr-q7")) + ap.add_argument("--out", default=str(HERE / "results" / "snapcompact-3d-activation-terrain.png")) + ap.add_argument("--bins", type=int, default=128) + args = ap.parse_args() + + result_dir = Path(args.result_dir) + summary = json.loads((result_dir / "summary.json").read_text()) + data = np.load(result_dir / "heatmaps.npz") + answer = norm_quantile(downsample(data["answer_binned"], args.bins)) + random = norm_quantile(downsample(data["random_binned"], args.bins)) + ratio = norm_quantile(downsample(data["ratio_binned"], args.bins), 0.97) + + panel = render_matplotlib_panel(answer, random, ratio) + w, h = 2200, 1320 + canvas = Image.new("RGB", (w, h), BG) + draw = ImageDraw.Draw(canvas) + for y in range(0, h, 16): + draw.line((0, y, w, y), fill=(8, 10 + y % 10, 14 + y % 12)) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-260, -160, 950, 600), fill=(255, 83, 62, 34)) + gd.ellipse((1100, 220, 2500, 1500), fill=(80, 220, 255, 28)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(82))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + draw.text((64, 42), "SNAPCOMPACT WHITEBOX", fill=AMBER, font=font(24, True)) + draw.text((64, 82), "Activation terrain from a missing answer", fill=INK, font=font(68, True)) + draw.text((66, 168), "Actual decoder hidden states: layer × image-token bin × ||original − masked||. A blog-friendly 3D tensor slice, not a schematic.", fill=MUTED, font=font(26)) + + # Left evidence strip. + draw.rounded_rectangle((64, 238, 600, 1234), radius=30, fill=PANEL, outline=(31, 42, 50), width=1) + q = summary["question"] + draw.text((96, 270), "visual intervention", fill=INK, font=font(32, True)) + draw.text((96, 310), "answer cells are blanked", fill=MUTED, font=font(18)) + base = Image.open(result_dir / "images" / "original.png").convert("RGB") + masked = Image.open(result_dir / "images" / "answer-mask.png").convert("RGB") + cols = summary["geometry"]["cols"] + crop = crop_with_box(base, q["answer_start"], q["answer_end"], cols, 8, 13) + masked_crop = crop_with_box(masked, q["answer_start"], q["answer_end"], cols, 8, 13) + draw.text((96, 366), "ORIGINAL", fill=CYAN, font=font(17, True)) + draw.rounded_rectangle((96, 394, 568, 560), radius=14, fill=(244, 242, 230), outline=CYAN, width=3) + paste_fit(canvas, crop, (112, 410, 552, 544)) + draw.text((96, 618), "ANSWER ERASED", fill=RED, font=font(17, True)) + draw.rounded_rectangle((96, 646, 568, 812), radius=14, fill=(244, 242, 230), outline=RED, width=3) + paste_fit(canvas, masked_crop, (112, 662, 552, 796)) + question = q["q"] + if len(question) > 54: + question = question[:51] + "…" + draw.text((96, 890), "question", fill=MUTED, font=font(16, True)) + draw.text((96, 920), question, fill=INK, font=font(22)) + draw.text((96, 990), "gold answer", fill=MUTED, font=font(16, True)) + draw.text((96, 1024), str(q["answer_text"]), fill=AMBER, font=font(42, True)) + draw.text((96, 1110), f"{summary['layers']} layers", fill=MUTED, font=font(20)) + draw.text((96, 1142), f"{summary['image_tokens']} image tokens", fill=MUTED, font=font(20)) + draw.text((96, 1174), f"answer/random Δ = {summary['answer_over_random_delta']:.2f}×", fill=INK, font=font(22, True)) + + # Main 3D panel. + draw.rounded_rectangle((632, 238, 2134, 1234), radius=30, fill=PANEL, outline=(31, 42, 50), width=1) + panel = panel.resize((1450, 786), Image.Resampling.LANCZOS) + canvas.paste(panel, (660, 330)) + draw.text((672, 268), "3D residual-stream delta terrain", fill=INK, font=font(36, True)) + draw.text((672, 311), "Gold-mask spikes rise where the model’s image-token activations react to losing the answer glyphs.", fill=MUTED, font=font(20)) + + # Color scale. + cmap = cm.get_cmap("magma") + for i in range(260): + rgb = tuple(int(v * 255) for v in cmap(i / 259)[:3]) + draw.rectangle((1810 + i, 274, 1811 + i, 292), fill=rgb) + draw.text((1810, 246), "low Δ", fill=MUTED, font=font(14)) + draw.text((2018, 246), "high Δ", fill=MUTED, font=font(14)) + + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + canvas.save(out) + print(out) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_activation_probe.py b/packages/snapcompact/research/snapcompact_activation_probe.py new file mode 100644 index 000000000..b65036dfe --- /dev/null +++ b/packages/snapcompact/research/snapcompact_activation_probe.py @@ -0,0 +1,225 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy", "torch", "transformers", "sentencepiece", "protobuf", "einops"] +# /// +"""White-box snapcompact activation pilot for a local Hugging Face VLM. + +Runs a tiny dense-bitmap/text paired corpus through a local VLM and compares +hidden states across carriers. The default targets the PaddleOCR-VL snapshot +available on spark.internal because the served Qwen2.5-VL NVFP4 checkpoint is a +vLLM/modelopt artifact that Transformers cannot load directly. +""" + +from __future__ import annotations + +import argparse +import json +import math +import random +import sys +from pathlib import Path +from typing import Any + +import numpy as np +from PIL import Image, ImageDraw + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import capacity, render # noqa: E402 +from run import CACHE, FONTS, load_prompt # noqa: E402 +from snapcompact_blackbox_occlusion import mask_cells, random_span, sample_answer_questions # noqa: E402 + +DEFAULT_MODEL_DIR = ( + "/home/can/.cache/huggingface/hub/models--PaddlePaddle--PaddleOCR-VL/" + "snapshots/2b77538ef936207f60c16b45082841068987d08c" +) + + +def centered_gram(x: np.ndarray) -> np.ndarray: + gram = x @ x.T + row_mean = gram.mean(axis=1, keepdims=True) + col_mean = gram.mean(axis=0, keepdims=True) + return gram - row_mean - col_mean + gram.mean() + + +def linear_cka(x: np.ndarray, y: np.ndarray) -> float: + if x.shape[0] < 2 or y.shape[0] < 2: + return float("nan") + x = x - x.mean(axis=0, keepdims=True) + y = y - y.mean(axis=0, keepdims=True) + k = centered_gram(x) + l = centered_gram(y) + denom = math.sqrt(float((k * k).sum()) * float((l * l).sum())) + if denom == 0: + return float("nan") + return float((k * l).sum() / denom) + + +def paired_cosine(x: np.ndarray, y: np.ndarray) -> float: + dot = (x * y).sum(axis=1) + denom = np.linalg.norm(x, axis=1) * np.linalg.norm(y, axis=1) + valid = denom > 0 + if not valid.any(): + return float("nan") + return float((dot[valid] / denom[valid]).mean()) + + +def make_prompt(q: str, cols: int, rows: int) -> str: + return ( + load_prompt("qa-image.md").format(cols=cols, rows=rows) + + f"\n\nQuestion: {q}\nAnswer with only the shortest extractive answer." + ) + + +def to_device(batch: dict[str, Any], device: Any) -> dict[str, Any]: + return {k: (v.to(device) if hasattr(v, "to") else v) for k, v in batch.items()} + + +def hidden_features(model: Any, processor: Any, *, image: Image.Image | None, text: str, device: Any) -> list[np.ndarray]: + import torch + + if image is None: + messages = [{"role": "user", "content": [{"type": "text", "text": text}]}] + templated = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + batch = processor(text=templated, return_tensors="pt") + else: + messages = [{"role": "user", "content": [{"type": "image", "image": image}, {"type": "text", "text": text}]}] + templated = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + batch = processor(images=image, text=templated, return_tensors="pt") + batch = to_device(batch, device) + with torch.no_grad(): + out = model(**batch, output_hidden_states=True, output_attentions=False, use_cache=False) + feats: list[np.ndarray] = [] + for h in out.hidden_states: + # Mean-pool the prompt sequence. This avoids brittle alignment between + # image-token and text-token positions while preserving layer geometry. + pooled = h.float().mean(dim=1).detach().cpu().numpy()[0] + feats.append(pooled.astype(np.float32, copy=False)) + return feats + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--model-dir", default=DEFAULT_MODEL_DIR) + ap.add_argument("--font", default="5x8", choices=sorted(FONTS)) + ap.add_argument("--variant", default="bw") + ap.add_argument("--size", type=int, default=512) + ap.add_argument("--limit-paras", type=int, default=20) + ap.add_argument("--qpc", type=int, default=8) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--out", default="snapcompact-activation") + args = ap.parse_args() + + import torch + from transformers import AutoModel, AutoProcessor + + out_dir = HERE / "results" / args.out + img_dir = out_dir / "activation-images" + img_dir.mkdir(parents=True, exist_ok=True) + CACHE.mkdir(exist_ok=True) + + cfg = FONTS[args.font] + cols, rows, budget = capacity(cfg, args.size) + paras = squad.load_paragraphs(CACHE)[: args.limit_paras] + flow, offsets = squad.build_flow(paras) + start, end = 0, min(len(flow), budget) + chunk = flow[start:end] + questions = sample_answer_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + raise SystemExit("no sampled questions fit in the activation chunk") + + base_img = render(chunk, cfg, CACHE, args.size, args.variant) + base_path = img_dir / "base.png" + base_img.save(base_path) + fill = (255, 255, 255) if args.variant not in ("dark", "dark-sent") else (0, 0, 0) + + print(f"loading {args.model_dir}", flush=True) + processor = AutoProcessor.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True, use_fast=False) + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + dtype = torch.bfloat16 if device.type == "cuda" else torch.float32 + model = AutoModel.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True, dtype=dtype).to(device).eval() + + feature_sets: dict[str, list[list[np.ndarray]]] = {"text": [], "image": [], "answer_mask": [], "random_mask": []} + records: list[dict[str, Any]] = [] + for qi, q in enumerate(questions): + span_len = max(1, q["answer_end"] - q["answer_start"]) + rng = random.Random(args.seed * 31 + qi) + rand_start, rand_end = random_span(rng, len(chunk), span_len, q["answer_start"], q["answer_end"]) + answer_img = mask_cells(base_img, q["answer_start"], q["answer_end"], cols, cfg.adv, cfg.pitch, fill) + random_img = mask_cells(base_img, rand_start, rand_end, cols, cfg.adv, cfg.pitch, fill) + answer_path = img_dir / f"q{qi}-answer-mask.png" + random_path = img_dir / f"q{qi}-random-mask.png" + answer_img.save(answer_path) + random_img.save(random_path) + + img_prompt = make_prompt(q["q"], cols, rows) + text_prompt = ( + "Below is reference material. Answer the question using only it.\n\n" + f"{chunk}\n\nQuestion: {q['q']}\n" + "Answer with only the shortest extractive answer." + ) + feature_sets["text"].append(hidden_features(model, processor, image=None, text=text_prompt, device=device)) + feature_sets["image"].append(hidden_features(model, processor, image=base_img, text=img_prompt, device=device)) + feature_sets["answer_mask"].append(hidden_features(model, processor, image=answer_img, text=img_prompt, device=device)) + feature_sets["random_mask"].append(hidden_features(model, processor, image=random_img, text=img_prompt, device=device)) + records.append( + { + "question_index": qi, + "q": q["q"], + "golds": q["golds"], + "answer_text": q["answer_text"], + "answer_start": q["answer_start"], + "answer_end": q["answer_end"], + "random_start": rand_start, + "random_end": rand_end, + } + ) + print(f"captured {qi + 1}/{len(questions)}", flush=True) + + layer_count = len(feature_sets["image"][0]) + layers = [] + for layer in range(layer_count): + arrays = { + name: np.stack([sample[layer] for sample in samples], axis=0) + for name, samples in feature_sets.items() + } + img = arrays["image"] + ans = arrays["answer_mask"] + rnd = arrays["random_mask"] + answer_delta = np.linalg.norm(img - ans, axis=1) + random_delta = np.linalg.norm(img - rnd, axis=1) + layers.append( + { + "layer": layer, + "cka_text_image": linear_cka(arrays["text"], img), + "cka_image_answer_mask": linear_cka(img, ans), + "cka_image_random_mask": linear_cka(img, rnd), + "cos_text_image": paired_cosine(arrays["text"], img), + "cos_image_answer_mask": paired_cosine(img, ans), + "cos_image_random_mask": paired_cosine(img, rnd), + "answer_delta_norm": float(answer_delta.mean()), + "random_delta_norm": float(random_delta.mean()), + "answer_over_random_delta": float(answer_delta.mean() / random_delta.mean()) if random_delta.mean() else float("inf"), + } + ) + + summary = { + "args": vars(args), + "model_dir": args.model_dir, + "device": str(device), + "geometry": {"cols": cols, "rows": rows, "capacity": budget}, + "n": len(records), + "layers": layers, + } + with (out_dir / "records.jsonl").open("w") as fh: + for record in records: + fh.write(json.dumps(record) + "\n") + (out_dir / "summary.json").write_text(json.dumps(summary, indent=1)) + print(json.dumps(summary, indent=1)) + print(f"results -> {out_dir}") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_blackbox_occlusion.py b/packages/snapcompact/research/snapcompact_blackbox_occlusion.py new file mode 100644 index 000000000..15faac599 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_blackbox_occlusion.py @@ -0,0 +1,261 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""Black-box snapcompact occlusion probe against an OpenAI-compatible VLM endpoint. + +For sampled SQuAD questions, render the carrier as a dense bitmap, then compare +QA on the original image, an image with the gold answer cells masked, and an +image with an equal-sized random mask. A real visual-retrieval mechanism should +show a larger F1 drop for answer masks than random masks. +""" + +from __future__ import annotations + +import argparse +import base64 +import hashlib +import json +import random +import sys +import urllib.error +import urllib.request +from pathlib import Path +from typing import Any + +from PIL import Image, ImageDraw + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import capacity, render # noqa: E402 +from run import CACHE, FONTS, load_prompt, sha8 # noqa: E402 + + +def sample_answer_questions(paras: list[dict], offsets: list[int], start: int, end: int, n: int, seed: int) -> list[dict]: + """Sample questions like squad.sample_chunk_questions, preserving answer offsets.""" + rng = random.Random(seed * 1_000_003 + start) + eligible = [ + i + for i in range(len(offsets)) + if offsets[i] >= start and offsets[i] + len(paras[i]["ctx"]) <= end and paras[i].get("qas") + ] + if not eligible: + return [] + n = min(n, len(eligible)) + step = len(eligible) / n + picked: list[dict] = [] + for k in range(n): + pi = eligible[int(k * step)] + qa = rng.choice(paras[pi]["qas"]) + answers = qa.get("answers") or [] + if not answers: + continue + answer = answers[0] + picked.append( + { + "q": " ".join(qa["question"].split()), + "golds": sorted({a["text"] for a in answers}), + "answer_text": answer["text"], + "answer_start": offsets[pi] - start + int(answer["answer_start"]), + "answer_end": offsets[pi] - start + int(answer["answer_start"]) + len(answer["text"]), + "pos_rel": (offsets[pi] - start) / (end - start), + } + ) + return picked + + +def mask_cells(img: Image.Image, start: int, end: int, cols: int, adv: int, pitch: int, fill: tuple[int, int, int]) -> Image.Image: + out = img.copy() + draw = ImageDraw.Draw(out) + start = max(0, start) + end = max(start + 1, end) + first_row = start // cols + last_row = (end - 1) // cols + for row in range(first_row, last_row + 1): + c0 = start % cols if row == first_row else 0 + c1 = (end - 1) % cols if row == last_row else cols - 1 + x0 = max(0, c0 * adv - adv) + y0 = max(0, row * pitch - 1) + x1 = min(out.width, (c1 + 2) * adv) + y1 = min(out.height, (row + 1) * pitch + 1) + draw.rectangle((x0, y0, x1, y1), fill=fill) + return out + + +def random_span(rng: random.Random, text_len: int, span_len: int, avoid_start: int, avoid_end: int) -> tuple[int, int]: + if text_len <= span_len: + return 0, text_len + for _ in range(100): + start = rng.randrange(0, text_len - span_len) + end = start + span_len + if end < avoid_start - span_len or start > avoid_end + span_len: + return start, end + start = 0 if avoid_start > text_len // 2 else max(0, text_len - span_len) + return start, min(text_len, start + span_len) + + +def post_chat(endpoint: str, model: str, image_path: Path, prompt: str, max_tokens: int, cache_dir: Path, fresh: bool) -> tuple[str, dict]: + payload_key = sha8(model, prompt, hashlib.sha1(image_path.read_bytes()).hexdigest()) + cache_path = cache_dir / f"{payload_key}.json" + if cache_path.exists() and not fresh: + cached = json.loads(cache_path.read_text()) + return cached["text"], cached.get("usage", {}) + + image_b64 = base64.b64encode(image_path.read_bytes()).decode() + body = { + "model": model, + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": prompt}, + {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_b64}"}}, + ], + } + ], + "temperature": 0, + "max_tokens": max_tokens, + } + req = urllib.request.Request( + endpoint, + data=json.dumps(body).encode(), + headers={"content-type": "application/json"}, + method="POST", + ) + try: + with urllib.request.urlopen(req, timeout=180) as resp: + out = json.loads(resp.read().decode()) + except urllib.error.HTTPError as err: + raise RuntimeError(err.read().decode()) from err + choice = (out.get("choices") or [{}])[0] + text = ((choice.get("message") or {}).get("content") or "").strip() + usage = out.get("usage") or {} + cache_path.write_text(json.dumps({"text": text, "usage": usage}, indent=1)) + return text, usage + + +def aggregate(records: list[dict]) -> dict: + by_variant: dict[str, list[dict]] = {} + for record in records: + by_variant.setdefault(record["variant"], []).append(record) + out: dict[str, Any] = {"n": len(records) // 3, "variants": {}} + for name, rows in sorted(by_variant.items()): + out["variants"][name] = { + "n": len(rows), + "em": sum(r["em"] for r in rows) / max(1, len(rows)), + "f1": sum(r["f1"] for r in rows) / max(1, len(rows)), + "abstained": sum(1 for r in rows if "unreadable" in r["answer"].lower()), + "prompt_tokens": sum((r.get("usage") or {}).get("prompt_tokens", 0) for r in rows), + "completion_tokens": sum((r.get("usage") or {}).get("completion_tokens", 0) for r in rows), + } + base = out["variants"].get("original", {}).get("f1", 0.0) + out["drops"] = { + name: base - row["f1"] for name, row in out["variants"].items() if name != "original" + } + return out + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--endpoint", default="http://spark.internal:8000/v1/chat/completions") + ap.add_argument("--model", default="Qwen2.5-VL-7B-Instruct-NVFP4") + ap.add_argument("--font", default="5x8", choices=sorted(FONTS)) + ap.add_argument("--variant", default="bw") + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--limit-paras", type=int, default=50) + ap.add_argument("--qpc", type=int, default=12) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--max-tokens", type=int, default=48) + ap.add_argument("--out", default="snapcompact-occlusion") + ap.add_argument("--fresh", action="store_true") + args = ap.parse_args() + + out_dir = HERE / "results" / args.out + img_dir = out_dir / "images" + cache_dir = out_dir / "api-cache" + img_dir.mkdir(parents=True, exist_ok=True) + cache_dir.mkdir(parents=True, exist_ok=True) + CACHE.mkdir(exist_ok=True) + + cfg = FONTS[args.font] + cols, rows, budget = capacity(cfg, args.size) + paras = squad.load_paragraphs(CACHE)[: args.limit_paras] + flow, offsets = squad.build_flow(paras) + prompt_base = load_prompt("qa-image.md").format(cols=cols, rows=rows) + mask_fill = (255, 255, 255) if args.variant not in ("dark", "dark-sent") else (0, 0, 0) + + tasks = [] + for start in range(0, len(flow), budget): + end = min(start + budget, len(flow)) + chunk = flow[start:end] + questions = sample_answer_questions(paras, offsets, start, end, args.qpc, args.seed) + if questions: + tasks.append((start, end, chunk, questions)) + + records: list[dict] = [] + for chunk_index, (start, end, chunk, questions) in enumerate(tasks): + base_img = render(chunk, cfg, CACHE, args.size, args.variant) + base_path = img_dir / f"chunk-{start}-{args.font}-{args.variant}.png" + if not base_path.exists(): + base_img.save(base_path) + for qi, q in enumerate(questions): + span_len = max(1, q["answer_end"] - q["answer_start"]) + rng = random.Random(args.seed * 17 + start + qi) + rand_start, rand_end = random_span(rng, len(chunk), span_len, q["answer_start"], q["answer_end"]) + answer_path = img_dir / f"chunk-{start}-q{qi}-answer-mask.png" + random_path = img_dir / f"chunk-{start}-q{qi}-random-mask.png" + if not answer_path.exists(): + mask_cells(base_img, q["answer_start"], q["answer_end"], cols, cfg.adv, cfg.pitch, mask_fill).save(answer_path) + if not random_path.exists(): + mask_cells(base_img, rand_start, rand_end, cols, cfg.adv, cfg.pitch, mask_fill).save(random_path) + + prompt = ( + f"{prompt_base}\n\nQuestion: {q['q']}\n" + "Answer with only the shortest extractive answer copied from the image. " + "If the answer is unreadable, reply exactly UNREADABLE." + ) + for variant_name, path in ( + ("original", base_path), + ("answer_mask", answer_path), + ("random_mask", random_path), + ): + answer, usage = post_chat(args.endpoint, args.model, path, prompt, args.max_tokens, cache_dir, args.fresh) + records.append( + { + "chunk": start, + "chunk_index": chunk_index, + "question_index": qi, + "variant": variant_name, + "q": q["q"], + "answer": answer, + "golds": q["golds"], + "answer_text": q["answer_text"], + "answer_start": q["answer_start"], + "answer_end": q["answer_end"], + "random_start": rand_start, + "random_end": rand_end, + "pos_rel": q["pos_rel"], + "em": squad.exact_match(answer, q["golds"]), + "f1": squad.f1(answer, q["golds"]), + "usage": usage, + } + ) + print(f"{len(records):04d} {variant_name:<11} f1={records[-1]['f1']:.3f} answer={answer[:80]!r}", flush=True) + + summary = { + "args": vars(args), + "geometry": {"cols": cols, "rows": rows, "capacity": budget}, + **aggregate(records), + } + with (out_dir / "records.jsonl").open("w") as fh: + for record in records: + fh.write(json.dumps(record) + "\n") + (out_dir / "summary.json").write_text(json.dumps(summary, indent=1)) + print(json.dumps(summary, indent=1)) + print(f"results -> {out_dir}") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_blog_viz.py b/packages/snapcompact/research/snapcompact_blog_viz.py new file mode 100644 index 000000000..c431c2205 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_blog_viz.py @@ -0,0 +1,231 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""Render a blog-ready snapcompact white-box visualization from pilot outputs.""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path + +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent + +PALETTE = { + "bg": (8, 11, 14), + "panel": (18, 24, 30), + "panel2": (13, 18, 24), + "grid": (47, 61, 72), + "text": (234, 238, 229), + "muted": (139, 151, 156), + "accent": (255, 104, 72), + "accent2": (67, 210, 255), + "green": (158, 255, 121), + "amber": (255, 197, 74), + "red": (255, 70, 70), +} + + +def font(size: int, bold: bool = False) -> ImageFont.FreeTypeFont | ImageFont.ImageFont: + candidates = [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/System/Library/Fonts/Monaco.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ] + for path in candidates: + if path and Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def rounded(draw: ImageDraw.ImageDraw, xy: tuple[int, int, int, int], fill: tuple[int, int, int], outline=None, radius=24, width=1) -> None: + draw.rounded_rectangle(xy, radius=radius, fill=fill, outline=outline, width=width) + + +def draw_label(draw: ImageDraw.ImageDraw, xy: tuple[int, int], text: str, color: tuple[int, int, int], size: int = 24, bold: bool = False) -> None: + draw.text(xy, text, fill=color, font=font(size, bold=bold)) + + +def chart( + draw: ImageDraw.ImageDraw, + box: tuple[int, int, int, int], + series: list[tuple[str, list[float], tuple[int, int, int]]], + y_min: float, + y_max: float, + title: str, + subtitle: str, +) -> None: + x0, y0, x1, y1 = box + rounded(draw, box, PALETTE["panel"], outline=(35, 47, 56), radius=22) + draw_label(draw, (x0 + 28, y0 + 22), title, PALETTE["text"], 28, True) + draw_label(draw, (x0 + 28, y0 + 57), subtitle, PALETTE["muted"], 17) + gx0, gy0, gx1, gy1 = x0 + 58, y0 + 98, x1 - 30, y1 - 58 + for i in range(5): + y = gy0 + round((gy1 - gy0) * i / 4) + draw.line((gx0, y, gx1, y), fill=PALETTE["grid"], width=1) + value = y_max - (y_max - y_min) * i / 4 + draw.text((x0 + 18, y - 9), f"{value:.2f}", fill=PALETTE["muted"], font=font(13)) + n = len(series[0][1]) + for label, values, color in series: + pts = [] + for i, value in enumerate(values): + x = gx0 + (gx1 - gx0) * i / max(1, n - 1) + y = gy1 - (gy1 - gy0) * (value - y_min) / (y_max - y_min) + pts.append((round(x), round(y))) + draw.line(pts, fill=color, width=4, joint="curve") + for p in pts[:: max(1, n // 6)]: + draw.ellipse((p[0] - 4, p[1] - 4, p[0] + 4, p[1] + 4), fill=color) + lx = gx0 + ly = y1 - 36 + for label, _values, color in series: + draw.rounded_rectangle((lx, ly, lx + 20, ly + 10), radius=5, fill=color) + draw.text((lx + 28, ly - 5), label, fill=PALETTE["muted"], font=font(15)) + lx += 210 + + +def crop_with_box(img: Image.Image, start: int, end: int, cols: int, adv: int, pitch: int, pad_cells: int = 24) -> Image.Image: + row0 = max(0, start // cols - 4) + row1 = min(img.height // pitch, end // cols + 5) + col0 = max(0, start % cols - pad_cells) + col1 = min(cols, end % cols + pad_cells) + if row1 <= row0: + row1 = min(img.height // pitch, row0 + 8) + if col1 <= col0: + col1 = min(cols, col0 + 48) + x0, y0, x1, y1 = col0 * adv, row0 * pitch, col1 * adv, row1 * pitch + crop = img.crop((x0, y0, x1, y1)).convert("RGB") + draw = ImageDraw.Draw(crop) + bx0 = max(0, (start % cols - col0) * adv - adv) + bx1 = min(crop.width - 1, ((end - 1) % cols - col0 + 2) * adv) + by0 = max(0, (start // cols - row0) * pitch - 1) + by1 = min(crop.height - 1, ((end - 1) // cols - row0 + 1) * pitch + 1) + draw.rounded_rectangle((bx0, by0, bx1, by1), radius=3, outline=PALETTE["red"], width=3) + return crop + + +def paste_fit(canvas: Image.Image, img: Image.Image, box: tuple[int, int, int, int]) -> None: + x0, y0, x1, y1 = box + scale = min((x1 - x0) / img.width, (y1 - y0) / img.height) + resized = img.resize((max(1, round(img.width * scale)), max(1, round(img.height * scale))), Image.Resampling.NEAREST) + px = x0 + (x1 - x0 - resized.width) // 2 + py = y0 + (y1 - y0 - resized.height) // 2 + canvas.paste(resized, (px, py)) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--activation", default=str(HERE / "results" / "activation-paddleocr-8x13-n16")) + ap.add_argument("--occlusion", default=str(HERE / "results" / "snapcompact-occlusion-qwen-8x13")) + ap.add_argument("--out", default=str(HERE / "results" / "snapcompact-blog-whitebox.png")) + args = ap.parse_args() + + act_dir = Path(args.activation) + occ_dir = Path(args.occlusion) + summary = json.loads((act_dir / "summary.json").read_text()) + occ = json.loads((occ_dir / "summary.json").read_text()) + records = [json.loads(line) for line in (act_dir / "records.jsonl").read_text().splitlines() if line] + layers = summary["layers"] + + w, h = 1800, 1040 + canvas = Image.new("RGB", (w, h), PALETTE["bg"]) + draw = ImageDraw.Draw(canvas) + # Background texture. + for y in range(0, h, 18): + color = (10 + y % 17, 13 + y % 13, 17 + y % 11) + draw.line((0, y, w, y), fill=color) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-160, -220, 760, 520), fill=(255, 104, 72, 34)) + gd.ellipse((1100, 130, 2100, 1160), fill=(67, 210, 255, 28)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(70))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + draw_label(draw, (56, 38), "SNAPCOMPACT UNDER THE MICROSCOPE", PALETTE["amber"], 21, True) + draw_label(draw, (56, 78), "Dense text-images leave a white-box trace", PALETTE["text"], 54, True) + draw_label(draw, (58, 145), "Text and image prompts converge late; blanking the gold answer region perturbs hidden states far more than an equal random blank.", PALETTE["muted"], 24) + + # Big stat cards. + stats = [ + ("Qwen black-box F1", f"{occ['variants']['original']['f1']:.2f}", "original image"), + ("gold-mask drop", f"−{occ['drops']['answer_mask']:.2f}", "answer region blanked"), + ("random-mask drop", f"−{occ['drops']['random_mask']:.2f}", "same-size random blank"), + ] + sx = 56 + card_w = 258 + for title, value, caption in stats: + rounded(draw, (sx, 205, sx + card_w, 330), PALETTE["panel2"], outline=(36, 48, 56), radius=22) + draw_label(draw, (sx + 20, 226), title, PALETTE["muted"], 16) + draw_label(draw, (sx + 20, 252), value, PALETTE["text"], 42, True) + draw_label(draw, (sx + 20, 300), caption, PALETTE["muted"], 14) + sx += card_w + 22 + + chart( + draw, + (56, 368, 872, 668), + [ + ("text ↔ image CKA", [x["cka_text_image"] for x in layers], PALETTE["accent2"]), + ("answer-mask CKA", [x["cka_image_answer_mask"] for x in layers], PALETTE["accent"]), + ("random-mask CKA", [x["cka_image_random_mask"] for x in layers], PALETTE["green"]), + ], + 0.2, + 1.0, + "Layer geometry", + "PaddleOCR-VL hidden-state similarity across 19 decoder layers", + ) + chart( + draw, + (56, 698, 872, 990), + [ + ("answer / random perturbation", [x["answer_over_random_delta"] for x in layers], PALETTE["amber"]), + ], + 1.0, + 1.6, + "Causal-ish scar", + "Mean hidden-state shift: gold answer mask divided by random mask", + ) + + # Visual crop panel. + panel = (920, 205, 1744, 990) + rounded(draw, panel, PALETTE["panel"], outline=(35, 47, 56), radius=26) + draw_label(draw, (950, 232), "What the mask test looks like", PALETTE["text"], 34, True) + draw_label(draw, (950, 274), "Same question, same bitmap. Only the gold answer cells are erased.", PALETTE["muted"], 19) + + base = Image.open(act_dir / "activation-images" / "base.png").convert("RGB") + # Use a later question if possible because it gives a better-looking crop. + rec = records[min(7, len(records) - 1)] + ans = Image.open(act_dir / "activation-images" / f"q{rec['question_index']}-answer-mask.png").convert("RGB") + rnd = Image.open(act_dir / "activation-images" / f"q{rec['question_index']}-random-mask.png").convert("RGB") + cols = summary["geometry"]["cols"] + adv = 8 + pitch = 13 + crops = [ + ("original", crop_with_box(base, rec["answer_start"], rec["answer_end"], cols, adv, pitch), PALETTE["accent2"]), + ("answer masked", crop_with_box(ans, rec["answer_start"], rec["answer_end"], cols, adv, pitch), PALETTE["accent"]), + ("random masked", crop_with_box(rnd, rec["random_start"], rec["random_end"], cols, adv, pitch), PALETTE["green"]), + ] + y = 334 + for label, img, color in crops: + draw_label(draw, (950, y - 31), label.upper(), color, 17, True) + rounded(draw, (950, y, 1714, y + 150), (244, 242, 230), outline=color, radius=14, width=3) + paste_fit(canvas, img, (966, y + 16, 1698, y + 134)) + y += 198 + + q = rec["q"] + if len(q) > 92: + q = q[:89] + "…" + draw_label(draw, (950, 916), "sample question", PALETTE["muted"], 17, True) + draw_label(draw, (950, 940), q, PALETTE["text"], 20) + draw_label(draw, (950, 966), f"gold answer: {rec['answer_text']}", PALETTE["amber"], 18, True) + + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + canvas.save(out) + print(out) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_carrier_convergence.py b/packages/snapcompact/research/snapcompact_carrier_convergence.py new file mode 100644 index 000000000..cfdb21af2 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_carrier_convergence.py @@ -0,0 +1,212 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy", "torch", "transformers"] +# /// +"""Do text and image carriers converge to the same internal state? + +For N questions over the same chunk, capture the last-prompt-token hidden state +(the model's "about to answer" summary) per decoder layer, once with the chunk +as raw text and once as a snapcompact bitmap. Carrier-specific means are +subtracted per layer so prompt boilerplate and modality signatures cancel out. + +Evidence of convergence: +1. matched pairs (same question, different carrier) >> mismatched pairs +2. the question-by-question similarity geometry (RSA) is shared across carriers +3. both carriers generate the same answers +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path +from typing import Any + +import numpy as np +from PIL import Image + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import capacity, render # noqa: E402 +from run import CACHE, FONTS, load_prompt # noqa: E402 +from snapcompact_blackbox_occlusion import sample_answer_questions # noqa: E402 + + +def make_text_prompt(chunk: str, question: str) -> str: + return ( + "Below is reference material. Answer the question using only it.\n\n" + f"{chunk}\n\nQuestion: {question}\n" + "Answer with only the shortest extractive answer." + ) + + +def make_image_prompt(cols: int, rows: int, question: str) -> str: + return load_prompt("qa-image.md").format(cols=cols, rows=rows) + f"\n\nQuestion: {question}\nAnswer with only the shortest extractive answer." + + +def capture_last_token(model: Any, processor: Any, device: Any, text: str, image: Image.Image | None) -> tuple[np.ndarray, str]: + """Return per-layer hidden state at the final prompt position plus a short generation.""" + import torch + + content: list[dict[str, Any]] = [] + if image is not None: + content.append({"type": "image", "image": image}) + content.append({"type": "text", "text": text}) + templated = processor.apply_chat_template([{"role": "user", "content": content}], tokenize=False, add_generation_prompt=True) + if image is not None: + batch = processor(images=image, text=templated, return_tensors="pt") + else: + batch = processor(text=templated, return_tensors="pt") + batch = {k: (v.to(device) if hasattr(v, "to") else v) for k, v in batch.items()} + with torch.no_grad(): + out = model(**batch, output_hidden_states=True, use_cache=False) + generated = model.generate(**batch, max_new_tokens=16, do_sample=False) + states = np.stack([h[0, -1, :].float().detach().cpu().numpy() for h in out.hidden_states], axis=0) + answer = processor.batch_decode(generated[:, batch["input_ids"].shape[1] :], skip_special_tokens=True)[0].strip() + return states.astype(np.float32, copy=False), answer + + +def cosine_rows(a: np.ndarray, b: np.ndarray) -> np.ndarray: + a_n = a / np.maximum(np.linalg.norm(a, axis=-1, keepdims=True), 1e-6) + b_n = b / np.maximum(np.linalg.norm(b, axis=-1, keepdims=True), 1e-6) + return a_n @ b_n.T + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--model-dir", default="Qwen/Qwen2.5-VL-7B-Instruct") + ap.add_argument("--font", default="8x13", choices=sorted(FONTS)) + ap.add_argument("--variant", default="bw") + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--limit-paras", type=int, default=80) + ap.add_argument("--questions", type=int, default=12) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--out", default="qwen-carrier-convergence") + args = ap.parse_args() + + import torch + from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration + + out_dir = HERE / "results" / args.out + img_dir = out_dir / "images" + img_dir.mkdir(parents=True, exist_ok=True) + CACHE.mkdir(exist_ok=True) + + cfg = FONTS[args.font] + cols, rows, budget = capacity(cfg, args.size) + paras = squad.load_paragraphs(CACHE)[: args.limit_paras] + flow, offsets = squad.build_flow(paras) + chunk = flow[: min(len(flow), budget)] + questions = sample_answer_questions(paras, offsets, 0, len(chunk), args.questions * 2, args.seed) + # Deduplicate gold answers so the RSA geometry has distinct content per row. + seen: set[str] = set() + picked: list[dict[str, Any]] = [] + for q in questions: + key = q["answer_text"].lower() + if key not in seen: + seen.add(key) + picked.append(q) + if len(picked) >= args.questions: + break + if len(picked) < 4: + raise SystemExit("not enough distinct questions in chunk") + img = render(chunk, cfg, CACHE, args.size, args.variant) + img.save(img_dir / "image-carrier.png") + + print(f"loading {args.model_dir}", flush=True) + processor = AutoProcessor.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True, use_fast=False) + model = Qwen2_5_VLForConditionalGeneration.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True, dtype=torch.bfloat16, device_map="auto").eval() + device = next(model.parameters()).device + + text_states: list[np.ndarray] = [] + image_states: list[np.ndarray] = [] + records: list[dict[str, Any]] = [] + for qi, q in enumerate(picked): + t_states, t_answer = capture_last_token(model, processor, device, make_text_prompt(chunk, q["q"]), None) + i_states, i_answer = capture_last_token(model, processor, device, make_image_prompt(cols, rows, q["q"]), img) + text_states.append(t_states) + image_states.append(i_states) + records.append( + { + "question_index": qi, + "q": q["q"], + "gold": q["answer_text"], + "golds": q["golds"], + "text_answer": t_answer, + "image_answer": i_answer, + "text_em": squad.exact_match(t_answer, q["golds"]), + "image_em": squad.exact_match(i_answer, q["golds"]), + "agree": squad.f1(t_answer, [i_answer]) >= 0.99, + } + ) + print(f"{qi + 1}/{len(picked)} text={t_answer!r} image={i_answer!r} gold={q['answer_text']!r}", flush=True) + + text_arr = np.stack(text_states, axis=0) # [Q, L, D] + image_arr = np.stack(image_states, axis=0) + n_q, n_layers, _dim = text_arr.shape + + layers: list[dict[str, Any]] = [] + text_sim_by_layer = np.zeros((n_layers, n_q, n_q), dtype=np.float32) + image_sim_by_layer = np.zeros((n_layers, n_q, n_q), dtype=np.float32) + cross_sim_by_layer = np.zeros((n_layers, n_q, n_q), dtype=np.float32) + off_diag = ~np.eye(n_q, dtype=bool) + for layer in range(n_layers): + text_l = text_arr[:, layer, :] + image_l = image_arr[:, layer, :] + # Carrier-centering removes modality/prompt signature; what remains is + # per-question content variation within each carrier. + text_c = text_l - text_l.mean(axis=0, keepdims=True) + image_c = image_l - image_l.mean(axis=0, keepdims=True) + cross = cosine_rows(text_c, image_c) + text_sim = cosine_rows(text_c, text_c) + image_sim = cosine_rows(image_c, image_c) + matched = float(np.diag(cross).mean()) + mismatched = float(cross[off_diag].mean()) + rsa = float(np.corrcoef(text_sim[off_diag], image_sim[off_diag])[0, 1]) + layers.append( + { + "layer": layer, + "matched_cosine": matched, + "mismatched_cosine": mismatched, + "separation": matched - mismatched, + "rsa_pearson": rsa, + "match_rank_accuracy": float((np.argmax(cross, axis=1) == np.arange(n_q)).mean()), + } + ) + text_sim_by_layer[layer] = text_sim + image_sim_by_layer[layer] = image_sim + cross_sim_by_layer[layer] = cross + + best_layer = int(np.argmax([l["separation"] for l in layers])) + summary = { + "args": vars(args), + "geometry": {"cols": cols, "rows": rows, "capacity": budget}, + "n_questions": n_q, + "layers": len(layers), + "per_layer": layers, + "best_layer": best_layer, + "best": layers[best_layer], + "final": layers[-1], + "answer_agreement": float(np.mean([r["agree"] for r in records])), + "text_em": float(np.mean([r["text_em"] for r in records])), + "image_em": float(np.mean([r["image_em"] for r in records])), + "records": records, + } + np.savez_compressed( + out_dir / "carrier_convergence.npz", + text_states=text_arr, + image_states=image_arr, + text_sim=text_sim_by_layer, + image_sim=image_sim_by_layer, + cross_sim=cross_sim_by_layer, + ) + (out_dir / "summary.json").write_text(json.dumps(summary, indent=1)) + print(json.dumps({k: v for k, v in summary.items() if k not in ("per_layer", "records")}, indent=1)) + print(f"results -> {out_dir}") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_convergence_3d.py b/packages/snapcompact/research/snapcompact_convergence_3d.py new file mode 100644 index 000000000..ec4f8ed28 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_convergence_3d.py @@ -0,0 +1,187 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["matplotlib", "numpy", "pillow"] +# /// +"""3D convergence strands: text and image trajectories fusing through depth.""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path + +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +BG = (5, 7, 10) +PANEL = (12, 17, 23) +INK = (241, 239, 224) +MUTED = (143, 154, 160) +AMBER = (255, 196, 68) +CYAN = (75, 220, 255) +ORANGE = (255, 112, 72) + + +def ui_font(size: int, bold: bool = False) -> ImageFont.ImageFont: + for path in [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def question_hue(i: int, n: int) -> tuple[float, float, float]: + h = i / n + r = 0.5 + 0.5 * math.cos(2 * math.pi * (h + 0.00)) + g = 0.5 + 0.5 * math.cos(2 * math.pi * (h + 0.33)) + b = 0.5 + 0.5 * math.cos(2 * math.pi * (h + 0.67)) + return (0.28 + 0.72 * r, 0.28 + 0.72 * g, 0.28 + 0.72 * b) + + +def center(arr: np.ndarray) -> np.ndarray: + return arr - arr.mean(axis=0, keepdims=True) + + +def smooth_path(path: np.ndarray, passes: int = 2) -> np.ndarray: + out = path.copy() + for _ in range(passes): + mid = (out[:-2] + out[1:-1] * 2 + out[2:]) / 4 + out[1:-1] = mid + return out + + +def render_strands(text_arr: np.ndarray, image_arr: np.ndarray, best_layer: int) -> Image.Image: + n_q, n_layers, _ = text_arr.shape + ref = np.concatenate([center(text_arr[:, best_layer, :]), center(image_arr[:, best_layer, :])], axis=0) + _, _, vt = np.linalg.svd(ref, full_matrices=False) + basis = vt[:2].T + + # Per-layer projections, per-layer scale normalization so depth shows shape, + # not raw norm growth across layers. + t_proj = np.zeros((n_q, n_layers, 2), dtype=np.float64) + i_proj = np.zeros((n_q, n_layers, 2), dtype=np.float64) + for layer in range(n_layers): + t = center(text_arr[:, layer, :]) @ basis + i = center(image_arr[:, layer, :]) @ basis + scale = max(1e-6, float(np.abs(np.concatenate([t, i], axis=0)).max())) + t_proj[:, layer] = t / scale + i_proj[:, layer] = i / scale + + fig = plt.figure(figsize=(15.2, 9.4), dpi=170) + fig.patch.set_facecolor("#05070a") + ax = fig.add_subplot(111, projection="3d") + ax.set_facecolor((0.02, 0.025, 0.035, 1)) + for axis in (ax.xaxis, ax.yaxis, ax.zaxis): + axis.pane.set_facecolor((0.02, 0.025, 0.035, 0.0)) + axis._axinfo["grid"]["color"] = (0.32, 0.42, 0.48, 0.16) + ax.tick_params(colors="#8f9aa0", labelsize=8) + + layers_axis = np.arange(n_layers) + for qi in range(n_q): + color = question_hue(qi, n_q) + tp = smooth_path(np.column_stack([layers_axis, t_proj[qi, :, 0], t_proj[qi, :, 1]])) + ip = smooth_path(np.column_stack([layers_axis, i_proj[qi, :, 0], i_proj[qi, :, 1]])) + ax.plot(tp[:, 0], tp[:, 1], tp[:, 2], color=color, linewidth=2.6, alpha=0.95) + ax.plot(ip[:, 0], ip[:, 1], ip[:, 2], color=color, linewidth=2.6, alpha=0.55, linestyle=(0, (4, 2))) + # tie-lines every few layers showing the closing gap + for layer in range(1, n_layers, 4): + ax.plot( + [layer, layer], + [t_proj[qi, layer, 0], i_proj[qi, layer, 0]], + [t_proj[qi, layer, 1], i_proj[qi, layer, 1]], + color=color, + linewidth=0.9, + alpha=0.38, + ) + ax.scatter([0], [t_proj[qi, 0, 0]], [t_proj[qi, 0, 1]], color=color, s=26, marker="o", depthshade=False) + ax.scatter([0], [i_proj[qi, 0, 0]], [i_proj[qi, 0, 1]], color=color, s=30, marker="D", depthshade=False) + ax.scatter([best_layer], [t_proj[qi, best_layer, 0]], [t_proj[qi, best_layer, 1]], color=color, s=46, marker="o", edgecolors="white", linewidths=0.6, depthshade=False) + + # Peak-layer plane. + yy, zz = np.meshgrid(np.linspace(-1.05, 1.05, 2), np.linspace(-1.05, 1.05, 2)) + ax.plot_surface(np.full_like(yy, best_layer), yy, zz, color=(1.0, 0.77, 0.27, 0.10), shade=False) + + ax.set_xlim(0, n_layers - 1) + ax.set_ylim(-1.1, 1.1) + ax.set_zlim(-1.1, 1.1) + ax.set_xlabel("decoder layer →", color="#8f9aa0", labelpad=12) + ax.set_ylabel("content PC1", color="#8f9aa0", labelpad=10) + ax.set_zlabel("content PC2", color="#8f9aa0", labelpad=8) + ax.view_init(elev=18, azim=-66) + ax.set_box_aspect((2.9, 1.0, 0.9)) + tmp = HERE / "results" / ".convergence-3d-panel.png" + fig.subplots_adjust(left=0, right=1, top=1, bottom=0) + fig.savefig(tmp, facecolor=fig.get_facecolor(), transparent=False, bbox_inches="tight", pad_inches=0.05) + plt.close(fig) + img = Image.open(tmp).convert("RGB") + tmp.unlink(missing_ok=True) + return img + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--result-dir", default=str(HERE / "results" / "qwen-carrier-convergence-n12")) + ap.add_argument("--out", default=str(HERE / "results" / "qwen-carrier-convergence-n12" / "convergence-strands-3d.png")) + args = ap.parse_args() + result_dir = Path(args.result_dir) + summary = json.loads((result_dir / "summary.json").read_text()) + data = np.load(result_dir / "carrier_convergence.npz") + best_layer = summary["best_layer"] + best = summary["best"] + panel = render_strands(data["text_states"], data["image_states"], best_layer) + + w, h = 2200, 1300 + canvas = Image.new("RGB", (w, h), BG) + draw = ImageDraw.Draw(canvas) + for y in range(0, h, 16): + draw.line((0, y, w, y), fill=(7, 10 + y % 9, 15 + y % 11)) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-260, -220, 900, 700), fill=(75, 220, 255, 27)) + gd.ellipse((1240, 160, 2460, 1360), fill=(255, 112, 72, 25)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(84))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + draw.text((64, 42), "QWEN CARRIER CONVERGENCE — 3D STRANDS", fill=AMBER, font=ui_font(24, True)) + draw.text((64, 84), "Twelve thoughts, two doors, one room", fill=INK, font=ui_font(64, True)) + draw.text( + (66, 164), + "Each color is one question travelling through the decoder. Solid strand entered as text; dashed strand entered as pixels. Strand pairs braid together by depth.", + fill=MUTED, + font=ui_font(23), + ) + + draw.rounded_rectangle((64, 234, 2136, 1146), radius=30, fill=PANEL, outline=(35, 49, 59), width=1) + panel = panel.resize((1980, 832), Image.Resampling.LANCZOS) + canvas.paste(panel, (104, 286)) + draw.text((96, 252), f"PCA frame fixed at peak layer {best_layer}; per-layer scale normalized", fill=MUTED, font=ui_font(17)) + + stats = [ + ("matched cosine", f"{best['matched_cosine']:.2f}"), + ("mismatched", f"{best['mismatched_cosine']:.2f}"), + ("RSA geometry", f"{best['rsa_pearson']:.2f}"), + ("pair retrieval", f"{best['match_rank_accuracy'] * 100:.0f}%"), + ] + sx = 64 + for title, value in stats: + draw.rounded_rectangle((sx, 1170, sx + 320, 1262), radius=18, fill=PANEL, outline=(35, 49, 59), width=1) + draw.text((sx + 22, 1184), title, fill=MUTED, font=ui_font(16)) + draw.text((sx + 22, 1208), value, fill=INK, font=ui_font(34, True)) + sx += 344 + draw.text((sx + 20, 1196), "solid = text carrier dashed = image carrier thin rungs = pair gap", fill=MUTED, font=ui_font(18)) + + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + canvas.save(out) + print(out) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_convergence_extras.py b/packages/snapcompact/research/snapcompact_convergence_extras.py new file mode 100644 index 000000000..7ab101d6b --- /dev/null +++ b/packages/snapcompact/research/snapcompact_convergence_extras.py @@ -0,0 +1,207 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy"] +# /// +"""Extra convergence graphics: PCA funnel snapshots and an animated diagonal GIF.""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path +from typing import Any + +import numpy as np +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +PALETTE = { + "bg": (5, 7, 10), + "panel": (12, 17, 23), + "panel2": (8, 12, 17), + "ink": (241, 239, 224), + "muted": (143, 154, 160), + "amber": (255, 196, 68), + "cyan": (75, 220, 255), + "orange": (255, 112, 72), + "grid": (38, 49, 58), +} + + +def ui_font(size: int, bold: bool = False) -> ImageFont.ImageFont: + for path in [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def question_hue(i: int, n: int) -> tuple[int, int, int]: + """Distinct, bright hue per question.""" + h = i / n + r = 0.5 + 0.5 * math.cos(2 * math.pi * (h + 0.00)) + g = 0.5 + 0.5 * math.cos(2 * math.pi * (h + 0.33)) + b = 0.5 + 0.5 * math.cos(2 * math.pi * (h + 0.67)) + return (round(70 + 185 * r), round(70 + 185 * g), round(70 + 185 * b)) + + +def center(arr: np.ndarray) -> np.ndarray: + return arr - arr.mean(axis=0, keepdims=True) + + +def diverging_color(t: float) -> tuple[int, int, int]: + t = max(-1.0, min(1.0, t)) + if t < 0: + u = -t + return (round(8 + 12 * u), round(20 + 90 * u), round(34 + 190 * u)) + return (round(8 + 247 * t), round(20 + 130 * t), round(34 + 20 * t)) + + +def background(w: int, h: int) -> Image.Image: + canvas = Image.new("RGB", (w, h), PALETTE["bg"]) + draw = ImageDraw.Draw(canvas) + for y in range(0, h, 16): + draw.line((0, y, w, y), fill=(7, 10 + y % 9, 15 + y % 11)) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-240, -200, 880, 680), fill=(75, 220, 255, 25)) + gd.ellipse((w - 1000, h - 760, w + 240, h + 220), fill=(255, 112, 72, 25)) + return Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(84))).convert("RGB") + + +def render_funnel(out_path: Path, text_arr: np.ndarray, image_arr: np.ndarray, layers_meta: list[dict[str, Any]], best_layer: int, records: list[dict[str, Any]]) -> None: + n_q, n_layers, _ = text_arr.shape + snapshots = [1, max(2, best_layer // 2), best_layer] + # Shared PCA frame from the peak layer keeps the panels comparable. + ref = np.concatenate([center(text_arr[:, best_layer, :]), center(image_arr[:, best_layer, :])], axis=0) + _, _, vt = np.linalg.svd(ref, full_matrices=False) + basis = vt[:2].T # [D, 2] + + w, h = 2200, 1240 + canvas = background(w, h) + draw = ImageDraw.Draw(canvas) + draw.text((64, 42), "QWEN CARRIER CONVERGENCE — TRAJECTORY VIEW", fill=PALETTE["amber"], font=ui_font(24, True)) + draw.text((64, 84), "Watch the two carriers fuse", fill=PALETTE["ink"], font=ui_font(64, True)) + draw.text((66, 164), "Each color is one question; ● came in as text, ◆ came in as pixels. Same 2D projection at every depth. The tie-lines shrink as carriers converge.", fill=PALETTE["muted"], font=ui_font(23)) + + panel_w = 660 + titles = ["early (layer {})", "middle (layer {})", "peak (layer {})"] + for pi, (layer, title) in enumerate(zip(snapshots, titles)): + x0 = 64 + pi * (panel_w + 44) + box = (x0, 232, x0 + panel_w, 952) + draw.rounded_rectangle(box, radius=24, fill=PALETTE["panel"], outline=(35, 49, 59), width=1) + draw.text((x0 + 26, 252), title.format(layer), fill=PALETTE["ink"], font=ui_font(27, True)) + t_proj = center(text_arr[:, layer, :]) @ basis + i_proj = center(image_arr[:, layer, :]) @ basis + both = np.concatenate([t_proj, i_proj], axis=0) + lim = float(np.abs(both).max()) * 1.15 or 1.0 + gx0, gy0, gx1, gy1 = x0 + 36, 306, x0 + panel_w - 36, 912 + def to_px(p: np.ndarray) -> tuple[int, int]: + return ( + round(gx0 + (p[0] + lim) / (2 * lim) * (gx1 - gx0)), + round(gy0 + (1 - (p[1] + lim) / (2 * lim)) * (gy1 - gy0)), + ) + draw.line((gx0, (gy0 + gy1) // 2, gx1, (gy0 + gy1) // 2), fill=PALETTE["grid"], width=1) + draw.line(((gx0 + gx1) // 2, gy0, (gx0 + gx1) // 2, gy1), fill=PALETTE["grid"], width=1) + pair_dist = 0.0 + for qi in range(n_q): + color = question_hue(qi, n_q) + tp = to_px(t_proj[qi]) + ip = to_px(i_proj[qi]) + draw.line((tp, ip), fill=(*color, 0)[:3], width=3) + r = 11 + draw.ellipse((tp[0] - r, tp[1] - r, tp[0] + r, tp[1] + r), fill=color, outline=(8, 10, 12), width=2) + d = ImageDraw.Draw(canvas) + d.polygon([(ip[0], ip[1] - r - 2), (ip[0] + r + 2, ip[1]), (ip[0], ip[1] + r + 2), (ip[0] - r - 2, ip[1])], fill=color, outline=(8, 10, 12)) + pair_dist += float(np.linalg.norm(t_proj[qi] - i_proj[qi])) + pair_dist /= n_q + norm_dist = pair_dist / (2 * lim) + meta = layers_meta[layer] + draw.text((x0 + 26, 916), f"mean pair gap: {norm_dist * 100:.0f}% of frame · matched cos {meta['matched_cosine']:.2f}", fill=PALETTE["muted"], font=ui_font(17)) + + # Pair-distance by layer strip. + strip = (64, 996, 2136, 1190) + draw.rounded_rectangle(strip, radius=24, fill=PALETTE["panel"], outline=(35, 49, 59), width=1) + draw.text((96, 1014), "matched-pair separation by layer (lower = carriers agree)", fill=PALETTE["ink"], font=ui_font(22, True)) + gx0, gy0, gx1, gy1 = 110, 1062, 2100, 1162 + gaps = [] + for layer in range(n_layers): + t_proj = center(text_arr[:, layer, :]) + i_proj = center(image_arr[:, layer, :]) + t_n = t_proj / np.maximum(np.linalg.norm(t_proj, axis=1, keepdims=True), 1e-6) + i_n = i_proj / np.maximum(np.linalg.norm(i_proj, axis=1, keepdims=True), 1e-6) + gaps.append(1.0 - float((t_n * i_n).sum(axis=1).mean())) + hi = max(gaps) + bw = (gx1 - gx0) / n_layers + for layer, gap in enumerate(gaps): + xa = gx0 + layer * bw + 3 + xb = gx0 + (layer + 1) * bw - 3 + bh = (gy1 - gy0) * gap / hi + color = PALETTE["orange"] if layer == best_layer else (62, 86, 102) + draw.rounded_rectangle((round(xa), round(gy1 - bh), round(xb), gy1), radius=5, fill=color) + draw.text((gx0, gy1 + 6), "layer 0", fill=PALETTE["muted"], font=ui_font(13)) + draw.text((gx1 - 70, gy1 + 6), f"layer {n_layers - 1}", fill=PALETTE["muted"], font=ui_font(13)) + + out_path.parent.mkdir(parents=True, exist_ok=True) + canvas.save(out_path) + + +def render_gif(out_path: Path, cross_sim: np.ndarray, layers_meta: list[dict[str, Any]]) -> None: + n_layers, n_q, _ = cross_sim.shape + cell = 46 + pad = 36 + header = 132 + w = n_q * cell + pad * 2 + h = n_q * cell + header + pad + 64 + frames: list[Image.Image] = [] + for layer in range(n_layers): + frame = Image.new("RGB", (w, h), PALETTE["bg"]) + draw = ImageDraw.Draw(frame) + for y in range(0, h, 14): + draw.line((0, y, w, y), fill=(7, 10 + y % 9, 15 + y % 11)) + draw.text((pad, 22), "cross-carrier matching", fill=PALETTE["ink"], font=ui_font(30, True)) + draw.text((pad, 62), "text question i × image question j", fill=PALETTE["muted"], font=ui_font(17)) + meta = layers_meta[layer] + draw.text((pad, 92), f"layer {layer:02d} matched {meta['matched_cosine']:+.2f} others {meta['mismatched_cosine']:+.2f}", fill=PALETTE["amber"], font=ui_font(19, True)) + for r in range(n_q): + for c in range(n_q): + xa = pad + c * cell + ya = header + r * cell + draw.rounded_rectangle((xa, ya, xa + cell - 4, ya + cell - 4), radius=7, fill=diverging_color(float(cross_sim[layer, r, c]))) + # progress bar + bar_y = header + n_q * cell + 18 + draw.rounded_rectangle((pad, bar_y, w - pad, bar_y + 10), radius=5, fill=(30, 40, 48)) + draw.rounded_rectangle((pad, bar_y, pad + (w - 2 * pad) * (layer + 1) // n_layers, bar_y + 10), radius=5, fill=PALETTE["cyan"]) + frames.append(frame) + durations = [240] * n_layers + durations[-1] = 2200 + out_path.parent.mkdir(parents=True, exist_ok=True) + frames[0].save(out_path, save_all=True, append_images=frames[1:], duration=durations, loop=0) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--result-dir", default=str(HERE / "results" / "qwen-carrier-convergence-n12")) + args = ap.parse_args() + result_dir = Path(args.result_dir) + summary = json.loads((result_dir / "summary.json").read_text()) + data = np.load(result_dir / "carrier_convergence.npz") + text_arr = data["text_states"] + image_arr = data["image_states"] + cross_sim = data["cross_sim"] + layers_meta = summary["per_layer"] + best_layer = summary["best_layer"] + + funnel_path = result_dir / "convergence-funnel.png" + gif_path = result_dir / "diagonal-emerges.gif" + render_funnel(funnel_path, text_arr, image_arr, layers_meta, best_layer, summary["records"]) + render_gif(gif_path, cross_sim, layers_meta) + print(funnel_path) + print(gif_path) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_convergence_viz.py b/packages/snapcompact/research/snapcompact_convergence_viz.py new file mode 100644 index 000000000..b28b33088 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_convergence_viz.py @@ -0,0 +1,219 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy"] +# /// +"""Render the carrier-convergence geometry figure: same content, same nothings.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + +import numpy as np +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +PALETTE = { + "bg": (5, 7, 10), + "panel": (12, 17, 23), + "panel2": (8, 12, 17), + "ink": (241, 239, 224), + "muted": (143, 154, 160), + "cyan": (75, 220, 255), + "orange": (255, 112, 72), + "green": (148, 255, 117), + "amber": (255, 196, 68), + "red": (255, 76, 62), + "grid": (38, 49, 58), +} + + +def ui_font(size: int, bold: bool = False) -> ImageFont.ImageFont: + for path in [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def mono_font(size: int) -> ImageFont.ImageFont: + for path in ["/System/Library/Fonts/Monaco.ttf", "/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf"]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def diverging_color(t: float) -> tuple[int, int, int]: + """-1..1 → blue-black-orange diverging scale.""" + t = max(-1.0, min(1.0, t)) + if t < 0: + u = -t + return (round(8 + 12 * u), round(20 + 90 * u), round(34 + 190 * u)) + u = t + return (round(8 + 247 * u), round(20 + 130 * u), round(34 + 20 * u)) + + +def draw_matrix(draw: ImageDraw.ImageDraw, mat: np.ndarray, box: tuple[int, int, int, int], title: str, subtitle: str, color: tuple[int, int, int], highlight_diag: bool = False) -> None: + x0, y0, x1, y1 = box + draw.rounded_rectangle(box, radius=20, fill=PALETTE["panel2"], outline=(34, 48, 58), width=1) + draw.text((x0 + 20, y0 + 16), title, fill=color, font=ui_font(24, True)) + draw.text((x0 + 20, y0 + 48), subtitle, fill=PALETTE["muted"], font=ui_font(15)) + n = mat.shape[0] + gx0, gy0 = x0 + 34, y0 + 84 + side = min(x1 - 34 - gx0, y1 - 30 - gy0) + cell = side / n + for r in range(n): + for c in range(n): + xa = round(gx0 + c * cell) + xb = round(gx0 + (c + 1) * cell) - 2 + ya = round(gy0 + r * cell) + yb = round(gy0 + (r + 1) * cell) - 2 + draw.rounded_rectangle((xa, ya, xb, yb), radius=4, fill=diverging_color(float(mat[r, c]))) + if highlight_diag: + for r in range(n): + xa = round(gx0 + r * cell) + ya = round(gy0 + r * cell) + draw.rounded_rectangle((xa - 1, ya - 1, round(xa + cell) - 1, round(ya + cell) - 1), radius=5, outline=PALETTE["amber"], width=2) + draw.text((gx0, round(gy0 + side) + 6), "questions →", fill=PALETTE["muted"], font=ui_font(13)) + + +def draw_curves(draw: ImageDraw.ImageDraw, box: tuple[int, int, int, int], layers: list[dict[str, Any]]) -> None: + x0, y0, x1, y1 = box + draw.rounded_rectangle(box, radius=20, fill=PALETTE["panel2"], outline=(34, 48, 58), width=1) + draw.text((x0 + 22, y0 + 16), "convergence by depth", fill=PALETTE["ink"], font=ui_font(24, True)) + draw.text((x0 + 22, y0 + 48), "carrier-centered cosine: same question across carriers vs different questions", fill=PALETTE["muted"], font=ui_font(15)) + gx0, gy0, gx1, gy1 = x0 + 52, y0 + 92, x1 - 26, y1 - 56 + lo, hi = -0.15, 1.0 + for i in range(5): + y = gy0 + (gy1 - gy0) * i / 4 + draw.line((gx0, y, gx1, y), fill=PALETTE["grid"], width=1) + value = hi - (hi - lo) * i / 4 + draw.text((x0 + 12, y - 8), f"{value:.1f}", fill=PALETTE["muted"], font=ui_font(12)) + series = [ + ("matched_cosine", PALETTE["amber"], 6), + ("mismatched_cosine", PALETTE["muted"], 4), + ("rsa_pearson", PALETTE["cyan"], 4), + ] + n = len(layers) + for key, color, width in series: + pts = [] + for i, row in enumerate(layers): + value = float(row[key]) + if value != value: # NaN guard (layer-0 RSA is undefined) + continue + value = max(lo, min(hi, value)) + x = gx0 + (gx1 - gx0) * i / max(1, n - 1) + y = gy1 - (gy1 - gy0) * (value - lo) / (hi - lo) + pts.append((round(x), round(y))) + if len(pts) < 2: + continue + draw.line(pts, fill=color, width=width, joint="curve") + draw.text((gx0, gy1 + 14), "layer 0", fill=PALETTE["muted"], font=ui_font(13)) + draw.text((gx1 - 64, gy1 + 14), f"layer {n - 1}", fill=PALETTE["muted"], font=ui_font(13)) + legend = [("same question, text↔image", PALETTE["amber"]), ("different questions", PALETTE["muted"]), ("RSA geometry corr", PALETTE["cyan"])] + lx = gx0 + for label, color in legend: + draw.rounded_rectangle((lx, y0 + 70, lx + 18, y0 + 78), radius=4, fill=color) + draw.text((lx + 24, y0 + 62), label, fill=PALETTE["muted"], font=ui_font(13)) + lx += 232 + + +def draw_answers(draw: ImageDraw.ImageDraw, box: tuple[int, int, int, int], records: list[dict[str, Any]]) -> None: + x0, y0, x1, y1 = box + draw.rounded_rectangle(box, radius=20, fill=PALETTE["panel2"], outline=(34, 48, 58), width=1) + draw.text((x0 + 22, y0 + 16), "behavioral check: both carriers answer alike", fill=PALETTE["ink"], font=ui_font(24, True)) + draw.text((x0 + 240, y0 + 56), "text carrier", fill=PALETTE["cyan"], font=ui_font(15, True)) + draw.text((x0 + 470, y0 + 56), "image carrier", fill=PALETTE["orange"], font=ui_font(15, True)) + y = y0 + 84 + row_h = (y1 - y0 - 96) // len(records) + fnt = mono_font(15) + for r in records: + gold = r["gold"][:22] + draw.text((x0 + 22, y), gold, fill=PALETTE["muted"], font=fnt) + draw.text((x0 + 240, y), r["text_answer"][:22], fill=PALETTE["ink"], font=fnt) + draw.text((x0 + 470, y), r["image_answer"][:22], fill=PALETTE["ink"], font=fnt) + mark = "=" if r["agree"] else "≠" + draw.text((x1 - 44, y), mark, fill=PALETTE["green"] if r["agree"] else PALETTE["red"], font=ui_font(17, True)) + y += row_h + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--result-dir", default=str(HERE / "results" / "qwen-carrier-convergence-n12")) + ap.add_argument("--out", default=str(HERE / "results" / "qwen-carrier-convergence-n12" / "carrier-convergence.png")) + args = ap.parse_args() + + result_dir = Path(args.result_dir) + summary = json.loads((result_dir / "summary.json").read_text()) + data = np.load(result_dir / "carrier_convergence.npz") + layers = summary["per_layer"] + best_layer = summary["best_layer"] + text_sim = data["text_sim"][best_layer] + image_sim = data["image_sim"][best_layer] + cross_sim = data["cross_sim"][best_layer] + records = summary["records"] + + w, h = 2200, 1320 + canvas = Image.new("RGB", (w, h), PALETTE["bg"]) + draw = ImageDraw.Draw(canvas) + for y in range(0, h, 16): + draw.line((0, y, w, y), fill=(7, 10 + y % 9, 15 + y % 11)) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-240, -200, 900, 700), fill=(75, 220, 255, 26)) + gd.ellipse((1300, 180, 2480, 1380), fill=(255, 112, 72, 24)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(86))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + best = summary["best"] + draw.text((64, 42), "QWEN CARRIER CONVERGENCE", fill=PALETTE["amber"], font=ui_font(24, True)) + draw.text((64, 84), "Two carriers, one thought", fill=PALETTE["ink"], font=ui_font(66, True)) + draw.text( + (66, 166), + "Hidden state at the answer position, carrier means removed. Same question through text or bitmap lands in the same place; different questions do not.", + fill=PALETTE["muted"], + font=ui_font(23), + ) + + stats = [ + ("matched pairs", f"{best['matched_cosine']:.2f}", "same Q, text ↔ image"), + ("mismatched pairs", f"{best['mismatched_cosine']:.2f}", "different questions"), + ("RSA geometry corr", f"{best['rsa_pearson']:.2f}", f"layer {best['layer']}"), + ("pair retrieval", f"{best['match_rank_accuracy'] * 100:.0f}%", "nearest cross-carrier match"), + ("answer agreement", f"{summary['answer_agreement'] * 100:.0f}%", "text vs image generations"), + ] + sx = 64 + for title, value, caption in stats: + draw.rounded_rectangle((sx, 222, sx + 396, 332), radius=20, fill=PALETTE["panel"], outline=(35, 49, 59), width=1) + draw.text((sx + 22, 240), title, fill=PALETTE["muted"], font=ui_font(16)) + draw.text((sx + 22, 264), value, fill=PALETTE["ink"], font=ui_font(40, True)) + draw.text((sx + 226, 290), caption, fill=PALETTE["muted"], font=ui_font(13)) + sx += 420 + + n = text_sim.shape[0] + draw_matrix(draw, text_sim, (64, 376, 600, 952), "text-carrier geometry", f"{n}×{n} question similarity, layer {best_layer}", PALETTE["cyan"]) + draw_matrix(draw, image_sim, (628, 376, 1164, 952), "image-carrier geometry", "same questions through the bitmap — same shape", PALETTE["orange"]) + draw_matrix(draw, cross_sim, (1192, 376, 1728, 952), "cross-carrier matching", "text question i × image question j — bright diagonal", PALETTE["green"], highlight_diag=True) + + draw_curves(draw, (64, 996, 1164, 1264), layers) + draw_answers(draw, (1192, 996, 2136, 1264), records) + + # Color scale. + for i in range(240): + t = 1 - i / 239 * 2 + draw.rectangle((1816, 420 + i * 2, 1836, 422 + i * 2), fill=diverging_color(t)) + draw.text((1848, 412), "+1 similar", fill=PALETTE["muted"], font=ui_font(14)) + draw.text((1848, 884), "−1 opposite", fill=PALETTE["muted"], font=ui_font(14)) + + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + canvas.save(out) + print(out) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_lockon_anatomy_viz.py b/packages/snapcompact/research/snapcompact_lockon_anatomy_viz.py new file mode 100644 index 000000000..357c63173 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_lockon_anatomy_viz.py @@ -0,0 +1,288 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy"] +# /// +"""The lock-on instrument: how we decide WHERE the answer materializes. + +Renders the measurement methodology as a depth-gauge diagram: the answer patch +descends the decoder shaft; at every layer a logit-lens probe (final norm + LM +head) reads the vocabulary distribution; lock-on is the first layer whose top-1 +token is a BPE piece of the answer. All readouts are real sweep data. +""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path +from typing import Any + +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +P = { + "bg": (5, 7, 10), + "panel": (12, 17, 23), + "panel2": (8, 12, 17), + "ink": (241, 239, 224), + "muted": (143, 154, 160), + "faint": (90, 101, 108), + "cyan": (75, 220, 255), + "orange": (255, 112, 72), + "green": (148, 255, 117), + "amber": (255, 196, 68), + "purple": (188, 112, 255), + "grid": (38, 49, 58), +} +COND_COLORS = { + "base-8x13": (143, 154, 160), + "repeat2-color": (255, 196, 68), + "align-7x14": (148, 255, 117), + "align-14x28": (75, 220, 255), + "align-28x28": (255, 112, 72), + "repeat2-align-14x28": (188, 112, 255), +} + + +def ui_font(size: int, bold: bool = False) -> ImageFont.ImageFont: + for path in [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def mono_font(size: int) -> ImageFont.ImageFont: + for path in ["/System/Library/Fonts/Monaco.ttf", "/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf"]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def label_font(label: str, size: int) -> ImageFont.ImageFont: + """Monaco for ASCII; Arial Unicode for anything it cannot shape (CJK).""" + if all(ord(ch) < 0x2000 for ch in label): + return mono_font(size) + unicode_path = "/System/Library/Fonts/Supplemental/Arial Unicode.ttf" + if Path(unicode_path).exists(): + return ImageFont.truetype(unicode_path, size) + return mono_font(size) + + +def crosshair(draw: ImageDraw.ImageDraw, cx: int, cy: int, r: int, color: tuple[int, int, int], width: int = 4) -> None: + draw.ellipse((cx - r, cy - r, cx + r, cy + r), outline=color, width=width) + draw.ellipse((cx - r // 2, cy - r // 2, cx + r // 2, cy + r // 2), outline=color, width=2) + for dx, dy in ((-1, 0), (1, 0), (0, -1), (0, 1)): + draw.line((cx + dx * (r - 6), cy + dy * (r - 6), cx + dx * (r + 14), cy + dy * (r + 14)), fill=color, width=width) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--result-dir", default=str(HERE / "results" / "qwen-materialize-sweep-q3")) + ap.add_argument("--condition", default="base-8x13") + ap.add_argument("--out", default=str(HERE / "results" / "qwen-materialize-sweep-q3" / "lockon-anatomy.png")) + args = ap.parse_args() + result_dir = Path(args.result_dir) + summary = json.loads((result_dir / "summary.json").read_text()) + conditions = {c["name"]: c for c in summary["conditions"]} + cond = conditions[args.condition] + q = summary["question"] + answer_strs = summary["answer_token_strs"] + answer_set = {s.strip() for s in answer_strs} + layers = cond["layers"] + n_layers = len(layers) + lock_on = cond["lock_on_layer"] + + w, h = 2200, 1420 + canvas = Image.new("RGB", (w, h), P["bg"]) + draw = ImageDraw.Draw(canvas) + for y in range(0, h, 16): + draw.line((0, y, w, y), fill=(7, 10 + y % 9, 15 + y % 11)) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((520, 620, 1280, 1180), fill=(255, 196, 68, 36)) + gd.ellipse((-260, -240, 760, 560), fill=(75, 220, 255, 26)) + gd.ellipse((1500, -100, 2480, 700), fill=(255, 112, 72, 20)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(90))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + draw.text((64, 40), "THE LOCK-ON INSTRUMENT", fill=P["amber"], font=ui_font(24, True)) + draw.text((64, 80), "How we decide where the answer materializes", fill=P["ink"], font=ui_font(58, True)) + draw.text( + (66, 154), + "At every layer, a logit-lens probe taps the answer patch's residual stream: final RMSNorm → LM head → softmax over 152k vocabulary entries.", + fill=P["muted"], + font=ui_font(22), + ) + draw.text( + (66, 186), + "LOCK-ON = the first layer whose #1 vocabulary entry is a BPE piece of the answer. Past this depth the fact is settled — remaining layers are free for reasoning.", + fill=P["amber"], + font=ui_font(22, True), + ) + + # ---- Probe pipeline card (top left). + pipe = (64, 248, 700, 420) + draw.rounded_rectangle(pipe, radius=22, fill=P["panel"], outline=(35, 49, 59), width=1) + draw.text((92, 268), "the probe, applied at every layer ℓ", fill=P["ink"], font=ui_font(22, True)) + stages = ["h(patch)", "RMSNorm", "LM head", "softmax", "top-1?"] + sx = 92 + for si, stage in enumerate(stages): + color = P["amber"] if si == len(stages) - 1 else P["cyan"] + tw = int(draw.textlength(stage, font=mono_font(16))) + 24 + draw.rounded_rectangle((sx, 318, sx + tw, 356), radius=10, fill=P["panel2"], outline=color, width=2) + draw.text((sx + 12, 327), stage, fill=color, font=mono_font(16)) + if si < len(stages) - 1: + draw.text((sx + tw + 4, 327), "→", fill=P["faint"], font=ui_font(18)) + sx += tw + 28 + draw.text((92, 376), f"vocabulary = 152k entries · answer BPEs = {answer_strs}", fill=P["muted"], font=mono_font(14)) + + # ---- The patch under test (left). + patch_card = (64, 460, 380, 760) + draw.rounded_rectangle(patch_card, radius=22, fill=P["panel"], outline=(35, 49, 59), width=1) + draw.text((92, 480), "specimen", fill=P["orange"], font=ui_font(21, True)) + carrier = Image.open(result_dir / "images" / f"{args.condition}.png").convert("RGB") + rw = 1568 + grid = 56 + px = 28 + lock_entry = layers[lock_on] + tok_idx = lock_entry["best_token_index"] + r0, c0 = tok_idx // grid, tok_idx % grid + cell = carrier.resize((rw, rw), Image.Resampling.LANCZOS).crop((c0 * px, r0 * px, (c0 + 1) * px, (r0 + 1) * px)) + big = cell.resize((196, 196), Image.Resampling.NEAREST) + draw.rounded_rectangle((118, 516, 326, 724), radius=12, fill=(244, 242, 230), outline=P["orange"], width=4) + canvas.paste(big, (124, 522)) + draw.text((118, 730), f"visual token #{tok_idx} · 28×28 px", fill=P["muted"], font=mono_font(13)) + + # ---- Depth shaft. + shaft_x = 470 + shaft_top, shaft_bot = 470, 1340 + draw.rounded_rectangle((shaft_x - 7, shaft_top, shaft_x + 7, shaft_bot), radius=7, fill=(20, 28, 35), outline=(40, 54, 64), width=1) + + def layer_y(layer: int) -> int: + return round(shaft_top + (shaft_bot - shaft_top) * layer / (n_layers - 1)) + + # p(answer) trajectory along the shaft. + traj = [(shaft_x + 14 + 230 * min(1.0, e["best_answer_p"]), layer_y(e["layer"])) for e in layers] + for i in range(len(traj) - 1): + draw.line((traj[i], traj[i + 1]), fill=(120, 96, 40), width=3) + draw.text((shaft_x + 30, shaft_bot + 10), "p(answer BPE) →", fill=(150, 124, 60), font=ui_font(14)) + + for layer in range(n_layers): + y = layer_y(layer) + major = layer % 4 == 0 or layer == n_layers - 1 + draw.line((shaft_x - (16 if major else 10), y, shaft_x + (16 if major else 10), y), fill=P["faint"] if major else (52, 64, 73), width=2) + if major: + draw.text((shaft_x - 58, y - 9), f"L{layer:02d}", fill=P["muted"], font=mono_font(13)) + # Patch entering the shaft. + draw.line((326, 620, shaft_x - 18, shaft_top + 6), fill=P["orange"], width=3) + draw.polygon([(shaft_x - 14, shaft_top + 2), (shaft_x - 30, shaft_top - 4), (shaft_x - 26, shaft_top + 16)], fill=P["orange"]) + + # ---- Readout cards at sampled depths (real top-5). + samples = [0, 10, 18, lock_on, n_layers - 1] + card_x = 790 + card_w = 620 + card_h = 128 + gap = 14 + desired = [layer_y(layer) - card_h // 2 for layer in samples] + card_ys = [0] * len(samples) + # Bottom-up pass: clamp the last card into the canvas, then keep every + # earlier card fully above its successor; final top clamp at 440. + card_ys[-1] = min(desired[-1], h - card_h - 70) + for i in range(len(samples) - 2, -1, -1): + card_ys[i] = min(desired[i], card_ys[i + 1] - card_h - gap) + shift = max(0, 440 - card_ys[0]) + card_ys = [cy + shift for cy in card_ys] + for layer, cy in zip(samples, card_ys): + entry = layers[layer] + is_lock = layer == lock_on + accent = P["amber"] if is_lock else P["cyan"] if entry["best_answer_p"] > 0.01 else P["faint"] + # Connector. + ly = layer_y(layer) + draw.line((shaft_x + 16, ly, card_x - 18, cy + card_h // 2), fill=accent, width=3 if is_lock else 2) + draw.ellipse((shaft_x + 12, ly - 5, shaft_x + 22, ly + 5), fill=accent) + draw.rounded_rectangle((card_x, cy, card_x + card_w, cy + card_h), radius=16, fill=P["panel2"], outline=accent, width=3 if is_lock else 1) + title = f"L{layer:02d} readout" + (" LOCK-ON" if is_lock else "") + draw.text((card_x + 20, cy + 10), title, fill=accent, font=ui_font(19, True)) + if is_lock: + tx = card_x + 20 + draw.textlength(f"L{layer:02d} readout ", font=ui_font(19, True)) + draw.ellipse((tx - 8, cy + 14, tx + 4, cy + 26), outline=accent, width=3) + bx = card_x + 20 + by = cy + 44 + for k, t in enumerate(entry["best_token_top"]): + label = t["str"].strip() or "␣" + if len(label) > 9: + label = label[:8] + "…" + hit = t["str"].strip() in answer_set + pill_w = 108 + fill = (66, 92, 36) if hit else (16, 22, 28) + outline = P["green"] if hit else (38, 52, 61) + draw.rounded_rectangle((bx, by, bx + pill_w, by + 30), radius=8, fill=fill, outline=outline, width=2) + draw.text((bx + 8, by + 6), label, fill=(220, 255, 190) if hit else P["ink"], font=label_font(label, 13)) + bar = round(min(1.0, t["p"] / 0.4) * pill_w) + draw.rounded_rectangle((bx, by + 36, bx + max(3, bar), by + 42), radius=3, fill=P["amber"] if hit else (60, 76, 88)) + draw.text((bx, by + 46, ), f"{t['p']:.3f}", fill=P["muted"], font=mono_font(10)) + bx += pill_w + 12 + if is_lock: + crosshair(draw, shaft_x, ly, 26, P["amber"], 4) + draw.text((shaft_x + 44, ly + 26), f"first top-1 hit: “{entry['best_token_top'][0]['str'].strip()}” p={entry['best_token_top'][0]['p']:.2f}", fill=P["amber"], font=ui_font(16, True)) + + # ---- Why it matters (right column). + why = (1460, 248, 2136, 716) + draw.rounded_rectangle(why, radius=22, fill=P["panel"], outline=(35, 49, 59), width=1) + draw.text((1492, 270), "why lock-on is the metric", fill=P["ink"], font=ui_font(26, True)) + lines = [ + ("It separates decoding from reasoning.", P["ink"]), + ("Layers before lock-on are spent turning", P["muted"]), + ("pixels into words; layers after are free to", P["muted"]), + ("reason about them. Earlier, harder lock-on", P["muted"]), + ("= more of the network left for thinking.", P["muted"]), + ] + ty = 314 + for text, color in lines: + draw.text((1492, ty), text, fill=color, font=ui_font(19)) + ty += 30 + draw.line((1492, ty + 8, 2104, ty + 8), fill=P["grid"], width=1) + ty += 26 + draw.text((1492, ty), "reasoning budget after lock-on", fill=P["muted"], font=ui_font(16, True)) + ty += 30 + for name, color in COND_COLORS.items(): + c = conditions.get(name) + if not c or c["lock_on_layer"] is None: + continue + budget = n_layers - 1 - c["lock_on_layer"] + bw_px = round(budget / (n_layers - 1) * 430) + draw.text((1492, ty), name, fill=color, font=mono_font(13)) + draw.rounded_rectangle((1492, ty + 20, 1492 + bw_px, ty + 32), radius=6, fill=color) + draw.text((1492 + bw_px + 10, ty + 17), f"{budget} layers · p {c['max_answer_p']:.2f}", fill=P["muted"], font=mono_font(12)) + ty += 44 + + # ---- Rule plate (bottom right). + plate = (1460, 740, 2136, 1000) + draw.rounded_rectangle(plate, radius=22, fill=P["panel"], outline=(255, 196, 68), width=2) + draw.text((1492, 762), "the rule", fill=P["amber"], font=ui_font(24, True)) + rule_lines = [ + "lock_on(patch) = min L such that", + " argmax softmax(W * norm(h_L))", + f" in {{{answer_strs[0]!r}, {answer_strs[1]!r}}}", + "", + f"here: L = {lock_on}, decoded “{layers[lock_on]['best_token_top'][0]['str'].strip()}”", + ] + ry = 806 + for line in rule_lines: + draw.text((1492, ry), line, fill=P["ink"] if line else P["muted"], font=mono_font(17)) + ry += 32 + draw.text((1492, 1014), f"question: {q['q'][:60]}…", fill=P["muted"], font=ui_font(15)) + draw.text((1492, 1040), f"answer: “{q['answer_text']}” · condition: {args.condition} · generation: “{cond['generation']}”", fill=P["muted"], font=ui_font(15)) + + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + canvas.save(out) + print(out) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_logit_lens_dump.py b/packages/snapcompact/research/snapcompact_logit_lens_dump.py new file mode 100644 index 000000000..095c5de17 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_logit_lens_dump.py @@ -0,0 +1,143 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy", "torch", "transformers"] +# /// +"""Logit-lens dump: what vocabulary word does each visual token become, per layer? + +For the visual tokens covering the answer word in a snapcompact bitmap, decode +every layer's hidden state through the final norm + lm_head and record the +top-k vocabulary tokens. If the bitmap is truly read into text space, the +patches' hidden states should decode to the answer's BPE tokens mid-stack. +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path +from typing import Any + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import capacity, render # noqa: E402 +from run import CACHE, FONTS, load_prompt # noqa: E402 +from snapcompact_blackbox_occlusion import sample_answer_questions # noqa: E402 +from snapcompact_text_image_compare import image_answer_token_indices # noqa: E402 + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--model-dir", default="Qwen/Qwen2.5-VL-7B-Instruct") + ap.add_argument("--font", default="8x13", choices=sorted(FONTS)) + ap.add_argument("--variant", default="bw") + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--limit-paras", type=int, default=80) + ap.add_argument("--question-index", type=int, default=3) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--topk", type=int, default=5) + ap.add_argument("--control-tokens", type=int, default=2) + ap.add_argument("--out", default="qwen-logit-lens") + args = ap.parse_args() + + import torch + from transformers import AutoProcessor, AutoTokenizer, Qwen2_5_VLForConditionalGeneration + + out_dir = HERE / "results" / args.out + img_dir = out_dir / "images" + img_dir.mkdir(parents=True, exist_ok=True) + CACHE.mkdir(exist_ok=True) + + cfg = FONTS[args.font] + cols, rows, budget = capacity(cfg, args.size) + paras = squad.load_paragraphs(CACHE)[: args.limit_paras] + flow, offsets = squad.build_flow(paras) + chunk = flow[: min(len(flow), budget)] + questions = sample_answer_questions(paras, offsets, 0, len(chunk), 24, args.seed) + q = questions[min(args.question_index, len(questions) - 1)] + img = render(chunk, cfg, CACHE, args.size, args.variant) + img.save(img_dir / "image-carrier.png") + + print(f"loading {args.model_dir}", flush=True) + processor = AutoProcessor.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True, use_fast=False) + tokenizer = AutoTokenizer.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True) + model = Qwen2_5_VLForConditionalGeneration.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True, dtype=torch.bfloat16, device_map="auto").eval() + device = next(model.parameters()).device + + prompt = load_prompt("qa-image.md").format(cols=cols, rows=rows) + f"\n\nQuestion: {q['q']}\nAnswer with only the shortest extractive answer." + messages = [{"role": "user", "content": [{"type": "image", "image": img}, {"type": "text", "text": prompt}]}] + templated = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + batch = processor(images=img, text=templated, return_tensors="pt") + image_token_id = processor.tokenizer.convert_tokens_to_ids(processor.image_token) + ids = batch["input_ids"][0].tolist() + image_positions = [i for i, token_id in enumerate(ids) if token_id == image_token_id] + n_tokens = len(image_positions) + grid = int(round(n_tokens**0.5)) + answer_indices = image_answer_token_indices(q["answer_start"], q["answer_end"], cols, cfg.adv, cfg.pitch, img.width, img.height, n_tokens) + + # Controls: blank-region tokens far from any text row boundary effects. + control_indices = [] + if answer_indices: + row_far = (answer_indices[0] // grid + grid // 2) % grid + for k in range(args.control_tokens): + control_indices.append(row_far * grid + (answer_indices[0] % grid + k)) + track = [("answer", idx) for idx in answer_indices] + [("control", idx) for idx in control_indices] + track_positions = [image_positions[idx] for _kind, idx in track] + + batch = {k: (v.to(device) if hasattr(v, "to") else v) for k, v in batch.items()} + with torch.no_grad(): + out = model(**batch, output_hidden_states=True, use_cache=False) + + norm = model.model.language_model.norm + lm_head = model.lm_head + answer_token_ids = tokenizer(q["answer_text"], add_special_tokens=False)["input_ids"] + answer_token_strs = [tokenizer.decode([t]) for t in answer_token_ids] + + lens: list[dict[str, Any]] = [] + with torch.no_grad(): + for layer, hidden in enumerate(out.hidden_states): + states = hidden[0, track_positions, :] + logits = lm_head(norm(states)).float() + probs = torch.softmax(logits, dim=-1) + topv, topi = probs.topk(args.topk, dim=-1) + for ti, (kind, idx) in enumerate(track): + entry = { + "layer": layer, + "kind": kind, + "token_index": int(idx), + "grid_rc": [int(idx // grid), int(idx % grid)], + "top": [ + {"str": tokenizer.decode([int(topi[ti, k])]), "id": int(topi[ti, k]), "p": round(float(topv[ti, k]), 5)} + for k in range(args.topk) + ], + "answer_token_p": [round(float(probs[ti, t]), 6) for t in answer_token_ids], + } + lens.append(entry) + print(f"layer {layer} done", flush=True) + + dump = { + "args": vars(args), + "question": {"q": q["q"], "answer_text": q["answer_text"], "answer_start": q["answer_start"], "answer_end": q["answer_end"]}, + "geometry": {"cols": cols, "rows": rows, "image_w": img.width, "image_h": img.height}, + "image_tokens": n_tokens, + "image_grid": grid, + "token_pixel_size": 28, + "answer_token_ids": answer_token_ids, + "answer_token_strs": answer_token_strs, + "answer_indices": [int(i) for i in answer_indices], + "control_indices": [int(i) for i in control_indices], + "layers": len(out.hidden_states), + "lens": lens, + } + (out_dir / "logit_lens.json").write_text(json.dumps(dump, indent=1)) + # Quick console summary: best layer per answer token. + for kind, idx in track: + best = max((e for e in lens if e["token_index"] == idx), key=lambda e: max(e["answer_token_p"])) + print(kind, idx, "best layer", best["layer"], "p", max(best["answer_token_p"]), "top1", best["top"][0]["str"]) + print(f"results -> {out_dir}") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_logit_lens_viz.py b/packages/snapcompact/research/snapcompact_logit_lens_viz.py new file mode 100644 index 000000000..9c9b4eb15 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_logit_lens_viz.py @@ -0,0 +1,169 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy"] +# /// +"""Render the logit-lens grid: pixel patches morphing into BPE tokens by layer.""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path +from typing import Any + +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +PALETTE = { + "bg": (5, 7, 10), + "panel": (12, 17, 23), + "panel2": (8, 12, 17), + "ink": (241, 239, 224), + "muted": (143, 154, 160), + "cyan": (75, 220, 255), + "orange": (255, 112, 72), + "green": (148, 255, 117), + "amber": (255, 196, 68), + "grid": (38, 49, 58), +} + + +def ui_font(size: int, bold: bool = False) -> ImageFont.ImageFont: + for path in [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def mono_font(size: int) -> ImageFont.ImageFont: + for path in ["/System/Library/Fonts/Monaco.ttf", "/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf"]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def heat_fill(p: float, hit: bool) -> tuple[int, int, int]: + """Cell background: brightness by confidence, hue by answer-match.""" + t = min(1.0, max(0.0, math.log10(max(p, 1e-6)) / 3 + 1)) # p=1 -> 1, p=1e-3 -> 0 + if hit: + return (round(30 + 130 * t), round(48 + 130 * t), round(18 + 40 * t)) + return (round(14 + 26 * t), round(19 + 30 * t), round(26 + 36 * t)) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--result-dir", default=str(HERE / "results" / "qwen-logit-lens-q3")) + ap.add_argument("--out", default=str(HERE / "results" / "qwen-logit-lens-q3" / "logit-lens-grid.png")) + ap.add_argument("--layer-step", type=int, default=1) + args = ap.parse_args() + result_dir = Path(args.result_dir) + dump = json.loads((result_dir / "logit_lens.json").read_text()) + carrier = Image.open(result_dir / "images" / "image-carrier.png").convert("RGB") + + q = dump["question"] + answer = q["answer_text"] + answer_ids = set(dump["answer_token_ids"]) + grid = dump["image_grid"] + px = dump["token_pixel_size"] + rw = grid * px + resized = carrier.resize((rw, rw), Image.Resampling.LANCZOS) + + track_indices = dump["answer_indices"] + dump["control_indices"] + by_token: dict[int, list[dict[str, Any]]] = {} + for e in dump["lens"]: + by_token.setdefault(e["token_index"], []).append(e) + for entries in by_token.values(): + entries.sort(key=lambda e: e["layer"]) + n_layers = dump["layers"] + layer_rows = list(range(0, n_layers, args.layer_step)) + + cell_w, cell_h = 150, 34 + header_h = 210 + left_w = 120 + n_cols = len(track_indices) + grid_w = left_w + n_cols * cell_w + margin = 64 + title_h = 200 + w = max(1900, grid_w + margin * 2) + h = title_h + header_h + len(layer_rows) * cell_h + 160 + canvas = Image.new("RGB", (w, h), PALETTE["bg"]) + draw = ImageDraw.Draw(canvas) + for y in range(0, h, 16): + draw.line((0, y, w, y), fill=(7, 10 + y % 9, 15 + y % 11)) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-240, -200, 900, 700), fill=(75, 220, 255, 26)) + gd.ellipse((w - 1000, h - 800, w + 240, h + 200), fill=(255, 112, 72, 24)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(86))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + draw.text((margin, 42), "QWEN LOGIT LENS — PIXELS BECOMING WORDS", fill=PALETTE["amber"], font=ui_font(24, True)) + draw.text((margin, 84), "Watch each patch decode into vocabulary", fill=PALETTE["ink"], font=ui_font(58, True)) + draw.text( + (margin + 2, 156), + f"Each column is one 28×28px visual token; each row is a decoder layer projected through the LM head. Green cells decode to a BPE piece of “{answer}”.", + fill=PALETTE["muted"], + font=ui_font(22), + ) + + gx0 = margin + left_w + gy0 = title_h + header_h + # Column headers: actual patch pixels. + patch_size = 108 + for ci, idx in enumerate(track_indices): + r, c = idx // grid, idx % grid + cell = resized.crop((c * px, r * px, (c + 1) * px, (r + 1) * px)).resize((patch_size, patch_size), Image.Resampling.NEAREST) + cx = gx0 + ci * cell_w + (cell_w - patch_size) // 2 + is_control = idx in dump["control_indices"] + color = PALETTE["muted"] if is_control else PALETTE["orange"] + draw.rounded_rectangle((cx - 4, title_h + 26, cx + patch_size + 4, title_h + 34 + patch_size), radius=8, fill=(244, 242, 230), outline=color, width=3) + canvas.paste(cell, (cx, title_h + 30)) + label = "control" if is_control else f"tok[{idx}]" + tw = draw.textlength(label, font=mono_font(13)) + draw.text((cx + (patch_size - tw) / 2, title_h + 42 + patch_size), label, fill=color, font=mono_font(13)) + draw.text((margin, title_h + 30 + patch_size // 2 - 10), "input\npixels", fill=PALETTE["muted"], font=ui_font(15, True)) + + # Grid rows. + fnt = mono_font(14) + for ri, layer in enumerate(layer_rows): + y = gy0 + ri * cell_h + draw.text((margin + 24, y + 8), f"L{layer:02d}", fill=PALETTE["muted"], font=mono_font(13)) + for ci, idx in enumerate(track_indices): + e = by_token[idx][layer] + top = e["top"][0] + hit = top["id"] in answer_ids + p_ans = max(e["answer_token_p"]) + x = gx0 + ci * cell_w + fill = heat_fill(top["p"] if not hit else max(top["p"], p_ans), hit) + draw.rounded_rectangle((x + 2, y + 2, x + cell_w - 6, y + cell_h - 4), radius=6, fill=fill, outline=(32, 44, 53), width=1) + label = top["str"].replace("\n", "⏎").strip() or "·" + if len(label) > 12: + label = label[:11] + "…" + color = (220, 255, 190) if hit else PALETTE["ink"] if top["p"] > 0.05 else PALETTE["muted"] + draw.text((x + 10, y + 8), label, fill=color, font=fnt) + if hit: + draw.text((x + cell_w - 52, y + 9), f"{p_ans:.2f}", fill=PALETTE["green"], font=mono_font(11)) + + # Footer. + fy = gy0 + len(layer_rows) * cell_h + 22 + draw.rounded_rectangle((margin, fy, w - margin, fy + 88), radius=18, fill=PALETTE["panel2"], outline=(34, 48, 58), width=1) + draw.text((margin + 28, fy + 16), f"question: {q['q'][:88]}", fill=PALETTE["ink"], font=ui_font(19)) + draw.text( + (margin + 28, fy + 50), + f"gold answer “{answer}” = BPE {dump['answer_token_strs']} · logit lens = hidden state → final norm → LM head · {dump['image_tokens']:,} visual tokens total, showing the {len(dump['answer_indices'])} covering the answer + {len(dump['control_indices'])} blank-region controls", + fill=PALETTE["muted"], + font=ui_font(16), + ) + + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + canvas.save(out) + print(out) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_materialize_sweep.py b/packages/snapcompact/research/snapcompact_materialize_sweep.py new file mode 100644 index 000000000..68ab4cfc5 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_materialize_sweep.py @@ -0,0 +1,226 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy", "torch", "transformers"] +# /// +"""Can rendering choices make visual tokens decode to vocabulary EARLIER? + +Sweeps rendering conditions (baseline, line-repeat-in-color, patch-aligned +glyph grids) over the same content/question and measures, per condition, the +layer at which the answer word materializes in logit-lens vocabulary space. +""" + +from __future__ import annotations + +import argparse +import json +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import FontCfg, capacity, render # noqa: E402 +from run import CACHE, FONTS, load_prompt # noqa: E402 +from snapcompact_blackbox_occlusion import sample_answer_questions # noqa: E402 + +TOKEN_PX = 28 # Qwen2.5-VL merged visual token size at native resolution + + +@dataclass(frozen=True) +class Condition: + name: str + cfg: FontCfg + variant: str + repeat: int # each text line rendered this many times consecutively + note: str + + +CONDITIONS = [ + Condition("base-8x13", FONTS["8x13"], "bw", 1, "baseline: glyphs straddle token cells on both axes"), + Condition("repeat2-color", FONTS["8x13"], "color", 2, "every line twice, consecutive rows in different hues"), + Condition("align-7x14", FontCfg("7x14a", "7x13", 7, 14), "bw", 1, "4 chars x 2 rows per token, no straddling"), + Condition("align-14x28", FontCfg("14x28a", "7x13", 14, 28, native=(7, 14)), "bw", 1, "2 chars x 1 row per token"), + Condition("align-28x28", FontCfg("28x28a", "8x13", 28, 28, native=(8, 13)), "bw", 1, "1 char per token"), + Condition("repeat2-align-14x28", FontCfg("14x28a", "7x13", 14, 28, native=(7, 14)), "color", 2, "aligned + repeated lines in hues"), +] + + +def build_layout(chunk: str, cols: int, rows: int, repeat: int) -> tuple[str, int]: + """Row-major render string with each line repeated `repeat` times. + + Returns (render_text, usable_chars) where usable_chars is how much of + `chunk` actually fits. + """ + if repeat == 1: + usable = min(len(chunk), cols * rows) + return chunk[:usable], usable + lines = rows // repeat + usable = min(len(chunk), cols * lines) + out: list[str] = [] + for li in range(lines): + line = chunk[li * cols : (li + 1) * cols].ljust(cols) + out.append(line * repeat) + return "".join(out), usable + + +def answer_token_indices(start: int, end: int, cols: int, adv: int, pitch: int, repeat: int, image_size: int, grid: int) -> list[int]: + """Visual-token indices covering chars [start, end) under the layout.""" + indices: set[int] = set() + for i in range(start, end): + row = i // cols + col = i % cols + for copy in range(repeat): + render_row = row * repeat + copy + x0 = col * adv + x1 = min(image_size - 1, (col + 1) * adv - 1) + y0 = render_row * pitch + y1 = min(image_size - 1, (render_row + 1) * pitch - 1) + if y0 >= image_size: + continue + for x in (x0, x1): + for y in (y0, y1): + indices.add((y // TOKEN_PX) * grid + (x // TOKEN_PX)) + return sorted(indices) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--model-dir", default="Qwen/Qwen2.5-VL-7B-Instruct") + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--limit-paras", type=int, default=80) + ap.add_argument("--question-index", type=int, default=3) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--topk", type=int, default=5) + ap.add_argument("--out", default="qwen-materialize-sweep") + args = ap.parse_args() + + import torch + from transformers import AutoProcessor, AutoTokenizer, Qwen2_5_VLForConditionalGeneration + + out_dir = HERE / "results" / args.out + img_dir = out_dir / "images" + img_dir.mkdir(parents=True, exist_ok=True) + CACHE.mkdir(exist_ok=True) + + # The question is selected on the baseline layout; every condition renders a + # prefix of the same flow, so chunk-relative answer offsets are unchanged. + base_cfg = FONTS["8x13"] + base_cols, base_rows, base_budget = capacity(base_cfg, args.size) + paras = squad.load_paragraphs(CACHE)[: args.limit_paras] + flow, offsets = squad.build_flow(paras) + base_chunk = flow[: min(len(flow), base_budget)] + questions = sample_answer_questions(paras, offsets, 0, len(base_chunk), 24, args.seed) + q = questions[min(args.question_index, len(questions) - 1)] + print(f"question: {q['q']!r} answer: {q['answer_text']!r} @ {q['answer_start']}", flush=True) + + print(f"loading {args.model_dir}", flush=True) + processor = AutoProcessor.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True, use_fast=False) + tokenizer = AutoTokenizer.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True) + model = Qwen2_5_VLForConditionalGeneration.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True, dtype=torch.bfloat16, device_map="auto").eval() + device = next(model.parameters()).device + image_token_id = processor.tokenizer.convert_tokens_to_ids(processor.image_token) + answer_token_ids = tokenizer(q["answer_text"], add_special_tokens=False)["input_ids"] + answer_id_set = set(answer_token_ids) + answer_token_strs = [tokenizer.decode([t]) for t in answer_token_ids] + norm = model.model.language_model.norm + lm_head = model.lm_head + + conditions_out: list[dict[str, Any]] = [] + for cond in CONDITIONS: + cols, rows, _cap = capacity(cond.cfg, args.size) + render_text, usable = build_layout(flow[: cols * rows], cols, rows, cond.repeat) + if q["answer_end"] > usable: + print(f"SKIP {cond.name}: answer beyond capacity ({usable})", flush=True) + continue + img = render(render_text, cond.cfg, CACHE, args.size, cond.variant) + img.save(img_dir / f"{cond.name}.png") + + prompt = load_prompt("qa-image.md").format(cols=cols, rows=rows) + f"\n\nQuestion: {q['q']}\nAnswer with only the shortest extractive answer." + messages = [{"role": "user", "content": [{"type": "image", "image": img}, {"type": "text", "text": prompt}]}] + templated = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + batch = processor(images=img, text=templated, return_tensors="pt") + ids = batch["input_ids"][0].tolist() + image_positions = [i for i, token_id in enumerate(ids) if token_id == image_token_id] + grid = int(round(len(image_positions) ** 0.5)) + track = answer_token_indices(q["answer_start"], q["answer_end"], cols, cond.cfg.adv, cond.cfg.pitch, cond.repeat, args.size, grid) + track_positions = [image_positions[idx] for idx in track] + batch = {k: (v.to(device) if hasattr(v, "to") else v) for k, v in batch.items()} + + with torch.no_grad(): + fwd = model(**batch, output_hidden_states=True, use_cache=False) + generated = model.generate(**batch, max_new_tokens=16, do_sample=False) + answer_gen = processor.batch_decode(generated[:, batch["input_ids"].shape[1] :], skip_special_tokens=True)[0].strip() + + layers_data: list[dict[str, Any]] = [] + lock_on_layer: int | None = None + soft_layer: int | None = None + with torch.no_grad(): + for layer, hidden in enumerate(fwd.hidden_states): + states = hidden[0, track_positions, :] + logits = lm_head(norm(states)).float() + probs = torch.softmax(logits, dim=-1) + top1 = probs.argmax(dim=-1) + answer_p = probs[:, answer_token_ids] # [n_track, n_answer_tokens] + best_p = float(answer_p.max()) + top1_hit = any(int(t) in answer_id_set for t in top1) + best_idx = int(answer_p.max(dim=1).values.argmax()) + tv, ti = probs[best_idx].topk(args.topk) + layers_data.append( + { + "layer": layer, + "best_answer_p": round(best_p, 6), + "top1_hit": bool(top1_hit), + "best_token_index": track[best_idx], + "best_token_top": [ + {"str": tokenizer.decode([int(ti[k])]), "p": round(float(tv[k]), 5)} for k in range(args.topk) + ], + } + ) + if top1_hit and lock_on_layer is None: + lock_on_layer = layer + if best_p > 0.1 and soft_layer is None: + soft_layer = layer + del fwd + torch.cuda.empty_cache() + + result = { + "name": cond.name, + "note": cond.note, + "variant": cond.variant, + "repeat": cond.repeat, + "adv": cond.cfg.adv, + "pitch": cond.cfg.pitch, + "cols": cols, + "rows": rows, + "usable_chars": usable, + "chars_per_token": round(cols * rows / (cond.repeat * grid * grid), 2), + "tracked_tokens": track, + "generation": answer_gen, + "generation_correct": q["answer_text"].lower() in answer_gen.lower(), + "lock_on_layer": lock_on_layer, + "soft_layer_p10": soft_layer, + "max_answer_p": max(l["best_answer_p"] for l in layers_data), + "layers": layers_data, + } + conditions_out.append(result) + print( + f"{cond.name}: lock_on={lock_on_layer} soft={soft_layer} max_p={result['max_answer_p']:.3f} gen={answer_gen!r}", + flush=True, + ) + + summary = { + "args": vars(args), + "question": {"q": q["q"], "answer_text": q["answer_text"], "answer_start": q["answer_start"], "answer_end": q["answer_end"]}, + "answer_token_ids": answer_token_ids, + "answer_token_strs": answer_token_strs, + "conditions": conditions_out, + } + (out_dir / "summary.json").write_text(json.dumps(summary, indent=1)) + print(f"results -> {out_dir}") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_materialize_viz.py b/packages/snapcompact/research/snapcompact_materialize_viz.py new file mode 100644 index 000000000..0e5414f82 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_materialize_viz.py @@ -0,0 +1,171 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy"] +# /// +"""Render the materialization sweep: rendering choices vs logit-lens confidence.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +PALETTE = { + "bg": (5, 7, 10), + "panel": (12, 17, 23), + "panel2": (8, 12, 17), + "ink": (241, 239, 224), + "muted": (143, 154, 160), + "grid": (38, 49, 58), +} +SERIES = [ + ("base-8x13", (143, 154, 160)), + ("repeat2-color", (255, 196, 68)), + ("align-7x14", (148, 255, 117)), + ("align-14x28", (75, 220, 255)), + ("align-28x28", (255, 112, 72)), + ("repeat2-align-14x28", (188, 112, 255)), +] + + +def ui_font(size: int, bold: bool = False) -> ImageFont.ImageFont: + for path in [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def mono_font(size: int) -> ImageFont.ImageFont: + for path in ["/System/Library/Fonts/Monaco.ttf", "/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf"]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def crop_answer_region(img_path: Path, cond: dict[str, Any], answer_start: int, answer_end: int, image_size: int = 1568) -> Image.Image: + img = Image.open(img_path).convert("RGB") + cols = cond["cols"] + adv = cond["adv"] + pitch = cond["pitch"] + repeat = cond["repeat"] + row = answer_start // cols + c0 = answer_start % cols + c1 = min(cols - 1, (answer_end - 1) % cols) + y0 = max(0, row * repeat * pitch - pitch) + y1 = min(image_size, (row * repeat + repeat) * pitch + pitch) + x0 = max(0, c0 * adv - 10 * adv) + x1 = min(image_size, (c1 + 1) * adv + 10 * adv) + crop = img.crop((x0, y0, x1, y1)) + d = ImageDraw.Draw(crop) + d.rectangle((c0 * adv - x0 - 2, row * repeat * pitch - y0 - 1, (c1 + 1) * adv - x0 + 2, (row * repeat + repeat) * pitch - y0 + 1), outline=(255, 112, 72), width=3) + return crop + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--result-dir", default=str(HERE / "results" / "qwen-materialize-sweep-q3")) + ap.add_argument("--out", default=str(HERE / "results" / "qwen-materialize-sweep-q3" / "materialize-sweep.png")) + args = ap.parse_args() + result_dir = Path(args.result_dir) + summary = json.loads((result_dir / "summary.json").read_text()) + conditions = {c["name"]: c for c in summary["conditions"]} + q = summary["question"] + + w, h = 2200, 1380 + canvas = Image.new("RGB", (w, h), PALETTE["bg"]) + draw = ImageDraw.Draw(canvas) + for y in range(0, h, 16): + draw.line((0, y, w, y), fill=(7, 10 + y % 9, 15 + y % 11)) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-240, -200, 900, 700), fill=(75, 220, 255, 25)) + gd.ellipse((1240, 540, 2460, 1480), fill=(255, 112, 72, 24)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(86))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + draw.text((64, 42), "QWEN MATERIALIZATION SWEEP — CAN RENDERING MOVE THE LAYER?", fill=(255, 196, 68), font=ui_font(24, True)) + draw.text((64, 84), "The depth is the model's; the clarity is yours", fill=PALETTE["ink"], font=ui_font(58, True)) + draw.text( + (66, 158), + "Six renderings of the same passage. Logit-lens p(answer BPE) at the answer patch, by layer.\n" + "Alignment and repetition barely move WHERE it materializes — they transform HOW HARD.", + fill=PALETTE["muted"], + font=ui_font(22), + ) + + # Main curve panel. + panel = (64, 226, 1380, 900) + draw.rounded_rectangle(panel, radius=26, fill=PALETTE["panel"], outline=(35, 49, 59), width=1) + draw.text((96, 248), "p(answer BPE) at the best answer patch, per layer", fill=PALETTE["ink"], font=ui_font(26, True)) + gx0, gy0, gx1, gy1 = 150, 320, 1330, 800 + n_layers = len(conditions["base-8x13"]["layers"]) + for i in range(6): + yy = gy0 + (gy1 - gy0) * i / 5 + draw.line((gx0, yy, gx1, yy), fill=PALETTE["grid"], width=1) + draw.text((96, yy - 9), f"{1.0 - i / 5:.1f}", fill=PALETTE["muted"], font=ui_font(14)) + for name, color in SERIES: + cond = conditions.get(name) + if not cond: + continue + pts = [] + for e in cond["layers"]: + x = gx0 + (gx1 - gx0) * e["layer"] / (n_layers - 1) + y = gy1 - (gy1 - gy0) * min(1.0, e["best_answer_p"]) + pts.append((round(x), round(y))) + draw.line(pts, fill=color, width=5 if name != "base-8x13" else 4, joint="curve") + if cond["lock_on_layer"] is not None: + lx = gx0 + (gx1 - gx0) * cond["lock_on_layer"] / (n_layers - 1) + draw.ellipse((lx - 7, gy1 - (gy1 - gy0) * min(1.0, cond["layers"][cond["lock_on_layer"]]["best_answer_p"]) - 7, lx + 7, gy1 - (gy1 - gy0) * min(1.0, cond["layers"][cond["lock_on_layer"]]["best_answer_p"]) + 7), outline=color, width=3) + draw.text((gx0, gy1 + 16), "layer 0", fill=PALETTE["muted"], font=ui_font(15)) + draw.text((gx1 - 76, gy1 + 16), f"layer {n_layers - 1}", fill=PALETTE["muted"], font=ui_font(15)) + draw.text((gx0 + 320, gy1 + 16), "rings mark lock-on (top-1 becomes an answer BPE)", fill=PALETTE["muted"], font=ui_font(15)) + legend_box = (1420, 226, 2136, 900) + draw.rounded_rectangle(legend_box, radius=26, fill=PALETTE["panel"], outline=(35, 49, 59), width=1) + draw.text((1452, 248), "conditions", fill=PALETTE["ink"], font=ui_font(26, True)) + ly = 304 + for name, color in SERIES: + cond = conditions.get(name) + if not cond: + continue + draw.rounded_rectangle((1452, ly, 1452 + 26, ly + 10), radius=4, fill=color) + draw.text((1492, ly - 9), name, fill=PALETTE["ink"], font=ui_font(21, True)) + draw.text((1492, ly + 19), cond["note"], fill=PALETTE["muted"], font=ui_font(14)) + draw.text((1492, ly + 42), f"lock-on L{cond['lock_on_layer']} · peak p {cond['max_answer_p']:.2f} · {cond['chars_per_token']} chars/token", fill=color, font=mono_font(14)) + ly += 96 + + # Condition cards with real crops. + card_y = 938 + card_w = 660 + draw.text((64, card_y - 24), "what the model actually saw (answer region outlined)", fill=PALETTE["ink"], font=ui_font(22, True)) + positions = [(64, card_y + 10), (64 + card_w + 24, card_y + 10), (64 + 2 * (card_w + 24), card_y + 10)] + featured = ["base-8x13", "align-28x28", "repeat2-align-14x28"] + for (cx, cy), name in zip(positions, featured): + cond = conditions.get(name) + if not cond: + continue + color = dict(SERIES)[name] + draw.rounded_rectangle((cx, cy, cx + card_w, cy + 350), radius=20, fill=PALETTE["panel2"], outline=(34, 48, 58), width=1) + draw.text((cx + 22, cy + 14), name, fill=color, font=ui_font(23, True)) + draw.text((cx + 22, cy + 46), cond["note"], fill=PALETTE["muted"], font=ui_font(15)) + crop = crop_answer_region(result_dir / "images" / f"{name}.png", cond, q["answer_start"], q["answer_end"]) + scale = min((card_w - 44) / crop.width, 200 / crop.height) + crop_r = crop.resize((round(crop.width * scale), round(crop.height * scale)), Image.Resampling.NEAREST) + draw.rounded_rectangle((cx + 20, cy + 76, cx + card_w - 20, cy + 286), radius=12, fill=(244, 242, 230)) + canvas.paste(crop_r, (cx + 22 + (card_w - 44 - crop_r.width) // 2, cy + 78 + (206 - crop_r.height) // 2)) + draw.text((cx + 22, cy + 300), f"lock-on L{cond['lock_on_layer']} · peak p {cond['max_answer_p']:.2f} · {cond['chars_per_token']} chars/token · gen “{cond['generation']}”", fill=PALETTE["ink"], font=ui_font(16, True)) + + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + canvas.save(out) + print(out) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_pricing_viz.py b/packages/snapcompact/research/snapcompact_pricing_viz.py new file mode 100644 index 000000000..6e113791f --- /dev/null +++ b/packages/snapcompact/research/snapcompact_pricing_viz.py @@ -0,0 +1,137 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""Pricing graphic for the snapcompact post: what a PNG bills vs what it carries.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +P = { + "bg": (5, 7, 10), + "panel": (12, 17, 23), + "panel2": (8, 12, 17), + "ink": (241, 239, 224), + "muted": (143, 154, 160), + "cyan": (75, 220, 255), + "orange": (255, 112, 72), + "green": (148, 255, 117), + "amber": (255, 196, 68), + "grid": (38, 49, 58), +} + +# (label, text-token equivalent, note). Billed image tokens are constant per canvas. +CARRY = [ + ("8x13 font · 1568²", 5219, "23,520 chars · measured BPE count", 3279), + ("6x10 font · 1568²", 10000, "40,716 chars · ~4 chars/token", 3279), + ("6x10 font · 2576²", 25000, "102,000 chars · one whole corpus", 4950), +] + + +def ui_font(size: int, bold: bool = False) -> ImageFont.ImageFont: + for path in [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def mono_font(size: int) -> ImageFont.ImageFont: + for path in ["/System/Library/Fonts/Monaco.ttf", "/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf"]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--out", default=str(HERE / "results" / "snapcompact-pricing.png")) + args = ap.parse_args() + + w, h = 2200, 1000 + canvas = Image.new("RGB", (w, h), P["bg"]) + draw = ImageDraw.Draw(canvas) + for y in range(0, h, 16): + draw.line((0, y, w, y), fill=(7, 10 + y % 9, 15 + y % 11)) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-260, -240, 800, 560), fill=(75, 220, 255, 26)) + gd.ellipse((1400, 400, 2460, 1240), fill=(255, 196, 68, 26)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(88))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + draw.text((64, 40), "THE BILLING MATH", fill=P["amber"], font=ui_font(24, True)) + draw.text((64, 80), "A flat fee per canvas, no matter what's inside", fill=P["ink"], font=ui_font(56, True)) + draw.text((66, 152), "Anthropic bills images at width × height ÷ 750 tokens. Text tokens scale with content; image tokens scale with pixels. Dense fonts exploit the gap.", fill=P["muted"], font=ui_font(22)) + + # Formula card. + card = (64, 224, 700, 420) + draw.rounded_rectangle(card, radius=22, fill=P["panel"], outline=(35, 49, 59), width=1) + draw.text((96, 246), "flat fee per canvas", fill=P["cyan"], font=ui_font(21, True)) + draw.text((96, 286), "1568 × 1568 → 3,279 tokens", fill=P["ink"], font=mono_font(24)) + draw.text((96, 326), "2576 × 2576 → 4,950 tokens", fill=P["ink"], font=mono_font(24)) + draw.text((96, 372), "(2576 is silently downscaled 0.75x — still the best $/char)", fill=P["muted"], font=ui_font(15)) + + # Cache card. + card = (64, 452, 700, 660) + draw.rounded_rectangle(card, radius=22, fill=P["panel"], outline=(35, 49, 59), width=1) + draw.text((96, 474), "with prompt caching", fill=P["green"], font=ui_font(21, True)) + draw.text((96, 514), "marginal re-ask ≈ 333 tokens/turn", fill=P["ink"], font=mono_font(22)) + draw.text((96, 554), "measured: 753 in · 3,330 cache-write", fill=P["muted"], font=mono_font(17)) + draw.text((96, 584), "16,650 cache-read over six calls", fill=P["muted"], font=mono_font(17)) + draw.text((96, 620), "$0.18 cached vs $0.33 uncached", fill=P["amber"], font=mono_font(18)) + + # Fine-print card. + card = (64, 692, 700, 920) + draw.rounded_rectangle(card, radius=22, fill=P["panel"], outline=(255, 112, 72), width=1) + draw.text((96, 714), "the decode tax", fill=P["orange"], font=ui_font(21, True)) + draw.text((96, 754), "Models reason their way through dense", fill=P["muted"], font=ui_font(18)) + draw.text((96, 782), "pixels: 5–10x more thinking tokens than", fill=P["muted"], font=ui_font(18)) + draw.text((96, 810), "text. Input savings are real; total cost", fill=P["muted"], font=ui_font(18)) + draw.text((96, 838), "depends on output pricing. Cache + re-ask", fill=P["muted"], font=ui_font(18)) + draw.text((96, 866), "is where it always wins.", fill=P["ink"], font=ui_font(18, True)) + + # Carry bars. + panel = (760, 224, 2136, 920) + draw.rounded_rectangle(panel, radius=26, fill=P["panel"], outline=(35, 49, 59), width=1) + draw.text((796, 250), "text-token equivalent carried vs image tokens billed", fill=P["ink"], font=ui_font(26, True)) + draw.text((796, 290), "same content, two meters — the orange bar is what you'd pay as text; the cyan bar is what the PNG bills", fill=P["muted"], font=ui_font(17)) + bx0, bx1 = 1100, 1860 + max_tokens = 25000 + y = 360 + for label, text_tokens, note, billed in CARRY: + draw.text((796, y + 6), label, fill=P["ink"], font=mono_font(17)) + draw.text((796, y + 32), note, fill=P["muted"], font=ui_font(13)) + tw = round((bx1 - bx0) * text_tokens / max_tokens) + bw = round((bx1 - bx0) * billed / max_tokens) + draw.rounded_rectangle((bx0, y, bx0 + tw, y + 26), radius=9, fill=P["orange"]) + draw.text((bx0 + tw + 12, y + 2), f"{text_tokens:,} as text", fill=P["orange"], font=mono_font(15)) + draw.rounded_rectangle((bx0, y + 34, bx0 + bw, y + 60), radius=9, fill=P["cyan"]) + ratio = text_tokens / billed + draw.text((bx0 + bw + 12, y + 36), f"{billed:,} billed · {ratio:.1f}x", fill=P["cyan"], font=mono_font(15)) + y += 130 + # Cached marginal bar. + draw.text((796, y + 6), "any font · cached re-ask", fill=P["ink"], font=mono_font(17)) + draw.text((796, y + 32), "image as cached prefix block", fill=P["muted"], font=ui_font(13)) + bw = max(6, round((bx1 - bx0) * 333 / max_tokens)) + draw.rounded_rectangle((bx0, y + 14, bx0 + bw, y + 40), radius=9, fill=P["green"]) + draw.text((bx0 + bw + 12, y + 16), "≈ 333 tokens/turn · 30x", fill=P["green"], font=mono_font(15)) + y += 110 + draw.line((796, y, 2100, y), fill=P["grid"], width=1) + draw.text((796, y + 16), "10,000 tokens of text, carried by 3,279 image tokens, amortizing to ~333 — that's the whole pitch.", fill=P["amber"], font=ui_font(19, True)) + + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + canvas.save(out) + print(out) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_qwen_control_intervention.py b/packages/snapcompact/research/snapcompact_qwen_control_intervention.py new file mode 100644 index 000000000..5c6f1376a --- /dev/null +++ b/packages/snapcompact/research/snapcompact_qwen_control_intervention.py @@ -0,0 +1,363 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy", "torch", "transformers", "qwen-vl-utils"] +# /// +"""Qwen snapcompact controls: alternate prompt plus activation intervention.""" + +from __future__ import annotations + +import argparse +import json +import math +import random +import sys +from pathlib import Path +from typing import Any + +import numpy as np +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import capacity, render # noqa: E402 +from run import CACHE, FONTS, load_prompt # noqa: E402 +from snapcompact_blackbox_occlusion import sample_answer_questions # noqa: E402 +from snapcompact_text_image_compare import ( # noqa: E402 + cosine, + image_answer_token_indices, + normalize_heat, + run_image, + run_text, + to_device, +) + +PALETTE = { + "bg": (5, 7, 10), + "panel": (12, 17, 23), + "panel2": (8, 12, 17), + "ink": (241, 239, 224), + "muted": (143, 154, 160), + "cyan": (75, 220, 255), + "orange": (255, 112, 72), + "green": (148, 255, 117), + "amber": (255, 196, 68), + "purple": (188, 112, 255), + "red": (255, 76, 62), +} + + +def ui_font(size: int, bold: bool = False) -> ImageFont.ImageFont: + for path in [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def mono_font(size: int) -> ImageFont.ImageFont: + for path in ["/System/Library/Fonts/Monaco.ttf", "/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf"]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def heat_color(t: float) -> tuple[int, int, int]: + t = max(0.0, min(1.0, t)) + stops = [ + (0.00, (4, 7, 20)), + (0.22, (24, 28, 88)), + (0.45, (49, 120, 190)), + (0.65, (54, 226, 195)), + (0.82, (188, 255, 120)), + (1.00, (255, 236, 128)), + ] + for (a, ca), (b, cb) in zip(stops, stops[1:]): + if t <= b: + u = (t - a) / (b - a) + return tuple(round(ca[i] + (cb[i] - ca[i]) * u) for i in range(3)) + return stops[-1][1] + + +def make_text_prompt(chunk: str, q: dict[str, Any]) -> str: + return ( + "Below is reference material. Answer the question using only it.\n\n" + f"{chunk}\n\nQuestion: {q['q']}\n" + "Answer with only the shortest extractive answer." + ) + + +def make_image_prompt(cols: int, rows: int, q: dict[str, Any]) -> str: + return load_prompt("qa-image.md").format(cols=cols, rows=rows) + f"\n\nQuestion: {q['q']}\nAnswer with only the shortest extractive answer." + + +def carrier_map(model: Any, processor: Any, img: Image.Image, chunk: str, q: dict[str, Any], cols: int, rows: int, device: Any) -> tuple[np.ndarray, np.ndarray, dict[str, Any]]: + text_layers, text_pos, _ = run_text(model, processor, make_text_prompt(chunk, q), chunk, q["answer_start"], q["answer_end"], device) + image_layers, image_positions, image_meta, _ = run_image(model, processor, img, make_image_prompt(cols, rows, q), device) + image_count = len(image_positions) + answer_indices = image_answer_token_indices(q["answer_start"], q["answer_end"], cols, 8, 13, img.width, img.height, image_count) + sims = [] + answer_cos = [] + for text_h, image_h in zip(text_layers, image_layers): + text_ans = text_h[text_pos["answer_start"] : text_pos["answer_end"]].mean(axis=0) + image_tokens = image_h[image_positions] + image_ans = image_tokens[answer_indices] if answer_indices else image_tokens + sims.append(cosine(np.repeat(text_ans[None, :], image_tokens.shape[0], axis=0), image_tokens).astype(np.float32, copy=False)) + answer_cos.append(float(cosine(text_ans[None, :], image_ans.mean(axis=0, keepdims=True))[0])) + raw = np.stack(sims, axis=0) + excess = raw - np.median(raw, axis=1, keepdims=True) + norm, lo, hi = normalize_heat(excess) + meta = { + "image_tokens": image_count, + "image_grid": round(math.sqrt(image_count)), + "image_meta": image_meta, + "answer_indices": answer_indices, + "answer_cosine": answer_cos, + "peak_layer": int(np.argmax(answer_cos)), + "peak_cosine": float(max(answer_cos)), + "final_cosine": float(answer_cos[-1]), + "heat_lo": lo, + "heat_hi": hi, + } + return raw, norm, meta + + +def generate_with_intervention( + model: Any, + processor: Any, + img: Image.Image, + prompt: str, + device: Any, + layer: int, + answer_indices: list[int], + mode: str, + seed: int, +) -> str: + import torch + + messages = [{"role": "user", "content": [{"type": "image", "image": img}, {"type": "text", "text": prompt}]}] + templated = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + batch = processor(images=img, text=templated, return_tensors="pt") + image_token_id = processor.tokenizer.convert_tokens_to_ids(processor.image_token) + image_positions = [i for i, token_id in enumerate(batch["input_ids"][0].tolist()) if token_id == image_token_id] + rng = random.Random(seed) + random_indices = sorted(rng.sample([i for i in range(len(image_positions)) if i not in set(answer_indices)], len(answer_indices))) if answer_indices else [] + target_indices = ( + answer_indices + if mode == "answer_mean_patch" + else random_indices + if mode == "random_mean_patch" + else list(range(len(image_positions))) + if mode == "all_image_zero" + else [] + ) + target_positions = [image_positions[i] for i in target_indices] + batch = to_device(batch, device) + + handle = None + if target_positions: + def hook(_module: Any, inputs: tuple[Any, ...]) -> tuple[Any, ...]: + hidden = inputs[0] + if hidden.ndim == 3 and hidden.shape[1] > max(target_positions): + patched = hidden.clone() + if mode == "all_image_zero": + patched[:, target_positions, :] = 0 + else: + source_positions = [p for p in image_positions if p not in target_positions] + mean_vec = hidden[:, source_positions, :].mean(dim=1, keepdim=True) + patched[:, target_positions, :] = mean_vec + return (patched, *inputs[1:]) + return inputs + + handle = model.model.language_model.layers[layer].register_forward_pre_hook(hook) + try: + with torch.no_grad(): + generated = model.generate(**batch, max_new_tokens=24, do_sample=False) + finally: + if handle is not None: + handle.remove() + new_tokens = generated[:, batch["input_ids"].shape[1] :] + return processor.batch_decode(new_tokens, skip_special_tokens=True)[0].strip() + + +def crop_answer(img: Image.Image, q: dict[str, Any], cols: int, adv: int = 8, pitch: int = 13) -> Image.Image: + start = q["answer_start"] + end = q["answer_end"] + row0 = max(0, start // cols - 5) + row1 = min(img.height // pitch, end // cols + 6) + col0 = max(0, start % cols - 42) + col1 = min(cols, end % cols + 42) + crop = img.crop((col0 * adv, row0 * pitch, col1 * adv, row1 * pitch)).convert("RGB") + d = ImageDraw.Draw(crop) + bx0 = max(0, (start % cols - col0) * adv - adv) + bx1 = min(crop.width - 1, ((end - 1) % cols - col0 + 2) * adv) + by0 = max(0, (start // cols - row0) * pitch - 1) + by1 = min(crop.height - 1, ((end - 1) // cols - row0 + 1) * pitch + 1) + d.rounded_rectangle((bx0, by0, bx1, by1), radius=3, outline=PALETTE["orange"], width=3) + return crop + + +def paste_fit(canvas: Image.Image, img: Image.Image, box: tuple[int, int, int, int]) -> None: + x0, y0, x1, y1 = box + scale = min((x1 - x0) / img.width, (y1 - y0) / img.height) + resized = img.resize((max(1, round(img.width * scale)), max(1, round(img.height * scale))), Image.Resampling.NEAREST) + canvas.paste(resized, (x0 + (x1 - x0 - resized.width) // 2, y0 + (y1 - y0 - resized.height) // 2)) + + +def draw_grid(draw: ImageDraw.ImageDraw, grid_values: np.ndarray, answer_indices: list[int], box: tuple[int, int, int, int], title: str, subtitle: str, color: tuple[int, int, int]) -> None: + x0, y0, x1, y1 = box + draw.rounded_rectangle(box, radius=22, fill=PALETTE["panel2"], outline=(34, 48, 58), width=1) + draw.text((x0 + 20, y0 + 18), title, fill=color, font=ui_font(25, True)) + draw.text((x0 + 20, y0 + 50), subtitle, fill=PALETTE["muted"], font=ui_font(15)) + gx0, gy0, gx1, gy1 = x0 + 30, y0 + 84, x1 - 30, y1 - 28 + rows, cols = grid_values.shape + cw = (gx1 - gx0) / cols + ch = (gy1 - gy0) / rows + for r in range(rows): + for c in range(cols): + xa = round(gx0 + c * cw) + xb = round(gx0 + (c + 1) * cw) + ya = round(gy0 + r * ch) + yb = round(gy0 + (r + 1) * ch) + draw.rectangle((xa, ya, xb, yb), fill=heat_color(float(grid_values[r, c]))) + for idx in answer_indices: + r, c = divmod(idx, cols) + xa = round(gx0 + c * cw) + xb = round(gx0 + (c + 1) * cw) + ya = round(gy0 + r * ch) + yb = round(gy0 + (r + 1) * ch) + draw.rectangle((xa - 2, ya - 2, xb + 2, yb + 2), outline=PALETTE["orange"], width=2) + + +def render_figure(out_path: Path, img: Image.Image, primary: dict[str, Any], distractor: dict[str, Any], primary_norm: np.ndarray, distractor_norm: np.ndarray, primary_meta: dict[str, Any], distractor_meta: dict[str, Any], generations: dict[str, str], cols: int) -> None: + w, h = 2200, 1320 + canvas = Image.new("RGB", (w, h), PALETTE["bg"]) + draw = ImageDraw.Draw(canvas) + for y in range(0, h, 16): + draw.line((0, y, w, y), fill=(7, 10 + y % 9, 15 + y % 11)) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-260, -240, 900, 760), fill=(75, 220, 255, 28)) + gd.ellipse((1240, 120, 2480, 1380), fill=(255, 112, 72, 27)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(86))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + draw.text((64, 42), "QWEN SNAPCOMPACT CONTROL + INTERVENTION", fill=PALETTE["amber"], font=ui_font(24, True)) + draw.text((64, 84), "Ask a different thing; patch the hidden answer", fill=PALETTE["ink"], font=ui_font(62, True)) + draw.text((66, 166), "Same bitmap, two questions. Then patch answer-region image-token activations at the peak layer and watch generation change.", fill=PALETTE["muted"], font=ui_font(24)) + + draw.rounded_rectangle((64, 238, 616, 1234), radius=30, fill=PALETTE["panel"], outline=(35, 49, 59), width=1) + draw.text((96, 270), "same image carrier", fill=PALETTE["ink"], font=ui_font(32, True)) + draw.text((96, 312), "Qwen2.5-VL-7B, 1568px bitmap", fill=PALETTE["muted"], font=ui_font(18)) + for label, q, y, color in [("PRIMARY", primary, 374, PALETTE["orange"]), ("DISTRACTOR", distractor, 658, PALETTE["cyan"] )]: + draw.text((96, y), label, fill=color, font=ui_font(17, True)) + crop = crop_answer(img, q, cols) + draw.rounded_rectangle((96, y + 34, 584, y + 194), radius=14, fill=(244, 242, 230), outline=color, width=3) + paste_fit(canvas, crop, (112, y + 48, 568, y + 180)) + draw.text((96, y + 216), q["q"][:58], fill=PALETTE["ink"], font=ui_font(18)) + draw.text((96, y + 244), f"gold: {q['answer_text']}", fill=PALETTE["amber"], font=ui_font(22, True)) + draw.text((96, 1012), f"primary peak: L{primary_meta['peak_layer']} cosine {primary_meta['peak_cosine']:.3f}", fill=PALETTE["orange"], font=ui_font(20, True)) + draw.text((96, 1044), f"distractor peak: L{distractor_meta['peak_layer']} cosine {distractor_meta['peak_cosine']:.3f}", fill=PALETTE["cyan"], font=ui_font(20, True)) + draw.text((96, 1102), f"image tokens: {primary_meta['image_tokens']} ({primary_meta['image_grid']}×{primary_meta['image_grid']})", fill=PALETTE["muted"], font=ui_font(18)) + + grid = primary_meta["image_grid"] + draw_grid(draw, primary_norm[primary_meta["peak_layer"]].reshape(grid, grid), primary_meta["answer_indices"], (666, 238, 1386, 706), "primary question map", f"{primary['answer_text']} @ layer {primary_meta['peak_layer']} — orange box marks true answer", PALETTE["orange"]) + draw_grid(draw, distractor_norm[distractor_meta["peak_layer"]].reshape(grid, grid), distractor_meta["answer_indices"], (1420, 238, 2140, 706), "distractor question map", f"{distractor['answer_text']} @ layer {distractor_meta['peak_layer']} — map should move", PALETTE["cyan"]) + + draw.rounded_rectangle((666, 746, 2140, 1234), radius=30, fill=PALETTE["panel"], outline=(35, 49, 59), width=1) + draw.text((704, 780), "activation patch test", fill=PALETTE["ink"], font=ui_font(34, True)) + draw.text((704, 822), "Before decoder layer 0, replace selected image-token residuals. Local answer patches test specificity; all-image zero is the sanity check.", fill=PALETTE["muted"], font=ui_font(20)) + rows = [ + ("normal", generations["normal"], PALETTE["green"]), + ("patch random region", generations["random_mean_patch"], PALETTE["cyan"]), + ("patch answer region", generations["answer_mean_patch"], PALETTE["red"]), + ("zero all image tokens", generations["all_image_zero"], PALETTE["purple"]), + ] + y = 878 + for label, text, color in rows: + draw.rounded_rectangle((704, y, 2078, y + 74), radius=18, fill=PALETTE["panel2"], outline=(34, 48, 58), width=1) + draw.text((730, y + 16), label.upper(), fill=color, font=ui_font(17, True)) + draw.text((1002, y + 15), text[:115], fill=PALETTE["ink"], font=ui_font(23, True)) + y += 86 + + out_path.parent.mkdir(parents=True, exist_ok=True) + canvas.save(out_path) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--model-dir", default="Qwen/Qwen2.5-VL-7B-Instruct") + ap.add_argument("--font", default="8x13", choices=sorted(FONTS)) + ap.add_argument("--variant", default="bw") + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--limit-paras", type=int, default=80) + ap.add_argument("--qpc", type=int, default=24) + ap.add_argument("--question-index", type=int, default=12) + ap.add_argument("--distractor-index", type=int, default=3) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--out", default="qwen-control-intervention") + args = ap.parse_args() + + import torch + from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration + + out_dir = HERE / "results" / args.out + img_dir = out_dir / "images" + img_dir.mkdir(parents=True, exist_ok=True) + CACHE.mkdir(exist_ok=True) + + cfg = FONTS[args.font] + cols, rows, budget = capacity(cfg, args.size) + paras = squad.load_paragraphs(CACHE)[: args.limit_paras] + flow, offsets = squad.build_flow(paras) + chunk = flow[: min(len(flow), budget)] + questions = sample_answer_questions(paras, offsets, 0, len(chunk), args.qpc, args.seed) + if len(questions) < 2: + raise SystemExit("not enough questions fit in chunk") + primary = questions[min(args.question_index, len(questions) - 1)] + distractor = questions[min(args.distractor_index, len(questions) - 1)] + if distractor is primary: + distractor = questions[0 if args.question_index != 0 else 1] + img = render(chunk, cfg, CACHE, args.size, args.variant) + img.save(img_dir / "image-carrier.png") + + print(f"loading {args.model_dir}", flush=True) + processor = AutoProcessor.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True, use_fast=False) + model = Qwen2_5_VLForConditionalGeneration.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True, dtype=torch.bfloat16, device_map="auto").eval() + device = next(model.parameters()).device + + primary_raw, primary_norm, primary_meta = carrier_map(model, processor, img, chunk, primary, cols, rows, device) + distractor_raw, distractor_norm, distractor_meta = carrier_map(model, processor, img, chunk, distractor, cols, rows, device) + + peak_layer = primary_meta["peak_layer"] + prompt = make_image_prompt(cols, rows, primary) + patch_layer = 0 + generations = { + "normal": generate_with_intervention(model, processor, img, prompt, device, patch_layer, primary_meta["answer_indices"], "none", args.seed), + "random_mean_patch": generate_with_intervention(model, processor, img, prompt, device, patch_layer, primary_meta["answer_indices"], "random_mean_patch", args.seed), + "answer_mean_patch": generate_with_intervention(model, processor, img, prompt, device, patch_layer, primary_meta["answer_indices"], "answer_mean_patch", args.seed), + "all_image_zero": generate_with_intervention(model, processor, img, prompt, device, patch_layer, primary_meta["answer_indices"], "all_image_zero", args.seed), + } + + summary = { + "args": vars(args), + "geometry": {"cols": cols, "rows": rows, "capacity": budget}, + "primary": primary, + "distractor": distractor, + "primary_meta": primary_meta, + "distractor_meta": distractor_meta, + "intervention_layer": patch_layer, + "generations": generations, + } + np.savez_compressed(out_dir / "control_intervention.npz", primary_raw=primary_raw, primary_norm=primary_norm, distractor_raw=distractor_raw, distractor_norm=distractor_norm) + (out_dir / "summary.json").write_text(json.dumps(summary, indent=1)) + render_figure(out_dir / "control-intervention.png", img, primary, distractor, primary_norm, distractor_norm, primary_meta, distractor_meta, generations, cols) + print(json.dumps(summary, indent=1)) + print(f"results -> {out_dir}") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_qwen_spotlight_viz.py b/packages/snapcompact/research/snapcompact_qwen_spotlight_viz.py new file mode 100644 index 000000000..0e34ec882 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_qwen_spotlight_viz.py @@ -0,0 +1,278 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy"] +# /// +"""Render a cleaner prompt-specific spotlight figure for Qwen snapcompact controls.""" + +from __future__ import annotations + +import argparse +import json +import random +from pathlib import Path +from typing import Any + +import numpy as np +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +PALETTE = { + "bg": (5, 7, 10), + "panel": (12, 17, 23), + "panel2": (8, 12, 17), + "ink": (241, 239, 224), + "muted": (143, 154, 160), + "orange": (255, 112, 72), + "cyan": (75, 220, 255), + "green": (148, 255, 117), + "amber": (255, 196, 68), + "red": (255, 76, 62), + "purple": (188, 112, 255), +} + + +def ui_font(size: int, bold: bool = False) -> ImageFont.ImageFont: + for path in [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def heat_color(t: float, theme: str) -> tuple[int, int, int, int]: + t = max(0.0, min(1.0, t)) + if theme == "cyan": + rgb0, rgb1 = (10, 38, 62), (75, 220, 255) + else: + rgb0, rgb1 = (62, 25, 10), (255, 112, 72) + rgb = tuple(round(rgb0[i] + (rgb1[i] - rgb0[i]) * t) for i in range(3)) + alpha = round(20 + 210 * t) + return (*rgb, alpha) + + +def normalize_positive(arr: np.ndarray) -> np.ndarray: + arr = np.maximum(arr, 0) + hi = float(np.quantile(arr, 0.985)) if arr.size else 1.0 + if hi <= 0: + hi = 1.0 + return np.clip(arr / hi, 0, 1) + + +def smooth_map(grid: np.ndarray, radius: float = 1.45) -> np.ndarray: + g = normalize_positive(grid) + img = Image.fromarray(np.uint8(g * 255), mode="L").filter(ImageFilter.GaussianBlur(radius=radius)) + return np.asarray(img, dtype=np.float32) / 255.0 + + +def answer_bbox(indices: list[int], grid: int, image_w: int, image_h: int) -> tuple[int, int, int, int]: + rows = [idx // grid for idx in indices] + cols = [idx % grid for idx in indices] + x0 = int(min(cols) / grid * image_w) + x1 = int((max(cols) + 1) / grid * image_w) + y0 = int(min(rows) / grid * image_h) + y1 = int((max(rows) + 1) / grid * image_h) + return x0, y0, x1, y1 + + +def overlay_heat(base: Image.Image, heat: np.ndarray, bbox: tuple[int, int, int, int], theme: str) -> Image.Image: + base_rgba = base.convert("RGBA") + heat_img = Image.new("RGBA", base.size, (0, 0, 0, 0)) + # Upscale smoothed 56x56 field to bitmap size; threshold softens static. + up = Image.fromarray(np.uint8(heat * 255), mode="L").resize(base.size, Image.Resampling.BICUBIC) + vals = np.asarray(up, dtype=np.float32) / 255.0 + threshold = float(np.quantile(vals, 0.72)) + vals = np.clip((vals - threshold) / max(1e-6, 1.0 - threshold), 0, 1) + px = heat_img.load() + for y in range(0, heat_img.height, 2): + for x in range(0, heat_img.width, 2): + t = float(vals[y, x]) + if t <= 0: + continue + color = heat_color(t, theme) + px[x, y] = color + if x + 1 < heat_img.width: + px[x + 1, y] = color + if y + 1 < heat_img.height: + px[x, y + 1] = color + if x + 1 < heat_img.width and y + 1 < heat_img.height: + px[x + 1, y + 1] = color + out = Image.alpha_composite(base_rgba, heat_img).convert("RGB") + draw = ImageDraw.Draw(out) + color = PALETTE["cyan"] if theme == "cyan" else PALETTE["orange"] + draw.rounded_rectangle(bbox, radius=4, outline=color, width=6) + # spotlight ring around answer bbox + x0, y0, x1, y1 = bbox + pad = 24 + draw.rounded_rectangle((x0 - pad, y0 - pad, x1 + pad, y1 + pad), radius=16, outline=color, width=3) + return out + + +def crop_box(img: Image.Image, bbox: tuple[int, int, int, int], pad: int = 180) -> Image.Image: + x0, y0, x1, y1 = bbox + return img.crop((max(0, x0 - pad), max(0, y0 - pad), min(img.width, x1 + pad), min(img.height, y1 + pad))).convert("RGB") + + +def paste_fit(canvas: Image.Image, img: Image.Image, box: tuple[int, int, int, int], resample: int = Image.Resampling.LANCZOS) -> None: + x0, y0, x1, y1 = box + scale = min((x1 - x0) / img.width, (y1 - y0) / img.height) + resized = img.resize((max(1, round(img.width * scale)), max(1, round(img.height * scale))), resample) + canvas.paste(resized, (x0 + (x1 - x0 - resized.width) // 2, y0 + (y1 - y0 - resized.height) // 2)) + + +def region_score(grid_map: np.ndarray, indices: list[int]) -> float: + if not indices: + return 0.0 + flat = grid_map.ravel() + return float(np.mean([flat[i] for i in indices if i < len(flat)])) + + +def random_region_scores(grid_map: np.ndarray, region_size: int, count: int = 600, seed: int = 7) -> np.ndarray: + rng = random.Random(seed) + flat = grid_map.ravel() + scores = [] + for _ in range(count): + picks = rng.sample(range(len(flat)), min(region_size, len(flat))) + scores.append(float(np.mean(flat[picks]))) + return np.array(scores, dtype=np.float32) + + +def draw_score_card(draw: ImageDraw.ImageDraw, box: tuple[int, int, int, int], title: str, score: float, random_scores: np.ndarray, color: tuple[int, int, int]) -> None: + x0, y0, x1, y1 = box + draw.rounded_rectangle(box, radius=18, fill=PALETTE["panel2"], outline=(34, 48, 58), width=1) + draw.text((x0 + 20, y0 + 18), title, fill=color, font=ui_font(22, True)) + percentile = float((random_scores < score).mean() * 100) + draw.text((x0 + 20, y0 + 50), f"answer region beats {percentile:.0f}% of random same-size regions", fill=PALETTE["muted"], font=ui_font(16)) + gx0, gy0, gx1, gy1 = x0 + 28, y0 + 96, x1 - 28, y1 - 42 + lo = float(min(random_scores.min(), score)) + hi = float(max(random_scores.max(), score)) + if hi <= lo: + hi = lo + 1e-6 + bins = np.linspace(lo, hi, 30) + hist, _ = np.histogram(random_scores, bins=bins) + max_h = max(1, int(hist.max())) + bw = (gx1 - gx0) / len(hist) + for i, h in enumerate(hist): + xa = gx0 + i * bw + xb = gx0 + (i + 1) * bw - 1 + ya = gy1 - (gy1 - gy0) * int(h) / max_h + draw.rectangle((round(xa), round(ya), round(xb), gy1), fill=(37, 49, 58)) + sx = gx0 + (gx1 - gx0) * (score - lo) / (hi - lo) + draw.line((sx, gy0 - 8, sx, gy1 + 10), fill=color, width=5) + draw.text((round(sx) - 34, gy0 - 36), "answer", fill=color, font=ui_font(15, True)) + + +def draw_generation_rows(draw: ImageDraw.ImageDraw, box: tuple[int, int, int, int], generations: dict[str, str]) -> None: + x0, y0, x1, y1 = box + draw.rounded_rectangle(box, radius=24, fill=PALETTE["panel"], outline=(35, 49, 59), width=1) + draw.text((x0 + 28, y0 + 24), "causal patch check", fill=PALETTE["ink"], font=ui_font(30, True)) + draw.text((x0 + 28, y0 + 60), "Patch before decoder layer 0; only the true answer-region patch changes the answer.", fill=PALETTE["muted"], font=ui_font(18)) + rows = [ + ("normal", generations["normal"], PALETTE["green"]), + ("random region patch", generations["random_mean_patch"], PALETTE["cyan"]), + ("answer region patch", generations["answer_mean_patch"], PALETTE["red"]), + ("all image tokens zero", generations["all_image_zero"] or "∅", PALETTE["purple"]), + ] + y = y0 + 116 + for label, text, color in rows: + draw.rounded_rectangle((x0 + 28, y, x1 - 28, y + 62), radius=14, fill=PALETTE["panel2"], outline=(34, 48, 58), width=1) + draw.text((x0 + 48, y + 17), label.upper(), fill=color, font=ui_font(15, True)) + draw.text((x0 + 330, y + 13), text[:70], fill=PALETTE["ink"], font=ui_font(24, True)) + y += 78 + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--result-dir", default=str(HERE / "results" / "qwen-control-intervention-q3-d12-prehook")) + ap.add_argument("--out", default=str(HERE / "results" / "qwen-control-intervention-q3-d12-prehook" / "spotlight-control.png")) + args = ap.parse_args() + + result_dir = Path(args.result_dir) + summary = json.loads((result_dir / "summary.json").read_text()) + data = np.load(result_dir / "control_intervention.npz") + img = Image.open(result_dir / "images" / "image-carrier.png").convert("RGB") + + primary = summary["primary"] + distractor = summary["distractor"] + primary_meta = summary["primary_meta"] + distractor_meta = summary["distractor_meta"] + grid = primary_meta["image_grid"] + primary_layer = primary_meta["peak_layer"] + distractor_layer = distractor_meta["peak_layer"] + primary_norm = data["primary_norm"][primary_layer].reshape(grid, grid) + distractor_norm = data["distractor_norm"][distractor_layer].reshape(grid, grid) + + # Prompt-specific contrast removes shared visual texture: what lights up more + # for this question than for the other question? + primary_contrast = smooth_map(primary_norm - distractor_norm) + distractor_contrast = smooth_map(distractor_norm - primary_norm) + primary_bbox = answer_bbox(primary_meta["answer_indices"], grid, img.width, img.height) + distractor_bbox = answer_bbox(distractor_meta["answer_indices"], grid, img.width, img.height) + primary_overlay = overlay_heat(img, primary_contrast, primary_bbox, "orange") + distractor_overlay = overlay_heat(img, distractor_contrast, distractor_bbox, "cyan") + + primary_score = region_score(primary_contrast, primary_meta["answer_indices"]) + distractor_score = region_score(distractor_contrast, distractor_meta["answer_indices"]) + primary_random = random_region_scores(primary_contrast, len(primary_meta["answer_indices"]), seed=11) + distractor_random = random_region_scores(distractor_contrast, len(distractor_meta["answer_indices"]), seed=13) + + w, h = 2200, 1320 + canvas = Image.new("RGB", (w, h), PALETTE["bg"]) + draw = ImageDraw.Draw(canvas) + for y in range(0, h, 16): + draw.line((0, y, w, y), fill=(7, 10 + y % 9, 15 + y % 11)) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-240, -200, 900, 700), fill=(255, 112, 72, 28)) + gd.ellipse((1160, 120, 2460, 1340), fill=(75, 220, 255, 25)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(86))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + draw.text((64, 42), "QWEN SNAPCOMPACT SPOTLIGHT", fill=PALETTE["amber"], font=ui_font(24, True)) + draw.text((64, 84), "Subtract the other question; the signal stops looking like static", fill=PALETTE["ink"], font=ui_font(58, True)) + draw.text((66, 160), "These are not raw activation carpets. Each overlay is prompt-specific excess: this question’s map minus the other question’s map, smoothed and thresholded.", fill=PALETTE["muted"], font=ui_font(23)) + + # Overlay panels. + draw.rounded_rectangle((64, 230, 1068, 794), radius=28, fill=PALETTE["panel"], outline=(35, 49, 59), width=1) + draw.text((96, 260), "primary prompt spotlight", fill=PALETTE["orange"], font=ui_font(31, True)) + draw.text((96, 298), f"{primary['q']} → {primary['answer_text']}", fill=PALETTE["muted"], font=ui_font(18)) + primary_crop = crop_box(primary_overlay, primary_bbox, pad=300) + paste_fit(canvas, primary_crop, (96, 342, 694, 760), Image.Resampling.LANCZOS) + draw.rounded_rectangle((720, 342, 1036, 760), radius=18, fill=(244, 242, 230), outline=PALETTE["orange"], width=3) + paste_fit(canvas, crop_box(primary_overlay, primary_bbox, pad=90), (736, 358, 1020, 744), Image.Resampling.LANCZOS) + draw.text((736, 724), "zoom: answer region", fill=PALETTE["orange"], font=ui_font(15, True)) + + draw.rounded_rectangle((1132, 230, 2136, 794), radius=28, fill=PALETTE["panel"], outline=(35, 49, 59), width=1) + draw.text((1164, 260), "distractor prompt spotlight", fill=PALETTE["cyan"], font=ui_font(31, True)) + draw.text((1164, 298), f"{distractor['q']} → {distractor['answer_text']}", fill=PALETTE["muted"], font=ui_font(18)) + distractor_crop = crop_box(distractor_overlay, distractor_bbox, pad=300) + paste_fit(canvas, distractor_crop, (1164, 342, 1762, 760), Image.Resampling.LANCZOS) + draw.rounded_rectangle((1788, 342, 2104, 760), radius=18, fill=(244, 242, 230), outline=PALETTE["cyan"], width=3) + paste_fit(canvas, crop_box(distractor_overlay, distractor_bbox, pad=90), (1804, 358, 2088, 744), Image.Resampling.LANCZOS) + draw.text((1804, 724), "zoom: answer region", fill=PALETTE["cyan"], font=ui_font(15, True)) + + draw_score_card(draw, (64, 836, 610, 1236), "primary answer-region score", primary_score, primary_random, PALETTE["orange"]) + draw_score_card(draw, (642, 836, 1188, 1236), "distractor answer-region score", distractor_score, distractor_random, PALETTE["cyan"]) + draw_generation_rows(draw, (1220, 836, 2136, 1236), summary["generations"]) + + # Save metrics alongside figure for the caption. + metrics = { + "primary_score": primary_score, + "primary_percentile": float((primary_random < primary_score).mean() * 100), + "distractor_score": distractor_score, + "distractor_percentile": float((distractor_random < distractor_score).mean() * 100), + "primary_peak_layer": primary_layer, + "distractor_peak_layer": distractor_layer, + "generations": summary["generations"], + } + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + canvas.save(out) + (out.parent / "spotlight-metrics.json").write_text(json.dumps(metrics, indent=1)) + print(out) + print(json.dumps(metrics, indent=1)) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_r2_chord.py b/packages/snapcompact/research/snapcompact_r2_chord.py new file mode 100755 index 000000000..0c687c034 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_r2_chord.py @@ -0,0 +1,265 @@ +#!/usr/bin/env python3 +"""Chord/ribbon diagram of cross-carrier matching in Qwen2.5-VL-7B. + +Left arc: 12 questions answered via the TEXT carrier (cyan). +Right arc: the same 12 questions answered via the IMAGE carrier (orange). +Ribbons between every (text_i, image_j) pair are sized by the REAL cosine +similarity of the carrier-centered hidden states at layer 19 (negatives +clipped to 0). Matched pairs (i == j) glow amber and visibly dominate. + +Data: results/qwen-carrier-convergence-n12/{carrier_convergence.npz,summary.json} +Output: results/agent-r2-chord/chord.png (~2200 px wide) +""" + +import json +import os + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +from matplotlib.path import Path as MplPath +from matplotlib.patches import PathPatch + +HERE = os.path.dirname(os.path.abspath(__file__)) +SRC = os.path.join(HERE, "results", "qwen-carrier-convergence-n12") +OUT_DIR = os.path.join(HERE, "results", "agent-r2-chord") +LAYER = 19 # best layer per summary.json + +# ---------------------------------------------------------------- palette +BG = "#05070a" +PANEL = "#0c1117" +INK = "#f1efe0" +MUTED = "#8f9aa0" +AMBER = "#ffc444" +CYAN = "#4bdcff" +ORANGE = "#ff7048" +GREEN = "#94ff75" + +# ---------------------------------------------------------------- data +npz = np.load(os.path.join(SRC, "carrier_convergence.npz")) +cross = np.asarray(npz["cross_sim"][LAYER], dtype=np.float64) # [12,12] text x image +with open(os.path.join(SRC, "summary.json")) as fh: + summary = json.load(fh) +records = summary["records"] +best = summary["best"] +n = cross.shape[0] +assert len(records) == n == 12 + +labels = [r["gold"] for r in records] +matched_mean = float(np.trace(cross) / n) +mismatched_mean = float((cross.sum() - np.trace(cross)) / (n * n - n)) +# retrieval: for each text row, is the matched image column the argmax? +retrieved = int((cross.argmax(axis=1) == np.arange(n)).sum()) + +w = np.clip(cross, 0.0, None) # ribbon weights: clip negatives + +# ---------------------------------------------------------------- geometry +# Left arc (text carrier): 110deg -> 250deg, top-left to bottom-left. +# Right arc (image carrier): 70deg -> -70deg, mirrored so matched ribbons +# run roughly horizontally across the circle. +R_IN = 0.955 # inner radius where ribbons attach +R_OUT = 1.000 # outer radius of the node band +SEG_DEG = 8.6 # angular width of each node segment +ts = np.linspace(0.0, 1.0, n) +left_centers = 110.0 + ts * 140.0 +right_centers = 70.0 - ts * 140.0 + + +def seg_bounds(center_deg): + return center_deg - SEG_DEG / 2.0, center_deg + SEG_DEG / 2.0 + + +def pol(theta_deg, r): + a = np.deg2rad(theta_deg) + return np.array([r * np.cos(a), r * np.sin(a)]) + + +def allocate(centers, weights_per_node): + """Split each node's segment into sub-spans proportional to ribbon weight. + + weights_per_node: [n, n] -- weights_per_node[i, j] is the weight of the + ribbon to opposite-side node j, allocated within node i's segment. + Sub-spans are ordered by the opposite node index so ribbons fan smoothly. + Returns spans[i][j] = (a0, a1) in degrees (a0 < a1) or None if weight ~ 0. + """ + spans = [] + for i in range(n): + lo, hi = seg_bounds(centers[i]) + tot = weights_per_node[i].sum() + spans_i = [None] * n + if tot <= 1e-9: + spans.append(spans_i) + continue + cursor = lo + for j in range(n): + frac = weights_per_node[i, j] / tot + width = frac * (hi - lo) + if weights_per_node[i, j] > 1e-9: + spans_i[j] = (cursor, cursor + width) + cursor += width + spans.append(spans_i) + return spans + + +# Left node i sends ribbons to right nodes j with weight w[i, j]; +# right node j receives from left nodes i with weight w[i, j]. +left_spans = allocate(left_centers, w) +right_spans = allocate(right_centers, w.T) + + +def arc_points(a0, a1, r, steps=12): + th = np.linspace(a0, a1, steps) + return np.stack([r * np.cos(np.deg2rad(th)), r * np.sin(np.deg2rad(th))], axis=1) + + +def ribbon_path(la, lb, ra, rb, pull=0.18): + """Filled ribbon: arc(la->lb) on the left rim, cubic bezier to the right + rim, arc(ra->rb), bezier back. Control points pulled toward the center.""" + p_lb = pol(lb, R_IN) + p_ra = pol(ra, R_IN) + p_rb = pol(rb, R_IN) + p_la = pol(la, R_IN) + verts = [] + codes = [] + arc1 = arc_points(la, lb, R_IN) + verts.extend(arc1) + codes.extend([MplPath.MOVETO] + [MplPath.LINETO] * (len(arc1) - 1)) + # bezier left-edge-end -> right-edge-start + verts.extend([p_lb * pull, p_ra * pull, p_ra]) + codes.extend([MplPath.CURVE4] * 3) + arc2 = arc_points(ra, rb, R_IN)[1:] + verts.extend(arc2) + codes.extend([MplPath.LINETO] * len(arc2)) + verts.extend([p_rb * pull, p_la * pull, p_la]) + codes.extend([MplPath.CURVE4] * 3) + codes.append(MplPath.CLOSEPOLY) + verts.append(p_la) + return MplPath(verts, codes) + + +def center_bezier(a_deg, b_deg, pull=0.18, steps=60): + p0, p3 = pol(a_deg, R_IN), pol(b_deg, R_IN) + p1, p2 = p0 * pull, p3 * pull + t = np.linspace(0, 1, steps)[:, None] + return ((1 - t) ** 3 * p0 + 3 * (1 - t) ** 2 * t * p1 + + 3 * (1 - t) * t ** 2 * p2 + t ** 3 * p3) + + +# ---------------------------------------------------------------- figure +fig = plt.figure(figsize=(22, 16.5), dpi=100, facecolor=BG) +ax = fig.add_axes([0.02, 0.0, 0.96, 0.94]) +ax.set_facecolor(BG) +ax.set_xlim(-1.95, 1.95) +ax.set_ylim(-1.32, 1.30) +ax.set_aspect("equal") +ax.axis("off") + +w_max = w.max() + +# mismatched ribbons first (thin, dim), then matched (amber, glowing) on top +order = sorted(((i, j) for i in range(n) for j in range(n)), + key=lambda ij: (ij[0] == ij[1], w[ij])) +for i, j in order: + ls, rs = left_spans[i][j], right_spans[j][i] + if ls is None or rs is None: + continue + val = w[i, j] + matched = i == j + # right span traversed in reverse so the ribbon doesn't twist + path = ribbon_path(ls[0], ls[1], rs[1], rs[0]) + if matched: + # glow: layered strokes along the centerline under the fill + mid_l = 0.5 * (ls[0] + ls[1]) + mid_r = 0.5 * (rs[0] + rs[1]) + spine = center_bezier(mid_l, mid_r) + for lw, al in ((26, 0.045), (14, 0.075), (7, 0.12)): + ax.plot(spine[:, 0], spine[:, 1], color=AMBER, lw=lw, alpha=al, + solid_capstyle="round", zorder=4) + ax.add_patch(PathPatch(path, facecolor=AMBER, edgecolor=AMBER, + lw=0.7, alpha=0.78, zorder=5)) + else: + alpha = 0.10 + 0.45 * (val / w_max) + ax.add_patch(PathPatch(path, facecolor=MUTED, edgecolor="none", + alpha=alpha * 0.55, zorder=2)) + +# ---------------------------------------------------------------- node bands +for i in range(n): + for centers, color, side in ((left_centers, CYAN, "L"), + (right_centers, ORANGE, "R")): + a0, a1 = seg_bounds(centers[i]) + band = arc_points(a0, a1, R_OUT, 16) + band_in = arc_points(a1, a0, R_IN, 16) + poly = np.vstack([band, band_in]) + ax.add_patch(plt.Polygon(poly, closed=True, facecolor=color, + edgecolor="none", alpha=0.95, zorder=6)) + +# ---------------------------------------------------------------- labels +for i in range(n): + txt = labels[i] + for centers, color, ha in ((left_centers, CYAN, "right"), + (right_centers, ORANGE, "left")): + c = centers[i] + p = pol(c, 1.05) + ax.text(p[0], p[1], txt, color=INK, fontsize=15.5, ha=ha, va="center", + zorder=8, family="DejaVu Sans") + # small question index tick just inside the label + ax.text(p[0] + (0.018 if ha == "left" else -0.018), + p[1] - 0.052, f"Q{i + 1}", color=color, fontsize=10.5, + ha=ha, va="center", alpha=0.85, zorder=8) + +# arc side headers +ax.text(*pol(180, 1.62), "TEXT CARRIER", color=CYAN, fontsize=21, + ha="center", va="center", rotation=90, weight="bold", alpha=0.95) +ax.text(*pol(180, 1.69), "5,219 prose tokens", color=MUTED, fontsize=13, + ha="center", va="center", rotation=90) +ax.text(*pol(0, 1.62), "IMAGE CARRIER", color=ORANGE, fontsize=21, + ha="center", va="center", rotation=-90, weight="bold", alpha=0.95) +ax.text(*pol(0, 1.69), "same passage, rendered as pixels", color=MUTED, + fontsize=13, ha="center", va="center", rotation=-90) + +# ---------------------------------------------------------------- titles +fig.text(0.5, 0.965, "ONE MEMORY, TWO CARRIERS", color=INK, fontsize=34, + ha="center", va="center", weight="bold", family="DejaVu Sans") +fig.text(0.5, 0.932, + "Cross-carrier cosine of answer states at layer 19 -- every text" + " question finds its image twin (Qwen2.5-VL-7B)", + color=MUTED, fontsize=16.5, ha="center", va="center") + +# ---------------------------------------------------------------- stat plate +plate = fig.add_axes([0.035, 0.05, 0.215, 0.135]) +plate.set_facecolor(PANEL) +plate.set_xlim(0, 1) +plate.set_ylim(0, 1) +for s in plate.spines.values(): + s.set_color("#1c242e") +plate.set_xticks([]) +plate.set_yticks([]) +plate.text(0.5, 0.84, f"LAYER {LAYER} -- BEST SEPARATION", color=MUTED, + fontsize=12.5, ha="center", va="center") +stats = ( + (f"{matched_mean:+.2f}", "matched cosine", AMBER), + (f"{mismatched_mean:+.2f}", "mismatched", MUTED), + (f"{retrieved}/{n}", "retrieval", GREEN), +) +for k, (val, lab, color) in enumerate(stats): + x = 0.18 + 0.32 * k + plate.text(x, 0.48, val, color=color, fontsize=24, ha="center", + va="center", weight="bold") + plate.text(x, 0.18, lab, color=MUTED, fontsize=12, ha="center", + va="center") + +# footnote +fig.text(0.5, 0.012, + "Ribbon width/opacity = cosine(text_state_i, image_state_j)," + " negatives clipped; amber = matched pair (i = j)." + f" RSA r = {best['rsa_pearson']:.2f}.", + color=MUTED, fontsize=12.5, ha="center", va="center") + +os.makedirs(OUT_DIR, exist_ok=True) +out_path = os.path.join(OUT_DIR, "chord.png") +fig.savefig(out_path, dpi=100, facecolor=BG) +print(f"wrote {out_path}") +print(f"matched={matched_mean:.4f} mismatched={mismatched_mean:.4f} " + f"retrieval={retrieved}/{n}") diff --git a/packages/snapcompact/research/snapcompact_r2_crystal.py b/packages/snapcompact/research/snapcompact_r2_crystal.py new file mode 100755 index 000000000..fbe894bfc --- /dev/null +++ b/packages/snapcompact/research/snapcompact_r2_crystal.py @@ -0,0 +1,338 @@ +#!/usr/bin/env python3 +"""Crystallization: a pixel patch becomes a vocabulary token. + +Animated logit-lens GIF for Qwen2.5-VL-7B. Steps through layers L0..L28 for +visual token #310 of the image carrier (the 28x28 patch covering the rendered +word "spectacular"), showing the REAL top-5 decoded vocab tokens per layer and +the probability of the answer BPE 'acular', which stays ~0 until L24 then +climbs 0.14 -> 0.39 at L28. A control patch (token #1878) is shown alongside +and stays noise. All numbers are read from logit_lens.json; nothing is +fabricated. + +Output: results/agent-r2-crystal/crystal.gif (+ crystal_final.png) +""" + +import json +from pathlib import Path + +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +BASE = Path(__file__).resolve().parent +DATA = BASE / "results/qwen-logit-lens-q3/logit_lens.json" +CARRIER = BASE / "results/qwen-logit-lens-q3/images/image-carrier.png" +OUT_DIR = BASE / "results/agent-r2-crystal" + +# ---------------------------------------------------------------- palette +BG = (5, 7, 10) +PANEL = (12, 17, 23) +PANEL_EDGE = (28, 36, 46) +INK = (241, 239, 224) +MUTED = (143, 154, 160) +DIM = (80, 90, 98) +AMBER = (255, 196, 68) +CYAN = (75, 220, 255) +GREEN = (148, 255, 117) + +W, H = 1200, 720 +TOKEN_IDX = 310 +ANSWER_BPE = "acular" +ANSWER_SLOT = 1 # answer_token_p[1] == p('acular') + +# ---------------------------------------------------------------- fonts +MENLO = "/System/Library/Fonts/Menlo.ttc" +UNI = "/System/Library/Fonts/Supplemental/Arial Unicode.ttf" + + +def font(path: str, size: int, index: int = 0) -> ImageFont.FreeTypeFont: + try: + return ImageFont.truetype(path, size, index=index) + except OSError: + return ImageFont.truetype(MENLO, size) + + +F_TITLE = font(MENLO, 30, index=1) +F_SUB = font(MENLO, 15) +F_LABEL = font(MENLO, 13) +F_LABEL_B = font(MENLO, 13, index=1) +F_TINY = font(MENLO, 11) +F_LAYER = font(MENLO, 64, index=1) +F_STAGE = font(MENLO, 16, index=1) +F_NUM = font(MENLO, 14) +F_TOK = font(UNI, 18) +F_TOK_B = font(MENLO, 18, index=1) +F_TOK_S = font(UNI, 14) +F_BADGE = font(MENLO, 17, index=1) + + +def load() -> tuple[dict, list[dict], list[dict]]: + data = json.loads(DATA.read_text()) + assert data["layers"] == 29 and data["image_grid"] == 56 + assert data["answer_token_strs"][ANSWER_SLOT] == ANSWER_BPE + + def by_layer(idx: int) -> list[dict]: + return sorted( + (e for e in data["lens"] if e["token_index"] == idx), + key=lambda e: e["layer"], + ) + + target = by_layer(TOKEN_IDX) + control = by_layer(data["control_indices"][0]) + assert len(target) == 29 and len(control) == 29 + return data, target, control + + +def crop_cell(carrier: Image.Image, idx: int) -> Image.Image: + r, c = idx // 56, idx % 56 + return carrier.crop((c * 28, r * 28, (c + 1) * 28, (r + 1) * 28)) + + +def sanitize(s: str) -> str: + out = "".join(ch if ch.isprintable() else "\ufffd" for ch in s) + return out.replace(" ", "\u2423", 1) if s.startswith(" ") else out + + +def text_w(f: ImageFont.FreeTypeFont, s: str) -> float: + return f.getlength(s) + + +STAGES = [ + (0, "static"), + (8, "gibberish"), + (16, "morphemes"), + (23, "warming up"), + (24, "lock-on"), + (28, "crystallized"), +] + + +def stage_for(layer: int) -> str: + name = STAGES[0][1] + for lo, label in STAGES: + if layer >= lo: + name = label + return name + + +def render_frame( + layer: int, + data: dict, + target: list[dict], + control: list[dict], + patch: Image.Image, + ctrl_patch: Image.Image, + strip: Image.Image, + strip_cell: int, +) -> Image.Image: + e = target[layer] + ce = control[layer] + p_ans = e["answer_token_p"][ANSWER_SLOT] + p_ctrl = ce["answer_token_p"][ANSWER_SLOT] + locked = layer >= 24 + final = layer == 28 + heat = min(1.0, p_ans / 0.40) # glow scales with the real probability + + img = Image.new("RGB", (W, H), BG) + d = ImageDraw.Draw(img) + + # ---------------------------------------------------------- header + d.text((40, 22), "CRYSTALLIZATION", font=F_TITLE, fill=INK) + tw = text_w(F_TITLE, "CRYSTALLIZATION") + d.text( + (40 + tw + 18, 36), + "\u2014 28\u00d728 pixels become the token '" + ANSWER_BPE + "'", + font=F_SUB, + fill=MUTED, + ) + d.text( + (40, 60), + "logit lens on Qwen2.5-VL-7B \u00b7 image carrier 1568px \u00b7 visual token " + f"#{TOKEN_IDX} (grid {e['grid_rc'][0]},{e['grid_rc'][1]}) \u00b7 " + "answer \u201cspectacular\u201d = [spect][acular]", + font=F_SUB, + fill=MUTED, + ) + if final: # celebratory badge in the free top-right corner + badge = f"'{ANSWER_BPE}' \u00b7 p={p_ans:.2f} \u00b7 CRYSTALLIZED" + bw = text_w(F_BADGE, badge) + bx = W - bw - 76 + d.rounded_rectangle((bx, 22, bx + bw + 36, 56), radius=17, fill=(38, 30, 10), outline=AMBER, width=2) + d.text((bx + 18, 29), badge, font=F_BADGE, fill=AMBER) + + # ---------------------------------------------------------- context strip + sx, sy = (W - strip.width) // 2, 96 + img.paste(strip, (sx, sy)) + d.rectangle((sx, sy, sx + strip.width - 1, sy + strip.height - 1), outline=PANEL_EDGE) + col = AMBER if locked else CYAN + d.rectangle( + (sx + strip_cell, sy, sx + strip_cell + strip.height, sy + strip.height - 1), + outline=col, + width=3, + ) + cap = "carrier row 5 \u00b7 the model never sees glyphs \u2014 only these pixels" + d.text(((W - text_w(F_TINY, cap)) / 2, sy + strip.height + 6), cap, font=F_TINY, fill=DIM) + + # ---------------------------------------------------------- main row + top_y = 232 + # patch panel ------------------------------------------------------ + px, py, ps = 40, top_y, 280 + img.paste(patch.resize((ps, ps), Image.NEAREST), (px, py)) + d.rectangle((px - 1, py - 1, px + ps, py + ps), outline=col, width=2) + d.text((px, py + ps + 10), f"visual token #{TOKEN_IDX}", font=F_LABEL_B, fill=INK) + d.text((px, py + ps + 28), "28\u00d728 px \u00b7 reads: \u2018\"sp\u2019 / \u2018and\u2019", font=F_TINY, fill=MUTED) + + # layer counter ---------------------------------------------------- + cx = 392 + d.text((cx, top_y - 2), "LAYER", font=F_LABEL, fill=MUTED) + num = f"{layer:02d}" + d.text((cx, top_y + 16), num, font=F_LAYER, fill=AMBER if locked else INK) + d.text((cx + text_w(F_LAYER, num) + 8, top_y + 58), "/28", font=F_STAGE, fill=DIM) + stage = stage_for(layer) + d.text((cx, top_y + 98), stage.upper(), font=F_STAGE, fill=GREEN if final else (AMBER if locked else MUTED)) + ry = top_y + 134 # mini rail of 29 ticks + for i in range(29): + tx = cx + i * 5 + d.rectangle((tx, ry, tx + 3, ry + 10), fill=AMBER if i <= layer else (40, 48, 58)) + d.text((cx, ry + 18), f"p('{ANSWER_BPE}') = {p_ans:.4f}", font=F_NUM, fill=AMBER if p_ans > 0.01 else DIM) + + # top-5 panel ------------------------------------------------------ + tx0, ty0, tx1, ty1 = 580, top_y - 12, 1160, top_y + 318 + d.rounded_rectangle((tx0, ty0, tx1, ty1), radius=8, fill=PANEL, outline=PANEL_EDGE) + d.text((tx0 + 18, ty0 + 12), "TOP-5 DECODED VOCAB TOKENS \u00b7 what this patch \u201cmeans\u201d so far", font=F_LABEL, fill=MUTED) + bar_x = tx0 + 230 + bar_max = tx1 - bar_x - 86 + scale = 0.45 # fixed probability scale across all frames + for i, t in enumerate(e["top"]): + yy = ty0 + 48 + i * 54 + is_ans = t["id"] in data["answer_token_ids"] + tok_s = sanitize(t["str"]) + if len(tok_s) > 16: + tok_s = tok_s[:15] + "\u2026" + d.text((tx0 + 18, yy), f"'{tok_s}'", font=F_TOK_B if is_ans else F_TOK, fill=AMBER if is_ans else INK) + bw = max(2, int(min(t["p"] / scale, 1.0) * bar_max)) + d.rectangle((bar_x, yy + 4, bar_x + bw, yy + 18), fill=AMBER if is_ans else (58, 70, 84)) + if is_ans and heat > 0.3: + d.rectangle((bar_x, yy + 4, bar_x + bw, yy + 18), outline=INK) + d.text((bar_x + bw + 10, yy + 3), f"{t['p']:.3f}", font=F_NUM, fill=AMBER if is_ans else MUTED) + d.text((tx0 + 18, yy + 24), f"id {t['id']}", font=F_TINY, fill=DIM) + + # ---------------------------------------------------------- bottom row + by0, by1 = 596, 708 + # confidence meter for 'acular' + mx0, mx1 = 40, 730 + d.rounded_rectangle((mx0, by0, mx1, by1), radius=8, fill=PANEL, outline=PANEL_EDGE) + d.text((mx0 + 16, by0 + 8), f"CONFIDENCE \u00b7 p('{ANSWER_BPE}') across layers", font=F_LABEL, fill=MUTED) + leg_x = mx1 - 130 + d.rectangle((leg_x, by0 + 12, leg_x + 14, by0 + 15), fill=AMBER) + d.text((leg_x + 20, by0 + 6), "answer", font=F_TINY, fill=AMBER) + d.rectangle((leg_x + 74, by0 + 12, leg_x + 88, by0 + 15), fill=(60, 70, 80)) + d.text((leg_x + 94, by0 + 6), "ctrl", font=F_TINY, fill=(96, 108, 118)) + ch_x0, ch_x1 = mx0 + 52, mx1 - 64 + ch_y0, ch_y1 = by0 + 32, by1 - 22 + p_max = 0.45 + d.line((ch_x0, ch_y1, ch_x1, ch_y1), fill=PANEL_EDGE) + for gv in (0.2, 0.4): + gy = ch_y1 - gv / p_max * (ch_y1 - ch_y0) + d.line((ch_x0, gy, ch_x1, gy), fill=(22, 28, 36)) + d.text((mx0 + 16, gy - 6), f"{gv:.1f}", font=F_TINY, fill=DIM) + + def xs(l: int) -> float: + return ch_x0 + l / 28 * (ch_x1 - ch_x0) + + def ys(p: float) -> float: + return ch_y1 - min(p, p_max) / p_max * (ch_y1 - ch_y0) + + pts = [(xs(l), ys(target[l]["answer_token_p"][ANSWER_SLOT])) for l in range(layer + 1)] + cpts = [(xs(l), ys(control[l]["answer_token_p"][ANSWER_SLOT])) for l in range(layer + 1)] + if len(cpts) > 1: + d.line(cpts, fill=(60, 70, 80), width=2) + if len(pts) > 1: + poly = pts + [(pts[-1][0], ch_y1), (pts[0][0], ch_y1)] + d.polygon(poly, fill=(76, 56, 16) if locked else (46, 36, 14)) + d.line(pts, fill=AMBER, width=3) + hx, hy = pts[-1] + d.ellipse((hx - 5, hy - 5, hx + 5, hy + 5), fill=AMBER if p_ans > 0.01 else MUTED) + head = f"{p_ans:.2f}" if p_ans >= 0.005 else f"{p_ans:.4f}" + d.text((min(hx + 8, ch_x1 - 8), hy - 18), head, font=F_NUM, fill=AMBER if p_ans > 0.01 else MUTED) + for ml in (24, 28): + if layer >= ml: + mlx = xs(ml) + d.line((mlx, ch_y1, mlx, ys(target[ml]["answer_token_p"][ANSWER_SLOT])), fill=(90, 72, 30)) + d.text((mlx - 10, ch_y1 + 6), f"L{ml}", font=F_TINY, fill=AMBER) + d.text((ch_x0, ch_y1 + 6), "L0", font=F_TINY, fill=DIM) + + # control panel ---------------------------------------------------- + kx0, kx1 = 760, 1160 + d.rounded_rectangle((kx0, by0, kx1, by1), radius=8, fill=PANEL, outline=PANEL_EDGE) + d.text((kx0 + 16, by0 + 8), "CONTROL \u00b7 token #" + str(ce["token_index"]), font=F_LABEL, fill=MUTED) + cps = 60 + img.paste(ctrl_patch.resize((cps, cps), Image.NEAREST), (kx0 + 16, by0 + 30)) + d.rectangle((kx0 + 15, by0 + 29, kx0 + 16 + cps, by0 + 30 + cps), outline=PANEL_EDGE) + ct = ce["top"][0] + ct_s = sanitize(ct["str"]) + if len(ct_s) > 12: + ct_s = ct_s[:11] + "\u2026" + lx = kx0 + 16 + cps + 14 + d.text((lx, by0 + 30), "top-1: ", font=F_NUM, fill=INK) + tx = lx + text_w(F_NUM, "top-1: ") + d.text((tx, by0 + 30), f"'{ct_s}'", font=F_TOK_S, fill=INK) + d.text((tx + text_w(F_TOK_S, f"'{ct_s}'") + 10, by0 + 30), f"{ct['p']:.3f}", font=F_NUM, fill=INK) + d.text((lx, by0 + 52), f"p('{ANSWER_BPE}') = {p_ctrl:.5f}", font=F_NUM, fill=MUTED) + d.text((lx, by0 + 74), "still noise \u2713" if p_ctrl < 0.01 else "?!", font=F_LABEL_B, fill=GREEN) + d.text((kx0 + 224, by0 + 8), "never converges to the answer", font=F_TINY, fill=DIM) + + # ---------------------------------------------------------- glow + if locked and heat > 0: + glow = Image.new("RGB", (W, H), (0, 0, 0)) + gd = ImageDraw.Draw(glow) + a = int(70 + 110 * heat) + gd.rectangle((px - 6, py - 6, px + ps + 5, py + ps + 5), outline=(a, int(a * 0.77), int(a * 0.27)), width=10) + if final: + gd.rectangle((px - 14, py - 14, px + ps + 13, py + ps + 13), outline=(a, int(a * 0.77), int(a * 0.27)), width=8) + glow = glow.filter(ImageFilter.GaussianBlur(12 if final else 8)) + img = Image.composite(Image.new("RGB", (W, H), AMBER), img, glow.convert("L").point(lambda v: min(v, 140))) + return img + + +def main() -> None: + OUT_DIR.mkdir(parents=True, exist_ok=True) + data, target, control = load() + carrier = Image.open(CARRIER).convert("RGB") + if carrier.size != (1568, 1568): + carrier = carrier.resize((1568, 1568), Image.LANCZOS) + patch = crop_cell(carrier, TOKEN_IDX) + ctrl_patch = crop_cell(carrier, data["control_indices"][0]) + # context strip: row 5, cols 26..39 (14 cells), scaled x3 -> 1176x84 + c0, c1 = 26, 40 + raw = carrier.crop((c0 * 28, 5 * 28, c1 * 28, 6 * 28)) + strip = raw.resize((raw.width * 3, raw.height * 3), Image.NEAREST) + strip_cell = (TOKEN_IDX % 56 - c0) * 84 + + frames, durations = [], [] + for layer in range(29): + fr = render_frame(layer, data, target, control, patch, ctrl_patch, strip, strip_cell) + frames.append(fr.quantize(colors=256, method=Image.MEDIANCUT, dither=Image.Dither.NONE)) + if layer < 23: + durations.append(220) + elif layer < 28: + durations.append(520) + else: + durations.append(2500) # celebratory hold on the final frame + + gif = OUT_DIR / "crystal.gif" + frames[0].save( + gif, + save_all=True, + append_images=frames[1:], + duration=durations, + loop=0, + disposal=2, + optimize=False, + ) + final_png = OUT_DIR / "crystal_final.png" + render_frame(28, data, target, control, patch, ctrl_patch, strip, strip_cell).save(final_png) + print(f"wrote {gif} ({gif.stat().st_size / 1024:.0f} KB, {len(frames)} frames)") + print(f"wrote {final_png}") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_r2_filmstrip.py b/packages/snapcompact/research/snapcompact_r2_filmstrip.py new file mode 100755 index 000000000..5f88b31fc --- /dev/null +++ b/packages/snapcompact/research/snapcompact_r2_filmstrip.py @@ -0,0 +1,256 @@ +#!/usr/bin/env python3 +"""Twin reels: a filmstrip poster of text-carrier vs image-carrier similarity geometry. + +Renders ~7 layers of the carrier-centered 12x12 cosine matrices (`text_sim`, +`image_sim` from carrier_convergence.npz) as paired frames on two parallel +film reels, with the REAL per-layer RSA Pearson from summary.json as a match +meter under each frame, plus the cross-carrier matched cosine as a secondary +tick. Closes with a callout frame for the best layer (RSA 0.85 @ L19). + +Output: results/agent-r2-filmstrip/filmstrip.png (2200 px wide). +""" + +import json +import os + +import matplotlib + +matplotlib.use("Agg") + +import matplotlib.pyplot as plt +import numpy as np +from matplotlib.colors import LinearSegmentedColormap +from matplotlib.patches import FancyBboxPatch, Rectangle + +HERE = os.path.dirname(os.path.abspath(__file__)) +DATA_DIR = os.path.join(HERE, "results", "qwen-carrier-convergence-n12") +OUT_DIR = os.path.join(HERE, "results", "agent-r2-filmstrip") + +# ---------------------------------------------------------------- palette +BG = "#05070a" +PANEL = "#0c1117" +FILM = "#0a0e14" +INK = "#f1efe0" +MUTED = "#8f9aa0" +AMBER = "#ffc444" +CYAN = "#4bdcff" +ORANGE = "#ff7048" +GREEN = "#94ff75" +EDGE = "#1d2630" + +DIVERGING = LinearSegmentedColormap.from_list( + "carrier_div", + [(0.0, CYAN), (0.30, "#16384a"), (0.50, "#0b1016"), (0.72, "#5c2c18"), (0.90, ORANGE), (1.0, AMBER)], +) + +LAYERS = [1, 5, 9, 13, 17, 19, 28] + +# ---------------------------------------------------------------- data +npz = np.load(os.path.join(DATA_DIR, "carrier_convergence.npz")) +text_sim = npz["text_sim"] # [29, 12, 12] +image_sim = npz["image_sim"] # [29, 12, 12] +assert text_sim.shape == image_sim.shape == (29, 12, 12) + +with open(os.path.join(DATA_DIR, "summary.json")) as fh: + summary = json.load(fh) +per_layer = {row["layer"]: row for row in summary["per_layer"]} +best = summary["best"] # layer 19: rsa 0.85, matched 0.66, mismatched -0.06 + +# ---------------------------------------------------------------- layout (pixel space) +W, H = 2200, 1000 +fig = plt.figure(figsize=(W / 100, H / 100), dpi=100) +fig.patch.set_facecolor(BG) +ax = fig.add_axes([0, 0, 1, 1]) +ax.set_xlim(0, W) +ax.set_ylim(H, 0) # y grows downward +ax.axis("off") +ax.set_facecolor(BG) + +X0, X1 = 150, 2136 +N_COLS = len(LAYERS) + 1 # 7 layer frames + closing callout +CW = (X1 - X0) / N_COLS +FS = 186 # matrix tile size + +TEXT_BAND_Y, BAND_H = 192, 262 +IMAGE_BAND_Y = TEXT_BAND_Y + BAND_H + 26 +METER_Y = IMAGE_BAND_Y + BAND_H + 36 +METER_H = 64 +FOOT_Y = METER_Y + METER_H + 52 + + +def col_cx(i: int) -> float: + return X0 + (i + 0.5) * CW + + +def sprockets(y: float, x_start: float, x_end: float) -> None: + x = x_start + 14 + while x + 20 < x_end: + ax.add_patch( + FancyBboxPatch( + (x, y), 20, 13, + boxstyle="round,pad=0,rounding_size=4", + facecolor=BG, edgecolor="#27313d", linewidth=1.0, zorder=6, + ) + ) + x += 49 + + +def film_band(y0: float, x_end: float) -> None: + ax.add_patch( + Rectangle((X0 - 26, y0), x_end - X0 + 26, BAND_H, + facecolor=FILM, edgecolor=EDGE, linewidth=1.2, zorder=2) + ) + sprockets(y0 + 11, X0 - 26, x_end) + sprockets(y0 + BAND_H - 24, X0 - 26, x_end) + + +BAND_X_END = X0 + (N_COLS - 1) * CW - 14 # bands stop before the callout column +film_band(TEXT_BAND_Y, BAND_X_END) +film_band(IMAGE_BAND_Y, BAND_X_END) + +# reel labels on the left edge +for y0, label, color in ( + (TEXT_BAND_Y, "TEXT REEL", CYAN), + (IMAGE_BAND_Y, "IMAGE REEL", ORANGE), +): + ax.text(X0 - 56, y0 + BAND_H / 2, label, color=color, fontsize=13, + fontweight="bold", rotation=90, ha="center", va="center", zorder=8) + ax.text(X0 - 84, y0 + BAND_H / 2, "12 \u00d7 12 carrier cosine", color=MUTED, + fontsize=8, rotation=90, ha="center", va="center", zorder=8) + +# ---------------------------------------------------------------- frames +VLIM = 0.75 # diagonal (cos=1) clips to amber, off-diagonal structure fills the range + + +def draw_matrix(mat: np.ndarray, cx: float, band_y: float) -> None: + x0m, y0m = cx - FS / 2, band_y + 36 + ax.imshow( + mat, cmap=DIVERGING, vmin=-VLIM, vmax=VLIM, + extent=(x0m, x0m + FS, y0m + FS, y0m), origin="upper", + interpolation="nearest", zorder=4, + ) + ax.add_patch(Rectangle((x0m, y0m), FS, FS, fill=False, + edgecolor=EDGE, linewidth=1.1, zorder=5)) + + +for i, layer in enumerate(LAYERS): + cx = col_cx(i) + draw_matrix(text_sim[layer], cx, TEXT_BAND_Y) + draw_matrix(image_sim[layer], cx, IMAGE_BAND_Y) + + # frame numbering, film style + ax.text(cx, TEXT_BAND_Y - 12, f"FRAME {i + 1:02d}", color=MUTED, + fontsize=8.5, ha="center", va="bottom", zorder=8) + for band_y in (TEXT_BAND_Y, IMAGE_BAND_Y): + ax.text(cx, band_y + 36 + FS + 14, f"LAYER {layer}", color=INK, + fontsize=10, fontweight="bold", ha="center", va="center", zorder=8) + + # dotted connector between the paired frames + ax.plot([cx, cx], [TEXT_BAND_Y + BAND_H + 3, IMAGE_BAND_Y - 3], + color="#3a4754", linewidth=1.2, linestyle=(0, (1, 3)), zorder=3) + +# ---------------------------------------------------------------- match meters +ax.text(X0 - 26, METER_Y - 14, "GEOMETRY MATCH", color=INK, fontsize=10, + fontweight="bold", ha="left", va="bottom", zorder=8) +ax.text(X0 + 152, METER_Y - 14, + "amber bar \u2014 RSA: Pearson r of the two reels' off-diagonal structure" + " cyan tick \u2014 matched cross-carrier cosine", + color=MUTED, fontsize=8.5, ha="left", va="bottom", zorder=8) + +BAR_W = FS +for i, layer in enumerate(LAYERS): + row = per_layer[layer] + rsa = row["rsa_pearson"] + matched = row["matched_cosine"] + cx = col_cx(i) + bx = cx - BAR_W / 2 + + ax.add_patch(Rectangle((bx, METER_Y), BAR_W, 12, facecolor=PANEL, + edgecolor=EDGE, linewidth=0.8, zorder=4)) + ax.add_patch(Rectangle((bx, METER_Y), BAR_W * rsa, 12, facecolor=AMBER, + edgecolor="none", zorder=5)) + ax.plot([bx + BAR_W * matched] * 2, [METER_Y - 4, METER_Y + 16], + color=CYAN, linewidth=2.0, zorder=6) + + ax.text(cx, METER_Y + 32, f"RSA {rsa:.2f}", color=AMBER, fontsize=10.5, + fontweight="bold", ha="center", va="center", zorder=8) + ax.text(cx, METER_Y + 50, f"matched cos {matched:.2f}", color=CYAN, + fontsize=8.5, ha="center", va="center", zorder=8) + +# ---------------------------------------------------------------- closing callout frame +cb_x = X0 + (N_COLS - 1) * CW + 2 +cb_w = X1 - cb_x +cb_y0, cb_y1 = TEXT_BAND_Y, METER_Y + METER_H +ax.add_patch( + FancyBboxPatch( + (cb_x, cb_y0), cb_w, cb_y1 - cb_y0, + boxstyle="round,pad=0,rounding_size=10", + facecolor=PANEL, edgecolor=AMBER, linewidth=1.6, zorder=4, + ) +) +ccx = cb_x + cb_w / 2 +ax.text(ccx, cb_y0 + 46, "THE SPLICE", color=MUTED, fontsize=10, + ha="center", va="center", zorder=8) +ax.text(ccx, cb_y0 + 122, f"RSA {best['rsa_pearson']:.2f}", color=AMBER, + fontsize=33, fontweight="bold", ha="center", va="center", zorder=8) +ax.text(ccx, cb_y0 + 168, f"@ LAYER {best['layer']}", color=INK, fontsize=14, + fontweight="bold", ha="center", va="center", zorder=8) + +ax.plot([cb_x + 28, cb_x + cb_w - 28], [cb_y0 + 206] * 2, + color=EDGE, linewidth=1.0, zorder=5) + +facts = [ + (f"matched cosine {best['matched_cosine']:.2f}", CYAN), + (f"mismatched {best['mismatched_cosine']:.2f}", MUTED), + (f"retrieval {int(round(best['match_rank_accuracy'] * 12))}/12", GREEN), +] +for j, (line, color) in enumerate(facts): + ax.text(ccx, cb_y0 + 244 + j * 34, line, color=color, fontsize=11.5, + fontweight="bold", ha="center", va="center", zorder=8) + +ax.text(ccx, cb_y0 + 380, "Read it as text or look at\nthe picture \u2014 by layer 19\nthe model files both under\nthe same geometry.", + color=INK, fontsize=10.5, ha="center", va="center", linespacing=1.6, zorder=8) + +# the actual L19 splice: the twin pair, miniaturized +MINI = 78 +for mat, mx, tag, tcol in ( + (text_sim[best["layer"]], ccx - MINI - 9, "text", CYAN), + (image_sim[best["layer"]], ccx + 9, "image", ORANGE), +): + ax.imshow(mat, cmap=DIVERGING, vmin=-VLIM, vmax=VLIM, + extent=(mx, mx + MINI, cb_y0 + 444 + MINI, cb_y0 + 444), + origin="upper", interpolation="nearest", zorder=6) + ax.add_patch(Rectangle((mx, cb_y0 + 444), MINI, MINI, fill=False, + edgecolor=EDGE, linewidth=1.0, zorder=7)) + ax.text(mx + MINI / 2, cb_y0 + 444 + MINI + 14, tag, color=tcol, + fontsize=9, ha="center", va="center", zorder=8) +ax.text(ccx, cb_y1 - 36, "two carriers,\none geometry", color=AMBER, fontsize=11, + fontweight="bold", fontstyle="italic", ha="center", va="center", zorder=8) + +# ---------------------------------------------------------------- title & footer +ax.text(X0 - 26, 64, "TWIN REELS", color=INK, fontsize=34, fontweight="bold", + ha="left", va="center", zorder=8) +ax.text(X0 + 318, 64, "\u2014 the same 12 facts, shot twice", color=AMBER, + fontsize=16, ha="left", va="center", zorder=8) +ax.text( + X0 - 26, 118, + "Twelve question\u2013answer pairs enter Qwen2.5-VL-7B twice: once as text, once rendered into pixels. " + "Each frame is the 12\u00d712 cosine similarity between carrier states at one layer \u2014 " + "the two reels print the same relational structure from the very first frames.", + color=MUTED, fontsize=11.5, ha="left", va="center", zorder=8, +) + +ax.text( + X0 - 26, FOOT_Y, + "data: results/qwen-carrier-convergence-n12 (carrier_convergence.npz \u00b7 summary.json) \u00b7 " + "carrier-centered cosine of hidden states, d = 3584, 29 layers \u00b7 " + "RSA = Pearson r over the 66 off-diagonal pairs \u00b7 " + "diverging scale \u2212%.2f \u2026 +%.2f (cyan \u2192 dark \u2192 orange)" % (VLIM, VLIM), + color=MUTED, fontsize=9, ha="left", va="center", zorder=8, +) + +# ---------------------------------------------------------------- save +os.makedirs(OUT_DIR, exist_ok=True) +out_path = os.path.join(OUT_DIR, "filmstrip.png") +fig.savefig(out_path, dpi=100, facecolor=BG) +print("wrote", out_path) diff --git a/packages/snapcompact/research/snapcompact_r2_hero.py b/packages/snapcompact/research/snapcompact_r2_hero.py new file mode 100755 index 000000000..49370267d --- /dev/null +++ b/packages/snapcompact/research/snapcompact_r2_hero.py @@ -0,0 +1,497 @@ +#!/usr/bin/env python3 +"""Hero / OG social card: "Two carriers, one thought." + +A single poster composition for the snapcompact blog post: a fragment of the +real text carrier (BPE tokens `spect`+`acular`) on the left, the real bitmap +patch pixels of the same word on the right, both flowing into one glowing +shared core annotated with the real layer-19 convergence stats. + +Outputs: + results/agent-r2-hero/hero-1200x630.png + results/agent-r2-hero/hero-2400x1260.png + +All numbers are read from: + results/qwen-carrier-convergence-n12/summary.json + results/qwen-token-entry-q3/token_entry.json + results/qwen-logit-lens-q3/images/image-carrier.png +""" + +from __future__ import annotations + +import json +import math +from pathlib import Path + +import numpy as np +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +ROOT = Path(__file__).resolve().parent +OUT_DIR = ROOT / "results" / "agent-r2-hero" + +# Layout is specified in 2400x1260 coordinates; the master canvas renders at +# MS x that for antialiasing, then is downsampled to both deliverables. +BASE_W, BASE_H = 2400, 1260 +MS = 2 +W, H = BASE_W * MS, BASE_H * MS + +# Palette (brief). +BG = (5, 7, 10) +PANEL = (12, 17, 23) +INK = (241, 239, 224) +MUTED = (143, 154, 160) +AMBER = (255, 196, 68) +CYAN = (75, 220, 255) +DIVIDER = (26, 34, 44) + +CORE_WORD = '"spectacular"' + + +def u(v: float) -> int: + return int(round(v * MS)) + + +def font_at(path: str, size: float, index: int = 0) -> ImageFont.FreeTypeFont | None: + p = Path(path) + if not p.exists(): + return None + try: + return ImageFont.truetype(str(p), u(size), index=index) + except OSError: + return None + + +def display_font(size: float) -> ImageFont.FreeTypeFont: + """Heavy display face for the title and big stats.""" + for path, index in [ + ("/System/Library/Fonts/Avenir Next.ttc", 8), # Heavy + ("/System/Library/Fonts/Supplemental/Arial Black.ttf", 0), + ("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 0), + ]: + f = font_at(path, size, index) + if f is not None: + return f + return ImageFont.load_default() + + +def label_font(size: float) -> ImageFont.FreeTypeFont: + for path, index in [ + ("/System/Library/Fonts/Avenir Next.ttc", 2), # Demi Bold + ("/System/Library/Fonts/Supplemental/Arial Bold.ttf", 0), + ("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 0), + ]: + f = font_at(path, size, index) + if f is not None: + return f + return ImageFont.load_default() + + +def body_font(size: float) -> ImageFont.FreeTypeFont: + for path, index in [ + ("/System/Library/Fonts/Avenir Next.ttc", 5), # Medium + ("/System/Library/Fonts/Supplemental/Arial.ttf", 0), + ("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 0), + ]: + f = font_at(path, size, index) + if f is not None: + return f + return ImageFont.load_default() + + +def mono_font(size: float) -> ImageFont.FreeTypeFont: + for path in ["/System/Library/Fonts/Monaco.ttf", "/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf"]: + f = font_at(path, size) + if f is not None: + return f + return ImageFont.load_default() + + +def tracked(draw: ImageDraw.ImageDraw, xy: tuple[int, int], text: str, font, fill, tracking: float = 0.0) -> int: + """Draw text with letterspacing; returns end x.""" + x, y = xy + t = u(tracking) + for ch in text: + draw.text((x, y), ch, font=font, fill=fill) + x += draw.textlength(ch, font=font) + t + return int(x) + + +def tracked_width(draw: ImageDraw.ImageDraw, text: str, font, tracking: float = 0.0) -> float: + t = u(tracking) + return sum(draw.textlength(ch, font=font) + t for ch in text) - (t if text else 0) + + +def bezier(p0, p1, p2, n=64): + pts = [] + for i in range(n + 1): + t = i / n + x = (1 - t) ** 2 * p0[0] + 2 * (1 - t) * t * p1[0] + t**2 * p2[0] + y = (1 - t) ** 2 * p0[1] + 2 * (1 - t) * t * p1[1] + t**2 * p2[1] + pts.append((x, y)) + return pts + + +def load_data(): + conv = json.loads((ROOT / "results" / "qwen-carrier-convergence-n12" / "summary.json").read_text()) + entry = json.loads((ROOT / "results" / "qwen-token-entry-q3" / "token_entry.json").read_text()) + carrier = Image.open(ROOT / "results" / "qwen-logit-lens-q3" / "images" / "image-carrier.png").convert("RGB") + if carrier.size != (1568, 1568): + carrier = carrier.resize((1568, 1568), Image.LANCZOS) + + best = conv["best"] + layer = conv["best_layer"] + n_q = conv["n_questions"] + stats = { + "layer": layer, + "n_layers": conv["layers"], + "matched": best["matched_cosine"], + "rsa": best["rsa_pearson"], + "retrieved": int(round(best["match_rank_accuracy"] * n_q)), + "n": n_q, + } + assert stats["layer"] == 19 and abs(stats["matched"] - 0.66) < 0.01 + assert abs(stats["rsa"] - 0.85) < 0.01 and stats["retrieved"] == 12 and stats["n"] == 12 + + toks = {t["i"]: t for t in entry["tokens"]} + answer = [t for t in entry["tokens"] if t["answer"]] + assert [t["str"] for t in answer] == ["spect", "acular"] + assert [t["id"] for t in answer] == [67082, 23006] + ctx_before = "…" + "".join(toks[i]["str"] for i in range(23, 32)) # " make the 50th Super Bowl \"" + ctx_after = "".join(toks[i]["str"] for i in range(34, 39)) + "…" # "\" and that it would" + + grid = entry["image_grid"] # 56 + word_idx = entry["image_answer_token_indices"][:4] # [310, 311, 312, 313] + assert word_idx == [310, 311, 312, 313] + assert word_idx[0] // grid == 5 and word_idx[0] % grid == 30 + + counts = { + "chars": entry["chunk_chars"], + "text_tokens": entry["chunk_text_tokens"], + "image_tokens": entry["image_tokens"], + "grid": grid, + "patch_px": entry["token_pixel_size"], # 28 + "embed_dim": entry["embed_dim"], + "visual_dim": entry["visual_out_dim"], + } + heads = { + "text": entry["text_entry"][0], # id 67082 "spect": 10-dim head + norm + "image": entry["image_entry"][0], # patch 310: 10-dim head + norm + } + assert heads["text"]["id"] == 67082 and heads["image"]["token_index"] == 310 + return stats, answer, (ctx_before, ctx_after), word_idx, counts, heads, carrier + + +# --------------------------------------------------------------------------- +# Composition geometry (2400x1260 space) +# --------------------------------------------------------------------------- +CORE = (1200, 660) +PANEL_TOP, PANEL_BOT = 332, 938 +LP = (92, PANEL_TOP, 782, PANEL_BOT) # left panel +RP = (1618, PANEL_TOP, 2308, PANEL_BOT) # right panel + + +def additive_base() -> np.ndarray: + """Background + ambient glow + orb fields + streamlines, all additive.""" + img = np.zeros((H, W, 3), dtype=np.float32) + img[:] = BG + + yy, xx = np.mgrid[0:H, 0:W].astype(np.float32) + cx, cy = u(CORE[0]), u(CORE[1]) + d = np.sqrt((xx - cx) ** 2 + (yy - cy) ** 2) + + # Ambient: amber wash on the left, cyan on the right, strongest near core. + side = np.clip((xx - cx) / u(900), -1.0, 1.0) + amber = np.array(AMBER, np.float32) / 255.0 + cyan = np.array(CYAN, np.float32) / 255.0 + mix = (1 - side[..., None]) / 2 * amber + (1 + side[..., None]) / 2 * cyan + img += 26.0 * mix * np.exp(-((d / u(760)) ** 2))[..., None] + + # Orb halo and warm nucleus (kept below blowout so the word stays legible). + img += 92.0 * mix * np.exp(-((d / u(250)) ** 2))[..., None] + warm = np.array((255, 240, 205), np.float32) / 255.0 + img += 96.0 * warm * np.exp(-((d / u(118)) ** 2))[..., None] + + # Subtle vignette. + ex = ((xx / W) - 0.5) ** 2 + ((yy / H) - 0.5) ** 2 + img *= (1.0 - 0.55 * ex)[..., None] + + # Streamlines: each carrier feeds the core. + glow = Image.new("RGB", (W, H), (0, 0, 0)) + gd = ImageDraw.Draw(glow) + sharp = Image.new("RGB", (W, H), (0, 0, 0)) + sd = ImageDraw.Draw(sharp) + rng = np.random.default_rng(19) + + def streams(x0: float, x_sign: float, color): + n = 9 + for k in range(n): + f = k / (n - 1) + y0 = PANEL_TOP + 120 + f * (PANEL_BOT - PANEL_TOP - 240) + rng.uniform(-14, 14) + ang = (f - 0.5) * 1.45 + rng.uniform(-0.07, 0.07) + r = 168 + x2 = CORE[0] - x_sign * r * math.cos(ang) + y2 = CORE[1] + r * math.sin(ang) + mid_x = (x0 + x2) / 2 + x_sign * 36 + mid_y = y0 + (CORE[1] - y0) * 0.62 + pts = bezier((x0, y0), (mid_x, mid_y), (x2, y2), 72) + spts = [(u(px), u(py)) for px, py in pts] + fade = 1.0 - abs(f - 0.5) * 0.9 + gd.line(spts, fill=tuple(int(c * 0.62 * fade) for c in color), width=u(7)) + sd.line(spts, fill=tuple(int(c * 0.5 * fade) for c in color), width=u(1.6)) + # Energy particles along the stream. + for t in (0.3, 0.55, 0.8): + i = int(t * len(pts)) + px, py = u(pts[i][0]), u(pts[i][1]) + rr = u(3.2) + gd.ellipse([px - rr, py - rr, px + rr, py + rr], fill=tuple(int(c * fade) for c in color)) + + streams(LP[2], 1.0, AMBER) + streams(RP[0], -1.0, CYAN) + + img += np.asarray(glow.filter(ImageFilter.GaussianBlur(u(11))), np.float32) * 0.9 + img += np.asarray(sharp.filter(ImageFilter.GaussianBlur(u(0.7))), np.float32) + + # Soft glow behind the core word (the sharp pass is drawn later, on top). + f = display_font(78) + layer = Image.new("RGB", (W, H), (0, 0, 0)) + ld = ImageDraw.Draw(layer) + tw = ld.textlength(CORE_WORD, font=f) + ld.text((u(CORE[0]) - tw / 2, u(CORE[1] - 54)), CORE_WORD, font=f, fill=(255, 232, 170)) + img += np.asarray(layer.filter(ImageFilter.GaussianBlur(u(9))), np.float32) * 0.8 + return img + + +def head_bars(ov: ImageDraw.ImageDraw, x: float, y_mid: float, values, color, label: str): + """Tiny bar strip of a real 10-dim vector head, centered on its axis.""" + vmax = max(abs(v) for v in values) + bw, gap, amp = 24, 11, 26 + total = len(values) * (bw + gap) - gap + ov.line([u(x), u(y_mid), u(x + total), u(y_mid)], fill=(*MUTED, 80), width=u(1)) + for i, v in enumerate(values): + bx = x + i * (bw + gap) + h = (v / vmax) * amp + y0, y1 = sorted((y_mid, y_mid - h)) + a = 120 + int(135 * abs(v) / vmax) + ov.rectangle([u(bx), u(y0), u(bx + bw), u(y1)], fill=(*color, a)) + f = mono_font(17) + tracked(ov, (u(x), u(y_mid - amp - 36)), label, f, (*MUTED, 255), tracking=0.5) + + +def draw_left_panel(ov: ImageDraw.ImageDraw, answer, ctx, counts, heads): + x0, y0, x1, _ = LP + pad = 44 + ctx_before, ctx_after = ctx + + tracked(ov, (u(x0 + pad), u(y0 + 34)), "TEXT CARRIER", label_font(30), AMBER, tracking=5) + sub = f"{counts['text_tokens']:,} BPE TOKENS" + f_sub = label_font(21) + tracked(ov, (int(u(x1 - pad) - tracked_width(ov, sub, f_sub, 2)), u(y0 + 42)), sub, f_sub, MUTED, tracking=2) + ov.line([u(x0 + pad), u(y0 + 88), u(x1 - pad), u(y0 + 88)], fill=(*DIVIDER, 255), width=u(1.2)) + + f_ctx = mono_font(23) + ov.text((u(x0 + pad), u(y0 + 116)), ctx_before, font=f_ctx, fill=(110, 118, 126)) + + # The two answer-token pills. + f_tok = mono_font(58) + f_id = mono_font(20) + px = x0 + pad + py = y0 + 184 + for t in answer: + s = t["str"] + wpx = ov.textlength(s, font=f_tok) / MS + ov.rounded_rectangle( + [u(px), u(py), u(px + wpx + 40), u(py + 96)], + radius=u(14), + fill=(38, 29, 10, 235), + outline=(*AMBER, 165), + width=u(1.6), + ) + ov.text((u(px + 20), u(py + 14)), s, font=f_tok, fill=(255, 224, 150)) + ov.text((u(px + 20), u(py + 110)), f"id {t['id']}", font=f_id, fill=(196, 156, 72)) + px += wpx + 40 + 22 + + ov.text((u(x0 + pad), u(y0 + 330)), ctx_after, font=f_ctx, fill=(110, 118, 126)) + + th = heads["text"] + head_bars( + ov, + x0 + pad, + y0 + 452, + th["vector_head"], + AMBER, + f"embedding row {th['id']} · dims 0-9 of {counts['embed_dim']:,} · norm {th['norm']:.2f}", + ) + + fy = y0 + 500 + f_fact = body_font(23) + ov.text((u(x0 + pad), u(fy)), f"{counts['chars']:,} characters of one SQuAD passage,", font=f_fact, fill=MUTED) + ov.text( + (u(x0 + pad), u(fy + 36)), + f"tokenized into {counts['text_tokens']:,} ids, each a {counts['embed_dim']:,}-dim row", + font=f_fact, + fill=MUTED, + ) + + +def draw_right_panel(base_img: Image.Image, ov: ImageDraw.ImageDraw, carrier: Image.Image, word_idx, counts, heads): + x0, y0, x1, _ = RP + pad = 44 + + tracked(ov, (u(x0 + pad), u(y0 + 34)), "IMAGE CARRIER", label_font(30), CYAN, tracking=5) + sub = f"{counts['image_tokens']:,} PATCHES" + f_sub = label_font(21) + tracked(ov, (int(u(x1 - pad) - tracked_width(ov, sub, f_sub, 2)), u(y0 + 42)), sub, f_sub, MUTED, tracking=2) + ov.line([u(x0 + pad), u(y0 + 88), u(x1 - pad), u(y0 + 88)], fill=(*DIVIDER, 255), width=u(1.2)) + + # Crop: patch rows 4..8, cols 27..38 of the 56x56 grid (28px cells). + pp = counts["patch_px"] + c0, c1, r0, r1 = 27, 38, 4, 8 + crop = carrier.crop((c0 * pp, r0 * pp, c1 * pp, r1 * pp)) # 308 x 112 + + scale = 2.0 # 28px cell -> 56px on the 2400 canvas + disp_w, disp_h = int(crop.width * scale), int(crop.height * scale) + big = crop.resize((u(disp_w), u(disp_h)), Image.NEAREST) + big = Image.eval(big, lambda v: int(v * 0.84)) # dim so the highlight pops + bx, by = x0 + pad, y0 + 122 + base_img.paste(big, (u(bx), u(by))) + + cell = pp * scale # 56 in 2400-space + grid_color = (CYAN[0], CYAN[1], CYAN[2], 46) + for c in range(c1 - c0 + 1): + ov.line([u(bx + c * cell), u(by), u(bx + c * cell), u(by + disp_h)], fill=grid_color, width=u(1)) + for r in range(r1 - r0 + 1): + ov.line([u(bx), u(by + r * cell), u(bx + disp_w), u(by + r * cell)], fill=grid_color, width=u(1)) + + # Highlight the four answer patches (grid row 5, cols 30..33) as one run. + grid = counts["grid"] + gr, gc = word_idx[0] // grid, word_idx[0] % grid + hx, hy = bx + (gc - c0) * cell, by + (gr - r0) * cell + hw = len(word_idx) * cell + ov.rectangle([u(hx), u(hy), u(hx + hw), u(hy + cell)], outline=(*CYAN, 240), width=u(2.4)) + for k in range(1, len(word_idx)): + ov.line([u(hx + k * cell), u(hy), u(hx + k * cell), u(hy + cell)], fill=(*CYAN, 130), width=u(1.2)) + + ov.text( + (u(bx), u(by + disp_h + 18)), + f"patches {word_idx[0]}-{word_idx[-1]} · grid row 5, cols 30-33 of {grid}×{grid}", + font=body_font(23), + fill=MUTED, + ) + + ih = heads["image"] + head_bars( + ov, + bx, + y0 + 460, + ih["vector_head"], + CYAN, + f"patch {ih['token_index']} vector · dims 0-9 of {counts['visual_dim']:,} · norm {ih['norm']:.1f}", + ) + + fy = y0 + 500 + f_fact = body_font(23) + ov.text((u(bx), u(fy)), "the same passage, rendered to a 1568 × 1568 px bitmap,", font=f_fact, fill=MUTED) + ov.text( + (u(bx), u(fy + 36)), + f"seen as {counts['image_tokens']:,} patches of {counts['patch_px']} px, each a {counts['visual_dim']:,}-dim vector", + font=f_fact, + fill=MUTED, + ) + + +def draw_title(ov: ImageDraw.ImageDraw): + kicker = "INSIDE QWEN2.5-VL · ONE FACT, TWO ENCODINGS" + f_k = label_font(24) + kw = tracked_width(ov, kicker, f_k, 7) + tracked(ov, (int((W - kw) / 2), u(64)), kicker, f_k, MUTED, tracking=7) + + f_t = display_font(96) + a, b = "TWO CARRIERS, ", "ONE THOUGHT." + wa, wb = ov.textlength(a, font=f_t), ov.textlength(b, font=f_t) + x = (W - wa - wb) / 2 + y = u(122) + ov.text((x, y), a, font=f_t, fill=INK) + ov.text((x + wa, y), b, font=f_t, fill=AMBER) + + +def draw_core(ov: ImageDraw.ImageDraw, stats): + f = display_font(78) + tw = ov.textlength(CORE_WORD, font=f) + ov.text( + (u(CORE[0]) - tw / 2, u(CORE[1] - 54)), + CORE_WORD, + font=f, + fill=(255, 248, 226), + stroke_width=u(1.4), + stroke_fill=(64, 44, 12, 160), + ) + cap = f"BY LAYER {stats['layer']} OF {stats['n_layers'] - 1}, ONE SHARED STATE" + f_c = label_font(23) + cw = tracked_width(ov, cap, f_c, 4) + tracked(ov, (int(u(CORE[0]) - cw / 2), u(CORE[1] + 96)), cap, f_c, (228, 222, 196), tracking=4) + + +def draw_stats_strip(ov: ImageDraw.ImageDraw, stats): + y_div = 992 + ov.line([u(92), u(y_div), u(2308), u(y_div)], fill=(*DIVIDER, 255), width=u(1.4)) + + groups = [ + (f"{stats['matched']:.2f}", "MATCHED COSINE · TEXT VS IMAGE", AMBER), + (f"{stats['rsa']:.2f}", "RSA · SAME RELATIONAL GEOMETRY", INK), + (f"{stats['retrieved']}/{stats['n']}", "CROSS-CARRIER RETRIEVAL", CYAN), + ] + centers = [500, 1200, 1900] + f_num = display_font(66) + f_cap = label_font(21) + for (num, cap, color), cx in zip(groups, centers): + nw = ov.textlength(num, font=f_num) + ov.text((u(cx) - nw / 2, u(1024)), num, font=f_num, fill=color) + cw = tracked_width(ov, cap, f_cap, 3) + tracked(ov, (int(u(cx) - cw / 2), u(1136)), cap, f_cap, MUTED, tracking=3) + for dx in (850, 1550): + ov.line([u(dx), u(1040), u(dx), u(1170)], fill=(*DIVIDER, 255), width=u(1.2)) + + foot = f"measured at layer {stats['layer']} · {stats['n']} SQuAD questions · carrier-convergence" + f_f = body_font(19) + fw = ov.textlength(foot, font=f_f) + ov.text(((W - fw) / 2, u(1206)), foot, font=f_f, fill=(92, 101, 108)) + + +def rounded_panel(overlay: ImageDraw.ImageDraw, box, accent, alpha_fill=216): + x0, y0, x1, y1 = (u(v) for v in box) + r = u(22) + overlay.rounded_rectangle([x0, y0, x1, y1], radius=r, fill=(*PANEL, alpha_fill)) + overlay.rounded_rectangle([x0, y0, x1, y1], radius=r, outline=(*accent, 70), width=u(1.4)) + + +def main() -> None: + OUT_DIR.mkdir(parents=True, exist_ok=True) + stats, answer, ctx, word_idx, counts, heads, carrier = load_data() + + base = Image.fromarray(np.clip(additive_base(), 0, 255).astype(np.uint8), "RGB") + + overlay = Image.new("RGBA", (W, H), (0, 0, 0, 0)) + ov = ImageDraw.Draw(overlay) + rounded_panel(ov, LP, AMBER) + rounded_panel(ov, RP, CYAN) + base = Image.alpha_composite(base.convert("RGBA"), overlay).convert("RGB") + + # Panel contents drawn on a fresh overlay so the bitmap paste sits beneath grids. + overlay2 = Image.new("RGBA", (W, H), (0, 0, 0, 0)) + ov2 = ImageDraw.Draw(overlay2) + draw_title(ov2) + draw_left_panel(ov2, answer, ctx, counts, heads) + draw_right_panel(base, ov2, carrier, word_idx, counts, heads) + draw_core(ov2, stats) + draw_stats_strip(ov2, stats) + final = Image.alpha_composite(base.convert("RGBA"), overlay2).convert("RGB") + + retina = final.resize((2400, 1260), Image.LANCZOS) + og = final.resize((1200, 630), Image.LANCZOS) + retina.save(OUT_DIR / "hero-2400x1260.png") + og.save(OUT_DIR / "hero-1200x630.png") + print(f"wrote {OUT_DIR / 'hero-2400x1260.png'} {retina.size}") + print(f"wrote {OUT_DIR / 'hero-1200x630.png'} {og.size}") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_r2_metro.py b/packages/snapcompact/research/snapcompact_r2_metro.py new file mode 100755 index 000000000..d7da1477a --- /dev/null +++ b/packages/snapcompact/research/snapcompact_r2_metro.py @@ -0,0 +1,348 @@ +#!/usr/bin/env python3 +"""Snapcompact R2 — "The Convergence Line". + +A transit/metro-map diagram of carrier convergence in Qwen2.5-VL-7B. +Two metro lines (cyan = text carrier, orange = image carrier) run through +29 stations (decoder layers L0..L28). The vertical gap between the lines at +each station is driven by real per-layer data: + + gap ~ 1 - matched_cosine (results/qwen-carrier-convergence-n12/summary.json) + +Named stations are grounded in the same summary.json plus +results/qwen-logit-lens-q3/logit_lens.json (visual tok[310] -> 'acular'). + +Output: results/agent-r2-metro/metro.png (~2200 px wide). +""" + +import json +import os + +import numpy as np +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 +from matplotlib.patches import Circle # noqa: E402 + +HERE = os.path.dirname(os.path.abspath(__file__)) +SUMMARY_PATH = os.path.join(HERE, "results", "qwen-carrier-convergence-n12", "summary.json") +LENS_PATH = os.path.join(HERE, "results", "qwen-logit-lens-q3", "logit_lens.json") +OUT_DIR = os.path.join(HERE, "results", "agent-r2-metro") +OUT_PNG = os.path.join(OUT_DIR, "metro.png") + +# ---------------------------------------------------------------- palette +BG = "#05070a" +PANEL = "#0c1117" +INK = "#f1efe0" +MUTED = "#8f9aa0" +AMBER = "#ffc444" +CYAN = "#4bdcff" +ORANGE = "#ff7048" +GRID = "#0e141b" +MONO = "DejaVu Sans Mono" +SANS = "DejaVu Sans" + +MINUS = "\u2212" + + +def load_data(): + with open(SUMMARY_PATH) as f: + summary = json.load(f) + with open(LENS_PATH) as f: + lens = json.load(f) + + per = summary["per_layer"] + assert len(per) == summary["layers"] == 29 + assert summary["best_layer"] == 19 + + cos = np.array([p["matched_cosine"] for p in per]) + rsa = np.array([p["rsa_pearson"] for p in per]) + acc = np.array([p["match_rank_accuracy"] for p in per]) + mism = np.array([p["mismatched_cosine"] for p in per]) + + # logit lens: visual token 310, answer piece 'acular' + tok310 = {e["layer"]: e for e in lens["lens"] if e["token_index"] == 310} + assert tok310[24]["top"][0]["str"] == "acular" + p_acular_24 = tok310[24]["answer_token_p"][1] + p_acular_28 = tok310[28]["answer_token_p"][1] + answer = "".join(lens["answer_token_strs"]) # 'spectacular' + + return { + "summary": summary, + "cos": cos, + "rsa": rsa, + "acc": acc, + "mism": mism, + "p24": p_acular_24, + "p28": p_acular_28, + "answer": answer, + "n_q": summary["n_questions"], + "geometry": summary["geometry"], + "size_px": summary["args"]["size"], + "text_em": summary["text_em"], + "image_em": summary["image_em"], + } + + +def catmull_rom(xs, ys, samples=26): + """Centripetal-ish Catmull-Rom through all points (uniform parameter).""" + pts = np.column_stack([xs, ys]).astype(float) + ext = np.vstack([pts[0], pts, pts[-1]]) + out = [] + t = np.linspace(0.0, 1.0, samples, endpoint=False)[:, None] + for i in range(len(pts) - 1): + p0, p1, p2, p3 = ext[i], ext[i + 1], ext[i + 2], ext[i + 3] + a = 2.0 * p1 + b = p2 - p0 + c = 2.0 * p0 - 5.0 * p1 + 4.0 * p2 - p3 + d = -p0 + 3.0 * p1 - 3.0 * p2 + p3 + out.append(0.5 * (a + b * t + c * t**2 + d * t**3)) + out.append(pts[-1][None]) + return np.vstack(out) + + +def fmt2(v): + s = f"{v:.2f}" + return s.replace("-", MINUS) + + +def main(): + d = load_data() + cos = d["cos"] + n_layers = len(cos) + + # ---- track geometry: gap shrinks as matched cosine rises ------------- + c_min, c_max = float(cos.min()), float(cos.max()) # 0.0 (L0) .. 0.658 (L19) + GAP_MAX, GAP_MIN = 5.6, 0.78 + t = (cos - c_min) / (c_max - c_min) + gap = GAP_MAX + (GAP_MIN - GAP_MAX) * t + xs = np.arange(n_layers, dtype=float) + y_text = gap / 2.0 + y_img = -gap / 2.0 + + # depot stubs before L0 + xs_t = np.concatenate([[-1.5], xs]) + xs_i = np.concatenate([[-1.5], xs]) + yt = np.concatenate([[y_text[0]], y_text]) + yi = np.concatenate([[y_img[0]], y_img]) + + path_t = catmull_rom(xs_t, yt) + path_i = catmull_rom(xs_i, yi) + + # ---- figure ---------------------------------------------------------- + X0, X1 = -2.6, 32.2 + Y0, Y1 = -7.1, 8.3 + W_IN = 22.0 + H_IN = W_IN * (Y1 - Y0) / (X1 - X0) # equal data aspect + fig = plt.figure(figsize=(W_IN, H_IN), dpi=100) + fig.patch.set_facecolor(BG) + ax = fig.add_axes([0, 0, 1, 1]) + ax.set_facecolor(BG) + ax.set_xlim(X0, X1) + ax.set_ylim(Y0, Y1) + ax.set_aspect("equal", adjustable="box") + ax.axis("off") + + pt_per_unit = (W_IN / (X1 - X0)) * 72.0 # ~45.5 pt per data unit + + # faint vertical guides at every station + for i in range(n_layers): + ax.plot([i, i], [-3.55, 3.95], color=GRID, lw=1.0, zorder=1) + + # convergence axis + ax.plot([-1.5, 28.0], [0, 0], color=MUTED, lw=1.0, alpha=0.28, + linestyle=(0, (1, 3)), zorder=1) + ax.text(8.0, 0.16, "convergence axis", color=MUTED, alpha=0.55, + fontsize=8.5, style="italic", family=SANS, ha="left", zorder=2) + + # ---- metro lines: glow, casing, stroke ------------------------------- + for path, col in ((path_t, CYAN), (path_i, ORANGE)): + px, py = path[:, 0], path[:, 1] + ax.plot(px, py, color=col, lw=24, alpha=0.05, solid_capstyle="round", zorder=2) + ax.plot(px, py, color=col, lw=16, alpha=0.07, solid_capstyle="round", zorder=2) + ax.plot(px, py, color=BG, lw=14, solid_capstyle="round", + solid_joinstyle="round", zorder=3) + ax.plot(px, py, color=col, lw=9.5, solid_capstyle="round", + solid_joinstyle="round", zorder=4) + + # ---- stations --------------------------------------------------------- + named = {19, 24, 27, 28} + for i in range(n_layers): + for y, col in ((y_text[i], CYAN), (y_img[i], ORANGE)): + if i in named: + continue + ax.scatter([i], [y], s=115, facecolor=BG, edgecolor=col, + linewidths=2.1, zorder=6) + + # L19 interchange capsule (the two lines meet in one station) + cap_lw_outer = 0.56 * pt_per_unit + ax.plot([19, 19], [y_img[19], y_text[19]], color=INK, + lw=cap_lw_outer, solid_capstyle="round", zorder=5) + ax.plot([19, 19], [y_img[19], y_text[19]], color=PANEL, + lw=cap_lw_outer - 7.5, solid_capstyle="round", zorder=5) + ax.scatter([19, 19], [y_text[19], y_img[19]], s=92, + c=[CYAN, ORANGE], edgecolor=BG, linewidths=1.2, zorder=6) + + # L24 interchange ring on the image line (pixels decode to vocabulary) + ax.scatter([24], [y_img[24]], s=300, facecolor=PANEL, edgecolor=INK, + linewidths=2.8, zorder=6) + ax.scatter([24], [y_img[24]], s=58, facecolor=ORANGE, edgecolor="none", zorder=6) + ax.scatter([24], [y_text[24]], s=115, facecolor=BG, edgecolor=CYAN, + linewidths=2.1, zorder=6) + + # L27 white-ring stations on both lines (terminal approach) + for y, col in ((y_text[27], CYAN), (y_img[27], ORANGE)): + ax.scatter([27], [y], s=170, facecolor=PANEL, edgecolor=INK, + linewidths=2.3, zorder=6) + ax.scatter([27], [y], s=34, facecolor=col, edgecolor="none", zorder=6) + + # L28 terminus: double ring over both tracks + ax.add_patch(Circle((28, 0), 1.02, facecolor=PANEL, edgecolor=INK, + lw=3.2, zorder=5)) + ax.add_patch(Circle((28, 0), 0.66, facecolor="none", edgecolor=INK, + lw=1.3, alpha=0.85, zorder=5)) + ax.scatter([27.78, 28.22], [0, 0], s=120, c=[CYAN, ORANGE], + edgecolor=BG, linewidths=1.4, zorder=6) + ax.text(28, -0.42, "TERMINUS", color=MUTED, fontsize=6.8, family=MONO, + ha="center", va="center", zorder=7) + + # ---- carrier labels (depots) ------------------------------------------ + geo = d["geometry"] + ax.text(-1.55, y_text[0] + 0.95, "TEXT CARRIER", color=CYAN, fontsize=12.5, + family=SANS, fontweight="bold", ha="left", zorder=7) + ax.text(-1.55, y_text[0] + 0.48, + f"the page as typed tokens \u00b7 {geo['capacity']:,} chars", + color=MUTED, fontsize=9, family=SANS, ha="left", zorder=7) + ax.text(-1.55, y_img[0] - 0.62, "IMAGE CARRIER", color=ORANGE, fontsize=12.5, + family=SANS, fontweight="bold", ha="left", zorder=7) + ax.text(-1.55, y_img[0] - 1.09, + f"the same page as a {d['size_px']} px bitmap \u00b7 " + f"{geo['cols']}\u00d7{geo['rows']} cell grid", + color=MUTED, fontsize=9, family=SANS, ha="left", zorder=7) + + # ---- named-station callouts ------------------------------------------- + def leader(x, y_from, y_to, color=MUTED, alpha=0.65): + ax.plot([x, x], [y_from, y_to], color=color, lw=1.1, alpha=alpha, zorder=6) + + # L1: instant alignment + leader(1, y_img[1] - 0.18, -2.18) + ax.text(1.7, -2.35, "L1 \u00b7 INSTANT ALIGNMENT", color=INK, fontsize=10.5, + family=SANS, fontweight="bold", ha="left", zorder=7) + ax.text(1.7, -2.78, + f"matched cos {fmt2(cos[1])} \u00b7 RSA {fmt2(d['rsa'][1])}", + color=MUTED, fontsize=8.8, family=MONO, ha="left", zorder=7) + ax.text(1.7, -3.14, + f"retrieval {int(round(d['acc'][1] * 12))}/12 \u2014 12/12 from L2 onward", + color=MUTED, fontsize=8.8, family=MONO, ha="left", zorder=7) + + # L13: first close pass + leader(13, y_text[13] + 0.18, 1.62) + ax.text(13, 1.84, f"L13 \u00b7 first close pass \u00b7 cos {fmt2(cos[13])}", + color=MUTED, fontsize=8.8, family=MONO, ha="center", zorder=7) + + # L19: geometry locks (star station) + leader(19, y_text[19] + 0.62, 2.42, color=AMBER, alpha=0.8) + ax.text(19, 3.42, "L19 \u00b7 GEOMETRY LOCKS", color=AMBER, fontsize=14, + family=SANS, fontweight="bold", ha="center", zorder=7) + ax.text(19, 2.96, + f"matched cos {fmt2(cos[19])} \u00b7 mismatched {fmt2(d['mism'][19])}", + color=INK, fontsize=9.6, family=MONO, ha="center", zorder=7) + ax.text(19, 2.58, + f"RSA {fmt2(d['rsa'][19])} \u00b7 retrieval 12/12 \u2014 closest approach", + color=MUTED, fontsize=9.6, family=MONO, ha="center", zorder=7) + + # L23: small drift + leader(23, y_text[23] + 0.18, 1.30) + ax.text(23, 1.52, f"L23 \u00b7 small drift \u00b7 cos {fmt2(cos[23])}", + color=MUTED, fontsize=8.8, family=MONO, ha="center", zorder=7) + + # L24: pixels decode to vocabulary + leader(24, y_img[24] - 0.32, -1.92, color=ORANGE, alpha=0.8) + ax.text(24, -2.18, "L24 \u00b7 PIXELS DECODE TO VOCABULARY", color=ORANGE, + fontsize=12.5, family=SANS, fontweight="bold", ha="center", zorder=7) + ax.text(24, -2.62, + f"visual tok[310] top-1 \u2192 'acular' \u00b7 p {d['p24']:.2f}", + color=INK, fontsize=9.4, family=MONO, ha="center", zorder=7) + ax.text(24, -3.00, + f"rising to p {d['p28']:.2f} by L28 \u2014 " + "the answer's second BPE piece", + color=MUTED, fontsize=9.4, family=MONO, ha="center", zorder=7) + + # L27-L28 terminal (block above the terminus circle) + leader(28, 1.18, 1.86, color=AMBER, alpha=0.8) + ax.text(28, 3.00, "L27\u2013L28 \u00b7 TERMINAL", color=INK, fontsize=12.5, + family=SANS, fontweight="bold", ha="center", zorder=7) + ax.text(28, 2.56, f"SAME ANSWER: \u201c{d['answer']}\u201d", color=AMBER, + fontsize=10.5, family=SANS, fontweight="bold", ha="center", zorder=7) + ax.text(28, 2.18, + f"matched cos {fmt2(cos[27])} \u2192 {fmt2(cos[28])} \u00b7 retrieval 12/12", + color=MUTED, fontsize=8.8, family=MONO, ha="center", zorder=7) + + # ---- station index + matched-cosine gauge rows ------------------------- + hl = {19: AMBER, 24: ORANGE, 27: INK, 28: INK} + ax.text(-0.55, -4.45, "layer", color=MUTED, fontsize=8, style="italic", + family=SANS, ha="right", va="center", zorder=7) + ax.text(-0.55, -5.02, "matched cos", color=MUTED, fontsize=8, style="italic", + family=SANS, ha="right", va="center", zorder=7) + for i in range(n_layers): + col = hl.get(i, MUTED) + w = "bold" if i in hl else "normal" + ax.text(i, -4.45, f"L{i}", color=col, fontsize=7.6, family=MONO, + ha="center", va="center", fontweight=w, zorder=7) + val = f"{cos[i]:.2f}".lstrip("0") + ax.text(i, -5.02, val, color=col, fontsize=7.6, family=MONO, + ha="center", va="center", fontweight=w, zorder=7) + + # ---- title ------------------------------------------------------------- + ax.text(-1.9, 8.05, "THE CONVERGENCE LINE", color=INK, fontsize=29, + family=SANS, fontweight="bold", ha="left", va="top", zorder=7) + ax.text(-1.9, 6.92, + "One Wikipedia page, two carriers: typed tokens (cyan) and a " + f"{d['size_px']} px screenshot (orange) ride Qwen2.5-VL-7B's 29 decoder layers.", + color=MUTED, fontsize=12, family=SANS, ha="left", va="top", zorder=7) + ax.text(-1.9, 6.42, + "The closer the tracks, the more the two internal representations agree " + f"\u2014 track gap \u221d 1 {MINUS} matched cosine, n = {d['n_q']} questions.", + color=MUTED, fontsize=12, family=SANS, ha="left", va="top", zorder=7) + + # ---- legend (top right) ------------------------------------------------- + lx = 22.9 + ax.plot([lx, lx + 1.3], [7.95, 7.95], color=CYAN, lw=8, + solid_capstyle="round", zorder=7) + ax.text(lx + 1.65, 7.95, "TEXT CARRIER", color=INK, fontsize=10, + family=SANS, fontweight="bold", ha="left", va="center", zorder=7) + ax.plot([lx, lx + 1.3], [7.32, 7.32], color=ORANGE, lw=8, + solid_capstyle="round", zorder=7) + ax.text(lx + 1.65, 7.32, "IMAGE CARRIER", color=INK, fontsize=10, + family=SANS, fontweight="bold", ha="left", va="center", zorder=7) + ax.text(lx, 6.62, f"track gap \u221d 1 {MINUS} matched cosine(text, image)", + color=MUTED, fontsize=9, family=MONO, ha="left", va="center", zorder=7) + # wide pair = L0 + ax.plot([lx, lx + 1.0], [6.18, 6.18], color=CYAN, lw=4, solid_capstyle="round", zorder=7) + ax.plot([lx, lx + 1.0], [5.74, 5.74], color=ORANGE, lw=4, solid_capstyle="round", zorder=7) + ax.text(lx + 1.65, 5.96, f"cos {fmt2(cos[0])} \u2014 far apart (L0)", + color=MUTED, fontsize=9, family=MONO, ha="left", va="center", zorder=7) + # tight pair = L19 + ax.plot([lx, lx + 1.0], [5.28, 5.28], color=CYAN, lw=4, solid_capstyle="round", zorder=7) + ax.plot([lx, lx + 1.0], [5.14, 5.14], color=ORANGE, lw=4, solid_capstyle="round", zorder=7) + ax.text(lx + 1.65, 5.21, f"cos {fmt2(cos[19])} \u2014 almost touching (L19)", + color=MUTED, fontsize=9, family=MONO, ha="left", va="center", zorder=7) + + # ---- footer -------------------------------------------------------------- + ax.text(-1.9, -6.05, + "Across the same 12 questions the image carrier matches gold answers as often as the text carrier " + f"\u2014 image EM {d['image_em'] * 100:.1f}% vs text EM {d['text_em'] * 100:.1f}%.", + color=MUTED, fontsize=9.5, family=SANS, ha="left", zorder=7) + ax.text(-1.9, -6.58, + "Data: results/qwen-carrier-convergence-n12/summary.json (29 layers \u00b7 12 SQuAD questions) " + "+ results/qwen-logit-lens-q3/logit_lens.json \u00b7 Qwen2.5-VL-7B-Instruct \u00b7 agent r2-metro", + color=MUTED, alpha=0.7, fontsize=8.5, family=MONO, ha="left", zorder=7) + + os.makedirs(OUT_DIR, exist_ok=True) + fig.savefig(OUT_PNG, dpi=100, facecolor=BG) + plt.close(fig) + print(f"wrote {OUT_PNG}") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_tensor_heatmap.py b/packages/snapcompact/research/snapcompact_tensor_heatmap.py new file mode 100644 index 000000000..0cc441e30 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_tensor_heatmap.py @@ -0,0 +1,354 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy", "torch", "transformers", "sentencepiece", "protobuf", "einops"] +# /// +"""Capture and render token/layer hidden-state heatmaps for snapcompact masks. + +This is the blog-visual version of the white-box probe: it compares the same +prompt with the original bitmap, a gold-answer-region mask, and an equal random +mask. For every decoder layer and every image placeholder token, it plots +||hidden(original) - hidden(masked)||. +""" + +from __future__ import annotations + +import argparse +import json +import math +import random +import sys +from pathlib import Path +from typing import Any + +import numpy as np +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import capacity, render # noqa: E402 +from run import CACHE, FONTS, load_prompt # noqa: E402 +from snapcompact_blackbox_occlusion import mask_cells, random_span, sample_answer_questions # noqa: E402 + +DEFAULT_MODEL_DIR = ( + "/home/can/.cache/huggingface/hub/models--PaddlePaddle--PaddleOCR-VL/" + "snapshots/2b77538ef936207f60c16b45082841068987d08c" +) + +PALETTE = { + "bg": (5, 7, 10), + "panel": (13, 18, 23), + "ink": (239, 239, 224), + "muted": (132, 147, 154), + "cyan": (77, 218, 255), + "red": (255, 83, 62), + "green": (145, 255, 112), + "amber": (255, 194, 65), + "grid": (34, 45, 53), +} + + +def ui_font(size: int, bold: bool = False) -> ImageFont.FreeTypeFont | ImageFont.ImageFont: + candidates = [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ] + for path in candidates: + if path and Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def heat_color(t: float) -> tuple[int, int, int]: + t = max(0.0, min(1.0, t)) + stops = [ + (0.00, (6, 8, 18)), + (0.20, (28, 20, 70)), + (0.43, (118, 29, 97)), + (0.67, (222, 72, 69)), + (0.85, (255, 164, 75)), + (1.00, (255, 243, 164)), + ] + for (a, ca), (b, cb) in zip(stops, stops[1:]): + if t <= b: + u = (t - a) / (b - a) + return tuple(round(ca[i] + (cb[i] - ca[i]) * u) for i in range(3)) + return stops[-1][1] + + +def downsample_cols(arr: np.ndarray, bins: int) -> np.ndarray: + if arr.shape[1] <= bins: + return arr + edges = np.linspace(0, arr.shape[1], bins + 1).round().astype(int) + out = np.zeros((arr.shape[0], bins), dtype=np.float32) + for i in range(bins): + lo, hi = edges[i], max(edges[i] + 1, edges[i + 1]) + out[:, i] = arr[:, lo:hi].mean(axis=1) + return out + + +def normalize(arr: np.ndarray, scale: float | None = None) -> tuple[np.ndarray, float]: + if scale is None: + scale = float(np.quantile(arr, 0.98)) if arr.size else 1.0 + if scale <= 0: + scale = 1.0 + return np.clip(arr / scale, 0, 1), scale + + +def draw_heatmap(draw: ImageDraw.ImageDraw, arr: np.ndarray, box: tuple[int, int, int, int], title: str, subtitle: str, color: tuple[int, int, int]) -> None: + x0, y0, x1, y1 = box + draw.rounded_rectangle(box, radius=22, fill=PALETTE["panel"], outline=(31, 42, 50), width=1) + draw.text((x0 + 24, y0 + 18), title, fill=color, font=ui_font(26, True)) + draw.text((x0 + 24, y0 + 50), subtitle, fill=PALETTE["muted"], font=ui_font(15)) + hx0, hy0, hx1, hy1 = x0 + 58, y0 + 84, x1 - 28, y1 - 44 + rows, cols = arr.shape + cw = (hx1 - hx0) / cols + ch = (hy1 - hy0) / rows + for r in range(rows): + y_a = round(hy0 + r * ch) + y_b = round(hy0 + (r + 1) * ch) + for c in range(cols): + x_a = round(hx0 + c * cw) + x_b = round(hx0 + (c + 1) * cw) + draw.rectangle((x_a, y_a, x_b, y_b), fill=heat_color(float(arr[r, c]))) + for r in range(0, rows, 4): + y = round(hy0 + (r + 0.5) * ch) + draw.text((x0 + 18, y - 8), str(r), fill=PALETTE["muted"], font=ui_font(12)) + draw.text((x0 + 16, hy0 - 4), "layer", fill=PALETTE["muted"], font=ui_font(12)) + draw.text((hx0, y1 - 31), "image token sequence →", fill=PALETTE["muted"], font=ui_font(13)) + + +def crop_with_box(img: Image.Image, start: int, end: int, cols: int, adv: int, pitch: int, pad_cells: int = 34) -> Image.Image: + row0 = max(0, start // cols - 5) + row1 = min(img.height // pitch, end // cols + 6) + col0 = max(0, start % cols - pad_cells) + col1 = min(cols, end % cols + pad_cells) + if col1 <= col0: + col1 = min(cols, col0 + 72) + x0, y0, x1, y1 = col0 * adv, row0 * pitch, col1 * adv, row1 * pitch + crop = img.crop((x0, y0, x1, y1)).convert("RGB") + d = ImageDraw.Draw(crop) + bx0 = max(0, (start % cols - col0) * adv - adv) + bx1 = min(crop.width - 1, ((end - 1) % cols - col0 + 2) * adv) + by0 = max(0, (start // cols - row0) * pitch - 1) + by1 = min(crop.height - 1, ((end - 1) // cols - row0 + 1) * pitch + 1) + d.rounded_rectangle((bx0, by0, bx1, by1), radius=3, outline=PALETTE["red"], width=3) + return crop + + +def paste_fit(canvas: Image.Image, img: Image.Image, box: tuple[int, int, int, int]) -> None: + x0, y0, x1, y1 = box + scale = min((x1 - x0) / img.width, (y1 - y0) / img.height) + resized = img.resize((max(1, round(img.width * scale)), max(1, round(img.height * scale))), Image.Resampling.NEAREST) + canvas.paste(resized, (x0 + (x1 - x0 - resized.width) // 2, y0 + (y1 - y0 - resized.height) // 2)) + + +def make_prompt(q: str, cols: int, rows: int) -> str: + return load_prompt("qa-image.md").format(cols=cols, rows=rows) + f"\n\nQuestion: {q}\nAnswer with only the shortest extractive answer." + + +def to_device(batch: dict[str, Any], device: Any) -> dict[str, Any]: + return {k: (v.to(device) if hasattr(v, "to") else v) for k, v in batch.items()} + + +def hidden_token_matrix(model: Any, processor: Any, image: Image.Image, prompt_text: str, device: Any) -> tuple[list[np.ndarray], list[int], dict[str, Any]]: + import torch + + messages = [{"role": "user", "content": [{"type": "image", "image": image}, {"type": "text", "text": prompt_text}]}] + templated = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + batch = processor(images=image, text=templated, return_tensors="pt") + image_token_id = processor.tokenizer.convert_tokens_to_ids(processor.image_token) + ids = batch["input_ids"][0].tolist() + image_positions = [i for i, token_id in enumerate(ids) if token_id == image_token_id] + meta = {k: (v.tolist() if hasattr(v, "tolist") else v) for k, v in batch.items() if k in ("image_grid_thw",)} + batch = to_device(batch, device) + with torch.no_grad(): + out = model(**batch, output_hidden_states=True, output_attentions=False, use_cache=False) + matrices: list[np.ndarray] = [] + for hidden in out.hidden_states: + token_hidden = hidden[0, image_positions, :].float().detach().cpu().numpy() + matrices.append(token_hidden.astype(np.float32, copy=False)) + return matrices, image_positions, meta + + +def render_tensor_card( + out_path: Path, + answer_heat: np.ndarray, + random_heat: np.ndarray, + ratio_heat: np.ndarray, + base_img: Image.Image, + answer_img: Image.Image, + record: dict[str, Any], + cols: int, + adv: int, + pitch: int, + summary: dict[str, Any], +) -> None: + w, h = 1900, 1180 + canvas = Image.new("RGB", (w, h), PALETTE["bg"]) + draw = ImageDraw.Draw(canvas) + for y in range(0, h, 14): + draw.line((0, y, w, y), fill=(8, 11 + (y % 9), 15 + (y % 13))) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-260, -180, 850, 640), fill=(255, 83, 62, 30)) + gd.ellipse((1080, 110, 2240, 1320), fill=(77, 218, 255, 30)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(80))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + draw.text((58, 38), "SNAPCOMPACT WHITEBOX", fill=PALETTE["amber"], font=ui_font(22, True)) + draw.text((58, 76), "The hidden-state scar of a missing answer", fill=PALETTE["ink"], font=ui_font(58, True)) + draw.text((60, 148), "Each pixel below is a decoder layer × image-token bin. Bright = larger ||hidden(original) − hidden(masked)||.", fill=PALETTE["muted"], font=ui_font(24)) + + # Left evidence panel. + draw.rounded_rectangle((58, 205, 700, 1098), radius=28, fill=PALETTE["panel"], outline=(31, 42, 50), width=1) + draw.text((90, 236), "the visual intervention", fill=PALETTE["ink"], font=ui_font(30, True)) + draw.text((90, 274), "same prompt, same bitmap; only answer cells blanked", fill=PALETTE["muted"], font=ui_font(17)) + crop = crop_with_box(base_img, record["answer_start"], record["answer_end"], cols, adv, pitch) + masked_crop = crop_with_box(answer_img, record["answer_start"], record["answer_end"], cols, adv, pitch) + draw.text((90, 326), "ORIGINAL", fill=PALETTE["cyan"], font=ui_font(16, True)) + draw.rounded_rectangle((90, 352, 668, 528), radius=14, fill=(244, 242, 230), outline=PALETTE["cyan"], width=3) + paste_fit(canvas, crop, (108, 368, 650, 512)) + draw.text((90, 568), "ANSWER ERASED", fill=PALETTE["red"], font=ui_font(16, True)) + draw.rounded_rectangle((90, 594, 668, 770), radius=14, fill=(244, 242, 230), outline=PALETTE["red"], width=3) + paste_fit(canvas, masked_crop, (108, 610, 650, 754)) + question = record["q"] + if len(question) > 72: + question = question[:69] + "…" + draw.text((90, 828), "question", fill=PALETTE["muted"], font=ui_font(16, True)) + draw.text((90, 856), question, fill=PALETTE["ink"], font=ui_font(21)) + draw.text((90, 914), "gold answer", fill=PALETTE["muted"], font=ui_font(16, True)) + draw.text((90, 942), str(record["answer_text"]), fill=PALETTE["amber"], font=ui_font(32, True)) + draw.text((90, 1014), f"{summary['layers']} hidden layers × {summary['image_tokens']} image tokens", fill=PALETTE["muted"], font=ui_font(18)) + + draw_heatmap(draw, answer_heat, (742, 205, 1818, 488), "gold answer mask", "activation delta when the true answer is blanked", PALETTE["red"]) + draw_heatmap(draw, random_heat, (742, 520, 1818, 803), "random equal-size mask", "control: blank the same number of glyph cells elsewhere", PALETTE["green"]) + draw_heatmap(draw, ratio_heat, (742, 835, 1818, 1098), "answer / random ratio", "bright bands mark layers/tokens more sensitive to the answer region", PALETTE["amber"]) + + # Color scale. + for i in range(220): + draw.rectangle((1588 + i, 158, 1589 + i, 174), fill=heat_color(i / 219)) + draw.text((1588, 133), "low", fill=PALETTE["muted"], font=ui_font(13)) + draw.text((1758, 133), "high", fill=PALETTE["muted"], font=ui_font(13)) + + out_path.parent.mkdir(parents=True, exist_ok=True) + canvas.save(out_path) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--model-dir", default=DEFAULT_MODEL_DIR) + ap.add_argument("--font", default="8x13", choices=sorted(FONTS)) + ap.add_argument("--variant", default="bw") + ap.add_argument("--size", type=int, default=768) + ap.add_argument("--limit-paras", type=int, default=40) + ap.add_argument("--qpc", type=int, default=16) + ap.add_argument("--question-index", type=int, default=7) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--bins", type=int, default=180) + ap.add_argument("--out", default="tensor-heatmap-paddleocr") + args = ap.parse_args() + + import torch + from transformers import AutoModel, AutoProcessor + + out_dir = HERE / "results" / args.out + img_dir = out_dir / "images" + img_dir.mkdir(parents=True, exist_ok=True) + CACHE.mkdir(exist_ok=True) + + cfg = FONTS[args.font] + cols, rows, budget = capacity(cfg, args.size) + paras = squad.load_paragraphs(CACHE)[: args.limit_paras] + flow, offsets = squad.build_flow(paras) + start, end = 0, min(len(flow), budget) + chunk = flow[start:end] + questions = sample_answer_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + raise SystemExit("no sampled questions fit in chunk") + q = questions[min(args.question_index, len(questions) - 1)] + + base_img = render(chunk, cfg, CACHE, args.size, args.variant) + fill = (255, 255, 255) if args.variant not in ("dark", "dark-sent") else (0, 0, 0) + span_len = max(1, q["answer_end"] - q["answer_start"]) + rng = random.Random(args.seed * 101 + args.question_index) + rand_start, rand_end = random_span(rng, len(chunk), span_len, q["answer_start"], q["answer_end"]) + answer_img = mask_cells(base_img, q["answer_start"], q["answer_end"], cols, cfg.adv, cfg.pitch, fill) + random_img = mask_cells(base_img, rand_start, rand_end, cols, cfg.adv, cfg.pitch, fill) + base_img.save(img_dir / "original.png") + answer_img.save(img_dir / "answer-mask.png") + random_img.save(img_dir / "random-mask.png") + + print(f"loading {args.model_dir}", flush=True) + processor = AutoProcessor.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True, use_fast=False) + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + dtype = torch.bfloat16 if device.type == "cuda" else torch.float32 + model = AutoModel.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True, dtype=dtype).to(device).eval() + + prompt = make_prompt(q["q"], cols, rows) + original, positions, meta = hidden_token_matrix(model, processor, base_img, prompt, device) + answer, answer_positions, _ = hidden_token_matrix(model, processor, answer_img, prompt, device) + random_mask, random_positions, _ = hidden_token_matrix(model, processor, random_img, prompt, device) + if positions != answer_positions or positions != random_positions: + raise SystemExit("image token positions changed across variants") + + answer_delta = np.stack([np.linalg.norm(a - b, axis=1) for a, b in zip(original, answer)], axis=0) + random_delta = np.stack([np.linalg.norm(a - b, axis=1) for a, b in zip(original, random_mask)], axis=0) + ratio = answer_delta / np.maximum(random_delta, 1e-6) + + answer_binned = downsample_cols(answer_delta, args.bins) + random_binned = downsample_cols(random_delta, args.bins) + ratio_binned = downsample_cols(ratio, args.bins) + common_scale = float(np.quantile(np.concatenate([answer_binned.ravel(), random_binned.ravel()]), 0.98)) + answer_norm, _ = normalize(answer_binned, common_scale) + random_norm, _ = normalize(random_binned, common_scale) + ratio_norm, ratio_scale = normalize(ratio_binned, float(np.quantile(ratio_binned, 0.98))) + + record = { + "q": q["q"], + "golds": q["golds"], + "answer_text": q["answer_text"], + "answer_start": q["answer_start"], + "answer_end": q["answer_end"], + "random_start": rand_start, + "random_end": rand_end, + } + summary = { + "args": vars(args), + "device": str(device), + "geometry": {"cols": cols, "rows": rows, "capacity": budget}, + "question": record, + "image_tokens": len(positions), + "layers": len(original), + "image_token_positions": {"first": positions[0], "last": positions[-1]}, + "processor_meta": meta, + "answer_delta_mean": float(answer_delta.mean()), + "random_delta_mean": float(random_delta.mean()), + "answer_over_random_delta": float(answer_delta.mean() / max(random_delta.mean(), 1e-6)), + "common_delta_scale_p98": common_scale, + "ratio_scale_p98": ratio_scale, + "max_ratio_layer": int(np.argmax(ratio.mean(axis=1))), + "mean_ratio_by_layer": [float(x) for x in ratio.mean(axis=1)], + } + + np.savez_compressed( + out_dir / "heatmaps.npz", + answer_delta=answer_delta, + random_delta=random_delta, + ratio=ratio, + answer_binned=answer_binned, + random_binned=random_binned, + ratio_binned=ratio_binned, + answer_norm=answer_norm, + random_norm=random_norm, + ratio_norm=ratio_norm, + ) + (out_dir / "summary.json").write_text(json.dumps(summary, indent=1)) + render_tensor_card(out_dir / "tensor-heatmap.png", answer_norm, random_norm, ratio_norm, base_img, answer_img, record, cols, cfg.adv, cfg.pitch, summary) + print(json.dumps(summary, indent=1)) + print(f"results -> {out_dir}") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_text_image_3d_viz.py b/packages/snapcompact/research/snapcompact_text_image_3d_viz.py new file mode 100644 index 000000000..12723a4f4 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_text_image_3d_viz.py @@ -0,0 +1,211 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["matplotlib", "numpy", "pillow"] +# /// +"""Render a 3D text-vs-image activation comparison from paired carrier data.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +from matplotlib import cm +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +BG = (5, 7, 10) +PANEL = (12, 17, 23) +INK = (241, 239, 224) +MUTED = (143, 154, 160) +CYAN = (75, 220, 255) +ORANGE = (255, 112, 72) +AMBER = (255, 196, 68) +GREEN = (148, 255, 117) + + +def font(size: int, bold: bool = False): + for path in [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def mono(size: int): + for path in ["/System/Library/Fonts/Monaco.ttf", "/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf"]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def downsample(arr: np.ndarray, cols: int) -> np.ndarray: + if arr.shape[1] <= cols: + return arr + edges = np.linspace(0, arr.shape[1], cols + 1).round().astype(int) + out = np.zeros((arr.shape[0], cols), dtype=np.float32) + for i in range(cols): + lo = edges[i] + hi = max(lo + 1, edges[i + 1]) + out[:, i] = arr[:, lo:hi].mean(axis=1) + return out + + +def normalize(arr: np.ndarray) -> np.ndarray: + lo = float(np.quantile(arr, 0.03)) + hi = float(np.quantile(arr, 0.985)) + if hi <= lo: + hi = lo + 1e-6 + return np.clip((arr - lo) / (hi - lo), 0, 1) + + +def render_surface(z: np.ndarray, answer_bins: list[int]) -> Image.Image: + fig = plt.figure(figsize=(14.5, 8.2), dpi=180) + fig.patch.set_facecolor("#05070a") + ax = fig.add_subplot(111, projection="3d") + ax.set_facecolor((0.02, 0.025, 0.035, 1)) + y = np.arange(z.shape[0]) + x = np.arange(z.shape[1]) + X, Y = np.meshgrid(x, y) + cmap = plt.colormaps["turbo"] + ax.plot_surface(X, Y, z, facecolors=cmap(z), linewidth=0, antialiased=True, shade=False, alpha=0.98) + ax.contour(X, Y, z, zdir="z", offset=-0.08, levels=12, cmap=cmap, linewidths=0.9, alpha=0.75) + for b in answer_bins: + if 0 <= b < z.shape[1]: + ax.plot([b, b], [0, z.shape[0] - 1], [1.08, 1.08], color="#ff7048", linewidth=2.6, alpha=0.78) + ax.plot([b, b], [0, z.shape[0] - 1], [-0.06, -0.06], color="#ff7048", linewidth=1.6, alpha=0.55) + ax.view_init(elev=32, azim=-58) + ax.set_box_aspect((3.2, 0.8, 0.72)) + ax.set_zlim(-0.08, 1.08) + ax.set_ylim(z.shape[0] - 1, 0) + ax.set_xlabel("image-token bins", color="#8f9aa0", labelpad=10) + ax.set_ylabel("decoder layer", color="#8f9aa0", labelpad=10) + ax.set_zlabel("excess cosine", color="#8f9aa0", labelpad=8) + ax.tick_params(colors="#8f9aa0", labelsize=8) + for axis in (ax.xaxis, ax.yaxis, ax.zaxis): + axis.pane.set_facecolor((0.02, 0.025, 0.035, 0.0)) + axis._axinfo["grid"]["color"] = (0.35, 0.45, 0.50, 0.18) + ax.set_title("text-answer vector ↔ image-token field", color="#efeede", fontsize=24, fontweight="bold", loc="left", pad=18) + tmp = HERE / "results" / ".text-image-3d-panel.png" + fig.savefig(tmp, facecolor=fig.get_facecolor(), transparent=False) + plt.close(fig) + img = Image.open(tmp).convert("RGB") + tmp.unlink(missing_ok=True) + return img + + +def crop_answer(img: Image.Image, q: dict, cols: int, adv: int = 8, pitch: int = 13) -> Image.Image: + start = q["answer_start"] + end = q["answer_end"] + row0 = max(0, start // cols - 5) + row1 = min(img.height // pitch, end // cols + 6) + col0 = max(0, start % cols - 34) + col1 = min(cols, end % cols + 34) + crop = img.crop((col0 * adv, row0 * pitch, col1 * adv, row1 * pitch)).convert("RGB") + d = ImageDraw.Draw(crop) + bx0 = max(0, (start % cols - col0) * adv - adv) + bx1 = min(crop.width - 1, ((end - 1) % cols - col0 + 2) * adv) + by0 = max(0, (start // cols - row0) * pitch - 1) + by1 = min(crop.height - 1, ((end - 1) // cols - row0 + 1) * pitch + 1) + d.rounded_rectangle((bx0, by0, bx1, by1), radius=3, outline=ORANGE, width=3) + return crop + + +def paste_fit(canvas: Image.Image, img: Image.Image, box: tuple[int, int, int, int]) -> None: + x0, y0, x1, y1 = box + scale = min((x1 - x0) / img.width, (y1 - y0) / img.height) + resized = img.resize((max(1, round(img.width * scale)), max(1, round(img.height * scale))), Image.Resampling.NEAREST) + canvas.paste(resized, (x0 + (x1 - x0 - resized.width) // 2, y0 + (y1 - y0 - resized.height) // 2)) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--result-dir", default=str(HERE / "results" / "text-image-compare-paddleocr-q7")) + ap.add_argument("--out", default=str(HERE / "results" / "text-image-compare-paddleocr-q7" / "text-vs-image-3d.png")) + ap.add_argument("--bins", type=int, default=150) + args = ap.parse_args() + result_dir = Path(args.result_dir) + summary = json.loads((result_dir / "summary.json").read_text()) + data = np.load(result_dir / "text_image_compare.npz") + raw = data["text_answer_to_image_excess"] if "text_answer_to_image_excess" in data else data["text_answer_to_image_cosine"] + z = normalize(downsample(raw, args.bins)) + token_count = summary["image_tokens"] + answer_bins = sorted({round(idx / max(1, token_count - 1) * (args.bins - 1)) for idx in summary["image_answer_token_indices"]}) + panel = render_surface(z, answer_bins) + + w, h = 2200, 1320 + canvas = Image.new("RGB", (w, h), BG) + draw = ImageDraw.Draw(canvas) + for y in range(0, h, 16): + draw.line((0, y, w, y), fill=(7, 10 + y % 9, 15 + y % 11)) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-260, -220, 860, 680), fill=(75, 220, 255, 30)) + gd.ellipse((1160, 80, 2440, 1320), fill=(255, 112, 72, 28)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(84))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + q = summary["question"] + draw.text((64, 42), "TEXT ↔ IMAGE WHITEBOX", fill=AMBER, font=font(24, True)) + draw.text((64, 84), "Same input, different carrier, shared hidden space", fill=INK, font=font(61, True)) + draw.text((66, 166), "For every decoder layer, compare the raw-text answer state against all bitmap image-token states. Peaks = image regions whose hidden state becomes text-like.", fill=MUTED, font=font(24)) + + draw.rounded_rectangle((64, 238, 618, 1234), radius=30, fill=PANEL, outline=(35, 49, 59), width=1) + draw.text((96, 270), "two carriers", fill=INK, font=font(34, True)) + draw.text((96, 312), "same chunk + same question", fill=MUTED, font=font(18)) + draw.text((96, 366), "RAW TEXT", fill=CYAN, font=font(18, True)) + y = 402 + draw.text((96, y), "Question:", fill=MUTED, font=font(17, True)) + y += 30 + for line in [q["q"][i : i + 46] for i in range(0, len(q["q"]), 46)]: + draw.text((96, y), line, fill=INK, font=font(19)) + y += 26 + y += 24 + draw.text((96, y), "Gold answer token span:", fill=MUTED, font=font(17, True)) + y += 32 + answer_text = str(q["answer_text"]) + draw.rounded_rectangle((96, y, 108 + max(72, len(answer_text) * 24), y + 40), radius=7, fill=AMBER) + draw.text((108, y + 7), answer_text, fill=(5, 7, 10), font=mono(22)) + y += 66 + draw.text((96, y), "The raw-text run receives the same", fill=INK, font=font(18)) + draw.text((96, y + 28), "SQuAD passage as ordinary tokens;", fill=INK, font=font(18)) + draw.text((96, y + 56), "the image run receives the passage", fill=INK, font=font(18)) + draw.text((96, y + 84), "only through the bitmap carrier.", fill=INK, font=font(18)) + draw.text((96, 674), f"text reference: {summary['text_reference_tokens']} tokens", fill=MUTED, font=font(18)) + draw.text((96, 704), f"answer span: {summary['text_answer_tokens']} text tokens", fill=MUTED, font=font(18)) + + draw.text((96, 774), "SNAPCOMPACT IMAGE", fill=ORANGE, font=font(18, True)) + img = Image.open(result_dir / "images" / "image-carrier.png").convert("RGB") + crop = crop_answer(img, q, summary["geometry"]["cols"]) + draw.rounded_rectangle((96, 812, 586, 1052), radius=16, fill=(244, 242, 230), outline=ORANGE, width=3) + paste_fit(canvas, crop, (112, 828, 570, 1036)) + draw.text((96, 1092), f"image field: {summary['image_tokens']} tokens ({summary['image_grid']}×{summary['image_grid']})", fill=MUTED, font=font(18)) + draw.text((96, 1138), f"peak alignment: {summary['answer_region_cosine_max']:.3f} @ layer {summary['answer_region_cosine_argmax']}", fill=AMBER, font=font(22, True)) + draw.text((96, 1172), f"final alignment: {summary['answer_region_cosine_final']:.3f}", fill=MUTED, font=font(19)) + + draw.rounded_rectangle((650, 238, 2134, 1234), radius=30, fill=PANEL, outline=(35, 49, 59), width=1) + draw.text((686, 270), "3D cross-carrier resonance terrain", fill=INK, font=font(36, True)) + draw.text((686, 314), "z-axis = excess cosine after subtracting each layer's median image-token similarity; orange rails mark the bitmap answer region", fill=MUTED, font=font(20)) + panel = panel.resize((1408, 794), Image.Resampling.LANCZOS) + canvas.paste(panel, (692, 378)) + cmap = plt.colormaps["turbo"] + for i in range(280): + rgb = tuple(int(v * 255) for v in cmap(i / 279)[:3]) + draw.rectangle((1790 + i, 282, 1791 + i, 300), fill=rgb) + draw.text((1790, 254), "low excess", fill=MUTED, font=font(14)) + draw.text((1992, 254), "high excess", fill=MUTED, font=font(14)) + + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + canvas.save(out) + print(out) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_text_image_compare.py b/packages/snapcompact/research/snapcompact_text_image_compare.py new file mode 100644 index 000000000..000e46d81 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_text_image_compare.py @@ -0,0 +1,489 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy", "torch", "transformers", "sentencepiece", "protobuf", "einops", "matplotlib"] +# /// +"""Compare raw-text vs snapcompact-image activations for the same input. + +The experiment feeds the same SQuAD chunk/question through a local VLM twice: +1. as ordinary raw text in a block +2. as a snapcompact bitmap plus the same question + +It then compares the text-carrier answer vector against every image-token vector +by decoder layer, producing a blog visual of cross-modal alignment. +""" + +from __future__ import annotations + +import argparse +import json +import math +import random +import sys +from pathlib import Path +from typing import Any + +import numpy as np +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import capacity, render # noqa: E402 +from run import CACHE, FONTS, load_prompt # noqa: E402 +from snapcompact_blackbox_occlusion import sample_answer_questions # noqa: E402 + +DEFAULT_MODEL_DIR = ( + "/home/can/.cache/huggingface/hub/models--PaddlePaddle--PaddleOCR-VL/" + "snapshots/2b77538ef936207f60c16b45082841068987d08c" +) + +PALETTE = { + "bg": (5, 7, 10), + "panel": (13, 18, 24), + "panel2": (9, 13, 18), + "ink": (241, 239, 224), + "muted": (143, 154, 160), + "cyan": (75, 220, 255), + "orange": (255, 112, 72), + "green": (148, 255, 117), + "amber": (255, 196, 68), + "purple": (180, 96, 255), + "grid": (38, 49, 58), +} + + +def ui_font(size: int, bold: bool = False) -> ImageFont.FreeTypeFont | ImageFont.ImageFont: + candidates = [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/System/Library/Fonts/Monaco.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ] + for path in candidates: + if path and Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def mono_font(size: int) -> ImageFont.FreeTypeFont | ImageFont.ImageFont: + candidates = [ + "/System/Library/Fonts/Monaco.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf", + ] + for path in candidates: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def heat_color(t: float) -> tuple[int, int, int]: + t = max(0.0, min(1.0, t)) + stops = [ + (0.00, (4, 7, 20)), + (0.18, (22, 24, 80)), + (0.38, (62, 68, 168)), + (0.58, (38, 183, 208)), + (0.78, (160, 250, 145)), + (1.00, (255, 245, 166)), + ] + for (a, ca), (b, cb) in zip(stops, stops[1:]): + if t <= b: + u = (t - a) / (b - a) + return tuple(round(ca[i] + (cb[i] - ca[i]) * u) for i in range(3)) + return stops[-1][1] + + +def cosine(a: np.ndarray, b: np.ndarray) -> np.ndarray: + a_norm = np.linalg.norm(a, axis=-1, keepdims=True) + b_norm = np.linalg.norm(b, axis=-1, keepdims=True) + return (a * b).sum(axis=-1) / np.maximum((a_norm * b_norm).squeeze(-1), 1e-6) + + +def normalize_heat(arr: np.ndarray, lo: float | None = None, hi: float | None = None) -> tuple[np.ndarray, float, float]: + if lo is None: + lo = float(np.quantile(arr, 0.03)) + if hi is None: + hi = float(np.quantile(arr, 0.98)) + if hi <= lo: + hi = lo + 1e-6 + return np.clip((arr - lo) / (hi - lo), 0, 1), lo, hi + + +def apply_template(processor: Any, content: list[dict[str, Any]]) -> str: + return processor.apply_chat_template([{"role": "user", "content": content}], tokenize=False, add_generation_prompt=True) + + +def text_spans(processor: Any, templated: str, chunk: str, answer_start: int, answer_end: int) -> dict[str, int]: + tokenizer = processor.tokenizer + chunk_at = templated.index(chunk) + prefix = templated[:chunk_at] + def n_tokens(s: str) -> int: + return len(tokenizer(s, add_special_tokens=False)["input_ids"]) + ref_start = n_tokens(prefix) + ref_end = n_tokens(prefix + chunk) + answer_tok_start = n_tokens(prefix + chunk[:answer_start]) + answer_tok_end = max(answer_tok_start + 1, n_tokens(prefix + chunk[:answer_end])) + return { + "ref_start": ref_start, + "ref_end": ref_end, + "answer_start": answer_tok_start, + "answer_end": answer_tok_end, + } + + +def to_device(batch: dict[str, Any], device: Any) -> dict[str, Any]: + return {k: (v.to(device) if hasattr(v, "to") else v) for k, v in batch.items()} + + +def run_text(model: Any, processor: Any, text_prompt: str, chunk: str, answer_start: int, answer_end: int, device: Any) -> tuple[list[np.ndarray], dict[str, int], str]: + import torch + + templated = apply_template(processor, [{"type": "text", "text": text_prompt}]) + spans = text_spans(processor, templated, chunk, answer_start, answer_end) + batch = processor(text=templated, return_tensors="pt") + batch = to_device(batch, device) + with torch.no_grad(): + out = model(**batch, output_hidden_states=True, output_attentions=False, use_cache=False) + layers = [h[0].float().detach().cpu().numpy().astype(np.float32, copy=False) for h in out.hidden_states] + return layers, spans, templated + + +def run_image(model: Any, processor: Any, img: Image.Image, img_prompt: str, device: Any) -> tuple[list[np.ndarray], list[int], dict[str, Any], str]: + import torch + + templated = apply_template(processor, [{"type": "image", "image": img}, {"type": "text", "text": img_prompt}]) + batch = processor(images=img, text=templated, return_tensors="pt") + image_token_id = processor.tokenizer.convert_tokens_to_ids(processor.image_token) + image_positions = [i for i, token_id in enumerate(batch["input_ids"][0].tolist()) if token_id == image_token_id] + meta = {k: (v.tolist() if hasattr(v, "tolist") else v) for k, v in batch.items() if k in ("image_grid_thw",)} + batch = to_device(batch, device) + with torch.no_grad(): + out = model(**batch, output_hidden_states=True, output_attentions=False, use_cache=False) + layers = [h[0].float().detach().cpu().numpy().astype(np.float32, copy=False) for h in out.hidden_states] + return layers, image_positions, meta, templated + + +def image_answer_token_indices(answer_start: int, answer_end: int, text_cols: int, adv: int, pitch: int, image_w: int, image_h: int, image_token_count: int) -> list[int]: + grid = round(math.sqrt(image_token_count)) + if grid * grid != image_token_count: + return [] + row0 = max(0, answer_start // text_cols) + row1 = max(row0, (answer_end - 1) // text_cols) + col0 = max(0, answer_start % text_cols) + col1 = max(col0, (answer_end - 1) % text_cols) + x0 = max(0, col0 * adv - adv) + x1 = min(image_w, (col1 + 2) * adv) + y0 = max(0, row0 * pitch - 1) + y1 = min(image_h, (row1 + 1) * pitch + 1) + gx0 = max(0, min(grid - 1, int(x0 / image_w * grid))) + gx1 = max(0, min(grid - 1, int(math.ceil(x1 / image_w * grid)))) + gy0 = max(0, min(grid - 1, int(y0 / image_h * grid))) + gy1 = max(0, min(grid - 1, int(math.ceil(y1 / image_h * grid)))) + out: list[int] = [] + for gy in range(gy0, gy1 + 1): + for gx in range(gx0, gx1 + 1): + out.append(gy * grid + gx) + return sorted(set(out)) + + +def crop_answer(img: Image.Image, start: int, end: int, cols: int, adv: int, pitch: int, pad_cells: int = 34) -> Image.Image: + row0 = max(0, start // cols - 5) + row1 = min(img.height // pitch, end // cols + 6) + col0 = max(0, start % cols - pad_cells) + col1 = min(cols, end % cols + pad_cells) + if col1 <= col0: + col1 = min(cols, col0 + 72) + crop = img.crop((col0 * adv, row0 * pitch, col1 * adv, row1 * pitch)).convert("RGB") + d = ImageDraw.Draw(crop) + bx0 = max(0, (start % cols - col0) * adv - adv) + bx1 = min(crop.width - 1, ((end - 1) % cols - col0 + 2) * adv) + by0 = max(0, (start // cols - row0) * pitch - 1) + by1 = min(crop.height - 1, ((end - 1) // cols - row0 + 1) * pitch + 1) + d.rounded_rectangle((bx0, by0, bx1, by1), radius=3, outline=PALETTE["orange"], width=3) + return crop + + +def draw_wrapped(draw: ImageDraw.ImageDraw, xy: tuple[int, int], text: str, width_chars: int, line_height: int, fill: tuple[int, int, int], fnt: ImageFont.ImageFont) -> int: + words = text.split() + lines: list[str] = [] + current = "" + for word in words: + cand = word if not current else current + " " + word + if len(cand) <= width_chars: + current = cand + else: + if current: + lines.append(current) + current = word + if current: + lines.append(current) + x, y = xy + for line in lines: + draw.text((x, y), line, fill=fill, font=fnt) + y += line_height + return y + + +def render_heat_grid(draw: ImageDraw.ImageDraw, grid: np.ndarray, box: tuple[int, int, int, int], title: str, layer: int, answer_indices: list[int], color: tuple[int, int, int]) -> None: + x0, y0, x1, y1 = box + draw.rounded_rectangle(box, radius=18, fill=PALETTE["panel2"], outline=(35, 49, 59), width=1) + draw.text((x0 + 18, y0 + 14), title, fill=color, font=ui_font(22, True)) + draw.text((x0 + 18, y0 + 43), f"decoder layer {layer}", fill=PALETTE["muted"], font=ui_font(15)) + gx0, gy0, gx1, gy1 = x0 + 26, y0 + 76, x1 - 26, y1 - 24 + rows, cols = grid.shape + cw = (gx1 - gx0) / cols + ch = (gy1 - gy0) / rows + for r in range(rows): + for c in range(cols): + xa = round(gx0 + c * cw) + xb = round(gx0 + (c + 1) * cw) + ya = round(gy0 + r * ch) + yb = round(gy0 + (r + 1) * ch) + draw.rectangle((xa, ya, xb, yb), fill=heat_color(float(grid[r, c]))) + for idx in answer_indices: + r, c = divmod(idx, cols) + xa = round(gx0 + c * cw) + xb = round(gx0 + (c + 1) * cw) + ya = round(gy0 + r * ch) + yb = round(gy0 + (r + 1) * ch) + draw.rectangle((xa - 2, ya - 2, xb + 2, yb + 2), outline=PALETTE["orange"], width=2) + + +def render_visual(out_path: Path, summary: dict[str, Any], arrays: dict[str, np.ndarray], original_img: Image.Image, chunk: str) -> None: + w, h = 2100, 1260 + canvas = Image.new("RGB", (w, h), PALETTE["bg"]) + draw = ImageDraw.Draw(canvas) + for y in range(0, h, 16): + draw.line((0, y, w, y), fill=(7, 10 + y % 9, 15 + y % 11)) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-240, -220, 860, 660), fill=(75, 220, 255, 28)) + gd.ellipse((1160, 80, 2420, 1320), fill=(255, 112, 72, 26)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(80))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + q = summary["question"] + draw.text((62, 42), "SNAPCOMPACT CARRIER COMPARISON", fill=PALETTE["amber"], font=ui_font(24, True)) + draw.text((62, 82), "Same input, two internal languages", fill=PALETTE["ink"], font=ui_font(68, True)) + draw.text((64, 166), "Raw text tokens vs bitmap image tokens. Bright fields show where the text-carrier answer vector resonates with the image-carrier hidden state.", fill=PALETTE["muted"], font=ui_font(25)) + + # Carrier cards. + draw.rounded_rectangle((62, 236, 620, 760), radius=28, fill=PALETTE["panel"], outline=(35, 49, 59), width=1) + draw.text((94, 270), "raw text carrier", fill=PALETTE["cyan"], font=ui_font(30, True)) + start = max(0, q["answer_start"] - 230) + end = min(len(chunk), q["answer_end"] + 230) + snippet = chunk[start:end].replace("\n", " ") + rel_a = q["answer_start"] - start + rel_b = q["answer_end"] - start + before = snippet[:rel_a] + answer = snippet[rel_a:rel_b] + after = snippet[rel_b:] + tx, ty = 94, 328 + ty = draw_wrapped(draw, (tx, ty), before[-260:], 52, 22, PALETTE["ink"], mono_font(15)) + draw.rounded_rectangle((tx, ty + 2, tx + 16 * max(3, len(answer)), ty + 27), radius=5, fill=(255, 196, 68)) + draw.text((tx + 4, ty + 5), answer, fill=(8, 10, 10), font=mono_font(16)) + ty += 36 + draw_wrapped(draw, (tx, ty), after[:260], 52, 22, PALETTE["ink"], mono_font(15)) + draw.text((94, 694), f"answer tokens: {summary['text_answer_tokens']}", fill=PALETTE["muted"], font=ui_font(18)) + draw.text((94, 724), f"reference tokens: {summary['text_reference_tokens']}", fill=PALETTE["muted"], font=ui_font(18)) + + draw.rounded_rectangle((62, 792, 620, 1192), radius=28, fill=PALETTE["panel"], outline=(35, 49, 59), width=1) + draw.text((94, 826), "image carrier", fill=PALETTE["orange"], font=ui_font(30, True)) + crop = crop_answer(original_img, q["answer_start"], q["answer_end"], summary["geometry"]["cols"], 8, 13) + scale = min(478 / crop.width, 218 / crop.height) + crop_r = crop.resize((round(crop.width * scale), round(crop.height * scale)), Image.Resampling.NEAREST) + draw.rounded_rectangle((94, 888, 588, 1134), radius=16, fill=(244, 242, 230), outline=PALETTE["orange"], width=3) + canvas.paste(crop_r, (94 + (494 - crop_r.width) // 2, 888 + (246 - crop_r.height) // 2)) + draw.text((94, 1150), f"image tokens: {summary['image_tokens']} ({summary['image_grid']}×{summary['image_grid']})", fill=PALETTE["muted"], font=ui_font(18)) + + # Layer grids. + sim = arrays["text_answer_to_image_excess_norm"] + grid = summary["image_grid"] + answer_indices = summary["image_answer_token_indices"] + layers = summary["selected_layers"] + boxes = [(672, 236, 1088, 626), (1118, 236, 1534, 626), (1564, 236, 1980, 626)] + names = ["input layer", "middle layer", "peak alignment"] + colors = [PALETTE["cyan"], PALETTE["purple"], PALETTE["green"]] + for layer, box, name, color in zip(layers, boxes, names, colors): + render_heat_grid(draw, sim[layer].reshape(grid, grid), box, name, layer, answer_indices, color) + + # Cosine bridge panel. + draw.rounded_rectangle((672, 672, 1980, 1192), radius=28, fill=PALETTE["panel"], outline=(35, 49, 59), width=1) + draw.text((704, 704), "cross-carrier convergence bridge", fill=PALETTE["ink"], font=ui_font(34, True)) + draw.text((704, 744), "Cosine similarity between pooled raw-text answer states and pooled bitmap answer-region states by layer", fill=PALETTE["muted"], font=ui_font(19)) + x0, y0, x1, y1 = 730, 820, 1908, 1096 + for i in range(5): + y = y0 + round((y1 - y0) * i / 4) + draw.line((x0, y, x1, y), fill=PALETTE["grid"], width=1) + local = arrays["answer_region_cosine"] + global_mean = arrays["global_mean_cosine"] + lo = float(min(local.min(), global_mean.min())) + hi = float(max(local.max(), global_mean.max())) + if hi <= lo: + hi = lo + 1e-6 + def pts(vals: np.ndarray) -> list[tuple[int, int]]: + out = [] + for i, v in enumerate(vals): + x = x0 + round((x1 - x0) * i / max(1, len(vals) - 1)) + y = y1 - round((y1 - y0) * (float(v) - lo) / (hi - lo)) + out.append((x, y)) + return out + p_local = pts(local) + p_global = pts(global_mean) + draw.line(p_global, fill=PALETTE["muted"], width=4) + draw.line(p_local, fill=PALETTE["amber"], width=6) + for x, y in p_local: + draw.ellipse((x - 5, y - 5, x + 5, y + 5), fill=PALETTE["amber"]) + draw.text((x0, y1 + 22), "layer 0", fill=PALETTE["muted"], font=ui_font(16)) + draw.text((x1 - 70, y1 + 22), f"layer {len(local) - 1}", fill=PALETTE["muted"], font=ui_font(16)) + peak_layer = int(np.argmax(local)) + draw.rounded_rectangle((1502, 790, 1938, 900), radius=18, fill=(9, 13, 18), outline=(38, 51, 60), width=1) + draw.text((1526, 812), f"answer cosine peaks: {local[peak_layer]:.3f} @L{peak_layer}", fill=PALETTE["amber"], font=ui_font(21, True)) + draw.text((1526, 842), f"final answer cosine: {local[-1]:.3f}", fill=PALETTE["muted"], font=ui_font(18)) + draw.text((1526, 868), f"final global carrier cosine: {global_mean[-1]:.3f}", fill=PALETTE["muted"], font=ui_font(18)) + draw.rounded_rectangle((704, 1120, 1238, 1168), radius=13, fill=(9, 13, 18), outline=(38, 51, 60), width=1) + draw.rectangle((724, 1138, 768, 1148), fill=PALETTE["amber"]) + draw.text((784, 1129), "answer region: text vector ↔ image region", fill=PALETTE["muted"], font=ui_font(17)) + draw.rectangle((1260, 1138, 1304, 1148), fill=PALETTE["muted"]) + draw.text((1320, 1129), "global carrier means", fill=PALETTE["muted"], font=ui_font(17)) + + out_path.parent.mkdir(parents=True, exist_ok=True) + canvas.save(out_path) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--model-dir", default=DEFAULT_MODEL_DIR) + ap.add_argument("--font", default="8x13", choices=sorted(FONTS)) + ap.add_argument("--variant", default="bw") + ap.add_argument("--size", type=int, default=768) + ap.add_argument("--limit-paras", type=int, default=40) + ap.add_argument("--qpc", type=int, default=16) + ap.add_argument("--question-index", type=int, default=7) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--out", default="text-image-compare-paddleocr-q7") + args = ap.parse_args() + + import torch + from transformers import AutoConfig, AutoModel, AutoProcessor + + out_dir = HERE / "results" / args.out + img_dir = out_dir / "images" + img_dir.mkdir(parents=True, exist_ok=True) + CACHE.mkdir(exist_ok=True) + + cfg = FONTS[args.font] + cols, rows, budget = capacity(cfg, args.size) + paras = squad.load_paragraphs(CACHE)[: args.limit_paras] + flow, offsets = squad.build_flow(paras) + start, end = 0, min(len(flow), budget) + chunk = flow[start:end] + questions = sample_answer_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + raise SystemExit("no sampled questions fit in chunk") + q = questions[min(args.question_index, len(questions) - 1)] + img = render(chunk, cfg, CACHE, args.size, args.variant) + img.save(img_dir / "image-carrier.png") + + print(f"loading {args.model_dir}", flush=True) + config = AutoConfig.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True) + processor = AutoProcessor.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True, use_fast=False) + target_device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + dtype = torch.bfloat16 if target_device.type == "cuda" else torch.float32 + if getattr(config, "model_type", "") == "qwen2_5_vl": + from transformers import Qwen2_5_VLForConditionalGeneration + + model = Qwen2_5_VLForConditionalGeneration.from_pretrained( + args.model_dir, + local_files_only=True, + trust_remote_code=True, + dtype=dtype, + device_map="auto" if target_device.type == "cuda" else None, + ).eval() + device = next(model.parameters()).device + else: + model = AutoModel.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True, dtype=dtype).to(target_device).eval() + device = target_device + + text_prompt = ( + "Below is reference material. Answer the question using only it.\n\n" + f"{chunk}\n\nQuestion: {q['q']}\n" + "Answer with only the shortest extractive answer." + ) + img_prompt = load_prompt("qa-image.md").format(cols=cols, rows=rows) + f"\n\nQuestion: {q['q']}\nAnswer with only the shortest extractive answer." + + text_layers, text_pos, text_template = run_text(model, processor, text_prompt, chunk, q["answer_start"], q["answer_end"], device) + image_layers, image_positions, image_meta, image_template = run_image(model, processor, img, img_prompt, device) + image_token_count = len(image_positions) + image_grid = round(math.sqrt(image_token_count)) + answer_image_indices = image_answer_token_indices(q["answer_start"], q["answer_end"], cols, cfg.adv, cfg.pitch, img.width, img.height, image_token_count) + + answer_cos = [] + global_cos = [] + text_answer_to_image = [] + for text_h, image_h in zip(text_layers, image_layers): + text_ref = text_h[text_pos["ref_start"] : text_pos["ref_end"]] + text_ans = text_h[text_pos["answer_start"] : text_pos["answer_end"]] + image_tokens = image_h[image_positions] + image_ans = image_tokens[answer_image_indices] if answer_image_indices else image_tokens + text_ans_mean = text_ans.mean(axis=0) + image_ans_mean = image_ans.mean(axis=0) + text_ref_mean = text_ref.mean(axis=0) + image_mean = image_tokens.mean(axis=0) + answer_cos.append(float(cosine(text_ans_mean[None, :], image_ans_mean[None, :])[0])) + global_cos.append(float(cosine(text_ref_mean[None, :], image_mean[None, :])[0])) + sims = cosine(np.repeat(text_ans_mean[None, :], image_tokens.shape[0], axis=0), image_tokens) + text_answer_to_image.append(sims.astype(np.float32, copy=False)) + + text_answer_to_image_arr = np.stack(text_answer_to_image, axis=0) + layer_baseline = np.median(text_answer_to_image_arr, axis=1, keepdims=True) + text_answer_to_image_excess = text_answer_to_image_arr - layer_baseline + normed, heat_lo, heat_hi = normalize_heat(text_answer_to_image_excess) + answer_cos_arr = np.array(answer_cos, dtype=np.float32) + global_cos_arr = np.array(global_cos, dtype=np.float32) + selected_layers = [0, len(text_layers) // 2, int(answer_cos_arr.argmax())] + + summary = { + "args": vars(args), + "device": str(device), + "geometry": {"cols": cols, "rows": rows, "capacity": budget}, + "question": { + "q": q["q"], + "golds": q["golds"], + "answer_text": q["answer_text"], + "answer_start": q["answer_start"], + "answer_end": q["answer_end"], + }, + "layers": len(text_layers), + "image_tokens": image_token_count, + "image_grid": image_grid, + "image_answer_token_indices": answer_image_indices, + "image_meta": image_meta, + "text_positions": text_pos, + "text_reference_tokens": text_pos["ref_end"] - text_pos["ref_start"], + "text_answer_tokens": text_pos["answer_end"] - text_pos["answer_start"], + "selected_layers": selected_layers, + "answer_region_cosine_final": float(answer_cos_arr[-1]), + "global_mean_cosine_final": float(global_cos_arr[-1]), + "answer_region_cosine_max": float(answer_cos_arr.max()), + "answer_region_cosine_argmax": int(answer_cos_arr.argmax()), + "heat_normalization": {"lo_p03": heat_lo, "hi_p98": heat_hi}, + "text_template_prefix": text_template[:240], + "image_template_prefix": image_template[:240], + } + arrays = { + "text_answer_to_image_cosine": text_answer_to_image_arr, + "text_answer_to_image_excess": text_answer_to_image_excess, + "text_answer_to_image_excess_norm": normed, + "answer_region_cosine": answer_cos_arr, + "global_mean_cosine": global_cos_arr, + } + np.savez_compressed(out_dir / "text_image_compare.npz", **arrays) + (out_dir / "summary.json").write_text(json.dumps(summary, indent=1)) + render_visual(out_dir / "text-vs-image.png", summary, arrays, img, chunk) + print(json.dumps(summary, indent=1)) + print(f"results -> {out_dir}") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_token_entry_dump.py b/packages/snapcompact/research/snapcompact_token_entry_dump.py new file mode 100644 index 000000000..5d9516192 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_token_entry_dump.py @@ -0,0 +1,158 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy", "torch", "transformers"] +# /// +"""Dump how the same content enters Qwen as text tokens vs visual tokens. + +Produces a JSON with real tokenizer output (token strings + ids), the real +embedding rows entering the decoder for the answer-word text tokens, and the +real visual-tower output vectors for the image tokens covering the same word. +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path +from typing import Any + +import numpy as np + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import capacity, render # noqa: E402 +from run import CACHE, FONTS # noqa: E402 +from snapcompact_blackbox_occlusion import sample_answer_questions # noqa: E402 +from snapcompact_text_image_compare import image_answer_token_indices # noqa: E402 + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--model-dir", default="Qwen/Qwen2.5-VL-7B-Instruct") + ap.add_argument("--font", default="8x13", choices=sorted(FONTS)) + ap.add_argument("--variant", default="bw") + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--limit-paras", type=int, default=80) + ap.add_argument("--question-index", type=int, default=3) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--context-chars", type=int, default=120) + ap.add_argument("--embed-dims", type=int, default=10) + ap.add_argument("--out", default="qwen-token-entry") + args = ap.parse_args() + + import torch + from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration + + out_dir = HERE / "results" / args.out + img_dir = out_dir / "images" + img_dir.mkdir(parents=True, exist_ok=True) + CACHE.mkdir(exist_ok=True) + + cfg = FONTS[args.font] + cols, rows, budget = capacity(cfg, args.size) + paras = squad.load_paragraphs(CACHE)[: args.limit_paras] + flow, offsets = squad.build_flow(paras) + chunk = flow[: min(len(flow), budget)] + questions = sample_answer_questions(paras, offsets, 0, len(chunk), 24, args.seed) + q = questions[min(args.question_index, len(questions) - 1)] + img = render(chunk, cfg, CACHE, args.size, args.variant) + img.save(img_dir / "image-carrier.png") + + print(f"loading {args.model_dir}", flush=True) + from transformers import AutoTokenizer + + processor = AutoProcessor.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True, use_fast=False) + model = Qwen2_5_VLForConditionalGeneration.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True, dtype=torch.bfloat16, device_map="auto").eval() + device = next(model.parameters()).device + tokenizer = AutoTokenizer.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True) # fast tokenizer for offsets + + # --- Text lane: real tokenization of the snippet around the answer. + snip_start = max(0, q["answer_start"] - args.context_chars) + snip_end = min(len(chunk), q["answer_end"] + args.context_chars) + snippet = chunk[snip_start:snip_end] + enc = tokenizer(snippet, add_special_tokens=False, return_offsets_mapping=True) + tokens = [] + answer_token_idx: list[int] = [] + rel_a = q["answer_start"] - snip_start + rel_b = q["answer_end"] - snip_start + for ti, (tok_id, (o0, o1)) in enumerate(zip(enc["input_ids"], enc["offset_mapping"])): + is_answer = o0 < rel_b and o1 > rel_a + if is_answer: + answer_token_idx.append(ti) + tokens.append({"i": ti, "id": int(tok_id), "str": tokenizer.decode([tok_id]), "answer": bool(is_answer)}) + + # Real embedding rows entering the decoder for the answer tokens. + embed = model.get_input_embeddings() + answer_ids = torch.tensor([tokens[i]["id"] for i in answer_token_idx], device=device) + with torch.no_grad(): + answer_embeds = embed(answer_ids).float().cpu().numpy() + text_entry = [ + { + "id": tokens[i]["id"], + "str": tokens[i]["str"], + "vector_head": [round(float(v), 4) for v in answer_embeds[k, : args.embed_dims]], + "norm": round(float(np.linalg.norm(answer_embeds[k])), 4), + } + for k, i in enumerate(answer_token_idx) + ] + chunk_token_count = len(tokenizer(chunk, add_special_tokens=False)["input_ids"]) + + # --- Image lane: real pixel patches and visual-tower output vectors. + batch = processor(images=img, text="<|vision_start|><|image_pad|><|vision_end|>", return_tensors="pt") + pixel_values = batch["pixel_values"] + grid_thw = batch["image_grid_thw"] + merge = int(getattr(processor.image_processor, "merge_size", 2)) + patch = int(getattr(processor.image_processor, "patch_size", 14)) + with torch.no_grad(): + visual_out = model.model.visual(pixel_values.to(device, dtype=torch.bfloat16), grid_thw=grid_thw.to(device)).float().cpu().numpy() + n_tokens = visual_out.shape[0] + grid = int(round(n_tokens**0.5)) + answer_img_indices = image_answer_token_indices(q["answer_start"], q["answer_end"], cols, cfg.adv, cfg.pitch, img.width, img.height, n_tokens) + image_entry = [ + { + "token_index": int(idx), + "grid_rc": [int(idx // grid), int(idx % grid)], + "vector_head": [round(float(v), 4) for v in visual_out[idx, : args.embed_dims]], + "norm": round(float(np.linalg.norm(visual_out[idx])), 4), + } + for idx in answer_img_indices + ] + # A few real normalized pixel values from the first answer patch (pre-visual-tower input). + patches_per_token = merge * merge + first_patch_row = answer_img_indices[0] * patches_per_token if answer_img_indices else 0 + pixel_head = [round(float(v), 4) for v in pixel_values[min(first_patch_row, pixel_values.shape[0] - 1), : args.embed_dims].tolist()] + + dump = { + "args": vars(args), + "question": {"q": q["q"], "answer_text": q["answer_text"], "answer_start": q["answer_start"], "answer_end": q["answer_end"]}, + "geometry": {"cols": cols, "rows": rows, "image_w": img.width, "image_h": img.height}, + "snippet": snippet, + "snippet_rel_answer": [rel_a, rel_b], + "tokens": tokens, + "text_entry": text_entry, + "chunk_chars": len(chunk), + "chunk_text_tokens": chunk_token_count, + "image_tokens": n_tokens, + "image_grid": grid, + "grid_thw": grid_thw.tolist(), + "patch_size": patch, + "merge_size": merge, + "token_pixel_size": patch * merge, + "processor_resized": [int(grid_thw[0][2]) * patch, int(grid_thw[0][1]) * patch], + "pixel_values_shape": list(pixel_values.shape), + "pixel_head_first_answer_patch": pixel_head, + "image_answer_token_indices": [int(i) for i in answer_img_indices], + "image_entry": image_entry, + "embed_dim": int(answer_embeds.shape[1]), + "visual_out_dim": int(visual_out.shape[1]), + } + (out_dir / "token_entry.json").write_text(json.dumps(dump, indent=1)) + print(json.dumps({k: v for k, v in dump.items() if k not in ("tokens", "snippet")}, indent=1)) + print(f"results -> {out_dir}") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_token_entry_viz.py b/packages/snapcompact/research/snapcompact_token_entry_viz.py new file mode 100644 index 000000000..3feaac0f6 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_token_entry_viz.py @@ -0,0 +1,225 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy"] +# /// +"""Render how the same word enters Qwen as BPE tokens vs 28px visual patches.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + +import numpy as np +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +PALETTE = { + "bg": (5, 7, 10), + "panel": (12, 17, 23), + "panel2": (8, 12, 17), + "ink": (241, 239, 224), + "muted": (143, 154, 160), + "cyan": (75, 220, 255), + "orange": (255, 112, 72), + "green": (148, 255, 117), + "amber": (255, 196, 68), + "grid": (38, 49, 58), +} + + +def ui_font(size: int, bold: bool = False) -> ImageFont.ImageFont: + for path in [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def mono_font(size: int) -> ImageFont.ImageFont: + for path in ["/System/Library/Fonts/Monaco.ttf", "/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf"]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def vector_text(head: list[float]) -> str: + return "[" + ", ".join(f"{v:+.2f}" for v in head[:6]) + ", …]" + + +def draw_vector_bar(draw: ImageDraw.ImageDraw, xy: tuple[int, int], head: list[float], color: tuple[int, int, int], width: int = 330) -> None: + x, y = xy + n = len(head) + bw = width // n + hi = max(0.001, max(abs(v) for v in head)) + mid = y + 22 + for i, v in enumerate(head): + bh = round(20 * abs(v) / hi) + xa = x + i * bw + if v >= 0: + draw.rounded_rectangle((xa, mid - bh, xa + bw - 4, mid), radius=3, fill=color) + else: + draw.rounded_rectangle((xa, mid, xa + bw - 4, mid + bh), radius=3, fill=tuple(c // 2 for c in color)) + draw.line((x, mid, x + width, mid), fill=PALETTE["grid"], width=1) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--result-dir", default=str(HERE / "results" / "qwen-token-entry-q3")) + ap.add_argument("--out", default=str(HERE / "results" / "qwen-token-entry-q3" / "token-entry.png")) + args = ap.parse_args() + result_dir = Path(args.result_dir) + dump = json.loads((result_dir / "token_entry.json").read_text()) + carrier = Image.open(result_dir / "images" / "image-carrier.png").convert("RGB") + + w, h = 2200, 1400 + canvas = Image.new("RGB", (w, h), PALETTE["bg"]) + draw = ImageDraw.Draw(canvas) + for y in range(0, h, 16): + draw.line((0, y, w, y), fill=(7, 10 + y % 9, 15 + y % 11)) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-240, -200, 940, 760), fill=(75, 220, 255, 27)) + gd.ellipse((1240, 540, 2460, 1480), fill=(255, 112, 72, 25)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(86))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + q = dump["question"] + answer = q["answer_text"] + draw.text((64, 42), "QWEN TOKEN ENTRY — SAME WORD, TWO ENCODINGS", fill=PALETTE["amber"], font=ui_font(24, True)) + draw.text((64, 84), f"How “{answer}” gets into the model", fill=PALETTE["ink"], font=ui_font(64, True)) + draw.text( + (66, 164), + "Real values, no schematic: actual BPE ids and embedding rows on the text path; actual 28×28 pixel patches and visual-tower output vectors on the image path.", + fill=PALETTE["muted"], + font=ui_font(23), + ) + + # ---- TEXT LANE ---- + lane = (64, 238, 2136, 700) + draw.rounded_rectangle(lane, radius=28, fill=PALETTE["panel"], outline=(35, 49, 59), width=1) + draw.text((96, 262), "text carrier — BPE tokens", fill=PALETTE["cyan"], font=ui_font(30, True)) + draw.text((96, 302), f"snippet around the answer · {dump['chunk_chars']:,} chars → {dump['chunk_text_tokens']:,} text tokens for the whole chunk", fill=PALETTE["muted"], font=ui_font(18)) + + # Token ribbon: show tokens around the answer. + tokens = dump["tokens"] + answer_positions = [t["i"] for t in tokens if t["answer"]] + mid_token = answer_positions[0] if answer_positions else len(tokens) // 2 + window = [t for t in tokens if mid_token - 7 <= t["i"] <= mid_token + 7] + fnt = mono_font(19) + fnt_id = mono_font(12) + x = 96 + y = 356 + for t in window: + label = t["str"].replace("\n", "⏎") + if not label.strip(): + label = "␣" * max(1, len(label)) + tw = max(int(draw.textlength(label, font=fnt)) + 22, 54) + if x + tw > 2100: + x = 96 + y += 96 + color = PALETTE["amber"] if t["answer"] else (30, 41, 50) + text_color = (8, 10, 12) if t["answer"] else PALETTE["ink"] + draw.rounded_rectangle((x, y, x + tw, y + 44), radius=9, fill=color, outline=(52, 68, 80), width=1) + draw.text((x + 11, y + 9), label, fill=text_color, font=fnt) + draw.text((x + 4, y + 50), f"id {t['id']}", fill=PALETTE["muted"], font=fnt_id) + x += tw + 8 + + draw.text((96, 500), "what actually enters the decoder (embedding row, first 6 of " + f"{dump['embed_dim']} dims):", fill=PALETTE["muted"], font=ui_font(18, True)) + ex = 96 + for entry in dump["text_entry"][:3]: + box = (ex, 536, ex + 470, 668) + draw.rounded_rectangle(box, radius=16, fill=PALETTE["panel2"], outline=(34, 48, 58), width=1) + draw.text((ex + 18, 548), f"“{entry['str']}” id {entry['id']}", fill=PALETTE["cyan"], font=ui_font(20, True)) + draw.text((ex + 18, 578), vector_text(entry["vector_head"]), fill=PALETTE["ink"], font=mono_font(15)) + draw_vector_bar(draw, (ex + 18, 606), entry["vector_head"], PALETTE["cyan"], width=430) + draw.text((ex + 360, 548), f"‖x‖={entry['norm']:.2f}", fill=PALETTE["muted"], font=ui_font(14)) + ex += 494 + + # ---- IMAGE LANE ---- + lane = (64, 736, 2136, 1336) + draw.rounded_rectangle(lane, radius=28, fill=PALETTE["panel"], outline=(35, 49, 59), width=1) + draw.text((96, 760), "image carrier — visual patch tokens", fill=PALETTE["orange"], font=ui_font(30, True)) + px = dump["token_pixel_size"] + draw.text( + (96, 800), + f"same word as pixels · bitmap resized to {dump['processor_resized'][0]}×{dump['processor_resized'][1]} → {dump['patch_size']}px patches, {dump['merge_size']}×{dump['merge_size']} merged → {dump['image_tokens']:,} tokens of {px}×{px}px", + fill=PALETTE["muted"], + font=ui_font(18), + ) + + # Zoomed answer region with the real patch grid. + grid = dump["image_grid"] + rw, rh = dump["processor_resized"] + resized = carrier.resize((rw, rh), Image.Resampling.LANCZOS) + indices = dump["image_answer_token_indices"] + rows = sorted({i // grid for i in indices}) + cols_ = sorted({i % grid for i in indices}) + pad = 3 + cx0 = max(0, (min(cols_) - pad) * px) + cx1 = min(rw, (max(cols_) + 1 + pad) * px) + cy0 = max(0, (min(rows) - pad) * px) + cy1 = min(rh, (max(rows) + 1 + pad) * px) + crop = resized.crop((cx0, cy0, cx1, cy1)) + scale = min(940 / crop.width, 225 / crop.height) + crop_big = crop.resize((round(crop.width * scale), round(crop.height * scale)), Image.Resampling.NEAREST) + ox, oy = 96, 852 + draw.rounded_rectangle((ox - 6, oy - 6, ox + crop_big.width + 6, oy + crop_big.height + 6), radius=10, fill=(244, 242, 230)) + canvas.paste(crop_big, (ox, oy)) + cd = ImageDraw.Draw(canvas) + for gx in range(cx0 // px, cx1 // px + 1): + lx = ox + (gx * px - cx0) * scale + cd.line((lx, oy, lx, oy + crop_big.height), fill=(150, 158, 162, 80), width=1) + for gy in range(cy0 // px, cy1 // px + 1): + ly = oy + (gy * px - cy0) * scale + cd.line((ox, ly, ox + crop_big.width, ly), fill=(150, 158, 162, 80), width=1) + for idx in indices: + r, c = idx // grid, idx % grid + xa = ox + (c * px - cx0) * scale + ya = oy + (r * px - cy0) * scale + cd.rectangle((xa, ya, xa + px * scale, ya + px * scale), outline=PALETTE["orange"], width=4) + draw.text((ox, oy + crop_big.height + 14), f"orange cells = the {len(indices)} visual tokens covering “{answer}” (token grid {grid}×{grid})", fill=PALETTE["muted"], font=ui_font(17)) + + # Magnified single patches. + sx = ox + crop_big.width + 60 + draw.text((sx, 852 - 26), "individual visual tokens (real input pixels):", fill=PALETTE["muted"], font=ui_font(18, True)) + for k, idx in enumerate(indices[:5]): + r, c = idx // grid, idx % grid + cell = resized.crop((c * px, r * px, (c + 1) * px, (r + 1) * px)).resize((132, 132), Image.Resampling.NEAREST) + bx = sx + k * 160 + draw.rounded_rectangle((bx - 4, 852 - 4, bx + 136, 852 + 136), radius=8, fill=(244, 242, 230), outline=PALETTE["orange"], width=3) + canvas.paste(cell, (bx, 852)) + draw.text((bx, 996), f"tok[{idx}]", fill=PALETTE["muted"], font=mono_font(13)) + draw.text((sx, 1030), f"pre-tower normalized pixels of first patch: {vector_text(dump['pixel_head_first_answer_patch'])}", fill=PALETTE["muted"], font=mono_font(14)) + + draw.text((96, 1106), f"what actually enters the decoder (visual-tower output, first 6 of {dump['visual_out_dim']} dims):", fill=PALETTE["muted"], font=ui_font(18, True)) + ex = 96 + for entry in dump["image_entry"][:4]: + box = (ex, 1142, ex + 470, 1274) + draw.rounded_rectangle(box, radius=16, fill=PALETTE["panel2"], outline=(34, 48, 58), width=1) + r, c = entry["grid_rc"] + draw.text((ex + 18, 1154), f"visual tok[{entry['token_index']}] (row {r}, col {c})", fill=PALETTE["orange"], font=ui_font(20, True)) + draw.text((ex + 18, 1184), vector_text(entry["vector_head"]), fill=PALETTE["ink"], font=mono_font(15)) + draw_vector_bar(draw, (ex + 18, 1212), entry["vector_head"], PALETTE["orange"], width=430) + draw.text((ex + 360, 1154), f"‖x‖={entry['norm']:.2f}", fill=PALETTE["muted"], font=ui_font(14)) + ex += 494 + + # Comparison strip. + text_tok_for_word = len(dump["text_entry"]) + draw.rounded_rectangle((1100, 536, 2104, 668), radius=16, fill=PALETTE["panel2"], outline=(34, 48, 58), width=1) + draw.text((1128, 556), f"“{answer}” = {text_tok_for_word} text token(s) · {len(indices)} visual tokens", fill=PALETTE["ink"], font=ui_font(22, True)) + draw.text((1128, 592), f"both end up as {dump['embed_dim']}-dim rows in the same decoder", fill=PALETTE["ink"], font=ui_font(19)) + draw.text((1128, 626), "text path: lookup table row. image path: ViT forward over 4 raw patches → merger MLP.", fill=PALETTE["muted"], font=ui_font(16)) + + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + canvas.save(out) + print(out) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_viz_atlas.py b/packages/snapcompact/research/snapcompact_viz_atlas.py new file mode 100644 index 000000000..29d05c57b --- /dev/null +++ b/packages/snapcompact/research/snapcompact_viz_atlas.py @@ -0,0 +1,377 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["matplotlib", "numpy", "pillow"] +# /// +"""Render an Activation Atlas-style 2D geography from snapcompact activations.""" + +from __future__ import annotations + +import argparse +import csv +import json +import math +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") + +import matplotlib.patheffects as pe +import matplotlib.pyplot as plt +import numpy as np +from matplotlib.colors import LinearSegmentedColormap, Normalize +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +DEFAULT_DATA_DIR = HERE / "results" / "tensor-heatmap-paddleocr-q7" +DEFAULT_OUT_DIR = HERE / "results" / "agent-viz-atlas" + +BG = (4, 7, 12) +PANEL = (10, 15, 23) +INK = (241, 239, 224) +MUTED = (139, 151, 160) +CYAN = (70, 216, 255) +RED = (255, 75, 61) +AMBER = (255, 198, 68) +GREEN = (135, 255, 139) +PURPLE = (183, 108, 255) + + +def font(size: int, bold: bool = False) -> ImageFont.FreeTypeFont | ImageFont.ImageFont: + candidates = [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/System/Library/Fonts/Supplemental/Avenir Next Condensed.ttc", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ] + for path in candidates: + if path and Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def quantile_norm(x: np.ndarray, lo_q: float = 0.02, hi_q: float = 0.985) -> np.ndarray: + lo = float(np.quantile(x, lo_q)) + hi = float(np.quantile(x, hi_q)) + if hi <= lo: + return np.zeros_like(x, dtype=np.float32) + return np.clip((x - lo) / (hi - lo), 0.0, 1.0).astype(np.float32) + + +def pca2(features: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + x = features.astype(np.float64, copy=True) + x -= x.mean(axis=0, keepdims=True) + scale = x.std(axis=0, keepdims=True) + scale[scale < 1e-9] = 1.0 + x /= scale + _u, s, vt = np.linalg.svd(x, full_matrices=False) + coords = x @ vt[:2].T + explained = (s[:2] ** 2) / np.maximum(np.sum(s**2), 1e-12) + return coords.astype(np.float32), explained.astype(np.float32) + + +def normalize_coords(coords: np.ndarray) -> np.ndarray: + out = coords.copy() + for axis in range(2): + lo = float(np.quantile(out[:, axis], 0.01)) + hi = float(np.quantile(out[:, axis], 0.99)) + if hi <= lo: + out[:, axis] = 0.5 + else: + out[:, axis] = np.clip((out[:, axis] - lo) / (hi - lo), 0, 1) + out[:, 1] = 1.0 - out[:, 1] + return out + + +def kmeans(points: np.ndarray, k: int = 5, iters: int = 32) -> tuple[np.ndarray, np.ndarray]: + # Deterministic farthest-point seeding avoids random output drift. + centers = [points[np.argmax(points[:, 0] + points[:, 1])]] + for _ in range(1, k): + dist = np.min(np.sum((points[:, None, :] - np.asarray(centers)[None, :, :]) ** 2, axis=2), axis=1) + centers.append(points[int(np.argmax(dist))]) + c = np.asarray(centers, dtype=np.float32) + labels = np.zeros(points.shape[0], dtype=np.int32) + for _ in range(iters): + d = np.sum((points[:, None, :] - c[None, :, :]) ** 2, axis=2) + new_labels = np.argmin(d, axis=1).astype(np.int32) + if np.array_equal(new_labels, labels): + break + labels = new_labels + for i in range(k): + mask = labels == i + if np.any(mask): + c[i] = points[mask].mean(axis=0) + return labels, c + + +def crop_answer_strip(img: Image.Image, start: int, end: int, cols: int, adv: int = 8, pitch: int = 13) -> Image.Image: + row0 = max(0, start // cols - 4) + row1 = min(img.height // pitch, end // cols + 5) + col0 = max(0, start % cols - 32) + col1 = min(cols, end % cols + 34) + crop = img.crop((col0 * adv, row0 * pitch, col1 * adv, row1 * pitch)).convert("RGB") + d = ImageDraw.Draw(crop) + bx0 = max(0, (start % cols - col0) * adv - adv) + bx1 = min(crop.width - 1, ((end - 1) % cols - col0 + 2) * adv) + by0 = max(0, (start // cols - row0) * pitch - 2) + by1 = min(crop.height - 1, ((end - 1) // cols - row0 + 1) * pitch + 2) + d.rounded_rectangle((bx0, by0, bx1, by1), radius=4, outline=RED, width=3) + return crop + + +def paste_fit(canvas: Image.Image, img: Image.Image, box: tuple[int, int, int, int]) -> None: + x0, y0, x1, y1 = box + scale = min((x1 - x0) / img.width, (y1 - y0) / img.height) + resized = img.resize((max(1, round(img.width * scale)), max(1, round(img.height * scale))), Image.Resampling.NEAREST) + canvas.paste(resized, (x0 + (x1 - x0 - resized.width) // 2, y0 + (y1 - y0 - resized.height) // 2)) + + +def render_atlas_panel( + points: np.ndarray, + labels: np.ndarray, + centers: np.ndarray, + ratio_strength: np.ndarray, + answer_strength: np.ndarray, + peak_layers: np.ndarray, + explained: np.ndarray, + summary: dict, + out_dir: Path, +) -> Image.Image: + cmap = LinearSegmentedColormap.from_list("scar", ["#182132", "#245d7a", "#48d8ff", "#ffd04e", "#ff493d"]) + fig = plt.figure(figsize=(15.8, 10.6), dpi=170) + fig.patch.set_facecolor("#04070c") + ax = fig.add_axes((0.045, 0.06, 0.91, 0.88), facecolor="#07101a") + + x = points[:, 0] + y = points[:, 1] + hb = ax.hexbin(x, y, C=ratio_strength, gridsize=46, reduce_C_function=np.mean, cmap=cmap, mincnt=1, linewidths=0, alpha=0.64) + hb.set_clim(0.0, 1.0) + + cluster_colors = ["#46d8ff", "#ff4b3d", "#ffc644", "#87ff8b", "#b76cff"] + for i, color in enumerate(cluster_colors): + mask = labels == i + if np.count_nonzero(mask) < 4: + continue + ax.scatter(x[mask], y[mask], s=28 + answer_strength[mask] * 150, c=color, alpha=0.24, linewidths=0) + ax.scatter( + x[mask], + y[mask], + s=10 + answer_strength[mask] * 52, + c=ratio_strength[mask], + cmap=cmap, + norm=Normalize(0, 1), + alpha=0.93, + edgecolors=color, + linewidths=0.45, + ) + + hot = np.argsort(ratio_strength + answer_strength * 0.55)[-9:] + ax.scatter(x[hot], y[hot], s=210, facecolors="none", edgecolors="#fff0a8", linewidths=1.5, alpha=0.95) + for rank, idx in enumerate(hot[-5:][::-1], 1): + ax.text( + x[idx] + 0.012, + y[idx] + 0.010, + f"T{idx} · L{int(peak_layers[idx])}", + color="#fff3b0", + fontsize=8, + weight="bold", + path_effects=[pe.withStroke(linewidth=2.5, foreground="#05070a")], + ) + + names = ["answer ridge", "control basin", "early glyph shore", "late-context upland", "ratio reef"] + cluster_scores = [] + for i in range(len(centers)): + mask = labels == i + cluster_scores.append((float(ratio_strength[mask].mean()) if np.any(mask) else 0.0, i)) + order = {old: new for new, (_score, old) in enumerate(sorted(cluster_scores, reverse=True))} + for i, c in enumerate(centers): + mask = labels == i + if np.count_nonzero(mask) < 5: + continue + label = names[order[i] % len(names)] + ax.text( + c[0], + c[1], + label.upper(), + color=cluster_colors[i], + fontsize=11, + weight="bold", + ha="center", + va="center", + alpha=0.96, + path_effects=[pe.withStroke(linewidth=4, foreground="#05070a")], + ) + + ax.text( + 0.015, + 0.982, + "Activation Atlas projection", + transform=ax.transAxes, + color="#f1efe0", + fontsize=24, + weight="bold", + va="top", + ) + ax.text( + 0.017, + 0.942, + "Each island is one image token; geography = PCA of 19-layer answer-vs-random activation delta vectors.", + transform=ax.transAxes, + color="#8b97a0", + fontsize=11, + va="top", + ) + ax.text( + 0.017, + 0.905, + f"Question: {summary['question']['q']} · gold answer: {summary['question']['answer_text']} · answer/random mean Δ {summary['answer_over_random_delta']:.2f}×", + transform=ax.transAxes, + color="#ffc644", + fontsize=10, + weight="bold", + va="top", + ) + ax.text( + 0.99, + 0.02, + f"PCA variance: PC1 {explained[0] * 100:.1f}% · PC2 {explained[1] * 100:.1f}% color: answer/random scar size: answer-mask Δ labels: peak layer depth", + transform=ax.transAxes, + color="#8b97a0", + fontsize=9, + ha="right", + va="bottom", + ) + + cax = fig.add_axes((0.83, 0.865, 0.12, 0.014)) + cb = fig.colorbar(hb, cax=cax, orientation="horizontal") + cb.outline.set_visible(False) + cb.set_ticks([0, 1]) + cb.set_ticklabels(["random-like", "answer scar"]) + cb.ax.tick_params(colors="#cfd6d0", labelsize=8, length=0) + cax.set_facecolor("#07101a") + + for spine in ax.spines.values(): + spine.set_color("#223140") + spine.set_linewidth(1.0) + ax.set_xticks([]) + ax.set_yticks([]) + ax.set_xlim(-0.04, 1.04) + ax.set_ylim(-0.04, 1.04) + ax.grid(color="#2d4255", alpha=0.14, linewidth=0.7) + + tmp = out_dir / ".atlas-panel.png" + fig.savefig(tmp, facecolor=fig.get_facecolor(), transparent=False) + plt.close(fig) + img = Image.open(tmp).convert("RGB") + tmp.unlink(missing_ok=True) + return img + + +def draw_shell(panel: Image.Image, summary: dict, data_dir: Path, out: Path) -> None: + w, h = 2400, 1500 + canvas = Image.new("RGB", (w, h), BG) + draw = ImageDraw.Draw(canvas) + for y in range(0, h, 18): + draw.line((0, y, w, y), fill=(6, 10 + y % 11, 16 + y % 9)) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-360, -240, 1000, 760), fill=(70, 216, 255, 32)) + gd.ellipse((1270, 210, 2740, 1610), fill=(255, 75, 61, 32)) + gd.ellipse((690, 920, 1740, 1780), fill=(255, 198, 68, 18)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(90))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + draw.text((74, 48), "SNAPCOMPACT WHITEBOX", fill=AMBER, font=font(24, True)) + draw.text((74, 88), "Activation Atlas of the missing answer", fill=INK, font=font(72, True)) + draw.text((78, 178), "A PCA geography of image-token residual scars: where blanking the gold answer ‘2003’ moves the model differently than a random blank.", fill=MUTED, font=font(27)) + + draw.rounded_rectangle((74, 252, 590, 1390), radius=32, fill=PANEL, outline=(32, 45, 58), width=1) + q = summary["question"] + cols = summary["geometry"]["cols"] + original = Image.open(data_dir / "images" / "original.png").convert("RGB") + answer_mask = Image.open(data_dir / "images" / "answer-mask.png").convert("RGB") + random_mask = Image.open(data_dir / "images" / "random-mask.png").convert("RGB") + strips = [ + ("ORIGINAL", original, CYAN), + ("ANSWER MASK", answer_mask, RED), + ("RANDOM MASK", random_mask, GREEN), + ] + y = 314 + for title, img, color in strips: + draw.text((110, y), title, fill=color, font=font(18, True)) + draw.rounded_rectangle((110, y + 28, 554, y + 166), radius=16, fill=(242, 241, 229), outline=color, width=3) + paste_fit(canvas, crop_answer_strip(img, q["answer_start"], q["answer_end"], cols), (124, y + 42, 540, y + 152)) + y += 226 + + draw.rounded_rectangle((110, 1002, 554, 1300), radius=24, fill=(7, 11, 18), outline=(35, 51, 66), width=1) + metrics = [ + ("answer", q["answer_text"], AMBER, 48), + ("layers", str(summary["layers"]), CYAN, 38), + ("image tokens", str(summary["image_tokens"]), GREEN, 38), + ("answer/random Δ", f"{summary['answer_over_random_delta']:.2f}×", RED, 38), + ] + yy = 1038 + for label, value, color, size in metrics: + draw.text((142, yy), label, fill=MUTED, font=font(16, True)) + draw.text((142, yy + 26), value, fill=color, font=font(size, True)) + yy += 68 + draw.text((110, 1336), "Actual heatmaps.npz + summary.json; no schematic points.", fill=MUTED, font=font(18)) + + draw.rounded_rectangle((622, 252, 2326, 1390), radius=32, fill=PANEL, outline=(32, 45, 58), width=1) + panel = panel.resize((1640, 1098), Image.Resampling.LANCZOS) + canvas.paste(panel, (654, 272)) + + out.parent.mkdir(parents=True, exist_ok=True) + canvas.save(out, quality=95) + + +def write_source_data(out_dir: Path, points: np.ndarray, labels: np.ndarray, ratio_strength: np.ndarray, answer_strength: np.ndarray, peak_layers: np.ndarray, explained: np.ndarray) -> None: + np.savez_compressed( + out_dir / "atlas_source.npz", + points=points, + cluster=labels, + ratio_strength=ratio_strength, + answer_strength=answer_strength, + peak_layer=peak_layers, + pca_explained=explained, + ) + with (out_dir / "atlas_points.csv").open("w", newline="") as f: + writer = csv.writer(f) + writer.writerow(["token", "atlas_x", "atlas_y", "cluster", "ratio_strength", "answer_strength", "peak_layer"]) + for i in range(points.shape[0]): + writer.writerow([i, f"{points[i, 0]:.6f}", f"{points[i, 1]:.6f}", int(labels[i]), f"{ratio_strength[i]:.6f}", f"{answer_strength[i]:.6f}", int(peak_layers[i])]) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--data-dir", default=str(DEFAULT_DATA_DIR)) + ap.add_argument("--out-dir", default=str(DEFAULT_OUT_DIR)) + args = ap.parse_args() + + data_dir = Path(args.data_dir) + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + summary = json.loads((data_dir / "summary.json").read_text()) + heatmaps = np.load(data_dir / "heatmaps.npz") + answer = heatmaps["answer_delta"].astype(np.float32) + random = heatmaps["random_delta"].astype(np.float32) + ratio = heatmaps["ratio"].astype(np.float32) + + contrast = np.log1p(answer) - np.log1p(random) + features = np.concatenate([contrast.T, np.log1p(ratio).T, np.log1p(answer).T], axis=1) + raw_coords, explained = pca2(features) + points = normalize_coords(raw_coords) + labels, centers = kmeans(points, k=5) + + ratio_strength = quantile_norm(np.log1p(ratio).mean(axis=0), 0.02, 0.99) + answer_strength = quantile_norm(np.log1p(answer).mean(axis=0), 0.02, 0.99) + peak_layers = np.argmax(ratio, axis=0).astype(np.int32) + + write_source_data(out_dir, points, labels, ratio_strength, answer_strength, peak_layers, explained) + panel = render_atlas_panel(points, labels, centers, ratio_strength, answer_strength, peak_layers, explained, summary, out_dir) + out = out_dir / "atlas.png" + draw_shell(panel, summary, data_dir, out) + print(out) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_viz_circuit.py b/packages/snapcompact/research/snapcompact_viz_circuit.py new file mode 100644 index 000000000..f78c70df3 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_viz_circuit.py @@ -0,0 +1,427 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["numpy", "pillow"] +# /// +"""Render a circuit-graph visualization from snapcompact activation deltas.""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path +from typing import Iterable + +import numpy as np +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +DEFAULT_RESULT_DIR = HERE / "results" / "tensor-heatmap-paddleocr-q7" +DEFAULT_OUT_DIR = HERE / "results" / "agent-viz-circuit" + +BG = (4, 6, 10) +PANEL = (13, 17, 24) +PANEL_2 = (8, 12, 18) +INK = (242, 239, 225) +MUTED = (132, 146, 153) +BLUE = (83, 218, 255) +GOLD = (255, 199, 74) +ORANGE = (255, 122, 54) +RED = (255, 72, 82) +GREEN = (127, 245, 148) + + +def font(size: int, bold: bool = False) -> ImageFont.FreeTypeFont | ImageFont.ImageFont: + candidates = [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/System/Library/Fonts/Supplemental/Helvetica Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Helvetica.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ] + for path in candidates: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def clamp01(v: float) -> float: + return max(0.0, min(1.0, v)) + + +def mix(a: tuple[int, int, int], b: tuple[int, int, int], t: float) -> tuple[int, int, int]: + t = clamp01(t) + return tuple(round(x + (y - x) * t) for x, y in zip(a, b)) + + +def quantile_norm(values: np.ndarray, q: float = 0.97) -> np.ndarray: + scale = float(np.quantile(values, q)) + if not math.isfinite(scale) or scale <= 0: + scale = 1.0 + return np.clip(values / scale, 0, 1) + + +def rounded_panel(draw: ImageDraw.ImageDraw, box: tuple[int, int, int, int], radius: int = 34) -> None: + draw.rounded_rectangle(box, radius=radius, fill=PANEL, outline=(31, 41, 51), width=1) + + +def multiline(draw: ImageDraw.ImageDraw, xy: tuple[int, int], text: str, *, fill: tuple[int, int, int], fnt: ImageFont.ImageFont, max_width: int, line_gap: int = 8) -> int: + words = text.split() + lines: list[str] = [] + cur = "" + for word in words: + trial = word if not cur else f"{cur} {word}" + if draw.textlength(trial, font=fnt) <= max_width: + cur = trial + else: + if cur: + lines.append(cur) + cur = word + if cur: + lines.append(cur) + x, y = xy + step = fnt.size + line_gap if hasattr(fnt, "size") else 20 + for line in lines: + draw.text((x, y), line, fill=fill, font=fnt) + y += step + return y + + +def crop_answer_region(img: Image.Image, summary: dict, pad_cells: int = 42) -> Image.Image: + q = summary["question"] + cols = int(summary["geometry"]["cols"]) + rows = int(summary["geometry"]["rows"]) + adv = max(1, img.width // cols) + pitch = max(1, img.height // rows) + start = int(q["answer_start"]) + end = int(q["answer_end"]) + row0 = max(0, start // cols - 5) + row1 = min(rows, (end - 1) // cols + 7) + col0 = max(0, start % cols - pad_cells) + col1 = min(cols, (end - 1) % cols + pad_cells) + if col1 <= col0 + 8: + col1 = min(cols, col0 + 90) + crop = img.crop((col0 * adv, row0 * pitch, col1 * adv, row1 * pitch)).convert("RGB") + d = ImageDraw.Draw(crop) + bx0 = max(0, (start % cols - col0) * adv - adv) + bx1 = min(crop.width - 1, ((end - 1) % cols - col0 + 2) * adv) + by0 = max(0, (start // cols - row0) * pitch - 2) + by1 = min(crop.height - 1, ((end - 1) // cols - row0 + 1) * pitch + 2) + d.rounded_rectangle((bx0, by0, bx1, by1), radius=4, outline=RED, width=4) + return crop + + +def paste_fit(canvas: Image.Image, img: Image.Image, box: tuple[int, int, int, int], *, resample: int = Image.Resampling.LANCZOS) -> None: + x0, y0, x1, y1 = box + scale = min((x1 - x0) / img.width, (y1 - y0) / img.height) + size = (max(1, round(img.width * scale)), max(1, round(img.height * scale))) + resized = img.resize(size, resample) + canvas.paste(resized, (x0 + (x1 - x0 - size[0]) // 2, y0 + (y1 - y0 - size[1]) // 2)) + + +def draw_bezier(draw: ImageDraw.ImageDraw, points: tuple[tuple[float, float], tuple[float, float], tuple[float, float], tuple[float, float]], *, fill: tuple[int, int, int, int], width: int) -> None: + p0, p1, p2, p3 = points + coords: list[tuple[float, float]] = [] + for i in range(46): + t = i / 45 + mt = 1 - t + x = mt**3 * p0[0] + 3 * mt**2 * t * p1[0] + 3 * mt * t**2 * p2[0] + t**3 * p3[0] + y = mt**3 * p0[1] + 3 * mt**2 * t * p1[1] + 3 * mt * t**2 * p2[1] + t**3 * p3[1] + coords.append((x, y)) + draw.line(coords, fill=fill, width=width, joint="curve") + + +def token_groups(grid_side: int = 27, tiles: int = 3) -> list[dict[str, int | str]]: + groups: list[dict[str, int | str]] = [] + names = ["upper-left", "upper", "upper-right", "left", "center", "right", "lower-left", "lower", "lower-right"] + idx = 0 + for gy in range(tiles): + y0 = round(gy * grid_side / tiles) + y1 = round((gy + 1) * grid_side / tiles) + for gx in range(tiles): + x0 = round(gx * grid_side / tiles) + x1 = round((gx + 1) * grid_side / tiles) + groups.append({"name": names[idx], "x0": x0, "x1": x1, "y0": y0, "y1": y1}) + idx += 1 + return groups + + +def group_indices(group: dict[str, int | str], grid_side: int = 27) -> np.ndarray: + ids: list[int] = [] + for y in range(int(group["y0"]), int(group["y1"])): + for x in range(int(group["x0"]), int(group["x1"])): + ids.append(y * grid_side + x) + return np.asarray(ids, dtype=np.int64) + + +def build_metrics(answer: np.ndarray, random: np.ndarray, ratio: np.ndarray) -> tuple[list[dict], list[dict], np.ndarray, np.ndarray]: + layers, tokens = answer.shape + grid_side = int(round(math.sqrt(tokens))) + if grid_side * grid_side != tokens: + raise ValueError(f"expected square image-token grid, got {tokens} tokens") + + groups = token_groups(grid_side) + layer_score = np.maximum(answer - random, 0.0) * np.log1p(np.maximum(ratio, 0.0)) + layer_norm = quantile_norm(layer_score, 0.975) + + group_rows: list[dict] = [] + for i, group in enumerate(groups): + ids = group_indices(group, grid_side) + a = answer[:, ids] + r = random[:, ids] + rr = ratio[:, ids] + raw = layer_score[:, ids] + group_rows.append( + { + "id": i, + "name": group["name"], + "x0": group["x0"], + "x1": group["x1"], + "y0": group["y0"], + "y1": group["y1"], + "answer_delta_mean": float(a.mean()), + "random_delta_mean": float(r.mean()), + "ratio_mean": float(rr.mean()), + "answer_minus_random_mean": float((a - r).mean()), + "edge_score": float(raw.mean()), + "layer_scores": [float(raw[j].mean()) for j in range(layers)], + "layer_ratios": [float(rr[j].mean()) for j in range(layers)], + } + ) + + layer_rows: list[dict] = [] + for layer in range(layers): + layer_rows.append( + { + "layer": layer, + "answer_delta_mean": float(answer[layer].mean()), + "random_delta_mean": float(random[layer].mean()), + "ratio_mean": float(ratio[layer].mean()), + "answer_minus_random_mean": float((answer[layer] - random[layer]).mean()), + "edge_score": float(layer_score[layer].mean()), + } + ) + return group_rows, layer_rows, layer_score, layer_norm + + +def draw_token_grid(draw: ImageDraw.ImageDraw, box: tuple[int, int, int, int], token_strength: np.ndarray, group_rows: list[dict]) -> list[tuple[int, int]]: + x0, y0, x1, y1 = box + grid = token_strength.reshape(27, 27) + norm = quantile_norm(grid, 0.985) + cell = min((x1 - x0) // 27, (y1 - y0) // 27) + gx = x0 + ((x1 - x0) - 27 * cell) // 2 + gy = y0 + ((y1 - y0) - 27 * cell) // 2 + for y in range(27): + for x in range(27): + v = float(norm[y, x]) + color = mix((13, 24, 34), ORANGE, v) + if v > 0.72: + color = mix(color, GOLD, (v - 0.72) / 0.28) + draw.rectangle((gx + x * cell, gy + y * cell, gx + (x + 1) * cell - 1, gy + (y + 1) * cell - 1), fill=color) + draw.rectangle((gx - 1, gy - 1, gx + 27 * cell, gy + 27 * cell), outline=(70, 88, 101), width=2) + + centers: list[tuple[int, int]] = [] + scores = np.asarray([g["edge_score"] for g in group_rows], dtype=np.float32) + score_norm = quantile_norm(scores, 0.92) + for g, s in zip(group_rows, score_norm): + cx = gx + round((int(g["x0"]) + int(g["x1"])) * 0.5 * cell) + cy = gy + round((int(g["y0"]) + int(g["y1"])) * 0.5 * cell) + centers.append((cx, cy)) + rad = round(8 + 19 * float(s)) + draw.ellipse((cx - rad, cy - rad, cx + rad, cy + rad), outline=mix(BLUE, GOLD, float(s)), width=3) + return centers + + +def draw_layer_bands(draw: ImageDraw.ImageDraw, box: tuple[int, int, int, int], layer_rows: list[dict]) -> list[tuple[int, int]]: + x0, y0, x1, y1 = box + scores = np.asarray([r["edge_score"] for r in layer_rows], dtype=np.float32) + ratios = np.asarray([r["ratio_mean"] for r in layer_rows], dtype=np.float32) + score_norm = quantile_norm(scores, 0.96) + ratio_norm = quantile_norm(ratios, 0.96) + centers: list[tuple[int, int]] = [] + gap = 7 + h = ((y1 - y0) - gap * (len(layer_rows) - 1)) / len(layer_rows) + for i, (row, s, rr) in enumerate(zip(layer_rows, score_norm, ratio_norm)): + yy0 = round(y0 + i * (h + gap)) + yy1 = round(yy0 + h) + inset = round(26 * (1 - float(s))) + color = mix((17, 27, 38), GOLD, float(rr) * 0.80) + outline = mix((54, 71, 83), RED, float(s)) + draw.rounded_rectangle((x0 + inset, yy0, x1 - inset, yy1), radius=8, fill=color, outline=outline, width=2) + draw.text((x0 - 74, yy0 + max(0, (yy1 - yy0 - 18) // 2)), f"L{int(row['layer']):02d}", fill=mix(MUTED, INK, float(s)), font=font(16, True)) + centers.append(((x0 + x1) // 2, (yy0 + yy1) // 2)) + return centers + + +def render(summary: dict, answer: np.ndarray, random: np.ndarray, ratio: np.ndarray, result_dir: Path, out_dir: Path) -> dict: + group_rows, layer_rows, layer_score, layer_norm = build_metrics(answer, random, ratio) + w, h = 2400, 1350 + canvas = Image.new("RGB", (w, h), BG) + draw = ImageDraw.Draw(canvas) + for y in range(0, h, 18): + shade = 8 + (y // 18) % 4 + draw.line((0, y, w, y), fill=(shade, shade + 2, shade + 7)) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-380, -260, 980, 640), fill=(255, 72, 82, 36)) + gd.ellipse((780, 70, 2320, 1420), fill=(83, 218, 255, 22)) + gd.ellipse((1440, -120, 2760, 860), fill=(255, 199, 74, 26)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(90))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + draw.text((70, 44), "SNAPCOMPACT CIRCUIT TRACE", fill=GOLD, font=font(25, True)) + draw.text((70, 82), "The answer glyphs light a decoder circuit", fill=INK, font=font(66, True)) + multiline( + draw, + (72, 166), + "Edges are computed from actual hidden-state deltas: max(answer − random, 0) × log(1 + answer/random ratio), averaged by image-token region and decoder layer.", + fill=MUTED, + fnt=font(24), + max_width=1470, + line_gap=7, + ) + + rounded_panel(draw, (62, 258, 520, 1238)) + rounded_panel(draw, (568, 258, 1002, 1238)) + rounded_panel(draw, (1126, 258, 1632, 1238)) + rounded_panel(draw, (1816, 258, 2338, 1238)) + + q = summary["question"] + base = Image.open(result_dir / "images" / "original.png").convert("RGB") + masked = Image.open(result_dir / "images" / "answer-mask.png").convert("RGB") + crop = crop_answer_region(base, summary) + masked_crop = crop_answer_region(masked, summary) + + draw.text((94, 294), "1. bitmap intervention", fill=INK, font=font(30, True)) + draw.text((94, 334), "question targets one visible year", fill=MUTED, font=font(18)) + draw.rounded_rectangle((94, 386, 488, 560), radius=16, fill=(240, 238, 224), outline=BLUE, width=3) + paste_fit(canvas, crop, (108, 400, 474, 546), resample=Image.Resampling.NEAREST) + draw.text((94, 574), "original answer region", fill=BLUE, font=font(18, True)) + draw.rounded_rectangle((94, 654, 488, 828), radius=16, fill=(240, 238, 224), outline=RED, width=3) + paste_fit(canvas, masked_crop, (108, 668, 474, 814), resample=Image.Resampling.NEAREST) + draw.text((94, 842), "blanked answer mask", fill=RED, font=font(18, True)) + draw.text((94, 930), "question", fill=MUTED, font=font(15, True)) + multiline(draw, (94, 956), str(q["q"]), fill=INK, fnt=font(23), max_width=370, line_gap=8) + draw.text((94, 1070), "gold answer", fill=MUTED, font=font(15, True)) + draw.text((94, 1098), str(q["answer_text"]), fill=GOLD, font=font(52, True)) + draw.text((94, 1172), f"global Δ ratio {summary['answer_over_random_delta']:.2f}×", fill=INK, font=font(22, True)) + + draw.text((600, 294), "2. image-token regions", fill=INK, font=font(30, True)) + draw.text((600, 334), "27×27 token lattice, colored by circuit score", fill=MUTED, font=font(18)) + token_strength = layer_score.mean(axis=0) + token_centers = draw_token_grid(draw, (616, 392, 954, 730), token_strength, group_rows) + top_groups = sorted(group_rows, key=lambda g: g["edge_score"], reverse=True)[:4] + draw.text((600, 794), "strongest token groups", fill=MUTED, font=font(16, True)) + y = 826 + group_score_norm = quantile_norm(np.asarray([g["edge_score"] for g in group_rows], dtype=np.float32), 0.92) + for g in top_groups: + s = float(group_score_norm[int(g["id"])]) + draw.rounded_rectangle((600, y, 970, y + 62), radius=14, fill=PANEL_2, outline=mix((44, 58, 68), GOLD, s), width=2) + draw.text((620, y + 12), str(g["name"]), fill=INK, font=font(20, True)) + draw.text((820, y + 12), f"{g['ratio_mean']:.2f}×", fill=mix(BLUE, GOLD, s), font=font(21, True)) + draw.text((620, y + 38), f"Δ {g['answer_delta_mean']:.2f} vs {g['random_delta_mean']:.2f}", fill=MUTED, font=font(14)) + y += 78 + + draw.text((1158, 294), "3. decoder layer bands", fill=INK, font=font(30, True)) + draw.text((1158, 334), "band width/color follows per-layer answer specificity", fill=MUTED, font=font(18)) + layer_centers = draw_layer_bands(draw, (1246, 394, 1566, 1122), layer_rows) + + draw.text((1848, 294), "4. output answer", fill=INK, font=font(30, True)) + draw.text((1848, 334), "residual stream converges on text", fill=MUTED, font=font(18)) + draw.rounded_rectangle((1880, 462, 2274, 730), radius=34, fill=(10, 13, 18), outline=(73, 82, 92), width=2) + draw.text((1918, 500), "PaddleOCR-VL", fill=MUTED, font=font(20, True)) + draw.text((1918, 558), "answers", fill=INK, font=font(32, True)) + draw.text((1918, 606), str(q["answer_text"]), fill=GOLD, font=font(82, True)) + draw.rounded_rectangle((1880, 820, 2274, 1034), radius=28, fill=PANEL_2, outline=(47, 62, 73), width=2) + draw.text((1918, 858), f"{summary['layers']} decoder layers", fill=INK, font=font(26, True)) + draw.text((1918, 900), f"{summary['image_tokens']} image tokens", fill=MUTED, font=font(21)) + draw.text((1918, 938), "edge thickness = grouped delta score", fill=MUTED, font=font(21)) + draw.text((1918, 976), "edge color = answer/random ratio", fill=MUTED, font=font(21)) + + # Edges live in a transparent layer so glow can sit behind node labels. + edges = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + ed = ImageDraw.Draw(edges) + all_group_layer = np.asarray([g["layer_scores"] for g in group_rows], dtype=np.float32) + group_layer_norm = quantile_norm(all_group_layer, 0.965) + group_layer_ratio = np.asarray([g["layer_ratios"] for g in group_rows], dtype=np.float32) + ratio_norm = quantile_norm(group_layer_ratio, 0.955) + + selected_groups = [int(g["id"]) for g in sorted(group_rows, key=lambda g: g["edge_score"], reverse=True)[:7]] + selected_layers = [0, 1, 2, 3, 4, 5, 7, 9, 12, 15, 18] + for gi in selected_groups: + sx, sy = token_centers[gi] + for li in selected_layers: + strength = float(group_layer_norm[gi, li]) + if strength < 0.10: + continue + ex, ey = layer_centers[li] + col = mix(BLUE, RED, float(ratio_norm[gi, li])) + alpha = round(54 + 156 * strength) + width = max(1, round(1 + 9 * strength)) + draw_bezier(ed, ((sx + 18, sy), (1046, sy), (1110, ey), (ex - 162, ey)), fill=(*col, alpha), width=width) + + layer_edge_norm = quantile_norm(np.asarray([r["edge_score"] for r in layer_rows], dtype=np.float32), 0.96) + layer_ratio_norm = quantile_norm(np.asarray([r["ratio_mean"] for r in layer_rows], dtype=np.float32), 0.96) + out_anchor = (1880, 596) + for li in selected_layers: + sx, sy = layer_centers[li] + strength = float(layer_edge_norm[li]) + col = mix(GOLD, RED, float(layer_ratio_norm[li])) + width = max(2, round(2 + 11 * strength)) + alpha = round(76 + 160 * strength) + draw_bezier(ed, ((sx + 162, sy), (1668, sy), (1748, out_anchor[1] + (sy - 760) * 0.18), out_anchor), fill=(*col, alpha), width=width) + + edges = edges.filter(ImageFilter.GaussianBlur(0.18)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), edges).convert("RGB") + draw = ImageDraw.Draw(canvas) + + # Connector arrows and legend drawn after glowing edges. + draw.line((520, 748, 568, 748), fill=(70, 84, 96), width=3) + draw.polygon([(568, 748), (550, 738), (550, 758)], fill=(70, 84, 96)) + draw.line((1002, 748, 1126, 748), fill=(70, 84, 96), width=3) + draw.polygon([(1126, 748), (1108, 738), (1108, 758)], fill=(70, 84, 96)) + draw.line((1632, 748, 1816, 748), fill=(70, 84, 96), width=3) + draw.polygon([(1816, 748), (1798, 738), (1798, 758)], fill=(70, 84, 96)) + + legend_x, legend_y = 590, 1168 + draw.text((legend_x, legend_y), "edge encoding", fill=INK, font=font(18, True)) + for i, (lab, val, col) in enumerate([("weak", 0.20, BLUE), ("medium", 0.55, GOLD), ("answer-specific", 0.95, RED)]): + yy = legend_y + 38 + i * 32 + draw.line((legend_x, yy, legend_x + 122, yy), fill=col, width=round(2 + 9 * val)) + draw.text((legend_x + 146, yy - 12), lab, fill=MUTED if i < 2 else INK, font=font(16)) + + draw.text((1158, 1164), "Data: heatmaps.npz answer_delta, random_delta, ratio. No schematic edges: every width/color is grouped from observed tensors.", fill=MUTED, font=font(17)) + + out_dir.mkdir(parents=True, exist_ok=True) + out_png = out_dir / "circuit.png" + canvas.save(out_png, optimize=True) + + source = { + "question": q, + "layers": summary["layers"], + "image_tokens": summary["image_tokens"], + "global_answer_over_random_delta": summary["answer_over_random_delta"], + "edge_formula": "max(answer_delta - random_delta, 0) * log1p(ratio)", + "token_groups": group_rows, + "layers_metrics": layer_rows, + "selected_token_groups": selected_groups, + "selected_layers": selected_layers, + } + (out_dir / "circuit-source-data.json").write_text(json.dumps(source, indent=2)) + return {"png": str(out_png), "source": str(out_dir / "circuit-source-data.json")} + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--result-dir", default=str(DEFAULT_RESULT_DIR)) + ap.add_argument("--out-dir", default=str(DEFAULT_OUT_DIR)) + args = ap.parse_args() + + result_dir = Path(args.result_dir) + out_dir = Path(args.out_dir) + summary = json.loads((result_dir / "summary.json").read_text()) + data = np.load(result_dir / "heatmaps.npz") + paths = render(summary, data["answer_delta"], data["random_delta"], data["ratio"], result_dir, out_dir) + print(paths["png"]) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_viz_city.py b/packages/snapcompact/research/snapcompact_viz_city.py new file mode 100644 index 000000000..5b34f9fb9 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_viz_city.py @@ -0,0 +1,237 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy"] +# /// +"""Render an isometric activation city from snapcompact heatmap tensors.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Iterable + +import numpy as np +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +DATA_DIR = HERE / "results" / "tensor-heatmap-paddleocr-q7" +OUT_DIR = HERE / "results" / "agent-viz-city" + +BG = (8, 11, 24) +INK = (232, 238, 255) +MUTED = (136, 148, 184) +ANSWER = (255, 94, 117) +ANSWER_HI = (255, 198, 97) +RANDOM = (72, 201, 255) +RANDOM_HI = (127, 246, 213) +ROAD = (21, 28, 52) +GRID = (39, 48, 83) + + +def font(size: int, bold: bool = False) -> ImageFont.FreeTypeFont | ImageFont.ImageFont: + candidates = [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/System/Library/Fonts/Supplemental/Helvetica Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Helvetica.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ] + for path in candidates: + try: + return ImageFont.truetype(path, size) + except OSError: + pass + return ImageFont.load_default() + + +def clamp255(v: float) -> int: + return max(0, min(255, int(round(v)))) + + +def mix(a: tuple[int, int, int], b: tuple[int, int, int], t: float) -> tuple[int, int, int]: + t = max(0.0, min(1.0, t)) + return tuple(clamp255(x + (y - x) * t) for x, y in zip(a, b)) + + +def shade(c: tuple[int, int, int], factor: float) -> tuple[int, int, int]: + return tuple(clamp255(x * factor) for x in c) + + +def iso(x: float, y: float, origin: tuple[float, float], tile_w: float, tile_h: float) -> tuple[float, float]: + ox, oy = origin + return ox + (x - y) * tile_w * 0.5, oy + (x + y) * tile_h * 0.5 + + +def diamond(cx: float, cy: float, tile_w: float, tile_h: float) -> list[tuple[float, float]]: + return [ + (cx, cy - tile_h * 0.5), + (cx + tile_w * 0.5, cy), + (cx, cy + tile_h * 0.5), + (cx - tile_w * 0.5, cy), + ] + + +def building_faces(cx: float, cy: float, h: float, tile_w: float, tile_h: float) -> tuple[list[tuple[float, float]], list[tuple[float, float]], list[tuple[float, float]]]: + top = diamond(cx, cy - h, tile_w, tile_h) + right = [top[1], (cx + tile_w * 0.5, cy), (cx, cy + tile_h * 0.5), top[2]] + left = [top[3], top[2], (cx, cy + tile_h * 0.5), (cx - tile_w * 0.5, cy)] + return top, right, left + + +def draw_soft_line(draw: ImageDraw.ImageDraw, pts: Iterable[tuple[float, float]], fill: tuple[int, int, int], width: int = 1) -> None: + draw.line([(int(x), int(y)) for x, y in pts], fill=fill, width=width) + + +def draw_district( + draw: ImageDraw.ImageDraw, + values: np.ndarray, + ratios: np.ndarray, + origin: tuple[float, float], + base_color: tuple[int, int, int], + high_color: tuple[int, int, int], + label: str, + scale: float, + ratio_scale: float, +) -> None: + layers, bins = values.shape + tile_w = 9.0 + tile_h = 5.0 + max_h = 245.0 + + # Foundation grid and layer streets. + for layer in range(layers): + left = iso(0, layer, origin, tile_w, tile_h) + right = iso(bins - 1, layer, origin, tile_w, tile_h) + draw_soft_line(draw, [left, right], GRID if layer % 3 else (66, 76, 118), 1) + for token in range(0, bins, 10): + near = iso(token, 0, origin, tile_w, tile_h) + far = iso(token, layers - 1, origin, tile_w, tile_h) + draw_soft_line(draw, [near, far], (31, 39, 70), 1) + + # Draw far blocks first, near blocks last. + for layer in range(layers - 1, -1, -1): + for token in range(bins - 1, -1, -1): + v = float(values[layer, token]) + r = float(ratios[layer, token]) + intensity = min(1.0, np.log1p(v) / np.log1p(scale)) + h = 8.0 + (intensity**1.65) * max_h + if v <= 0.0: + h = 3.0 + cx, cy = iso(token, layer, origin, tile_w, tile_h) + top, right, left = building_faces(cx, cy, h, tile_w * 0.92, tile_h * 0.92) + ratio_t = min(1.0, np.log1p(max(r, 0.0)) / np.log1p(ratio_scale)) + c = mix(base_color, high_color, max(intensity * 0.55, ratio_t * 0.85)) + draw.polygon(left, fill=shade(c, 0.42)) + draw.polygon(right, fill=shade(c, 0.62)) + draw.polygon(top, fill=mix(shade(c, 0.95), (255, 255, 255), intensity * 0.20)) + if ratio_t > 0.80 or intensity > 0.90: + draw.line([(int(x), int(y)) for x, y in top + [top[0]]], fill=mix(c, (255, 255, 255), 0.25), width=1) + + # District label plaque. + x0, y0 = iso(-2, layers + 4, origin, tile_w, tile_h) + x1, y1 = iso(54, layers + 4, origin, tile_w, tile_h) + draw.rounded_rectangle((x0 - 26, y0 + 18, x1 + 26, y1 + 64), radius=14, fill=(13, 18, 37), outline=shade(base_color, 0.75), width=2) + draw.text((x0 - 8, y0 + 27), label, font=font(26, True), fill=mix(base_color, high_color, 0.55)) + + +def draw_legend(draw: ImageDraw.ImageDraw, summary: dict[str, object], scale: float) -> None: + draw.text((88, 70), "Snapcompact Activation City", font=font(54, True), fill=INK) + draw.text( + (92, 136), + "729 image tokens → 180 token-bin city blocks · 19 transformer layers → depth streets · building height = activation spike magnitude", + font=font(21), + fill=MUTED, + ) + question = str(summary.get("question", {}).get("q", "")) if isinstance(summary.get("question"), dict) else "" + answer = str(summary.get("question", {}).get("answer_text", "")) if isinstance(summary.get("question"), dict) else "" + draw.text((92, 173), f"Question: {question} Gold answer: {answer}", font=font(20), fill=(180, 190, 220)) + + ratio = float(summary.get("answer_over_random_delta", 0.0)) + draw.rounded_rectangle((1738, 72, 2286, 206), radius=24, fill=(12, 17, 36), outline=(50, 60, 99), width=2) + draw.text((1772, 96), "Answer-mask / random-mask mean delta", font=font(18), fill=MUTED) + draw.text((1772, 125), f"{ratio:.2f}×", font=font(52, True), fill=ANSWER_HI) + draw.text((1906, 149), f"common p98 height scale {scale:.1f}", font=font(17), fill=(176, 186, 218)) + + y = 1400 + draw.rounded_rectangle((88, y, 772, y + 96), radius=18, fill=(12, 17, 36), outline=(44, 54, 92), width=1) + draw.text((116, y + 18), "How to read it", font=font(22, True), fill=INK) + draw.text((116, y + 52), "Tall towers mark token/layer bins where masking changed hidden states most.", font=font(18), fill=MUTED) + draw.rounded_rectangle((836, y, 1520, y + 96), radius=18, fill=(12, 17, 36), outline=(44, 54, 92), width=1) + draw.text((864, y + 18), "Districts", font=font(22, True), fill=INK) + draw.text((864, y + 52), "Warm city = answer mask around “2003”; cool city = same-size random mask.", font=font(18), fill=MUTED) + draw.rounded_rectangle((1584, y, 2268, y + 96), radius=18, fill=(12, 17, 36), outline=(44, 54, 92), width=1) + draw.text((1612, y + 18), "Color halos", font=font(22, True), fill=INK) + draw.text((1612, y + 52), "Bright caps emphasize bins with high answer/random activation ratio.", font=font(18), fill=MUTED) + + +def render() -> None: + OUT_DIR.mkdir(parents=True, exist_ok=True) + heatmaps = np.load(DATA_DIR / "heatmaps.npz") + summary = json.loads((DATA_DIR / "summary.json").read_text()) + + answer = np.asarray(heatmaps["answer_binned"], dtype=np.float32) + random = np.asarray(heatmaps["random_binned"], dtype=np.float32) + ratio = np.asarray(heatmaps["ratio_binned"], dtype=np.float32) + if answer.shape != random.shape or answer.shape != ratio.shape: + raise ValueError(f"expected matching binned shapes, got {answer.shape}, {random.shape}, {ratio.shape}") + + scale = float(summary.get("common_delta_scale_p98") or np.percentile(np.concatenate([answer.ravel(), random.ravel()]), 98)) + ratio_scale = float(summary.get("ratio_scale_p98") or np.percentile(ratio, 98)) + + w, h = 2400, 1600 + canvas = Image.new("RGB", (w, h), BG) + + # Atmospheric glow layer. + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((90, 250, 1090, 1260), fill=(255, 80, 95, 36)) + gd.ellipse((1220, 250, 2250, 1260), fill=(62, 190, 255, 34)) + gd.rectangle((0, 1240, w, h), fill=(2, 4, 11, 90)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(58))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + # Basemap plates. + draw.rounded_rectangle((52, 244, 1134, 1330), radius=42, fill=(9, 14, 30), outline=(42, 31, 55), width=2) + draw.rounded_rectangle((1234, 244, 2316, 1330), radius=42, fill=(8, 15, 31), outline=(24, 50, 70), width=2) + for y in range(312, 1300, 70): + draw.line((70, y, 1116, y), fill=ROAD, width=1) + draw.line((1252, y, 2298, y), fill=ROAD, width=1) + + draw_district(draw, answer, ratio, (252.0, 870.0), ANSWER, ANSWER_HI, "answer-mask district", scale, ratio_scale) + draw_district(draw, random, ratio, (1434.0, 870.0), RANDOM, RANDOM_HI, "random-mask district", scale, ratio_scale) + draw_legend(draw, summary, scale) + + # Fine vignette frame. + vignette = Image.new("L", (w, h), 0) + vd = ImageDraw.Draw(vignette) + vd.rectangle((0, 0, w, h), fill=255) + vignette = vignette.filter(ImageFilter.GaussianBlur(42)) + frame = Image.new("RGB", (w, h), (0, 0, 0)) + canvas = Image.composite(canvas, frame, vignette) + + np.savez_compressed( + OUT_DIR / "city_source_data.npz", + answer_binned=answer, + random_binned=random, + ratio_binned=ratio, + common_delta_scale_p98=np.array(scale, dtype=np.float32), + ratio_scale_p98=np.array(ratio_scale, dtype=np.float32), + ) + (OUT_DIR / "city_summary.json").write_text( + json.dumps( + { + "source_heatmaps": str(DATA_DIR / "heatmaps.npz"), + "source_summary": str(DATA_DIR / "summary.json"), + "shape": list(answer.shape), + "height_encoding": "log1p(delta) scaled by common_delta_scale_p98", + "districts": {"answer": "answer_binned", "random": "random_binned"}, + "ratio_encoding": "bright caps use ratio_binned / ratio_scale_p98", + "answer_over_random_delta": summary.get("answer_over_random_delta"), + }, + indent=2, + ) + + "\n" + ) + canvas.save(OUT_DIR / "city.png") + + +if __name__ == "__main__": + render() diff --git a/packages/snapcompact/research/snapcompact_viz_explainer.py b/packages/snapcompact/research/snapcompact_viz_explainer.py new file mode 100644 index 000000000..bc89e27f3 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_viz_explainer.py @@ -0,0 +1,444 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy"] +# /// +"""Render a single-frame snapcompact white-box explainer composite. + +The figure uses the saved tensor heatmaps from the PaddleOCR-VL run. It lays out +four linked stages: the source bitmap, the answer-region mask intervention, the +layer/token hidden-state tensor, and the interpretation stats. +""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path +from typing import Iterable + +import numpy as np +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +DEFAULT_DATA = HERE / "results" / "tensor-heatmap-paddleocr-q7" +DEFAULT_OUT = HERE / "results" / "agent-viz-explainer" + +BG = (4, 7, 13) +PANEL = (13, 18, 27) +PANEL_2 = (17, 24, 36) +INK = (244, 243, 231) +MUTED = (137, 153, 166) +DIM = (72, 84, 98) +CYAN = (78, 219, 255) +RED = (255, 82, 65) +AMBER = (255, 197, 78) +GREEN = (129, 255, 136) +VIOLET = (172, 116, 255) +LINE = (37, 50, 64) + + +def font(size: int, bold: bool = False) -> ImageFont.FreeTypeFont | ImageFont.ImageFont: + names = [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/System/Library/Fonts/Supplemental/Helvetica Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Helvetica.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ] + for name in names: + if name and Path(name).exists(): + return ImageFont.truetype(name, size) + return ImageFont.load_default() + + +F10 = font(10) +F12 = font(12) +F14 = font(14) +F16 = font(16) +F18 = font(18) +F20 = font(20) +F22 = font(22, True) +F26 = font(26, True) +F30 = font(30, True) +F38 = font(38, True) +F56 = font(56, True) +F72 = font(72, True) + + +def lerp(a: int, b: int, t: float) -> int: + return round(a + (b - a) * t) + + +def mix(a: tuple[int, int, int], b: tuple[int, int, int], t: float) -> tuple[int, int, int]: + return tuple(lerp(a[i], b[i], t) for i in range(3)) + + +def heat_color(t: float) -> tuple[int, int, int]: + t = max(0.0, min(1.0, float(t))) + stops = [ + (0.00, (5, 8, 20)), + (0.16, (22, 28, 75)), + (0.34, (81, 40, 125)), + (0.56, (208, 54, 101)), + (0.76, (255, 130, 69)), + (0.91, (255, 210, 94)), + (1.00, (255, 252, 200)), + ] + for (x0, c0), (x1, c1) in zip(stops, stops[1:]): + if t <= x1: + return mix(c0, c1, (t - x0) / (x1 - x0)) + return stops[-1][1] + + +def blue_color(t: float) -> tuple[int, int, int]: + t = max(0.0, min(1.0, float(t))) + return mix((8, 13, 28), CYAN, t**0.75) + + +def paste_round(base: Image.Image, img: Image.Image, box: tuple[int, int, int, int], radius: int = 24) -> None: + x0, y0, x1, y1 = box + img = img.convert("RGB") + scale = min((x1 - x0) / img.width, (y1 - y0) / img.height) + resized = img.resize((max(1, round(img.width * scale)), max(1, round(img.height * scale))), Image.Resampling.LANCZOS) + px = x0 + (x1 - x0 - resized.width) // 2 + py = y0 + (y1 - y0 - resized.height) // 2 + mask = Image.new("L", resized.size, 0) + ImageDraw.Draw(mask).rounded_rectangle((0, 0, resized.width - 1, resized.height - 1), radius=radius, fill=255) + base.paste(resized, (px, py), mask) + + +def draw_panel(draw: ImageDraw.ImageDraw, box: tuple[int, int, int, int], title: str, subtitle: str, accent: tuple[int, int, int]) -> None: + x0, y0, x1, y1 = box + draw.rounded_rectangle(box, radius=30, fill=PANEL, outline=LINE, width=2) + draw.rectangle((x0 + 28, y0 + 22, x0 + 84, y0 + 28), fill=accent) + draw.text((x0 + 28, y0 + 43), title, fill=INK, font=F26) + draw.text((x0 + 28, y0 + 78), subtitle, fill=MUTED, font=F16) + + +def draw_arrow(draw: ImageDraw.ImageDraw, start: tuple[int, int], end: tuple[int, int], color: tuple[int, int, int], label: str) -> None: + sx, sy = start + ex, ey = end + draw.line((sx, sy, ex - 18, ey), fill=color, width=5) + draw.polygon([(ex, ey), (ex - 22, ey - 14), (ex - 22, ey + 14)], fill=color) + if not label: + return + tw = round(draw.textlength(label, font=F14)) + draw.rounded_rectangle((sx + 20, sy - 31, sx + 42 + tw, sy - 6), radius=12, fill=(9, 14, 24), outline=mix(color, LINE, 0.35)) + draw.text((sx + 31, sy - 29), label, fill=color, font=F14) + + +def cell_box(summary: dict) -> tuple[int, int, int, int]: + q = summary["question"] + g = summary["geometry"] + cols = int(g["cols"]) + rows = int(g["rows"]) + start = int(q["answer_start"]) + end = int(q["answer_end"]) + cw = 768 / cols + ch = 768 / rows + r0, c0 = divmod(start, cols) + r1, c1 = divmod(max(start, end - 1), cols) + return (math.floor(c0 * cw), math.floor(r0 * ch), math.ceil((c1 + 1) * cw), math.ceil((r1 + 1) * ch)) + + +def answer_crop(img: Image.Image, summary: dict, pad_cells: int = 31) -> tuple[Image.Image, tuple[int, int, int, int]]: + g = summary["geometry"] + cols = int(g["cols"]) + rows = int(g["rows"]) + q = summary["question"] + start = int(q["answer_start"]) + end = int(q["answer_end"]) + row = start // cols + col0 = start % cols + col1 = (end - 1) % cols + 1 + cw = img.width / cols + ch = img.height / rows + x0 = max(0, math.floor((col0 - pad_cells) * cw)) + x1 = min(img.width, math.ceil((col1 + pad_cells) * cw)) + y0 = max(0, math.floor((row - 5) * ch)) + y1 = min(img.height, math.ceil((row + 6) * ch)) + crop = img.crop((x0, y0, x1, y1)).convert("RGB") + local = (round(col0 * cw - x0), round(row * ch - y0), round(col1 * cw - x0), round((row + 1) * ch - y0)) + return crop, local + + +def draw_crop_card(canvas: Image.Image, box: tuple[int, int, int, int], img: Image.Image, local_box: tuple[int, int, int, int], title: str, accent: tuple[int, int, int]) -> None: + draw = ImageDraw.Draw(canvas) + x0, y0, x1, y1 = box + draw.text((x0, y0 - 28), title, fill=accent, font=F16) + draw.rounded_rectangle(box, radius=18, fill=(236, 234, 219), outline=accent, width=3) + pad = 14 + scale = min((x1 - x0 - 2 * pad) / img.width, (y1 - y0 - 2 * pad) / img.height) + resized = img.resize((round(img.width * scale), round(img.height * scale)), Image.Resampling.NEAREST) + px = x0 + (x1 - x0 - resized.width) // 2 + py = y0 + (y1 - y0 - resized.height) // 2 + canvas.paste(resized, (px, py)) + bx = tuple(round(v * scale) for v in local_box) + draw.rounded_rectangle((px + bx[0] - 4, py + bx[1] - 4, px + bx[2] + 4, py + bx[3] + 4), radius=6, outline=accent, width=4) + + +def draw_heatmap(draw: ImageDraw.ImageDraw, arr: np.ndarray, box: tuple[int, int, int, int]) -> None: + x0, y0, x1, y1 = box + rows, cols = arr.shape + cw = (x1 - x0) / cols + ch = (y1 - y0) / rows + for r in range(rows): + ya = round(y0 + r * ch) + yb = round(y0 + (r + 1) * ch) + for c in range(cols): + xa = round(x0 + c * cw) + xb = round(x0 + (c + 1) * cw) + draw.rectangle((xa, ya, xb, yb), fill=heat_color(float(arr[r, c]))) + for r in range(rows + 1): + y = round(y0 + r * ch) + draw.line((x0, y, x1, y), fill=(0, 0, 0, 90) if False else (20, 26, 36), width=1) + draw.rectangle(box, outline=(83, 101, 118), width=1) + + +def draw_tensor_ribbons(draw: ImageDraw.ImageDraw, answer: np.ndarray, random: np.ndarray, box: tuple[int, int, int, int]) -> None: + x0, y0, x1, y1 = box + rows, cols = answer.shape + lane_h = (y1 - y0) / rows + for r in range(rows): + ya = y0 + r * lane_h + yb = y0 + (r + 0.68) * lane_h + for c in range(cols): + xa = x0 + c * (x1 - x0) / cols + xb = x0 + (c + 1) * (x1 - x0) / cols + a = float(answer[r, c]) + rr = float(random[r, c]) + color = heat_color(a) + if rr > a * 0.86: + color = mix(color, (45, 88, 77), min(0.50, rr * 0.45)) + draw.rectangle((round(xa), round(ya), round(xb), round(yb)), fill=color) + if r % 3 == 0: + draw.text((x0 - 33, round(ya + 2)), f"L{r}", fill=MUTED, font=F12) + draw.rectangle(box, outline=(91, 106, 122), width=1) + + +def draw_token_grid(draw: ImageDraw.ImageDraw, ratio: np.ndarray, box: tuple[int, int, int, int]) -> None: + x0, y0, x1, y1 = box + grid = ratio.mean(axis=0).reshape(27, 27) + q98 = float(np.quantile(grid, 0.98)) or 1.0 + norm = np.clip(grid / q98, 0, 1) + cell = min((x1 - x0) / 27, (y1 - y0) / 27) + ox = x0 + ((x1 - x0) - 27 * cell) / 2 + oy = y0 + ((y1 - y0) - 27 * cell) / 2 + for r in range(27): + for c in range(27): + xa = round(ox + c * cell) + ya = round(oy + r * cell) + xb = round(ox + (c + 1) * cell - 1) + yb = round(oy + (r + 1) * cell - 1) + draw.rounded_rectangle((xa, ya, xb, yb), radius=3, fill=blue_color(float(norm[r, c]))) + top = np.unravel_index(np.argsort(grid, axis=None)[-6:], grid.shape) + for r, c in zip(top[0], top[1]): + xa = round(ox + c * cell) + ya = round(oy + r * cell) + draw.rounded_rectangle((xa - 2, ya - 2, round(xa + cell + 1), round(ya + cell + 1)), radius=4, outline=AMBER, width=2) + + +def polyline(draw: ImageDraw.ImageDraw, values: Iterable[float], box: tuple[int, int, int, int], color: tuple[int, int, int], width: int = 4) -> None: + vals = list(values) + x0, y0, x1, y1 = box + lo = min(vals) + hi = max(vals) + span = hi - lo if hi > lo else 1.0 + points = [] + for i, v in enumerate(vals): + x = x0 + i * (x1 - x0) / max(1, len(vals) - 1) + y = y1 - ((v - lo) / span) * (y1 - y0) + points.append((round(x), round(y))) + for i in range(1, len(points)): + draw.line((points[i - 1], points[i]), fill=color, width=width) + for x, y in points: + draw.ellipse((x - 3, y - 3, x + 3, y + 3), fill=color) + + +def metric(draw: ImageDraw.ImageDraw, box: tuple[int, int, int, int], label: str, value: str, sub: str, accent: tuple[int, int, int]) -> None: + x0, y0, x1, y1 = box + draw.rounded_rectangle(box, radius=20, fill=PANEL_2, outline=mix(accent, LINE, 0.35), width=2) + draw.text((x0 + 18, y0 + 16), label, fill=MUTED, font=F14) + draw.text((x0 + 18, y0 + 41), value, fill=accent, font=F38) + draw.text((x0 + 18, y1 - 32), sub, fill=INK, font=F14) + + +def wrap_text(draw: ImageDraw.ImageDraw, text: str, max_width: int, fnt: ImageFont.ImageFont) -> list[str]: + words = text.split() + lines: list[str] = [] + cur = "" + for word in words: + trial = word if not cur else f"{cur} {word}" + if draw.textlength(trial, font=fnt) <= max_width: + cur = trial + else: + if cur: + lines.append(cur) + cur = word + if cur: + lines.append(cur) + return lines + + +def save_source_metrics(out_dir: Path, summary: dict, arrays: dict[str, np.ndarray]) -> None: + ratio = arrays["ratio"] + answer = arrays["answer_delta"] + random = arrays["random_delta"] + top_cell = np.unravel_index(int(np.argmax(ratio)), ratio.shape) + metrics = { + "question": summary["question"]["q"], + "answer": summary["question"]["answer_text"], + "layers": int(summary["layers"]), + "image_tokens": int(summary["image_tokens"]), + "answer_delta_mean": float(summary["answer_delta_mean"]), + "random_delta_mean": float(summary["random_delta_mean"]), + "answer_over_random_delta": float(summary["answer_over_random_delta"]), + "max_ratio_layer": int(top_cell[0]), + "max_ratio_token": int(top_cell[1]), + "max_ratio": float(ratio[top_cell]), + "mean_answer_by_layer": [float(x) for x in answer.mean(axis=1)], + "mean_random_by_layer": [float(x) for x in random.mean(axis=1)], + "mean_ratio_by_layer": [float(x) for x in ratio.mean(axis=1)], + } + out_dir.mkdir(parents=True, exist_ok=True) + (out_dir / "explainer_metrics.json").write_text(json.dumps(metrics, indent=2) + "\n") + + +def render(data_dir: Path, out_dir: Path) -> Path: + summary = json.loads((data_dir / "summary.json").read_text()) + npz = np.load(data_dir / "heatmaps.npz") + arrays = {name: npz[name] for name in npz.files} + original = Image.open(data_dir / "images" / "original.png").convert("RGB") + answer_mask = Image.open(data_dir / "images" / "answer-mask.png").convert("RGB") + + w, h = 2400, 1500 + canvas = Image.new("RGB", (w, h), BG) + draw = ImageDraw.Draw(canvas) + + # Subtle technical-paper background and activation glows. + for y in range(0, h, 16): + draw.line((0, y, w, y), fill=(6, 10 + (y % 5), 18 + (y % 7))) + for x in range(0, w, 32): + draw.line((x, 0, x, h), fill=(5, 8, 15)) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-300, -210, 980, 760), fill=(78, 219, 255, 35)) + gd.ellipse((690, 200, 1850, 1370), fill=(255, 82, 65, 32)) + gd.ellipse((1550, -120, 2660, 980), fill=(255, 197, 78, 24)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(90))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + draw.text((72, 52), "SNAPCOMPACT ACTIVATION EXPLAINER", fill=AMBER, font=F22) + draw.text((72, 86), "From erased pixels to a hidden-state scar", fill=INK, font=F72) + q = summary["question"]["q"] + subtitle = f"Question: {q} · gold answer: {summary['question']['answer_text']}" + draw.text((75, 172), subtitle, fill=MUTED, font=F22) + + # Stage panels. + p1 = (72, 240, 585, 1178) + p2 = (630, 240, 1143, 1178) + p3 = (1188, 240, 1770, 1178) + p4 = (1815, 240, 2328, 1178) + draw_panel(draw, p1, "1 · input bitmap", "Rendered context before intervention", CYAN) + draw_panel(draw, p2, "2 · mask intervention", "Only the true answer cells are blanked", RED) + draw_panel(draw, p3, "3 · hidden-state tensor", "Layer × image-token response", VIOLET) + draw_panel(draw, p4, "4 · interpretation", "Where the answer mattered most", AMBER) + draw_arrow(draw, (585, 715), (630, 715), CYAN, "") + draw_arrow(draw, (1143, 715), (1188, 715), RED, "") + draw_arrow(draw, (1770, 715), (1815, 715), AMBER, "") + + # Input / intervention panels. + paste_round(canvas, original, (108, 352, 548, 792), 24) + full_box = cell_box(summary) + scale = 440 / 768 + ox, oy = 108, 352 + draw.rounded_rectangle((ox + round(full_box[0] * scale), oy + round(full_box[1] * scale), ox + round(full_box[2] * scale), oy + round(full_box[3] * scale)), radius=5, outline=AMBER, width=4) + ocrop, local = answer_crop(original, summary) + mcrop, mlocal = answer_crop(answer_mask, summary) + draw_crop_card(canvas, (108, 885, 548, 1042), ocrop, local, "magnified answer glyphs", AMBER) + draw.text((108, 1083), "The OCR input is a fixed bitmap. The answer span", fill=MUTED, font=F16) + draw.text((108, 1108), f"occupies character cells {summary['question']['answer_start']}–{summary['question']['answer_end'] - 1}.", fill=MUTED, font=F16) + + paste_round(canvas, answer_mask, (666, 352, 1106, 792), 24) + draw.rounded_rectangle((666 + round(full_box[0] * scale), 352 + round(full_box[1] * scale), 666 + round(full_box[2] * scale), 352 + round(full_box[3] * scale)), radius=5, outline=RED, width=4) + draw_crop_card(canvas, (666, 885, 1106, 1042), mcrop, mlocal, "same crop after masking", RED) + draw.text((666, 1083), "Same prompt, same rendered page. Difference:", fill=MUTED, font=F16) + draw.text((666, 1108), "the four answer glyphs are removed before inference.", fill=MUTED, font=F16) + + # Tensor panel. + draw.text((1226, 332), "answer-mask delta", fill=RED, font=F18) + draw.text((1600, 332), "random control mixed in green", fill=GREEN, font=F14) + draw_tensor_ribbons(draw, arrays["answer_norm"], arrays["random_norm"], (1240, 372, 1718, 665)) + draw.text((1238, 686), "answer / random ratio", fill=AMBER, font=F18) + draw.text((1238, 712), "bright = answer-region deletion moves hidden states more than an equal random mask", fill=MUTED, font=F14) + draw_heatmap(draw, arrays["ratio_norm"], (1240, 748, 1718, 1000)) + max_layer = int(summary["max_ratio_layer"]) + draw.line((1240, 748 + round((max_layer + 0.5) * 252 / 19), 1718, 748 + round((max_layer + 0.5) * 252 / 19)), fill=AMBER, width=3) + for i in range(240): + draw.rectangle((1240 + i, 1046, 1241 + i, 1063), fill=heat_color(i / 239)) + draw.text((1240, 1022), "low", fill=MUTED, font=F12) + draw.text((1446, 1022), "high", fill=MUTED, font=F12) + draw.text((1240, 1094), f"{summary['layers']} decoder layers × {summary['image_tokens']} image tokens", fill=INK, font=F20) + draw.text((1240, 1124), "Each cell uses the saved heatmaps.npz tensor values.", fill=MUTED, font=F16) + + # Interpretation panel. + metric(draw, (1850, 344, 2075, 478), "mean delta ratio", f"{summary['answer_over_random_delta']:.2f}×", "answer mask vs control", AMBER) + metric(draw, (2086, 344, 2293, 478), "strongest layer", f"L{summary['max_ratio_layer']}", "mean ratio peak", VIOLET) + metric(draw, (1850, 500, 2075, 634), "answer delta", f"{summary['answer_delta_mean']:.2f}", "mean ||Δh||", RED) + metric(draw, (2086, 500, 2293, 634), "control delta", f"{summary['random_delta_mean']:.2f}", "mean ||Δh||", GREEN) + + draw.text((1852, 684), "layer sensitivity curve", fill=INK, font=F20) + curve_box = (1862, 725, 2290, 858) + draw.rounded_rectangle((1850, 704, 2304, 884), radius=20, fill=PANEL_2, outline=LINE, width=2) + for i in range(5): + y = curve_box[1] + i * (curve_box[3] - curve_box[1]) / 4 + draw.line((curve_box[0], round(y), curve_box[2], round(y)), fill=(31, 42, 54)) + polyline(draw, summary["mean_ratio_by_layer"], curve_box, AMBER, 4) + draw.text((1862, 862), "L0", fill=MUTED, font=F12) + draw.text((2262, 862), f"L{summary['layers'] - 1}", fill=MUTED, font=F12) + + draw.text((1852, 927), "image-token sensitivity field", fill=INK, font=F20) + draw.rounded_rectangle((1850, 955, 2067, 1150), radius=20, fill=PANEL_2, outline=LINE, width=2) + draw_token_grid(draw, arrays["ratio"], (1868, 970, 2049, 1132)) + explanation = "Answer deletion creates a high-ratio band in early layers; later layers diffuse it into surrounding context." + for i, line in enumerate(wrap_text(draw, explanation, 195, F14)): + draw.text((2092, 968 + i * 24), line, fill=INK if i == 0 else MUTED, font=F14) + draw.text((2092, 1090), "Interpretation:", fill=AMBER, font=F16) + draw.text((2092, 1118), "the answer glyphs are not just OCR text;", fill=MUTED, font=F14) + draw.text((2092, 1142), "they perturb the multimodal residual stream.", fill=MUTED, font=F14) + + # Footer with provenance. + footer = (72, 1228, 2328, 1422) + draw.rounded_rectangle(footer, radius=30, fill=(8, 12, 20), outline=LINE, width=2) + draw.text((108, 1266), "Reading the composite", fill=INK, font=F30) + bullets = [ + (CYAN, "Input bitmap", "is the rendered evidence page passed to PaddleOCR-VL."), + (RED, "Mask intervention", "removes only the gold answer span: 2003."), + (VIOLET, "Hidden-state tensor", "plots ||hidden(original) − hidden(masked)|| over saved layer/token arrays."), + (AMBER, "Interpretation", "compares that scar to an equal-size random mask: 2.52× stronger on average."), + ] + x = 108 + for color, head, text in bullets: + draw.rounded_rectangle((x, 1320, x + 500, 1384), radius=18, fill=PANEL_2, outline=mix(color, LINE, 0.35), width=2) + draw.ellipse((x + 18, 1343, x + 36, 1361), fill=color) + draw.text((x + 50, 1330), head, fill=color, font=F16) + draw.text((x + 50, 1355), text, fill=MUTED, font=F14) + x += 545 + + save_source_metrics(out_dir, summary, arrays) + out_dir.mkdir(parents=True, exist_ok=True) + out_path = out_dir / "explainer.png" + canvas.save(out_path, optimize=True) + return out_path + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--data-dir", type=Path, default=DEFAULT_DATA) + parser.add_argument("--out-dir", type=Path, default=DEFAULT_OUT) + args = parser.parse_args() + out = render(args.data_dir, args.out_dir) + print(out) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_viz_glass_stack.py b/packages/snapcompact/research/snapcompact_viz_glass_stack.py new file mode 100644 index 000000000..29b206203 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_viz_glass_stack.py @@ -0,0 +1,255 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy"] +# /// +"""Render a perspective glass-stack view of snapcompact activation deltas.""" + +from __future__ import annotations + +import json +import math +from pathlib import Path +from typing import Iterable + +import numpy as np +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +DATA_DIR = HERE / "results" / "tensor-heatmap-paddleocr-q7" +OUT_DIR = HERE / "results" / "agent-viz-glass-stack" + +INK = (238, 248, 255) +MUTED = (133, 158, 174) +CYAN = (82, 226, 255) +GOLD = (255, 206, 93) +RED = (255, 78, 93) +PANEL = (9, 16, 26) + + +def font(size: int, bold: bool = False) -> ImageFont.FreeTypeFont | ImageFont.ImageFont: + candidates = [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/System/Library/Fonts/Helvetica.ttc", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ] + for candidate in candidates: + if candidate and Path(candidate).exists(): + return ImageFont.truetype(candidate, size) + return ImageFont.load_default() + + +def mix(a: tuple[int, int, int], b: tuple[int, int, int], t: float) -> tuple[int, int, int]: + t = max(0.0, min(1.0, t)) + return tuple(round(a[i] + (b[i] - a[i]) * t) for i in range(3)) + + +def glass_heat(t: float) -> tuple[int, int, int]: + t = max(0.0, min(1.0, t)) + stops = [ + (0.00, (11, 24, 48)), + (0.24, (20, 67, 105)), + (0.50, (47, 183, 214)), + (0.72, (255, 91, 116)), + (0.88, (255, 178, 87)), + (1.00, (255, 252, 197)), + ] + for (ta, ca), (tb, cb) in zip(stops, stops[1:]): + if t <= tb: + return mix(ca, cb, (t - ta) / (tb - ta)) + return stops[-1][1] + + +def plane_corners(layer: int) -> tuple[tuple[float, float], tuple[float, float], tuple[float, float], tuple[float, float]]: + x = 225 + layer * 25.0 + y = 920 - layer * 34.0 + width = 930.0 + dx, dy = 190.0, -78.0 + return (x, y), (x + width, y), (x + width + dx, y + dy), (x + dx, y + dy) + + +def bilerp(corners: tuple[tuple[float, float], tuple[float, float], tuple[float, float], tuple[float, float]], u: float, v: float) -> tuple[float, float]: + fl, fr, br, bl = corners + ax = fl[0] + (fr[0] - fl[0]) * u + ay = fl[1] + (fr[1] - fl[1]) * u + bx = bl[0] + (br[0] - bl[0]) * u + by = bl[1] + (br[1] - bl[1]) * u + return ax + (bx - ax) * v, ay + (by - ay) * v + + +def poly(points: Iterable[tuple[float, float]]) -> list[tuple[int, int]]: + return [(round(x), round(y)) for x, y in points] + + +def select_scars(ratio_norm: np.ndarray, answer_binned: np.ndarray, random_binned: np.ndarray, count: int = 7) -> list[int]: + advantage = np.maximum(answer_binned - random_binned, 0.0) + if float(advantage.max(initial=0.0)) > 0: + advantage = advantage / float(np.quantile(advantage, 0.985)) + score = ratio_norm.mean(axis=0) * 0.68 + np.clip(advantage, 0, 1).mean(axis=0) * 0.32 + order = np.argsort(score)[::-1] + chosen: list[int] = [] + for idx in order: + i = int(idx) + if all(abs(i - old) >= 11 for old in chosen): + chosen.append(i) + if len(chosen) == count: + break + return sorted(chosen) + + +def draw_background(canvas: Image.Image) -> None: + draw = ImageDraw.Draw(canvas) + width, height = canvas.size + for y in range(height): + t = y / max(1, height - 1) + draw.line((0, y, width, y), fill=mix((2, 6, 13), (10, 18, 32), t)) + grid = Image.new("RGBA", canvas.size, (0, 0, 0, 0)) + gd = ImageDraw.Draw(grid) + for x in range(-240, width + 240, 48): + gd.line((x, height, x + 620, 0), fill=(65, 145, 190, 16), width=1) + for y in range(92, height, 46): + gd.line((0, y, width, y - 138), fill=(65, 145, 190, 12), width=1) + gd.ellipse((-220, -240, 850, 560), fill=(42, 188, 255, 34)) + gd.ellipse((1120, 420, 2040, 1380), fill=(255, 65, 112, 35)) + canvas.alpha_composite(grid.filter(ImageFilter.GaussianBlur(0.4))) + + +def draw_plane(canvas: Image.Image, values: np.ndarray, layer: int) -> None: + corners = plane_corners(layer) + overlay = Image.new("RGBA", canvas.size, (0, 0, 0, 0)) + draw = ImageDraw.Draw(overlay, "RGBA") + cols = values.shape[0] + + # Blue glass substrate. + draw.polygon(poly(corners), fill=(28, 85, 122, 24), outline=(108, 222, 255, 54)) + + for c, raw in enumerate(values): + u0 = c / cols + u1 = (c + 1) / cols + shade = float(np.clip(raw, 0, 1)) + rgb = glass_heat(shade) + alpha = round(26 + 96 * math.pow(shade, 0.82)) + draw.polygon( + poly((bilerp(corners, u0, 0.03), bilerp(corners, u1, 0.03), bilerp(corners, u1, 0.97), bilerp(corners, u0, 0.97))), + fill=(*rgb, alpha), + ) + + for u in np.linspace(0, 1, 13): + draw.line(poly((bilerp(corners, float(u), 0), bilerp(corners, float(u), 1))), fill=(190, 242, 255, 28), width=1) + for v in np.linspace(0, 1, 5): + draw.line(poly((bilerp(corners, 0, float(v)), bilerp(corners, 1, float(v)))), fill=(190, 242, 255, 24), width=1) + draw.line(poly((corners[0], corners[1], corners[2], corners[3], corners[0])), fill=(174, 241, 255, 70), width=2) + if layer in (0, 6, 12, 18): + x, y = corners[0] + draw.text((round(x - 64), round(y - 10)), f"L{layer:02d}", fill=(178, 226, 239, 150), font=font(16, True)) + canvas.alpha_composite(overlay) + + +def draw_scars(canvas: Image.Image, scar_bins: list[int], ratio_norm: np.ndarray) -> None: + glow = Image.new("RGBA", canvas.size, (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow, "RGBA") + cols = ratio_norm.shape[1] + scar_colors = [(255, 238, 164), (255, 102, 132), (87, 236, 255), (255, 190, 80), (205, 111, 255), (255, 255, 255), (72, 255, 190)] + for n, c in enumerate(scar_bins): + u = (c + 0.5) / cols + pts = [bilerp(plane_corners(layer), u, 0.46) for layer in range(ratio_norm.shape[0])] + color = scar_colors[n % len(scar_colors)] + gd.line(poly(pts), fill=(*color, 120), width=12) + for layer, pt in enumerate(pts): + r = 5 + 12 * float(ratio_norm[layer, c]) + x, y = pt + gd.ellipse((x - r, y - r, x + r, y + r), fill=(*color, 90)) + canvas.alpha_composite(glow.filter(ImageFilter.GaussianBlur(13))) + + draw = ImageDraw.Draw(canvas, "RGBA") + for n, c in enumerate(scar_bins): + u = (c + 0.5) / cols + pts = [bilerp(plane_corners(layer), u, 0.46) for layer in range(ratio_norm.shape[0])] + color = scar_colors[n % len(scar_colors)] + draw.line(poly(pts), fill=(*color, 235), width=3) + top = pts[-1] + draw.text((round(top[0] + 10), round(top[1] - 14)), f"bin {c}", fill=(*color, 220), font=font(13, True)) + for layer, pt in enumerate(pts): + r = 2.0 + 4.5 * float(ratio_norm[layer, c]) + x, y = pt + draw.ellipse((x - r, y - r, x + r, y + r), fill=(255, 255, 230, 225), outline=(*color, 255), width=1) + + +def draw_labels(canvas: Image.Image, summary: dict, scar_bins: list[int], ratio_binned: np.ndarray) -> None: + draw = ImageDraw.Draw(canvas, "RGBA") + q = summary["question"]["q"] + answer = summary["question"]["answer_text"] + draw.text((70, 54), "SNAPCOMPACT GLASS STACK", fill=GOLD, font=font(22, True)) + draw.text((70, 88), "Answer-mask scars through decoder depth", fill=INK, font=font(54, True)) + draw.text((73, 154), f"Question: {q} · gold answer: {answer}", fill=MUTED, font=font(22)) + + x0, y0, x1, y1 = 70, 960, 770, 1110 + draw.rounded_rectangle((x0, y0, x1, y1), radius=24, fill=(7, 13, 22, 205), outline=(115, 217, 255, 72), width=1) + ratio = summary["answer_over_random_delta"] + draw.text((x0 + 26, y0 + 22), f"{ratio:.2f}×", fill=GOLD, font=font(48, True)) + draw.text((x0 + 170, y0 + 31), "mean answer-mask / random-mask delta", fill=INK, font=font(22, True)) + draw.text((x0 + 28, y0 + 86), f"{summary['layers']} semi-transparent decoder planes · {summary['image_tokens']} image tokens binned to {ratio_binned.shape[1]} columns", fill=MUTED, font=font(18)) + + lx0, ly0 = 1240, 930 + draw.rounded_rectangle((lx0, ly0, lx0 + 475, ly0 + 182), radius=24, fill=(7, 13, 22, 210), outline=(115, 217, 255, 70), width=1) + draw.text((lx0 + 24, ly0 + 22), "encoding", fill=INK, font=font(25, True)) + draw.text((lx0 + 24, ly0 + 61), "plane color = answer/random ratio", fill=MUTED, font=font(18)) + draw.text((lx0 + 24, ly0 + 92), "vertical scar = high-ratio token bin", fill=MUTED, font=font(18)) + draw.text((lx0 + 24, ly0 + 124), "selected bins: " + ", ".join(map(str, scar_bins)), fill=(203, 231, 240), font=font(17)) + + # Color ramp. + for i in range(220): + draw.rectangle((lx0 + 230 + i, ly0 + 30, lx0 + 231 + i, ly0 + 49), fill=(*glass_heat(i / 219), 255)) + draw.text((lx0 + 230, ly0 + 54), "low", fill=MUTED, font=font(13)) + draw.text((lx0 + 417, ly0 + 54), "high", fill=MUTED, font=font(13)) + + +def main() -> None: + OUT_DIR.mkdir(parents=True, exist_ok=True) + with (DATA_DIR / "summary.json").open("r", encoding="utf-8") as f: + summary = json.load(f) + heat = np.load(DATA_DIR / "heatmaps.npz") + ratio_norm = heat["ratio_norm"].astype(np.float32, copy=False) + ratio_binned = heat["ratio_binned"].astype(np.float32, copy=False) + answer_binned = heat["answer_binned"].astype(np.float32, copy=False) + random_binned = heat["random_binned"].astype(np.float32, copy=False) + scar_bins = select_scars(ratio_norm, answer_binned, random_binned) + + canvas = Image.new("RGBA", (1800, 1200), (0, 0, 0, 255)) + draw_background(canvas) + + # Paint upper layers first, lower layers last, so the stack reads as transparent sheets in perspective. + for layer in range(ratio_norm.shape[0] - 1, -1, -1): + draw_plane(canvas, ratio_norm[layer], layer) + draw_scars(canvas, scar_bins, ratio_norm) + draw_labels(canvas, summary, scar_bins, ratio_binned) + + ImageDraw.Draw(canvas).rounded_rectangle((42, 36, 1760, 1142), radius=38, outline=(128, 225, 255, 44), width=2) + out_path = OUT_DIR / "glass-stack.png" + canvas.convert("RGB").save(out_path, quality=95) + + np.savez_compressed( + OUT_DIR / "glass-stack-source-data.npz", + ratio_binned=ratio_binned, + ratio_norm=ratio_norm, + answer_minus_random_binned=answer_binned - random_binned, + scar_bins=np.array(scar_bins, dtype=np.int16), + ) + with (OUT_DIR / "glass-stack-source-summary.json").open("w", encoding="utf-8") as f: + json.dump( + { + "source_heatmaps": str(DATA_DIR / "heatmaps.npz"), + "source_summary": str(DATA_DIR / "summary.json"), + "output": str(out_path), + "scar_bins": scar_bins, + "answer_over_random_delta": summary["answer_over_random_delta"], + "layers": summary["layers"], + "image_tokens": summary["image_tokens"], + }, + f, + indent=2, + ) + print(out_path) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_viz_glyph_matrix.py b/packages/snapcompact/research/snapcompact_viz_glyph_matrix.py new file mode 100644 index 000000000..77dd63f19 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_viz_glyph_matrix.py @@ -0,0 +1,397 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy"] +# /// +"""Render a glyph-to-activation matrix for the snapcompact answer scar. + +The figure keeps the original OCR bitmap visible, then projects the 27x27 image-token +activation field from heatmaps.npz back onto that bitmap. The answer glyphs are +outlined in text-cell coordinates; high-scar image tokens are outlined in model-token +coordinates; side bars show how the same region changes across decoder layers. +""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path +from typing import Iterable + +import numpy as np +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +DEFAULT_SOURCE = HERE / "results" / "tensor-heatmap-paddleocr-q7" +DEFAULT_OUT = HERE / "results" / "agent-viz-glyph-matrix" + +PALETTE = { + "bg": (3, 6, 12), + "panel": (10, 15, 22), + "panel2": (13, 21, 30), + "ink": (244, 245, 232), + "muted": (139, 154, 166), + "grid": (37, 53, 66), + "cyan": (86, 224, 255), + "red": (255, 78, 69), + "amber": (255, 202, 82), + "green": (135, 255, 159), + "violet": (173, 116, 255), +} + + +def font(size: int, bold: bool = False) -> ImageFont.FreeTypeFont | ImageFont.ImageFont: + candidates = [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/System/Library/Fonts/Supplemental/Helvetica.ttc", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ] + for path in candidates: + if path and Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def mix(a: tuple[int, int, int], b: tuple[int, int, int], t: float) -> tuple[int, int, int]: + t = max(0.0, min(1.0, t)) + return tuple(round(a[i] + (b[i] - a[i]) * t) for i in range(3)) + + +def heat_color(t: float) -> tuple[int, int, int]: + t = max(0.0, min(1.0, t)) + stops = [ + (0.00, (8, 10, 26)), + (0.22, (33, 25, 81)), + (0.46, (117, 41, 117)), + (0.68, (224, 67, 79)), + (0.86, (255, 165, 73)), + (1.00, (255, 243, 174)), + ] + for (ta, ca), (tb, cb) in zip(stops, stops[1:]): + if t <= tb: + return mix(ca, cb, (t - ta) / (tb - ta)) + return stops[-1][1] + + +def quantile_norm(values: np.ndarray, q: float = 0.98) -> np.ndarray: + scale = float(np.quantile(values, q)) if values.size else 1.0 + if scale <= 0 or not math.isfinite(scale): + scale = 1.0 + return np.clip(values / scale, 0.0, 1.0) + + +def rounded_panel(draw: ImageDraw.ImageDraw, box: tuple[int, int, int, int], title: str | None = None, subtitle: str | None = None) -> None: + draw.rounded_rectangle(box, radius=24, fill=PALETTE["panel"], outline=(34, 48, 61), width=1) + if title: + draw.text((box[0] + 24, box[1] + 18), title, fill=PALETTE["ink"], font=font(28, True)) + if subtitle: + draw.text((box[0] + 24, box[1] + 54), subtitle, fill=PALETTE["muted"], font=font(16)) + + +def token_boxes(side: int, grid: int) -> list[tuple[int, int, int, int]]: + boxes = [] + for idx in range(grid * grid): + r, c = divmod(idx, grid) + x0 = round(c * side / grid) + y0 = round(r * side / grid) + x1 = round((c + 1) * side / grid) + y1 = round((r + 1) * side / grid) + boxes.append((x0, y0, x1, y1)) + return boxes + + +def intersect_area(a: tuple[float, float, float, float], b: tuple[float, float, float, float]) -> float: + x0 = max(a[0], b[0]) + y0 = max(a[1], b[1]) + x1 = min(a[2], b[2]) + y1 = min(a[3], b[3]) + return max(0.0, x1 - x0) * max(0.0, y1 - y0) + + +def answer_bbox(start: int, end: int, cols: int, adv: int, pitch: int) -> tuple[int, int, int, int]: + row0, col0 = divmod(start, cols) + row1, col1 = divmod(max(start, end - 1), cols) + x0 = max(0, col0 * adv) + y0 = max(0, row0 * pitch) + x1 = (col1 + 1) * adv + y1 = (row1 + 1) * pitch + return x0, y0, x1, y1 + + +def draw_text_wrapped(draw: ImageDraw.ImageDraw, xy: tuple[int, int], text: str, width: int, fill: tuple[int, int, int], size: int, bold: bool = False, line_gap: int = 4) -> int: + words = text.split() + lines: list[str] = [] + current = "" + f = font(size, bold) + for word in words: + candidate = word if not current else f"{current} {word}" + if draw.textbbox((0, 0), candidate, font=f)[2] <= width or not current: + current = candidate + else: + lines.append(current) + current = word + if current: + lines.append(current) + x, y = xy + for line in lines: + draw.text((x, y), line, fill=fill, font=f) + y += size + line_gap + return y + + +def paste_shadowed(canvas: Image.Image, img: Image.Image, xy: tuple[int, int]) -> None: + shadow = Image.new("RGBA", img.size, (0, 0, 0, 0)) + alpha = Image.new("L", img.size, 180) + shadow.putalpha(alpha) + canvas.alpha_composite(shadow.filter(ImageFilter.GaussianBlur(12)), (xy[0] + 8, xy[1] + 10)) + canvas.alpha_composite(img, xy) + + +def draw_activation_overlay( + original: Image.Image, + answer_mask: Image.Image, + token_score: np.ndarray, + top_tokens: Iterable[int], + answer_tokens: Iterable[int], + bbox: tuple[int, int, int, int], +) -> Image.Image: + side = original.width + grid = int(round(math.sqrt(token_score.size))) + if grid * grid != token_score.size: + raise ValueError(f"expected square image token grid, got {token_score.size}") + base = original.convert("RGBA") + tint = Image.new("RGBA", base.size, (0, 0, 0, 0)) + td = ImageDraw.Draw(tint) + boxes = token_boxes(side, grid) + for idx, box in enumerate(boxes): + t = float(token_score[idx]) + if t <= 0: + continue + r, g, b = heat_color(t) + alpha = round(25 + 142 * t) + td.rectangle(box, fill=(r, g, b, alpha)) + composite = Image.alpha_composite(base, tint) + draw = ImageDraw.Draw(composite) + + for idx in top_tokens: + box = boxes[int(idx)] + draw.rounded_rectangle(box, radius=3, outline=PALETTE["amber"] + (235,), width=3) + for idx in answer_tokens: + box = boxes[int(idx)] + draw.rounded_rectangle(box, radius=4, outline=PALETTE["cyan"] + (245,), width=4) + + glow = Image.new("RGBA", composite.size, (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + for w, a in ((16, 46), (9, 80), (4, 235)): + gd.rounded_rectangle((bbox[0] - 8, bbox[1] - 7, bbox[2] + 8, bbox[3] + 8), radius=8, outline=PALETTE["red"] + (a,), width=w) + composite = Image.alpha_composite(composite, glow.filter(ImageFilter.GaussianBlur(4))) + draw = ImageDraw.Draw(composite) + draw.rounded_rectangle((bbox[0] - 8, bbox[1] - 7, bbox[2] + 8, bbox[3] + 8), radius=8, outline=PALETTE["red"] + (255,), width=3) + + mask_delta = Image.blend(original.convert("RGB"), answer_mask.convert("RGB"), 0.42).convert("RGBA") + crop = mask_delta.crop((max(0, bbox[0] - 76), max(0, bbox[1] - 42), min(side, bbox[2] + 154), min(side, bbox[3] + 48))) + crop = crop.resize((crop.width * 3, crop.height * 3), Image.Resampling.NEAREST) + crop_draw = ImageDraw.Draw(crop) + scale = 3 + cx0 = (bbox[0] - max(0, bbox[0] - 76)) * scale + cy0 = (bbox[1] - max(0, bbox[1] - 42)) * scale + cx1 = (bbox[2] - max(0, bbox[0] - 76)) * scale + cy1 = (bbox[3] - max(0, bbox[1] - 42)) * scale + crop_draw.rounded_rectangle((cx0 - 4, cy0 - 4, cx1 + 4, cy1 + 4), radius=8, outline=PALETTE["red"] + (255,), width=5) + composite.alpha_composite(crop, (side - crop.width - 20, 20)) + draw = ImageDraw.Draw(composite) + draw.text((side - crop.width - 16, 20 + crop.height + 8), "answer glyph crop: original → masked", fill=PALETTE["ink"] + (235,), font=font(18, True)) + return composite + + +def draw_layer_bars( + draw: ImageDraw.ImageDraw, + box: tuple[int, int, int, int], + answer_layer: np.ndarray, + random_layer: np.ndarray, + answer_region_layer: np.ndarray, + ratio_layer: np.ndarray, +) -> None: + rounded_panel(draw, box, "layer-by-layer scar", "red = answer mask, green = equal random mask, cyan = answer glyph tokens") + x0, y0, x1, y1 = box + chart = (x0 + 74, y0 + 103, x1 - 34, y1 - 72) + rows = answer_layer.size + row_h = (chart[3] - chart[1]) / rows + scale = float(np.quantile(np.concatenate([answer_layer, random_layer, answer_region_layer]), 0.96)) + scale = max(scale, 1e-6) + for i in range(rows): + y = chart[1] + i * row_h + draw.text((x0 + 28, round(y + row_h * 0.18)), f"L{i:02d}", fill=PALETTE["muted"], font=font(12)) + max_w = chart[2] - chart[0] + aw = round(max_w * min(1.0, float(answer_layer[i]) / scale)) + rw = round(max_w * min(1.0, float(random_layer[i]) / scale)) + gw = round(max_w * min(1.0, float(answer_region_layer[i]) / scale)) + yy = round(y) + draw.rounded_rectangle((chart[0], yy + 2, chart[0] + aw, yy + 8), radius=3, fill=PALETTE["red"]) + draw.rounded_rectangle((chart[0], yy + 11, chart[0] + rw, yy + 17), radius=3, fill=PALETTE["green"]) + draw.rounded_rectangle((chart[0], yy + 20, chart[0] + gw, yy + 27), radius=3, fill=PALETTE["cyan"]) + ratio = float(ratio_layer[i]) + draw.text((chart[2] - 58, yy + 8), f"{ratio:4.1f}×", fill=PALETTE["amber"], font=font(13, True)) + draw.text((chart[0], y1 - 45), "Mean delta per decoder layer. Ratio labels compare answer-mask vs random-mask deltas.", fill=PALETTE["muted"], font=font(14)) + + +def draw_scar_strip(draw: ImageDraw.ImageDraw, box: tuple[int, int, int, int], ratio_norm: np.ndarray, answer_tokens: list[int], top_tokens: list[int]) -> None: + rounded_panel(draw, box, "token scar matrix", "decoder layers × image tokens; vertical lines locate answer glyphs and top scar bins") + x0, y0, x1, y1 = box + hx0, hy0, hx1, hy1 = x0 + 58, y0 + 90, x1 - 28, y1 - 54 + rows, cols = ratio_norm.shape + cw = (hx1 - hx0) / cols + ch = (hy1 - hy0) / rows + for r in range(rows): + ya = round(hy0 + r * ch) + yb = round(hy0 + (r + 1) * ch) + for c in range(cols): + xa = round(hx0 + c * cw) + xb = round(hx0 + (c + 1) * cw) + draw.rectangle((xa, ya, xb, yb), fill=heat_color(float(ratio_norm[r, c]))) + for tok in answer_tokens: + x = round(hx0 + (tok + 0.5) * (hx1 - hx0) / 729) + draw.line((x, hy0 - 8, x, hy1 + 8), fill=PALETTE["cyan"], width=2) + for tok in top_tokens[:12]: + x = round(hx0 + (tok + 0.5) * (hx1 - hx0) / 729) + draw.line((x, hy0, x, hy1), fill=PALETTE["amber"], width=1) + for r in range(0, rows, 4): + y = round(hy0 + (r + 0.5) * ch) + draw.text((x0 + 22, y - 7), str(r), fill=PALETTE["muted"], font=font(12)) + draw.text((hx0, y1 - 32), "image-token sequence →", fill=PALETTE["muted"], font=font(13)) + + +def draw_top_token_table(draw: ImageDraw.ImageDraw, box: tuple[int, int, int, int], top_tokens: list[int], ratio_mean: np.ndarray, answer_mean: np.ndarray, grid: int) -> None: + rounded_panel(draw, box, "highest-scar token bins", "actual heatmaps.npz token IDs") + x0, y0, _, y1 = box + y = y0 + 90 + row_gap = 26 + max_rows = max(1, min(8, (y1 - y - 18) // row_gap)) + bar_scale = max(1e-6, float(np.quantile(answer_mean, 0.98))) + for rank, tok in enumerate(top_tokens[:max_rows], start=1): + r, c = divmod(tok, grid) + draw.text((x0 + 28, y), f"{rank:02d}", fill=PALETTE["amber"], font=font(13, True)) + draw.text((x0 + 68, y), f"token {tok:03d}", fill=PALETTE["ink"], font=font(14, True)) + draw.text((x0 + 164, y), f"grid r{r:02d} c{c:02d}", fill=PALETTE["muted"], font=font(13)) + draw.text((x0 + 292, y), f"ratio {ratio_mean[tok]:.2f}×", fill=PALETTE["cyan"], font=font(13, True)) + bar_w = round(112 * min(1.0, float(answer_mean[tok]) / bar_scale)) + draw.rounded_rectangle((x0 + 408, y + 4, x0 + 408 + bar_w, y + 14), radius=4, fill=PALETTE["red"]) + y += row_gap + + +def render(source: Path, out_dir: Path) -> None: + summary = json.loads((source / "summary.json").read_text()) + data = np.load(source / "heatmaps.npz") + original = Image.open(source / "images" / "original.png").convert("RGB") + answer_mask = Image.open(source / "images" / "answer-mask.png").convert("RGB") + + answer_delta = data["answer_delta"].astype(np.float32) + random_delta = data["random_delta"].astype(np.float32) + ratio = data["ratio"].astype(np.float32) + ratio_binned = data["ratio_binned"].astype(np.float32) + ratio_norm_binned = data["ratio_norm"].astype(np.float32) + + q = summary["question"] + geom = summary["geometry"] + cols = int(geom["cols"]) + rows = int(geom["rows"]) + adv = original.width // cols + pitch = max(1, original.height // rows) + bbox = answer_bbox(int(q["answer_start"]), int(q["answer_end"]), cols, adv, pitch) + + token_count = answer_delta.shape[1] + grid = int(round(math.sqrt(token_count))) + boxes = token_boxes(original.width, grid) + answer_area = (bbox[0], bbox[1], bbox[2], bbox[3]) + answer_tokens = [i for i, b in enumerate(boxes) if intersect_area(answer_area, b) > 0] + if not answer_tokens: + center_x = (bbox[0] + bbox[2]) / 2 + center_y = (bbox[1] + bbox[3]) / 2 + answer_tokens = [min(token_count - 1, max(0, int(center_y / original.height * grid) * grid + int(center_x / original.width * grid)))] + + ratio_mean = ratio.mean(axis=0) + answer_mean = answer_delta.mean(axis=0) + token_score = quantile_norm(ratio_mean, 0.985) + answer_set = set(answer_tokens) + top_tokens = [int(i) for i in np.argsort(ratio_mean)[::-1] if int(i) not in answer_set][:24] + answer_region_layer = answer_delta[:, answer_tokens].mean(axis=1) + answer_layer = answer_delta.mean(axis=1) + random_layer = random_delta.mean(axis=1) + ratio_layer = answer_layer / np.maximum(random_layer, 1e-6) + + out_dir.mkdir(parents=True, exist_ok=True) + overlay = draw_activation_overlay(original, answer_mask, token_score, top_tokens[:18], answer_tokens, bbox) + overlay = overlay.resize((760, 760), Image.Resampling.LANCZOS) + + W, H = 1900, 1260 + canvas = Image.new("RGBA", (W, H), PALETTE["bg"] + (255,)) + draw = ImageDraw.Draw(canvas) + for y in range(0, H, 16): + draw.line((0, y, W, y), fill=(6, 11 + y % 17, 18 + y % 11, 255)) + glow = Image.new("RGBA", (W, H), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-220, 120, 900, 1040), fill=PALETTE["red"] + (28,)) + gd.ellipse((820, -260, 2100, 820), fill=PALETTE["cyan"] + (26,)) + gd.ellipse((980, 650, 2050, 1510), fill=PALETTE["violet"] + (20,)) + canvas = Image.alpha_composite(canvas, glow.filter(ImageFilter.GaussianBlur(78))) + draw = ImageDraw.Draw(canvas) + + draw.text((64, 42), "SNAPCOMPACT GLYPH MATRIX", fill=PALETTE["amber"], font=font(22, True)) + draw.text((64, 78), "The answer glyphs leave a hidden activation scar", fill=PALETTE["ink"], font=font(56, True)) + subtitle = "Original dense text bitmap, overlaid with answer/random activation ratios from 19 decoder layers × 729 image tokens." + draw.text((68, 145), subtitle, fill=PALETTE["muted"], font=font(22)) + + rounded_panel(draw, (52, 205, 862, 1066), "visible glyphs ↔ hidden tokens", "red box = actual answer cells; cyan = intersecting image tokens; amber = top scar bins") + paste_shadowed(canvas, overlay, (78, 282)) + draw.text((82, 1085), f"Question: {q['q']}", fill=PALETTE["ink"], font=font(21, True)) + draw.text((82, 1120), f"Gold answer: {q['answer_text']} · cells {q['answer_start']}–{q['answer_end'] - 1}", fill=PALETTE["amber"], font=font(24, True)) + draw.text((82, 1160), f"Answer/random mean delta: {summary['answer_over_random_delta']:.2f}×", fill=PALETTE["cyan"], font=font(22, True)) + + draw_layer_bars(draw, (900, 205, 1838, 628), answer_layer, random_layer, answer_region_layer, ratio_layer) + draw_scar_strip(draw, (900, 662, 1838, 930), ratio_norm_binned, answer_tokens, top_tokens) + draw_top_token_table(draw, (900, 964, 1838, 1196), top_tokens, ratio_mean, answer_mean, grid) + + for i in range(240): + draw.rectangle((1568 + i, 156, 1569 + i, 174), fill=heat_color(i / 239)) + draw.text((1568, 132), "activation ratio", fill=PALETTE["muted"], font=font(13)) + draw.text((1568, 179), "low", fill=PALETTE["muted"], font=font(12)) + draw.text((1776, 179), "high", fill=PALETTE["muted"], font=font(12)) + + png = out_dir / "glyph-matrix.png" + canvas.convert("RGB").save(png, quality=96) + + source_data = { + "source": str(source), + "question": q, + "geometry": {"text_cols": cols, "text_rows": rows, "glyph_adv": adv, "glyph_pitch": pitch, "image_token_grid": [grid, grid]}, + "answer_bbox_pixels": list(map(int, bbox)), + "answer_image_tokens": [int(x) for x in answer_tokens], + "top_scar_tokens": [ + { + "token": int(tok), + "row": int(tok // grid), + "col": int(tok % grid), + "mean_ratio": float(ratio_mean[tok]), + "mean_answer_delta": float(answer_mean[tok]), + } + for tok in top_tokens[:24] + ], + "mean_answer_delta": float(answer_delta.mean()), + "mean_random_delta": float(random_delta.mean()), + "answer_over_random_delta": float(summary["answer_over_random_delta"]), + "max_ratio_binned": float(ratio_binned.max()), + } + (out_dir / "glyph-matrix-data.json").write_text(json.dumps(source_data, indent=2)) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--source", type=Path, default=DEFAULT_SOURCE) + parser.add_argument("--out-dir", type=Path, default=DEFAULT_OUT) + args = parser.parse_args() + render(args.source, args.out_dir) + print(args.out_dir / "glyph-matrix.png") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_viz_radial.py b/packages/snapcompact/research/snapcompact_viz_radial.py new file mode 100644 index 000000000..a5f477d8c --- /dev/null +++ b/packages/snapcompact/research/snapcompact_viz_radial.py @@ -0,0 +1,219 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["matplotlib", "numpy"] +# /// +"""Render a radial sonar view of snapcompact answer/random activation echoes.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") + +import matplotlib.colors as mcolors +import matplotlib.patheffects as pe +import matplotlib.pyplot as plt +import numpy as np +from matplotlib.patches import Wedge + +HERE = Path(__file__).resolve().parent +DATA_DIR = HERE / "results" / "tensor-heatmap-paddleocr-q7" +OUT_DIR = HERE / "results" / "agent-viz-radial" + +BG = "#02060a" +GRID = "#2cf5d044" +CYAN = "#38f4ff" +GREEN = "#81ffb4" +AMBER = "#ffc247" +RED = "#ff4d42" +INK = "#f3f1dd" +MUTED = "#8da0a8" + + +def robust_norm(values: np.ndarray, q: float = 0.975) -> np.ndarray: + scale = float(np.quantile(values, q)) + if not np.isfinite(scale) or scale <= 0: + scale = float(np.max(values)) or 1.0 + return np.clip(values / scale, 0.0, 1.0) + + +def polar_edges(cols: int, layers: int) -> tuple[np.ndarray, np.ndarray]: + theta = np.linspace(0.0, 2.0 * np.pi, cols + 1) + radius = np.arange(layers + 1, dtype=np.float32) + 1.0 + return theta, radius + + +def radar_cmap() -> mcolors.LinearSegmentedColormap: + colors = [ + (0.00, "#02060a"), + (0.10, "#03241f"), + (0.32, "#08705e"), + (0.55, "#16f0be"), + (0.76, "#fff06a"), + (1.00, "#fff8e0"), + ] + return mcolors.LinearSegmentedColormap.from_list("snapcompact_radar", colors) + + +def top_echoes(ratio: np.ndarray, answer: np.ndarray, random: np.ndarray, limit: int = 18) -> list[dict[str, float | int]]: + flat = np.argpartition(ratio.ravel(), -limit)[-limit:] + flat = flat[np.argsort(ratio.ravel()[flat])[::-1]] + rows: list[dict[str, float | int]] = [] + for idx in flat: + layer, bin_idx = np.unravel_index(int(idx), ratio.shape) + rows.append( + { + "rank": len(rows) + 1, + "layer": int(layer), + "bin": int(bin_idx), + "angle_degrees": round(float((bin_idx + 0.5) * 360.0 / ratio.shape[1]), 2), + "answer_delta": round(float(answer[layer, bin_idx]), 4), + "random_delta": round(float(random[layer, bin_idx]), 4), + "answer_random_ratio": round(float(ratio[layer, bin_idx]), 4), + } + ) + return rows + + +def add_glow_spikes(ax: plt.Axes, ratio: np.ndarray, norm_ratio: np.ndarray) -> None: + layers, bins = ratio.shape + theta_centers = (np.arange(bins) + 0.5) * 2.0 * np.pi / bins + threshold = float(np.quantile(norm_ratio, 0.91)) + for layer in range(layers): + active = np.flatnonzero(norm_ratio[layer] >= threshold) + if active.size == 0: + active = np.argpartition(norm_ratio[layer], -3)[-3:] + for idx in active: + v = float(norm_ratio[layer, idx]) + base_r = layer + 1.18 + tip_r = base_r + 0.12 + 0.58 * v + theta = float(theta_centers[idx]) + color = AMBER if v > 0.78 else CYAN + ax.plot([theta, theta], [base_r, tip_r], color=color, linewidth=0.7 + 1.8 * v, alpha=0.30 + 0.55 * v) + ax.scatter([theta], [tip_r], s=5 + 28 * v, color=color, alpha=0.26 + 0.55 * v, linewidths=0) + + +def draw_radial(summary: dict, answer: np.ndarray, random: np.ndarray, ratio: np.ndarray) -> plt.Figure: + layers, bins = ratio.shape + norm_ratio = robust_norm(ratio, 0.972) + theta_edges, radius_edges = polar_edges(bins, layers) + theta_grid, radius_grid = np.meshgrid(theta_edges, radius_edges) + + fig = plt.figure(figsize=(16, 10), dpi=180, facecolor=BG) + ax = fig.add_axes([0.04, 0.04, 0.68, 0.90], projection="polar", facecolor=BG) + ax.set_theta_zero_location("N") + ax.set_theta_direction(-1) + ax.set_ylim(0, layers + 2.05) + ax.set_xticks(np.deg2rad(np.arange(0, 360, 30))) + ax.set_xticklabels([f"{d}°" for d in range(0, 360, 30)], color=MUTED, fontsize=8) + ax.set_yticks(np.arange(1, layers + 1) + 0.5) + ax.set_yticklabels([str(i) for i in range(layers)], color="#8da0a888", fontsize=7) + ax.grid(color=GRID, linewidth=0.6, alpha=0.55) + ax.spines["polar"].set_color("#38f4ff66") + ax.spines["polar"].set_linewidth(1.2) + + ax.pcolormesh(theta_grid, radius_grid, norm_ratio, cmap=radar_cmap(), shading="flat", alpha=0.96) + + # Soft trace underneath the hottest angular bearings, like phosphor persistence. + bearing_strength = norm_ratio.mean(axis=0) + norm_ratio.max(axis=0) * 0.42 + sweep_bin = int(np.argmax(bearing_strength)) + sweep_angle = float((sweep_bin + 0.5) * 360.0 / bins) + sweep_theta = np.deg2rad(sweep_angle) + for width, alpha in ((38, 0.055), (22, 0.075), (8, 0.14)): + half_width = np.deg2rad(width / 2) + theta = np.linspace(sweep_theta - half_width, sweep_theta + half_width, 80) + ax.fill_between(theta, 0.0, layers + 1.75, color=GREEN, alpha=alpha, linewidth=0) + + add_glow_spikes(ax, ratio, norm_ratio) + + for r in range(1, layers + 2): + ax.plot(np.linspace(0, 2 * np.pi, 360), np.full(360, r), color="#6fffe522", linewidth=0.55) + for deg in range(0, 360, 15): + th = np.deg2rad(deg) + ax.plot([th, th], [1, layers + 1.4], color="#6fffe516", linewidth=0.45) + + ax.text(0.5, 0.5, "ECHO\nCORE", color="#dff", fontsize=13, fontweight="bold", ha="center", va="center", transform=ax.transAxes) + ax.text(np.deg2rad(sweep_angle), layers + 1.35, "strongest bearing", color=GREEN, fontsize=8, ha="center", va="center") + + side = fig.add_axes([0.70, 0.06, 0.27, 0.86], facecolor=BG) + side.axis("off") + side.set_xlim(0, 1) + side.set_ylim(0, 1) + q = summary["question"] + ratio_mean = float(summary["answer_over_random_delta"]) + max_layer = int(summary.get("max_ratio_layer", int(np.argmax(ratio.mean(axis=1))))) + top = top_echoes(ratio, answer, random, 7) + max_echo = top[0] + + title_fx = [pe.withStroke(linewidth=4, foreground="#0b1918")] + side.text(0.00, 0.98, "SNAPCOMPACT RADAR", color=GREEN, fontsize=12, fontweight="bold", va="top") + side.text(0.00, 0.925, "Where the missing\nanswer echoes", color=INK, fontsize=27, fontweight="bold", va="top", linespacing=0.92, path_effects=title_fx) + side.text(0.00, 0.765, "Concentric rings are decoder layers. Angles are image-token bins. Bright spikes are answer-mask residuals divided by the random-mask control.", color=MUTED, fontsize=9.5, va="top", wrap=True) + + metrics = [ + ("gold answer", str(q["answer_text"]), AMBER), + ("question", q["q"], INK), + ("image tokens", f"{summary['image_tokens']:,}", CYAN), + ("layers", f"{summary['layers']}", CYAN), + ("mean answer/random Δ", f"{ratio_mean:.2f}×", AMBER), + ("max-ratio layer", f"L{max_layer}", GREEN), + ("loudest echo", f"L{max_echo['layer']} · bin {max_echo['bin']} · {max_echo['answer_random_ratio']:.1f}×", RED), + ] + y = 0.655 + for label, value, color in metrics: + side.text(0.00, y, label.upper(), color=MUTED, fontsize=7.2, fontweight="bold", va="top") + value_size = 12.6 if len(value) < 34 else 8.7 + side.text(0.00, y - 0.026, value, color=color, fontsize=value_size, fontweight="bold" if label != "question" else "normal", va="top", wrap=True) + y -= 0.075 if label != "question" else 0.105 + + side.text(0.00, y - 0.006, "TOP ECHOES", color=GREEN, fontsize=7.6, fontweight="bold", va="top") + y -= 0.040 + for row in top[:4]: + intensity = min(1.0, float(row["answer_random_ratio"]) / float(max_echo["answer_random_ratio"])) + side.plot([0.00, 0.36 * intensity], [y - 0.004, y - 0.004], color=AMBER, linewidth=3.2, alpha=0.35 + 0.55 * intensity, solid_capstyle="round") + side.text(0.40, y - 0.014, f"L{row['layer']:02d} bin {row['bin']:03d} {row['answer_random_ratio']:>5.1f}×", color=INK, fontsize=7.4, va="bottom", family="monospace") + y -= 0.032 + + # Tiny color scale and data provenance line. + grad_ax = fig.add_axes([0.708, 0.048, 0.19, 0.014], facecolor=BG) + grad_ax.imshow(np.linspace(0, 1, 512)[None, :], cmap=radar_cmap(), aspect="auto") + grad_ax.set_axis_off() + side.text(0.00, 0.006, "low ratio", color=MUTED, fontsize=7, va="bottom") + side.text(0.59, 0.006, "high answer echo", color=MUTED, fontsize=7, va="bottom") + fig.text(0.045, 0.018, "Actual heatmaps.npz arrays: ratio_binned, answer_binned, random_binned", color="#8da0a888", fontsize=8) + return fig + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--data-dir", default=str(DATA_DIR)) + parser.add_argument("--out-dir", default=str(OUT_DIR)) + args = parser.parse_args() + + data_dir = Path(args.data_dir) + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + summary = json.loads((data_dir / "summary.json").read_text()) + heatmaps = np.load(data_dir / "heatmaps.npz") + answer = np.asarray(heatmaps["answer_binned"], dtype=np.float32) + random = np.asarray(heatmaps["random_binned"], dtype=np.float32) + ratio = np.asarray(heatmaps["ratio_binned"], dtype=np.float32) + + fig = draw_radial(summary, answer, random, ratio) + out_png = out_dir / "radial.png" + fig.savefig(out_png, facecolor=BG) + plt.close(fig) + + echoes = top_echoes(ratio, answer, random, 24) + (out_dir / "radial_top_echoes.json").write_text(json.dumps({"source": str(data_dir / "heatmaps.npz"), "top_echoes": echoes}, indent=2) + "\n") + np.savez_compressed(out_dir / "radial_source.npz", answer_binned=answer, random_binned=random, ratio_binned=ratio, ratio_norm=robust_norm(ratio, 0.972)) + print(out_png) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_viz_token_grid.py b/packages/snapcompact/research/snapcompact_viz_token_grid.py new file mode 100644 index 000000000..7087390ad --- /dev/null +++ b/packages/snapcompact/research/snapcompact_viz_token_grid.py @@ -0,0 +1,396 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy"] +# /// +"""Render a spatial image-token map for the snapcompact white-box run. + +The PaddleOCR-VL processor reports a 1 x 54 x 54 visual patch grid, while the +recorded hidden states contain 729 image tokens. This script folds the token +axis back to 27 x 27 (2 x 2 patch merge) and projects answer-mask delta / ratio +onto the original bitmap so the scar is visible in image space. +""" + +from __future__ import annotations + +import json +import math +from pathlib import Path +from typing import Iterable + +import numpy as np +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +SOURCE = HERE / "results" / "tensor-heatmap-paddleocr-q7" +OUT_DIR = HERE / "results" / "agent-viz-token-grid" +OUT = OUT_DIR / "token-grid.png" + +PALETTE = { + "bg": (4, 6, 10), + "panel": (12, 17, 24), + "panel2": (17, 23, 31), + "ink": (244, 241, 225), + "muted": (139, 153, 163), + "grid": (49, 64, 75), + "cyan": (75, 218, 255), + "red": (255, 80, 66), + "amber": (255, 194, 72), + "green": (148, 255, 126), +} + + +def font(size: int, bold: bool = False) -> ImageFont.FreeTypeFont | ImageFont.ImageFont: + candidates = [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/System/Library/Fonts/Supplemental/Helvetica Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Helvetica.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ] + for candidate in candidates: + if candidate and Path(candidate).exists(): + return ImageFont.truetype(candidate, size) + return ImageFont.load_default() + + +def lerp(a: int, b: int, t: float) -> int: + return round(a + (b - a) * t) + + +def mix(a: tuple[int, int, int], b: tuple[int, int, int], t: float) -> tuple[int, int, int]: + return (lerp(a[0], b[0], t), lerp(a[1], b[1], t), lerp(a[2], b[2], t)) + + +def heat_color(t: float) -> tuple[int, int, int]: + t = max(0.0, min(1.0, float(t))) + stops = [ + (0.00, (7, 12, 25)), + (0.18, (28, 24, 91)), + (0.38, (113, 31, 112)), + (0.62, (220, 61, 72)), + (0.82, (255, 152, 67)), + (1.00, (255, 242, 158)), + ] + for (pa, ca), (pb, cb) in zip(stops, stops[1:]): + if t <= pb: + return mix(ca, cb, (t - pa) / (pb - pa)) + return stops[-1][1] + + +def normalize(arr: np.ndarray, q: float = 0.985) -> tuple[np.ndarray, float]: + scale = float(np.quantile(arr, q)) if arr.size else 1.0 + if not math.isfinite(scale) or scale <= 0: + scale = 1.0 + return np.clip(arr / scale, 0, 1), scale + + +def token_side(summary: dict, token_count: int) -> int: + side = math.isqrt(token_count) + if side * side == token_count: + return side + grid = summary.get("processor_meta", {}).get("image_grid_thw", [[1, 0, 0]])[0] + _, gh, gw = grid + merge = math.isqrt(max(1, (gh * gw) // token_count)) + if merge and gh % merge == 0 and gw % merge == 0 and (gh // merge) * (gw // merge) == token_count: + return gh // merge + raise ValueError(f"cannot fold {token_count} image tokens into a square grid") + + +def fold_tokens(arr: np.ndarray, side: int) -> np.ndarray: + if arr.ndim == 1: + return arr.reshape(side, side) + return arr.reshape(arr.shape[0], side, side) + + +def heat_overlay(base: Image.Image, heat: np.ndarray, alpha_floor: int = 28, alpha_peak: int = 220) -> Image.Image: + norm, _ = normalize(heat) + small = Image.new("RGBA", (heat.shape[1], heat.shape[0]), (0, 0, 0, 0)) + pix = small.load() + for y in range(heat.shape[0]): + for x in range(heat.shape[1]): + t = float(norm[y, x]) + r, g, b = heat_color(t) + pix[x, y] = (r, g, b, round(alpha_floor + (alpha_peak - alpha_floor) * (t ** 0.85))) + overlay = small.resize(base.size, Image.Resampling.BICUBIC).filter(ImageFilter.GaussianBlur(1.0)) + dim = Image.blend(base.convert("RGB"), Image.new("RGB", base.size, (5, 8, 13)), 0.28).convert("RGBA") + return Image.alpha_composite(dim, overlay).convert("RGB") + + +def draw_token_grid(draw: ImageDraw.ImageDraw, box: tuple[int, int, int, int], side: int, color: tuple[int, int, int] = (255, 255, 255)) -> None: + x0, y0, x1, y1 = box + for i in range(side + 1): + x = round(x0 + (x1 - x0) * i / side) + y = round(y0 + (y1 - y0) * i / side) + fill = (*color, 36) if hasattr(draw, "mode") else color + draw.line((x, y0, x, y1), fill=fill, width=1) + draw.line((x0, y, x1, y), fill=fill, width=1) + + +def paste_fit(canvas: Image.Image, img: Image.Image, box: tuple[int, int, int, int], resample: int = Image.Resampling.LANCZOS) -> tuple[int, int, int, int]: + x0, y0, x1, y1 = box + scale = min((x1 - x0) / img.width, (y1 - y0) / img.height) + w = max(1, round(img.width * scale)) + h = max(1, round(img.height * scale)) + resized = img.resize((w, h), resample) + px = x0 + (x1 - x0 - w) // 2 + py = y0 + (y1 - y0 - h) // 2 + canvas.paste(resized, (px, py)) + return (px, py, px + w, py + h) + + +def crop_answer(img: Image.Image, start: int, end: int, cols: int, adv: int, pitch: int, pad_cells: int = 34) -> Image.Image: + rows = img.height // pitch + row0 = max(0, start // cols - 5) + row1 = min(rows, end // cols + 6) + col0 = max(0, start % cols - pad_cells) + col1 = min(cols, end % cols + pad_cells) + if col1 <= col0: + col1 = min(cols, col0 + 72) + crop = img.crop((col0 * adv, row0 * pitch, col1 * adv, row1 * pitch)).convert("RGB") + d = ImageDraw.Draw(crop) + bx0 = max(0, (start % cols - col0) * adv - adv) + bx1 = min(crop.width - 1, ((end - 1) % cols - col0 + 2) * adv) + by0 = max(0, (start // cols - row0) * pitch - 1) + by1 = min(crop.height - 1, ((end - 1) // cols - row0 + 1) * pitch + 1) + d.rounded_rectangle((bx0, by0, bx1, by1), radius=3, outline=PALETTE["red"], width=3) + return crop + + +def answer_bbox(start: int, end: int, cols: int, adv: int, pitch: int) -> tuple[int, int, int, int]: + return ( + max(0, (start % cols) * adv - adv), + max(0, (start // cols) * pitch - 2), + min(cols * adv, ((end - 1) % cols + 2) * adv), + ((end - 1) // cols + 1) * pitch + 2, + ) + + +def draw_panel(draw: ImageDraw.ImageDraw, box: tuple[int, int, int, int], title: str, subtitle: str | None = None) -> None: + draw.rounded_rectangle(box, radius=26, fill=PALETTE["panel"], outline=(32, 43, 55), width=1) + x0, y0, _, _ = box + draw.text((x0 + 24, y0 + 20), title, fill=PALETTE["ink"], font=font(28, True)) + if subtitle: + draw.text((x0 + 24, y0 + 56), subtitle, fill=PALETTE["muted"], font=font(17)) + + +def draw_micro_grid(canvas: Image.Image, heat: np.ndarray, box: tuple[int, int, int, int], title: str, subtitle: str) -> None: + draw = ImageDraw.Draw(canvas) + draw_panel(draw, box, title, subtitle) + x0, y0, x1, y1 = box + gx0, gy0, gx1, gy1 = x0 + 32, y0 + 96, x1 - 32, y1 - 42 + side = heat.shape[0] + norm, _ = normalize(heat) + cw = (gx1 - gx0) / side + ch = (gy1 - gy0) / side + for r in range(side): + for c in range(side): + xa = round(gx0 + c * cw) + ya = round(gy0 + r * ch) + xb = round(gx0 + (c + 1) * cw) + yb = round(gy0 + (r + 1) * ch) + draw.rectangle((xa, ya, xb, yb), fill=heat_color(float(norm[r, c]))) + for i in range(0, side + 1, 3): + x = round(gx0 + (gx1 - gx0) * i / side) + y = round(gy0 + (gy1 - gy0) * i / side) + draw.line((x, gy0, x, gy1), fill=(255, 255, 255, 34)) + draw.line((gx0, y, gx1, y), fill=(255, 255, 255, 34)) + + +def label(draw: ImageDraw.ImageDraw, xy: tuple[int, int], text: str, color: tuple[int, int, int], size: int = 18, bold: bool = True) -> None: + x, y = xy + pad = 8 + f = font(size, bold) + box = draw.textbbox((x, y), text, font=f) + draw.rounded_rectangle((box[0] - pad, box[1] - 4, box[2] + pad, box[3] + 5), radius=9, fill=(4, 6, 10), outline=color, width=1) + draw.text((x, y), text, fill=color, font=f) + + +def draw_hotspots(draw: ImageDraw.ImageDraw, box: tuple[int, int, int, int], heat: np.ndarray, count: int = 9) -> None: + x0, y0, x1, y1 = box + side = heat.shape[0] + flat = heat.ravel() + # Suppress immediate duplicates by greedily keeping separated cells. + chosen: list[int] = [] + for idx in np.argsort(flat)[::-1]: + r, c = divmod(int(idx), side) + if all(abs(r - divmod(j, side)[0]) + abs(c - divmod(j, side)[1]) >= 3 for j in chosen): + chosen.append(int(idx)) + if len(chosen) == count: + break + for rank, idx in enumerate(chosen, start=1): + r, c = divmod(idx, side) + cx = round(x0 + (c + 0.5) * (x1 - x0) / side) + cy = round(y0 + (r + 0.5) * (y1 - y0) / side) + rad = 11 if rank <= 3 else 8 + draw.ellipse((cx - rad, cy - rad, cx + rad, cy + rad), outline=PALETTE["amber"], width=3) + if rank <= 5: + draw.text((cx + 10, cy - 16), str(rank), fill=PALETTE["amber"], font=font(16, True)) + + +def text_block(draw: ImageDraw.ImageDraw, xy: tuple[int, int], lines: Iterable[str], fill: tuple[int, int, int], size: int = 20, gap: int = 8) -> None: + x, y = xy + f = font(size) + for line in lines: + draw.text((x, y), line, fill=fill, font=f) + y += size + gap + + +def render() -> None: + OUT_DIR.mkdir(parents=True, exist_ok=True) + summary = json.loads((SOURCE / "summary.json").read_text()) + arrays = np.load(SOURCE / "heatmaps.npz") + original = Image.open(SOURCE / "images" / "original.png").convert("RGB") + masked = Image.open(SOURCE / "images" / "answer-mask.png").convert("RGB") + + answer_delta = arrays["answer_delta"].astype(np.float32, copy=False) + random_delta = arrays["random_delta"].astype(np.float32, copy=False) + ratio = arrays["ratio"].astype(np.float32, copy=False) + side = token_side(summary, answer_delta.shape[1]) + answer_grid = fold_tokens(answer_delta, side) + random_grid = fold_tokens(random_delta, side) + ratio_grid = fold_tokens(ratio, side) + + answer_mean = answer_grid.mean(axis=0) + random_mean = random_grid.mean(axis=0) + ratio_mean = ratio_grid.mean(axis=0) + early_ratio = ratio_grid[:4].mean(axis=0) + mid_delta = answer_grid[6:13].mean(axis=0) + late_delta = answer_grid[-4:].mean(axis=0) + + np.savez_compressed( + OUT_DIR / "token_grid_source.npz", + answer_mean=answer_mean, + random_mean=random_mean, + ratio_mean=ratio_mean, + early_ratio=early_ratio, + mid_answer_delta=mid_delta, + late_answer_delta=late_delta, + image_grid_thw=np.array(summary["processor_meta"]["image_grid_thw"][0], dtype=np.int32), + ) + (OUT_DIR / "token_grid_summary.json").write_text( + json.dumps( + { + "source": str(SOURCE), + "image_grid_thw": summary["processor_meta"]["image_grid_thw"][0], + "image_tokens": int(summary["image_tokens"]), + "rendered_token_grid": [side, side], + "patch_merge": int(summary["processor_meta"]["image_grid_thw"][0][1] // side), + "answer_over_random_delta": float(summary["answer_over_random_delta"]), + "question": summary["question"]["q"], + "answer_text": summary["question"]["answer_text"], + }, + indent=2, + ) + ) + + W, H = 2200, 1500 + canvas = Image.new("RGB", (W, H), PALETTE["bg"]) + draw = ImageDraw.Draw(canvas) + for y in range(0, H, 18): + draw.line((0, y, W, y), fill=(7, 10 + (y % 11), 17 + (y % 13))) + glow = Image.new("RGBA", (W, H), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-320, -240, 960, 780), fill=(255, 80, 66, 34)) + gd.ellipse((920, -120, 2350, 1100), fill=(75, 218, 255, 30)) + gd.ellipse((760, 860, 1810, 1760), fill=(255, 194, 72, 18)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(95))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + draw.text((64, 48), "SNAPCOMPACT TOKEN FIELD", fill=PALETTE["amber"], font=font(24, True)) + draw.text((64, 86), "Where the hidden-state scar lands on the bitmap", fill=PALETTE["ink"], font=font(62, True)) + draw.text( + (66, 164), + "PaddleOCR-VL reports a 1×54×54 image patch grid; 729 hidden-state image tokens fold back to 27×27 spatial cells.", + fill=PALETTE["muted"], + font=font(24), + ) + + # Main spatial map. + main_panel = (545, 225, 1455, 1340) + draw_panel(draw, main_panel, "answer-mask delta projected onto image tokens", "mean ||hidden(original) − hidden(answer-mask)|| across 19 layers") + map_box = (610, 330, 1390, 1110) + projected = heat_overlay(original, answer_mean) + pasted = paste_fit(canvas, projected, map_box, Image.Resampling.LANCZOS) + # Grid + answer box sit over the pasted square. + overlay = Image.new("RGBA", canvas.size, (0, 0, 0, 0)) + od = ImageDraw.Draw(overlay) + draw_token_grid(od, pasted, side, (255, 255, 255)) + bbox = answer_bbox(summary["question"]["answer_start"], summary["question"]["answer_end"], summary["geometry"]["cols"], 8, 13) + sx = (pasted[2] - pasted[0]) / original.width + sy = (pasted[3] - pasted[1]) / original.height + answer_rect = ( + round(pasted[0] + bbox[0] * sx), + round(pasted[1] + bbox[1] * sy), + round(pasted[0] + bbox[2] * sx), + round(pasted[1] + bbox[3] * sy), + ) + od.rounded_rectangle(answer_rect, radius=8, outline=(*PALETTE["red"], 255), width=5) + draw_hotspots(od, pasted, answer_mean) + canvas = Image.alpha_composite(canvas.convert("RGBA"), overlay).convert("RGB") + draw = ImageDraw.Draw(canvas) + label(draw, (pasted[0] + 18, pasted[1] + 18), "27×27 reconstructed image-token grid", PALETTE["cyan"], 19) + label(draw, (answer_rect[2] + 14, answer_rect[1] - 5), "erased answer text", PALETTE["red"], 18) + text_block( + draw, + (620, 1162), + [ + f"Q: {summary['question']['q']}", + f"Gold answer: {summary['question']['answer_text']} answer-mask mean delta: {summary['answer_delta_mean']:.2f} random-mask mean delta: {summary['random_delta_mean']:.2f}", + f"Answer/random delta ratio: {summary['answer_over_random_delta']:.2f}×. Bright cells are token locations most perturbed by hiding the answer span.", + ], + PALETTE["muted"], + 20, + 7, + ) + + # Evidence crops. + left = (64, 225, 505, 1340) + draw_panel(draw, left, "bitmap intervention", "original crop vs. answer erased") + crop = crop_answer(original, summary["question"]["answer_start"], summary["question"]["answer_end"], summary["geometry"]["cols"], 8, 13) + mcrop = crop_answer(masked, summary["question"]["answer_start"], summary["question"]["answer_end"], summary["geometry"]["cols"], 8, 13) + draw.text((96, 332), "ORIGINAL", fill=PALETTE["cyan"], font=font(17, True)) + draw.rounded_rectangle((94, 360, 475, 525), radius=16, fill=(240, 238, 226), outline=PALETTE["cyan"], width=3) + paste_fit(canvas, crop, (108, 374, 461, 511), Image.Resampling.NEAREST) + draw.text((96, 572), "ANSWER MASK", fill=PALETTE["red"], font=font(17, True)) + draw.rounded_rectangle((94, 600, 475, 765), radius=16, fill=(240, 238, 226), outline=PALETTE["red"], width=3) + paste_fit(canvas, mcrop, (108, 614, 461, 751), Image.Resampling.NEAREST) + draw.text((96, 822), "source arrays", fill=PALETTE["muted"], font=font(17, True)) + text_block( + draw, + (96, 858), + [ + "heatmaps.npz:", + "answer_delta[19,729]", + "random_delta[19,729]", + "ratio[19,729]", + "", + "fold rule:", + "54×54 patches / 2×2 merge", + "→ 27×27 visual tokens", + ], + PALETTE["ink"], + 21, + 8, + ) + draw.rounded_rectangle((96, 1110, 472, 1268), radius=18, fill=PALETTE["panel2"], outline=(38, 51, 64), width=1) + draw.text((118, 1132), "scar strength", fill=PALETTE["amber"], font=font(18, True)) + draw.text((118, 1170), f"{summary['answer_over_random_delta']:.2f}×", fill=PALETTE["ink"], font=font(54, True)) + draw.text((120, 1232), "answer-mask / random-mask mean delta", fill=PALETTE["muted"], font=font(17)) + + # Right analytical small multiples. + draw_micro_grid(canvas, ratio_mean, (1495, 225, 2136, 590), "ratio field", "mean answer_delta / random_delta") + draw_micro_grid(canvas, early_ratio, (1495, 620, 1810, 975), "early layers", "ratio, layers 0–3") + draw_micro_grid(canvas, mid_delta, (1820, 620, 2136, 975), "middle layers", "answer delta, layers 6–12") + draw_micro_grid(canvas, late_delta, (1495, 1005, 1810, 1340), "late layers", "answer delta, last 4") + draw_micro_grid(canvas, random_mean, (1820, 1005, 2136, 1340), "random control", "random-mask delta") + + # Color legend. + lx0, ly0, lx1, ly1 = 1530, 530, 2100, 552 + for x in range(lx0, lx1): + draw.line((x, ly0, x, ly1), fill=heat_color((x - lx0) / (lx1 - lx0))) + draw.text((lx0, ly1 + 10), "low", fill=PALETTE["muted"], font=font(15)) + draw.text((lx1 - 34, ly1 + 10), "high", fill=PALETTE["muted"], font=font(15)) + + canvas.save(OUT) + + +if __name__ == "__main__": + render() diff --git a/packages/snapcompact/research/snapcompact_viz_volume.py b/packages/snapcompact/research/snapcompact_viz_volume.py new file mode 100644 index 000000000..7089b3608 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_viz_volume.py @@ -0,0 +1,235 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["matplotlib", "numpy"] +# /// +"""Render a volumetric tensor-cube visualization for snapcompact activations.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") + +import matplotlib.pyplot as plt +import numpy as np +from matplotlib import cm +from matplotlib.colors import LinearSegmentedColormap +from mpl_toolkits.mplot3d.art3d import Line3DCollection + +HERE = Path(__file__).resolve().parent +DATA_DIR = HERE / "results" / "tensor-heatmap-paddleocr-q7" +OUT_DIR = HERE / "results" / "agent-viz-volume" + +BG = "#05070b" +PANEL = "#0b1017" +GRID = "#33424c" +INK = "#efe9d5" +MUTED = "#87959b" +CYAN = "#44d9ff" +RED = "#ff5d4c" +AMBER = "#ffc84a" +GREEN = "#87ff80" + + +def robust01(values: np.ndarray, q: float = 0.985) -> np.ndarray: + """Quantile-normalize positive activation magnitudes without copying when possible.""" + scale = float(np.nanquantile(values, q)) + if not np.isfinite(scale) or scale <= 0.0: + scale = 1.0 + return np.clip(values / scale, 0.0, 1.0).astype(np.float32, copy=False) + + +def tinted_cmap(name: str, low: str, high: str) -> LinearSegmentedColormap: + return LinearSegmentedColormap.from_list(name, [(0.0, BG), (0.24, low), (1.0, high)], N=256) + + +def load_volume(data_dir: Path) -> tuple[np.ndarray, dict, list[str]]: + npz = np.load(data_dir / "heatmaps.npz") + summary = json.loads((data_dir / "summary.json").read_text()) + + answer = robust01(npz["answer_binned"]) + random = robust01(npz["random_binned"]) + ratio = robust01(npz["ratio_binned"]) + volume = np.stack([answer, random, ratio], axis=0) + labels = ["ANSWER Δ", "RANDOM Δ", "ANSWER/RANDOM"] + return volume, summary, labels + + +def cube_edges(x0: float, x1: float, y0: float, y1: float, z0: float, z1: float) -> list[list[tuple[float, float, float]]]: + p = { + "000": (x0, y0, z0), + "100": (x1, y0, z0), + "010": (x0, y1, z0), + "110": (x1, y1, z0), + "001": (x0, y0, z1), + "101": (x1, y0, z1), + "011": (x0, y1, z1), + "111": (x1, y1, z1), + } + return [ + [p["000"], p["100"]], [p["010"], p["110"]], [p["001"], p["101"]], [p["011"], p["111"]], + [p["000"], p["010"]], [p["100"], p["110"]], [p["001"], p["011"]], [p["101"], p["111"]], + [p["000"], p["001"]], [p["100"], p["101"]], [p["010"], p["011"]], [p["110"], p["111"]], + ] + + +def style_3d(ax) -> None: + ax.set_facecolor(BG) + for axis in (ax.xaxis, ax.yaxis, ax.zaxis): + axis.pane.set_facecolor((0.02, 0.03, 0.04, 0.0)) + axis._axinfo["grid"]["color"] = (0.35, 0.48, 0.56, 0.12) + axis._axinfo["tick"]["color"] = (0.75, 0.82, 0.82, 0.55) + ax.tick_params(colors=MUTED, labelsize=8, pad=0) + ax.set_xlabel("image-token bins (729 → 180)", color=MUTED, labelpad=9) + ax.set_ylabel("condition", color=MUTED, labelpad=7) + ax.set_zlabel("decoder layer", color=MUTED, labelpad=7) + ax.set_xlim(0, 179) + ax.set_ylim(-0.38, 2.38) + ax.set_zlim(0, 18) + ax.set_yticks([0, 1, 2]) + ax.set_yticklabels(["answer", "random", "ratio"], color=INK) + ax.set_xticks([0, 45, 90, 135, 179]) + ax.set_zticks([0, 4, 9, 14, 18]) + ax.view_init(elev=23, azim=-61) + ax.set_box_aspect((3.9, 1.0, 1.15)) + + +def add_volume(ax, volume: np.ndarray) -> None: + cmaps = [tinted_cmap("answer_ct", "#063842", CYAN), tinted_cmap("random_ct", "#461813", RED), tinted_cmap("ratio_ct", "#3c2b05", AMBER)] + edge_colors = [CYAN, RED, AMBER] + layers = np.arange(volume.shape[1]) + bins = np.arange(volume.shape[2]) + x, z = np.meshgrid(bins, layers) + + # Translucent CT slices: condition is depth, layer is vertical, image-token bin is horizontal. + for cond, cmap in enumerate(cmaps): + vals = volume[cond] + rgba = cmap(vals) + rgba[..., 3] = 0.08 + 0.68 * np.power(vals, 1.55) + y = np.full_like(x, cond, dtype=np.float32) + ax.plot_surface(x, y, z, facecolors=rgba, rstride=1, cstride=1, linewidth=0, antialiased=False, shade=False) + + # Bright activation voxels above each condition's 98th percentile. + threshold = float(np.quantile(vals, 0.982)) + zz, xx = np.where(vals >= threshold) + yy = np.full(xx.shape, cond, dtype=np.float32) + strength = vals[zz, xx] + ax.scatter(xx, yy, zz, s=10 + 90 * strength, c=edge_colors[cond], marker="s", alpha=0.58, depthshade=False, linewidths=0) + + ax.add_collection3d(Line3DCollection(cube_edges(0, 179, -0.23, 2.23, 0, 18), colors=(0.42, 0.72, 0.82, 0.22), linewidths=0.9)) + + # Crosshair slices through the strongest answer/random separation. + ratio = volume[2] + layer_profile = ratio.mean(axis=1) + bin_profile = ratio.mean(axis=0) + peak_layer = int(layer_profile.argmax()) + peak_bin = int(bin_profile.argmax()) + ax.plot([peak_bin, peak_bin], [-0.28, 2.28], [peak_layer, peak_layer], color=GREEN, alpha=0.9, linewidth=1.5) + ax.plot([0, 179], [2.28, 2.28], [peak_layer, peak_layer], color=GREEN, alpha=0.45, linewidth=1.1) + ax.text(peak_bin + 3, 2.35, peak_layer + 0.2, "hottest ratio slice", color=GREEN, fontsize=8) + + +def add_projection_panel(ax, volume: np.ndarray, labels: list[str]) -> None: + ax.set_facecolor(PANEL) + cmap = tinted_cmap("small_ct", "#10252f", "#f2d87b") + strip = np.vstack([volume[0], np.full((2, volume.shape[2]), np.nan), volume[1], np.full((2, volume.shape[2]), np.nan), volume[2]]) + masked = np.ma.masked_invalid(strip) + cmap.set_bad(PANEL) + ax.imshow(masked, aspect="auto", interpolation="nearest", cmap=cmap, vmin=0, vmax=1) + ax.set_xticks([0, 45, 90, 135, 179]) + ax.set_yticks([9, 30, 51]) + ax.set_yticklabels(labels, color=INK, fontsize=8) + ax.tick_params(colors=MUTED, labelsize=8, length=0) + ax.set_title("unwrapped tensor volume", color=INK, fontsize=12, loc="left", pad=8) + for spine in ax.spines.values(): + spine.set_color("#27323a") + + +def add_layer_panel(ax, volume: np.ndarray) -> None: + ax.set_facecolor(PANEL) + colors = [CYAN, RED, AMBER] + names = ["answer", "random", "ratio"] + for cond, color in enumerate(colors): + profile = volume[cond].mean(axis=1) + ax.plot(np.arange(profile.size), profile, color=color, linewidth=2.0, label=names[cond]) + ax.fill_between(np.arange(profile.size), profile, 0, color=color, alpha=0.08) + ax.set_xlim(0, 18) + ax.set_ylim(0, 1.0) + ax.set_xlabel("layer", color=MUTED, fontsize=8) + ax.set_ylabel("mean normalized intensity", color=MUTED, fontsize=8) + ax.tick_params(colors=MUTED, labelsize=8) + ax.grid(color=GRID, alpha=0.18, linewidth=0.7) + ax.legend(frameon=False, labelcolor=INK, fontsize=8, loc="upper right") + ax.set_title("layer dose curve", color=INK, fontsize=12, loc="left", pad=8) + for spine in ax.spines.values(): + spine.set_color("#27323a") + + +def render(volume: np.ndarray, summary: dict, labels: list[str], out_path: Path) -> None: + fig = plt.figure(figsize=(18, 11), dpi=180, facecolor=BG) + gs = fig.add_gridspec(3, 5, width_ratios=[1.35, 1.35, 1.35, 0.95, 0.95], height_ratios=[0.12, 1.0, 0.42], wspace=0.22, hspace=0.24) + + title_ax = fig.add_subplot(gs[0, :]) + title_ax.axis("off") + title_ax.text(0.0, 0.70, "SNAPCOMPACT ACTIVATION CT", color=INK, fontsize=27, fontweight="bold", transform=title_ax.transAxes) + title_ax.text(0.0, 0.24, "volumetric tensor cube: 19 layers × 180 image-token bins × 3 conditions", color=MUTED, fontsize=11, transform=title_ax.transAxes) + title_ax.text(0.985, 0.58, f"PaddleOCR-VL · Q: {summary['question']['q']}", color=MUTED, fontsize=9, ha="right", transform=title_ax.transAxes) + title_ax.text(0.985, 0.24, f"gold answer {summary['question']['answer_text']} · answer/random mean Δ {summary['answer_over_random_delta']:.2f}×", color=AMBER, fontsize=10, ha="right", transform=title_ax.transAxes) + + ax3d = fig.add_subplot(gs[1:, :3], projection="3d") + style_3d(ax3d) + add_volume(ax3d, volume) + ax3d.set_title("MRI-style scan of hidden-state deltas", color=INK, fontsize=15, loc="left", pad=12) + + ax_proj = fig.add_subplot(gs[1, 3:]) + add_projection_panel(ax_proj, volume, labels) + + ax_layer = fig.add_subplot(gs[2, 3:]) + add_layer_panel(ax_layer, volume) + + fig.text(0.055, 0.055, "source: heatmaps.npz arrays answer_binned, random_binned, ratio_binned · quantile normalized per condition", color="#617078", fontsize=8) + fig.text(0.055, 0.033, "cyan=answer evidence · red=random control · gold=answer/random amplification · green=crosshair at peak ratio slice", color="#617078", fontsize=8) + + out_path.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(out_path, facecolor=BG, bbox_inches="tight", pad_inches=0.22) + plt.close(fig) + + +def write_source_data(volume: np.ndarray, summary: dict, out_dir: Path) -> None: + out_dir.mkdir(parents=True, exist_ok=True) + np.savez_compressed( + out_dir / "volume_source.npz", + volume=volume, + answer_norm=volume[0], + random_norm=volume[1], + ratio_norm=volume[2], + ) + payload = { + "description": "Quantile-normalized tensor cube used by snapcompact_viz_volume.py.", + "shape": {"condition": 3, "layers": int(volume.shape[1]), "image_token_bins": int(volume.shape[2])}, + "conditions": ["answer_delta", "random_delta", "answer_over_random_ratio"], + "question": summary["question"], + "answer_over_random_delta": summary["answer_over_random_delta"], + "source": str(DATA_DIR / "heatmaps.npz"), + } + (out_dir / "volume_source.json").write_text(json.dumps(payload, indent=2) + "\n") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--data-dir", type=Path, default=DATA_DIR) + parser.add_argument("--out-dir", type=Path, default=OUT_DIR) + parser.add_argument("--out", type=str, default="volume.png") + args = parser.parse_args() + + volume, summary, labels = load_volume(args.data_dir) + write_source_data(volume, summary, args.out_dir) + render(volume, summary, labels, args.out_dir / args.out) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_viz_waterfall.py b/packages/snapcompact/research/snapcompact_viz_waterfall.py new file mode 100755 index 000000000..de4bf49f4 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_viz_waterfall.py @@ -0,0 +1,301 @@ +#!/usr/bin/env python3 +"""Layered snapcompact activation waterfall. + +Builds a seismic/ridgeline rendering from the PaddleOCR-VL white-box +activation deltas in results/tensor-heatmap-paddleocr-q7/heatmaps.npz. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import matplotlib.patheffects as pe +import matplotlib.pyplot as plt +import numpy as np +from matplotlib.colors import LinearSegmentedColormap +from matplotlib.patches import Rectangle + + +SCRIPT_DIR = Path(__file__).resolve().parent +DATA_DIR = SCRIPT_DIR / "results" / "tensor-heatmap-paddleocr-q7" +OUT_DIR = SCRIPT_DIR / "results" / "agent-viz-waterfall" + + +def smooth_rows(values: np.ndarray, radius: int = 3) -> np.ndarray: + """Small separable bin smoother; preserves shape and avoids scipy.""" + if radius <= 0: + return values.copy() + x = np.arange(-radius, radius + 1, dtype=np.float32) + kernel = np.exp(-(x * x) / (2.0 * (radius / 1.8) ** 2)) + kernel /= kernel.sum() + padded = np.pad(values, ((0, 0), (radius, radius)), mode="edge") + out = np.empty_like(values, dtype=np.float32) + for row in range(values.shape[0]): + out[row] = np.convolve(padded[row], kernel, mode="valid") + return out + + +def robust_unit(values: np.ndarray, high: float) -> np.ndarray: + scaled = np.log1p(np.maximum(values, 0.0)) / np.log1p(high) + return np.clip(scaled, 0.0, 1.0).astype(np.float32) + + +def load_source() -> tuple[dict, dict[str, np.ndarray]]: + with (DATA_DIR / "summary.json").open("r", encoding="utf-8") as handle: + summary = json.load(handle) + with np.load(DATA_DIR / "heatmaps.npz") as npz: + arrays = {name: npz[name].astype(np.float32) for name in npz.files} + return summary, arrays + + +def build_waterfall_data(summary: dict, arrays: dict[str, np.ndarray]) -> dict[str, np.ndarray]: + answer = smooth_rows(arrays["answer_binned"], radius=3) + random = smooth_rows(arrays["random_binned"], radius=3) + ratio = smooth_rows(arrays["ratio_binned"], radius=2) + + # Use one common robust scale so answer/random amplitudes are visually comparable. + common_high = float(summary.get("common_delta_scale_p98") or np.percentile(np.r_[answer, random], 98)) + answer_u = robust_unit(answer, common_high) + random_u = robust_unit(random, common_high) + contrast = np.tanh((answer_u - random_u) * 2.8).astype(np.float32) + ratio_u = robust_unit(ratio, float(summary.get("ratio_scale_p98") or np.percentile(ratio, 98))) + + layers, bins = answer.shape + x = np.linspace(0.0, 1.0, bins, dtype=np.float32) + baselines = np.arange(layers, dtype=np.float32)[::-1] + return { + "x": x, + "baselines": baselines, + "answer": answer, + "random": random, + "ratio": ratio, + "answer_unit": answer_u, + "random_unit": random_u, + "contrast": contrast, + "ratio_unit": ratio_u, + } + + +def draw_glow_line(ax, x, y, color, lw=1.4, z=5, alpha=1.0): + line, = ax.plot(x, y, color=color, lw=lw, alpha=alpha, zorder=z, solid_joinstyle="round") + line.set_path_effects( + [ + pe.Stroke(linewidth=lw + 8.5, foreground=color, alpha=0.055), + pe.Stroke(linewidth=lw + 4.5, foreground=color, alpha=0.12), + pe.Normal(), + ] + ) + return line + + +def render(summary: dict, data: dict[str, np.ndarray]) -> None: + OUT_DIR.mkdir(parents=True, exist_ok=True) + out_png = OUT_DIR / "waterfall.png" + + x = data["x"] + baselines = data["baselines"] + answer_u = data["answer_unit"] + random_u = data["random_unit"] + contrast = data["contrast"] + ratio_u = data["ratio_unit"] + layers, bins = answer_u.shape + + fig = plt.figure(figsize=(15.5, 10.5), dpi=210, facecolor="#05070d") + ax = fig.add_axes([0.055, 0.09, 0.89, 0.80], facecolor="#05070d") + + # Background ratio field: a dim spectrogram behind the ridges. + cmap = LinearSegmentedColormap.from_list( + "snap_seismic", + ["#05070d", "#0b1831", "#133f59", "#2d6f75", "#f3a44c", "#fff0bd"], + ) + extent = (0.0, 1.0, -0.78, layers - 0.22) + ax.imshow( + ratio_u[::-1], + extent=extent, + aspect="auto", + cmap=cmap, + interpolation="bicubic", + alpha=0.38, + zorder=0, + ) + + # Seismic paper grid and scanlines. + for gx in np.linspace(0, 1, 13): + ax.axvline(gx, color="#7cc8ff", lw=0.45, alpha=0.10, zorder=1) + for y in range(layers): + ax.axhline(y, color="#d4eaff", lw=0.38, alpha=0.085, zorder=1) + for yy in np.linspace(-0.6, layers - 0.35, 78): + ax.axhline(yy, color="#ffffff", lw=0.2, alpha=0.018, zorder=1) + + answer_color = "#ffc05a" + random_color = "#33d7ff" + gain_color = "#ff4d8d" + amplitude = 0.72 + + for layer_index, base in enumerate(baselines): + ans = base + answer_u[layer_index] * amplitude + rnd = base - random_u[layer_index] * amplitude * 0.82 + mid = base + contrast[layer_index] * amplitude * 0.62 + + ax.fill_between(x, base, ans, color=answer_color, alpha=0.075, zorder=2) + ax.fill_between(x, base, rnd, color=random_color, alpha=0.045, zorder=2) + ax.fill_between( + x, + rnd, + ans, + where=answer_u[layer_index] >= random_u[layer_index], + interpolate=True, + color=gain_color, + alpha=0.055, + zorder=2, + ) + + # Double-trace each layer: random-mask lower trace, answer-mask upper trace, + # plus a magenta differential tremor to make the comparison readable. + draw_glow_line(ax, x, rnd, random_color, lw=0.9, z=4, alpha=0.74) + draw_glow_line(ax, x, ans, answer_color, lw=1.16, z=5, alpha=0.92) + ax.plot(x, mid, color=gain_color, lw=0.48, alpha=0.55, zorder=3) + + if layer_index in {0, 4, 9, 14, layers - 1}: + ax.text( + -0.018, + base, + f"L{layer_index:02d}", + ha="right", + va="center", + color="#b9dfff", + fontsize=9, + family="monospace", + alpha=0.82, + ) + + # Highlight the strongest answer-vs-random bin per layer with tiny hot pips. + gain = answer_u - random_u + strongest = np.argmax(gain, axis=1) + ax.scatter( + x[strongest], + baselines + answer_u[np.arange(layers), strongest] * amplitude + 0.045, + s=8 + 32 * np.clip(gain[np.arange(layers), strongest], 0, 1), + c="#fff6cf", + alpha=0.72, + edgecolors="none", + zorder=6, + ) + + # Framing labels. + answer_mean = float(summary["answer_delta_mean"]) + random_mean = float(summary["random_delta_mean"]) + ratio = float(summary["answer_over_random_delta"]) + question = summary["question"]["q"] + answer_text = summary["question"]["answer_text"] + image_tokens = int(summary["image_tokens"]) + + ax.text( + 0.0, + layers + 0.62, + "SNAPCOMPACT ACTIVATION WATERFALL", + color="#f7fbff", + fontsize=22, + weight="bold", + family="monospace", + ha="left", + va="bottom", + ) + ax.text( + 0.0, + layers + 0.20, + f"PaddleOCR-VL · {layers} decoder layers · {image_tokens} image tokens binned into {bins} traces · answer '{answer_text}' vs random mask", + color="#9cc7e5", + fontsize=10.5, + family="monospace", + ha="left", + va="bottom", + ) + ax.text( + 0.0, + layers - 0.19, + f"Q: {question}", + color="#d8ecff", + fontsize=9.5, + family="monospace", + ha="left", + va="top", + alpha=0.84, + ) + ax.text( + 1.0, + layers + 0.27, + f"Δmean {answer_mean:.2f} / {random_mean:.2f} = {ratio:.2f}×", + color="#ffd37a", + fontsize=13, + family="monospace", + weight="bold", + ha="right", + va="bottom", + ) + + # Legend built as luminous calibration bars. + legend_y = -1.35 + ax.plot([0.02, 0.10], [legend_y, legend_y], color=answer_color, lw=2.2) + ax.text(0.112, legend_y, "answer-mask ridge", color="#ffdca0", fontsize=9, va="center", family="monospace") + ax.plot([0.32, 0.40], [legend_y, legend_y], color=random_color, lw=2.2) + ax.text(0.412, legend_y, "random-mask ridge", color="#9ff0ff", fontsize=9, va="center", family="monospace") + ax.plot([0.62, 0.70], [legend_y, legend_y], color=gain_color, lw=1.4) + ax.text(0.712, legend_y, "answer excess tremor", color="#ff9bbb", fontsize=9, va="center", family="monospace") + + # Outer phosphor frame. + ax.add_patch(Rectangle((0, -0.78), 1, layers - 0.44, fill=False, lw=0.9, edgecolor="#5fb7ff", alpha=0.34, zorder=10)) + ax.set_xlim(-0.055, 1.02) + ax.set_ylim(-1.62, layers + 0.98) + ax.set_xticks(np.linspace(0, 1, 7)) + ax.set_xticklabels([f"{int(t * image_tokens):03d}" for t in np.linspace(0, 1, 7)], color="#8fbede", fontsize=8, family="monospace") + ax.set_yticks([]) + ax.set_xlabel("image-token bin →", color="#9cc7e5", fontsize=10, family="monospace", labelpad=12) + for spine in ax.spines.values(): + spine.set_visible(False) + ax.tick_params(axis="x", length=0) + + # Save source arrays used by this rendering for reproducibility. + np.savez_compressed( + OUT_DIR / "waterfall_source.npz", + x=x, + baselines=baselines, + answer_binned=data["answer"], + random_binned=data["random"], + ratio_binned=data["ratio"], + answer_unit=answer_u, + random_unit=random_u, + ratio_unit=ratio_u, + contrast=contrast, + ) + with (OUT_DIR / "waterfall_source.json").open("w", encoding="utf-8") as handle: + json.dump( + { + "question": question, + "answer_text": answer_text, + "layers": layers, + "bins": bins, + "image_tokens": image_tokens, + "answer_delta_mean": answer_mean, + "random_delta_mean": random_mean, + "answer_over_random_delta": ratio, + "source_npz": str(DATA_DIR / "heatmaps.npz"), + }, + handle, + indent=2, + ) + + fig.savefig(out_png, facecolor=fig.get_facecolor(), bbox_inches="tight", pad_inches=0.14) + plt.close(fig) + print(out_png) + + +def main() -> None: + summary, arrays = load_source() + data = build_waterfall_data(summary, arrays) + render(summary, data) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/squad.py b/packages/snapcompact/research/squad.py new file mode 100644 index 000000000..39d206f4d --- /dev/null +++ b/packages/snapcompact/research/squad.py @@ -0,0 +1,120 @@ +"""SQuAD v1.1 dev: passage flow, question sampling, official EM/F1 scoring.""" + +import json +import random +import re +import string +import urllib.request +from collections import Counter +from pathlib import Path + +SQUAD_URL = "https://rajpurkar.github.io/SQuAD-explorer/dataset/dev-v1.1.json" + + +def load_paragraphs(cache: Path) -> list[dict]: + """Flattened [{ctx, qas, title}] in deterministic dataset order.""" + path = cache / "squad-dev-v1.1.json" + if not path.exists(): + urllib.request.urlretrieve(SQUAD_URL, path) + data = json.loads(path.read_text())["data"] + out = [] + for art in data: + for p in art["paragraphs"]: + out.append({"ctx": " ".join(p["context"].split()), "qas": p["qas"], "title": art["title"]}) + return out + + +def build_flow(paras: list[dict], max_chars: int | None = None) -> tuple[str, list[int]]: + """Space-joined passage stream + start offset of each passage.""" + flow, offsets = "", [] + for p in paras: + offsets.append(len(flow)) + flow += p["ctx"] + " " + if max_chars is not None and len(flow) >= max_chars: + break + return flow, offsets + + +def sample_chunk_questions( + paras: list[dict], offsets: list[int], start: int, end: int, n: int, seed: int +) -> list[dict]: + """Up to n questions from passages fully inside [start, end), evenly spread. + + Passages straddling a chunk boundary are skipped (their answers may be cut). + Each question records pos_rel: passage start relative to the chunk, 0..1. + """ + rng = random.Random(seed * 1_000_003 + start) + eligible = [ + i + for i in range(len(offsets)) + if offsets[i] >= start and offsets[i] + len(paras[i]["ctx"]) <= end + ] + if not eligible: + return [] + n = min(n, len(eligible)) + step = len(eligible) / n + picked = [] + for k in range(n): + pi = eligible[int(k * step)] + qa = rng.choice(paras[pi]["qas"]) + picked.append( + { + "q": " ".join(qa["question"].split()), + "golds": sorted({a["text"] for a in qa["answers"]}), + "pos_rel": (offsets[pi] - start) / (end - start), + } + ) + return picked + + +# --- official SQuAD normalization / metrics --- + + +def _normalize(s: str) -> str: + s = s.lower() + s = "".join(ch for ch in s if ch not in string.punctuation) + s = re.sub(r"\b(a|an|the)\b", " ", s) + return " ".join(s.split()) + + +def exact_match(pred: str, golds: list[str]) -> float: + return float(any(_normalize(pred) == _normalize(g) for g in golds)) + + +def f1(pred: str, golds: list[str]) -> float: + best = 0.0 + for g in golds: + p_tok, g_tok = _normalize(pred).split(), _normalize(g).split() + common = Counter(p_tok) & Counter(g_tok) + overlap = sum(common.values()) + if overlap == 0: + continue + prec, rec = overlap / len(p_tok), overlap / len(g_tok) + best = max(best, 2 * prec * rec / (prec + rec)) + return best + + +def parse_numbered(text: str, n: int) -> list[str]: + """Extract answers from a numbered list; missing entries become ''. """ + answers = [""] * n + for line in text.splitlines(): + m = re.match(r"\s*(\d+)[.):]\s*(.*\S)?\s*$", line) + if m and m.group(2): + idx = int(m.group(1)) - 1 + if 0 <= idx < n and not answers[idx]: + answers[idx] = m.group(2).strip() + return answers + + +def score(answers: list[str], questions: list[dict]) -> dict: + ems = [exact_match(a, q["golds"]) for a, q in zip(answers, questions)] + f1s = [f1(a, q["golds"]) for a, q in zip(answers, questions)] + return { + "em": sum(ems) / len(ems), + "f1": sum(f1s) / len(f1s), + "abstained": sum("unreadable" in a.lower() for a in answers), + "per_question": [ + {"answer": a, "golds": q["golds"], "em": e, "f1": f} + for a, q, e, f in zip(answers, questions, ems, f1s) + ], + } diff --git a/packages/snapcompact/src/index.ts b/packages/snapcompact/src/index.ts new file mode 100644 index 000000000..efd6a6990 --- /dev/null +++ b/packages/snapcompact/src/index.ts @@ -0,0 +1 @@ +export * from "./snapcompact"; diff --git a/packages/snapcompact/src/prompts/file-operations.md b/packages/snapcompact/src/prompts/file-operations.md new file mode 100644 index 000000000..053b76839 --- /dev/null +++ b/packages/snapcompact/src/prompts/file-operations.md @@ -0,0 +1,10 @@ +{{#if readFiles.length}} +{{#xml "read-files"}} +{{join readFiles "\n"}} +{{/xml}} +{{/if}} +{{#if modifiedFiles.length}} +{{#xml "modified-files"}} +{{join modifiedFiles "\n"}} +{{/xml}} +{{/if}} diff --git a/packages/snapcompact/src/prompts/snapcompact-summary.md b/packages/snapcompact/src/prompts/snapcompact-summary.md new file mode 100644 index 000000000..0d8d22323 --- /dev/null +++ b/packages/snapcompact/src/prompts/snapcompact-summary.md @@ -0,0 +1,17 @@ +Prior conversation history has been archived verbatim onto {{frameCount}} snapcompact frame{{#if multipleFrames}}s{{/if}} — the bitmap image{{#if multipleFrames}}s{{/if}} attached below{{#if multipleFrames}}, ordered oldest to newest{{/if}}. + +Reading a frame: monospace {{fontCell}} pixel font on a white background, {{cols}} characters per row, {{rows}} text rows per frame; read left to right, top to bottom. Text flows continuously with no word wrap, so words may break across row ends. Whitespace runs (including newlines) were collapsed to single spaces. {{#if sentenceInk}}Ink color cycles through six colors, advancing at sentence boundaries — a color change marks a new sentence.{{else}}Glyphs are plain black ink.{{/if}}{{#if lineRepeated}} Every text line is printed twice in a row — first on the white background, then repeated on a pale yellow band. The copies are identical: read each line once and use the duplicate only to double-check hard glyphs.{{/if}} Roles are tagged inline as [User]:, [Assistant]:, [Assistant thinking]:, [Assistant tool calls]:, and [Tool result]:. +{{#if mixedShapes}} + +Older frames may use a different font, grid, or ink coloring than described above; the reading order is always the same (left to right, top to bottom, oldest frame first). +{{/if}} +{{#if includedPreviousSummary}} + +The earliest frame begins with "[Summary of earlier history]" — a condensed digest of context that predates the archived conversation. +{{/if}} +{{#if truncatedChars}} + +{{truncatedChars}} characters of older history were dropped to respect the frame budget. The first frame (session start) is always kept, so the missing span sits between the first frame and the next. +{{/if}} + +Total archived: {{totalChars}} characters. Consult the frames whenever you need exact earlier details (user wording, decisions, file paths, tool output). If a region is hard to read, re-derive the fact from the workspace (re-read files, re-run commands) rather than guessing. diff --git a/packages/snapcompact/src/snapcompact.ts b/packages/snapcompact/src/snapcompact.ts new file mode 100644 index 000000000..8e688bc8e --- /dev/null +++ b/packages/snapcompact/src/snapcompact.ts @@ -0,0 +1,638 @@ +/** + * Snapcompact compaction: archive conversation history as dense bitmap images. + * + * Instead of asking an LLM to summarize discarded history, the serialized + * conversation is rendered into square PNG frames of pixel-font text that + * vision models read back directly, like an archivist at a snapcompact frame + * reader. + * + * The frame shape is provider-aware, following the snapcompact SQuAD evals + * (`packages/snapcompact`, 200k-token monolithic runs): + * + * - **Anthropic** (`8x8r-bw`): unscii-8 square cells, black ink, every line + * printed twice with the copy on a pale highlight band. Read at F1 parity + * with raw text at ~2x lower cost; the colored variants drew refusals at + * scale, the repeated plain shape did not. + * - **Google** (`8x8r-sent`): same repeated grid with six-hue sentence + * coloring (0.90 F1 at ~2.9x lower cost on gemini-3.5-flash). + * - **OpenAI** (`6x6u-sent`): OpenAI bills a flat ~2.9k tokens per image, so + * image count is the only cost lever — unscii-8 Lanczos-stretched to 6x6 + * cells packs the most readable chars per frame. Frames request + * `detail: "original"`; the default `auto` downscale destroys 6px glyphs. + * - **Unknown providers** default to the Anthropic shape (most + * refusal-robust). Gateways that resize images (e.g. OpenRouter normalizes + * visual payloads to a fixed token budget) defeat any shape — optical + * context fails silently there. + * + * The whole pass is local and deterministic — no LLM call, no API key, no + * latency beyond rendering. Rasterization and PNG encoding happen in native + * code (`renderSnapcompactPng` in `crates/pi-natives/src/snapcompact.rs`). + * Frames persist in the compaction entry's `preserveData` and are + * re-attached to the compaction summary message on every context rebuild. + */ + +import type { Api, ImageContent, Message, Model } from "@oh-my-pi/pi-ai"; +import { renderSnapcompactPng } from "@oh-my-pi/pi-natives"; +import { prompt } from "@oh-my-pi/pi-utils"; +import fileOperationsTemplate from "./prompts/file-operations.md" with { type: "text" }; +import snapcompactSummaryPrompt from "./prompts/snapcompact-summary.md" with { type: "text" }; + +// ============================================================================ +// Shapes +// ============================================================================ + +/** One eval-validated frame shape: font, cell, ink, repetition, and size. */ +export interface SnapcompactShape { + /** Bundled font in the native renderer. */ + font: "5x8" | "8x8"; + /** Target cell advance in pixels; differing from the font's natural cell + * renders via Lanczos stretch (anti-aliased RGB frame). */ + cellWidth: number; + /** Target cell pitch in pixels. */ + cellHeight: number; + /** Ink: `sent` cycles six hues at sentence boundaries; `bw` is black. */ + variant: "sent" | "bw"; + /** Each text line is printed this many times; copies after the first sit + * on a pale highlight band (redundancy coding). */ + lineRepeat: number; + /** Frame edge in pixels. */ + frameSize: number; + /** Per-frame billed-token estimate for the shape's target provider. */ + frameTokenEstimate: number; + /** Resolution hint attached to frame images (OpenAI-only). */ + imageDetail?: ImageContent["detail"]; +} + +/** Eval-validated shapes, keyed by the provider family they won on. */ +export const SNAPCOMPACT_SHAPES = { + /** `8x8r-bw`: unscii square, black ink, lines doubled on highlight bands. */ + anthropic: { + font: "8x8", + cellWidth: 8, + cellHeight: 8, + variant: "bw", + lineRepeat: 2, + frameSize: 1568, + frameTokenEstimate: 3300, + }, + /** `8x8r-sent`: the repeated grid with sentence-hue ink. */ + google: { + font: "8x8", + cellWidth: 8, + cellHeight: 8, + variant: "sent", + lineRepeat: 2, + frameSize: 1568, + frameTokenEstimate: 1100, + }, + /** `6x6u-sent`: unscii stretched to 6x6 — densest readable cell, fewest + * frames (OpenAI bills per image, ~2.9k tokens flat). */ + openaiDense: { + font: "8x8", + cellWidth: 6, + cellHeight: 6, + variant: "sent", + lineRepeat: 1, + frameSize: 1568, + frameTokenEstimate: 2900, + imageDetail: "original", + }, + /** Original 5x8 X.org shape (pre-shape-table sessions rendered this). */ + legacy: { + font: "5x8", + cellWidth: 5, + cellHeight: 8, + variant: "sent", + lineRepeat: 1, + frameSize: 2576, + frameTokenEstimate: 3300, + }, +} as const satisfies Record; + +/** Pick the eval-optimal frame shape for a provider API. */ +export function resolveSnapcompactShape(api?: Api): SnapcompactShape { + switch (api) { + case "openai-completions": + case "openai-responses": + case "openai-codex-responses": + case "azure-openai-responses": + return SNAPCOMPACT_SHAPES.openaiDense; + case "google-generative-ai": + case "google-gemini-cli": + case "google-vertex": + return SNAPCOMPACT_SHAPES.google; + default: + // anthropic-messages, bedrock-converse-stream, and anything unknown: + // the plain repeated grid is the most refusal-robust reader shape. + return SNAPCOMPACT_SHAPES.anthropic; + } +} + +// ============================================================================ +// Constants +// ============================================================================ + +/** Legacy frame edge in pixels (the 5x8 shape's eval-validated size). New + * shapes carry their own `frameSize`. */ +export const SNAPCOMPACT_FRAME_SIZE = 2576; + +/** Maximum frames carried on a compaction entry. Oldest frames are dropped + * first once the budget is exceeded (mirrors how iterative text summaries + * fade the oldest detail). */ +export const SNAPCOMPACT_MAX_FRAMES = 8; + +/** Conservative per-frame token estimate used for context budgeting + * (upper bound across shapes: Anthropic bills 1568*1568/750 ≈ 3,278). */ +export const SNAPCOMPACT_FRAME_TOKEN_ESTIMATE = 3300; + +/** Key under `CompactionEntry.preserveData` holding the frame archive. */ +export const SNAPCOMPACT_PRESERVE_KEY = "snapcompact"; + +// ============================================================================ +// Types +// ============================================================================ + +/** One developed snapcompact frame: a base64 PNG plus its reading geometry. */ +export interface SnapcompactFrame { + /** Base64-encoded PNG. */ + data: string; + mimeType: string; + /** Characters per row in the frame grid. */ + cols: number; + /** Text rows in the frame grid (unique lines, not repeated copies). */ + rows: number; + /** Characters actually printed onto this frame. */ + chars: number; + /** Shape metadata (absent on legacy frames, which are 5x8 `sent`). */ + font?: SnapcompactShape["font"]; + variant?: SnapcompactShape["variant"]; + lineRepeat?: number; + /** Resolution hint forwarded to the provider when re-attaching. */ + detail?: ImageContent["detail"]; +} + +/** Frame archive persisted under `preserveData[SNAPCOMPACT_PRESERVE_KEY]`. */ +export interface SnapcompactArchive { + /** Frames ordered oldest to newest. */ + frames: SnapcompactFrame[]; + /** Characters currently readable across all frames. */ + totalChars: number; + /** Characters dropped so far to respect the frame budget. */ + truncatedChars: number; +} + +export interface SnapcompactGeometry { + cols: number; + rows: number; + /** Characters that fit one frame (cols * rows). */ + capacity: number; +} + +export interface SnapcompactOptions { + /** App-level message transformer (same contract as agent-core's `SummaryOptions.convertToLlm`). */ + convertToLlm?: SnapcompactConvertToLlm; + /** Model whose provider API selects the frame shape. */ + model?: Pick; + /** Explicit shape override; wins over `model`. */ + shape?: SnapcompactShape; + /** Frame edge in pixels. Defaults to the shape's `frameSize`. */ + frameSize?: number; + /** Frame budget. Defaults to {@link SNAPCOMPACT_MAX_FRAMES}. */ + maxFrames?: number; +} + +/** Result of rendering one frame, before base64 packing. */ +export interface RenderedFrame { + png: Uint8Array; + cols: number; + rows: number; + /** Characters printed (input may be shorter than capacity). */ + chars: number; +} + +// ============================================================================ +// Compaction data contracts +// ============================================================================ + +export interface SnapcompactFileOperations { + read: Set; + written: Set; + edited: Set; +} + +export interface SnapcompactCompactionDetails { + readFiles: string[]; + modifiedFiles: string[]; +} + +export interface SnapcompactCompactionPreparation { + /** UUID of first entry to keep. */ + firstKeptEntryId: string; + /** Messages that will be archived and discarded. */ + messagesToSummarize: TMessage[]; + /** Messages that will be archived as the split-turn prefix, if any. */ + turnPrefixMessages: TMessage[]; + tokensBefore: number; + /** Summary from previous compaction, for continuity when no prior snapcompact archive exists. */ + previousSummary?: string; + /** Preserved opaque compaction payload from the previous compaction, if any. */ + previousPreserveData?: Record; + /** File operations extracted by the host agent. */ + fileOps: SnapcompactFileOperations; +} + +export interface SnapcompactCompactionResult { + summary: string; + shortSummary?: string; + firstKeptEntryId: string; + tokensBefore: number; + details?: T; + preserveData?: Record; +} + +export type SnapcompactConvertToLlm = (messages: TMessage[]) => Message[]; + +function defaultConvertToLlm(messages: TMessage[]): Message[] { + return messages as unknown as Message[]; +} + +// ============================================================================ +// File operation helpers +// ============================================================================ + +export function createSnapcompactFileOps(): SnapcompactFileOperations { + return { + read: new Set(), + written: new Set(), + edited: new Set(), + }; +} + +export function computeSnapcompactFileLists(fileOps: SnapcompactFileOperations): SnapcompactCompactionDetails { + const modified = new Set([...fileOps.edited, ...fileOps.written]); + const readFiles = [...fileOps.read].filter(file => !modified.has(file)).sort(); + const modifiedFiles = [...modified].sort(); + return { readFiles, modifiedFiles }; +} + +const FILE_OPERATION_SUMMARY_LIMIT = 20; + +function truncateFileList(files: string[]): string[] { + if (files.length <= FILE_OPERATION_SUMMARY_LIMIT) return files; + const omitted = files.length - FILE_OPERATION_SUMMARY_LIMIT; + return [...files.slice(0, FILE_OPERATION_SUMMARY_LIMIT), `… (${omitted} more files omitted)`]; +} + +function stripFileOperationTags(summary: string): string { + const withoutReadFiles = summary.replace(/[\s\S]*?<\/read-files>\s*/g, ""); + const withoutModifiedFiles = withoutReadFiles.replace(/[\s\S]*?<\/modified-files>\s*/g, ""); + return withoutModifiedFiles.trimEnd(); +} + +function formatFileOperations(readFiles: string[], modifiedFiles: string[]): string { + if (readFiles.length === 0 && modifiedFiles.length === 0) return ""; + return prompt.render(fileOperationsTemplate, { + readFiles: truncateFileList(readFiles), + modifiedFiles: truncateFileList(modifiedFiles), + }); +} + +export function upsertSnapcompactFileOperations(summary: string, readFiles: string[], modifiedFiles: string[]): string { + const baseSummary = stripFileOperationTags(summary); + const fileOperations = formatFileOperations(readFiles, modifiedFiles); + if (!fileOperations) return baseSummary; + if (!baseSummary) return fileOperations; + return `${baseSummary}\n\n${fileOperations}`; +} + +// ============================================================================ +// Message serialization +// ============================================================================ + +const TOOL_RESULT_MAX_CHARS = 2000; + +function truncateForSummary(text: string, maxChars: number): string { + if (text.length <= maxChars) return text; + const truncatedChars = text.length - maxChars; + return `${text.slice(0, maxChars)}\n\n[... ${truncatedChars} more characters truncated]`; +} + +export function serializeSnapcompactConversation(messages: Message[]): string { + const parts: string[] = []; + + for (const msg of messages) { + if (msg.role === "user") { + const content = + typeof msg.content === "string" + ? msg.content + : msg.content + .filter((content): content is { type: "text"; text: string } => content.type === "text") + .map(content => content.text) + .join(""); + if (content) parts.push(`[User]: ${content}`); + } else if (msg.role === "assistant") { + const textParts: string[] = []; + const thinkingParts: string[] = []; + const toolCalls: string[] = []; + + for (const block of msg.content) { + if (block.type === "text") { + textParts.push(block.text); + } else if (block.type === "thinking") { + thinkingParts.push(block.thinking); + } else if (block.type === "toolCall") { + const args = block.arguments as Record; + const argsStr = Object.entries(args) + .map(([key, value]) => `${key}=${JSON.stringify(value)}`) + .join(", "); + toolCalls.push(`${block.name}(${argsStr})`); + } + } + + if (thinkingParts.length > 0) { + parts.push(`[Assistant thinking]: ${thinkingParts.join("\n")}`); + } + if (textParts.length > 0) { + parts.push(`[Assistant]: ${textParts.join("\n")}`); + } + if (toolCalls.length > 0) { + parts.push(`[Assistant tool calls]: ${toolCalls.join("; ")}`); + } + } else if (msg.role === "toolResult") { + const content = msg.content + .filter((block): block is { type: "text"; text: string } => block.type === "text") + .map(block => block.text) + .join(""); + if (content) { + parts.push(`[Tool result]: ${truncateForSummary(content, TOOL_RESULT_MAX_CHARS)}`); + } + } + } + + return parts.join("\n\n"); +} + +// ============================================================================ +// Preserve-data helpers +// ============================================================================ + +const OPENAI_REMOTE_COMPACTION_PRESERVE_KEY = "openaiRemoteCompaction"; + +function stripOpenAiRemoteCompactionPreserveData( + preserveData: Record | undefined, +): Record | undefined { + if (!preserveData || !(OPENAI_REMOTE_COMPACTION_PRESERVE_KEY in preserveData)) { + return preserveData; + } + const { [OPENAI_REMOTE_COMPACTION_PRESERVE_KEY]: _removed, ...rest } = preserveData; + return Object.keys(rest).length > 0 ? rest : undefined; +} + +// ============================================================================ +// Text normalization +// ============================================================================ + +/** Folds for common non-Latin-1 characters the bundled fonts cannot draw. */ +const CHAR_FOLD: Record = { + "\u2018": "'", + "\u2019": "'", + "\u201a": "'", + "\u201b": "'", + "\u201c": '"', + "\u201d": '"', + "\u201e": '"', + "\u2013": "-", + "\u2014": "-", + "\u2015": "-", + "\u2212": "-", + "\u2026": "...", + "\u2022": "*", + "\u25cf": "*", + "\u25a0": "*", + "\u25aa": "*", + "\u2190": "<-", + "\u2192": "->", + "\u21d2": "=>", + "\u2713": "v", + "\u2714": "v", + "\u2717": "x", + "\u2718": "x", +}; + +/** + * Prepare text for printing: collapse whitespace runs (incl. newlines) to + * single spaces — the eval's "paragraph breaks collapsed to spaces" format — + * then fold everything outside the fonts' ASCII + Latin-1 coverage to ASCII + * approximations (`?` as the last resort). + */ +export function normalizeForSnapcompact(text: string): string { + const collapsed = text.replace(/\s+/g, " ").trim(); + let out = ""; + for (const ch of collapsed) { + const cp = ch.codePointAt(0) as number; + if (cp < 0x7f || (cp >= 0xa0 && cp <= 0xff)) { + out += ch; + continue; + } + const fold = CHAR_FOLD[ch]; + if (fold !== undefined) { + out += fold; + } else if (cp >= 0x2500 && cp <= 0x257f) { + // Box drawing: keep table skeletons legible. + out += cp === 0x2502 || cp === 0x2503 ? "|" : cp === 0x2500 || cp === 0x2501 ? "-" : "+"; + } else { + out += "?"; + } + } + return out; +} + +// ============================================================================ +// Rendering +// ============================================================================ + +export function snapcompactGeometry(shape: SnapcompactShape, size: number = shape.frameSize): SnapcompactGeometry { + const cols = Math.floor(size / shape.cellWidth); + const rows = Math.floor(size / shape.cellHeight / shape.lineRepeat); + return { cols, rows, capacity: cols * rows }; +} + +/** Render one snapcompact frame from already-normalized text. */ +export function renderSnapcompactFrame( + text: string, + shape: SnapcompactShape, + size: number = shape.frameSize, +): RenderedFrame { + const { cols, rows, capacity } = snapcompactGeometry(shape, size); + const chars = Math.min(text.length, capacity); + const png = renderSnapcompactPng(text, { + size, + font: shape.font, + cellWidth: shape.cellWidth, + cellHeight: shape.cellHeight, + variant: shape.variant, + lineRepeat: shape.lineRepeat, + }); + return { png, cols, rows, chars }; +} + +// ============================================================================ +// Archive helpers +// ============================================================================ + +/** Validate and extract a persisted frame archive from `preserveData`. */ +export function getPreservedSnapcompactArchive( + preserveData: Record | undefined, +): SnapcompactArchive | undefined { + const candidate = preserveData?.[SNAPCOMPACT_PRESERVE_KEY]; + if (!candidate || typeof candidate !== "object") return undefined; + const archive = candidate as SnapcompactArchive; + if (!Array.isArray(archive.frames)) return undefined; + const frames = archive.frames.filter( + frame => + !!frame && + typeof frame.data === "string" && + frame.data.length > 0 && + typeof frame.mimeType === "string" && + typeof frame.cols === "number" && + typeof frame.rows === "number" && + typeof frame.chars === "number", + ); + if (frames.length === 0) return undefined; + return { + frames, + totalChars: typeof archive.totalChars === "number" ? archive.totalChars : 0, + truncatedChars: typeof archive.truncatedChars === "number" ? archive.truncatedChars : 0, + }; +} + +/** Convert archive frames into LLM image blocks (oldest first). */ +export function snapcompactImages(archive: SnapcompactArchive): ImageContent[] { + return archive.frames.map(frame => ({ + type: "image", + data: frame.data, + mimeType: frame.mimeType, + ...(frame.detail ? { detail: frame.detail } : {}), + })); +} + +// ============================================================================ +// Compaction entry point +// ============================================================================ + +/** + * Run a snapcompact compaction over prepared messages. Fully local: serializes + * the discarded history, prints it onto PNG frames in the provider-optimal + * shape, merges previously archived frames (oldest dropped beyond the + * budget), and produces a deterministic summary explaining how to read the + * frames. + * + * Frames archived under a different shape (provider switches, legacy 5x8 + * sessions) are kept as-is — each frame carries its own geometry, and the + * summary describes the newest shape while noting that older frames may + * differ. + * + * If the previous compaction was text-based, its summary is printed at the + * head of the frame archive as `[Summary of earlier history]` so no continuity is lost. + */ +export async function snapcompactCompact( + preparation: SnapcompactCompactionPreparation, + options?: SnapcompactOptions, +): Promise { + const { firstKeptEntryId, tokensBefore, previousSummary, previousPreserveData, fileOps } = preparation; + if (!firstKeptEntryId) { + throw new Error("First kept entry has no ID - session may need migration"); + } + const shape = options?.shape ?? resolveSnapcompactShape(options?.model?.api); + const frameSize = options?.frameSize ?? shape.frameSize; + const maxFrames = Math.max(1, options?.maxFrames ?? SNAPCOMPACT_MAX_FRAMES); + const geometry = snapcompactGeometry(shape, frameSize); + + const messages = preparation.messagesToSummarize.concat(preparation.turnPrefixMessages); + const llmMessages = (options?.convertToLlm ?? defaultConvertToLlm)(messages); + let archiveText = normalizeForSnapcompact(serializeSnapcompactConversation(llmMessages)); + + const previousArchive = getPreservedSnapcompactArchive(previousPreserveData); + const includedPreviousSummary = !previousArchive && !!previousSummary; + if (includedPreviousSummary && previousSummary) { + const head = `[Summary of earlier history] ${normalizeForSnapcompact(previousSummary)}`; + archiveText = archiveText.length > 0 ? `${head} [Recent conversation] ${archiveText}` : head; + } + + let truncatedChars = previousArchive?.truncatedChars ?? 0; + + const newFrames: SnapcompactFrame[] = []; + for (let offset = 0; offset < archiveText.length; offset += geometry.capacity) { + const chunk = archiveText.slice(offset, offset + geometry.capacity); + const rendered = renderSnapcompactFrame(chunk, shape, frameSize); + newFrames.push({ + data: Buffer.from(rendered.png).toBase64(), + mimeType: "image/png", + cols: rendered.cols, + rows: rendered.rows, + chars: rendered.chars, + font: shape.font, + variant: shape.variant, + lineRepeat: shape.lineRepeat, + ...(shape.imageDetail ? { detail: shape.imageDetail } : {}), + }); + // Keep the event loop responsive between native render passes. + await Bun.sleep(0); + } + + const frames = [...(previousArchive?.frames ?? []), ...newFrames]; + if (frames.length > maxFrames) { + // Pin the earliest frame: it anchors the session head (the original + // request, or the filmed summary of even older history) the way the + // LLM-summary strategies keep the original goal alive across rounds. + // Eviction removes the oldest *unpinned* frames, so the archive fades + // from the middle out — head and tail survive. With a budget of one + // frame the pin is moot; keep the newest frame instead. + const evictStart = maxFrames >= 2 ? 1 : 0; + const dropped = frames.splice(evictStart, frames.length - maxFrames); + for (const frame of dropped) truncatedChars += frame.chars; + } + const totalChars = frames.reduce((sum, frame) => sum + frame.chars, 0); + const mixedShapes = frames.some( + frame => + frame.cols !== geometry.cols || + frame.rows !== geometry.rows || + (frame.variant ?? "sent") !== shape.variant || + (frame.lineRepeat ?? 1) !== shape.lineRepeat, + ); + + let summary: string; + if (frames.length === 0) { + summary = "No prior history."; + } else { + summary = prompt.render(snapcompactSummaryPrompt, { + frameCount: frames.length, + multipleFrames: frames.length > 1, + fontCell: `${shape.cellWidth}x${shape.cellHeight}`, + cols: geometry.cols, + rows: geometry.rows, + sentenceInk: shape.variant === "sent", + lineRepeated: shape.lineRepeat > 1, + mixedShapes, + totalChars, + truncatedChars, + includedPreviousSummary, + }); + } + const { readFiles, modifiedFiles } = computeSnapcompactFileLists(fileOps); + summary = upsertSnapcompactFileOperations(summary, readFiles, modifiedFiles); + + // A snapcompact pass replaces any provider-side replacement history; strip the + // OpenAI remote-compaction payload like the default summarizer path does. + const basePreserve = stripOpenAiRemoteCompactionPreserveData(previousPreserveData) ?? {}; + const archive: SnapcompactArchive = { frames, totalChars, truncatedChars }; + + return { + summary, + shortSummary: `Archived ${totalChars.toLocaleString()} chars of history onto ${frames.length} snapcompact frame${frames.length === 1 ? "" : "s"}`, + firstKeptEntryId, + tokensBefore, + details: { readFiles, modifiedFiles }, + preserveData: { ...basePreserve, [SNAPCOMPACT_PRESERVE_KEY]: archive }, + }; +} diff --git a/packages/agent/test/snapcompact.test.ts b/packages/snapcompact/test/snapcompact.test.ts similarity index 64% rename from packages/agent/test/snapcompact.test.ts rename to packages/snapcompact/test/snapcompact.test.ts index fe224304c..316389ebb 100644 --- a/packages/agent/test/snapcompact.test.ts +++ b/packages/snapcompact/test/snapcompact.test.ts @@ -1,24 +1,60 @@ import { describe, expect, it } from "bun:test"; -import type { ImageContent } from "@oh-my-pi/pi-ai"; -import { type CompactionPreparation, type CompactionResult, estimateTokens } from "../src/compaction/compaction"; -import { createCompactionSummaryMessage, defaultConvertToLlm } from "../src/compaction/messages"; +import type { AssistantMessage, Message, Usage } from "@oh-my-pi/pi-ai"; import { + createSnapcompactFileOps, getPreservedSnapcompactArchive, normalizeForSnapcompact, renderSnapcompactFrame, - SNAPCOMPACT_FRAME_TOKEN_ESTIMATE, + resolveSnapcompactShape, SNAPCOMPACT_PRESERVE_KEY, + SNAPCOMPACT_SHAPES, type SnapcompactArchive, + type SnapcompactCompactionPreparation, + type SnapcompactCompactionResult, snapcompactCompact, snapcompactGeometry, -} from "../src/compaction/snapcompact"; -import { createFileOps } from "../src/compaction/utils"; -import { createAssistantMessage, createUserMessage } from "./helpers"; + snapcompactImages, +} from "../src"; -// Small frames keep render time negligible: 320px → 64 cols x 40 rows = 2560 chars. +// Small frames keep render time negligible. Legacy 5x8 shape: 320px → 64 cols +// x 40 rows = 2560 chars. Default (anthropic 8x8r-bw): 40 cols x 20 rows = 800. const TEST_FRAME_SIZE = 320; -function makePreparation(overrides: Partial = {}): CompactionPreparation { +function createUserMessage(content: string): Message { + return { role: "user", content, timestamp: 0 }; +} + +const ZERO_USAGE: Usage = { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + total: 0, + }, +}; + +function createAssistantMessage(content: AssistantMessage["content"]): Message { + return { + role: "assistant", + content, + api: "mock", + provider: "mock", + model: "mock", + usage: ZERO_USAGE, + stopReason: "stop", + timestamp: 0, + }; +} + +function makePreparation( + overrides: Partial> = {}, +): SnapcompactCompactionPreparation { return { firstKeptEntryId: "kept-1", messagesToSummarize: [ @@ -26,13 +62,10 @@ function makePreparation(overrides: Partial = {}): Compac createAssistantMessage([{ type: "text", text: "Fixed the TTL comparison in src/login.ts." }]), ], turnPrefixMessages: [], - recentMessages: [], - isSplitTurn: false, tokensBefore: 99000, previousSummary: undefined, previousPreserveData: undefined, - fileOps: createFileOps(), - settings: { enabled: true, reserveTokens: 16384, keepRecentTokens: 20000 }, + fileOps: createSnapcompactFileOps(), ...overrides, }; } @@ -106,12 +139,42 @@ describe("normalizeForSnapcompact", () => { }); }); +describe("shape resolution", () => { + it("maps provider APIs to their eval-winning shapes", () => { + expect(resolveSnapcompactShape("anthropic-messages")).toBe(SNAPCOMPACT_SHAPES.anthropic); + expect(resolveSnapcompactShape("openai-responses")).toBe(SNAPCOMPACT_SHAPES.openaiDense); + expect(resolveSnapcompactShape("azure-openai-responses")).toBe(SNAPCOMPACT_SHAPES.openaiDense); + expect(resolveSnapcompactShape("google-generative-ai")).toBe(SNAPCOMPACT_SHAPES.google); + // Unknown and absent APIs fall back to the refusal-robust plain shape. + expect(resolveSnapcompactShape("some-future-api")).toBe(SNAPCOMPACT_SHAPES.anthropic); + expect(resolveSnapcompactShape(undefined)).toBe(SNAPCOMPACT_SHAPES.anthropic); + }); + + it("snapcompactImages forwards the per-frame detail hint", () => { + const archive: SnapcompactArchive = { + frames: [ + { data: "ZmFrZQ==", mimeType: "image/png", cols: 10, rows: 10, chars: 5, detail: "original" }, + { data: "ZmFrZTI=", mimeType: "image/png", cols: 10, rows: 10, chars: 5 }, + ], + totalChars: 10, + truncatedChars: 0, + }; + const [withDetail, without] = snapcompactImages(archive); + expect(withDetail.detail).toBe("original"); + expect("detail" in without).toBe(false); + }); +}); + describe("renderSnapcompactFrame", () => { - it("produces an indexed PNG of the declared geometry with sentence-cycled ink", () => { - const geometry = snapcompactGeometry(TEST_FRAME_SIZE); + it("produces an indexed PNG of the declared geometry with sentence-cycled ink (legacy 5x8)", () => { + const geometry = snapcompactGeometry(SNAPCOMPACT_SHAPES.legacy, TEST_FRAME_SIZE); expect(geometry).toEqual({ cols: 64, rows: 40, capacity: 2560 }); - const frame = renderSnapcompactFrame("First sentence here. Second one differs.", TEST_FRAME_SIZE); + const frame = renderSnapcompactFrame( + "First sentence here. Second one differs.", + SNAPCOMPACT_SHAPES.legacy, + TEST_FRAME_SIZE, + ); expect(frame.cols).toBe(64); expect(frame.rows).toBe(40); expect(frame.chars).toBe(40); @@ -128,24 +191,46 @@ describe("renderSnapcompactFrame", () => { expect(used.has(3)).toBe(false); }); + it("renders the anthropic shape with doubled lines, black ink, and highlight bands", () => { + const geometry = snapcompactGeometry(SNAPCOMPACT_SHAPES.anthropic, TEST_FRAME_SIZE); + expect(geometry).toEqual({ cols: 40, rows: 20, capacity: 800 }); + + const frame = renderSnapcompactFrame("Hello world. Again.", SNAPCOMPACT_SHAPES.anthropic, TEST_FRAME_SIZE); + const decoded = decodePng(frame.png); + expect(decoded.colorType).toBe(3); + const used = new Set(decoded.pixels); + expect(used.has(7)).toBe(true); // black bw ink + expect(used.has(8)).toBe(true); // repeat highlight band + expect(used.has(1)).toBe(false); // no sentence hues in bw + }); + + it("renders the openai stretch shape as truecolor RGB", () => { + const frame = renderSnapcompactFrame("Hello world.", SNAPCOMPACT_SHAPES.openaiDense, TEST_FRAME_SIZE); + // IHDR color type byte: 2 = truecolor RGB (anti-aliased stretch output). + expect(frame.png[25]).toBe(2); + expect(frame.cols).toBe(Math.floor(TEST_FRAME_SIZE / 6)); + }); + it("caps printed characters at frame capacity", () => { - const { capacity } = snapcompactGeometry(TEST_FRAME_SIZE); - const frame = renderSnapcompactFrame("x".repeat(capacity + 500), TEST_FRAME_SIZE); + const { capacity } = snapcompactGeometry(SNAPCOMPACT_SHAPES.legacy, TEST_FRAME_SIZE); + const frame = renderSnapcompactFrame("x".repeat(capacity + 500), SNAPCOMPACT_SHAPES.legacy, TEST_FRAME_SIZE); expect(frame.chars).toBe(capacity); }); }); describe("snapcompactCompact", () => { it("archives history onto frames with a self-describing summary", async () => { - const fileOps = createFileOps(); + const fileOps = createSnapcompactFileOps(); fileOps.read.add("src/auth.ts"); fileOps.edited.add("src/login.ts"); const result = await snapcompactCompact(makePreparation({ fileOps }), { frameSize: TEST_FRAME_SIZE }); expect(result.firstKeptEntryId).toBe("kept-1"); expect(result.tokensBefore).toBe(99000); - // Reading instructions reflect the actual grid geometry. - expect(result.summary).toContain("64 characters per row"); + // Reading instructions reflect the default (anthropic 8x8r-bw) shape. + expect(result.summary).toContain("40 characters per row"); + expect(result.summary).toContain("printed twice"); + expect(result.summary).toContain("plain black ink"); expect(result.summary).toContain("snapcompact frame"); // File operations are upserted like every other compaction summary. expect(result.summary).toContain(""); @@ -157,6 +242,9 @@ describe("snapcompactCompact", () => { expect(archive?.frames.length).toBe(1); expect(archive?.frames[0].mimeType).toBe("image/png"); expect(archive?.frames[0].chars).toBe(archive?.totalChars); + expect(archive?.frames[0].font).toBe("8x8"); + expect(archive?.frames[0].variant).toBe("bw"); + expect(archive?.frames[0].lineRepeat).toBe(2); expect(archive?.truncatedChars).toBe(0); // Frame data round-trips as a decodable PNG. const decoded = decodePng(Buffer.from(archive?.frames[0].data ?? "", "base64")); @@ -164,7 +252,7 @@ describe("snapcompactCompact", () => { }); it("splits oversized history across frames and evicts beyond the budget", async () => { - const { capacity } = snapcompactGeometry(TEST_FRAME_SIZE); + const { capacity } = snapcompactGeometry(SNAPCOMPACT_SHAPES.anthropic, TEST_FRAME_SIZE); // Sentences avoid whitespace collapse shrinking the payload below 2.5 frames. const longText = "Important fact number one. ".repeat(Math.ceil((capacity * 2.5) / 28)); const result = await snapcompactCompact(makePreparation({ messagesToSummarize: [createUserMessage(longText)] }), { @@ -178,7 +266,7 @@ describe("snapcompactCompact", () => { }); it("evicts the oldest unpinned frames, keeping the session-head frame alive", async () => { - let previous: CompactionResult | undefined; + let previous: SnapcompactCompactionResult | undefined; let headFrameData = ""; let secondFrameData = ""; for (let pass = 1; pass <= 4; pass++) { @@ -237,46 +325,27 @@ describe("snapcompactCompact", () => { expect(second.preserveData?.openaiRemoteCompaction).toBeUndefined(); expect(second.preserveData?.appKey).toBe("kept"); }); + + it("flags mixed shapes when merged frames disagree with the active shape", async () => { + const first = await snapcompactCompact(makePreparation(), { + frameSize: TEST_FRAME_SIZE, + shape: SNAPCOMPACT_SHAPES.legacy, + }); + const second = await snapcompactCompact( + makePreparation({ + messagesToSummarize: [createUserMessage("Another turn after a provider switch.")], + previousSummary: first.summary, + previousPreserveData: first.preserveData, + }), + { frameSize: TEST_FRAME_SIZE, model: { api: "anthropic-messages" } }, + ); + expect(second.summary).toContain("Older frames may use a different font"); + // Same-shape merges stay silent. + expect(first.summary).not.toContain("Older frames may use a different font"); + }); }); -describe("compaction summary message with snapcompact frames", () => { - const images: ImageContent[] = [ - { type: "image", data: "ZmFrZQ==", mimeType: "image/png" }, - { type: "image", data: "ZmFrZTI=", mimeType: "image/png" }, - ]; - - it("estimateTokens charges per attached frame", () => { - const bare = createCompactionSummaryMessage("summary text", 1000, new Date().toISOString()); - const withFrames = createCompactionSummaryMessage( - "summary text", - 1000, - new Date().toISOString(), - undefined, - undefined, - images, - ); - expect(estimateTokens(withFrames) - estimateTokens(bare)).toBe(2 * SNAPCOMPACT_FRAME_TOKEN_ESTIMATE); - }); - - it("defaultConvertToLlm appends frames as image blocks after the summary text", () => { - const message = createCompactionSummaryMessage( - "the snapcompact archive", - 1000, - new Date().toISOString(), - undefined, - undefined, - images, - ); - const [converted] = defaultConvertToLlm([message]); - expect(converted.role).toBe("user"); - const content = converted.content as Array<{ type: string; text?: string; data?: string }>; - expect(content.length).toBe(3); - expect(content[0].type).toBe("text"); - expect(content[0].text).toContain("the snapcompact archive"); - expect(content[1]).toEqual(images[0]); - expect(content[2]).toEqual(images[1]); - }); - +describe("archive helpers", () => { it("getPreservedSnapcompactArchive rejects malformed payloads", () => { expect(getPreservedSnapcompactArchive(undefined)).toBeUndefined(); expect(getPreservedSnapcompactArchive({ [SNAPCOMPACT_PRESERVE_KEY]: "nope" })).toBeUndefined(); diff --git a/packages/snapcompact/tsconfig.json b/packages/snapcompact/tsconfig.json new file mode 100644 index 000000000..08130e07c --- /dev/null +++ b/packages/snapcompact/tsconfig.json @@ -0,0 +1,7 @@ +{ + "extends": "../tsconfig.workspace.json", + "include": [ + "src", + "test" + ] +} diff --git a/packages/snapcompact/tsconfig.publish.json b/packages/snapcompact/tsconfig.publish.json new file mode 100644 index 000000000..c216426de --- /dev/null +++ b/packages/snapcompact/tsconfig.publish.json @@ -0,0 +1,25 @@ +{ + "extends": "./tsconfig.json", + "compilerOptions": { + "noEmit": false, + "emitDeclarationOnly": true, + "declaration": true, + "declarationMap": false, + "sourceMap": false, + "inlineSources": false, + "rootDir": "src", + "outDir": "dist/types", + "noCheck": true + }, + "include": [ + "src" + ], + "exclude": [ + "dist", + "node_modules", + "test", + "tests", + "scripts", + "bench" + ] +} diff --git a/packages/typescript-edit-benchmark/package.json b/packages/typescript-edit-benchmark/package.json index 9ad9028fe..21c8553b7 100644 --- a/packages/typescript-edit-benchmark/package.json +++ b/packages/typescript-edit-benchmark/package.json @@ -4,6 +4,7 @@ "name": "@oh-my-pi/typescript-edit-benchmark", "version": "0.0.1", "description": "Edit benchmark suite using TypeScript source code mutations", + "homepage": "https://omp.sh", "author": "Can Boluk", "license": "MIT", "repository": { diff --git a/packages/utils/test/prompt.test.ts b/packages/utils/test/prompt.test.ts index 7dd5e2126..7b0700bf3 100644 --- a/packages/utils/test/prompt.test.ts +++ b/packages/utils/test/prompt.test.ts @@ -96,7 +96,7 @@ describe("compile cache", () => { }); describe("helpers: join", () => { - it('unescapes \\n and \\t in the separator (Handlebars string literals carry no escapes)', () => { + it("unescapes \\n and \\t in the separator (Handlebars string literals carry no escapes)", () => { // Regression: `{{join files "\n"}}` used to emit the literal two-char `\n` // between entries (visible in compaction lists). expect(prompt.render('{{join files "\\n"}}', { files: ["a.ts", "b.ts"] })).toBe("a.ts\nb.ts"); diff --git a/python/robomp/web/package.json b/python/robomp/web/package.json index feb6caeaa..681d6471c 100644 --- a/python/robomp/web/package.json +++ b/python/robomp/web/package.json @@ -4,6 +4,7 @@ "version": "0.1.0", "type": "module", "description": "Glassmorphic SolidJS dashboard bundled by Vite and served by robomp's FastAPI app.", + "homepage": "https://omp.sh", "scripts": { "dev": "vite", "build": "vite build", diff --git a/scripts/ci-release-publish.ts b/scripts/ci-release-publish.ts index b5b9af951..87b7d36e1 100644 --- a/scripts/ci-release-publish.ts +++ b/scripts/ci-release-publish.ts @@ -91,6 +91,7 @@ export const packages: PublishPackage[] = [ { dir: "packages/tui", kind: "typescript" }, { dir: "packages/hashline", kind: "typescript" }, { dir: "packages/mnemopi", kind: "typescript" }, + { dir: "packages/snapcompact", kind: "typescript" }, { dir: "packages/stats", kind: "typescript", diff --git a/scripts/install-tests/run-ci.sh b/scripts/install-tests/run-ci.sh index 62f50b140..f401b1895 100755 --- a/scripts/install-tests/run-ci.sh +++ b/scripts/install-tests/run-ci.sh @@ -92,7 +92,7 @@ cp "$natives_pkg_backup" "$ROOT_DIR/packages/natives/package.json" # 3. Pack the remaining workspace packages (natives core and coding-agent # handled separately). -for pkg in utils hashline catalog ai mnemopi agent tui stats; do +for pkg in utils hashline catalog ai mnemopi snapcompact agent tui stats; do ( cd "$ROOT_DIR/packages/$pkg" bun pm pack --destination "$TARBALL_DIR" --quiet >/dev/null @@ -122,6 +122,7 @@ hashline_tgz="$(find_tarball "$TARBALL_DIR"/oh-my-pi-hashline-*.tgz)" catalog_tgz="$(find_tarball "$TARBALL_DIR"/oh-my-pi-pi-catalog-*.tgz)" ai_tgz="$(find_tarball "$TARBALL_DIR"/oh-my-pi-pi-ai-*.tgz)" mnemopi_tgz="$(find_tarball "$TARBALL_DIR"/oh-my-pi-pi-mnemopi-*.tgz)" +snapcompact_tgz="$(find_tarball "$TARBALL_DIR"/oh-my-pi-snapcompact-*.tgz)" agent_tgz="$(find_tarball "$TARBALL_DIR"/oh-my-pi-pi-agent-core-*.tgz)" tui_tgz="$(find_tarball "$TARBALL_DIR"/oh-my-pi-pi-tui-*.tgz)" stats_tgz="$(find_tarball "$TARBALL_DIR"/oh-my-pi-omp-stats-*.tgz)" @@ -145,6 +146,7 @@ mkdir -p "$TARBALL_APP_DIR" '@oh-my-pi/pi-ai': '$ai_tgz', '@oh-my-pi/pi-catalog': '$catalog_tgz', '@oh-my-pi/pi-mnemopi': '$mnemopi_tgz', + '@oh-my-pi/snapcompact': '$snapcompact_tgz', '@oh-my-pi/pi-agent-core': '$agent_tgz', '@oh-my-pi/pi-tui': '$tui_tgz', '@oh-my-pi/omp-stats': '$stats_tgz', @@ -153,7 +155,7 @@ mkdir -p "$TARBALL_APP_DIR" require('fs').writeFileSync('package.json', JSON.stringify(pkg, null, 2)); " - bun add "$utils_tgz" "$natives_tgz" "$hashline_tgz" "$catalog_tgz" "$ai_tgz" "$mnemopi_tgz" "$agent_tgz" "$tui_tgz" "$stats_tgz" "$coding_agent_tgz" + bun add "$utils_tgz" "$natives_tgz" "$hashline_tgz" "$catalog_tgz" "$ai_tgz" "$mnemopi_tgz" "$snapcompact_tgz" "$agent_tgz" "$tui_tgz" "$stats_tgz" "$coding_agent_tgz" # The platform leaf must arrive through the core's optionalDependencies + # override, not as a direct dependency — assert it landed before smoking so a # resolution regression is distinguishable from a runtime loader bug. From 0fe75212b46a478da1e10ff886f89853f7e61a4a Mon Sep 17 00:00:00 2001 From: can1357 Date: Wed, 10 Jun 2026 21:50:19 +0200 Subject: [PATCH 27/27] ux(hashline): appended context previews to unresolved block-edit errors - Added a shared `formatAnchoredContext` helper and switched mismatch formatting to use it. - Extended unresolved block-edit errors to append nearby numbered context with `*`-marked anchor lines. - Updated unit and integration tests and docs to reflect the new unresolved-block preview behavior. --- docs/tools/edit.md | 2 +- .../test/core/block-replace.test.ts | 3 +- packages/hashline/CHANGELOG.md | 4 ++ packages/hashline/src/block.ts | 4 +- packages/hashline/src/messages.ts | 49 ++++++++++++++++--- packages/hashline/src/mismatch.ts | 30 ++---------- packages/hashline/test/block.test.ts | 29 +++++++++++ 7 files changed, 86 insertions(+), 35 deletions(-) diff --git a/docs/tools/edit.md b/docs/tools/edit.md index 401edf035..dd4b6231a 100644 --- a/docs/tools/edit.md +++ b/docs/tools/edit.md @@ -172,7 +172,7 @@ delete 20 - `line N: \`insert\` needs at least one \`+TEXT\` body row.` - `line N: \`replace block N:\` needs at least one \`+TEXT\` body row. To delete a block, use \`delete N..M\` with the block's line range.` - Unresolvable `replace block N:` (apply / final-preview path only): - - `line N: \`replace block X:\` could not resolve a syntactic block beginning on line X. The language may be unsupported, the line may be blank or a closing delimiter, or the block may not parse. Use \`replace X..M:\` with the block's explicit end line instead.` + - `line N: \`replace block X:\` could not resolve a syntactic block beginning on line X. The language may be unsupported, the line may be blank or a closing delimiter, or the block may not parse. Use \`replace X..M:\` with the block's explicit end line instead.` — followed by a blank line and numbered `*`-marked context rows around line X (same shape as the mismatch preview). - Delete with body: - `line N: \`delete N..M\` does not take body rows. Remove the body, or use \`replace N..M:\`.` - `line N: \`delete block N\` does not take body rows. Remove the body, or use \`replace block N:\` to replace the block.` diff --git a/packages/coding-agent/test/core/block-replace.test.ts b/packages/coding-agent/test/core/block-replace.test.ts index 769e63eea..c247d17e6 100644 --- a/packages/coding-agent/test/core/block-replace.test.ts +++ b/packages/coding-agent/test/core/block-replace.test.ts @@ -148,8 +148,9 @@ describe("replace block — native tree-sitter resolution end-to-end", () => { // Line 3 is ` }` — a closing delimiter, not a block opener. const input = `${header}\nreplace block 3:\n+ }`; + // Steers to the concrete form and previews the file around the anchor (`*`-marked). await expect(executeHashlineSingle(executeOptions(tempDir, input, session))).rejects.toThrow( - /could not resolve a syntactic block beginning on line 3.*replace 3\.\.M:/s, + /could not resolve a syntactic block beginning on line 3.*replace 3\.\.M:.*^ 1:function x\(\) \{$.*^\*3: {2}\}$/ms, ); // Disk untouched — refusal never leaves a partial write. expect(await Bun.file(filePath).text()).toBe(TS_SOURCE); diff --git a/packages/hashline/CHANGELOG.md b/packages/hashline/CHANGELOG.md index 98f72e212..dafea4c1b 100644 --- a/packages/hashline/CHANGELOG.md +++ b/packages/hashline/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Changed + +- Block-unresolved errors (`replace block N:` / `delete block N` / `insert after block N:` failing to resolve a syntactic block) now append a numbered preview of the file around the anchor line — same `*`-marked context rows the hash-mismatch error shows — so the offending line is visible without a re-read + ## [15.10.11] - 2026-06-10 ### Breaking Changes diff --git a/packages/hashline/src/block.ts b/packages/hashline/src/block.ts index d4b44cb75..8f7ae87ab 100644 --- a/packages/hashline/src/block.ts +++ b/packages/hashline/src/block.ts @@ -69,7 +69,9 @@ export function resolveBlockEdits( if (span === null) { if (onUnresolved === "drop") continue; throw new Error( - `line ${edit.lineNum}: ${resolver ? blockUnresolvedMessage(edit.anchor.line, op) : BLOCK_RESOLVER_UNAVAILABLE}`, + `line ${edit.lineNum}: ${ + resolver ? blockUnresolvedMessage(edit.anchor.line, op, text.split("\n")) : BLOCK_RESOLVER_UNAVAILABLE + }`, ); } options.onResolved?.({ diff --git a/packages/hashline/src/messages.ts b/packages/hashline/src/messages.ts index 4cc6493ea..73f43b15d 100644 --- a/packages/hashline/src/messages.ts +++ b/packages/hashline/src/messages.ts @@ -5,11 +5,37 @@ * them. */ -import { HL_FILE_HASH_SEP, HL_FILE_PREFIX, HL_FILE_SUFFIX } from "./format"; +import { formatNumberedLine, HL_FILE_HASH_SEP, HL_FILE_PREFIX, HL_FILE_SUFFIX } from "./format"; /** Lines of context shown either side of a hash mismatch. */ export const MISMATCH_CONTEXT = 2; +/** + * Render numbered `LINE:TEXT` context rows around `anchorLines` + * (±{@link MISMATCH_CONTEXT} lines each), `*`-marking the anchored lines and + * separating non-adjacent runs with `...`. Out-of-range anchors contribute no + * rows; returns an empty array when every anchor is out of range. + */ +export function formatAnchoredContext(anchorLines: readonly number[], fileLines: readonly string[]): string[] { + const displayLines = new Set(); + for (const line of anchorLines) { + if (line < 1 || line > fileLines.length) continue; + const lo = Math.max(1, line - MISMATCH_CONTEXT); + const hi = Math.min(fileLines.length, line + MISMATCH_CONTEXT); + for (let lineNum = lo; lineNum <= hi; lineNum++) displayLines.add(lineNum); + } + const anchorSet = new Set(anchorLines); + const rows: string[] = []; + let previous = -1; + for (const lineNum of [...displayLines].sort((a, b) => a - b)) { + if (previous !== -1 && lineNum > previous + 1) rows.push("..."); + previous = lineNum; + const marker = anchorSet.has(lineNum) ? "*" : " "; + rows.push(`${marker}${formatNumberedLine(lineNum, fileLines[lineNum - 1] ?? "")}`); + } + return rows; +} + /** Optional patch envelope start marker; silently consumed when present. */ export const BEGIN_PATCH_MARKER = "*** Begin Patch"; @@ -50,10 +76,15 @@ export const EMPTY_BLOCK = * Error text emitted when a block-anchored op cannot be resolved to a * syntactic block (unrecognized language, blank/out-of-range line, no node * begins on line N such as a lone closing delimiter, or the resolved block has - * a syntax error). Names the offending line and steers back to an explicit - * concrete-line form. + * a syntax error). Names the offending line, steers back to an explicit + * concrete-line form, and — when `fileLines` is provided — appends a + * {@link formatAnchoredContext} preview of the file around the anchor line. */ -export function blockUnresolvedMessage(line: number, op: "replace" | "delete" | "insert_after" = "replace"): string { +export function blockUnresolvedMessage( + line: number, + op: "replace" | "delete" | "insert_after" = "replace", + fileLines?: readonly string[], +): string { const phrase = op === "delete" ? `delete block ${line}` @@ -66,11 +97,15 @@ export function blockUnresolvedMessage(line: number, op: "replace" | "delete" | : op === "insert_after" ? `\`insert after M:\` with the block's explicit last line` : `\`replace ${line}..M:\` with the block's explicit end line`; - return ( + let message = `\`${phrase}\` could not resolve a syntactic block beginning on line ${line}. ` + `The language may be unsupported, the line may be blank or a closing delimiter, or the block may not parse. ` + - `Use ${fallback} instead.` - ); + `Use ${fallback} instead.`; + if (fileLines) { + const context = formatAnchoredContext([line], fileLines); + if (context.length > 0) message += `\n\n${context.join("\n")}`; + } + return message; } /** diff --git a/packages/hashline/src/mismatch.ts b/packages/hashline/src/mismatch.ts index 5e9bed476..c1e171ac7 100644 --- a/packages/hashline/src/mismatch.ts +++ b/packages/hashline/src/mismatch.ts @@ -6,8 +6,8 @@ * plus a couple of lines of surrounding context. The {@link MismatchError} * formats this into a message at construction time. */ -import { formatNumberedLine, HL_FILE_HASH_EXAMPLES, HL_FILE_HASH_SEP, HL_FILE_PREFIX, HL_FILE_SUFFIX } from "./format"; -import { MISMATCH_CONTEXT } from "./messages"; +import { HL_FILE_HASH_EXAMPLES, HL_FILE_HASH_SEP, HL_FILE_PREFIX, HL_FILE_SUFFIX } from "./format"; +import { formatAnchoredContext } from "./messages"; const LINE_REF_RE = /^\s*[>+\-*]*\s*(\d+)(?::.*)?\s*$/; /** Format the required-shape diagnostic shown when a line reference is malformed. */ @@ -46,17 +46,6 @@ export interface MismatchDetails { hashRecognized?: boolean; } -function getMismatchDisplayLines(anchorLines: readonly number[], fileLines: string[]): number[] { - const displayLines = new Set(); - for (const line of anchorLines) { - if (line < 1 || line > fileLines.length) continue; - const lo = Math.max(1, line - MISMATCH_CONTEXT); - const hi = Math.min(fileLines.length, line + MISMATCH_CONTEXT); - for (let lineNum = lo; lineNum <= hi; lineNum++) displayLines.add(lineNum); - } - return [...displayLines].sort((a, b) => a - b); -} - /** * Raised when a hashline section's snapshot tag doesn't match the live file's * content (and recovery, if configured, declined the merge). Carries the @@ -113,19 +102,10 @@ export class MismatchError extends Error { } static formatMessage(details: MismatchDetails): string { - const anchorSet = new Set(details.anchorLines ?? []); const lines = MismatchError.rejectionHeader(details); - const displayLines = getMismatchDisplayLines(details.anchorLines ?? [], details.fileLines); - if (displayLines.length === 0) return lines.join("\n"); - lines.push(""); - let previous = -1; - for (const lineNum of displayLines) { - if (previous !== -1 && lineNum > previous + 1) lines.push("..."); - previous = lineNum; - const text = details.fileLines[lineNum - 1] ?? ""; - const marker = anchorSet.has(lineNum) ? "*" : " "; - lines.push(`${marker}${formatNumberedLine(lineNum, text)}`); - } + const context = formatAnchoredContext(details.anchorLines ?? [], details.fileLines); + if (context.length === 0) return lines.join("\n"); + lines.push("", ...context); return lines.join("\n"); } } diff --git a/packages/hashline/test/block.test.ts b/packages/hashline/test/block.test.ts index 548fa3010..2e404f7ed 100644 --- a/packages/hashline/test/block.test.ts +++ b/packages/hashline/test/block.test.ts @@ -87,6 +87,35 @@ describe("resolveBlockEdits", () => { ); }); + it("includes a nearby-context preview in the block-unresolved error", () => { + const edits = parsePatch("replace block 3:\n+X").edits; + const text = "alpha\nbravo\ncharlie\ndelta\necho\nfoxtrot"; + let error: Error | undefined; + try { + resolveBlockEdits(edits, text, PATH, () => null); + } catch (err) { + error = err as Error; + } + expect(error?.message).toContain("could not resolve a syntactic block beginning on line 3"); + // ±2 lines of context around the anchor, anchor `*`-marked. + expect(error?.message).toContain(" 1:alpha"); + expect(error?.message).toContain("*3:charlie"); + expect(error?.message).toContain(" 5:echo"); + expect(error?.message).not.toContain("foxtrot"); + }); + + it("omits the context preview when the anchor line is out of range", () => { + const edits = parsePatch("replace block 9:\n+X").edits; + let error: Error | undefined; + try { + resolveBlockEdits(edits, "only\ntwo", PATH, () => null); + } catch (err) { + error = err as Error; + } + expect(error?.message).toContain("could not resolve a syntactic block beginning on line 9"); + expect(error?.message).not.toContain("\n\n"); + }); + it("fires onResolved with the resolved span for replace and delete blocks", () => { const seen: BlockResolution[] = []; // stubResolver maps line N → span [N, N+1].