diff --git a/extensions/subagents/docs/design-plan.md b/extensions/subagents/docs/design-plan.md index c9bcaf77..3ea2ace6 100644 --- a/extensions/subagents/docs/design-plan.md +++ b/extensions/subagents/docs/design-plan.md @@ -32,7 +32,7 @@ Source: `extensions/subagents/` (`index.ts`, `manager.ts`, `prompt.ts`, | Tool | Parameters | Behavior | |---|---|---| | `subagent_spawn` | `prompt`, `title`, `working_dir?`, `model?`, `provider?`, `reasoning_effort?` | Fire-and-forget spawn. Returns immediately with an id (`sa-N`). Enforces `MAX_RUNNING = 4` with a synchronous reservation so parallel tool calls can't race past the cap. Validates `working_dir`, resolves model against the registry (inherit parent model/thinking level by default), truncates title to 160 chars. | -| `subagent_wait` | `ids[]` (max 64) | Blocks until all listed subagents settle; respects the tool `AbortSignal`; streams `Waiting for ...` via `onUpdate`. Marks the awaited results "consumed" so they are not also auto-delivered. Output budgets: 48KB total, 16KB per agent, with per-section fallbacks (`[omitted: ...]`). Errors on unknown ids (lists known ids). | +| `subagent_wait` | `ids[]` (max 64) | Blocks until all listed subagents settle; respects the tool `AbortSignal`; streams `Waiting for ...` via `onUpdate`. Marks the awaited results "consumed" so they are not also auto-delivered. The hard output ceiling is 48 KiB, with a 16 KiB static per-agent cap. When Pi reports authoritative parent context usage, projections may narrow to 50% of the remaining headroom after fixed wrapper text; short results yield unused bytes to longer siblings. Unknown or invalid usage falls back to the static caps. Errors on unknown ids (lists known ids). | | `subagent_cancel` | `ids[]` | Aborts running subagents (marks consumed first to avoid duplicate delivery), waits for settlement, reports per-id `Cancelled ...` / `was already `. Partial transcripts remain on disk. | | `subagent_check` | `id` | Non-blocking peek: status line, turn count, error text, up to 2KB/20 lines of latest output (includes the live streaming assistant message). Does not consume the result. | | `subagent_list` | — | One `describeSubagent()` line per agent: `id [status] "title" (provider/model, ctx%, elapsed, cwd)`. | @@ -81,8 +81,13 @@ the parent conversation. { deliverAs: "followUp", triggerTurn: true })`; a separate session entry renders the report at its actual completion point. Content is built by `buildSubagentResultMessage` (`Subagent sa-N "title" - finished/failed.` + optional `Error:` line + output truncated to 24KB/600 lines with a - pointer to the child session file for the full transcript). + finished/failed.` + optional `Error:` line). Automatic batches have a 48 KiB + hard ceiling and a 24 KiB static per-result cap. As with `subagent_wait`, Pi's + authoritative context usage can narrow the projection budget dynamically; the + fixed headers, separators, and guidance are charged before result bytes are + allocated. Oversized results retain roughly 75% head + 25% tail and point to an exact, + content-addressed final-answer artifact below the Pi agent cache; the parent + can page it with Pi's native `read` instead of parsing the child session JSONL). ### 1.4 UI (carried over into v2 essentially as-is) @@ -144,7 +149,9 @@ Common denominator all three can supply: token usage, errors; - a way to send a follow-up/steering user message into a live session; - an interrupt operation; -- a final result text per run; +- a final result text per run; oversized parent projections preserve both the + head and tail, while the exact final text remains available as a plain-text + artifact for native `read` pagination; - metadata: backend name, model identifier, session/log file path (pi session file, Claude session id + projects dir JSONL, Codex rollout path), working dir. @@ -564,8 +571,14 @@ Recommendation: (a) during development, rename to final names when v2 replaces v 7. **Binary/SDK discovery + failure UX.** When `codex`/`claude` isn't installed or has no credentials, should `subagent_spawn` fail fast with a clear tool error (proposed), or should the backends be hidden from the `agent` enum dynamically? -8. **Result truncation budgets.** Keep v1's numbers (24KB result message, 48KB wait - total, 16KB per agent, 2KB check preview) unchanged? +8. **Result projection budgets (resolved).** Keep the 24 KiB automatic per-result, + 48 KiB batch, 16 KiB wait per-agent, and 2 KiB check ceilings as static safety + caps. For automatic delivery and explicit waits, narrow the projection when Pi's + authoritative parent-context reading shows less headroom: spend at most 50% of + the remaining tokens (estimated at four UTF-8 bytes each), subtract fixed wrapper + text first, and distribute the remainder across the batch. The runtime does not + maintain a second token counter; missing or stale Pi usage falls back to static + caps. Exact content remains in the artifact regardless of projection size. 9. **Effect version pinning.** Effect v4 is beta — pin an exact `4.0.0-beta.x` and accept manual bumps, or track the beta dist-tag? 10. **Persistence across reloads.** v1 loses all subagents on `session_shutdown` diff --git a/extensions/subagents/index.test.ts b/extensions/subagents/index.test.ts index 9d037509..b1e46f81 100644 --- a/extensions/subagents/index.test.ts +++ b/extensions/subagents/index.test.ts @@ -9,7 +9,11 @@ import type { ExtensionContext, } from "@earendil-works/pi-coding-agent"; import { PLAN_MODE_CHANNEL } from "../shared/plan-mode-state.ts"; -import subagents, { createSubagentResultDispatcher } from "./index.ts"; +import subagents, { + createSubagentResultDispatcher, + truncatedOutput, +} from "./index.ts"; +import { projectResult } from "./src/result-artifact.ts"; const emptySessionManager = { getBranch: () => [] }; @@ -78,6 +82,129 @@ test("subagent results render before the hidden wake-up message", () => { ]); }); +test("automatic result projection keeps both ends and persists the exact final answer", () => { + const finalText = `BEGIN\n${"evidence\n".repeat(100)}FINAL-VERDICT`; + let persisted = ""; + const text = truncatedOutput( + { + id: "sa-3", + origin: "model", + backend: "pi", + title: "inspect", + prompt: "inspect", + cwd: process.cwd(), + status: "done", + createdAt: 0, + settledAt: 1_000, + meta: { backend: "pi" }, + usage: {}, + transcript: [], + liveTools: [], + queued: [], + finalText, + turns: 1, + }, + 120, + (content) => { + persisted = content; + return "/tmp/subagent-final.txt"; + }, + ); + + assert.equal(persisted, finalText); + assert.match(text, /^BEGIN/); + assert.match(text, /FINAL-VERDICT/); + assert.match(text, /Full final answer: "\/tmp\/subagent-final\.txt"/); +}); + +test("automatic result delivery shrinks a batch against authoritative parent headroom", () => { + const budgets: number[] = []; + const pi = { + appendEntry() {}, + sendMessage() {}, + } as unknown as ExtensionAPI; + const dispatch = createSubagentResultDispatcher( + pi, + (_snap, maxBytes) => { + budgets.push(maxBytes); + return "projected"; + }, + () => ({ tokens: 98_000, contextWindow: 100_000 }), + ); + const snapshot = (id: string) => ({ + id, + origin: "model" as const, + backend: "pi" as const, + title: id, + prompt: "inspect", + cwd: process.cwd(), + status: "done" as const, + createdAt: 0, + settledAt: 1_000, + meta: { backend: "pi" as const }, + usage: {}, + transcript: [], + liveTools: [], + queued: [], + finalText: "x".repeat(40 * 1024), + turns: 1, + }); + + dispatch([snapshot("sa-1"), snapshot("sa-2")]); + + assert.deepEqual(budgets, [2048, 2048]); +}); + +test("automatic result wrappers and projections stay inside the shared batch cap", () => { + let delivered = ""; + const pi = { + appendEntry(_customType: string, data: { content: string }) { + delivered = data.content; + }, + sendMessage() {}, + } as unknown as ExtensionAPI; + const dispatch = createSubagentResultDispatcher( + pi, + (snap, maxBytes) => + projectResult(snap.finalText, { + maxBytes, + maxLines: 600, + writeArtifact: () => `/tmp/${snap.id}.txt`, + }).text, + ); + const snapshot = (id: string) => ({ + id, + origin: "model" as const, + backend: "pi" as const, + title: `long report ${id}`, + prompt: "inspect", + cwd: process.cwd(), + status: "done" as const, + createdAt: 0, + settledAt: 1_000, + meta: { backend: "pi" as const }, + usage: {}, + transcript: [], + liveTools: [], + queued: [], + finalText: `BEGIN-${id}\n${"evidence\n".repeat(10_000)}END-${id}`, + turns: 1, + }); + + dispatch([ + snapshot("sa-1"), + snapshot("sa-2"), + snapshot("sa-3"), + snapshot("sa-4"), + ]); + + assert.ok(Buffer.byteLength(delivered, "utf8") <= 48 * 1024); + for (const id of ["sa-1", "sa-2", "sa-3", "sa-4"]) { + assert.match(delivered, new RegExp(`BEGIN-${id}`)); + assert.match(delivered, new RegExp(`END-${id}`)); + } +}); + test("the visible subagent result entry renders the completed report", () => { const renderers = new Map(); const pi = { diff --git a/extensions/subagents/index.ts b/extensions/subagents/index.ts index 8f714248..6d2eb763 100644 --- a/extensions/subagents/index.ts +++ b/extensions/subagents/index.ts @@ -36,7 +36,6 @@ import { DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES, defineTool, - formatSize, getAgentDir, getMarkdownTheme, keyHint, @@ -45,11 +44,50 @@ import { import { Markdown, Text } from "@earendil-works/pi-tui"; import { Type } from "typebox"; import { + formatActivityStatus, + hasActivity, + unreadActivityCounts, +} from "../shared/activity-status.ts"; +import { + BelowEditorNavigationEditor, + BelowEditorStripState, +} from "../shared/below-editor-navigation.ts"; +import { + effectiveChildToolAllowlist, + resolveStandaloneChildProjectTrust, +} from "../shared/child-session.ts"; +import { formatContextUtilization } from "../shared/context-utilization.ts"; +import { + registerEditorLayer, + removeEditorLayer, +} from "../shared/editor-layers.ts"; +import { + PLAN_MODE_CHANNEL, + type PlanModeState, + planModeAllowsDeclaredTools, + planModeChildTools, +} from "../shared/plan-mode-state.ts"; +import { loadSetupConfig } from "../shared/setup-config.ts"; +import { + OPENPI_TOOL_SURFACE, + patchOwnedTools, +} from "../shared/tool-surface.ts"; +import { + createWorktree, + reclaimWorktree, + type Worktree, +} from "../shared/worktree.ts"; +import { + normalizeSubagentTitle, + SubagentStripWidget, + selectSubagentStripEntry, +} from "./navigation.ts"; +import { + type AgentType, formatAgentTypeDiagnostics, loadAgentTypes, roleModelForAgentType, selectSubagentModel, - type AgentType, } from "./src/agent-types.ts"; import { deriveBtwTitle, isModelVisible } from "./src/by-the-way.ts"; import { @@ -59,19 +97,11 @@ import { type SubagentSnapshot, } from "./src/domain.ts"; import { - formatActivityStatus, - hasActivity, - unreadActivityCounts, -} from "../shared/activity-status.ts"; -import { - OPENPI_TOOL_SURFACE, - patchOwnedTools, -} from "../shared/tool-surface.ts"; -import { - registerEditorLayer, - removeEditorLayer, -} from "../shared/editor-layers.ts"; -import { formatContextUtilization } from "../shared/context-utilization.ts"; + restoreSubagentIdCounters, + SUBAGENT_ID_WATERMARK_ENTRY_TYPE, + type SubagentIdCounters, + subagentIdWatermark, +} from "./src/id-sequence.ts"; import { SubagentManager, type SubagentManagerShape } from "./src/manager.ts"; import { buildSubagentResultMessage, @@ -90,43 +120,17 @@ import { SUBAGENT_WAIT_PARAMETER_DESCRIPTIONS, SUBAGENT_WAIT_TOOL_DESCRIPTION, } from "./src/prompt.ts"; -import { createSubagentResultDelivery } from "./src/result-delivery.ts"; -import { - effectiveChildToolAllowlist, - resolveStandaloneChildProjectTrust, -} from "../shared/child-session.ts"; +import { persistResultArtifact, projectResult } from "./src/result-artifact.ts"; import { - BelowEditorNavigationEditor, - BelowEditorStripState, -} from "../shared/below-editor-navigation.ts"; -import { loadSetupConfig } from "../shared/setup-config.ts"; -import { - PLAN_MODE_CHANNEL, - planModeAllowsDeclaredTools, - planModeChildTools, - type PlanModeState, -} from "../shared/plan-mode-state.ts"; -import { - createWorktree, - reclaimWorktree, - type Worktree, -} from "../shared/worktree.ts"; + allocateResultBudgets, + type ParentContextUsage, +} from "./src/result-budget.ts"; +import { createSubagentResultDelivery } from "./src/result-delivery.ts"; import { createSubagentRuntime, runTool, type SubagentRuntime, } from "./src/runtime.ts"; -import { - restoreSubagentIdCounters, - SUBAGENT_ID_WATERMARK_ENTRY_TYPE, - subagentIdWatermark, - type SubagentIdCounters, -} from "./src/id-sequence.ts"; -import { - normalizeSubagentTitle, - selectSubagentStripEntry, - SubagentStripWidget, -} from "./navigation.ts"; import { openSubagentPicker, openSubagentTakeover } from "./src/ui/takeover.ts"; import { buildWaitResultPreview, @@ -135,8 +139,13 @@ import { } from "./src/ui/wait-result.ts"; const SUBAGENT_OUTPUT_MAX_BYTES = 24 * 1024; +const AUTOMATIC_OUTPUT_MAX_BYTES = 48 * 1024; +const AUTOMATIC_MIN_RESULT_BYTES = 2 * 1024; const WAIT_OUTPUT_MAX_BYTES = 48 * 1024; const WAIT_PER_AGENT_MAX_BYTES = 16 * 1024; +const WAIT_MIN_RESULT_BYTES = 512; +const RESULT_HEADROOM_SHARE = 0.5; +const ESTIMATED_BYTES_PER_TOKEN = 4; interface SpawnResultDetails { readonly id?: string; @@ -190,36 +199,71 @@ function describeSubagent(snap: SubagentSnapshot) { return `${snap.id} [${snap.status}] "${snap.title}" (${details.join(", ")})`; } -function truncatedOutput( +export function truncatedOutput( snap: SubagentSnapshot, maxBytes = SUBAGENT_OUTPUT_MAX_BYTES, + writeArtifact: (content: string) => string = (content) => + persistResultArtifact(getAgentDir(), content), ): string { const output = snap.finalText || "(no output)"; - const truncation = truncateHead(output, { + return projectResult(output, { maxBytes: Math.min(maxBytes, DEFAULT_MAX_BYTES), maxLines: Math.min(600, DEFAULT_MAX_LINES), - }); - let text = truncation.content; - if (truncation.truncated) { - text += `\n\n[Output truncated: ${formatSize(truncation.outputBytes)} of ${formatSize(truncation.totalBytes)} shown. Full transcript in session file: ${snap.meta.sessionFilePath ?? "?"}]`; - } - return text; + writeArtifact, + }).text; } export function createSubagentResultDispatcher( pi: ExtensionAPI, - outputFor: (snap: SubagentSnapshot) => string = truncatedOutput, + outputFor: ( + snap: SubagentSnapshot, + maxBytes: number, + ) => string = truncatedOutput, + getContextUsage: () => ParentContextUsage | undefined = () => undefined, ) { return (snaps: readonly SubagentSnapshot[]) => { if (snaps.length === 0) return; + const emptyMessages = snaps.map((snap) => + buildSubagentResultMessage({ + id: snap.id, + title: snap.title, + status: snap.status, + errorText: snap.errorText, + output: "", + }), + ); + const wrapperBytes = + emptyMessages.reduce( + (sum, message) => sum + Buffer.byteLength(message, "utf8"), + 0, + ) + + Math.max(0, snaps.length - 1) * 2; + const projectionBatchBytes = Math.max( + AUTOMATIC_MIN_RESULT_BYTES * snaps.length, + AUTOMATIC_OUTPUT_MAX_BYTES - wrapperBytes, + ); + const allocation = allocateResultBudgets( + snaps.map((snap) => + Buffer.byteLength(snap.finalText || "(no output)", "utf8"), + ), + getContextUsage(), + { + maxBatchBytes: projectionBatchBytes, + maxResultBytes: SUBAGENT_OUTPUT_MAX_BYTES, + minResultBytes: AUTOMATIC_MIN_RESULT_BYTES, + headroomShare: RESULT_HEADROOM_SHARE, + estimatedBytesPerToken: ESTIMATED_BYTES_PER_TOKEN, + fixedBytes: wrapperBytes, + }, + ); const content = snaps - .map((snap) => + .map((snap, index) => buildSubagentResultMessage({ id: snap.id, title: snap.title, status: snap.status, errorText: snap.errorText, - output: outputFor(snap), + output: outputFor(snap, allocation.budgets[index]!), }), ) .join("\n\n"); @@ -326,7 +370,11 @@ export default function (pi: ExtensionAPI) { let requestWidgetRender: (() => void) | undefined; let navigationLayerRegistered = false; let dashboardOpen = false; - const dispatchResults = createSubagentResultDispatcher(pi); + const dispatchResults = createSubagentResultDispatcher( + pi, + truncatedOutput, + () => sessionContext?.getContextUsage(), + ); const resultDelivery = createSubagentResultDelivery({ isIdle: () => sessionContext?.isIdle() === true, // Every unconsumed fire-and-forget result must reach the parent. The @@ -815,7 +863,7 @@ export default function (pi: ExtensionAPI) { description: SUBAGENT_WAIT_PARAMETER_DESCRIPTIONS.ids, }), }), - async execute(_toolCallId, params, signal, onUpdate) { + async execute(_toolCallId, params, signal, onUpdate, ctx) { const manager = await getManager(); const ids = [...new Set(params.ids)]; if (ids.length === 0) @@ -851,33 +899,66 @@ export default function (pi: ExtensionAPI) { // deferred automatic delivery now that the tool is returning the result. resultDelivery.consume(ids); - const sections: string[] = []; - let remainingBytes = WAIT_OUTPUT_MAX_BYTES; - for (const id of ids) { + const entries: Array< + | { readonly id: string; readonly section: string } + | { + readonly id: string; + readonly snap: SubagentSnapshot; + readonly header: string; + } + > = ids.map((id) => { const snap = manager.view.get(id); - if (!snap) { - sections.push(`## ${id}\n\n(no longer tracked)`); - continue; - } + if (!snap) return { id, section: `## ${id}\n\n(no longer tracked)` }; const verb = snap.status === "error" ? "failed" : "finished"; - let section = `## ${snap.id} "${snap.title}" ${verb}`; - if (snap.errorText) section += `\nError: ${snap.errorText}`; - const headerBytes = Buffer.byteLength(section, "utf8") + 2; - const outputBudget = Math.max( - 512, - Math.min(WAIT_PER_AGENT_MAX_BYTES, remainingBytes - headerBytes), + let header = `## ${snap.id} "${snap.title}" ${verb}`; + if (snap.errorText) header += `\nError: ${snap.errorText}`; + return { id, snap, header }; + }); + const separatorsBytes = Math.max(0, entries.length - 1) * 7; + const fixedBytes = + separatorsBytes + + entries.reduce( + (sum, entry) => + sum + + Buffer.byteLength( + "section" in entry ? entry.section : `${entry.header}\n\n`, + "utf8", + ), + 0, ); - section += `\n\n${truncatedOutput(snap, outputBudget)}`; - const sectionBytes = Buffer.byteLength(section, "utf8"); - if (sectionBytes > remainingBytes) { - sections.push( - `## ${snap.id} "${snap.title}"\n\n[omitted: total wait output limit reached]`, - ); - break; - } - sections.push(section); - remainingBytes -= sectionBytes; - } + const resultEntries = entries.filter( + ( + entry, + ): entry is { + readonly id: string; + readonly snap: SubagentSnapshot; + readonly header: string; + } => "snap" in entry, + ); + const projectionBatchBytes = Math.max( + WAIT_MIN_RESULT_BYTES * resultEntries.length, + WAIT_OUTPUT_MAX_BYTES - fixedBytes, + ); + const allocation = allocateResultBudgets( + resultEntries.map(({ snap }) => + Buffer.byteLength(snap.finalText || "(no output)", "utf8"), + ), + ctx.getContextUsage(), + { + maxBatchBytes: projectionBatchBytes, + maxResultBytes: WAIT_PER_AGENT_MAX_BYTES, + minResultBytes: WAIT_MIN_RESULT_BYTES, + headroomShare: RESULT_HEADROOM_SHARE, + estimatedBytesPerToken: ESTIMATED_BYTES_PER_TOKEN, + fixedBytes, + }, + ); + let resultIndex = 0; + const sections = entries.map((entry) => { + if ("section" in entry) return entry.section; + const outputBudget = allocation.budgets[resultIndex++]!; + return `${entry.header}\n\n${truncatedOutput(entry.snap, outputBudget)}`; + }); const combined = sections.join("\n\n---\n\n"); const bounded = truncateHead(combined, { diff --git a/extensions/subagents/result-artifact.test.ts b/extensions/subagents/result-artifact.test.ts new file mode 100644 index 00000000..91f7187c --- /dev/null +++ b/extensions/subagents/result-artifact.test.ts @@ -0,0 +1,181 @@ +import assert from "node:assert/strict"; +import { + lstat, + mkdtemp, + readFile, + rm, + symlink, + writeFile, +} from "node:fs/promises"; +import { tmpdir } from "node:os"; +import path from "node:path"; +import test from "node:test"; +import { persistResultArtifact, projectResult } from "./src/result-artifact.ts"; + +test("short results pass through without creating an artifact", () => { + let writes = 0; + const result = projectResult("short report", { + maxBytes: 100, + maxLines: 10, + writeArtifact: () => { + writes++; + return "/unused"; + }, + }); + + assert.deepEqual(result, { text: "short report", truncated: false }); + assert.equal(writes, 0); +}); + +test("byte truncation keeps head and tail and points to the exact artifact", () => { + const content = `BEGIN\n${"middle-data\n".repeat(40)}FINAL-VERDICT`; + let persisted = ""; + const result = projectResult(content, { + maxBytes: 120, + maxLines: 100, + writeArtifact: (value) => { + persisted = value; + return "/tmp/final.txt"; + }, + }); + + assert.equal(result.truncated, true); + assert.equal(result.artifactPath, "/tmp/final.txt"); + assert.equal(persisted, content); + assert.match(result.text, /^BEGIN/); + assert.match(result.text, /FINAL-VERDICT/); + assert.match(result.text, /\[\.\.\. middle omitted \.\.\.\]/); + assert.match(result.text, /Full final answer: "\/tmp\/final\.txt"/); + assert.match(result.text, /offset=\d+, limit=200/); + assert.match(result.text, /\d+ total lines/); +}); + +test("the complete projection stays within its byte budget", () => { + const content = `BEGIN\n${"middle-data\n".repeat(400)}FINAL-VERDICT`; + const result = projectResult(content, { + maxBytes: 512, + maxLines: 100, + writeArtifact: () => "/tmp/final.txt", + }); + + assert.equal(result.truncated, true); + assert.ok(Buffer.byteLength(result.text, "utf8") <= 512); + assert.match(result.text, /^BEGIN/); + assert.match(result.text, /FINAL-VERDICT/); + assert.match(result.text, /Full final answer:/); +}); + +test("projection budgets remain hard caps across UTF-8 sizes", () => { + const content = `开头\n${"中间证据-abcdef\n".repeat(1000)}最终结论`; + for (const maxBytes of [512, 768, 1024, 2048, 16 * 1024]) { + const result = projectResult(content, { + maxBytes, + maxLines: 600, + writeArtifact: () => "/tmp/final.txt", + }); + assert.ok( + Buffer.byteLength(result.text, "utf8") <= maxBytes, + `projection exceeded ${maxBytes} bytes`, + ); + assert.match(result.text, /^开头/); + assert.match(result.text, /最终结论/); + } +}); + +test("line truncation keeps both ends even below the byte ceiling", () => { + const content = Array.from({ length: 20 }, (_, i) => `line-${i}`).join("\n"); + const result = projectResult(content, { + maxBytes: 10_000, + maxLines: 8, + writeArtifact: () => "/tmp/lines.txt", + }); + + assert.equal(result.truncated, true); + assert.match(result.text, /^line-0/); + assert.match(result.text, /line-19/); +}); + +test("a long UTF-8 line keeps valid characters at both ends", () => { + const content = `开头-${"中".repeat(100)}-结尾`; + const result = projectResult(content, { + maxBytes: 80, + maxLines: 10, + writeArtifact: () => "/tmp/chinese.txt", + }); + + assert.equal(result.truncated, true); + assert.match(result.text, /^开头-/); + assert.match(result.text, /-结尾/); + assert.doesNotMatch(result.text, /�/); +}); + +test("artifact failure is explicit and never advertises a false path", () => { + const result = projectResult("start\n" + "x\n".repeat(100) + "end", { + maxBytes: 80, + maxLines: 10, + writeArtifact: () => { + throw new Error("disk full"); + }, + }); + + assert.equal(result.truncated, true); + assert.equal(result.artifactPath, undefined); + assert.match(result.text, /could not be saved/); + assert.doesNotMatch(result.text, /Full final answer:/); +}); + +test("content-addressed artifacts are exact, private, and reusable", async () => { + const agentDir = await mkdtemp( + path.join(tmpdir(), "openpi-result-artifact-"), + ); + try { + const content = "complete final answer\nwith verdict"; + const first = persistResultArtifact(agentDir, content); + const second = persistResultArtifact(agentDir, content); + + assert.equal(first, second); + assert.equal(await readFile(first, "utf8"), content); + assert.equal((await lstat(first)).mode & 0o777, 0o600); + assert.equal(path.basename(first).length, 68); + } finally { + await rm(agentDir, { recursive: true, force: true }); + } +}); + +test("artifact persistence refuses a symlinked cache component", async () => { + const agentDir = await mkdtemp(path.join(tmpdir(), "openpi-result-symlink-")); + const outside = await mkdtemp(path.join(tmpdir(), "openpi-result-outside-")); + try { + await symlink(outside, path.join(agentDir, "cache")); + assert.throws( + () => persistResultArtifact(agentDir, "do not write outside"), + /Unsafe result artifact directory/, + ); + assert.deepEqual( + await readFile(path.join(outside, "sentinel"), "utf8").catch( + () => undefined, + ), + undefined, + ); + } finally { + await rm(agentDir, { recursive: true, force: true }); + await rm(outside, { recursive: true, force: true }); + } +}); + +test("artifact persistence refuses an existing file with the wrong content", async () => { + const agentDir = await mkdtemp( + path.join(tmpdir(), "openpi-result-collision-"), + ); + try { + const content = "original final answer"; + const artifactPath = persistResultArtifact(agentDir, content); + await writeFile(artifactPath, "tampered", "utf8"); + assert.throws( + () => persistResultArtifact(agentDir, content), + /Result artifact collision/, + ); + } finally { + await rm(agentDir, { recursive: true, force: true }); + } +}); diff --git a/extensions/subagents/result-budget.test.ts b/extensions/subagents/result-budget.test.ts new file mode 100644 index 00000000..1c67855a --- /dev/null +++ b/extensions/subagents/result-budget.test.ts @@ -0,0 +1,115 @@ +import assert from "node:assert/strict"; +import test from "node:test"; +import { + allocateResultBudgets, + type ResultBudgetPolicy, +} from "./src/result-budget.ts"; + +const policy: ResultBudgetPolicy = { + maxBatchBytes: 48 * 1024, + maxResultBytes: 24 * 1024, + minResultBytes: 2 * 1024, + headroomShare: 0.5, + estimatedBytesPerToken: 4, +}; + +test("unknown context usage falls back to the deterministic static caps", () => { + const allocation = allocateResultBudgets( + [40 * 1024, 40 * 1024], + undefined, + policy, + ); + + assert.deepEqual(allocation, { + budgets: [24 * 1024, 24 * 1024], + batchBytes: 48 * 1024, + source: "static", + }); +}); + +test("valid parent headroom dynamically shrinks a result batch", () => { + const allocation = allocateResultBudgets( + [40 * 1024, 40 * 1024], + { tokens: 98_000, contextWindow: 100_000 }, + policy, + ); + + assert.deepEqual(allocation, { + budgets: [2 * 1024, 2 * 1024], + batchBytes: 4 * 1024, + source: "dynamic", + }); +}); + +test("dynamic headroom includes fixed wrapper metadata", () => { + const allocation = allocateResultBudgets( + [40 * 1024, 40 * 1024], + { tokens: 96_000, contextWindow: 100_000 }, + { ...policy, fixedBytes: 4 * 1024 }, + ); + + assert.deepEqual(allocation.budgets, [2048, 2048]); + assert.equal(allocation.batchBytes, 4096); + assert.equal(allocation.source, "dynamic"); +}); + +test("short results yield their unused share to longer siblings", () => { + const allocation = allocateResultBudgets( + [1024, 40 * 1024, 40 * 1024], + undefined, + policy, + ); + + assert.deepEqual(allocation.budgets, [1024, 24_064, 24_064]); + assert.equal( + allocation.budgets.reduce((sum, budget) => sum + budget, 0), + 48 * 1024, + ); +}); + +test("dynamic batches preserve a readable floor for every result", () => { + const allocation = allocateResultBudgets( + [40 * 1024, 40 * 1024, 40 * 1024, 40 * 1024], + { tokens: 100_000, contextWindow: 100_000 }, + policy, + ); + + assert.deepEqual(allocation.budgets, [2048, 2048, 2048, 2048]); + assert.equal(allocation.batchBytes, 8192); + assert.equal(allocation.source, "dynamic"); +}); + +test("invalid or stale-looking usage never narrows the static fallback", () => { + for (const usage of [ + { tokens: null, contextWindow: 100_000 }, + { tokens: Number.NaN, contextWindow: 100_000 }, + { tokens: 1_000, contextWindow: 0 }, + ]) { + assert.equal( + allocateResultBudgets([40 * 1024], usage, policy).source, + "static", + ); + } +}); + +test("empty batches allocate nothing", () => { + assert.deepEqual(allocateResultBudgets([], undefined, policy), { + budgets: [], + batchBytes: 0, + source: "static", + }); +}); + +test("large batches keep their readable floors inside the static batch cap", () => { + const allocation = allocateResultBudgets( + Array.from({ length: 64 }, () => 40 * 1024), + { tokens: 100_000, contextWindow: 100_000 }, + policy, + ); + + assert.equal(allocation.batchBytes, 48 * 1024); + assert.equal( + allocation.budgets.reduce((sum, budget) => sum + budget, 0), + 48 * 1024, + ); +}); diff --git a/extensions/subagents/src/result-artifact.ts b/extensions/subagents/src/result-artifact.ts new file mode 100644 index 00000000..795a8bd0 --- /dev/null +++ b/extensions/subagents/src/result-artifact.ts @@ -0,0 +1,142 @@ +import { createHash } from "node:crypto"; +import { lstatSync, mkdirSync, readFileSync, writeFileSync } from "node:fs"; +import path from "node:path"; +import { + formatSize, + truncateHead, + truncateTail, +} from "@earendil-works/pi-coding-agent"; + +const HEAD_SHARE = 0.75; +const RESULT_ARTIFACT_DIR = ["cache", "openpi", "subagent-results"]; + +export interface ResultProjectionOptions { + readonly maxBytes: number; + readonly maxLines: number; + readonly writeArtifact: (content: string) => string; +} + +export interface ResultProjection { + readonly text: string; + readonly truncated: boolean; + readonly artifactPath?: string; +} + +function sliceStartToUtf8Bytes(content: string, maxBytes: number) { + const bytes = Buffer.from(content, "utf8"); + if (bytes.length <= maxBytes) return content; + let end = maxBytes; + while (end > 0 && (bytes[end] & 0xc0) === 0x80) end--; + return bytes.subarray(0, end).toString("utf8"); +} + +function ensureDirectory(parent: string, name: string) { + const directory = path.join(parent, name); + try { + mkdirSync(directory, { mode: 0o700 }); + } catch (error) { + if ((error as NodeJS.ErrnoException).code !== "EEXIST") throw error; + } + const stat = lstatSync(directory); + if (!stat.isDirectory() || stat.isSymbolicLink()) { + throw new Error(`Unsafe result artifact directory: ${directory}`); + } + return directory; +} + +/** + * Persist one immutable, content-addressed final answer below Pi's cache. + * Model-authored titles and paths never participate in the filename. + */ +export function persistResultArtifact(agentDir: string, content: string) { + let directory = path.resolve(agentDir); + for (const segment of RESULT_ARTIFACT_DIR) { + directory = ensureDirectory(directory, segment); + } + + const digest = createHash("sha256").update(content).digest("hex"); + const artifactPath = path.join(directory, `${digest}.txt`); + try { + writeFileSync(artifactPath, content, { + encoding: "utf8", + flag: "wx", + mode: 0o600, + }); + } catch (error) { + if ((error as NodeJS.ErrnoException).code !== "EEXIST") throw error; + const stat = lstatSync(artifactPath); + if ( + !stat.isFile() || + stat.isSymbolicLink() || + readFileSync(artifactPath, "utf8") !== content + ) { + throw new Error(`Result artifact collision: ${artifactPath}`); + } + } + return artifactPath; +} + +/** + * Build the single model-visible projection used by automatic delivery and + * explicit waits. Short answers pass through byte-for-byte. Long answers keep + * both decision context at the start and verdict/evidence at the end, while a + * plain-text artifact preserves the exact final answer for Pi's native read. + */ +export function projectResult( + content: string, + options: ResultProjectionOptions, +): ResultProjection { + const probe = truncateHead(content, { + maxBytes: options.maxBytes, + maxLines: options.maxLines, + }); + if (!probe.truncated) return { text: content, truncated: false }; + + const headLines = Math.max(1, Math.floor(options.maxLines * HEAD_SHARE)); + const tailLines = Math.max(1, options.maxLines - headLines); + + let artifactPath: string | undefined; + try { + artifactPath = options.writeArtifact(content); + } catch { + // Delivery is more important than the optional recovery cache. The footer + // below stays explicit so a failed write never advertises a false path. + } + + let bodyBudget = options.maxBytes; + let text = ""; + for (let attempt = 0; attempt < 8; attempt++) { + const headBytes = Math.max(1, Math.floor(bodyBudget * HEAD_SHARE)); + const tailBytes = Math.max(1, bodyBudget - headBytes); + const headResult = truncateHead(content, { + maxBytes: headBytes, + maxLines: headLines, + }); + const tailResult = truncateTail(content, { + maxBytes: tailBytes, + maxLines: tailLines, + }); + const head = + headResult.content || sliceStartToUtf8Bytes(content, headBytes); + const tail = tailResult.content; + const shownBytes = + Buffer.byteLength(head, "utf8") + Buffer.byteLength(tail, "utf8"); + const recovery = artifactPath + ? `Full final answer: ${JSON.stringify(artifactPath)}\nUse Pi's read tool with path=${JSON.stringify(artifactPath)}, offset=${Math.max(1, headResult.outputLines + 1)}, limit=200 to inspect the omitted middle; adjust offset to continue.` + : "Full final answer could not be saved; only the head and tail above are available."; + const footer = + `[Output truncated: showing ${formatSize(shownBytes)} of ${formatSize(probe.totalBytes)} ` + + `across the head and tail (${probe.totalLines} total lines).\n${recovery}]`; + text = `${head}\n\n[... middle omitted ...]\n\n${tail}\n\n${footer}`; + + const overflow = Buffer.byteLength(text, "utf8") - options.maxBytes; + if (overflow <= 0 || bodyBudget <= overflow + 2) break; + bodyBudget -= overflow; + } + + return { + text, + truncated: true, + ...(artifactPath ? { artifactPath } : {}), + }; +} diff --git a/extensions/subagents/src/result-budget.ts b/extensions/subagents/src/result-budget.ts new file mode 100644 index 00000000..687a06ac --- /dev/null +++ b/extensions/subagents/src/result-budget.ts @@ -0,0 +1,134 @@ +export interface ParentContextUsage { + readonly tokens: number | null; + readonly contextWindow: number; +} + +export interface ResultBudgetPolicy { + readonly maxBatchBytes: number; + readonly maxResultBytes: number; + readonly minResultBytes: number; + readonly headroomShare: number; + readonly estimatedBytesPerToken: number; + /** Batch metadata that consumes the same parent-context headroom. */ + readonly fixedBytes?: number; +} + +export interface ResultBudgetAllocation { + readonly budgets: readonly number[]; + readonly batchBytes: number; + readonly source: "static" | "dynamic"; +} + +function validUsage( + usage: ParentContextUsage | null | undefined, +): usage is { readonly tokens: number; readonly contextWindow: number } { + return Boolean( + usage && + typeof usage.tokens === "number" && + Number.isFinite(usage.tokens) && + usage.tokens >= 0 && + Number.isFinite(usage.contextWindow) && + usage.contextWindow > 0, + ); +} + +function distribute( + desired: readonly number[], + batchBytes: number, + minResultBytes: number, +) { + const budgets = desired.map((bytes) => Math.min(bytes, minResultBytes)); + let remaining = Math.max( + 0, + batchBytes - budgets.reduce((sum, bytes) => sum + bytes, 0), + ); + let active = desired + .map((bytes, index) => ({ bytes, index })) + .filter(({ bytes, index }) => bytes > budgets[index]!); + + while (remaining > 0 && active.length > 0) { + const share = Math.floor(remaining / active.length); + if (share === 0) { + for (const { bytes, index } of active) { + if (remaining === 0) break; + if (budgets[index]! >= bytes) continue; + budgets[index]! += 1; + remaining--; + } + break; + } + + const satisfied = active.filter( + ({ bytes, index }) => bytes - budgets[index]! <= share, + ); + if (satisfied.length > 0) { + for (const { bytes, index } of satisfied) { + const addition = bytes - budgets[index]!; + budgets[index]! = bytes; + remaining -= addition; + } + const settled = new Set(satisfied.map(({ index }) => index)); + active = active.filter(({ index }) => !settled.has(index)); + continue; + } + + for (const { index } of active) budgets[index]! += share; + remaining -= share * active.length; + } + + return budgets; +} + +/** + * Allocate one bounded projection budget per result. Small results yield their + * unused share to larger siblings. Parent context usage is optional: Pi marks + * it unknown after compaction until a fresh response, in which case the + * deterministic static batch cap remains the source of truth. + */ +export function allocateResultBudgets( + resultBytes: readonly number[], + usage: ParentContextUsage | null | undefined, + policy: ResultBudgetPolicy, +): ResultBudgetAllocation { + if (resultBytes.length === 0) { + return { budgets: [], batchBytes: 0, source: "static" }; + } + + const desired = resultBytes.map((bytes) => + Math.min(policy.maxResultBytes, Math.max(0, Math.floor(bytes))), + ); + const perResultFloor = Math.min( + policy.minResultBytes, + Math.floor(policy.maxBatchBytes / resultBytes.length), + ); + const desiredTotal = desired.reduce((sum, bytes) => sum + bytes, 0); + const staticBatchBytes = Math.min(policy.maxBatchBytes, desiredTotal); + const minimumBatchBytes = desired.reduce( + (sum, bytes) => sum + Math.min(bytes, perResultFloor), + 0, + ); + + let batchBytes = staticBatchBytes; + let source: ResultBudgetAllocation["source"] = "static"; + if (validUsage(usage)) { + const headroomTokens = Math.max(0, usage.contextWindow - usage.tokens); + const dynamicBytes = Math.max( + 0, + Math.floor( + headroomTokens * policy.headroomShare * policy.estimatedBytesPerToken, + ) - Math.max(0, policy.fixedBytes ?? 0), + ); + const narrowed = Math.max( + minimumBatchBytes, + Math.min(staticBatchBytes, dynamicBytes), + ); + if (narrowed < staticBatchBytes) source = "dynamic"; + batchBytes = narrowed; + } + + return { + budgets: distribute(desired, batchBytes, perResultFloor), + batchBytes, + source, + }; +}