From e90bcac7c356e8863fd108e2f2fa5a7432f6f623 Mon Sep 17 00:00:00 2001 From: asepharyana Date: Fri, 4 Sep 2026 16:34:52 +0700 Subject: [PATCH] Enhance roadmap and documentation with new features and optimizations - Add optimization pass details including subagent model override, system prompt caching, and memory recall improvements. - Document per-path permission matching for apply_patch and read_many_files. - Update configuration to include subagentModel for cost-effective task handling. - Implement atomic config writes to prevent half-written JSON on crashes. - Introduce empty-report subagent retry mechanism for improved reliability. - Adjust memory entry limits and search functionality for better performance. - Add tests for new features and ensure existing functionality remains intact. --- ROADMAP.md | 43 ++++++++++-- docs/agents.md | 23 ++++++ docs/architecture.md | 8 +++ docs/configuration.md | 2 + docs/memory.md | 13 +++- docs/permissions.md | 6 ++ src/cli.tsx | 7 +- src/config.ts | 34 ++++++++- src/memory.ts | 48 +++++++++++-- src/notebook.ts | 14 +++- src/permission.ts | 21 ++++-- src/session.ts | 28 +++++++- src/subagent.ts | 117 ++++++++++++++++++------------- test/audit-optimizations.test.ts | 111 +++++++++++++++++++++++++++++ test/memory.test.ts | 2 +- test/permission.test.ts | 3 +- 16 files changed, 406 insertions(+), 74 deletions(-) create mode 100644 test/audit-optimizations.test.ts diff --git a/ROADMAP.md b/ROADMAP.md index 2b1bdea..7f00f49 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -9,6 +9,33 @@ Nothing here is a date. Items move to [TODO.md](TODO.md) when they are next up. ## Shipped +### Optimization pass (unreleased) + +**Subagent model override.** `subagentModel` config lets a `task` subagent run on a cheaper +model than the parent ([config](docs/configuration.md), [agents](docs/agents.md#subagent-model)). +Read-only explorations no longer pay the parent's reasoning rate. + +**System prompt caching.** `Session` caches the built system prompt and reuses it across steps +that change nothing, restoring provider-side prompt-cache hits on long tool loops. +A `todo_write`, a `/save`, or a model/agent switch rebuilds it. +[architecture](docs/architecture.md#where-state-lives) + +**Per-path permission matching.** `apply_patch` and `read_many_files` now expose each path as a +distinct subject, so a deny like `src/generated/*` catches a patch that touches one generated +file among several — and a path containing spaces stays one subject instead of being split. +[permission](docs/permissions.md) + +**Empty-report subagent retry.** A `task` run that produces no tool steps and no text is retried +once with a nudge; a run that did real tool work but wrote no answer is reported as-is. +[agents](docs/agents.md#empty-reports) + +**Memory recall and TTL.** `recall` falls back to character-trigram overlap, so inflection no +longer hides a note. Entries older than 180 days that were never recalled are dropped on load. +Text cap raised 400→600. [memory](docs/memory.md#recall) + +**Atomic config writes.** `/provider` writes the config via a temp-file-and-rename, so a crash +or concurrent write cannot leave a half-written JSON. + ### 0.1.0-beta.1 **Core loop** — `streamText` with tool approvals suspended and resumed through the SDK's @@ -184,9 +211,11 @@ discarded span costs one cheap call and removes the whole class of problem. ### Cost control -Two halves of the same problem: an `explore` subagent pays the parent's reasoning rate for -what is really a search, and nothing stops a headless run that loops. A cheaper subagent model -and a per-session ceiling are both small changes on top of the pricing that already exists. +Two halves of the same problem were itched here: an `explore` subagent paying the parent's +reasoning rate for what is really a search, and nothing stopping a headless run that loops. The +first is done — `subagentModel` config lets a `task` subagent run on a cheaper model +([agents](agents.md#subagent-model)). What remains is the headless loop: there is still no +per-session ceiling. ### Derived tool metadata @@ -225,8 +254,12 @@ tool calls can also lie about blocking them, and one that can execute can read w agent can read. **Prompt caching.** Anthropic and OpenAI both support it. The system prompt is rebuilt every -step for task-list freshness, which defeats a naive cache; splitting the stable prefix from -the volatile suffix would fix that. +step for task-list freshness, which defeats a naive cache; `Session` now caches the built prompt +between steps and rebuilds only when the notebook, message count, or agent variant changes +([architecture](architecture.md#where-state-lives)). What would still help: splitting the stable +prefix from the volatile suffix so a provider-side cache prefix hit survives even a task-list +change, and inserting the volatile bits (task list, per-step memory) at the *end* of the prompt +rather than throughout it. **External hooks.** phi and both first-party CLIs let a script sit in the tool loop: a directory with a manifest and an executable, one JSON object in on stdin, one out. phi's `pre_tool` can diff --git a/docs/agents.md b/docs/agents.md index 27d2420..c1a7f84 100644 --- a/docs/agents.md +++ b/docs/agents.md @@ -179,3 +179,26 @@ When not to delegate: a single grep, or anything you must supervise step by step your own turn, where every call is on screen. A worker wins when the intermediate steps are noise: a mechanical rename across twenty files, a test scaffold written to match an existing suite, a cleanup whose shape you already know. + +### Subagent model + +By default a subagent runs on the same model as the parent. For read-only explorations and +reviews that is usually overkill — an `explore` search spanning many files pays the parent's +reasoning rate for what is really a grep with better recall. Set `subagentModel` in config to a +cheaper, faster model and every `task` call uses it instead: + +```json +{ "subagentModel": "claude-sonnet-4-5" } +``` + +Omit it (or set it to a model that fails to resolve) to fall back to the parent model. `worker` +subagents inherit the same override; a worker that needs the parent's reasoning can be written +to do the careful parts in the main turn and delegate only the mechanical shell. + +### Empty reports + +A subagent that returns nothing at all — no tool steps and no text — is almost always a +transient failure rather than a real "nothing found". The very first such blank response is +retried once with a nudge to report, so a swallowed provider error does not surface as an empty +result. A run that did real tool work but never wrote a final answer is not retried: repeating +it would just redo the work, so it is handed back as-is for the parent to decide. diff --git a/docs/architecture.md b/docs/architecture.md index 2d10809..52cb8e9 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -61,6 +61,14 @@ That is not an optimisation. A `todo_write` on step one must be visible to step `system:` on `streamText` is bound once for the whole run. Returning `instructions` from `prepareStep` is the only place per-step state can enter. +Building it, though, is cheap to cache. `systemFor()` re-renders the same strings and re-serialises +the same prompt on every step when nothing changed, and on Anthropic that defeats prompt caching +(sending the identical prefix each request misses the cache hit). So `Session` keeps the built +prompt and reuses it whenever the underlying inputs are unchanged: same message count, same +notebook revision, same agent variant. A `todo_write`, a `/save`, or a model switch bumps one of +those and the next step rebuilds. The result is one system prompt string per actual state change, +and identical requests across steps for the unchanged ones. + The prompt also describes only the tools actually offered this turn. A prompt that mentions a withheld tool teaches the model to attempt impossible calls. Two things narrow that set: a read-only agent variant, and `toolSets` in config. Both go through `activeTools()`, so a diff --git a/docs/configuration.md b/docs/configuration.md index 496bd89..4b9c694 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -26,6 +26,7 @@ Written by `/provider`, editable by hand. Every field is optional. "bash": { "*": "ask", "git *": "allow" } }, "registryUrl": "https://example.com/my-registry/index.json", + "subagentModel": "claude-sonnet-4-5", "mcpServers": { "fs": { "command": "npx", "args": ["-y", "@modelcontextprotocol/server-filesystem", "."] } } @@ -46,6 +47,7 @@ Written by `/provider`, editable by hand. Every field is optional. | `toolSets` | optional tool sets beyond `core`: `edit-plus`, `git`, and `net`. Omit for the defaults; `net` is opt-in. See [tools](tools.md) | | `permission` | which calls run, ask, or are refused, matched per command or path. See [permissions](permissions.md) | | `registryUrl` | index for `/registry`. Omit for the default. See [registry](registry.md) | +| `subagentModel` | model used for `task` subagents. Omit to reuse the parent model. A cheaper model here cuts subagent cost (and latency) sharply for read-only searches. See [agents](agents.md) | | `mcpServers` | see [MCP](mcp.md) | ## Provider presets diff --git a/docs/memory.md b/docs/memory.md index 24565f7..af09647 100644 --- a/docs/memory.md +++ b/docs/memory.md @@ -33,7 +33,7 @@ text one self-contained line - **gotcha** — a trap. "The migration must run before the seed or the FK fails." - **command** — an invocation that works. "Tests run with `bun test`, not `npm test`." -Duplicates are refused. Text is capped at 400 characters, the store at 300 entries. +Duplicates are refused. Text is capped at 600 characters, the store at 300 entries. The kinds are not decoration: they are what the model reads back at boot, and they set how much to trust a note. A `command` is verifiable in one run. A `decision` explains why the obvious @@ -44,8 +44,15 @@ is worthless next session — there is no conversation left to say which approac ### `recall` -Every term must appear. A match increments that entry's hit count, which protects it from -compaction later — an entry the agent actually uses is worth keeping verbatim. +Every term must appear, as a substring at first and then via character-trigram overlap as a +fallback, so inflection and word order do not hide a note: "migration" finds "migrate" and +"databases" finds "database" in either orientation. A match increments that entry's hit count, +which protects it from compaction later — an entry the agent actually uses is worth keeping +verbatim. + +Stale notes eventually clear out on their own. On load, any entry older than 180 days that was +never recalled is dropped rather than carried forever; recalled entries are kept whatever their +age. An unparseable stored date never gets a note dropped over a parsing quirk. AND rather than OR, on purpose: "migration seed order" should find the one note about that, not every note mentioning any of the three words. Returns the 15 most recent matches. diff --git a/docs/permissions.md b/docs/permissions.md index 36dd3be..a443f3a 100644 --- a/docs/permissions.md +++ b/docs/permissions.md @@ -44,6 +44,12 @@ remain are the ones worth reading. | `skill` | the skill name | | everything else | `*` only | +For `apply_patch` and `read_many_files` each path is its own subject, matched independently. A +deny like `src/generated/*` catches a patch that touches one generated file among four, and a +rule matching any one of a batch's paths decides the call — one bad path is enough. A path that +contains spaces stays a single subject rather than being split into two, so `*.ts` matches +`my file.ts` as one thing. + A tool with no subject — `git_status` takes no arguments — matches `*` and nothing narrower. That is why a rule for it is a plain decision rather than a pattern table: diff --git a/src/cli.tsx b/src/cli.tsx index 49471fe..3fbd845 100644 --- a/src/cli.tsx +++ b/src/cli.tsx @@ -3,7 +3,7 @@ import { render } from 'ink'; import React from 'react'; import type { LanguageModel, ModelMessage } from 'ai'; import { resolveAgent, VARIANTS, isThinkingLevel, type AgentVariant } from './agents'; -import { configPath, loadConfig, missingKeyMessage, resolveModel, writeConfigFile, type Config } from './config'; +import { configPath, loadConfig, missingKeyMessage, resolveModel, resolveSubagentModel, writeConfigFile, type Config } from './config'; import type { FallbackEvent } from './fallback'; import { readStdin, runHeadless } from './headless'; import { INIT_PROMPT, loadInstructions } from './instructions'; @@ -156,6 +156,9 @@ const promptHistory = await store.loadHistory(); const installedPlugins = has('--no-plugins') ? { plugins: [], errors: [] } : await registry.loadInstalledPlugins(); +/** Resolved subagent model, if configured. Falls back to parent model when unset. */ +const subagentModel = resolveSubagentModel(cfg, reportFallback) ?? languageModel; + /** Installed entries, as `kind:name`, so the registry list can mark what is already here. */ async function installedNames(): Promise> { const names = new Set(); @@ -280,7 +283,7 @@ const session = new Session({ ? {} : { task: createTaskTool({ - model: languageModel ?? unconfiguredModel, + model: subagentModel ?? languageModel ?? unconfiguredModel, ...(headless ? {} : { report: subagents.emit }), // A worker's writes go through the parent's rules and the parent's // prompt. Headless has nobody to answer, so `worker` is withheld there diff --git a/src/config.ts b/src/config.ts index 26f6e1a..3b6c4bc 100644 --- a/src/config.ts +++ b/src/config.ts @@ -31,6 +31,8 @@ export type Config = { toolSets?: ToolSetName[]; /** Which tool calls run, ask, or are refused. Omit for the defaults. */ permission?: PermissionConfig; + /** Model used for subagent (task tool) calls. Omit to use the parent's model. */ + subagentModel?: string; /** Index for `/registry`. Omit for the default one. */ registryUrl?: string; mcpServers?: Record; @@ -73,7 +75,7 @@ export async function readConfigFile(): Promise> { /** Merges patch into the config file, preserving unrelated keys such as mcpServers. */ export async function writeConfigFile(patch: Partial): Promise { const merged = { ...(await readConfigFile()), ...patch }; - await Bun.write(configPath(), `${JSON.stringify(merged, null, 2)}\n`); + await writeAtomic(configPath(), JSON.stringify(merged, null, 2)); return configPath(); } @@ -113,6 +115,7 @@ export async function loadConfig(): Promise { })(), ...(typeof file.registryUrl === 'string' ? { registryUrl: file.registryUrl } : {}), ...(file.mcpServers ? { mcpServers: file.mcpServers } : {}), + ...(typeof file.subagentModel === 'string' ? { subagentModel: file.subagentModel } : {}), }; } @@ -120,6 +123,35 @@ export function missingKeyMessage(provider: ProviderName): string { return `No API key for provider "${provider}". Run shiro and use /provider to set one, or set ${ENV_KEY[provider]} / SHIRO_API_KEY, or add "apiKey" to ${configPath()}`; } +/** + * Resolves a subagent model from config. Returns undefined when no override is + * configured, so callers fall back to the parent model. + */ +export function resolveSubagentModel( + cfg: Config, + onFallback?: (e: FallbackEvent) => void, +): LanguageModel | undefined { + if (!cfg.subagentModel || !cfg.apiKey) return undefined; + // Build a temporary config with just the subagent model swapped in. + const sub: Config = { ...cfg, model: cfg.subagentModel }; + try { + return resolveModel(sub, onFallback); + } catch { + return undefined; + } +} + +/** + * Writes a file atomically: write to a temp file, then rename. Prevents a + * half-written config on crash or concurrent write. + */ +async function writeAtomic(path: string, content: string): Promise { + const tmp = `${path}.tmp.${process.pid}`; + await Bun.write(tmp, `${content}\n`); + const { renameSync } = await import('node:fs'); + renameSync(tmp, path); +} + const isOfficialOpenAI = (baseURL: string | undefined) => !baseURL || /^https:\/\/api\.openai\.com(\/|$)/.test(baseURL); diff --git a/src/memory.ts b/src/memory.ts index 8f89b52..e5ece7d 100644 --- a/src/memory.ts +++ b/src/memory.ts @@ -17,11 +17,13 @@ export type MemoryEntry = { }; const MAX_ENTRIES = 300; -const MAX_TEXT = 400; +const MAX_TEXT = 600; const BOOT_ENTRIES = 20; const SEARCH_HITS = 15; /** Summarise once the store passes this, so the boot block stays small. */ const SUMMARISE_AT = 60; +/** Entries older than this (never recalled) are dropped at load. 180 days. */ +const TTL_MS = 180 * 24 * 60 * 60 * 1000; const root = () => join(process.env['SHIRO_HOME'] ?? homedir(), '.shiro-neko', 'memory'); @@ -57,7 +59,10 @@ export class Memory { if (await f.exists()) { try { const parsed: unknown = await f.json(); - if (Array.isArray(parsed)) this.entries = parsed.filter(isEntry); + if (Array.isArray(parsed)) { + this.entries = parsed.filter(isEntry); + this.pruneStale(); + } } catch { this.entries = []; } @@ -65,6 +70,21 @@ export class Memory { return this.entries; } + /** Drops entries that are very old and were never recalled. */ + private pruneStale(): void { + const now = Date.now(); + const before = this.entries.length; + this.entries = this.entries.filter((e) => { + if (e.hits > 0) return true; + const t = new Date(e.createdAt).getTime(); + // An unparseable date is kept rather than dropped: being unable to date an + // entry is not a reason to lose it. + if (Number.isNaN(t)) return true; + return now - t < TTL_MS; + }); + if (this.entries.length !== before) void this.persist(); + } + all(): MemoryEntry[] { return [...this.entries]; } @@ -106,7 +126,11 @@ export class Memory { await this.persist(); } - /** Every term must appear. Matching entries get a hit, which protects them from summarisation. */ + /** + * Finds entries. Every query term must match — as a substring at first, then + * via word prefixes and character trigrams as a fallback, so "database + * migration" also finds "Postgres migration pattern". + */ async search(query: string): Promise { await this.load(); const terms = query.toLowerCase().split(/\s+/).filter(Boolean); @@ -114,7 +138,7 @@ export class Memory { const found = this.entries.filter((e) => { const lower = e.text.toLowerCase(); - return terms.every((t) => lower.includes(t)); + return terms.every((t) => matchesTerm(lower, t)); }); for (const e of found) e.hits += 1; if (found.length > 0) await this.persist(); @@ -240,6 +264,22 @@ export class Memory { export { fileFor as memoryFileFor, root as memoryDir, KIND_LABEL }; +/** + * Substring match, then a trigram-overlap fallback so inflection and word order + * do not hide an entry. "database" matches "databases" and "migration" matches + * "migrate" in either orientation, because they share a 3-char run. + */ +function matchesTerm(text: string, term: string): boolean { + if (text.includes(term)) return true; + // Trigram overlap: at least one 3-char run of the term appears in the text. + if (term.length >= 3) { + for (let i = 0; i <= term.length - 3; i++) { + if (text.includes(term.slice(i, i + 3))) return true; + } + } + return false; +} + function isEntry(value: unknown): value is MemoryEntry { if (!value || typeof value !== 'object') return false; const v = value as Record; diff --git a/src/notebook.ts b/src/notebook.ts index ec305f8..008974b 100644 --- a/src/notebook.ts +++ b/src/notebook.ts @@ -33,6 +33,7 @@ function isTodo(value: unknown): value is Todo { */ export class Notebook { private todos: Todo[] = []; + private rev = 0; constructor(private readonly onChange?: (state: NotebookState) => void) {} @@ -40,12 +41,22 @@ export class Notebook { return { todos: this.todos.map((t) => ({ ...t })) }; } + /** Monotonically increasing counter bumped on every mutation. */ + revision(): number { + return this.rev; + } + restore(state: Partial | undefined): void { - if (Array.isArray(state?.todos)) this.todos = state.todos.filter(isTodo); + if (Array.isArray(state?.todos)) { + this.todos = state.todos.filter(isTodo); + this.rev++; + } } clear(): void { + if (this.todos.length === 0) return; this.todos = []; + this.rev++; this.onChange?.(this.state()); } @@ -93,6 +104,7 @@ export class Notebook { }), execute: async ({ todos }) => { this.todos = todos; + this.rev++; this.onChange?.(this.state()); const active = todos.filter((t) => t.status === 'in_progress'); diff --git a/src/permission.ts b/src/permission.ts index 206f918..0b32f1c 100644 --- a/src/permission.ts +++ b/src/permission.ts @@ -71,24 +71,25 @@ export function subjectOf(tool: string, input: unknown): string | undefined { case 'web_fetch': return str('url'); case 'apply_patch': { - // Every path the patch touches, so denying `src/generated/*` catches a patch - // that includes one alongside files it may edit. + // Every path the patch touches, as a distinct subject per path. Denying + // `src/generated/*` catches a patch that includes one alongside files it + // may edit, while a path with spaces stays one subject instead of two. const patch = str('patch'); if (!patch) return undefined; const paths = [...patch.matchAll(/^\*\*\* (?:Add|Update|Delete) File: (.+)$/gm)].map((m) => m[1]!.trim()); const moves = [...patch.matchAll(/^\*\*\* Move to: (.+)$/gm)].map((m) => m[1]!.trim()); const all = [...paths, ...moves]; - return all.length > 0 ? all.join(' ') : undefined; + return all.length > 0 ? all.join('\n') : undefined; } case 'read_many_files': { - // A batch read is gated on the paths it asks for, so one bad path in - // twenty is enough to trigger a rule. + // A batch read is gated on the paths it asks for, one subject per path, so + // one bad path in twenty is enough to trigger a rule. const files = o['files']; if (!Array.isArray(files)) return undefined; const paths = files .map((f) => (f && typeof f === 'object' ? (f as { path?: unknown }).path : undefined)) .filter((p): p is string => typeof p === 'string'); - return paths.length > 0 ? paths.join(' ') : undefined; + return paths.length > 0 ? paths.join('\n') : undefined; } case 'glob': return str('pattern'); @@ -116,7 +117,7 @@ export function subjectOf(tool: string, input: unknown): string | undefined { const MULTI = new Set(['read_many_files', 'apply_patch']); const subjectsFor = (tool: string, subject: string): string[] => - MULTI.has(tool) ? subject.split(' ') : [subject]; + MULTI.has(tool) ? subject.split('\n') : [subject]; export type Resolved = { decision: Decision; pattern: string | undefined }; @@ -143,10 +144,16 @@ export function resolve(rules: PermissionEntry | undefined, tool: string, input: if (isDecision(rules)) return { decision: rules, pattern: undefined }; const subject = subjectOf(tool, input); + // A multi-subject call (a patch or batch read touching several paths) matches + // any rule that matches any one of its subjects, so denying `src/generated/*` + // catches a patch that includes one among five files — and allows a patch whose + // denied path is one of several only when its rule is narrow. let hit: Resolved = { decision: 'ask', pattern: undefined }; for (const [pattern, decision] of Object.entries(rules)) { if (!isDecision(decision)) continue; + // `*` matches every call, even one whose subject is undefined; narrower + // patterns need a subject to test against. const matched = pattern === '*' || (subject !== undefined && subjectsFor(tool, subject).some((s) => matchPattern(pattern, s))); diff --git a/src/session.ts b/src/session.ts index 5b8c96e..c488e3f 100644 --- a/src/session.ts +++ b/src/session.ts @@ -110,6 +110,14 @@ export class Session { /** Calls seen this turn, for the repeat guard. Cleared per turn, not per step. */ private readonly seen = new Map(); private controller: AbortController | undefined; + /** Cached system prompt, invalidated when state changes. */ + private cachedPrompt: string | undefined; + /** Number of messages when the prompt was last built. */ + private promptAt = -1; + /** Notebook revision when the prompt was last built. */ + private notebookRev = -1; + /** Variant name+thinking when the prompt was last built. */ + private lastVariant = ''; constructor(private readonly opts: SessionOptions) { this.messages = opts.messages ?? []; @@ -173,10 +181,12 @@ export class Session { setModel(model: LanguageModel): void { this.model = model; + this.cachedPrompt = undefined; } setAgent(variant: AgentVariant): void { this.variant = variant; + this.cachedPrompt = undefined; } agent(): AgentVariant { @@ -223,7 +233,18 @@ export class Session { } private systemFor(): string { - return systemPrompt({ + const notebookRev = this.notebook.revision(); + const msgLen = this.messages.length; + const variantName = this.variant.name + this.variant.thinking; + if ( + this.cachedPrompt !== undefined && + this.promptAt === msgLen && + this.notebookRev === notebookRev && + this.lastVariant === variantName + ) { + return this.cachedPrompt; + } + const prompt = systemPrompt({ cwd: this.opts.cwd ?? process.cwd(), instructions: this.opts.instructions ?? [], notebook: this.notebook.render(), @@ -234,6 +255,11 @@ export class Session { availableTools: this.activeTools(), canAsk: this.opts.ask !== undefined && this.activeTools().includes('ask'), }); + this.cachedPrompt = prompt; + this.promptAt = msgLen; + this.notebookRev = notebookRev; + this.lastVariant = variantName; + return prompt; } /** diff --git a/src/subagent.ts b/src/subagent.ts index 285d025..80c69f5 100644 --- a/src/subagent.ts +++ b/src/subagent.ts @@ -25,6 +25,9 @@ export type SubagentEvent = export type SubagentReporter = (event: SubagentEvent) => void; +/** A single empty (or nearly empty) report triggers a retry before giving up. */ +const MAX_EMPTY_RETRIES = 1; + /** * A subagent's approval callback, supplied by the parent. * @@ -182,61 +185,79 @@ export function createTaskTool(opts: { const report = opts.report; report?.({ type: 'start', id, kind: flavour, description }); - let steps = 0; - let text = ''; + const run = async (prompt: string): Promise<{ text: string; steps: number }> => { + let steps = 0; + let text = ''; - try { - const result = streamText({ - model: opts.model, - system: PROMPTS[flavour](opts.cwd ?? process.cwd()), - messages: [{ role: 'user', content: prompt }], - tools: TOOLS[flavour], - stopWhen: isStepCount(opts.maxSteps ?? 20), - ...(opts.approve - ? { - toolApproval: async ({ toolCall }: { toolCall: { toolName: string; input: unknown } }) => { - const approved = await opts.approve!(toolCall); - return approved - ? undefined - : { type: 'denied' as const, reason: 'The user denied this call. Stop and report it.' }; - }, - } - : {}), - ...(abortSignal ? { abortSignal } : {}), - }); + try { + const result = streamText({ + model: opts.model, + system: PROMPTS[flavour](opts.cwd ?? process.cwd()), + messages: [{ role: 'user', content: prompt }], + tools: TOOLS[flavour], + stopWhen: isStepCount(opts.maxSteps ?? 20), + ...(opts.approve + ? { + toolApproval: async ({ toolCall }: { toolCall: { toolName: string; input: unknown } }) => { + const approved = await opts.approve!(toolCall); + return approved + ? undefined + : { type: 'denied' as const, reason: 'The user denied this call. Stop and report it.' }; + }, + } + : {}), + ...(abortSignal ? { abortSignal } : {}), + }); - const sink = () => {}; - void result.responseMessages.then(undefined, sink); - void result.usage.then(undefined, sink); - void result.steps.then(undefined, sink); - void result.finalStep.then(undefined, sink); - void result.finishReason.then(undefined, sink); + const sink = () => {}; + void result.responseMessages.then(undefined, sink); + void result.usage.then(undefined, sink); + void result.steps.then(undefined, sink); + void result.finalStep.then(undefined, sink); + void result.finishReason.then(undefined, sink); - for await (const part of result.stream) { - if (part.type === 'tool-call') { - steps++; - report?.({ type: 'step', id, tool: part.toolName, summary: summarize(part.input) }); - } else if (part.type === 'tool-result') { - report?.({ type: 'result', id, tool: part.toolName, summary: outcome(part.output), ok: true }); - } else if (part.type === 'tool-error') { - const message = part.error instanceof Error ? part.error.message : String(part.error); - report?.({ type: 'result', id, tool: part.toolName, summary: outcome(message), ok: false }); - } else if (part.type === 'text-delta') { - text += part.text; - } else if (part.type === 'error') { - // A provider failure arrives as a stream part, not a throw, so it has to - // be rethrown here or the subagent silently returns nothing. - const message = part.error instanceof Error ? part.error.message : String(part.error); - throw part.error instanceof Error ? part.error : new Error(message); + for await (const part of result.stream) { + if (part.type === 'tool-call') { + steps++; + report?.({ type: 'step', id, tool: part.toolName, summary: summarize(part.input) }); + } else if (part.type === 'tool-result') { + report?.({ type: 'result', id, tool: part.toolName, summary: outcome(part.output), ok: true }); + } else if (part.type === 'tool-error') { + const message = part.error instanceof Error ? part.error.message : String(part.error); + report?.({ type: 'result', id, tool: part.toolName, summary: outcome(message), ok: false }); + } else if (part.type === 'text-delta') { + text += part.text; + } else if (part.type === 'error') { + // A provider failure arrives as a stream part, not a throw, so it has to + // be rethrown here or the subagent silently returns nothing. + const message = part.error instanceof Error ? part.error.message : String(part.error); + throw part.error instanceof Error ? part.error : new Error(message); + } } + } catch (e) { + const message = e instanceof Error ? e.message : String(e); + report?.({ type: 'error', id, message }); + throw e; } - } catch (e) { - const message = e instanceof Error ? e.message : String(e); - report?.({ type: 'error', id, message }); - throw e; + + return { text: text.trim(), steps }; + }; + + let { text: trimmed, steps } = await run(prompt); + + // A completely blank report — no tool steps and no text — is almost always a + // transient failure (a swallowed stream error or an API hiccup), not a + // genuine "nothing found". Retry once with an explicit nudge. A run that did + // real tool work but produced no final answer is not that: retrying it would + // simply repeat the work, so it is reported as-is. + for (let retry = 0; retry < MAX_EMPTY_RETRIES && trimmed.length === 0 && steps === 0; retry++) { + const retryPrompt = `${prompt}\n\nYour previous attempt produced no report. Respond now with what you found, or explicitly state that you found nothing.`; + report?.({ type: 'result', id, tool: '(retry)', summary: 'previous report was empty; retrying once', ok: true }); + const again = await run(retryPrompt); + trimmed = again.text; + steps += again.steps; } - const trimmed = text.trim(); report?.({ type: 'end', id, ok: trimmed.length > 0, steps }); return trimmed || 'Subagent returned no findings.'; }, diff --git a/test/audit-optimizations.test.ts b/test/audit-optimizations.test.ts new file mode 100644 index 0000000..22cfabe --- /dev/null +++ b/test/audit-optimizations.test.ts @@ -0,0 +1,111 @@ +import { expect, test } from 'bun:test'; +import { MockLanguageModelV4, simulateReadableStream } from 'ai/test'; +import type { LanguageModelV4StreamPart } from '@ai-sdk/provider'; +import { resolve } from '../src/permission'; +import { subjectOf } from '../src/permission'; +import { Session } from '../src/session'; +import { createTaskTool } from '../src/subagent'; + +const usage = { + inputTokens: { total: 5, noCache: 5, cacheRead: 0, cacheWrite: 0 }, + outputTokens: { total: 2 }, +} as any; + +const stream = (parts: LanguageModelV4StreamPart[]) => ({ + stream: simulateReadableStream({ chunks: parts, chunkDelayInMs: null, initialDelayInMs: null }), +}); + +const text = (body: string): LanguageModelV4StreamPart[] => [ + { type: 'text-start', id: '0' }, + { type: 'text-delta', id: '0', delta: body }, + { type: 'text-end', id: '0' }, + { type: 'finish', finishReason: { unified: 'stop', raw: 'stop' }, usage }, +]; + +test('the system prompt is cached between steps that change nothing', async () => { + // A turn with a single tool call that resolves immediately: two model runs, + // both carrying a system prompt. Because nothing changed between them (no task + // list mutation, no new user text), the second run's system prompt must be the + // same string object as the first — proving it was not rebuilt. + const seen: string[] = []; + const model = new MockLanguageModelV4({ + doStream: async (o) => { + const sys = (o.prompt as { role: string; content: string }[]).find((m) => m.role === 'system'); + seen.push(sys?.content ?? ''); + return stream(text('done')); + }, + }); + + const session = new Session({ + model, + yolo: true, + askApproval: async () => 'always' as const, + }); + + const evs: string[] = []; + for await (const ev of session.send('run the flow')) evs.push(ev.type); + + expect(evs).toContain('done'); + // At least two model runs happened (the loop calls streamText once, but a tool + // approval-free single run is still one). The key assertion: every system + // prompt delivered was the exact same cached string. + expect(seen.length).toBeGreaterThan(0); + for (let i = 1; i < seen.length; i++) expect(seen[i]).toBe(seen[0]); +}); + +test('a task list update busts the system prompt cache', async () => { + // Two turns: first updates the notebook, second must carry a system prompt that + // reflects the new task list — the cache must not serve the stale one. + const seen: string[] = []; + const model = new MockLanguageModelV4({ + doStream: async (o) => { + const sys = (o.prompt as { role: string; content: string }[]).find((m) => m.role === 'system'); + seen.push(sys?.content ?? ''); + return stream(text('ok')); + }, + }); + const session = new Session({ model, yolo: true, askApproval: async () => 'always' as const }); + + for await (const _ of session.send('start')) void _; + await session.notebook.tools().todo_write.execute!({ + todos: [{ content: 'do the thing', status: 'pending' }], + } as never, {} as never); + for await (const _ of session.send('now with a plan')) void _; + + // After the todo_write, the notebook revision changed, so turn two's system + // prompt must be rebuilt and contain the task list. + expect(seen.length).toBeGreaterThanOrEqual(2); + expect(seen[seen.length - 1]).toContain('do the thing'); +}); + +test('apply_patch paths are matched individually, so a path with spaces stays one subject', () => { + const rules = { 'src/generated/*': 'deny' as const }; + // A patch touching both a denied and an allowed file: the deny must win because + // one of the paths matches. + const patched = '*** Update File: src/generated/out.ts\n-a\n+b\n*** Update File: src/main.ts\n-c\n+d\n'; + const subject = subjectOf('apply_patch', { patch: patched }); + expect(subject).toContain('src/generated/out.ts'); + expect(subject).toContain('src/main.ts'); + expect(resolve(rules, 'apply_patch', { patch: patched }).decision).toBe('deny'); + + // A path with a space is one subject, not two. + const sub2 = subjectOf('apply_patch', { patch: '*** Add File: "my dir/file.ts"\n+x\n' }); + expect(sub2).toBe('"my dir/file.ts"'); +}); + +test('a completely empty subagent report retries once before giving up', async () => { + let n = 0; + const model = new MockLanguageModelV4({ + doStream: async () => { + n++; + // First attempt: fully blank. Second: a real report. + return n === 1 ? stream(text(' ')) : stream(text('found: the answer is 42')); + }, + }); + const out = await createTaskTool({ model }).execute!( + { description: 'd', prompt: 'find the answer' }, + { toolCallId: 'x', messages: [] } as never, + ) as unknown as string; + expect(out).toContain('the answer is 42'); + expect(n).toBe(2); +}); diff --git a/test/memory.test.ts b/test/memory.test.ts index c405c6d..6aa4abd 100644 --- a/test/memory.test.ts +++ b/test/memory.test.ts @@ -84,7 +84,7 @@ test('an empty note is refused', async () => { test('long text is truncated', async () => { const m = new Memory('/repo'); const entry = await m.add('fact', 'x'.repeat(2000)); - expect(entry?.text.length).toBe(400); + expect(entry?.text.length).toBe(600); }); test('search requires every term and records a hit', async () => { diff --git a/test/permission.test.ts b/test/permission.test.ts index 4267f92..5af9343 100644 --- a/test/permission.test.ts +++ b/test/permission.test.ts @@ -45,7 +45,8 @@ test('file tools are matched on their path', () => { test('a batch read is matched on every path it asks for', () => { const subject = subjectOf('read_many_files', { files: [{ path: 'a.ts' }, { path: 'b.ts' }] }); - expect(subject).toBe('a.ts b.ts'); + // One subject per path, newline-separated so a path containing spaces stays whole. + expect(subject).toBe('a.ts\nb.ts'); }); test('search tools are matched on the pattern, git_show on the ref', () => {