diff --git a/DESIGN.md b/DESIGN.md index 53526a7476..524a968a3e 100644 --- a/DESIGN.md +++ b/DESIGN.md @@ -118,6 +118,7 @@ Index of the design notes for the harper core: one line per note, grouped by the - [A recorded decision is committed before its id is returned, and its facts are separate rows](resources/models/DESIGN.md#a-recorded-decision-is-committed-before-its-id-is-returned-and-its-facts-are-separate-rows) — Recording is opt-in (`persist: true`); a recorded `decide` commits its `hdb_model_decisions` row in its own transaction before returning; truth and action are separate `hdb_model_outcomes` rows that copy the decision's expiry, so facts from different nodes never overwrite each other. - [No-match is a separate score, never a value](resources/models/DESIGN.md#no-match-is-a-separate-score-never-a-value) — An opted-in leaf gets a no-match score beside the unchanged distribution; no sentinel, capability-gated per candidate, dropped where not asked for. - [A `@decide` decision is recorded only when the record can reach it](resources/models/DESIGN.md#a-decide-decision-is-recorded-only-when-the-record-can-reach-it) — A directive records its decision only with a `decision` attribute, which is server-owned; association is bounded, and read-only nodes run no model hook and stage no cache fill. +- [A generative call is priced by the entry that served it](resources/models/DESIGN.md#a-generative-call-is-priced-by-the-entry-that-served-it) — `costUsd` comes from the serving backend's `pricing`, needs both token counts, is computed outside the fallback catch, and never changes calibration identity. - [A fitted calibration never crosses a score source, and never waits on storage or fails a call](resources/models/DESIGN.md#a-fitted-calibration-never-crosses-a-score-source-and-never-waits-on-storage-or-fails-a-call) — Populations are keyed by served-entry fingerprints and signatures; the decide path reads a cache only and fails open; immutable versions are ranked by the evidence they saw; the job is internal to the scheduler. ## server/ — HTTP stacks, threads, operation dispatch @@ -190,6 +191,8 @@ Index of the design notes for the harper core: one line per note, grouped by the - [The fs tools never return key material, and `config` is the config file (`agent/tools/fsTools.ts`)](agent/DESIGN.md#the-fs-tools-never-return-key-material-and-config-is-the-config-file-agenttoolsfstoolsts) — `config` defaults to the config file and is never enumerated; key directories, key file names and PEM private-key text are refused in every scope, against canonical paths. - [A turn that did not finish is rejected before anything is recorded (`agent/loop.ts`)](agent/DESIGN.md#a-turn-that-did-not-finish-is-rejected-before-anything-is-recorded-agentloopts) — a turn the backend reports as cut short, an unparsed tool call or an empty reply ends the run `error` and appends nothing. - [Tool results are capped where they are stored; a context-window rejection shrinks the newest once (`agent/loop.ts`)](agent/DESIGN.md#tool-results-are-capped-where-they-are-stored-a-context-window-rejection-shrinks-the-newest-once-agentloopts) — no observation over `agent.maxToolResultBytes` reaches the transcript; backends flag a context-window rejection and the loop shrinks the newest results once and retries. +- [Spend caps are folded from the transcript and fail closed (`agent/budget.ts`, `agent/loop.ts`)](agent/DESIGN.md#spend-caps-are-folded-from-the-transcript-and-fail-closed-agentbudgetts-agentloopts) — Usage rides on each assistant message; caps are per session, checked before approvals and after every call; unmeasured stops; a stop answers every tool call in one write. +- [Every tool call gets one response, and a cancel is final (`agent/loop.ts`, `agent/session.ts`)](agent/DESIGN.md#every-tool-call-gets-one-response-and-a-cancel-is-final-agentloopts-agentsessionts) — approvals are claimed before they run; a dead run's unanswered calls are answered `interrupted` in place; every late response goes to its turn; nothing replaces `aborted`; a cancel fences out runs admitted before it. ## config/ diff --git a/agent/DESIGN.md b/agent/DESIGN.md index 328b2bac4a..7e9a9ca1eb 100644 --- a/agent/DESIGN.md +++ b/agent/DESIGN.md @@ -2,7 +2,7 @@ The built-in Harper agent component. -**Read this when:** changing the agent's toolset, the `http_fetch` tool, the filesystem tools or their scopes, what `set_agent_config` can patch, how the loop ends a run, or how tool results reach the transcript. +**Read this when:** changing the agent's toolset, the `http_fetch` tool, the filesystem tools or their scopes, what `set_agent_config` can patch, how the loop ends a run, how tool results reach the transcript, or how a run records and caps model usage. Index of every design note: [DESIGN.md](../DESIGN.md). @@ -28,7 +28,7 @@ The rules bound the fs tools only: an operations tool returns its own text unfil ## A turn that did not finish is rejected before anything is recorded (`agent/loop.ts`) -Callers poll `get_agent_session` and read `completed` as success, so `completed` with no `lastError` must mean a natural, non-empty final answer (harper#3039). The `maxTurns` ceiling is the one documented exception: it ends `completed` and sets `lastError`. `rejectTurn` runs right after each `generate` and ends the run `error` when the backend reports `length` or `content_filter`, when it reports `tool_calls` but no call parsed, or when a reply has neither text nor tool calls. +Callers poll `get_agent_session` and read `completed` as success, so `completed` with no `lastError` must mean a natural, non-empty final answer (harper#3039). Two documented exceptions end `completed` and set `lastError`: the `maxTurns` ceiling, and a final answer that crosses a spend cap, which is still delivered (below). `rejectTurn` runs right after each `generate` and ends the run `error` when the backend reports `length` or `content_filter`, when it reports `tool_calls` but no call parsed, or when a reply has neither text nor tool calls. - **Nothing from a rejected turn is appended.** Partial text would read as the answer. Tool calls recorded without their responses would make the session's next request invalid, and a truncated call carries truncated arguments, so none is executed or queued for approval. The rejection is one status write, so there is no partial state to recover. - **Tool calls still dispatch on `stop`.** Only the explicit cut-short reasons block dispatch. `synthesizeGenerateFromStream` (`resources/models/backendRegistry.ts`) reports `stop` for a stream-only backend that emits no finish reason, so gating dispatch on `tool_calls` would stop such a backend from calling tools at all. The cost: truncation from such a backend is invisible here, and the guarantee holds only for backends that report it. The four built-in backends implement `generate` and do. @@ -43,3 +43,24 @@ Every request replays the whole transcript, so one oversized tool result fails e - **A context-window rejection shrinks, then retries once.** Backends own provider wording. They set `contextWindowExceeded` on their error (`isContextWindowRejection`, `resources/models/backendHelpers.ts`), and the loop reads only `isContextWindowExceeded`. It cuts each result over 2 KiB in the newest run of tool messages that has one down to a 2 KiB head, then retries. Selecting by size makes the step idempotent and reaches a result stored before a later prompt. It is best effort: a long prompt or history ends the run `error` with a `lastError` saying so. `Models.generate` surfaces the first candidate's error, so a rejection from a fallback candidate behind an unrelated primary failure is not recovered. The system prompt is sent only as `input.system`. The backends hoist or prepend it, so an inline copy would be sent twice. + +## Spend caps are folded from the transcript and fail closed (`agent/budget.ts`, `agent/loop.ts`) + +`agent.maxSessionTokens` and `agent.maxCostUsd` bound a session, not a run: `approve_agent_action` and a follow-up `agent_prompt` each start a new `runAgent` or join the live one, so a per-run counter (which is what `maxTurns` is) would reset on every resume (harper#3040). Three rules keep the caps honest: + +- **Usage lives on the call it describes.** Each assistant message carries the `usage` its `generate` reported and its `costUsd`, in the same write as the message. Session totals are folded from the transcript by `summarizeSessionUsage`, never stored, so they cannot drift from the calls. A rejected turn (above) never reaches the transcript, but it was billed, so `endRun` records its usage on `rejectedCalls`, which the fold counts as a call. A total stays present only while every assistant message reported it, so a backend that reports nothing, an unpriced model, or a message written before usage was recorded makes that total unmeasured, not zero. +- **Unmeasured stops the session.** `budgetViolation` treats a cap whose total is unmeasured as violated. Checked before `consumeResolvedApprovals` (so a cap lowered during a pause stops an approved action from running) and after every call. +- **A stop leaves a sendable transcript.** `endRun` appends the assistant turn, a response for every tool call and every open approval, resolved or not (`budget_exceeded`, or `denied_by_operator` for a denial), and the terminal status in one write, so a stop leaves no tool call without the response the provider requires on the next call, and no approval stays pending once the session leaves `awaiting_approval`. A cancel that landed first keeps its `aborted` status; the stop still completes the transcript. + +A final answer that crosses a cap is delivered; the run ends `completed` with `lastError` naming the cap. What a cap does not bound, by construction: one call can cross it (it is a stop-after-call threshold), a failed attempt before a model fallback is not counted (failure records carry no usage), and an unpriced model is detected only after its first call. Pinned by `unitTests/agent/loop.test.js` and `integrationTests/server/agent-budget.test.ts`. + +## Every tool call gets one response, and a cancel is final (`agent/loop.ts`, `agent/session.ts`) + +A provider rejects a request in which a tool call has no response, or has two, or in which a response follows a later prompt; a session in that state cannot continue. Four rules keep every transcript sendable while cancels, stops, approvals and crashes race each other: + +- **An approved action is claimed before it runs, and runs at most once.** `claimApproval` marks it taken under the session lock, and every writer that answers approvals (`cancelSession`, `endRun`) decides which are open on the row it read under that lock, so a cancel or a stop never answers an action already running. The run checks the cancel again after the claim. A crash between the claim and the result leaves that call unanswered rather than running the action twice. +- **Calls a dead run left unanswered are repaired, in place.** Inline results are appended one at a time as each call finishes, so a crash mid-turn can leave calls without a response too. The next run starts by answering every such call of the latest turn `interrupted`, outcome unknown (`answerInterruptedCalls`), without running it again; it looks only at the latest turn, as a provider may reuse a tool-call id. A run that fails answers its open approvals and its turn's remaining calls `run_failed` in its final write. Every response written after its turn has moved on (an approval's result, a cancel's or stop's answer, `run_failed`, the repair) is inserted right after that turn's existing responses (`responseIndex`), so a prompt appended meanwhile, for example after a crash between resolving an approval and starting its run, stays after them. +- **A cancel is final.** `cancelSession` answers every open approval `cancelled_by_operator` in the write that marks the session `aborted`, and decides under the lock, so it leaves a session a run already finished as it is. Every status the loop writes after `running` goes through `endRun`, which keeps `aborted`, and `addPendingApproval` answers a gated call `cancelled_by_operator` on an `aborted` session. A cancel between a turn's tool calls answers the rest of the turn in one write. Every cancel bumps a per-session count (`fenceRuns`); a request that starts a run takes an `admitRun` token before its first await, and `runAgent` starts nothing for a token a later cancel fenced out, so a cancel that arrives while `agent_prompt` or `approve_agent_action` is still writing is not undone by the run it would start. `agent_prompt` answers 409 while a cancelled run is still finishing (`isRunInFlight`), because a prompt appended then would land between that run's tool calls and their results; each of its 409s carries a `code` and `retryable` (false only when the prompt was saved but not run), and session reads carry `runInFlight`, since a cancelled session reads `aborted` before its run is done. +- **An approval resolved while a run is live joins that run.** That covers one resolved during the drain or during the turn's later inline calls. So the run drains after run start and after any turn that parked an approval, repeats until nothing resolved is left, decides the pause under the lock (`pauseWhileApprovalsOpen`), and writes `running` back over the `idle` that resolving the last approval leaves. Such an approval is checked against the caps the live run started with, like everything else in that run; a cap lowered meanwhile applies from the next run. A scheduled followup that fires while a run is in flight or approvals are pending retries until neither holds, carrying its admission token so a cancel drops it. + +Run control is per process: the session lock, the in-flight runs and the cancel count live in memory, while `hdb_agent_session` replicates, so two nodes acting on one session are not serialized. Pinned by `unitTests/agent/loop.test.js` and `unitTests/agent/operations.test.js`. diff --git a/agent/agent.ts b/agent/agent.ts index be0ed0d2a6..afa678393d 100644 --- a/agent/agent.ts +++ b/agent/agent.ts @@ -30,24 +30,30 @@ import { composeRegistryTools, ensureOperationsToolsRegistered } from './registr import { buildOperations } from './operations.ts'; import { registerAgentMcpTools } from './mcpTools.ts'; import { - runAgent, - _resetInFlightForTests, + admitRun, DEFAULT_MAX_TOOL_RESULT_BYTES, + fenceRuns, + isRunInFlight, isValidMaxToolResultBytes, MAX_MAX_TOOL_RESULT_BYTES, MIN_MAX_TOOL_RESULT_BYTES, + runAgent, + _resetInFlightForTests, } from './loop.ts'; +import { parseBudgetCap, parseMaxTurns } from './budget.ts'; import { appendMessage, getSession } from './session.ts'; import type { AgentConfig, AgentScopes, AgentTool, HttpFetchConfig } from './types.ts'; const log = harperLogger.loggerWithTag('agent'); +const FOLLOWUP_RETRY_MS = 1000; +const FOLLOWUP_RETRY_MAX_MS = 30_000; + const DEFAULT_CONFIG: AgentConfig = { enabled: false, maxTurns: 50, maxTokens: 16384, maxToolResultBytes: DEFAULT_MAX_TOOL_RESULT_BYTES, - maxCostUsd: 5, autoApprove: false, allowDestructive: false, user: 'hdb_agent', @@ -64,10 +70,11 @@ interface StartOpts { enabled?: boolean; provider?: string; model?: string; - maxTurns?: number; + maxTurns?: unknown; maxTokens?: unknown; maxToolResultBytes?: unknown; - maxCostUsd?: number; + maxSessionTokens?: unknown; + maxCostUsd?: unknown; autoApprove?: boolean; allowDestructive?: boolean; user?: string; @@ -83,6 +90,10 @@ export async function startOnMainThread(opts: StartOpts): Promise { log.info?.('Agent component disabled (agent.enabled=false); skipping registration'); return; } + // A malformed limit fails the agent's start rather than running it uncapped or turnless. + if (opts.maxTurns !== undefined) config.maxTurns = parseMaxTurns(opts.maxTurns); + config.maxSessionTokens = parseBudgetCap('maxSessionTokens', opts.maxSessionTokens); + config.maxCostUsd = parseBudgetCap('maxCostUsd', opts.maxCostUsd); // Lazily imported to avoid pulling configUtils at module-eval time for tests. const { getConfigPath, getConfigFilePath } = await import('../config/configUtils.ts'); @@ -128,7 +139,7 @@ export async function startOnMainThread(opts: StartOpts): Promise { return composeToolset({ allowDestructive: liveConfig.allowDestructive, httpFetchTool, - onFollowup: handleFollowup, + onFollowup: (sessionId, prompt) => handleFollowup(sessionId, prompt), inspectorTools, registryTools, extraTools, @@ -152,35 +163,48 @@ export async function startOnMainThread(opts: StartOpts): Promise { } } - // Only warn when the operator explicitly configured `maxCostUsd`. Logging on the default - // every boot would flood the log without telling anyone anything actionable. - if (opts.maxCostUsd !== undefined) { - log.warn?.( - `agent.maxCostUsd=${liveConfig.maxCostUsd} is advertised but not yet enforced; cost-cap wiring depends on #612 telemetry.` - ); - } - - async function handleFollowup(sessionId: string, prompt: string): Promise { + async function handleFollowup( + sessionId: string, + prompt: string, + admitted = admitRun(sessionId), + retryMs = FOLLOWUP_RETRY_MS + ): Promise { const session = await getSession(sessionId); if (!session) { log.warn?.(`schedule_followup target session ${sessionId} not found`); return; } + if (admitRun(sessionId) !== admitted) return; + // A prompt appended mid-run or mid-approval would land between tool calls and their results, so it + // waits; the token makes a cancel drop the retry, as it clears the scheduled timer. + if (isRunInFlight(sessionId) || session.status === 'awaiting_approval') { + const next = Math.min(retryMs * 2, FOLLOWUP_RETRY_MAX_MS); + setTimeout( + () => + handleFollowup(sessionId, prompt, admitted, next).catch((err) => + log.error?.(`schedule_followup for ${sessionId} failed: ${(err as Error)?.message ?? err}`) + ), + retryMs + ).unref?.(); + return; + } await appendMessage(sessionId, { role: 'user', content: prompt, createdAt: Date.now() }); - startRun(sessionId); + startRun(sessionId, admitted); } function currentTools(): AgentTool[] { return composed.tools; } - function startRun(sessionId: string): void { + /** Returns false when a cancel since the request was admitted fenced the run out. */ + function startRun(sessionId: string, admitted?: number): boolean { + if (admitted !== undefined && admitted !== admitRun(sessionId)) return false; // A run is already active for this session. `runAgent` coalesces concurrent starts onto the // existing in-flight promise (bound to the existing controller), so creating a second // controller here would orphan it — `cancelRun` would then abort a controller nothing is // listening to, leaving the live run uncancellable. Any message appended before this call // (e.g. by a scheduled followup) is picked up by the in-flight loop on its next turn. - if (abortControllers.has(sessionId)) return; + if (abortControllers.has(sessionId)) return true; const controller = new AbortController(); abortControllers.set(sessionId, controller); runAgent({ @@ -189,8 +213,11 @@ export async function startOnMainThread(opts: StartOpts): Promise { tools: currentTools(), scopes, maxTurns: liveConfig.maxTurns, + maxSessionTokens: liveConfig.maxSessionTokens, + maxCostUsd: liveConfig.maxCostUsd, autoApprove: liveConfig.autoApprove, signal: controller.signal, + admitted, generateOpts: { model: liveConfig.model, maxTokens: liveConfig.maxTokens }, systemPrompt: composeSystemPrompt(staticSystemPrompt, liveConfig.systemPromptAppend), maxToolResultBytes: liveConfig.maxToolResultBytes, @@ -199,9 +226,11 @@ export async function startOnMainThread(opts: StartOpts): Promise { .finally(() => { if (abortControllers.get(sessionId) === controller) abortControllers.delete(sessionId); }); + return true; } function cancelRun(sessionId: string): boolean { + fenceRuns(sessionId); // Clear any scheduled followups first. Without this, a timer set via `schedule_followup` // would fire after the operator cancelled, silently re-injecting a user prompt and kicking // the loop off again — surprising behavior and avoidable LLM cost. @@ -237,6 +266,8 @@ export async function startOnMainThread(opts: StartOpts): Promise { setConfig, startRun, cancelRun, + admitRun, + isRunInFlight, }); for (const op of operations) opts.server.registerOperation(op); @@ -382,12 +413,10 @@ export function mergeConfig(opts: StartOpts): AgentConfig { ...(opts.enabled !== undefined && { enabled: !!opts.enabled }), ...(opts.provider !== undefined && { provider: String(opts.provider) }), ...(opts.model !== undefined && { model: String(opts.model) }), - ...(opts.maxTurns !== undefined && { maxTurns: Number(opts.maxTurns) }), ...(opts.maxTokens !== undefined && { maxTokens: resolveMaxTokens(opts.maxTokens) }), ...(opts.maxToolResultBytes !== undefined && { maxToolResultBytes: resolveMaxToolResultBytes(opts.maxToolResultBytes), }), - ...(opts.maxCostUsd !== undefined && { maxCostUsd: Number(opts.maxCostUsd) }), ...(opts.autoApprove !== undefined && { autoApprove: !!opts.autoApprove }), ...(opts.allowDestructive !== undefined && { allowDestructive: !!opts.allowDestructive }), ...(opts.user !== undefined && { user: String(opts.user) }), diff --git a/agent/budget.ts b/agent/budget.ts new file mode 100644 index 0000000000..0612b5c2ae --- /dev/null +++ b/agent/budget.ts @@ -0,0 +1,107 @@ +/** + * Per-session spend caps for the built-in agent (#3040), and validation of the agent's numeric limits. + * Usage is never stored as a running total: + * each assistant message carries what its model call reported, and the session's usage is folded + * from the transcript, so a total cannot drift from the calls it summarizes. + */ + +import { ClientError } from '../utility/errors/hdbError.ts'; +import type { AgentCallUsage, AgentMessage, AgentSessionUsage } from './types.ts'; + +export interface BudgetCaps { + maxSessionTokens?: number; + maxCostUsd?: number; +} + +export function hasBudget(caps: BudgetCaps): boolean { + return caps.maxSessionTokens !== undefined || caps.maxCostUsd !== undefined; +} + +export function summarizeSessionUsage( + messages: Iterable, + rejectedCalls: Iterable = [] +): AgentSessionUsage { + let calls = 0; + let promptTokens: number | undefined = 0; + let completionTokens: number | undefined = 0; + let costUsd: number | undefined = 0; + const add = (usage: AgentCallUsage | undefined) => { + calls++; + promptTokens = addMeasured(promptTokens, usage?.promptTokens); + completionTokens = addMeasured(completionTokens, usage?.completionTokens); + costUsd = addMeasured(costUsd, usage?.costUsd); + }; + for (const message of messages) { + if (message.role === 'assistant') add(message.usage); + } + for (const usage of rejectedCalls) add(usage); + const usage: AgentSessionUsage = { calls }; + if (promptTokens !== undefined) usage.promptTokens = promptTokens; + if (completionTokens !== undefined) usage.completionTokens = completionTokens; + if (costUsd !== undefined) usage.costUsd = costUsd; + return usage; +} + +function addMeasured(total: number | undefined, value: unknown): number | undefined { + if (total === undefined || typeof value !== 'number' || !Number.isFinite(value) || value < 0) return undefined; + return total + value; +} + +export function callUsage( + usage: { promptTokens?: number; completionTokens?: number } | undefined, + costUsd: number | undefined +): AgentCallUsage | undefined { + const reported: AgentCallUsage = {}; + if (usage?.promptTokens !== undefined) reported.promptTokens = usage.promptTokens; + if (usage?.completionTokens !== undefined) reported.completionTokens = usage.completionTokens; + if (costUsd !== undefined) reported.costUsd = costUsd; + return Object.keys(reported).length > 0 ? reported : undefined; +} + +/** An unmeasured total violates its cap: running uncapped under a configured cap is what this prevents. */ +export function budgetViolation(usage: AgentSessionUsage, caps: BudgetCaps): string | undefined { + if (caps.maxSessionTokens !== undefined) { + if (usage.promptTokens === undefined || usage.completionTokens === undefined) { + return `agent.maxSessionTokens=${caps.maxSessionTokens} cannot be enforced: a model call in this session reported no token usage; start a new session or clear the cap`; + } + const used = usage.promptTokens + usage.completionTokens; + if (used >= caps.maxSessionTokens) + return `agent.maxSessionTokens=${caps.maxSessionTokens} reached: this session used ${used} tokens`; + } + if (caps.maxCostUsd !== undefined) { + if (usage.costUsd === undefined) { + return `agent.maxCostUsd=${caps.maxCostUsd} cannot be enforced: a model call in this session was not priced; declare pricing on the models.generative entry the agent uses, then start a new session or clear the cap`; + } + if (usage.costUsd >= caps.maxCostUsd) { + return `agent.maxCostUsd=${caps.maxCostUsd} reached: this session cost $${usage.costUsd.toFixed(4)}`; + } + } + return undefined; +} + +/** Refuses rather than coerces: a cap coerced to `NaN` would compare false forever and silently uncap. */ +export function parseBudgetCap(name: keyof BudgetCaps, value: unknown): number | undefined { + if (value === undefined || value === null) return undefined; + const cap = configNumber(value); + const valid = + typeof cap === 'number' && + Number.isFinite(cap) && + (name === 'maxSessionTokens' ? Number.isInteger(cap) && cap >= 1 : cap >= 0); + if (!valid) { + const expected = name === 'maxSessionTokens' ? 'a positive integer' : 'a non-negative number'; + throw new ClientError(`agent.${name} must be ${expected} or null, got ${JSON.stringify(value)}`, 400); + } + return cap; +} + +/** `maxTurns` has a default rather than an "off", so `null`, which a bare config key reads as, is refused like any other non-integer. */ +export function parseMaxTurns(value: unknown): number { + const turns = configNumber(value); + if (typeof turns === 'number' && Number.isInteger(turns) && turns >= 1) return turns; + throw new ClientError(`agent.maxTurns must be a positive integer, got ${JSON.stringify(value)}`, 400); +} + +/** Env vars and CLI flags arrive as strings. */ +function configNumber(value: unknown): unknown { + return typeof value === 'string' && value.trim() !== '' ? Number(value) : value; +} diff --git a/agent/loop.ts b/agent/loop.ts index ae6c9efacd..58fa71f939 100644 --- a/agent/loop.ts +++ b/agent/loop.ts @@ -21,14 +21,19 @@ import { isContextWindowExceeded } from '../resources/models/backendHelpers.ts'; import harperLogger from '../utility/logging/harper_logger.ts'; import { addPendingApproval, + answerInterruptedCalls, + answerToolCall, appendMessage, + cancelledObservation, + claimApproval, + endRun, getSession, - markApprovalConsumed, setStatus, shrinkNewestToolResults, } from './session.ts'; -import type { AgentMessage, AgentScopes, AgentTool, AgentToolContext } from './types.ts'; +import type { AgentMessage, AgentScopes, AgentTool, AgentToolContext, ApprovalRequest } from './types.ts'; import { toolMapByName } from './toolset.ts'; +import { budgetViolation, callUsage, hasBudget, summarizeSessionUsage } from './budget.ts'; const log = harperLogger.loggerWithTag('agent'); @@ -57,6 +62,11 @@ export interface RunAgentOpts { tools: AgentTool[]; scopes: AgentScopes; maxTurns: number; + /** Per-session caps, measured from the transcript, so an approval resume or a new prompt cannot reset them. */ + maxSessionTokens?: number; + maxCostUsd?: number; + /** From `admitRun` when the request that starts this run was admitted; a cancel since then fences it out. */ + admitted?: number; /** When false, destructive tools pause the loop with a pending approval instead of executing. */ autoApprove?: boolean; signal?: AbortSignal; @@ -68,8 +78,23 @@ export interface RunAgentOpts { } const inFlight = new Map>(); +const cancels = new Map(); + +/** A token for a run admitted now. Take it before the request's first await. */ +export function admitRun(sessionId: string): number { + return cancels.get(sessionId) ?? 0; +} + +/** + * Called on every cancel, including one that finds no live run to abort: a request admitted before it, + * still writing when the cancel arrived, must not start a run that writes `running` over `aborted`. + */ +export function fenceRuns(sessionId: string): void { + cancels.set(sessionId, admitRun(sessionId) + 1); +} export function runAgent(opts: RunAgentOpts): Promise { + if (opts.admitted !== undefined && opts.admitted !== admitRun(opts.sessionId)) return Promise.resolve(); const existing = inFlight.get(opts.sessionId); if (existing) return existing; const run = doRun(opts).finally(() => { @@ -79,9 +104,17 @@ export function runAgent(opts: RunAgentOpts): Promise { return run; } +/** Includes a cancelled run still finishing a model call or a handler. */ +export function isRunInFlight(sessionId: string): boolean { + return inFlight.has(sessionId); +} + async function doRun(opts: RunAgentOpts): Promise { const toolMap = toolMapByName(opts.tools); const toolDefs: ToolDef[] = opts.tools.map((t) => t.def); + // The cancel path already wrote `aborted`; `running` must not replace it. Every later status write + // goes through `endRun`, which keeps `aborted` under the session lock. + if (opts.signal?.aborted) return; await setStatus(opts.sessionId, 'running'); const ctx: AgentToolContext = { sessionId: opts.sessionId, @@ -91,26 +124,25 @@ async function doRun(opts: RunAgentOpts): Promise { }; try { - // First, drain any resolved-but-unconsumed approvals from a prior pause. Either execute - // or refuse each saved call, recording an observation, so the next model turn sees the - // result of the operator decision. - await consumeResolvedApprovals(opts.sessionId, toolMap, ctx); - - // If a turn produced multiple gated tool calls and the operator has only resolved some of - // them, the remaining approvals are still pending — meaning the assistant's tool_calls do - // not yet all have tool responses. Re-entering the generate loop now would send an - // incomplete tool-response set and the provider would 400. Stay paused until every gated - // call for this turn is resolved (each `approve_agent_action` re-runs this path). - const afterConsume = await getSession(opts.sessionId); - if (afterConsume?.pendingApprovals.some((a) => !a.resolved)) { - await setStatus(opts.sessionId, 'awaiting_approval'); - return; - } + // Calls a dead run left unanswered are repaired first, so no later path sends them unanswered. + await answerInterruptedCalls(opts.sessionId); + // Before the approvals below can execute anything: a cap lowered while a run sat awaiting + // approval must stop the approved action, not just the next model call. + if (await stopIfOverBudget(opts)) return; + if (!(await drainApprovals(opts, toolMap, ctx))) return; for (let turn = 0; turn < opts.maxTurns; turn++) { if (opts.signal?.aborted) return; // status was already set to `aborted` by cancelRun const result = await generateTurn(opts, toolDefs); - if (opts.signal?.aborted) return; + if (opts.signal?.aborted) { + // The call returned after the cancel; its reply is discarded but it was billed, so it is counted. + await endRun(opts.sessionId, { + status: 'aborted', + lastError: 'Cancelled by operator', + rejectedCall: callUsage(result.usage, result.costUsd) ?? {}, + }); + return; + } log.debug?.( `Session ${opts.sessionId} turn ${turn}: finishReason=${result.finishReason} toolCalls=${result.toolCalls?.length ?? 0} contentChars=${result.content?.length ?? 0} completionTokens=${result.usage?.completionTokens ?? 'unknown'}` ); @@ -118,33 +150,64 @@ async function doRun(opts: RunAgentOpts): Promise { // Must run before appendMessage: nothing from a rejected turn may reach the transcript (agent/DESIGN.md). const rejection = rejectTurn(result, opts.generateOpts?.maxTokens); if (rejection) { - await setStatus(opts.sessionId, 'error', rejection); + // The rejected call was still billed, so its usage is kept on the session, outside the transcript. + await endRun(opts.sessionId, { + status: 'error', + lastError: rejection, + rejectedCall: callUsage(result.usage, result.costUsd) ?? {}, + }); return; } - await appendMessage(opts.sessionId, { + const assistantMessage: AgentMessage = { role: 'assistant', content: result.content ?? '', toolCalls: result.toolCalls, createdAt: Date.now(), - }); + }; + const usage = callUsage(result.usage, result.costUsd); + if (usage) assistantMessage.usage = usage; + + const session = hasBudget(opts) ? await getSession(opts.sessionId) : undefined; + const violation = session + ? budgetViolation(summarizeSessionUsage([...session.messages, assistantMessage], session.rejectedCalls), opts) + : undefined; + if (violation) { + const refusals = (result.toolCalls ?? []).map((call) => budgetRefusal(call.id, violation)); + await endRun(opts.sessionId, { + append: [assistantMessage, ...refusals], + status: refusals.length > 0 ? 'error' : 'completed', + lastError: violation, + }); + return; + } + + await appendMessage(opts.sessionId, assistantMessage); if (!result.toolCalls || result.toolCalls.length === 0) { - if (opts.signal?.aborted) return; - await setStatus(opts.sessionId, 'completed'); + await endRun(opts.sessionId, { status: 'completed' }); return; } const paused = await dispatchToolCalls(result.toolCalls, toolMap, ctx, opts); - if (paused || opts.signal?.aborted) return; + if (opts.signal?.aborted) return; + if (paused && !(await drainApprovals(opts, toolMap, ctx))) return; } - await setStatus(opts.sessionId, 'completed', `Reached maxTurns=${opts.maxTurns} without a final answer.`); + await endRun(opts.sessionId, { + status: 'completed', + lastError: `Reached maxTurns=${opts.maxTurns} without a final answer.`, + }); } catch (err) { // If the abort signal fired, the cancel path already set the session to `aborted` — // don't clobber that with `error`. The rejection here is just the awaited generate/tool // honoring the signal, not a real failure. if (opts.signal?.aborted) return; - await setStatus(opts.sessionId, 'error', err instanceof Error ? err.message : String(err)); + await endRun(opts.sessionId, { + status: 'error', + lastError: err instanceof Error ? err.message : String(err), + answerOpenApprovals: (approval) => JSON.stringify({ ok: false, error: 'run_failed', tool: approval.toolName }), + answerUnansweredCalls: (toolName) => JSON.stringify({ ok: false, error: 'run_failed', tool: toolName }), + }); throw err; } } @@ -217,6 +280,68 @@ function contextWindowError(cause: unknown, detail: string): Error { return new Error(`The conversation no longer fits the model's context window: ${detail}. (${reason})`, { cause }); } +/** + * Drain resolved approvals until none is left; false when the run must stop here (cancelled, or paused on + * an unresolved approval). An approval resolved while this run is live, during the drain or during the + * turn's later inline calls, joins this run (`startRun`), so it is drained here, not left unanswered for + * the next model call. Generating while any approval is unresolved would send tool calls without their + * responses, which the provider rejects, so the run pauses; the pause is decided under the lock. + */ +async function drainApprovals( + opts: RunAgentOpts, + toolMap: Map, + ctx: AgentToolContext +): Promise { + for (;;) { + if (opts.signal?.aborted) return false; + await consumeResolvedApprovals(opts.sessionId, toolMap, ctx); + const session = await getSession(opts.sessionId); + const approvals = session?.pendingApprovals ?? []; + if (approvals.some((a) => a.resolved && !a.consumed)) continue; + if (!approvals.some((a) => !a.resolved)) { + // `resolveApproval` writes `idle` when it resolves the last approval, even into a live run. + if (session && session.status !== 'running') await endRun(opts.sessionId, { status: 'running' }); + return true; + } + const paused = await endRun(opts.sessionId, { status: 'awaiting_approval', pauseWhileApprovalsOpen: true }); + if (paused.status === 'aborted' || paused.pendingApprovals.some((a) => !a.resolved)) return false; + } +} + +/** + * Every open approval, resolved or not, is answered without executing it, so no tool call is left + * without a response once the session leaves `awaiting_approval`. + */ +async function stopIfOverBudget(opts: RunAgentOpts): Promise { + if (!hasBudget(opts)) return false; + const session = await getSession(opts.sessionId); + if (!session) return false; + const violation = budgetViolation(summarizeSessionUsage(session.messages, session.rejectedCalls), opts); + if (!violation) return false; + await endRun(opts.sessionId, { + answerOpenApprovals: (approval) => approvalRefusal(approval, violation), + status: 'error', + lastError: violation, + }); + return true; +} + +function budgetRefusal(toolCallId: string, reason: string): AgentMessage { + return { role: 'tool', content: budgetExceeded(reason), toolCallId, createdAt: Date.now() }; +} + +function budgetExceeded(reason: string): string { + return JSON.stringify({ ok: false, error: 'budget_exceeded', reason }); +} + +function approvalRefusal(approval: ApprovalRequest, reason: string): string { + return !approval.resolved || approval.approved ? budgetExceeded(reason) : deniedObservation(approval); +} + +function deniedObservation(approval: ApprovalRequest): string { + return JSON.stringify({ ok: false, error: 'denied_by_operator', tool: approval.toolName }); +} + /** * Returns `true` when the loop should pause (any destructive tool call required approval). * Non-destructive calls execute inline and their observations are appended. Destructive calls @@ -232,8 +357,21 @@ async function dispatchToolCalls( opts: RunAgentOpts ): Promise { let needsApproval = false; - for (const call of calls) { - if (opts.signal?.aborted) return true; + for (const [index, call] of calls.entries()) { + if (opts.signal?.aborted) { + // Answer the rest of the turn in one write, so the transcript stays sendable after the cancel. + await endRun(opts.sessionId, { + append: calls.slice(index).map((unanswered) => ({ + role: 'tool', + content: cancelledObservation(unanswered.name), + toolCallId: unanswered.id, + createdAt: Date.now(), + })), + status: 'aborted', + lastError: 'Cancelled by operator', + }); + return true; + } const tool = toolMap.get(call.name); const destructiveAndGated = tool?.destructive && !opts.autoApprove; if (destructiveAndGated) { @@ -269,20 +407,22 @@ async function consumeResolvedApprovals( if (!session) return; const toConsume = session.pendingApprovals.filter((a) => a.resolved && !a.consumed); for (const approval of toConsume) { + if (ctx.signal?.aborted) return; + if (!(await claimApproval(sessionId, approval.id))) continue; + // The claim makes this run the only writer that can answer the call; a cancel that landed during + // the claim gets that answer instead of the action. + if (ctx.signal?.aborted) { + await answerToolCall(sessionId, approval.toolCallId, cancelledObservation(approval.toolName)); + return; + } const observation = approval.approved ? await invokeTool( { id: approval.toolCallId, name: approval.toolName, arguments: approval.arguments }, toolMap, ctx ) - : JSON.stringify({ ok: false, error: 'denied_by_operator', tool: approval.toolName }); - await appendMessage(sessionId, { - role: 'tool', - content: observation, - toolCallId: approval.toolCallId, - createdAt: Date.now(), - }); - await markApprovalConsumed(sessionId, approval.id); + : deniedObservation(approval); + await answerToolCall(sessionId, approval.toolCallId, observation); } } @@ -313,4 +453,5 @@ function toModelMessages(items: AgentMessage[]): Message[] { /** Test-only: clear the in-flight tracking map. */ export function _resetInFlightForTests(): void { inFlight.clear(); + cancels.clear(); } diff --git a/agent/operations.ts b/agent/operations.ts index 49198f81aa..fbaf00a874 100644 --- a/agent/operations.ts +++ b/agent/operations.ts @@ -13,15 +13,27 @@ import type { OperationDefinition } from '../server/serverHelpers/serverUtilities.ts'; import { OPERATIONS_ENUM } from '../utility/hdbTerms.ts'; import { ClientError } from '../utility/errors/hdbError.ts'; -import { createSession, getSession, listSessions, appendMessage, resolveApproval, setStatus } from './session.ts'; +import { + AGENT_SESSION_ATTRIBUTES, + appendMessage, + cancelSession, + createSession, + getSession, + listSessions, + resolveApproval, +} from './session.ts'; +import { parseBudgetCap, parseMaxTurns, summarizeSessionUsage } from './budget.ts'; +import type { AgentConfig, AgentMessage, AgentRunStatus, AgentSessionRow, AgentSessionUsage } from './types.ts'; import { isValidMaxToolResultBytes, MAX_MAX_TOOL_RESULT_BYTES, MIN_MAX_TOOL_RESULT_BYTES } from './loop.ts'; -import type { AgentConfig, AgentMessage, AgentRunStatus } from './types.ts'; export interface OperationDeps { getConfig: () => AgentConfig; setConfig: (patch: Partial) => AgentConfig; - startRun: (sessionId: string) => void; + /** `admitted` is the `admitRun` token taken before the request's first await; false when a cancel fenced it out. */ + startRun: (sessionId: string, admitted?: number) => boolean; cancelRun: (sessionId: string) => boolean; + admitRun: (sessionId: string) => number; + isRunInFlight: (sessionId: string) => boolean; } export function buildOperations(deps: OperationDeps): OperationDefinition[] { @@ -33,12 +45,12 @@ export function buildOperations(deps: OperationDeps): OperationDefinition[] { }, { name: OPERATIONS_ENUM.GET_AGENT_SESSION, - execute: async (op) => getAgentSession(op), + execute: async (op) => getAgentSession(op, deps), requiresSuperUser: true, }, { name: OPERATIONS_ENUM.LIST_AGENT_SESSIONS, - execute: async (op) => listAgentSessions(op), + execute: async (op) => listAgentSessions(op, deps), requiresSuperUser: true, }, { @@ -68,13 +80,31 @@ async function agentPrompt(op: any, deps: OperationDeps) { if (!message) throw new ClientError('message is required', 400); let sessionId: string = op?.session_id; + const admitted = sessionId ? deps.admitRun(sessionId) : undefined; if (sessionId) { const existing = await getSession(sessionId); if (!existing) throw new ClientError(`Unknown session ${sessionId}`, 404); if (existing.status === 'running' || existing.status === 'awaiting_approval') { - throw new ClientError( + throw promptConflict( `Session ${sessionId} is ${existing.status}; resolve or cancel before sending another prompt`, - 409 + 'AGENT_SESSION_BUSY', + true + ); + } + // A cancelled run can still be finishing a model call or a handler; a prompt appended now would land + // between that call's tool calls and their results, or be folded into the cancelled run and lost. + if (deps.isRunInFlight(sessionId)) { + throw promptConflict( + `Session ${sessionId} still has a run finishing; send the prompt again shortly`, + 'AGENT_RUN_FINISHING', + true + ); + } + if (deps.admitRun(sessionId) !== admitted) { + throw promptConflict( + `A cancel for session ${sessionId} arrived while this prompt was being accepted`, + 'AGENT_PROMPT_CANCELLED', + true ); } await appendMessage(sessionId, asUserMessage(message)); @@ -88,21 +118,49 @@ async function agentPrompt(op: any, deps: OperationDeps) { sessionId = created.session_id; } - deps.startRun(sessionId); + if (!deps.startRun(sessionId, admitted)) { + throw promptConflict( + `A cancel for session ${sessionId} arrived while this prompt was being accepted; the prompt was recorded but not run`, + 'AGENT_PROMPT_NOT_RUN', + false + ); + } return { session_id: sessionId, status: 'running' satisfies AgentRunStatus }; } -async function getAgentSession(op: any) { +/** `retryable` tells a client whether the prompt was left unrecorded, so sending it again cannot duplicate it. */ +function promptConflict(message: string, code: string, retryable: boolean): ClientError { + return Object.assign(new ClientError(message, 409), { code, retryable }); +} + +async function getAgentSession(op: any, deps: OperationDeps) { const sessionId = String(op?.session_id ?? ''); if (!sessionId) throw new ClientError('session_id is required', 400); const session = await getSession(sessionId); if (!session) throw new ClientError(`Unknown session ${sessionId}`, 404); - return session; + return withUsage(session, deps); } -async function listAgentSessions(op: any) { +async function listAgentSessions(op: any, deps: OperationDeps) { const limit = Number.isFinite(op?.limit) ? Number(op.limit) : undefined; - return { sessions: await listSessions({ limit }) }; + return { sessions: (await listSessions({ limit })).map((session) => withUsage(session, deps)) }; +} + +/** Copies fields by name, never by spread: a stored record may expose fields as accessors a spread skips. */ +function withUsage( + session: AgentSessionRow, + deps: OperationDeps +): Record & { usage: AgentSessionUsage; runInFlight: boolean } { + const view: Record = {}; + for (const { name } of AGENT_SESSION_ATTRIBUTES) { + const value = session[name as keyof AgentSessionRow]; + if (value !== undefined) view[name] = value; + } + // A cancelled session reads `aborted` before its run has finished; this is what says a prompt can be sent. + return Object.assign(view, { + usage: summarizeSessionUsage(session.messages, session.rejectedCalls), + runInFlight: deps.isRunInFlight(session.session_id), + }); } async function cancelAgentRun(op: any, deps: OperationDeps) { @@ -114,9 +172,8 @@ async function cancelAgentRun(op: any, deps: OperationDeps) { // in `awaiting_approval` or sitting `idle` between turns). Always update the persisted // status so a paused session can still be terminated by the operator. const signalledLiveRun = deps.cancelRun(sessionId); - const wasTerminal = session.status === 'completed' || session.status === 'aborted' || session.status === 'error'; - if (!wasTerminal) await setStatus(sessionId, 'aborted', 'Cancelled by operator'); - return { cancelled: !wasTerminal, signalledLiveRun }; + const cancelled = await cancelSession(sessionId); + return { cancelled, signalledLiveRun }; } async function approveAgentAction(op: any, deps: OperationDeps) { @@ -126,10 +183,11 @@ async function approveAgentAction(op: any, deps: OperationDeps) { throw new ClientError('session_id and approval_id are required', 400); } const approved = op?.approved !== false; + const admitted = deps.admitRun(sessionId); const resolved = await resolveApproval(sessionId, approvalId, approved); // Either decision should resume the loop: an approval means execute the saved tool call; // a denial means hand the operator-rejection observation back to the model so it can adjust. - deps.startRun(sessionId); + deps.startRun(sessionId, admitted); return resolved; } @@ -153,16 +211,18 @@ async function setAgentConfig(op: any, deps: OperationDeps) { 'enabled', 'provider', 'model', - 'maxTurns', 'maxTokens', 'maxToolResultBytes', - 'maxCostUsd', 'autoApprove', 'allowDestructive', 'systemPromptAppend', ]) { if (op?.[key] !== undefined) (patch as any)[key] = op[key]; } + if (op?.maxTurns !== undefined) patch.maxTurns = parseMaxTurns(op.maxTurns); + for (const key of ['maxSessionTokens', 'maxCostUsd'] as const) { + if (op?.[key] !== undefined) patch[key] = parseBudgetCap(key, op[key]); + } return deps.setConfig(patch); } diff --git a/agent/session.ts b/agent/session.ts index 9e6e4f11f3..0e847fc5a6 100644 --- a/agent/session.ts +++ b/agent/session.ts @@ -12,7 +12,8 @@ import { randomUUID } from 'node:crypto'; import { table } from '../resources/databases.ts'; import { SYSTEM_SCHEMA_NAME, SYSTEM_TABLE_NAMES } from '../utility/hdbTerms.ts'; -import type { AgentMessage, AgentRunStatus, AgentSessionRow, ApprovalRequest } from './types.ts'; +import type { ToolCall } from '../resources/models/types.ts'; +import type { AgentCallUsage, AgentMessage, AgentRunStatus, AgentSessionRow, ApprovalRequest } from './types.ts'; let cachedTable: any; @@ -28,6 +29,7 @@ export const AGENT_SESSION_ATTRIBUTES = [ { name: 'createdAt', type: 'number', indexed: true }, { name: 'updatedAt', type: 'number', indexed: true }, { name: 'lastError', type: 'string' }, + { name: 'rejectedCalls' }, ]; export function getAgentSessionTable(): any { @@ -166,31 +168,221 @@ export function shrinkNewestToolResults( }); } +/** + * Every status the loop writes after `running` goes through here, in one write with its closing messages + * and any open approvals answered under the lock: an ending never leaves a tool call without the response + * a provider requires before the next call, and a cancel that landed first keeps its `aborted` status. + */ +export function endRun( + sessionId: string, + end: { + append?: AgentMessage[]; + answerOpenApprovals?: (approval: ApprovalRequest) => string; + rejectedCall?: AgentCallUsage; + /** Answers each tool call of the last assistant turn that is still without a response. */ + answerUnansweredCalls?: (toolName: string) => string; + /** Leave the status alone when no approval is still unresolved under the lock. */ + pauseWhileApprovalsOpen?: boolean; + status: AgentRunStatus; + lastError?: string; + } +): Promise { + return withSessionLock(sessionId, async () => { + const session = await requireSession(sessionId); + if (end.append) session.messages.push(...end.append); + if (end.rejectedCall) (session.rejectedCalls ??= []).push(end.rejectedCall); + const now = Date.now(); + if (end.answerOpenApprovals) closeOpenApprovals(session, end.answerOpenApprovals, now); + if (end.answerUnansweredCalls) answerUnansweredCalls(session, end.answerUnansweredCalls, now); + const nothingToPauseFor = end.pauseWhileApprovalsOpen && !session.pendingApprovals.some((a) => !a.resolved); + if (session.status !== 'aborted' && !nothingToPauseFor) { + session.status = end.status; + session.lastError = end.lastError; + } + session.updatedAt = now; + await getAgentSessionTable().put(session); + return session; + }); +} + +/** + * Cancel a session and answer every open approval in the same write, so the cancel is acknowledged only + * once no approved action is left to run on a later prompt, and no tool call is left without a response. + * Decided under the lock: a session a run already finished is left as it is. Returns whether it cancelled. + */ +export function cancelSession(sessionId: string): Promise { + return withSessionLock(sessionId, async () => { + const session = await requireSession(sessionId); + if (session.status === 'completed' || session.status === 'aborted' || session.status === 'error') return false; + const now = Date.now(); + closeOpenApprovals(session, (approval) => cancelledObservation(approval.toolName), now); + session.status = 'aborted'; + session.lastError = 'Cancelled by operator'; + session.updatedAt = now; + await getAgentSessionTable().put(session); + return true; + }); +} + +/** + * Inserted right after the turn's existing responses, so the transcript stays sendable even when a prompt + * was appended after the turn. Returns whether it answered anything. + */ +function answerUnansweredCalls(session: AgentSessionRow, answer: (toolName: string) => string, now: number): boolean { + const unanswered = unansweredCalls(session); + if (!unanswered) return false; + session.messages.splice( + unanswered.at, + 0, + ...unanswered.calls.map((call) => ({ + role: 'tool' as const, + content: answer(call.name), + toolCallId: call.id, + createdAt: now, + })) + ); + return true; +} + +/** Decided on the row read under the lock, so an approval another writer already claimed is never answered twice. */ +function closeOpenApprovals( + session: AgentSessionRow, + answer: (approval: ApprovalRequest) => string, + now: number +): void { + for (const approval of session.pendingApprovals) { + if (approval.consumed) continue; + session.messages.splice(responseIndex(session, approval.toolCallId), 0, { + role: 'tool', + content: answer(approval), + toolCallId: approval.toolCallId, + createdAt: now, + }); + approval.approved ??= false; + approval.resolvedAt ??= now; + approval.resolved = true; + approval.consumed = true; + } +} + +/** + * Runs are serialized per session, so a tool call of the latest turn still without a response belongs to + * a run that died before writing it: an approval claimed but never given its result, or an inline call. + * Each is answered `interrupted`. Nothing is run again. + */ +export async function answerInterruptedCalls(sessionId: string): Promise { + const current = await getSession(sessionId); + if (!current || !unansweredCalls(current)) return; + await withSessionLock(sessionId, async () => { + const session = await requireSession(sessionId); + const now = Date.now(); + const interrupted = (toolName: string) => + JSON.stringify({ ok: false, error: 'interrupted', outcome: 'unknown', tool: toolName }); + if (!answerUnansweredCalls(session, interrupted, now)) return; + session.updatedAt = now; + await getAgentSessionTable().put(session); + }); +} + +/** + * Matched on the latest turn only, as a provider may reuse a tool-call id. A call still waiting on its + * approval is left for the drain or `closeOpenApprovals` to answer. + */ +function unansweredCalls(session: AgentSessionRow): { calls: ToolCall[]; at: number } | undefined { + let turn = session.messages.length - 1; + while (turn >= 0 && session.messages[turn].role !== 'assistant') turn--; + if (turn < 0) return undefined; + const answered = new Set(session.messages.slice(turn + 1).map((message) => message.toolCallId)); + const awaitingApproval = new Set(session.pendingApprovals.filter((a) => !a.consumed).map((a) => a.toolCallId)); + const calls = (session.messages[turn].toolCalls ?? []).filter( + (call) => !answered.has(call.id) && !awaitingApproval.has(call.id) + ); + if (calls.length === 0) return undefined; + return { calls, at: endOfResponses(session, turn) }; +} + +/** + * Where a response to `toolCallId` goes: after the responses that already follow the latest turn calling + * it, so a prompt appended after that turn stays after its responses, as a provider requires. + */ +function responseIndex(session: AgentSessionRow, toolCallId: string): number { + let turn = session.messages.length - 1; + while (turn >= 0 && !session.messages[turn].toolCalls?.some((call) => call.id === toolCallId)) turn--; + return turn < 0 ? session.messages.length : endOfResponses(session, turn); +} + +function endOfResponses(session: AgentSessionRow, turn: number): number { + let at = turn + 1; + while (at < session.messages.length && session.messages[at].role === 'tool') at++; + return at; +} + +/** An approval's result, which a prompt appended while it waited must not precede. */ +export function answerToolCall(sessionId: string, toolCallId: string, content: string): Promise { + return withSessionLock(sessionId, async () => { + const session = await requireSession(sessionId); + const now = Date.now(); + session.messages.splice(responseIndex(session, toolCallId), 0, { + role: 'tool', + content, + toolCallId, + createdAt: now, + }); + session.updatedAt = now; + await getAgentSessionTable().put(session); + return session; + }); +} + +/** + * Park a gated tool call for the operator. On a session a cancel already marked `aborted`, the call is + * answered `cancelled_by_operator` instead, so the cancel is not undone and the call is not left open. + */ export function addPendingApproval( sessionId: string, approval: Omit -): Promise { +): Promise { return withSessionLock(sessionId, async () => { const session = await requireSession(sessionId); - const entry: ApprovalRequest = { ...approval, id: randomUUID(), createdAt: Date.now() }; - session.pendingApprovals.push(entry); - session.status = 'awaiting_approval'; - session.updatedAt = Date.now(); + const now = Date.now(); + let entry: ApprovalRequest | undefined; + if (session.status === 'aborted') { + session.messages.push({ + role: 'tool', + content: cancelledObservation(approval.toolName), + toolCallId: approval.toolCallId, + createdAt: now, + }); + } else { + entry = { ...approval, id: randomUUID(), createdAt: now }; + session.pendingApprovals.push(entry); + session.status = 'awaiting_approval'; + } + session.updatedAt = now; await getAgentSessionTable().put(session); return entry; }); } -export function markApprovalConsumed(sessionId: string, approvalId: string): Promise { +export function cancelledObservation(toolName: string): string { + return JSON.stringify({ ok: false, error: 'cancelled_by_operator', tool: toolName }); +} + +/** + * Claim a resolved approval before its action runs; only the first claim wins. A cancel or budget stop + * therefore closes only approvals whose action has not started, and no tool call gets two responses. + * The cost is at-most-once: a crash between the claim and the result write leaves that call unanswered + * rather than running an operator-approved action twice. + */ +export function claimApproval(sessionId: string, approvalId: string): Promise { return withSessionLock(sessionId, async () => { const session = await requireSession(sessionId); const entry = session.pendingApprovals.find((a) => a.id === approvalId); - if (!entry) throw new Error(`No pending approval ${approvalId} on session ${sessionId}`); - if (!entry.resolved) throw new Error(`Approval ${approvalId} not yet resolved`); - if (entry.consumed) return; + if (!entry?.resolved || entry.consumed) return false; entry.consumed = true; session.updatedAt = Date.now(); await getAgentSessionTable().put(session); + return true; }); } diff --git a/agent/types.ts b/agent/types.ts index 2676864bd6..d935d9109d 100644 --- a/agent/types.ts +++ b/agent/types.ts @@ -17,6 +17,26 @@ export interface AgentMessage { toolCallId?: string; /** Wall-clock timestamp of when this item entered the transcript. */ createdAt: number; + /** Assistant messages only: what the model call that produced this message reported. */ + usage?: AgentCallUsage; +} + +/** One model call's usage as reported; a field the backend or pricing did not supply is absent. */ +export interface AgentCallUsage { + promptTokens?: number; + completionTokens?: number; + costUsd?: number; +} + +/** + * A session's usage folded from its transcript (`summarizeSessionUsage`). A total is present only + * while every call in the session reported it, so an absent total with `calls > 0` means unmeasured. + */ +export interface AgentSessionUsage { + calls: number; + promptTokens?: number; + completionTokens?: number; + costUsd?: number; } export interface ApprovalRequest { @@ -46,6 +66,8 @@ export interface AgentSessionRow { createdAt: number; updatedAt: number; lastError?: string; + /** Usage of billed model calls whose reply was rejected and kept out of `messages`; folded into the session's usage. */ + rejectedCalls?: AgentCallUsage[]; } export interface AgentToolHandler { @@ -90,7 +112,10 @@ export interface AgentConfig { maxTokens: number; /** Largest tool result, in bytes of serialized JSON, the loop appends to the transcript. */ maxToolResultBytes: number; - maxCostUsd: number; + /** Per-session cap on prompt + completion tokens. Unset: uncapped. */ + maxSessionTokens?: number; + /** Per-session cap in USD, from the serving model entry's `pricing`. Unset: uncapped. */ + maxCostUsd?: number; autoApprove: boolean; allowDestructive: boolean; user: string; diff --git a/config-root.schema.json b/config-root.schema.json index f8a3315d92..623041f19c 100644 --- a/config-root.schema.json +++ b/config-root.schema.json @@ -734,10 +734,15 @@ "maximum": 1048576, "description": "Largest tool result, in bytes of JSON, the agent adds to its conversation. A larger result is cut to this size with a note telling the model how to ask for less, and the filesystem read tools return pages of at most half this size. Tunable at runtime via set_agent_config. Default: 65536" }, + "maxSessionTokens": { + "type": "integer", + "minimum": 1, + "description": "Per-session cap on prompt plus completion tokens, checked after each model call; a session at the cap makes no further calls. A model that reports no token usage stops the session instead of running uncapped. A call that fails is not counted. Default: none" + }, "maxCostUsd": { "type": "number", "minimum": 0, - "description": "Per-session cost budget. NOT YET ENFORCED — surfacing operator intent until #612 ships per-call cost telemetry. Default: 5.00" + "description": "Per-session cost cap in USD, checked after each model call and priced from the `pricing` on the models.generative entry that served it. A call with no pricing stops the session instead of running uncapped. A call that fails is not counted. Default: none" }, "autoApprove": { "type": "boolean", diff --git a/integrationTests/server/agent-budget.test.ts b/integrationTests/server/agent-budget.test.ts new file mode 100644 index 0000000000..1cd0f14861 --- /dev/null +++ b/integrationTests/server/agent-budget.test.ts @@ -0,0 +1,207 @@ +/** + * The built-in agent's per-session usage and caps, end to end (#3040): model pricing from config, the + * facade pricing each call, usage folded from the durable transcript, and the operations API. + */ +import { suite, test, before, after } from 'node:test'; +import { strictEqual, ok, deepStrictEqual } from 'node:assert'; +import { createServer, type IncomingMessage, type ServerResponse, type Server } from 'node:http'; +import type { AddressInfo } from 'node:net'; +import { startHarper, teardownHarper } from '@harperfast/integration-testing'; +// @ts-expect-error utils/client.mjs has no type declarations; runtime resolves fine +import { createApiClient } from '../apiTests/utils/client.mjs'; + +const PROMPT_TOKENS = 1000; +const COMPLETION_TOKENS = 200; +// 1000 prompt tokens at $10/Mtok plus 200 completion tokens at $50/Mtok. +const COST_PER_CALL = 0.02; +const TOOL_ROUNDS_BEFORE_ANSWER = 2; + +interface FakeModel { + baseUrl: string; + calls: () => number; + close: () => Promise; +} + +/** + * An OpenAI-compatible model that asks for a tool until the conversation holds two tool responses, + * then answers. The tool name is unknown to the agent, so each call runs nothing and only feeds the + * next model turn. + */ +async function startFakeModel(): Promise { + let calls = 0; + const server: Server = createServer((req: IncomingMessage, res: ServerResponse) => { + if (req.method !== 'POST' || req.url !== '/v1/chat/completions') { + res.writeHead(404); + res.end(); + return; + } + let raw = ''; + req.on('data', (chunk) => (raw += chunk)); + req.on('end', () => { + calls++; + const { messages } = JSON.parse(raw) as { messages: { role: string }[] }; + const toolResponses = messages.filter((message) => message.role === 'tool').length; + const answer = toolResponses >= TOOL_ROUNDS_BEFORE_ANSWER; + res.writeHead(200, { 'Content-Type': 'application/json' }); + res.end( + JSON.stringify({ + choices: [ + { + finish_reason: answer ? 'stop' : 'tool_calls', + message: answer + ? { role: 'assistant', content: 'done' } + : { + role: 'assistant', + content: null, + tool_calls: [ + { + id: `call-${calls}`, + type: 'function', + function: { name: 'fake_lookup', arguments: '{}' }, + }, + ], + }, + }, + ], + usage: { prompt_tokens: PROMPT_TOKENS, completion_tokens: COMPLETION_TOKENS }, + }) + ); + }); + }); + await new Promise((resolve) => server.listen(0, '127.0.0.1', resolve)); + const { port } = server.address() as AddressInfo; + return { + baseUrl: `http://127.0.0.1:${port}/v1`, + calls: () => calls, + close: () => new Promise((resolve, reject) => server.close((err) => (err ? reject(err) : resolve()))), + }; +} + +const TERMINAL = new Set(['completed', 'error', 'aborted', 'awaiting_approval']); + +suite('agent per-session usage and caps (#3040)', (ctx: any) => { + let fake: FakeModel; + let client: any; + + const op = async (body: object, status = 200) => { + const response = await client.req().send(body); + strictEqual(response.status, status, response.text); + return response.body; + }; + + async function runToEnd(message: string) { + const { session_id: sessionId } = await op({ operation: 'agent_prompt', message }); + const deadline = Date.now() + 30_000; + for (;;) { + const session = await op({ operation: 'get_agent_session', session_id: sessionId }); + if (TERMINAL.has(session.status)) return session; + ok(Date.now() < deadline, `session ${sessionId} still ${session.status} after 30s`); + await new Promise((resolve) => setTimeout(resolve, 100)); + } + } + + const toolResponsesAnswerEveryCall = (session: any) => { + const asked = session.messages.flatMap((m: any) => (m.toolCalls ?? []).map((call: any) => call.id)); + const answered = session.messages.filter((m: any) => m.role === 'tool').map((m: any) => m.toolCallId); + deepStrictEqual(answered.sort(), asked.sort()); + }; + + before(async () => { + fake = await startFakeModel(); + const entry = { backend: 'openai', baseUrl: fake.baseUrl, apiKey: 'fake-key', model: 'fake-model' }; + await startHarper(ctx, { + config: { + models: { + generative: { + default: { ...entry, pricing: { inputPerMtok: 10, outputPerMtok: 50 } }, + unpriced: entry, + }, + }, + agent: { enabled: true }, + }, + env: {}, + }); + client = createApiClient(ctx.harper); + }); + + after(async () => { + try { + await teardownHarper(ctx); + } finally { + await fake.close(); + } + }); + + test('an uncapped run reports usage that matches what the model reported', async () => { + const callsBefore = fake.calls(); + const session = await runToEnd('look something up'); + const calls = fake.calls() - callsBefore; + strictEqual(session.status, 'completed', session.lastError); + strictEqual(calls, TOOL_ROUNDS_BEFORE_ANSWER + 1); + strictEqual(session.usage.calls, calls); + strictEqual(session.usage.promptTokens, PROMPT_TOKENS * calls); + strictEqual(session.usage.completionTokens, COMPLETION_TOKENS * calls); + ok(Math.abs(session.usage.costUsd - COST_PER_CALL * calls) < 1e-9, String(session.usage.costUsd)); + + const { sessions } = await op({ operation: 'list_agent_sessions' }); + const listed = sessions.find((s: any) => s.session_id === session.session_id); + deepStrictEqual(listed.usage, session.usage); + }); + + test('maxSessionTokens stops the session after the call that reaches it, with every tool call answered', async () => { + await op({ operation: 'set_agent_config', maxSessionTokens: 2000 }); + try { + const callsBefore = fake.calls(); + const session = await runToEnd('look again'); + strictEqual(fake.calls() - callsBefore, 2, 'the second call reaches 2400 tokens'); + strictEqual(session.status, 'error'); + ok( + /agent\.maxSessionTokens=2000 reached: this session used 2400 tokens/.test(session.lastError), + session.lastError + ); + toolResponsesAnswerEveryCall(session); + + const before = fake.calls(); + await op({ operation: 'agent_prompt', session_id: session.session_id, message: 'continue' }); + // Asserting a non-event: the session was already `error`, so there is no state change to wait on. + await new Promise((resolve) => setTimeout(resolve, 1000)); + const after = await op({ operation: 'get_agent_session', session_id: session.session_id }); + strictEqual(after.status, 'error'); + strictEqual(after.messages.at(-1).content, 'continue'); + strictEqual(fake.calls(), before, 'a capped session makes no further model call'); + } finally { + await op({ operation: 'set_agent_config', maxSessionTokens: null }); + } + }); + + test('maxCostUsd stops the session from the priced cost', async () => { + await op({ operation: 'set_agent_config', maxCostUsd: 0.03 }); + try { + const session = await runToEnd('price check'); + strictEqual(session.status, 'error'); + strictEqual(session.usage.calls, 2); + ok(/agent\.maxCostUsd=0\.03 reached/.test(session.lastError), session.lastError); + toolResponsesAnswerEveryCall(session); + } finally { + await op({ operation: 'set_agent_config', maxCostUsd: null }); + } + }); + + test('maxCostUsd against an unpriced model stops after one call and names the missing pricing', async () => { + await op({ operation: 'set_agent_config', model: 'unpriced', maxCostUsd: 5 }); + try { + const session = await runToEnd('no prices here'); + strictEqual(session.status, 'error'); + strictEqual(session.usage.calls, 1); + strictEqual(session.usage.costUsd, undefined); + ok(/maxCostUsd=5 cannot be enforced: .*declare pricing/.test(session.lastError), session.lastError); + } finally { + await op({ operation: 'set_agent_config', model: 'default', maxCostUsd: null }); + } + }); + + test('set_agent_config refuses a malformed cap', async () => { + const body = await op({ operation: 'set_agent_config', maxSessionTokens: -1 }, 400); + ok(/agent\.maxSessionTokens must be a positive integer/.test(JSON.stringify(body)), JSON.stringify(body)); + }); +}); diff --git a/resources/models/DESIGN.md b/resources/models/DESIGN.md index e5678e104c..cddcfc857f 100644 --- a/resources/models/DESIGN.md +++ b/resources/models/DESIGN.md @@ -45,3 +45,7 @@ The built-in decider passes `persist: true` only when the directive names a `dec `applyFits` (`calibrationStore.ts`) reads the in-memory cache only. A key never loaded returns the decision raw and starts one load; an entry older than the freshness window keeps serving while it refreshes, so a population is uncorrected only before its first load or after a load times out or fails, at most `maxLoads` in flight, each counted until its read settles, so a store whose reads never settle holds the slots and leaves that node's decisions uncorrected rather than piling up reads. Faults are logged through `safeErrorMessage` and return the decision raw. Reliability scores the value a decision returned, which on tied scores need not be the schema-first one, so a fit is judged on the decisions it will actually serve. A reload that leaves the calibration block unchanged keeps the cache and the scheduled job, so frequent unrelated config changes never postpone the first run. Calibration runs after the analytics row, outside the candidate loop, so it never bills a second call. `requires: ['calibrated']` still routes on the backend's own claim. A schema with a no-match leaf is never fit-calibrated, because its no-match score would stay raw under the flag; its population gets reports whose reason is `no-match-schema`, never an eligible version. Versions in `hdb_model_calibrations` (`replicate: true`) are immutable. `fitId` covers the input digest, the policy and the result, so a repeated run writes nothing. A row is applied only when its `key`, its `rank` prefix and the key recomputed from its `population` and `field` all match the key it was read under, so a malformed replicated row cannot cross a population. The reader takes the greatest `rank` (`key | evidenceAt | fittedAt | fitId`) in one descending prefix read. `evidenceAt` is the newest decision or truth fact a version read, so the version that saw newer evidence wins whatever node wrote it and in whatever order rows arrive; two runs on different nodes that both step past the same previous version tie on `evidenceAt`, and the later `fittedAt` wins until a run that reads both writes a successor. A new version's `evidenceAt` is also at least one above the version it replaces, so a correction stamped by a node with a slower clock still supersedes it. A version is written whenever a population's input digest changes, even with no labels left, so a correction revokes an older eligible version. Rows expire after their decisions do. `calibrate()` can only lower a run's budgets, never raise them or change its policy, and a run invalidates only the cache entries of the keys it wrote. A population's decisions are read newest first through the indexed `populationRank` (`population | expiresAt | id`) with a limit, so a large population is never loaded in full. The limit is the newest decisions whose examples, at their cost across every field, fit in the configured `maxBytes` beside the population's own head: a size fixed by the configuration and the schema, so a population too large for the byte budget is still refit, on the same sample each run. A run that cannot afford that sample, because `calibrate()` lowered `maxBytes`, stops on the byte budget and writes nothing, because a smaller sample could revoke a fit its evidence still supports. Heads and a population's sample are budgeted separately, each against `maxBytes`, so other populations in a run never crowd a sample out, and a run holds at most about twice `maxBytes`. A head is renewed only by a run that read decisions for it, so a population whose decisions have all expired lapses instead of staying in the table. Discovery uses at most half of `maxDecisions`, split between the newest decisions and a saved cursor that keeps the ids on its boundary, and splits its population budget the same way. A row that hits a cap is left for the next run, a newly discovered population gets a head at once, and ties at one timestamp never stall the cursor. No discovery page asks for more rows than the scan may still take. Decisions read for fitting come out of the rest of the same budget, so one run reads at most `maxDecisions`, plus one decision per deferred population. A population is read up to one decision past what is left of the budget, and deferred to a later run when that extra decision exists, so a truncated read is never fitted and a population that exactly fits is not deferred forever; a population a run could not finish, deferred or stopped by the byte or time budget, moves to the back of the queue so it cannot hold up the ones behind it, and configuration requires `maxDecisions` of at least twice `maxExamplesPerKey` plus 2, so the first population fitted in a run always has room for a full sample, because a partial sample would write a version that revokes a fit its evidence still supports. A failed run fails its scheduled job. A run walks known populations least recently fitted first, then those discovered newest first, under global decision, population, byte and time budgets, and writes a population only after its report is complete. Its job is registered with `registerInternalJobs` (`resources/scheduler/engine.ts`), under a reserved prefix no component can register or unregister. Enforced by `unitTests/resources/models/calibration.test.js`, `calibrationStore.test.js`, `sourceIdentity.test.js` and `unitTests/resources/scheduler/engine.test.js`. + +## A generative call is priced by the entry that served it + +`GenerateResult.costUsd` is framework-owned, like `usage`: `Models.generate` computes it from the `pricing` of the backend instance that answered (`getBackendPricing`, set in `bootstrap.ts` `publishEntry`) and drops any `usage` or `costUsd` the backend put on its output. Pricing by the requested logical name would bill a fallback at the primary's rate (harper#3040). Both token counts must be non-negative integers (`isTokenCount`, shared with backend parsing and the auto loop) or the call is unpriced, so a partial report is never under-priced. The arithmetic runs after the candidate loop's try/catch, so an accounting fault cannot become a second, billed fallback call. `pricing` is accepted only on generative entries and is excluded from `sourceFingerprint`, because a price does not change what a model returns. A price edit rebuilds a built-in entry's backend on reload; a module-backed entry, like any change to one, takes it on restart. With `toolMode: 'auto'`, `costUsd` on the returned result is the final round's, like `usage`; the loop sums every round for `maxCostUsd`, and a round it cannot measure (no token counts under `maxToolTokens`, unpriced under `maxCostUsd`) stops it with a `BudgetExceededError`, as the agent's caps fail closed. Enforced by `unitTests/resources/models/pricing.test.js`. diff --git a/resources/models/Models.ts b/resources/models/Models.ts index f18d8ccef3..8115a474a6 100644 --- a/resources/models/Models.ts +++ b/resources/models/Models.ts @@ -3,6 +3,7 @@ import { contextStorage } from '../transaction.ts'; import { defineBackend, getBackend, + getBackendPricing, getBackendSource, ModelBackendNotFoundError, registerBackend as registerBackendImpl, @@ -27,7 +28,7 @@ import { getModelCallAnalyticsWriter, type ModelCallAnalyticsWriter, type ModelC import { recordAction } from '../analytics/write.ts'; import { ClientError, ServerError } from '../../utility/errors/hdbError.ts'; import { runAgentLoop, runAgentLoopStream } from './agentLoop.ts'; -import { assignFiniteTokenCount, ChoiceScoringUnsupportedError } from './backendHelpers.ts'; +import { assignFiniteTokenCount, ChoiceScoringUnsupportedError, isTokenCount } from './backendHelpers.ts'; import { DecisionContractError, DecisionInputError, @@ -71,6 +72,7 @@ import type { GenerateResult, ModelBackend, ModelKind, + ModelPricing, ModelRouter, ModelCallResult, Models as ModelsContract, @@ -230,7 +232,7 @@ export class Models implements ModelsContract { let firstError: unknown = undefined; let hasError = false; for (const backend of resolved.candidates) { - let completed: GenerateResult | undefined; + let completed: { output: GenerateResult; usage?: TokenUsage } | undefined; // Don't spend a backend call on an already-cancelled request — between // candidates the caller may have aborted (and at entry it may already be). signal?.throwIfAborted(); @@ -240,11 +242,9 @@ export class Models implements ModelsContract { const backendOpts = toBackendOpts(opts, signal, accounting); const result = await backend.generate!(input, backendOpts); if (result.status !== 'completed') throw new ModelPendingNotSupportedError(backend.name); + if (!result.output) throw new ServerError(`Model backend '${backend.name}' completed without output`); this.#record(backend, 'generate', opts.model, accounting, opts, result, attemptStart); - // Propagate usage onto the returned GenerateResult so callers (notably the - // `toolMode: 'auto'` loop's budget tracker) can read cumulative tokens without - // re-querying analytics. Pure pass-through — backend usage is the source of truth. - completed = result.usage ? { ...result.output, usage: result.usage } : result.output; + completed = { output: result.output, usage: result.usage }; } catch (err) { this.#recordFailure(backend, 'generate', opts.model, accounting, opts, attemptStart, err); if (!hasError) { @@ -255,7 +255,7 @@ export class Models implements ModelsContract { } if (completed) { reportServed(opts, backend); - return completed; + return finalizeGenerateResult(completed.output, completed.usage, getBackendPricing(backend)); } } throw firstError; @@ -681,6 +681,28 @@ function usageFromError(err: unknown): TokenUsage | undefined { return Object.keys(usage).length > 0 ? usage : undefined; } +/** Runs after the candidate try/catch, so an accounting fault cannot turn a billed success into a fallback call. */ +function finalizeGenerateResult( + output: GenerateResult, + usage: TokenUsage | undefined, + pricing: ModelPricing | undefined +): GenerateResult { + const { usage: _backendUsage, costUsd: _backendCost, ...result } = output; + const generateResult: GenerateResult = result; + if (usage) generateResult.usage = usage; + const costUsd = callCostUsd(usage, pricing); + if (costUsd !== undefined) generateResult.costUsd = costUsd; + return generateResult; +} + +/** Cost from both token counts; a missing or invalid count leaves the call unpriced rather than under-priced. */ +function callCostUsd(usage: TokenUsage | undefined, pricing: ModelPricing | undefined): number | undefined { + if (!pricing || !isTokenCount(usage?.promptTokens) || !isTokenCount(usage?.completionTokens)) return undefined; + const costUsd = + (usage.promptTokens * pricing.inputPerMtok + usage.completionTokens * pricing.outputPerMtok) / 1_000_000; + return Number.isFinite(costUsd) && costUsd >= 0 ? costUsd : undefined; +} + function buildRecord( backend: ModelBackend | undefined, method: CallMethod, diff --git a/resources/models/agentLoop.ts b/resources/models/agentLoop.ts index 7ccb08cbdb..0803fcd066 100644 --- a/resources/models/agentLoop.ts +++ b/resources/models/agentLoop.ts @@ -37,8 +37,8 @@ * supply a handler for is a caller config bug (hard `ClientError(400)`); an UNdeclared * name is a model hallucination, surfaced as a recoverable tool error (`toolErrorMode` * decides recover-vs-abort). - * - A token/cost budget set against a backend that reports no `usage` warns once that it - * is unenforceable rather than silently no-opping. + * - A token budget against a backend that reports no `usage`, or a cost budget against an + * entry with no `pricing`, stops the loop with a `BudgetExceededError` rather than running uncapped. * - An incomplete streamed tool call (id without a name) is recorded in the trace + warned, * not silently dropped. * - `toolErrorMode: 'abort'` throws BEFORE the conversation sink sees the round's tool @@ -54,6 +54,7 @@ */ import { ClientError, ServerError } from '../../utility/errors/hdbError.ts'; import { logger } from '../../utility/logging/logger.ts'; +import { isTokenCount } from './backendHelpers.ts'; import type { AccountingContext, GenerateChunk, @@ -109,9 +110,23 @@ export async function runAgentLoop(args: RunAgentLoopArgs): Promise { + if (maxToolTokens !== undefined && totalTokens >= maxToolTokens) { + throw new BudgetExceededError( + 'tokens', + `agent loop exceeded maxToolTokens=${maxToolTokens} (cumulative=${totalTokens})`, + trace + ); + } + if (maxCostUsd !== undefined && totalCostUsd >= maxCostUsd) { + throw new BudgetExceededError( + 'cost', + `agent loop exceeded maxCostUsd=${maxCostUsd} (cumulative=${totalCostUsd})`, + trace + ); + } + }; // `conversation` is a one-way SINK for NEW turns produced by this loop — the loop // does NOT re-append the caller's input messages, even when they include user @@ -140,6 +155,7 @@ export async function runAgentLoop(args: RunAgentLoopArgs): Promise= cap`, so `maxCostUsd: 0` blocks every call rather - // than dormantly admitting all of them until a real rate card lands and then - // abruptly blocking everything. Token trip uses `>=` for symmetry. + // Both caps trip at `>= cap`. // // Discarded-content asymmetry: a TERMINAL round that trips returns // BudgetExceededError instead of the final assistant content. The content is @@ -184,21 +201,8 @@ export async function runAgentLoop(args: RunAgentLoopArgs): Promise= maxToolTokens) { - throw new BudgetExceededError( - 'tokens', - `agent loop exceeded maxToolTokens=${maxToolTokens} (cumulative=${totalTokens})`, - trace - ); - } - if (maxCostUsd !== undefined && totalCostUsd >= maxCostUsd) { - throw new BudgetExceededError( - 'cost', - `agent loop exceeded maxCostUsd=${maxCostUsd} (cumulative=${totalCostUsd})`, - trace - ); - } + totalCostUsd += result.costUsd ?? 0; + throwIfBudgetReached(); const calls = result.toolCalls; // Backends (notably OpenAI) leave `content` as `null` / undefined on tool-call @@ -859,32 +863,9 @@ export class ToolHandlerError extends ServerError { } } -/** - * Cost computation hook. v1 returns 0 — no per-model rate card is wired today — - * so `maxCostUsd` never trips in production. The cap, the trip path, and the - * `BudgetExceededError({kind: 'cost'})` shape ARE wired and exercised by tests - * that inject a non-zero function via `_setComputeCallCostUsdForTests`. When the - * rate card lands, replace this implementation; no surface change needed. - */ -let computeCallCostUsd: (usage: TokenUsage | undefined, model: string | undefined) => number = () => 0; - -/** - * Test-only override. Public callers must not depend on this — it exists so unit - * tests can prove the `maxCostUsd` trip path works end-to-end before a real rate - * card lands. Leading underscore marks the intent. - */ -export function _setComputeCallCostUsdForTests( - fn: (usage: TokenUsage | undefined, model: string | undefined) => number -): void { - computeCallCostUsd = fn; -} - -/** - * Reset the cost function to the v1 stub. Pair with `_setComputeCallCostUsdForTests` - * in test `afterEach` so suites don't leak state into each other. - */ -export function _resetComputeCallCostUsdForTests(): void { - computeCallCostUsd = () => 0; +/** Both counts, as the facade requires for pricing: `{}` or a half report would silently under-count the cap. */ +function reportsTokenCounts(usage: TokenUsage | undefined): boolean { + return isTokenCount(usage?.promptTokens) && isTokenCount(usage?.completionTokens); } function sumTokens(usage: TokenUsage | undefined): number { diff --git a/resources/models/backendHelpers.ts b/resources/models/backendHelpers.ts index 4e50fe7ab1..522021852c 100644 --- a/resources/models/backendHelpers.ts +++ b/resources/models/backendHelpers.ts @@ -78,9 +78,11 @@ export function assignFiniteTokenCount( key: 'promptTokens' | 'completionTokens' | 'embeddingTokens', value: unknown ): void { - if (typeof value !== 'number') return; - if (!Number.isFinite(value) || value < 0 || !Number.isInteger(value)) return; - usage[key] = value; + if (isTokenCount(value)) usage[key] = value; +} + +export function isTokenCount(value: unknown): value is number { + return typeof value === 'number' && Number.isFinite(value) && value >= 0 && Number.isInteger(value); } // Body-read caps. A hostile or buggy upstream that returns a multi-GiB body diff --git a/resources/models/backendRegistry.ts b/resources/models/backendRegistry.ts index f564b27ad9..800943c820 100644 --- a/resources/models/backendRegistry.ts +++ b/resources/models/backendRegistry.ts @@ -6,6 +6,7 @@ import type { ModelBackend, ModelCapabilities, ModelKind, + ModelPricing, ToolCall, } from './types.ts'; @@ -68,6 +69,19 @@ export function getBackendSource(backend: ModelBackend): string | undefined { return sources.get(backend); } +// Keyed by instance, like `sources`, so a call is priced by the entry that built the backend that +// served it, and a reload that rebuilds the backend drops the old price with it. +const pricing = new WeakMap(); + +export function setBackendPricing(backend: ModelBackend, modelPricing: ModelPricing | undefined): void { + if (modelPricing) pricing.set(backend, modelPricing); + else pricing.delete(backend); +} + +export function getBackendPricing(backend: ModelBackend): ModelPricing | undefined { + return pricing.get(backend); +} + // Async-context scoped: a module-global slot would divert unrelated registrations and collide // concurrent constructions. const captureScope = new AsyncLocalStorage(); diff --git a/resources/models/bootstrap.ts b/resources/models/bootstrap.ts index c7ee884f82..d7c8a9d42d 100644 --- a/resources/models/bootstrap.ts +++ b/resources/models/bootstrap.ts @@ -46,13 +46,14 @@ import { getBackend, removeIfCurrent, replaceIfCurrent, + setBackendPricing, setBackendSource, type CapturedInstall, } from './backendRegistry.ts'; import { getSharedRootConfigWatcher, RootConfigWatcher } from '../../config/RootConfigWatcher.ts'; import { validateModelsBlock } from '../../validation/configValidator.ts'; import { registerGenerativeDecisionBackend, type GenerativeDecisionConfig } from './generativeDecision.ts'; -import type { ModelBackend, ModelKind } from './types.ts'; +import type { ModelBackend, ModelKind, ModelPricing } from './types.ts'; /** * Field names treated as credentials. When present in config as a literal @@ -82,6 +83,8 @@ interface ModelEntry { requireStructuredOutput?: boolean; /** Ordered fallback group: other logical names tried, in order, after this one (#1326). */ fallback?: string[]; + /** What the operator pays for this generative model; prices `Models.generate` calls it serves. */ + pricing?: ModelPricing; } interface ModelsConfig { @@ -293,6 +296,7 @@ function publishEntry( const key = slotKey(kind, logicalName); const sourceConfig = desiredEntry.configJson ? JSON.parse(desiredEntry.configJson) : entry; setBackendSource(backend, sourceFingerprint(kind, logicalName, sourceConfig)); + if (kind === 'generative') setBackendPricing(backend, entry.pricing); for (const extra of extras) setBackendSource(extra.backend, sourceFingerprint(extra.kind, extra.logicalName, sourceConfig)); // Boot overwrites occupants (the documented contract); a reload replaces only what this @@ -516,8 +520,12 @@ async function applyModels(block: ModelsConfig | null | undefined, isBoot: boole function installedConfiguration(): Record | undefined { const configuration: Record = {}; for (const [key, slot] of installedSlots) { - if (slot.backend) - configuration[key] = JSON.parse(slot.configJson ?? JSON.stringify(expandEnvVarsDeep(JSON.parse(slot.entryJson)))); + if (!slot.backend) continue; + // What a call costs is not part of what produced a decision. + const { pricing: _pricing, ...entry } = JSON.parse( + slot.configJson ?? JSON.stringify(expandEnvVarsDeep(JSON.parse(slot.entryJson))) + ); + configuration[key] = entry; } return Object.keys(configuration).length > 0 ? configuration : undefined; } diff --git a/resources/models/decisionStore.ts b/resources/models/decisionStore.ts index 7d7935cb75..4395046342 100644 --- a/resources/models/decisionStore.ts +++ b/resources/models/decisionStore.ts @@ -421,7 +421,7 @@ export function sourceFingerprint(kind: string, logicalName: string | undefined, canonicalJson({ kind, logicalName: logicalName ?? null, - config: withoutNamedCredentials(withoutFallback(config)), + config: withoutNamedCredentials(withoutRoutingAndPricing(config)), }) ) .digest('hex'); @@ -444,10 +444,10 @@ function withoutNamedCredentials(value: unknown): unknown { return kept; } -/** Which backend serves after this one is routing, not what produces this one's scores. */ -function withoutFallback(config: unknown): unknown { +/** Which backend serves after this one, and what a call costs, do not change what this one scores. */ +function withoutRoutingAndPricing(config: unknown): unknown { if (!config || typeof config !== 'object' || Array.isArray(config)) return config; - const { fallback: _fallback, ...rest } = config as Record; + const { fallback: _fallback, pricing: _pricing, ...rest } = config as Record; return rest; } diff --git a/resources/models/types.ts b/resources/models/types.ts index 07e871e481..67e1f0ef21 100644 --- a/resources/models/types.ts +++ b/resources/models/types.ts @@ -344,22 +344,19 @@ export type GenerateOpts = { maxToolIterations?: number; /** * `toolMode: 'auto'` only. Cumulative prompt+completion token cap across all - * iterations. Trips `BudgetExceededError({kind: 'tokens'})` when exceeded. - * - * Best-effort: enforcement depends on the backend reporting `usage`. Against a - * backend that returns no usage the cap is unmeasurable — the loop warns once and - * continues (it does NOT silently pretend the cap is in force); `maxToolIterations` - * remains the hard bound. Not yet supported on `generateStream` (throws 501). + * iterations. Throws `BudgetExceededError({kind: 'tokens'})` once the total reaches + * it, so `0` allows no call. A round whose backend reports no token counts makes the + * cap unmeasurable and throws the same error. Not yet supported on `generateStream` + * (throws 501). */ maxToolTokens?: number; /** - * `toolMode: 'auto'` only. Cumulative cost cap across all iterations. The v1 - * cost-per-call function returns 0 (no rate card yet) so this trips only when - * a test or follow-up wires a non-zero function — the seam is in place. - * - * Same best-effort caveat as `maxToolTokens`: cost is derived from backend usage, - * so a backend reporting no usage makes the cap unmeasurable (warn-once, continue). - * Not yet supported on `generateStream` (throws 501). + * `toolMode: 'auto'` only. Cumulative USD cap across all iterations, each priced + * from the `pricing` of the entry that served it (`GenerateResult.costUsd`). Throws + * `BudgetExceededError({kind: 'cost'})` once the total reaches it, so `0` allows no + * call. An unpriced round, or one without both token counts, makes the cap + * unmeasurable and throws the same error. Not yet supported on `generateStream` + * (throws 501). */ maxCostUsd?: number; /** @@ -448,6 +445,12 @@ export type ModelCallResult = | { status: 'completed'; output: T; usage?: TokenUsage } | { status: 'pending'; operationId: string; resumeAfter?: number }; +/** What an operator pays for a generative model, in USD per million tokens (`models.generative..pricing`). */ +export interface ModelPricing { + inputPerMtok: number; + outputPerMtok: number; +} + export interface TokenUsage { promptTokens?: number; completionTokens?: number; @@ -485,16 +488,22 @@ export interface GenerateResult { /** Why generation stopped. Backend-agnostic; backends map their native reasons. */ finishReason: 'stop' | 'length' | 'tool_calls' | 'content_filter'; /** - * Token / cost usage for this call. Pass-through from the backend's + * Token usage for this call. Pass-through from the backend's * `ModelCallResult.usage` so callers (and the `toolMode: 'auto'` loop) can * cap cumulative budgets without re-querying analytics. * - * Backends MUST NOT set this field on the `output` they return — the framework - * (`Models.generate`) injects it from `ModelCallResult.usage`, which is the - * authoritative source. A backend that sets both gets the framework's value; - * easy to miss but unambiguous: usage lives at the call-result layer. + * Set only by the framework (`Models.generate`) from `ModelCallResult.usage`, + * the authoritative source. A value a backend puts on its `output` is dropped, + * as is `costUsd`. */ usage?: TokenUsage; + /** + * USD cost of this call, from the `pricing` declared on the model entry that + * served it (a fallback is priced at its own rate). Absent when that entry + * declares no pricing or the backend did not report both token counts. With + * `toolMode: 'auto'` it covers the final round only, like `usage`. + */ + costUsd?: number; /** * `toolMode: 'auto'` only, and only when `includeToolTrace: true` (or when an * error path attaches `partialTrace` via `BudgetExceededError`). One entry per diff --git a/unitTests/agent/agent.test.js b/unitTests/agent/agent.test.js index 1c54fb3eec..a34474ebde 100644 --- a/unitTests/agent/agent.test.js +++ b/unitTests/agent/agent.test.js @@ -229,3 +229,25 @@ describe('agent/agent mergeConfig maxTokens', () => { } }); }); + +describe('agent/agent startOnMainThread budget caps (#3040)', () => { + const { startOnMainThread } = require('#src/agent/agent'); + + it('refuses to start an enabled agent with a malformed cap, before registering anything', async () => { + const registered = []; + const server = { registerOperation: (op) => registered.push(op.name) }; + await assert.rejects(startOnMainThread({ server, enabled: true, maxSessionTokens: 'lots' }), (err) => { + assert.match(err.message, /agent\.maxSessionTokens must be a positive integer/); + return true; + }); + await assert.rejects(startOnMainThread({ server, enabled: true, maxCostUsd: -1 }), /agent\.maxCostUsd must be/); + await assert.rejects(startOnMainThread({ server, enabled: true, maxTurns: 'lots' }), /agent\.maxTurns must be/); + // A bare `maxTurns:` key in the config file arrives as null. + await assert.rejects(startOnMainThread({ server, enabled: true, maxTurns: null }), /agent\.maxTurns must be/); + assert.deepStrictEqual(registered, []); + }); + + it('ignores a malformed cap while the agent is disabled', async () => { + await startOnMainThread({ server: { registerOperation: () => {} }, enabled: false, maxSessionTokens: 'lots' }); + }); +}); diff --git a/unitTests/agent/budget.test.js b/unitTests/agent/budget.test.js new file mode 100644 index 0000000000..7d49251b2e --- /dev/null +++ b/unitTests/agent/budget.test.js @@ -0,0 +1,160 @@ +'use strict'; + +const assert = require('node:assert'); +const { + summarizeSessionUsage, + budgetViolation, + parseBudgetCap, + parseMaxTurns, + callUsage, +} = require('#src/agent/budget'); + +const assistant = (usage) => ({ role: 'assistant', content: '', createdAt: 0, ...(usage && { usage }) }); +const user = { role: 'user', content: 'hi', createdAt: 0 }; +const tool = { role: 'tool', content: '{}', toolCallId: 'c1', createdAt: 0 }; + +describe('agent/budget', () => { + describe('summarizeSessionUsage', () => { + it('counts only assistant messages as calls and sums their figures', () => { + const usage = summarizeSessionUsage([ + user, + assistant({ promptTokens: 10, completionTokens: 2, costUsd: 0.5 }), + tool, + assistant({ promptTokens: 5, completionTokens: 1, costUsd: 0.25 }), + ]); + assert.deepStrictEqual(usage, { calls: 2, promptTokens: 15, completionTokens: 3, costUsd: 0.75 }); + }); + + it('reports zero totals for a session with no calls yet', () => { + assert.deepStrictEqual(summarizeSessionUsage([user]), { + calls: 0, + promptTokens: 0, + completionTokens: 0, + costUsd: 0, + }); + }); + + it('drops a total for good once any call did not report it', () => { + const usage = summarizeSessionUsage([ + assistant({ promptTokens: 10, completionTokens: 2 }), + assistant({ promptTokens: 5, completionTokens: 1, costUsd: 0.25 }), + assistant(undefined), + ]); + assert.deepStrictEqual(usage, { calls: 3 }); + }); + + it('counts rejected calls kept beside the transcript as calls', () => { + assert.deepStrictEqual( + summarizeSessionUsage( + [assistant({ promptTokens: 10, completionTokens: 2, costUsd: 0.5 })], + [{ promptTokens: 5, completionTokens: 1, costUsd: 0.25 }] + ), + { calls: 2, promptTokens: 15, completionTokens: 3, costUsd: 0.75 } + ); + assert.deepStrictEqual(summarizeSessionUsage([], [{}]), { calls: 1 }); + }); + + it('treats a non-finite or negative figure as unreported', () => { + assert.deepStrictEqual( + summarizeSessionUsage([assistant({ promptTokens: Number.NaN, completionTokens: -1, costUsd: Infinity })]), + { calls: 1 } + ); + }); + }); + + describe('callUsage', () => { + it('keeps only supplied fields and returns nothing for an empty report', () => { + assert.deepStrictEqual(callUsage({ promptTokens: 3 }, undefined), { promptTokens: 3 }); + assert.deepStrictEqual(callUsage({ promptTokens: 3, completionTokens: 0 }, 0), { + promptTokens: 3, + completionTokens: 0, + costUsd: 0, + }); + assert.strictEqual(callUsage({}, undefined), undefined); + assert.strictEqual(callUsage(undefined, undefined), undefined); + }); + }); + + describe('budgetViolation', () => { + const measured = { calls: 1, promptTokens: 60, completionTokens: 40, costUsd: 0.5 }; + + it('allows a session under every cap, and any session with no caps', () => { + assert.strictEqual(budgetViolation(measured, { maxSessionTokens: 101, maxCostUsd: 0.51 }), undefined); + assert.strictEqual(budgetViolation({ calls: 9 }, {}), undefined); + }); + + it('trips at the cap, not only past it', () => { + assert.match( + budgetViolation(measured, { maxSessionTokens: 100 }), + /maxSessionTokens=100 reached: this session used 100 tokens/ + ); + assert.match( + budgetViolation(measured, { maxCostUsd: 0.5 }), + /maxCostUsd=0.5 reached: this session cost \$0\.5000/ + ); + }); + + it('maxCostUsd 0 stops before the first call', () => { + assert.match(budgetViolation(summarizeSessionUsage([]), { maxCostUsd: 0 }), /maxCostUsd=0 reached/); + }); + + it('stops a session whose capped figure is unmeasured', () => { + assert.match( + budgetViolation({ calls: 1, costUsd: 0.1 }, { maxSessionTokens: 1000 }), + /maxSessionTokens=1000 cannot be enforced/ + ); + assert.match( + budgetViolation({ calls: 1, promptTokens: 1, completionTokens: 1 }, { maxCostUsd: 5 }), + /maxCostUsd=5 cannot be enforced: .*declare pricing/ + ); + }); + }); + + describe('parseBudgetCap', () => { + it('accepts valid caps, numeric strings from env-style config, and clears on null or absence', () => { + assert.strictEqual(parseBudgetCap('maxSessionTokens', 500000), 500000); + assert.strictEqual(parseBudgetCap('maxSessionTokens', '500000'), 500000); + assert.strictEqual(parseBudgetCap('maxCostUsd', 0), 0); + assert.strictEqual(parseBudgetCap('maxCostUsd', 2.5), 2.5); + assert.strictEqual(parseBudgetCap('maxCostUsd', null), undefined); + assert.strictEqual(parseBudgetCap('maxSessionTokens', undefined), undefined); + }); + + it('refuses a value that would silently uncap or mis-trip', () => { + for (const [name, value] of [ + ['maxSessionTokens', 0], + ['maxSessionTokens', 1.5], + ['maxSessionTokens', 'many'], + ['maxSessionTokens', ''], + ['maxCostUsd', -1], + ['maxCostUsd', Number.NaN], + ['maxCostUsd', Infinity], + ['maxCostUsd', true], + ['maxCostUsd', {}], + ]) { + assert.throws( + () => parseBudgetCap(name, value), + (err) => err.statusCode === 400 && err.message.startsWith(`agent.${name} must be`), + `${name}=${String(value)}` + ); + } + }); + }); + + describe('parseMaxTurns', () => { + it('accepts a positive integer, also as a config string', () => { + assert.strictEqual(parseMaxTurns(50), 50); + assert.strictEqual(parseMaxTurns('7'), 7); + }); + + it('refuses null and anything that would make every run turnless', () => { + for (const value of [null, 0, -3, 1.5, Number.NaN, 'many', '', true]) { + assert.throws( + () => parseMaxTurns(value), + (err) => err.statusCode === 400 && /agent\.maxTurns must be a positive integer/.test(err.message), + String(value) + ); + } + }); + }); +}); diff --git a/unitTests/agent/loop.test.js b/unitTests/agent/loop.test.js index f870705568..a55405fc83 100644 --- a/unitTests/agent/loop.test.js +++ b/unitTests/agent/loop.test.js @@ -1,7 +1,7 @@ 'use strict'; const assert = require('node:assert'); -const { runAgent, _resetInFlightForTests } = require('#src/agent/loop'); +const { runAgent, admitRun, fenceRuns, _resetInFlightForTests } = require('#src/agent/loop'); const session = require('#src/agent/session'); function deepFreeze(o) { @@ -1123,3 +1123,1036 @@ describe('agent/loop observation size and context window', () => { assert.equal(reloaded.messages[2].content, 'B'.repeat(10_000)); }); }); + +describe('agent/loop per-session budgets (#3040)', () => { + const { summarizeSessionUsage } = require('#src/agent/budget'); + + beforeEach(() => { + session._setTableForTests(makeMockTable()); + _resetInFlightForTests(); + }); + afterEach(() => { + session._setTableForTests(undefined); + }); + + const usage = (promptTokens, completionTokens, costUsd) => ({ + usage: { promptTokens, completionTokens }, + ...(costUsd !== undefined && { costUsd }), + }); + const toolTurn = (id, name, metered) => ({ + content: '', + finishReason: 'tool_calls', + toolCalls: [{ id, name, arguments: {} }], + ...metered, + }); + const finalTurn = (content, metered) => ({ content, finishReason: 'stop', ...metered }); + + function countingTool(name, { destructive = false } = {}) { + const tool = { + runs: 0, + def: { name, description: name, parameters: { type: 'object' } }, + destructive, + handler: async () => { + tool.runs++; + return { ok: true }; + }, + }; + return tool; + } + + async function newSession() { + const created = await session.createSession({ user: 'admin' }); + await session.appendMessage(created.session_id, { role: 'user', content: 'go', createdAt: Date.now() }); + return created.session_id; + } + + it('records each call on its assistant message, and the folded total matches the sum', async () => { + const sessionId = await newSession(); + const read = countingTool('read'); + const models = stubModels([toolTurn('c1', 'read', usage(100, 20, 0.01)), finalTurn('done', usage(150, 30, 0.02))]); + await runAgent({ sessionId, models, tools: [read], scopes, maxTurns: 5 }); + + const reloaded = await session.getSession(sessionId); + assert.equal(reloaded.status, 'completed'); + const assistants = reloaded.messages.filter((m) => m.role === 'assistant'); + assert.deepStrictEqual(assistants[0].usage, { promptTokens: 100, completionTokens: 20, costUsd: 0.01 }); + assert.deepStrictEqual(summarizeSessionUsage(reloaded.messages), { + calls: 2, + promptTokens: 250, + completionTokens: 50, + costUsd: 0.03, + }); + }); + + it('maxSessionTokens reached on a tool-call turn answers the calls without running them and ends the session error', async () => { + const sessionId = await newSession(); + const read = countingTool('read'); + const models = stubModels([toolTurn('c1', 'read', usage(80, 30))]); + await runAgent({ sessionId, models, tools: [read], scopes, maxTurns: 5, maxSessionTokens: 100 }); + + const reloaded = await session.getSession(sessionId); + assert.equal(read.runs, 0); + assert.equal(reloaded.status, 'error'); + assert.match(reloaded.lastError, /agent\.maxSessionTokens=100 reached: this session used 110 tokens/); + const toolMessages = reloaded.messages.filter((m) => m.role === 'tool'); + assert.equal(toolMessages.length, 1, 'every tool call keeps its response, so the transcript stays sendable'); + assert.equal(toolMessages[0].toolCallId, 'c1'); + assert.equal(JSON.parse(toolMessages[0].content).error, 'budget_exceeded'); + }); + + it('a final answer that crosses the cap is delivered, and the completed run names the cap', async () => { + const sessionId = await newSession(); + const models = stubModels([finalTurn('the answer', usage(80, 30))]); + await runAgent({ sessionId, models, tools: noTools, scopes, maxTurns: 5, maxSessionTokens: 100 }); + + const reloaded = await session.getSession(sessionId); + assert.equal(reloaded.status, 'completed'); + assert.match(reloaded.lastError, /agent\.maxSessionTokens=100 reached/); + assert.equal(reloaded.messages.at(-1).content, 'the answer'); + }); + + it('a session already at its cap makes no model call on the next prompt', async () => { + const sessionId = await newSession(); + await runAgent({ + sessionId, + models: stubModels([finalTurn('first', usage(80, 30))]), + tools: noTools, + scopes, + maxTurns: 5, + maxSessionTokens: 100, + }); + await session.appendMessage(sessionId, { role: 'user', content: 'more', createdAt: Date.now() }); + let generateCalls = 0; + const models = { + async generate() { + generateCalls++; + return finalTurn('unreached', usage(1, 1)); + }, + }; + await runAgent({ sessionId, models, tools: noTools, scopes, maxTurns: 5, maxSessionTokens: 100 }); + + const reloaded = await session.getSession(sessionId); + assert.equal(generateCalls, 0); + assert.equal(reloaded.status, 'error'); + assert.match(reloaded.lastError, /maxSessionTokens=100 reached/); + }); + + it('maxCostUsd stops after a call it cannot price, naming the missing pricing', async () => { + const sessionId = await newSession(); + const read = countingTool('read'); + const models = stubModels([toolTurn('c1', 'read', usage(10, 10)), finalTurn('unreached', usage(1, 1))]); + await runAgent({ sessionId, models, tools: [read], scopes, maxTurns: 5, maxCostUsd: 1 }); + + const reloaded = await session.getSession(sessionId); + assert.equal(read.runs, 0); + assert.equal(reloaded.status, 'error'); + assert.match(reloaded.lastError, /maxCostUsd=1 cannot be enforced: .*declare pricing/); + }); + + it('maxSessionTokens stops after a call that reported no token usage', async () => { + const sessionId = await newSession(); + const models = stubModels([toolTurn('c1', 'read', {}), finalTurn('unreached', usage(1, 1))]); + await runAgent({ sessionId, models, tools: [countingTool('read')], scopes, maxTurns: 5, maxSessionTokens: 1000 }); + + const reloaded = await session.getSession(sessionId); + assert.equal(reloaded.status, 'error'); + assert.match(reloaded.lastError, /maxSessionTokens=1000 cannot be enforced/); + }); + + it('priced calls under maxCostUsd keep running', async () => { + const sessionId = await newSession(); + const read = countingTool('read'); + const models = stubModels([toolTurn('c1', 'read', usage(10, 10, 0.2)), finalTurn('done', usage(10, 10, 0.2))]); + await runAgent({ sessionId, models, tools: [read], scopes, maxTurns: 5, maxCostUsd: 1 }); + + const reloaded = await session.getSession(sessionId); + assert.equal(read.runs, 1); + assert.equal(reloaded.status, 'completed'); + }); + + it('a cap lowered while an action awaits approval stops the approved action from running', async () => { + const sessionId = await newSession(); + const drop = countingTool('drop', { destructive: true }); + await runAgent({ + sessionId, + models: stubModels([toolTurn('c1', 'drop', usage(60, 10))]), + tools: [drop], + scopes, + maxTurns: 5, + autoApprove: false, + }); + const paused = await session.getSession(sessionId); + assert.equal(paused.status, 'awaiting_approval'); + await session.resolveApproval(sessionId, paused.pendingApprovals[0].id, true); + + await runAgent({ + sessionId, + models: stubModels([]), + tools: [drop], + scopes, + maxTurns: 5, + autoApprove: false, + maxSessionTokens: 50, + }); + + const reloaded = await session.getSession(sessionId); + assert.equal(drop.runs, 0, 'the approved destructive action must not run over the cap'); + assert.equal(reloaded.status, 'error'); + assert.equal(reloaded.pendingApprovals[0].consumed, true); + const toolMessages = reloaded.messages.filter((m) => m.role === 'tool'); + assert.equal(toolMessages.length, 1); + assert.equal(toolMessages[0].toolCallId, 'c1'); + assert.equal(JSON.parse(toolMessages[0].content).error, 'budget_exceeded'); + }); + + it('a stop on resume answers every open approval, resolved or not, so none is left pending', async () => { + const sessionId = await newSession(); + const drop = countingTool('drop', { destructive: true }); + await runAgent({ + sessionId, + models: stubModels([ + { + content: '', + finishReason: 'tool_calls', + toolCalls: [ + { id: 'c1', name: 'drop', arguments: {} }, + { id: 'c2', name: 'drop', arguments: {} }, + ], + ...usage(60, 10), + }, + ]), + tools: [drop], + scopes, + maxTurns: 5, + autoApprove: false, + }); + const paused = await session.getSession(sessionId); + assert.equal(paused.pendingApprovals.length, 2); + await session.resolveApproval(sessionId, paused.pendingApprovals[0].id, true); + + await runAgent({ + sessionId, + models: stubModels([]), + tools: [drop], + scopes, + maxTurns: 5, + autoApprove: false, + maxSessionTokens: 50, + }); + + const reloaded = await session.getSession(sessionId); + assert.equal(drop.runs, 0); + assert.equal(reloaded.status, 'error'); + assert.ok(reloaded.pendingApprovals.every((a) => a.resolved && a.consumed)); + assert.equal(reloaded.pendingApprovals[1].approved, false); + const answered = reloaded.messages.filter((m) => m.role === 'tool').map((m) => m.toolCallId); + assert.deepStrictEqual(answered.sort(), ['c1', 'c2']); + }); + + it('a run that sees a cancel runs no approved action, and the cancel closes it for later prompts', async () => { + for (const caps of [{}, { maxSessionTokens: 1_000_000 }]) { + _resetInFlightForTests(); + const sessionId = await newSession(); + const drop = countingTool('drop', { destructive: true }); + await runAgent({ + sessionId, + models: stubModels([toolTurn('c1', 'drop', usage(60, 10))]), + tools: [drop], + scopes, + maxTurns: 5, + autoApprove: false, + }); + const paused = await session.getSession(sessionId); + await session.resolveApproval(sessionId, paused.pendingApprovals[0].id, true); + const controller = new AbortController(); + controller.abort(); + await runAgent({ + sessionId, + models: stubModels([]), + tools: [drop], + scopes, + maxTurns: 5, + autoApprove: false, + signal: controller.signal, + ...caps, + }); + assert.equal(drop.runs, 0, JSON.stringify(caps)); + + await session.cancelSession(sessionId); + const cancelled = await session.getSession(sessionId); + assert.ok(cancelled.pendingApprovals.every((a) => a.consumed)); + assert.equal(JSON.parse(cancelled.messages.at(-1).content).error, 'cancelled_by_operator'); + + await session.appendMessage(sessionId, { role: 'user', content: 'summarize', createdAt: Date.now() }); + await runAgent({ + sessionId, + models: stubModels([finalTurn('summary', usage(1, 1))]), + tools: [drop], + scopes, + maxTurns: 5, + autoApprove: false, + ...caps, + }); + assert.equal(drop.runs, 0, 'the next prompt must not run the cancelled approval'); + } + }); + + it('a cancel that lands while an approved action runs leaves that call one response, its real result', async () => { + const sessionId = await newSession(); + let release; + let markStarted; + const started = new Promise((resolve) => (markStarted = resolve)); + const drop = { + runs: 0, + def: { name: 'drop', description: 'drop', parameters: { type: 'object' } }, + destructive: true, + handler: async () => { + drop.runs++; + markStarted(); + await new Promise((resolve) => (release = resolve)); + return { dropped: true }; + }, + }; + await runAgent({ + sessionId, + models: stubModels([toolTurn('c1', 'drop', usage(60, 10))]), + tools: [drop], + scopes, + maxTurns: 5, + autoApprove: false, + }); + const paused = await session.getSession(sessionId); + await session.resolveApproval(sessionId, paused.pendingApprovals[0].id, true); + + const run = runAgent({ + sessionId, + models: stubModels([finalTurn('done', usage(1, 1))]), + tools: [drop], + scopes, + maxTurns: 5, + autoApprove: false, + }); + await started; + await session.cancelSession(sessionId); + release(); + await run; + + const reloaded = await session.getSession(sessionId); + const responses = reloaded.messages.filter((m) => m.role === 'tool' && m.toolCallId === 'c1'); + assert.equal(responses.length, 1, 'one response per tool call'); + assert.deepStrictEqual(JSON.parse(responses[0].content), { ok: true, result: { dropped: true } }); + assert.equal(drop.runs, 1); + }); + + it('a cancel that lands during the claim answers the call instead of running the action', async () => { + const sessionId = await newSession(); + const drop = countingTool('drop', { destructive: true }); + await runAgent({ + sessionId, + models: stubModels([toolTurn('c1', 'drop', usage(60, 10))]), + tools: [drop], + scopes, + maxTurns: 5, + autoApprove: false, + }); + const paused = await session.getSession(sessionId); + const approvalId = paused.pendingApprovals[0].id; + await session.resolveApproval(sessionId, approvalId, true); + + const controller = new AbortController(); + const table = session.getAgentSessionTable(); + const put = table.put.bind(table); + table.put = async (row) => { + await put(row); + if (row.pendingApprovals.some((a) => a.id === approvalId && a.consumed)) controller.abort(); + }; + await runAgent({ + sessionId, + models: stubModels([]), + tools: [drop], + scopes, + maxTurns: 5, + autoApprove: false, + signal: controller.signal, + }); + + const reloaded = await session.getSession(sessionId); + assert.equal(drop.runs, 0); + const responses = reloaded.messages.filter((m) => m.role === 'tool' && m.toolCallId === 'c1'); + assert.equal(responses.length, 1); + assert.equal(JSON.parse(responses[0].content).error, 'cancelled_by_operator'); + }); + + it('endRun completes the transcript but keeps an aborted status a cancel already wrote', async () => { + const sessionId = await newSession(); + await session.setStatus(sessionId, 'aborted', 'Cancelled by operator'); + await session.endRun(sessionId, { + append: [{ role: 'tool', content: '{}', toolCallId: 'c1', createdAt: Date.now() }], + status: 'error', + lastError: 'agent.maxSessionTokens=1 reached', + }); + const reloaded = await session.getSession(sessionId); + assert.equal(reloaded.status, 'aborted'); + assert.equal(reloaded.lastError, 'Cancelled by operator'); + assert.equal(reloaded.messages.at(-1).toolCallId, 'c1'); + }); + + it("records a rejected reply's usage beside the transcript and counts it toward the caps", async () => { + const sessionId = await newSession(); + const cutOff = { content: 'partial', finishReason: 'length', ...usage(80, 30, 0.5) }; + await runAgent({ sessionId, models: stubModels([cutOff]), tools: noTools, scopes, maxTurns: 5 }); + + const rejected = await session.getSession(sessionId); + assert.equal(rejected.status, 'error'); + assert.equal( + rejected.messages.filter((m) => m.role === 'assistant').length, + 0, + 'the reply stays out of the transcript' + ); + assert.deepStrictEqual(rejected.rejectedCalls, [{ promptTokens: 80, completionTokens: 30, costUsd: 0.5 }]); + assert.deepStrictEqual(summarizeSessionUsage(rejected.messages, rejected.rejectedCalls), { + calls: 1, + promptTokens: 80, + completionTokens: 30, + costUsd: 0.5, + }); + + await session.appendMessage(sessionId, { role: 'user', content: 'again', createdAt: Date.now() }); + let generateCalls = 0; + const models = { + async generate() { + generateCalls++; + return finalTurn('unreached', usage(1, 1)); + }, + }; + await runAgent({ sessionId, models, tools: noTools, scopes, maxTurns: 5, maxSessionTokens: 100 }); + assert.equal(generateCalls, 0, 'the rejected call alone already reached the cap'); + assert.match( + (await session.getSession(sessionId)).lastError, + /maxSessionTokens=100 reached: this session used 110 tokens/ + ); + }); + + it('a rejected reply that reported no usage makes a capped session unmeasured', async () => { + const sessionId = await newSession(); + await runAgent({ + sessionId, + models: stubModels([{ content: '', finishReason: 'stop' }]), + tools: noTools, + scopes, + maxTurns: 5, + }); + const rejected = await session.getSession(sessionId); + assert.deepStrictEqual(rejected.rejectedCalls, [{}]); + assert.deepStrictEqual(summarizeSessionUsage(rejected.messages, rejected.rejectedCalls), { calls: 1 }); + }); + + it('an assistant turn written before usage was recorded makes a capped session unmeasured, not free', async () => { + const sessionId = await newSession(); + await session.appendMessage(sessionId, { role: 'assistant', content: 'legacy', createdAt: Date.now() }); + await session.appendMessage(sessionId, { role: 'user', content: 'again', createdAt: Date.now() }); + let generateCalls = 0; + const models = { + async generate() { + generateCalls++; + return finalTurn('unreached', usage(1, 1)); + }, + }; + await runAgent({ sessionId, models, tools: noTools, scopes, maxTurns: 5, maxSessionTokens: 1_000_000 }); + + const reloaded = await session.getSession(sessionId); + assert.equal(generateCalls, 0); + assert.equal(reloaded.status, 'error'); + assert.match(reloaded.lastError, /cannot be enforced/); + }); +}); + +describe('agent/loop never overwrites a cancel', () => { + beforeEach(() => { + session._setTableForTests(makeMockTable()); + _resetInFlightForTests(); + }); + afterEach(() => { + session._setTableForTests(undefined); + }); + + const cancel = (sessionId) => session.setStatus(sessionId, 'aborted', 'Cancelled by operator'); + const toolCall = (id, name) => ({ + content: '', + finishReason: 'tool_calls', + toolCalls: [{ id, name, arguments: {} }], + }); + const tool = (name, { destructive = false, handler } = {}) => ({ + def: { name, description: name, parameters: { type: 'object' } }, + destructive, + handler: handler ?? (async () => ({ ok: true })), + }); + + async function newSession() { + const created = await session.createSession({ user: 'admin' }); + await session.appendMessage(created.session_id, { role: 'user', content: 'go', createdAt: Date.now() }); + return created.session_id; + } + + async function expectStillAborted(sessionId) { + const reloaded = await session.getSession(sessionId); + assert.equal(reloaded.status, 'aborted'); + assert.equal(reloaded.lastError, 'Cancelled by operator'); + return reloaded; + } + + it('keeps a cancel that lands while the final answer is generated', async () => { + const sessionId = await newSession(); + const models = { + async generate() { + await cancel(sessionId); + return { content: 'done', finishReason: 'stop' }; + }, + }; + await runAgent({ sessionId, models, tools: noTools, scopes, maxTurns: 5 }); + await expectStillAborted(sessionId); + }); + + it('keeps a cancel that lands before the maxTurns ceiling is written', async () => { + const sessionId = await newSession(); + const models = { + async generate() { + await cancel(sessionId); + return toolCall('c1', 'read'); + }, + }; + await runAgent({ sessionId, models, tools: [tool('read')], scopes, maxTurns: 1 }); + await expectStillAborted(sessionId); + }); + + it('keeps a cancel that lands while a run pauses on another pending approval', async () => { + const sessionId = await newSession(); + const drop = tool('drop', { destructive: true, handler: async () => cancel(sessionId) }); + await runAgent({ + sessionId, + models: stubModels([ + { + content: '', + finishReason: 'tool_calls', + toolCalls: [ + { id: 'c1', name: 'drop', arguments: {} }, + { id: 'c2', name: 'drop', arguments: {} }, + ], + }, + ]), + tools: [drop], + scopes, + maxTurns: 5, + autoApprove: false, + }); + const paused = await session.getSession(sessionId); + await session.resolveApproval(sessionId, paused.pendingApprovals[0].id, true); + await runAgent({ sessionId, models: stubModels([]), tools: [drop], scopes, maxTurns: 5, autoApprove: false }); + await expectStillAborted(sessionId); + }); + + it('keeps a cancel that lands before a failed run writes its error', async () => { + const sessionId = await newSession(); + const models = { + async generate() { + await cancel(sessionId); + throw new Error('provider down'); + }, + }; + await assert.rejects(runAgent({ sessionId, models, tools: noTools, scopes, maxTurns: 5 }), /provider down/); + await expectStillAborted(sessionId); + }); + + it('answers a gated call as cancelled instead of reopening a cancelled session for approval', async () => { + const sessionId = await newSession(); + const models = { + async generate() { + await cancel(sessionId); + return toolCall('c1', 'drop'); + }, + }; + await runAgent({ + sessionId, + models, + tools: [tool('drop', { destructive: true })], + scopes, + maxTurns: 5, + autoApprove: false, + }); + const reloaded = await expectStillAborted(sessionId); + assert.equal(reloaded.pendingApprovals.length, 0); + const response = reloaded.messages.find((m) => m.role === 'tool' && m.toolCallId === 'c1'); + assert.equal(JSON.parse(response.content).error, 'cancelled_by_operator'); + }); + + it('leaves a session a run already finished as it is', async () => { + const sessionId = await newSession(); + await session.setStatus(sessionId, 'completed'); + assert.equal(await session.cancelSession(sessionId), false); + assert.equal((await session.getSession(sessionId)).status, 'completed'); + }); + + it('answers the rest of the turn when a cancel lands between its tool calls', async () => { + const sessionId = await newSession(); + const controller = new AbortController(); + const ran = []; + const first = tool('first', { + handler: async () => { + ran.push('first'); + controller.abort(); + await cancel(sessionId); + return { ok: true }; + }, + }); + const second = tool('second', { handler: async () => ran.push('second') }); + const models = stubModels([ + { + content: '', + finishReason: 'tool_calls', + toolCalls: [ + { id: 'c1', name: 'first', arguments: {} }, + { id: 'c2', name: 'second', arguments: {} }, + ], + }, + ]); + await runAgent({ sessionId, models, tools: [first, second], scopes, maxTurns: 5, signal: controller.signal }); + + assert.deepStrictEqual(ran, ['first']); + const reloaded = await expectStillAborted(sessionId); + const responses = reloaded.messages.filter((m) => m.role === 'tool'); + assert.deepStrictEqual( + responses.map((m) => m.toolCallId), + ['c1', 'c2'] + ); + assert.equal(JSON.parse(responses[1].content).error, 'cancelled_by_operator'); + }); + + it('starts no run for a request admitted before a cancel that arrived while it was writing', async () => { + const sessionId = await newSession(); + const admitted = admitRun(sessionId); + fenceRuns(sessionId); + await cancel(sessionId); + let generateCalls = 0; + const models = { + async generate() { + generateCalls++; + return { content: 'unreached', finishReason: 'stop' }; + }, + }; + await runAgent({ sessionId, models, tools: noTools, scopes, maxTurns: 5, admitted }); + assert.equal(generateCalls, 0); + await expectStillAborted(sessionId); + + await runAgent({ sessionId, models, tools: noTools, scopes, maxTurns: 5, admitted: admitRun(sessionId) }); + assert.equal(generateCalls, 1, 'a request admitted after the cancel runs'); + }); + + it('keeps a cancel that lands while a rejected reply is generated', async () => { + const sessionId = await newSession(); + const models = { + async generate() { + await cancel(sessionId); + return { content: 'partial', finishReason: 'length' }; + }, + }; + await runAgent({ sessionId, models, tools: noTools, scopes, maxTurns: 5 }); + await expectStillAborted(sessionId); + }); + + it('drains an approval resolved while the run is executing an earlier one', async () => { + const sessionId = await newSession(); + const runs = []; + let secondApprovalId; + const drop = tool('drop', { + destructive: true, + handler: async (args) => { + runs.push(args.n); + if (args.n === 1) await session.resolveApproval(sessionId, secondApprovalId, true); + return { ok: true }; + }, + }); + await runAgent({ + sessionId, + models: stubModels([ + { + content: '', + finishReason: 'tool_calls', + toolCalls: [ + { id: 'c1', name: 'drop', arguments: { n: 1 } }, + { id: 'c2', name: 'drop', arguments: { n: 2 } }, + ], + }, + ]), + tools: [drop], + scopes, + maxTurns: 5, + autoApprove: false, + }); + const paused = await session.getSession(sessionId); + secondApprovalId = paused.pendingApprovals[1].id; + await session.resolveApproval(sessionId, paused.pendingApprovals[0].id, true); + let sentToModel; + const models = { + async generate(input) { + sentToModel = input.messages; + return { content: 'done', finishReason: 'stop' }; + }, + }; + await runAgent({ sessionId, models, tools: [drop], scopes, maxTurns: 5, autoApprove: false }); + + assert.deepStrictEqual(runs, [1, 2]); + const answered = sentToModel.filter((m) => m.role === 'tool').map((m) => m.toolCallId); + assert.deepStrictEqual( + answered.sort(), + ['c1', 'c2'], + 'every tool call has its response before the next model call' + ); + assert.equal((await session.getSession(sessionId)).status, 'completed'); + }); + + it('does not pause a session whose approvals were all resolved before the pause was written', async () => { + const sessionId = await newSession(); + await session.addPendingApproval(sessionId, { + toolName: 'drop', + arguments: {}, + toolCallId: 'c1', + reason: 'destructive', + }); + const paused = await session.getSession(sessionId); + await session.resolveApproval(sessionId, paused.pendingApprovals[0].id, true); + const after = await session.endRun(sessionId, { status: 'awaiting_approval', pauseWhileApprovalsOpen: true }); + assert.equal(after.status, 'idle'); + }); + + it("drains an approval resolved while the same turn's later calls are still running", async () => { + const sessionId = await newSession(); + const runs = []; + const drop = tool('drop', { + destructive: true, + handler: async () => { + runs.push('drop'); + return { ok: true }; + }, + }); + // The operator approves while this inline call runs; `startRun` folds that request into the live run. + const slow = tool('slow', { + handler: async () => { + const parked = await session.getSession(sessionId); + await session.resolveApproval(sessionId, parked.pendingApprovals[0].id, true); + return { ok: true }; + }, + }); + const replies = [ + { + content: '', + finishReason: 'tool_calls', + toolCalls: [ + { id: 'c1', name: 'drop', arguments: {} }, + { id: 'c2', name: 'slow', arguments: {} }, + ], + }, + ]; + let sentToModel; + let statusDuringNextCall; + const models = { + async generate(input) { + if (replies.length > 0) return replies.shift(); + sentToModel = input.messages; + statusDuringNextCall = (await session.getSession(sessionId)).status; + return { content: 'done', finishReason: 'stop' }; + }, + }; + await runAgent({ sessionId, models, tools: [drop, slow], scopes, maxTurns: 5, autoApprove: false }); + + assert.deepStrictEqual(runs, ['drop']); + const answered = sentToModel.filter((m) => m.role === 'tool').map((m) => m.toolCallId); + assert.deepStrictEqual( + answered.sort(), + ['c1', 'c2'], + 'every tool call has its response before the next model call' + ); + assert.equal(statusDuringNextCall, 'running'); + assert.equal((await session.getSession(sessionId)).status, 'completed'); + }); + + it('answers inline calls a dead run left without results, in place and without running them again', async () => { + const sessionId = await newSession(); + await session.appendMessage(sessionId, { + role: 'assistant', + content: '', + toolCalls: [ + { id: 'c1', name: 'read', arguments: {} }, + { id: 'c2', name: 'read', arguments: {} }, + ], + createdAt: 1, + }); + await session.appendMessage(sessionId, { role: 'tool', content: '{}', toolCallId: 'c1', createdAt: 2 }); + // What an operator does after the crash: cancel the stuck run, then prompt again. + await session.cancelSession(sessionId); + await session.appendMessage(sessionId, { role: 'user', content: 'what happened?', createdAt: 3 }); + + const read = tool('read', { handler: async () => assert.fail('must not run again') }); + let sentToModel; + const models = { + async generate(input) { + sentToModel = input.messages; + return { content: 'done', finishReason: 'stop' }; + }, + }; + await runAgent({ sessionId, models, tools: [read], scopes, maxTurns: 5, autoApprove: false }); + + assert.deepStrictEqual( + sentToModel.map((m) => m.toolCallId ?? m.role), + ['user', 'assistant', 'c1', 'c2', 'user'] + ); + assert.deepStrictEqual(JSON.parse(sentToModel[3].content), { + ok: false, + error: 'interrupted', + outcome: 'unknown', + tool: 'read', + }); + }); + + // Harper stopped after the approval was saved and before its run started, then the operator prompted. + async function approvedThenPrompted() { + const sessionId = await newSession(); + await session.appendMessage(sessionId, { + role: 'assistant', + content: '', + toolCalls: [{ id: 'c1', name: 'drop', arguments: {} }], + createdAt: 1, + }); + const approval = await session.addPendingApproval(sessionId, { + toolName: 'drop', + arguments: {}, + toolCallId: 'c1', + reason: 'destructive', + }); + await session.resolveApproval(sessionId, approval.id, true); + await session.appendMessage(sessionId, { role: 'user', content: 'and then?', createdAt: 2 }); + return sessionId; + } + + it('places an approved result ahead of a prompt appended while the approval waited for its run', async () => { + const sessionId = await approvedThenPrompted(); + let sentToModel; + const models = { + async generate(input) { + sentToModel = input.messages; + return { content: 'done', finishReason: 'stop' }; + }, + }; + await runAgent({ sessionId, models, tools: [tool('drop', { destructive: true })], scopes, maxTurns: 5 }); + assert.deepStrictEqual( + sentToModel.map((m) => m.toolCallId ?? m.role), + ['user', 'assistant', 'c1', 'user'] + ); + }); + + it('places a cancelled approval answer ahead of a prompt appended while the approval waited', async () => { + const sessionId = await approvedThenPrompted(); + await session.cancelSession(sessionId); + const order = (await session.getSession(sessionId)).messages.map((m) => m.toolCallId ?? m.role); + assert.deepStrictEqual(order, ['user', 'assistant', 'c1', 'user']); + }); + + it('answers a claimed approval whose run died before its result, in place and without running it again', async () => { + const sessionId = await newSession(); + const drop = tool('drop', { destructive: true, handler: async () => assert.fail('must not run twice') }); + await runAgent({ + sessionId, + models: stubModels([toolCall('c1', 'drop')]), + tools: [drop], + scopes, + maxTurns: 5, + autoApprove: false, + }); + const paused = await session.getSession(sessionId); + await session.resolveApproval(sessionId, paused.pendingApprovals[0].id, true); + assert.equal(await session.claimApproval(sessionId, paused.pendingApprovals[0].id), true); + // What an operator does after the crash: cancel the stuck run, then prompt again. + await session.cancelSession(sessionId); + await session.appendMessage(sessionId, { role: 'user', content: 'what happened?', createdAt: Date.now() }); + + let sentToModel; + const models = { + async generate(input) { + sentToModel = input.messages; + return { content: 'done', finishReason: 'stop' }; + }, + }; + await runAgent({ sessionId, models, tools: [drop], scopes, maxTurns: 5, autoApprove: false }); + const reloaded = await session.getSession(sessionId); + const order = reloaded.messages.map((m) => m.toolCallId ?? m.role); + assert.deepStrictEqual(order.slice(1, 4), ['assistant', 'c1', 'user'], 'the answer sits right after its tool call'); + const response = reloaded.messages.find((m) => m.toolCallId === 'c1'); + assert.deepStrictEqual(JSON.parse(response.content), { + ok: false, + error: 'interrupted', + outcome: 'unknown', + tool: 'drop', + }); + assert.ok( + sentToModel.findIndex((m) => m.toolCallId === 'c1') < + sentToModel.findIndex((m) => m.role === 'user' && m.content === 'what happened?') + ); + assert.equal(reloaded.status, 'completed'); + }); + + it('answers an interrupted claim on the latest turn when a provider reused its tool-call id', async () => { + const sessionId = await newSession(); + await session.appendMessage(sessionId, { + role: 'assistant', + content: '', + toolCalls: [{ id: 'c1', name: 'read', arguments: {} }], + createdAt: 1, + }); + await session.appendMessage(sessionId, { role: 'tool', content: '{}', toolCallId: 'c1', createdAt: 2 }); + await session.appendMessage(sessionId, { + role: 'assistant', + content: '', + toolCalls: [{ id: 'c1', name: 'drop', arguments: {} }], + createdAt: 3, + }); + const approval = await session.addPendingApproval(sessionId, { + toolName: 'drop', + arguments: {}, + toolCallId: 'c1', + reason: 'destructive', + }); + await session.resolveApproval(sessionId, approval.id, true); + await session.claimApproval(sessionId, approval.id); + + await session.answerInterruptedCalls(sessionId); + const roles = (await session.getSession(sessionId)).messages.map((m) => m.role); + assert.deepStrictEqual(roles, ['user', 'assistant', 'tool', 'assistant', 'tool']); + }); + + it('a failed run answers its open approvals and unanswered tool calls in its final write', async () => { + const sessionId = await newSession(); + const table = session.getAgentSessionTable(); + const put = table.put.bind(table); + let parked = 0; + table.put = async (row) => { + if (row.pendingApprovals.length === 2 && ++parked === 1) throw new Error('storage fault'); + return put(row); + }; + await assert.rejects( + runAgent({ + sessionId, + models: stubModels([ + { + content: '', + finishReason: 'tool_calls', + toolCalls: [ + { id: 'c1', name: 'drop', arguments: {} }, + { id: 'c2', name: 'drop', arguments: {} }, + ], + }, + ]), + tools: [tool('drop', { destructive: true })], + scopes, + maxTurns: 5, + autoApprove: false, + }), + /storage fault/ + ); + const failed = await session.getSession(sessionId); + assert.equal(failed.status, 'error'); + assert.ok(failed.pendingApprovals.every((a) => a.consumed)); + const answers = failed.messages.filter((m) => m.role === 'tool'); + assert.deepStrictEqual(answers.map((m) => m.toolCallId).sort(), ['c1', 'c2']); + assert.ok(answers.every((m) => JSON.parse(m.content).error === 'run_failed')); + }); + + it("a failed run answers a dead run's unanswered call in place, ahead of a later prompt", async () => { + const sessionId = await newSession(); + await session.appendMessage(sessionId, { + role: 'assistant', + content: '', + toolCalls: [{ id: 'c1', name: 'read', arguments: {} }], + createdAt: 1, + }); + await session.cancelSession(sessionId); + await session.appendMessage(sessionId, { role: 'user', content: 'what happened?', createdAt: 2 }); + const table = session.getAgentSessionTable(); + const put = table.put.bind(table); + let faulted = false; + table.put = async (row) => { + if (!faulted && row.messages.some((m) => m.role === 'tool')) { + faulted = true; + throw new Error('storage fault'); + } + return put(row); + }; + await assert.rejects( + runAgent({ sessionId, models: stubModels([]), tools: [tool('read')], scopes, maxTurns: 5, autoApprove: false }), + /storage fault/ + ); + const failed = await session.getSession(sessionId); + assert.deepStrictEqual( + failed.messages.map((m) => m.toolCallId ?? m.role), + ['user', 'assistant', 'c1', 'user'] + ); + assert.equal(JSON.parse(failed.messages[2].content).error, 'run_failed'); + }); + + it('answers an interrupted claim even when a budget stop ends the run first', async () => { + const sessionId = await newSession(); + await runAgent({ + sessionId, + models: stubModels([ + { + content: '', + finishReason: 'tool_calls', + toolCalls: [{ id: 'c1', name: 'drop', arguments: {} }], + usage: { promptTokens: 60, completionTokens: 10 }, + }, + ]), + tools: [tool('drop', { destructive: true })], + scopes, + maxTurns: 5, + autoApprove: false, + }); + const paused = await session.getSession(sessionId); + await session.resolveApproval(sessionId, paused.pendingApprovals[0].id, true); + await session.claimApproval(sessionId, paused.pendingApprovals[0].id); + await runAgent({ sessionId, models: stubModels([]), tools: [], scopes, maxTurns: 5, maxSessionTokens: 50 }); + const reloaded = await session.getSession(sessionId); + assert.equal(reloaded.status, 'error'); + assert.equal(JSON.parse(reloaded.messages.find((m) => m.toolCallId === 'c1').content).error, 'interrupted'); + }); + + it('counts a billed call that returns after the cancel', async () => { + const sessionId = await newSession(); + const controller = new AbortController(); + const models = { + async generate() { + controller.abort(); + await cancel(sessionId); + return { content: 'late', finishReason: 'stop', usage: { promptTokens: 40, completionTokens: 5 } }; + }, + }; + await runAgent({ sessionId, models, tools: noTools, scopes, maxTurns: 5, signal: controller.signal }); + const reloaded = await expectStillAborted(sessionId); + assert.deepStrictEqual(reloaded.rejectedCalls, [{ promptTokens: 40, completionTokens: 5 }]); + assert.equal(reloaded.messages.filter((m) => m.role === 'assistant').length, 0); + }); + + it('does not mark a run that was cancelled before it started as running', async () => { + const sessionId = await newSession(); + await cancel(sessionId); + const controller = new AbortController(); + controller.abort(); + let generateCalls = 0; + const models = { + async generate() { + generateCalls++; + return { content: 'unreached', finishReason: 'stop' }; + }, + }; + await runAgent({ sessionId, models, tools: noTools, scopes, maxTurns: 5, signal: controller.signal }); + assert.equal(generateCalls, 0); + await expectStillAborted(sessionId); + }); +}); diff --git a/unitTests/agent/operations.test.js b/unitTests/agent/operations.test.js index 454f01ca86..659e30f152 100644 --- a/unitTests/agent/operations.test.js +++ b/unitTests/agent/operations.test.js @@ -12,8 +12,10 @@ describe('agent/operations set_agent_config', () => { patches.push(patch); return patch; }, - startRun: () => {}, + startRun: () => true, cancelRun: () => false, + admitRun: () => 0, + isRunInFlight: () => false, }; const operation = buildOperations(deps).find((op) => op.name === 'set_agent_config'); return { execute: operation.execute, patches }; @@ -75,3 +77,265 @@ describe('agent/operations set_agent_config', () => { assert.deepStrictEqual(patches, [{ maxToolResultBytes: 1024 }, { maxToolResultBytes: 1_048_576 }]); }); }); + +describe('agent/operations budgets (#3040)', () => { + const session = require('#src/agent/session'); + + function operation(name, deps = {}) { + const patches = []; + const all = buildOperations({ + getConfig: () => ({}), + setConfig: (patch) => { + patches.push(patch); + return patch; + }, + startRun: () => true, + cancelRun: () => false, + admitRun: () => 0, + isRunInFlight: () => false, + ...deps, + }); + return { execute: all.find((op) => op.name === name).execute, patches }; + } + + it('set_agent_config applies maxSessionTokens and maxCostUsd, and null clears a cap', async () => { + const { execute, patches } = operation('set_agent_config'); + await execute({ operation: 'set_agent_config', maxSessionTokens: 200000, maxCostUsd: '1.5' }); + await execute({ operation: 'set_agent_config', maxCostUsd: null }); + assert.deepStrictEqual(patches, [{ maxSessionTokens: 200000, maxCostUsd: 1.5 }, { maxCostUsd: undefined }]); + }); + + it('set_agent_config rejects a maxTurns that is not a positive integer, applying nothing', async () => { + const { execute, patches } = operation('set_agent_config'); + for (const maxTurns of [null, 0, -1, 2.5, 'many', '']) { + await assert.rejects(execute({ operation: 'set_agent_config', maxTurns, autoApprove: true }), (err) => { + assert.strictEqual(err.statusCode, 400); + assert.match(err.message, /agent\.maxTurns must be a positive integer/); + return true; + }); + } + await execute({ operation: 'set_agent_config', maxTurns: '20' }); + assert.deepStrictEqual(patches, [{ maxTurns: 20 }]); + }); + + it('set_agent_config rejects a malformed cap and applies nothing from that request', async () => { + const { execute, patches } = operation('set_agent_config'); + await assert.rejects(execute({ operation: 'set_agent_config', maxTurns: 3, maxSessionTokens: -5 }), (err) => { + assert.strictEqual(err.statusCode, 400); + assert.match(err.message, /agent\.maxSessionTokens must be a positive integer/); + return true; + }); + assert.deepStrictEqual(patches, []); + }); + + describe('session reads carry folded usage', () => { + let store; + beforeEach(() => { + store = new Map(); + // Rows expose fields only as accessors, like a stored record, so a spread would copy nothing. + const asRecord = (row) => { + const record = {}; + for (const [key, value] of Object.entries(structuredClone(row))) { + Object.defineProperty(record, key, { get: () => value, enumerable: false }); + } + return record; + }; + session._setTableForTests({ + async put(row) { + store.set(row.session_id, structuredClone(row)); + }, + primaryStore: { + async get(key) { + return store.has(key) ? structuredClone(store.get(key)) : undefined; + }, + }, + async *search() { + for (const row of store.values()) yield asRecord(row); + }, + }); + }); + afterEach(() => session._setTableForTests(undefined)); + + async function seed() { + const created = await session.createSession({ user: 'admin' }); + await session.appendMessage(created.session_id, { role: 'user', content: 'hi', createdAt: 1 }); + await session.appendMessage(created.session_id, { + role: 'assistant', + content: 'ok', + createdAt: 2, + usage: { promptTokens: 7, completionTokens: 3, costUsd: 0.001 }, + }); + return created.session_id; + } + + it('cancel_agent_run answers and closes every open approval in the write that marks the session aborted', async () => { + const created = await session.createSession({ user: 'admin' }); + await session.appendMessage(created.session_id, { + role: 'assistant', + content: '', + createdAt: 1, + toolCalls: [ + { id: 'c1', name: 'drop', arguments: {} }, + { id: 'c2', name: 'drop', arguments: {} }, + ], + }); + const first = await session.addPendingApproval(created.session_id, { + toolName: 'drop', + arguments: {}, + toolCallId: 'c1', + reason: 'destructive', + }); + await session.addPendingApproval(created.session_id, { + toolName: 'drop', + arguments: {}, + toolCallId: 'c2', + reason: 'destructive', + }); + await session.resolveApproval(created.session_id, first.id, true); + await session.setStatus(created.session_id, 'awaiting_approval'); + + const result = await operation('cancel_agent_run').execute({ session_id: created.session_id }); + assert.deepStrictEqual(result, { cancelled: true, signalledLiveRun: false }); + const cancelled = await session.getSession(created.session_id); + assert.strictEqual(cancelled.status, 'aborted'); + assert.ok(cancelled.pendingApprovals.every((a) => a.resolved && a.consumed)); + const responses = cancelled.messages.filter((m) => m.role === 'tool'); + assert.deepStrictEqual(responses.map((m) => m.toolCallId).sort(), ['c1', 'c2']); + assert.ok(responses.every((m) => JSON.parse(m.content).error === 'cancelled_by_operator')); + }); + + it('agent_prompt refuses a session whose cancelled run is still finishing, appending nothing', async () => { + const created = await session.createSession({ user: 'admin' }); + await session.cancelSession(created.session_id); + const started = []; + const { execute } = operation('agent_prompt', { + getConfig: () => ({ enabled: true }), + startRun: (id) => started.push(id) > 0, + isRunInFlight: () => true, + }); + await assert.rejects(execute({ session_id: created.session_id, message: 'again' }), (err) => { + assert.strictEqual(err.statusCode, 409); + assert.match(err.message, /still has a run finishing/); + assert.strictEqual(err.code, 'AGENT_RUN_FINISHING'); + assert.strictEqual(err.retryable, true); + return true; + }); + assert.deepStrictEqual(started, []); + assert.strictEqual((await session.getSession(created.session_id)).messages.length, 0); + }); + + it('agent_prompt and approve_agent_action start the run with the token taken before their first await', async () => { + let cancels = 0; + const starts = []; + const deps = { + getConfig: () => ({ enabled: true }), + startRun: (id, admitted) => starts.push(admitted) > 0, + admitRun: () => cancels, + isRunInFlight: () => false, + }; + const created = await session.createSession({ user: 'admin' }); + const table = session.getAgentSessionTable(); + const put = table.put.bind(table); + table.put = async (row) => { + cancels++; + return put(row); + }; + await operation('agent_prompt', deps).execute({ session_id: created.session_id, message: 'again' }); + assert.deepStrictEqual(starts, [0], 'a cancel during the append fences out the run'); + + const approval = await session.addPendingApproval(created.session_id, { + toolName: 'drop', + arguments: {}, + toolCallId: 'c1', + reason: 'destructive', + }); + cancels = 0; + starts.length = 0; + await operation('approve_agent_action', deps).execute({ + session_id: created.session_id, + approval_id: approval.id, + }); + assert.deepStrictEqual(starts, [0], 'a cancel during the resolve fences out the run'); + }); + + it('agent_prompt answers 409 when a cancel fenced out the run it would start', async () => { + const created = await session.createSession({ user: 'admin' }); + const { execute } = operation('agent_prompt', { + getConfig: () => ({ enabled: true }), + startRun: () => false, + admitRun: () => 0, + isRunInFlight: () => false, + }); + await assert.rejects(execute({ session_id: created.session_id, message: 'again' }), (err) => { + assert.strictEqual(err.statusCode, 409); + assert.match(err.message, /recorded but not run/); + assert.strictEqual(err.code, 'AGENT_PROMPT_NOT_RUN'); + assert.strictEqual(err.retryable, false, 'the prompt was saved, so sending it again would duplicate it'); + return true; + }); + }); + + it('agent_prompt records nothing when a cancel arrives before the prompt is saved', async () => { + const created = await session.createSession({ user: 'admin' }); + let reads = 0; + const started = []; + const { execute } = operation('agent_prompt', { + getConfig: () => ({ enabled: true }), + startRun: (id) => started.push(id) > 0, + admitRun: () => (reads++ === 0 ? 0 : 1), + isRunInFlight: () => false, + }); + await assert.rejects(execute({ session_id: created.session_id, message: 'again' }), (err) => { + assert.strictEqual(err.statusCode, 409); + assert.match(err.message, /arrived while this prompt was being accepted$/); + assert.strictEqual(err.code, 'AGENT_PROMPT_CANCELLED'); + assert.strictEqual(err.retryable, true); + return true; + }); + assert.deepStrictEqual(started, []); + assert.strictEqual((await session.getSession(created.session_id)).messages.length, 0); + }); + + it('get_agent_session returns the row plus its usage', async () => { + const sessionId = await seed(); + const result = await operation('get_agent_session').execute({ session_id: sessionId }); + assert.strictEqual(result.session_id, sessionId); + assert.strictEqual(result.messages.length, 2); + assert.deepStrictEqual(result.usage, { calls: 1, promptTokens: 7, completionTokens: 3, costUsd: 0.001 }); + assert.strictEqual(result.runInFlight, false); + }); + + it('session reads say whether a run, even a cancelled one, is still in flight', async () => { + const sessionId = await seed(); + await session.cancelSession(sessionId); + const inFlight = { isRunInFlight: (id) => id === sessionId }; + const read = await operation('get_agent_session', inFlight).execute({ session_id: sessionId }); + assert.strictEqual(read.status, 'aborted'); + assert.strictEqual(read.runInFlight, true); + const { sessions } = await operation('list_agent_sessions', inFlight).execute({}); + assert.strictEqual(sessions[0].runInFlight, true); + }); + + it('agent_prompt marks a busy session retryable, since nothing was recorded', async () => { + const sessionId = await seed(); + await session.setStatus(sessionId, 'running'); + const { execute } = operation('agent_prompt', { getConfig: () => ({ enabled: true }) }); + await assert.rejects(execute({ session_id: sessionId, message: 'again' }), (err) => { + assert.strictEqual(err.statusCode, 409); + assert.strictEqual(err.code, 'AGENT_SESSION_BUSY'); + assert.strictEqual(err.retryable, true); + return true; + }); + assert.strictEqual((await session.getSession(sessionId)).messages.length, 2); + }); + + it('list_agent_sessions returns each row with usage, even when the row exposes fields as accessors', async () => { + const sessionId = await seed(); + const { sessions } = await operation('list_agent_sessions').execute({}); + assert.strictEqual(sessions.length, 1); + assert.strictEqual(sessions[0].session_id, sessionId); + assert.strictEqual(sessions[0].status, 'idle'); + assert.deepStrictEqual(sessions[0].usage, { calls: 1, promptTokens: 7, completionTokens: 3, costUsd: 0.001 }); + }); + }); +}); diff --git a/unitTests/resources/models/agentLoop.test.js b/unitTests/resources/models/agentLoop.test.js index 67950a9c04..720b8a83e6 100644 --- a/unitTests/resources/models/agentLoop.test.js +++ b/unitTests/resources/models/agentLoop.test.js @@ -4,15 +4,9 @@ const assert = require('node:assert'); // Prime module graph in the order other unit tests load it (avoids the ESM/CJS cycle // when transaction.ts is loaded ESM-first). require('#src/resources/databases'); -const { setGenerative, clearRegistry } = require('#src/resources/models/backendRegistry'); +const { setGenerative, setBackendPricing, clearRegistry } = require('#src/resources/models/backendRegistry'); const { Models } = require('#src/resources/models/Models'); -const { - BudgetExceededError, - ToolHandlerError, - _setComputeCallCostUsdForTests, - _resetComputeCallCostUsdForTests, - serializeToolResult, -} = require('#src/resources/models/agentLoop'); +const { BudgetExceededError, ToolHandlerError, serializeToolResult } = require('#src/resources/models/agentLoop'); const { logger } = require('#src/utility/logging/logger'); function makeMockWriter() { @@ -1214,10 +1208,6 @@ describe("agentLoop (toolMode: 'auto')", () => { }); describe('token + cost budgets', () => { - afterEach(() => { - _resetComputeCallCostUsdForTests(); - }); - // The default ScriptedBackend usage helpers set tokens to 1+1 per round; for these // tests we override per-call so the budget arithmetic is predictable. function tokenRound(content, tokens = { promptTokens: 50, completionTokens: 50 }) { @@ -1274,9 +1264,24 @@ describe("agentLoop (toolMode: 'auto')", () => { assert.strictEqual(result.content, 'over'); }); - it("maxCostUsd trips BudgetExceededError({kind: 'cost'}) when an injected cost function exceeds the cap", async () => { - // Inject a $0.01-per-call cost so two rounds = $0.02, capped at $0.015. - _setComputeCallCostUsdForTests(() => 0.01); + for (const [kind, cap] of [ + ['cost', { maxCostUsd: 0 }], + ['tokens', { maxToolTokens: 0 }], + ]) { + it(`a ${kind} cap of 0 makes no backend call`, async () => { + setBackendPricing(backend, { inputPerMtok: 100, outputPerMtok: 100 }); + backend.queue(tokenRound('unreached')); + await assert.rejects( + () => models.generate('q', { toolMode: 'auto', ...cap }), + (err) => err instanceof BudgetExceededError && err.kind === kind + ); + assert.strictEqual(backend.calls.length, 0); + }); + } + + it("maxCostUsd trips BudgetExceededError({kind: 'cost'}) from the serving entry's pricing", async () => { + // 100 tokens per round at $100/Mtok both ways = $0.01 per round; two rounds cross $0.015. + setBackendPricing(backend, { inputPerMtok: 100, outputPerMtok: 100 }); backend.queue(tokenToolRound('p', [tc('c1', 'echo', { i: 1 })]), tokenRound('done')); let caught; try { @@ -1293,30 +1298,16 @@ describe("agentLoop (toolMode: 'auto')", () => { assert.match(caught.message, /maxCostUsd=0.015/); }); - it('maxCostUsd does NOT trip with the v1 stub (computeCallCostUsd returns 0)', async () => { - // Default stub: cost = 0 per call, so any cap is unreachable today. Proves - // the v1 contract: cap is wired but doesn't fire in production until a - // real rate card lands. - backend.queue(tokenRound('done')); - const result = await models.generate('q', { toolMode: 'auto', maxCostUsd: 0.0001 }); - assert.strictEqual(result.content, 'done'); - }); - - it('cost function receives the usage object and model name', async () => { - const observations = []; - _setComputeCallCostUsdForTests((usage, model) => { - observations.push({ usage, model }); - return 0; + it('maxCostUsd under the priced total does not trip', async () => { + setBackendPricing(backend, { inputPerMtok: 100, outputPerMtok: 100 }); + backend.queue(tokenToolRound('p', [tc('c1', 'echo', { i: 1 })]), tokenRound('done')); + const result = await models.generate('q', { + toolMode: 'auto', + maxCostUsd: 0.05, + toolHandlers: { echo: (args) => args }, }); - backend.queue(tokenRound('done', { promptTokens: 5, completionTokens: 3 })); - // Use the default-registered backend; the model name passed in opts is what - // flows to computeCallCostUsd, independent of how the backend resolves. - await models.generate('q', { toolMode: 'auto' }); - assert.strictEqual(observations.length, 1); - assert.strictEqual(observations[0].usage.promptTokens, 5); - assert.strictEqual(observations[0].usage.completionTokens, 3); - // `opts.model` was undefined here; the cost function sees that. - assert.strictEqual(observations[0].model, undefined); + assert.strictEqual(result.content, 'done'); + assert.strictEqual(result.costUsd, 0.01, 'the terminal round carries its own cost'); }); it('caller abort during the in-flight backend round preempts the budget tally', async () => { @@ -1473,7 +1464,7 @@ describe("agentLoop (toolMode: 'auto')", () => { }; } - // Spy on logger.warn for the warn-once assertions, restoring the original after. + // Spy on logger.warn for the warn assertions, restoring the original after. let warnings; let origWarn; beforeEach(() => { @@ -1485,31 +1476,62 @@ describe("agentLoop (toolMode: 'auto')", () => { logger.warn = origWarn; }); - describe('budget warn-once when backend reports no usage', () => { - it('maxToolTokens set but no usage → warns once, does NOT trip, loop completes', async () => { - // Two tool rounds + terminal, all with usage: undefined. + describe('a budget the loop cannot measure stops it', () => { + const unenforceable = (kind) => (err) => + err instanceof BudgetExceededError && err.kind === kind && /cannot be enforced/.test(err.message); + + it('maxToolTokens with no usage reported stops after the first round, before its tools run', async () => { backend.queue( { output: { content: 'r1', finishReason: 'tool_calls', toolCalls: [tc('c1', 'echo', { i: 1 })] } }, - { output: { content: 'r2', finishReason: 'tool_calls', toolCalls: [tc('c2', 'echo', { i: 2 })] } }, { output: { content: 'final', finishReason: 'stop' } } ); - const result = await models.generate('q', { - toolMode: 'auto', - maxToolTokens: 1, // would trip instantly IF usage were measurable - toolHandlers: { echo: (a) => a }, + let ran = 0; + await assert.rejects( + () => + models.generate('q', { + toolMode: 'auto', + maxToolTokens: 1000, + toolHandlers: { echo: (a) => (ran++, a) }, + }), + unenforceable('tokens') + ); + assert.strictEqual(ran, 0); + }); + + for (const [label, usage] of [ + ['an empty', {}], + ['a half', { completionTokens: 5 }], + ['a negative', { promptTokens: -1, completionTokens: 5 }], + ]) { + it(`${label} token report counts as no usage for maxToolTokens`, async () => { + backend.queue( + { + output: { content: 'r1', finishReason: 'tool_calls', toolCalls: [tc('c1', 'echo', { i: 1 })] }, + usage, + }, + { output: { content: 'final', finishReason: 'stop' }, usage } + ); + await assert.rejects( + () => models.generate('q', { toolMode: 'auto', maxToolTokens: 1000, toolHandlers: { echo: (a) => a } }), + unenforceable('tokens') + ); }); - assert.strictEqual(result.content, 'final', 'unmeasurable budget must not falsely trip'); - const budgetWarns = warnings.filter((w) => /budget is unenforceable/.test(w)); - assert.strictEqual(budgetWarns.length, 1, 'warns exactly once across the run, not per-round'); + } + + it('maxCostUsd against an unpriced entry stops after the first round', async () => { + backend.queue(toolCallRound('r1', [tc('c1', 'echo', { i: 1 })]), final('final')); + await assert.rejects( + () => models.generate('q', { toolMode: 'auto', maxCostUsd: 5, toolHandlers: { echo: (a) => a } }), + unenforceable('cost') + ); }); it('maxToolTokens still trips normally when the backend DOES report usage', async () => { backend.queue(toolCallRound('r1', [tc('c1', 'echo', { i: 1 })]), final('unreached')); await assert.rejects( () => models.generate('q', { toolMode: 'auto', maxToolTokens: 1, toolHandlers: { echo: (a) => a } }), - (err) => err instanceof BudgetExceededError && err.kind === 'tokens' + (err) => err instanceof BudgetExceededError && err.kind === 'tokens' && /exceeded/.test(err.message) ); - assert.strictEqual(warnings.filter((w) => /budget is unenforceable/.test(w)).length, 0); }); }); diff --git a/unitTests/resources/models/pricing.test.js b/unitTests/resources/models/pricing.test.js new file mode 100644 index 0000000000..6022faff34 --- /dev/null +++ b/unitTests/resources/models/pricing.test.js @@ -0,0 +1,189 @@ +'use strict'; + +const assert = require('node:assert'); +// Prime the module graph in the same order the other models unit tests do. +require('#src/resources/databases'); +const { + setGenerative, + defineBackend, + clearRegistry, + getBackend, + getBackendPricing, + setBackendPricing, +} = require('#src/resources/models/backendRegistry'); +const { setFallbackGroup, clearRouting } = require('#src/resources/models/routing'); +const { bootstrapModels, applyModelsConfig, resetModelsProjection } = require('#src/resources/models/bootstrap'); +const { Models } = require('#src/resources/models/Models'); +const { getModelsConfigHash } = require('#src/resources/models/decisionStore'); + +const PRICING = { inputPerMtok: 3, outputPerMtok: 15 }; +const writer = () => ({ write: () => {} }); +const completedWith = (usage, output = { content: 'ok', finishReason: 'stop' }) => ({ + status: 'completed', + output, + usage, +}); + +function backendReturning(name, usage, output) { + return defineBackend({ name, generate: async () => completedWith(usage, output) }); +} + +describe('Models.generate pricing', () => { + beforeEach(() => { + clearRegistry(); + clearRouting(); + }); + afterEach(() => { + clearRegistry(); + clearRouting(); + }); + + it("prices a call from both token counts at the entry's rates", async () => { + const backend = backendReturning('priced', { promptTokens: 1_000_000, completionTokens: 200_000 }); + setGenerative('default', backend); + setBackendPricing(backend, PRICING); + const result = await new Models(writer()).generate('hi'); + assert.strictEqual(result.costUsd, 3 + 3); + assert.deepStrictEqual(result.usage, { promptTokens: 1_000_000, completionTokens: 200_000 }); + }); + + it('leaves an unpriced call without costUsd', async () => { + setGenerative('default', backendReturning('unpriced', { promptTokens: 10, completionTokens: 5 })); + const result = await new Models(writer()).generate('hi'); + assert.strictEqual('costUsd' in result, false); + assert.deepStrictEqual(result.usage, { promptTokens: 10, completionTokens: 5 }); + }); + + for (const [label, usage] of [ + ['an empty usage object', {}], + ['a missing completion count', { promptTokens: 10 }], + ['a non-finite count', { promptTokens: Number.NaN, completionTokens: 5 }], + ['a negative count', { promptTokens: -1, completionTokens: 5 }], + ['no usage at all', undefined], + ]) { + it(`does not price ${label}, rather than under-price it`, async () => { + const backend = backendReturning('partial', usage); + setGenerative('default', backend); + setBackendPricing(backend, PRICING); + const result = await new Models(writer()).generate('hi'); + assert.strictEqual(result.costUsd, undefined); + }); + } + + it('prices zero tokens at zero, not as unmeasured', async () => { + const backend = backendReturning('zero', { promptTokens: 0, completionTokens: 0 }); + setGenerative('default', backend); + setBackendPricing(backend, PRICING); + assert.strictEqual((await new Models(writer()).generate('hi')).costUsd, 0); + }); + + it('drops usage and costUsd a backend put on its own output', async () => { + setGenerative( + 'default', + backendReturning('smuggler', undefined, { + content: 'ok', + finishReason: 'stop', + usage: { promptTokens: 1, completionTokens: 1 }, + costUsd: 0, + }) + ); + const result = await new Models(writer()).generate('hi'); + assert.strictEqual('usage' in result, false); + assert.strictEqual('costUsd' in result, false); + assert.strictEqual(result.content, 'ok'); + }); + + it('prices a fallback-served call at the fallback entry rate, not the requested one', async () => { + const primary = defineBackend({ + name: 'primary', + generate: async () => { + throw new Error('provider down'); + }, + }); + const backup = backendReturning('backup', { promptTokens: 1_000_000, completionTokens: 0 }); + setGenerative('default', primary); + setGenerative('backup', backup); + setFallbackGroup('generative', 'default', ['backup']); + setBackendPricing(primary, { inputPerMtok: 100, outputPerMtok: 100 }); + setBackendPricing(backup, { inputPerMtok: 1, outputPerMtok: 1 }); + const result = await new Models(writer()).generate('hi'); + assert.strictEqual(result.costUsd, 1); + }); + + it('a completed result with no output is a failed attempt: the next candidate serves, and alone it rejects', async () => { + const records = []; + setGenerative('default', defineBackend({ name: 'hollow', generate: async () => ({ status: 'completed' }) })); + setGenerative('backup', backendReturning('backup', { promptTokens: 1, completionTokens: 1 })); + setFallbackGroup('generative', 'default', ['backup']); + const result = await new Models({ write: (record) => records.push(record) }).generate('hi'); + assert.strictEqual(result.content, 'ok'); + assert.deepStrictEqual( + records.map((r) => [r.backend, r.success]), + [ + ['hollow', false], + ['backup', true], + ] + ); + clearRouting(); + await assert.rejects(new Models(writer()).generate('hi'), /completed without output/); + }); + + it('an unpriced fallback leaves the call unpriced even when the requested entry is priced', async () => { + const primary = defineBackend({ + name: 'primary', + generate: async () => { + throw new Error('provider down'); + }, + }); + setGenerative('default', primary); + setGenerative('backup', backendReturning('backup', { promptTokens: 10, completionTokens: 10 })); + setFallbackGroup('generative', 'default', ['backup']); + setBackendPricing(primary, PRICING); + const result = await new Models(writer()).generate('hi'); + assert.strictEqual(result.costUsd, undefined); + }); +}); + +describe('models.generative..pricing from config', () => { + beforeEach(() => { + clearRegistry(); + clearRouting(); + resetModelsProjection(); + }); + afterEach(() => { + clearRegistry(); + clearRouting(); + resetModelsProjection(); + }); + + it('attaches the declared pricing to the generative backend the entry builds', async () => { + await bootstrapModels({ + models: { + generative: { default: { backend: 'ollama', model: 'm', pricing: PRICING } }, + embedding: { default: { backend: 'ollama', model: 'e' } }, + }, + }); + assert.deepStrictEqual(getBackendPricing(getBackend('generative', 'default')), PRICING); + assert.strictEqual(getBackendPricing(getBackend('embedding', 'default')), undefined); + }); + + it('leaves the decision config hash unchanged by a price edit', async () => { + await bootstrapModels({ models: { generative: { default: { backend: 'ollama', model: 'm', pricing: PRICING } } } }); + const priced = getModelsConfigHash(); + await applyModelsConfig({ + generative: { default: { backend: 'ollama', model: 'm', pricing: { inputPerMtok: 1, outputPerMtok: 2 } } }, + }); + assert.strictEqual(getModelsConfigHash(), priced); + await applyModelsConfig({ generative: { default: { backend: 'ollama', model: 'n', pricing: PRICING } } }); + assert.notStrictEqual(getModelsConfigHash(), priced, 'a real setting still changes it'); + }); + + it('a reload that changes or removes the price reprices the entry', async () => { + await bootstrapModels({ models: { generative: { default: { backend: 'ollama', model: 'm', pricing: PRICING } } } }); + const repriced = { inputPerMtok: 1, outputPerMtok: 2 }; + await applyModelsConfig({ generative: { default: { backend: 'ollama', model: 'm', pricing: repriced } } }); + assert.deepStrictEqual(getBackendPricing(getBackend('generative', 'default')), repriced); + await applyModelsConfig({ generative: { default: { backend: 'ollama', model: 'm' } } }); + assert.strictEqual(getBackendPricing(getBackend('generative', 'default')), undefined); + }); +}); diff --git a/unitTests/resources/models/sourceIdentity.test.js b/unitTests/resources/models/sourceIdentity.test.js index ac916b02da..c38706b7e1 100644 --- a/unitTests/resources/models/sourceIdentity.test.js +++ b/unitTests/resources/models/sourceIdentity.test.js @@ -122,12 +122,22 @@ describe('score-source identity for calibration (#2841)', () => { describe('configured entries carry a fingerprint of their settings', () => { const entry = (extra = {}) => ({ backend: FLAKY, answer: '{"value":"bug"}', ...extra }); - it('ignores credentials and the fallback list, and changes with revision or any other setting', () => { + it('ignores credentials, the fallback list and pricing, and changes with revision or any other setting', () => { const base = sourceFingerprint('generative', undefined, { backend: 'openai', model: 'm', apiKey: 'a' }); assert.strictEqual( sourceFingerprint('generative', undefined, { backend: 'openai', model: 'm', apiKey: 'b', fallback: ['x'] }), base ); + assert.strictEqual( + sourceFingerprint('generative', undefined, { + backend: 'openai', + model: 'm', + apiKey: 'a', + pricing: { inputPerMtok: 3, outputPerMtok: 15 }, + }), + base, + 'a price edit must not restart calibration' + ); assert.notStrictEqual( sourceFingerprint('generative', undefined, { backend: 'openai', model: 'm', apiKey: 'a', revision: 'r2' }), base diff --git a/unitTests/validation/configValidator.test.js b/unitTests/validation/configValidator.test.js index d604d93f52..176f231df3 100644 --- a/unitTests/validation/configValidator.test.js +++ b/unitTests/validation/configValidator.test.js @@ -1011,6 +1011,52 @@ describe('Test configValidator module', () => { expect(result.error).to.be.undefined; }); + describe('pricing', () => { + const PRICING = { inputPerMtok: 3, outputPerMtok: 15 }; + const validate = (models) => { + const config = baseConfig(); + config.models = models; + return configValidator(config, true).error; + }; + + it('accepts pricing on generative entries, built-in or module-backed, including zero rates', () => { + for (const backend of ['ollama', 'openai', 'anthropic', 'bedrock', '@acme/harper-llm']) { + const entry = { backend, model: 'm', pricing: PRICING }; + if (backend === 'openai' || backend === 'anthropic') entry.apiKey = 'k'; + expect(validate({ generative: { default: entry } }), backend).to.be.undefined; + } + expect( + validate({ + generative: { default: { backend: 'ollama', pricing: { inputPerMtok: 0, outputPerMtok: 0 } } }, + }) + ).to.be.undefined; + }); + + it('rejects a partial, negative, non-numeric or extended price', () => { + for (const pricing of [ + { inputPerMtok: 3 }, + { outputPerMtok: 15 }, + { inputPerMtok: -1, outputPerMtok: 15 }, + { inputPerMtok: 'three', outputPerMtok: 15 }, + { ...PRICING, cachedInputPerMtok: 0.3 }, + ]) { + expect(validate({ generative: { default: { backend: 'ollama', pricing } } }), JSON.stringify(pricing)).to.not + .be.undefined; + } + }); + + it('refuses pricing on embedding and decision entries, which nothing prices', () => { + for (const models of [ + { embedding: { default: { backend: 'ollama', pricing: PRICING } } }, + { embedding: { default: { backend: '@acme/harper-embed', pricing: PRICING } } }, + { decision: { default: { backend: 'generative', pricing: PRICING } } }, + { decision: { default: { backend: '@acme/harper-decision', pricing: PRICING } } }, + ]) { + expect(validate(models), JSON.stringify(models)).to.not.be.undefined; + } + }); + }); + it('accepts a generative entry with requestTimeoutMs', () => { const config = baseConfig(); config.models = { diff --git a/utility/hdbTerms.ts b/utility/hdbTerms.ts index eaf390f856..1a29833755 100644 --- a/utility/hdbTerms.ts +++ b/utility/hdbTerms.ts @@ -671,6 +671,7 @@ export const CONFIG_PARAMS = { AGENT_MODEL: 'agent_model', AGENT_MAXTURNS: 'agent_maxTurns', AGENT_MAXTOKENS: 'agent_maxTokens', + AGENT_MAXSESSIONTOKENS: 'agent_maxSessionTokens', AGENT_MAXTOOLRESULTBYTES: 'agent_maxToolResultBytes', AGENT_MAXCOSTUSD: 'agent_maxCostUsd', AGENT_AUTOAPPROVE: 'agent_autoApprove', diff --git a/validation/configValidator.ts b/validation/configValidator.ts index f75d82c7c8..1dffe0fb0a 100644 --- a/validation/configValidator.ts +++ b/validation/configValidator.ts @@ -206,6 +206,14 @@ const generativeDecisionEntrySchema = Joi.object({ const unknownBackendEntrySchema = Joi.object({ backend: string.required(), }).unknown(true); +// USD per million tokens. Only generative calls are priced, so every other kind refuses the field +// rather than accepting a price nothing reads. +const pricingSchema = Joi.object({ + inputPerMtok: number.min(0).required(), + outputPerMtok: number.min(0).required(), +}).unknown(false); +const withPricingFor = (kind, schema) => + schema.keys({ pricing: kind === 'generative' ? pricingSchema.optional() : Joi.forbidden() }); // A built-in that cannot serve the kind it is listed under fails here, with the kind named, rather // than passing as a module specifier and failing at boot. function rejectWrongKind(kind) { @@ -216,20 +224,20 @@ const wrongKindEntrySchema = (kind) => const providerEntrySchema = (kind) => Joi.alternatives().conditional('.backend', { switch: [ - { is: 'ollama', then: ollamaEntrySchema }, - { is: 'openai', then: openaiEntrySchema }, - { is: 'anthropic', then: anthropicEntrySchema }, - { is: 'bedrock', then: bedrockEntrySchema }, + { is: 'ollama', then: withPricingFor(kind, ollamaEntrySchema) }, + { is: 'openai', then: withPricingFor(kind, openaiEntrySchema) }, + { is: 'anthropic', then: withPricingFor(kind, anthropicEntrySchema) }, + { is: 'bedrock', then: withPricingFor(kind, bedrockEntrySchema) }, { is: 'generative', then: wrongKindEntrySchema(kind) }, ], - otherwise: unknownBackendEntrySchema, + otherwise: withPricingFor(kind, unknownBackendEntrySchema), }); const decisionEntrySchema = Joi.alternatives().conditional('.backend', { switch: [ { is: 'generative', then: generativeDecisionEntrySchema }, { is: Joi.valid('ollama', 'openai', 'anthropic', 'bedrock'), then: wrongKindEntrySchema('decision') }, ], - otherwise: unknownBackendEntrySchema, + otherwise: withPricingFor('decision', unknownBackendEntrySchema), }); const modelsSchema = Joi.object({ embedding: Joi.object().pattern(Joi.string(), providerEntrySchema('embedding')).optional(),