Repository Public
1/**2 * One turn's tool batch (CONTRACT §5.3): validate → policy → execute, each execution a journaled3 * `tool_call` step. Consecutive parallel-safe reads that policy allows run concurrently: their steps4 * are allocated in call order BEFORE anything is awaited, so step seqs are identical on replay.5 * Everything else runs sequentially, in the order the model issued the calls.6 */7import { z } from 'zod';8import { isControlError, putArtifact, type StepSpec } from '@aos/kernel';9import { invalidToolInputMessage, parseInvalidToolInput } from '@aos/providers';10import type { ContentBlock, EffectClass, ResultEnvelope, RiskClass, ToolResult, Trust } from '@aos/shared';11import {12 applyFileStateDelta, effectOf, finalizeToolOutput, isControlTool, parseControl, riskOf,13 type ControlPayload, type FileState, type FileStateDelta, type FinishPayload, type Tool, type ToolContext, type ToolRunResult,14} from '@aos/tools';15import { applyVerdict, emitDecision, verdictFor, type GateInput } from './approval.js';16import { runControl } from './control.js';17import { loadPolicySnapshot, type PolicySnapshot, type PolicyVerdict } from './policy.js';18import { sandboxDenyRead } from './sandbox.js';19import { toolResult, type RunState } from './state.js';2021export type ToolCall = Extract<ContentBlock, { type: 'tool_call' }>;2223/** What a `tool_call` step records: the model-facing result plus what the engine and the console need. */24export interface ToolStepResult {25 ok: boolean;26 content: string;27 artifact: { name: string; sha256: string; size: number } | null;28 trust: Trust;29 fileStateOut?: FileStateDelta;30 diff?: string;31}3233export interface BatchOutcome {34 results: ToolResult[];35 finish: { envelope: ResultEnvelope; payload: FinishPayload } | null;36}3738type Prepared =39 | { kind: 'error'; call: ToolCall; message: string }40 | { kind: 'control'; call: ToolCall; payload: ControlPayload }41 | { kind: 'tool'; call: ToolCall; tool: Tool; input: unknown; effect: EffectClass; risk: RiskClass; gate: GateInput; verdict: PolicyVerdict };4243const zodMessage = (e: z.ZodError): string => e.issues.map((i) => `${i.path.join('.') || 'input'}: ${i.message}`).join('; ');44const ASSUMED = 'The owner confirmed that this call had already run before a crash; its original output was lost. Inspect the current state before relying on it.';4546/** Validation and policy verdict of one call. Pure (no kernel call), so it can be evaluated ahead of execution. */47function prepare(st: RunState, call: ToolCall, snapshot: PolicySnapshot): Prepared {48 const tool = st.tools.get(call.name);49 if (!tool) return { kind: 'error', call, message: `Unknown tool "${call.name}". Available tools: ${[...st.tools.keys()].join(', ')}.` };50 if (parseInvalidToolInput(call.input)) return { kind: 'error', call, message: invalidToolInputMessage(call.input) ?? 'Invalid tool input.' };51 if (st.finalTurn && call.name !== 'finish') return { kind: 'error', call, message: 'Not executed: the budget is exhausted. Only `finish` is accepted now.' };52 let input: unknown;53 try {54 input = tool.input.parse(call.input);55 } catch (e) {56 if (!(e instanceof z.ZodError)) throw e;57 return { kind: 'error', call, message: `Invalid input for ${call.name}: ${zodMessage(e)}. Nothing was executed; fix the arguments and call it again.` };58 }59 if (isControlTool(call.name)) {60 try {61 return { kind: 'control', call, payload: parseControl(call.name, input) };62 } catch (e) {63 return { kind: 'error', call, message: `Invalid input for ${call.name}: ${e instanceof z.ZodError ? zodMessage(e) : String(e)}.` };64 }65 }66 const net = { network: st.meta.network };67 const risk = riskOf(tool, input, net);68 const gate: GateInput = { call, tool, input, risk };69 return { kind: 'tool', call, tool, input, effect: effectOf(tool, input, net), risk, gate, verdict: verdictFor(st, gate, snapshot) };70}7172async function toolContext(st: RunState, tool: Tool, idemKey: string, signal: AbortSignal, fileState: FileState): Promise<ToolContext> {73 const { ctx, deps } = st;74 const tctx: ToolContext = {75 runId: ctx.runId,76 missionId: ctx.missionId,77 workspaceAbs: ctx.workspaceAbs,78 network: st.meta.network,79 dataDir: deps.kernel.dataDir,80 signal,81 idemKey,82 fileState,83 putArtifact: (name, data, mime) => putArtifact(data, { name, mime, runId: ctx.runId }, deps.kernel),84 };85 if (tool.def.name === 'shell') {86 tctx.sandbox = {87 sandboxExecPath: deps.sandboxExecPath,88 extraDenyRead: await sandboxDenyRead({ workspaceAbs: ctx.workspaceAbs, workspaceRoot: deps.workspaceRoot, extra: deps.extraDenyRead, ownerFolders: deps.denyOwnerFolders }),89 };90 }91 return tctx;92}9394async function finalize(name: string, res: ToolRunResult, tctx: ToolContext): Promise<ToolStepResult> {95 const fin = await finalizeToolOutput(name, res.output, tctx);96 const out: ToolStepResult = { ok: res.ok, content: fin.content, artifact: fin.artifact, trust: res.trust ?? 'tool' };97 if (res.fileStateOut) out.fileStateOut = res.fileStateOut;98 if (res.diff) out.diff = res.diff;99 return out;100}101102/**103 * Starts the journaled step of one tool call. NOT async on purpose: `ctx.step` is invoked synchronously104 * so that a caller mapping over several calls allocates their seqs in call order.105 */106function startToolStep(st: RunState, p: Extract<Prepared, { kind: 'tool' }>, fileState: FileState): Promise<ToolStepResult> {107 const { tool, input, call } = p;108 const request = { name: call.name, input: call.input };109 const spec: StepSpec<typeof request, ToolStepResult> = {110 kind: 'tool_call', name: call.name, effect: p.effect, request,111 assumedResult: () => ({ ok: true, content: ASSUMED, artifact: null, trust: 'system' }),112 };113 if (tool.reconcile) {114 const probe = tool.reconcile.bind(tool);115 spec.reconcile = async (_req, idemKey) => {116 const tctx = await toolContext(st, tool, idemKey, st.ctx.signal, fileState);117 const r = await probe(input, tctx);118 if ('done' in r && r.done) return { done: true, result: await finalize(call.name, { ok: true, output: r.output }, tctx) };119 return 'done' in r ? { done: false } : { unknown: true };120 };121 }122 return st.ctx.step(spec, async (idemKey, signal) => {123 const tctx = await toolContext(st, tool, idemKey, signal, fileState);124 let res: ToolRunResult;125 try {126 res = await tool.run(input, tctx);127 } catch (e) {128 if (signal.aborted || isControlError(e)) throw e;129 res = { ok: false, output: `Tool ${call.name} failed unexpectedly: ${e instanceof Error ? e.message : String(e)}`.slice(0, 4000) };130 }131 // An execution interrupted by cancel / shutdown / lost lease is not a result: leave the step132 // 'started' so the next claim handles it by effect class instead of journaling "cancelled".133 if (signal.aborted) throw signal.reason instanceof Error ? signal.reason : new Error('tool execution aborted');134 return { result: await finalize(call.name, res, tctx) };135 });136}137138/** Folds one recorded tool result into the run state and returns what the model sees. */139function absorb(st: RunState, p: Extract<Prepared, { kind: 'tool' }>, r: ToolStepResult): ToolResult {140 st.fileState = applyFileStateDelta(st.fileState, r.fileStateOut);141 if (r.artifact) st.artifacts.push({ name: r.artifact.name, sha256: r.artifact.sha256 });142 const noop = p.call.name === 'git_commit' && /^(Nothing to commit|Already committed)/.test(r.content);143 if (r.ok && (p.effect === 'idempotent_write' || p.effect === 'non_idempotent') && p.call.name !== 'shell' && !noop) st.stall.progress();144 return toolResult(p.call, r.ok, r.content, { artifact: r.artifact, trust: r.trust });145}146147type PreparedTool = Extract<Prepared, { kind: 'tool' }>;148/** A read the policy allows outright: may run concurrently with its neighbours. (Plain boolean: a type guard would wrongly narrow the other branch.) */149const parallelEligible = (p: Prepared): boolean => p.kind === 'tool' && p.tool.parallelSafe && p.verdict.decision === 'allow';150151async function settleAll<T>(promises: Promise<T>[]): Promise<T[]> {152 const settled = await Promise.allSettled(promises);153 const failures = settled.filter((s): s is PromiseRejectedResult => s.status === 'rejected').map((s) => s.reason as unknown);154 if (failures.length > 0) throw failures.find(isControlError) ?? failures[0];155 return settled.map((s) => (s as PromiseFulfilledResult<T>).value);156}157158export async function runToolBatch(st: RunState, calls: ToolCall[]): Promise<BatchOutcome> {159 const { ctx, deps, meta } = st;160 // Rules and settings are read once per batch through a journaled step: decisions are replay-stable, and161 // a rule added from the console applies from the next batch on. 'pure' = a journaled observation.162 const snapshot = await ctx.step<Record<string, never>, PolicySnapshot>(163 { kind: 'policy', name: 'policy_rules', effect: 'pure', request: {} },164 async () => ({ result: await loadPolicySnapshot(deps.kernel, meta.ownerId, meta.projectId) }),165 );166 st.batchRules = [];167 const results: ToolResult[] = [];168 const record = (call: ToolCall, r: ToolResult): void => {169 results.push(r);170 st.stall.recordCall(call.name, call.input, r.ok, r.content);171 };172173 let i = 0;174 while (i < calls.length) {175 const p = prepare(st, calls[i] as ToolCall, snapshot);176 if (parallelEligible(p)) {177 const group: PreparedTool[] = [p as PreparedTool];178 while (i + group.length < calls.length) {179 const next = prepare(st, calls[i + group.length] as ToolCall, snapshot);180 if (!parallelEligible(next)) break;181 group.push(next as PreparedTool);182 }183 for (const g of group) await emitDecision(st, g.gate, g.verdict, g.verdict.ruleId ? { ruleId: g.verdict.ruleId } : {});184 const base = st.fileState;185 const done = await settleAll(group.map((g) => startToolStep(st, g, base)));186 group.forEach((g, n) => record(g.call, absorb(st, g, done[n] as ToolStepResult)));187 i += group.length;188 continue;189 }190 i++;191 if (p.kind === 'error') {192 record(p.call, toolResult(p.call, false, p.message, { trust: 'system' }));193 } else if (p.kind === 'control') {194 const out = await runControl(st, p.call, p.payload);195 record(p.call, out.result);196 if (out.progress) st.stall.progress();197 if (out.finish) return { results, finish: out.finish }; // calls issued after an accepted finish are not run198 } else {199 const gate = await applyVerdict(st, p.gate, p.verdict, snapshot);200 if (!gate.allowed) record(p.call, toolResult(p.call, false, gate.message, { trust: 'system' }));201 else record(p.call, absorb(st, p, await startToolStep(st, p, st.fileState)));202 }203 }204 return { results, finish: null };205}206