From e6fa8e0cb61f6256cf3b25d5d55ba89e731e4e87 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Mon, 14 Sep 2026 17:42:12 -0700 Subject: [PATCH] Let an operator bound what a turn and a session spend. The model decides what to ask for, so the thing that says no has to sit where the model does not reach: the supervisor, between steps. `OPENCODE_TEMPORAL_BUDGET_*` set a turn's tokens and wall clock, a deadline that stops the running step, and the same two bounds for the whole session; what a session spent is carried across a rollover and read back off its own row, so an idle session woken again does not get its allowance back. --- packages/core/src/session/runner/index.ts | 3 + packages/core/src/session/runner/llm.ts | 9 +- .../src/session/runner/publish-llm-event.ts | 20 +++ .../test/session-runner-model-call.test.ts | 40 ++++++ packages/opencode/src/cli/cmd/detached.ts | 6 +- packages/temporal/README.md | 4 + packages/temporal/src/config.ts | 45 +++++++ packages/temporal/src/drain.ts | 47 +++++++- packages/temporal/src/executor.ts | 5 + packages/temporal/src/stepped-drain.ts | 38 ++++-- packages/temporal/src/stepped-turn.ts | 31 ++++- packages/temporal/src/supervisor.ts | 109 ++++++++++++++++- packages/temporal/src/workflow.ts | 35 +++++- packages/temporal/test/config.test.ts | 42 +++++++ packages/temporal/test/drain-spend.test.ts | 113 +++++++++++++++++ .../test/session-supervisor-ceiling.test.ts | 114 +++++++++++++++++- packages/temporal/test/stepped-turn.test.ts | 13 +- 17 files changed, 639 insertions(+), 35 deletions(-) create mode 100644 packages/temporal/test/config.test.ts create mode 100644 packages/temporal/test/drain-spend.test.ts diff --git a/packages/core/src/session/runner/index.ts b/packages/core/src/session/runner/index.ts index f7347d60fe39..386bd34ce218 100644 --- a/packages/core/src/session/runner/index.ts +++ b/packages/core/src/session/runner/index.ts @@ -33,6 +33,9 @@ export interface StepResult { readonly continue: boolean readonly step: number readonly promotion: SessionInput.Delivery | undefined + /** What this step's provider attempt was billed for, for a caller keeping a budget. Absent from a + * step that made no attempt, and from one whose attempt reported no counts. */ + readonly spent?: { readonly tokens: number } } /** A tool call the provider asked for, recorded but not run, handed to the caller to dispatch. Every diff --git a/packages/core/src/session/runner/llm.ts b/packages/core/src/session/runner/llm.ts index 443c1031fdfb..ecc008310cf9 100644 --- a/packages/core/src/session/runner/llm.ts +++ b/packages/core/src/session/runner/llm.ts @@ -45,7 +45,7 @@ import { Service, } from "./index" import { SessionRunnerModel } from "./model" -import { createLLMEventPublisher, emitToolResult, record } from "./publish-llm-event" +import { billed, createLLMEventPublisher, emitToolResult, record } from "./publish-llm-event" import { toLLMMessages } from "./to-llm-message" import { MAX_STEPS_PROMPT } from "./max-steps" import { DEFAULT_MAX_STEPS, REPEAT_LIMIT, REPEATED_CALLS_PROMPT, trailingIdenticalToolSteps } from "./loop-guard" @@ -859,7 +859,12 @@ const layer = Layer.effect( const prologue = yield* stepPrologue(input) if (prologue.kind === "settled") return prologue.result const result = yield* runTurn(input.sessionID, prologue.promotion, input.step) - return yield* stepContinuation(input.sessionID, result.needsContinuation, result.step) + const closed = yield* stepContinuation(input.sessionID, result.needsContinuation, result.step) + // The attempt's own counts, carried out for a caller keeping a budget. The split path gets + // them from the settlement it hands back; this path publishes its own step ending, so + // without this a whole-step deployment has no bound on what a turn spends. + const spent = billed(result.settlement) + return spent ? { ...closed, spent } : closed }) const runModelCall = Effect.fn("SessionRunner.runModelCall")(function* (input: StepInput) { diff --git a/packages/core/src/session/runner/publish-llm-event.ts b/packages/core/src/session/runner/publish-llm-event.ts index 272074a1d34a..ef08b3a0dd07 100644 --- a/packages/core/src/session/runner/publish-llm-event.ts +++ b/packages/core/src/session/runner/publish-llm-event.ts @@ -38,6 +38,26 @@ export interface StepSettlement { readonly tokens: ReturnType } +/** + * What the attempt was billed for, as one number, for a budget to add up. Cache reads and writes + * are billed too, and at different rates, so they are counted rather than dropped: a bound that + * ignores them is one a long context walks straight through. + * + * Undefined when the attempt said nothing, which a budget must read as "unknown" rather than as + * nothing spent. + * + * The split path in `packages/temporal/src/l2-step.ts` writes this arithmetic out again, because + * that file is bundled into the workflow sandbox and may not import this one. Change one, change + * the other. + */ +export const billed = (settlement: StepSettlement | undefined) => { + const t = settlement?.tokens + if (!t) return undefined + return { + tokens: (t.input ?? 0) + (t.output ?? 0) + (t.reasoning ?? 0) + (t.cache?.read ?? 0) + (t.cache?.write ?? 0), + } +} + /** The shape a tool call's input is recorded in. Exported so a dispatcher publishing Tool.Called * for a deferred call records it exactly as the streaming path would. */ export const record = (value: unknown): Record => diff --git a/packages/core/test/session-runner-model-call.test.ts b/packages/core/test/session-runner-model-call.test.ts index 758decf3ac25..bed9a6cbe7ab 100644 --- a/packages/core/test/session-runner-model-call.test.ts +++ b/packages/core/test/session-runner-model-call.test.ts @@ -44,6 +44,25 @@ const callsTool: LLMClientShape["stream"] = () => LLMEvent.toolCall({ id: "call_probe", name: "probe_write", input: {} }), LLMEvent.stepFinish({ index: 0, reason: "tool-calls" }), ]) +// The same, with what the provider says it billed for. A budget outside the agent adds these up, +// so what the whole-step path carries out of a turn is as load-bearing as what it recorded. +const callsToolAndBills: LLMClientShape["stream"] = () => + Stream.fromIterable([ + LLMEvent.stepStart({ index: 0 }), + LLMEvent.toolCall({ id: "call_probe", name: "probe_write", input: {} }), + LLMEvent.stepFinish({ + index: 0, + reason: "tool-calls", + usage: { + inputTokens: 10, + nonCachedInputTokens: 8, + outputTokens: 4, + reasoningTokens: 1, + cacheReadInputTokens: 2, + cacheWriteInputTokens: 3, + }, + }), + ]) // Calls the tool that declines. const callsDecliningTool: LLMClientShape["stream"] = () => Stream.fromIterable([ @@ -310,6 +329,27 @@ describe("SessionRunner model-only attempt", () => { }), ) + runnerHarness(callsToolAndBills).effect("carries what the attempt was billed for out of a whole step", () => + Effect.gen(function* () { + yield* seedSession(sessionID) + yield* registerProbes(counters()) + const runner = yield* SessionRunner.Service + + const result = yield* runner.runStep({ + sessionID, + step: 2, + promotion: undefined, + first: false, + force: false, + }) + + // 8 not-cached in, 3 visible out (4 less the 1 of reasoning), 1 reasoning, 2 cache read and + // 3 cache write. A deployment that never splits a step has no other way to know what a turn + // cost, and every one of those five is billed. + expect(result.spent).toEqual({ tokens: 17 }) + }), + ) + // A step ending is not a turn ending: a steer or a queued prompt continues the same turn through // another step, so the one place that decides publishes the turn's own ending. runnerHarness(textOnly).effect("says the turn ended, once, when nothing follows it", () => diff --git a/packages/opencode/src/cli/cmd/detached.ts b/packages/opencode/src/cli/cmd/detached.ts index ad8c85f717cb..69d094156910 100644 --- a/packages/opencode/src/cli/cmd/detached.ts +++ b/packages/opencode/src/cli/cmd/detached.ts @@ -182,7 +182,11 @@ export const SessionScheduleCommand = cmd({ { sessionID, text: args.prompt, - session: { idleTimeout: config.idleTimeout, stepped: config.stepped === true }, + session: { + idleTimeout: config.idleTimeout, + stepped: config.stepped === true, + ...(config.budget ? { budget: config.budget } : {}), + }, }, ], }, diff --git a/packages/temporal/README.md b/packages/temporal/README.md index f6756948ceb7..2fc9d919d5ee 100644 --- a/packages/temporal/README.md +++ b/packages/temporal/README.md @@ -65,6 +65,9 @@ session runs as the workflow `session-exec-`. | `OPENCODE_TEMPORAL_PROFILE` | `local` | `fleet` sets the defaults a multi-host deployment needs and refuses combinations that cannot work. See below. | | `OPENCODE_TEMPORAL_API_KEY`, `OPENCODE_TEMPORAL_API_KEY_FILE` | | Temporal Cloud credentials, the second read from a file. | | `OPENCODE_TEMPORAL_TLS_CERT`, `OPENCODE_TEMPORAL_TLS_KEY`, `OPENCODE_TEMPORAL_TLS_CA` | | A certificate pair for a cluster with mTLS. `OPENCODE_TEMPORAL_TLS=1` for TLS without a client certificate. | +| `OPENCODE_TEMPORAL_BUDGET_TOKENS`, `OPENCODE_TEMPORAL_BUDGET_SECONDS` | | What one turn may spend before the supervisor stops driving it, checked between steps. | +| `OPENCODE_TEMPORAL_BUDGET_HARD_SECONDS` | | A deadline for one turn. It stops the step that is running, where the other bounds let it finish. | +| `OPENCODE_TEMPORAL_BUDGET_SESSION_TOKENS`, `OPENCODE_TEMPORAL_BUDGET_SESSION_SECONDS` | | The same two bounds for the whole session, across every turn it runs. | | `OPENCODE_EVENT_POLL_MS` | | How often a live subscriber re-reads the log for events another process appended. `0` turns the tick off. | | `OPENCODE_WATCH_POLL_MS` | 30000 | How often `session watch` asks the running set whether the turn ended in a gap it could not see. | | `OPENCODE_DB` | | One absolute path shared by every process on a host. | @@ -345,6 +348,7 @@ TEMPORAL_ADDRESS=temporal.internal:7233 \ | A fresh worker receives a project | Needs the directory | Packs rebuild the tracked files at the recorded path; ignored files do not travel | | A later turn reuses a directory an abandoned tool may still write | Not reachable | The directory is refused until the call can be shown to be over, then moved aside and rebuilt | | A turn never stops stepping | The coordinator's own loop | The supervisor stops after 200 steps and says so | +| A turn spends more than intended | Nothing bounds it | An operator's budget on the supervisor stops the turn on tokens or wall clock | An unknown tool outcome is a loss of evidence, not proof that the tool stopped or failed. The model can ask for a new call after reading it. A non-idempotent external effect needs an idempotency key, diff --git a/packages/temporal/src/config.ts b/packages/temporal/src/config.ts index 04e72dac9143..965c739b8f63 100644 --- a/packages/temporal/src/config.ts +++ b/packages/temporal/src/config.ts @@ -56,6 +56,18 @@ export interface Interface { /** Whether the store is shared across hosts (`OPENCODE_DB_URL`). A fleet without it is a set of * workers that cannot see each other's sessions. */ readonly sharedStore: boolean + /** What one turn, and the whole session, may spend before the supervisor stops driving it. Off + * unless set. Tokens are what the provider reported billed, added up as each attempt reports; + * seconds are wall clock on the workflow's own clock. `hardSeconds` is a deadline that stops the + * step that is running, where the other bounds wait for it to finish. Rides the workflow input, + * like the mode does, because the sandbox cannot read env and a rollover has to carry it. */ + readonly budget?: { + readonly tokens?: number + readonly seconds?: number + readonly hardSeconds?: number + readonly sessionTokens?: number + readonly sessionSeconds?: number + } } export class Service extends Context.Service()("@opencode/temporal/Config") {} @@ -64,6 +76,30 @@ const read = (path: string | undefined) => (path ? readFileSync(path, "utf8") : const given = (name: string) => process.env[name] !== undefined && process.env[name] !== "" const onOff = (name: string, fallback: boolean) => (given(name) ? process.env[name] === "1" : fallback) +const BUDGET_VARS = { + tokens: "OPENCODE_TEMPORAL_BUDGET_TOKENS", + seconds: "OPENCODE_TEMPORAL_BUDGET_SECONDS", + hardSeconds: "OPENCODE_TEMPORAL_BUDGET_HARD_SECONDS", + sessionTokens: "OPENCODE_TEMPORAL_BUDGET_SESSION_TOKENS", + sessionSeconds: "OPENCODE_TEMPORAL_BUDGET_SESSION_SECONDS", +} as const + +// A number a bound can be read as, or nothing. What is set but unreadable is left for `preflight` +// to refuse: a bound that silently reads as none is worse than a process that will not start. +const bound = (name: string) => { + if (!given(name)) return undefined + const value = Number(process.env[name]) + return Number.isFinite(value) && value > 0 ? value : undefined +} + +const budgetFromEnv = (): Interface["budget"] => { + const set = Object.entries(BUDGET_VARS).flatMap(([key, name]) => { + const value = bound(name) + return value === undefined ? [] : [[key, value] as const] + }) + return set.length === 0 ? undefined : Object.fromEntries(set) +} + export const fromEnv = (): Interface => { const stepAffinity = process.env.OPENCODE_TEMPORAL_STEP_AFFINITY !== "0" const worktreeAffinity = process.env.OPENCODE_TEMPORAL_WORKTREE_AFFINITY === "1" @@ -96,6 +132,7 @@ export const fromEnv = (): Interface => { // neither kind of affinity keeping the step's tools on one worker. serialTools: onOff("OPENCODE_TEMPORAL_SERIAL_TOOLS", !stepAffinity && !worktreeAffinity && sharedStore), sharedStore, + budget: budgetFromEnv(), } } @@ -146,6 +183,9 @@ export const preflight = (config: Interface): string[] => { problems.push("an API key is set but TEMPORAL_NAMESPACE is `default`, which is not a Cloud namespace") if (!!process.env.OPENCODE_TEMPORAL_TLS_CERT !== !!process.env.OPENCODE_TEMPORAL_TLS_KEY) problems.push("OPENCODE_TEMPORAL_TLS_CERT and OPENCODE_TEMPORAL_TLS_KEY come as a pair") + for (const name of Object.values(BUDGET_VARS)) + if (given(name) && bound(name) === undefined) + problems.push(`${name} is \`${process.env[name]}\`, which is not a positive number`) return problems } @@ -176,4 +216,9 @@ export const describe = (config: Interface): Record => ({ stepAffinity: String(config.stepAffinity !== false), serialTools: String(config.serialTools === true), credentials: config.apiKey ? "api key" : config.tls ? "certificate pair" : "none (plaintext)", + budget: config.budget + ? Object.entries(config.budget) + .map(([name, value]) => `${name}=${value}`) + .join(" ") + : "none", }) diff --git a/packages/temporal/src/drain.ts b/packages/temporal/src/drain.ts index 6597b6b106fd..a5aca3561035 100644 --- a/packages/temporal/src/drain.ts +++ b/packages/temporal/src/drain.ts @@ -32,14 +32,51 @@ export interface StepDrainResult { continue: boolean step: number promotion: string | null + /** What this step's provider attempt was billed for, as the attempt itself reported it. The + * supervisor adds these up for the turn's budget. Absent from a step that made no attempt. */ + spent?: { readonly tokens: number } + /** What the session has been billed in total, as its own row holds it. Read off the record rather + * than added up by the supervisor, so it survives everything the supervisor does not: a run that + * rolled over, a session that went idle and was woken again, a turn some other client ran. */ + session?: { readonly tokens: number } +} + +/** + * What the session has been billed in total, off its own row. Every count the record keeps, because + * a provider bills for all of them and a bound that drops the cached ones is one a long context + * walks straight through. + * + * Read before the step runs, so it lags by the step now running. A session's bound overshoots by + * the same step a turn's does, and for the same reason: what a step cost is not known until it is + * taken. + */ +export const sessionSpend = (session: SessionSchema.Info | undefined) => { + const t = session?.tokens + // Nothing rather than zero when the row has no counts: a bound that reads a missing number as + // nothing spent is one that never stops a session. + if (!t) return undefined + return { + tokens: (t.input ?? 0) + (t.output ?? 0) + (t.reasoning ?? 0) + (t.cache?.read ?? 0) + (t.cache?.write ?? 0), + } } /** The step result as it crosses the activity boundary. A missing session reads as a step with * nothing left to do, which is what a resume waiting on it should see. */ -export const toStepResult = (step: number, result?: SessionRunner.StepResult): StepDrainResult => - result === undefined - ? { continue: false, step, promotion: null } - : { continue: result.continue, step: result.step, promotion: result.promotion ?? null } +export const toStepResult = ( + step: number, + result?: SessionRunner.StepResult, + session?: SessionSchema.Info, +): StepDrainResult => { + if (result === undefined) return { continue: false, step, promotion: null } + const total = sessionSpend(session) + return { + continue: result.continue, + step: result.step, + promotion: result.promotion ?? null, + ...(result.spent ? { spent: result.spent } : {}), + ...(total ? { session: total } : {}), + } +} export interface DrainDeps { readonly store: SessionStore.Interface @@ -109,7 +146,7 @@ export const makeDrains = (deps: DrainDeps) => { first: input.first, force: input.force, }) - .pipe(Effect.map((result) => toStepResult(input.step, result))), + .pipe(Effect.map((result) => toStepResult(input.step, result, session))), ).pipe(Effect.map((result) => result ?? toStepResult(input.step))), // A whole step turns a declined permission into an interrupt to halt its own loop, so an // interrupt with nothing cancelling it is that refusal and nothing else. diff --git a/packages/temporal/src/executor.ts b/packages/temporal/src/executor.ts index d96f9cf20e1c..b42f79b6ee02 100644 --- a/packages/temporal/src/executor.ts +++ b/packages/temporal/src/executor.ts @@ -70,6 +70,9 @@ const layer = Layer.effect( // Only the client can read whether the store is shared, so whether a step's tools may overlap // is decided here and rides the workflow input. const SERIAL_TOOLS = config.serialTools === true + // The operator's bound on what a turn may spend. It rides the workflow input like the mode does, + // because the sandbox cannot read env and a rollover has to carry it. + const BUDGET = config.budget const AFFINITY = config.worktreeAffinity === true // The tree this process serves when affinity is on. A serve process with an embedded worker is // already sitting in it, so the process directory is the right default. @@ -226,6 +229,7 @@ const layer = Layer.effect( idleTimeout: IDLE_TIMEOUT, stepped: STEPPED, serialTools: SERIAL_TOOLS, + ...(BUDGET ? { budget: BUDGET } : {}), } satisfies WF.SessionTurnOptions, ], signal: WF.wake, @@ -292,6 +296,7 @@ const layer = Layer.effect( idleTimeout: IDLE_TIMEOUT, stepped: STEPPED, serialTools: SERIAL_TOOLS, + ...(BUDGET ? { budget: BUDGET } : {}), } satisfies WF.SessionTurnOptions, ], workflowIdConflictPolicy: "USE_EXISTING", diff --git a/packages/temporal/src/stepped-drain.ts b/packages/temporal/src/stepped-drain.ts index eac9084d7f2c..c67448e514ef 100644 --- a/packages/temporal/src/stepped-drain.ts +++ b/packages/temporal/src/stepped-drain.ts @@ -13,7 +13,7 @@ import type { StepSettlement } from "@opencode-ai/core/session/runner/publish-ll import type { SessionInput } from "@opencode-ai/core/session/input" import { WorktreeMaterializer } from "@opencode-ai/core/session/execution/worktree" import { runAtBoundary } from "./boundary" -import { type InSession, type StepDrainInput, type StepDrainResult, toStepResult } from "./drain" +import { type InSession, type StepDrainInput, type StepDrainResult, sessionSpend, toStepResult } from "./drain" /** The provider attempt of one step. Same shape as a whole-step drain: the difference is what it * does with the tool calls, not what it needs to start. */ @@ -36,6 +36,8 @@ export type ModelCallDrainResult = * worker was not given a queue of its own, and never required: the step falls back to the * shared queue and the tree is rebuilt there. */ readonly queue?: string + /** What the session has been billed in total, off its own row. See `sessionSpend`. */ + readonly session?: { readonly tokens: number } } export interface ToolCallDrainInput { @@ -90,19 +92,28 @@ export const makeSteppedDrains = ({ inSession, worktrees, stepQueue }: SteppedDr input.sessionID, signal, inSession(input.sessionID, input.owner, { claim: true }, (runner, session) => - runner.runModelCall({ - sessionID: session.id, - step: input.step, - promotion: (input.promotion ?? undefined) as SessionInput.Delivery | undefined, - first: input.first, - force: input.force, - }), + runner + .runModelCall({ + sessionID: session.id, + step: input.step, + promotion: (input.promotion ?? undefined) as SessionInput.Delivery | undefined, + first: input.first, + force: input.force, + }) + // The session's own totals, off the row this drain already loaded. The attempt says what + // it cost; the row says what the session has cost, which is what outlives this run. + .pipe(Effect.map((called) => ({ ...called, session: sessionSpend(session) }))), ).pipe( Effect.map( (result): ModelCallDrainResult => result === undefined || result.kind === "settled" ? { kind: "settled", result: toStepResult(input.step, result?.result) } - : { ...result, owner: input.owner, ...(stepQueue === undefined ? {} : { queue: stepQueue }) }, + : { + ...result, + owner: input.owner, + ...(stepQueue === undefined ? {} : { queue: stepQueue }), + ...(result.session ? { session: result.session } : {}), + }, ), ), ) @@ -156,7 +167,14 @@ export const makeSteppedDrains = ({ inSession, worktrees, stepQueue }: SteppedDr needsContinuation: input.needsContinuation, withoutTheTree: input.withoutTheTree, }) - .pipe(Effect.map((result) => toStepResult(input.step, result))), + // Only the loop decision: what the step spent was reported by its model call. + .pipe( + Effect.map((result) => ({ + continue: result.continue, + step: result.step, + promotion: result.promotion ?? null, + })), + ), ).pipe(Effect.map((result) => result ?? toStepResult(input.step))), ) diff --git a/packages/temporal/src/stepped-turn.ts b/packages/temporal/src/stepped-turn.ts index 8f57ae0dfa6b..47200a7036f3 100644 --- a/packages/temporal/src/stepped-turn.ts +++ b/packages/temporal/src/stepped-turn.ts @@ -175,20 +175,45 @@ export const makeSteppedTurn = } else { dispatched.push(...(await Promise.allSettled(model.calls.map(dispatch)))) } + // What the attempt cost, from the counts the provider returned with it. Cache reads and writes + // are billed too, and at different rates, so they are counted rather than dropped: a bound that + // ignores them is one a long context walks straight through. + // + // The same arithmetic as `billed` in `session/runner/publish-llm-event.ts`, which the whole-step + // path uses, and written out again here because this file is bundled into the workflow sandbox + // and may not import core at runtime. Change one, change the other. + const tokens = model.settlement?.tokens + const spent = tokens + ? { + tokens: + (tokens.input ?? 0) + + (tokens.output ?? 0) + + (tokens.reasoning ?? 0) + + (tokens.cache?.read ?? 0) + + (tokens.cache?.write ?? 0), + } + : undefined + // What the session has been billed in total, as its own row held it when this attempt opened it. + const session = model.session + const withSpend = (result: StepDrainResult): StepDrainResult => ({ + ...result, + ...(spent ? { spent } : {}), + ...(session ? { session } : {}), + }) const sealing = (stopped: boolean): SealDrainInput => ({ sessionID: input.sessionID, step: model.step, // A stopped step is not one that continues. The settlement carries the model's own finish // reason, and for a step that asked for tools that is `tool-calls`, which every follower - // reads as "another step follows". Passing it through on the way out would record a turn the + // reads as "another step follows". Passing it through on the way out recorded a turn the // user stopped as a turn still going. settlement: stopped && model.settlement ? { ...model.settlement, finish: "stop" } : model.settlement, assistantMessageID: model.assistantMessageID, needsContinuation: stopped ? false : model.needsContinuation, owner: model.owner, }) - const seal = (stopped: boolean) => viaPinned((on) => on.sealStep(sealing(stopped)), stopped) + const seal = (stopped: boolean) => viaPinned((on) => on.sealStep(sealing(stopped)), stopped).then(withSpend) for (const outcome of dispatched) { if (outcome.status !== "rejected") continue @@ -240,7 +265,7 @@ export const makeSteppedTurn = // inside one. Between the dispatch failing and the next step claiming the log there is a // window where nothing fences that host, and the only thing that makes the window harmless // is that nobody else publishes during it. - return activities.sealStep({ ...sealing(false), withoutTheTree: true }) + return activities.sealStep({ ...sealing(false), withoutTheTree: true }).then(withSpend) } const resumes = () => uncertain !== undefined if (resumes()) return closeElsewhere() diff --git a/packages/temporal/src/supervisor.ts b/packages/temporal/src/supervisor.ts index 9f2a17bd72b8..768e3beb27c9 100644 --- a/packages/temporal/src/supervisor.ts +++ b/packages/temporal/src/supervisor.ts @@ -44,12 +44,24 @@ export interface SupervisorRuntime { readonly allHandlersFinished?: () => boolean /** Restart the run with fresh history, carrying whether work is still pending. History-keeping * drivers only (Temporal). */ - readonly continueAsNew?: (sessionID: string, startWithWake: boolean) => Promise + readonly continueAsNew?: ( + sessionID: string, + startWithWake: boolean, + /** What the session has spent so far. Carried, or a session gets its allowance back every time + * it outgrows a run's history. */ + spent?: Spent, + ) => Promise /** Whether the driver says this run's history is large enough to roll over. A drain count cannot * answer this: one drain is a whole turn, and a stepped turn of 200 steps is thousands of events, * so a handful of drains can cross the server's limit long before the count does. Optional: * drivers without a history return false. */ readonly historyWantsRollover?: () => boolean + /** The workflow's own clock, so a turn's elapsed time reads the same on a replay. Optional: + * a driver without one is not asked for a wall-clock bound. */ + readonly now?: () => number + /** Wait, in the driver's own time. Only used for a turn's deadline, which is the one bound that + * has to fire while a step is running rather than between two of them. */ + readonly sleep?: (ms: number) => Promise /** Where a turn says it stopped because it ran out of steps rather than because it finished. */ readonly warn?: (message: string, attributes: Record) => void } @@ -61,6 +73,44 @@ export interface WorkflowOptions { readonly maxDrainsPerRun?: number /** Steps one turn may take before the supervisor stops driving it. */ readonly maxStepsPerTurn?: number + /** + * What one turn may spend before the supervisor stops driving it. Enforced between steps, never + * inside one: a step that has started is left to finish, because stopping it would leave a tool + * call the next attempt's transcript cannot be made from. + * + * This is what the workflow having the loop buys. The agent is not asked to keep to it and cannot + * be: the model decides what to ask for, so the thing that says no has to be somewhere the model + * does not reach. Off unless an operator sets it, because a bound that ends real work is worse + * than none and only the operator knows which is which. + */ + readonly budget?: { + // Tokens the turn's provider attempts may spend, added up as each one reports. + readonly tokens?: number + // Wall clock for the whole turn, in seconds. A number rather than a duration string: the parser + // for those is not something a workflow bundle should be reaching for, and a bound nobody can + // be sure of the units of is worse than one that says them. + readonly seconds?: number + // Wall clock again, but as a deadline rather than a bound checked between steps. At `seconds` + // the turn stops where it can; at this one it stops where it is, which is what a user pressing + // stop does: what has results keeps them, what is in flight comes back as an outcome nobody can + // vouch for, and the tools over there keep running until they are done. Opt-in on top. + readonly hardSeconds?: number + // And the same two for the session, across every turn it runs. A per-turn bound says what one + // answer may cost; this says what the whole session may, which is the number somebody is + // billed for. + readonly sessionTokens?: number + readonly sessionSeconds?: number + } + /** What the session had spent when this run started, from the run that rolled over. */ + readonly spent?: Spent +} + +/** What a session has spent so far, carried between runs of its workflow. */ +export interface Spent { + readonly tokens: number + // Wall clock its turns have used, rather than how long the session has existed: one sitting idle + // overnight has spent nothing. + readonly seconds: number } export const makeSupervisor = (rt: SupervisorRuntime, options?: WorkflowOptions) => { @@ -74,6 +124,16 @@ export const makeSupervisor = (rt: SupervisorRuntime, options?: WorkflowOptions) // the supervisor can see it, because each step is its own activity and each one succeeds. High // enough that real work never reaches it. const MAX_STEPS_PER_TURN = options?.maxStepsPerTurn ?? 200 + const BUDGET = options?.budget + // What this session has spent, across every turn it has run, carried in from the run that rolled + // over. A session does not get its allowance back by outgrowing a run's history. + const spent = { tokens: options?.spent?.tokens ?? 0, seconds: options?.spent?.seconds ?? 0 } + const outOfBudget = (turnTokens: number, turnSeconds: number, recorded?: number) => + BUDGET !== undefined && + ((BUDGET.tokens !== undefined && turnTokens > BUDGET.tokens) || + (BUDGET.seconds !== undefined && turnSeconds > BUDGET.seconds) || + (BUDGET.sessionTokens !== undefined && (recorded ?? spent.tokens + turnTokens) > BUDGET.sessionTokens) || + (BUDGET.sessionSeconds !== undefined && spent.seconds + turnSeconds > BUDGET.sessionSeconds)) // Each step (one provider attempt + its tools) is its own activity; the step loop is supervisor // control flow (step / promotion / first mirror SessionRunner.run's loop). `startWithWake` is the @@ -99,6 +159,30 @@ export const makeSupervisor = (rt: SupervisorRuntime, options?: WorkflowOptions) let step = 1 let promotion: string | null = null let first = true + // A deadline that ends the turn where it is. The wait lives in this drain's own scope, so + // a turn that finishes first cancels it on the way out, and the rejection that follows is + // the scope's own. + if (BUDGET?.hardSeconds !== undefined && rt.sleep) { + void rt.sleep(BUDGET.hardSeconds * 1000).then( + () => { + rt.warn?.("turn stopped where it was: its deadline passed", { + sessionID, + seconds: BUDGET.hardSeconds, + }) + rt.cancelCurrentScope() + }, + () => undefined, + ) + } + // What this turn has spent, and when it started spending it. Both are the supervisor's + // own state: the session's log holds what every turn spent, and this bound is about one. + const startedAt = rt.now?.() ?? 0 + let tokens = 0 + // What the record says the session has been billed, as of the last step that reported it. + // The count below is about a run; this one is about the session, and it is what a + // session's bound is measured against wherever the host can say it. + let recorded: { readonly tokens: number } | undefined + const turnSeconds = () => (rt.now ? (rt.now() - startedAt) / 1000 : 0) for (let taken = 0; ; taken++) { if (taken >= MAX_STEPS_PER_TURN) { rt.warn?.("turn hit the step ceiling and was left where it stopped", { @@ -107,7 +191,25 @@ export const makeSupervisor = (rt: SupervisorRuntime, options?: WorkflowOptions) }) break } + // Between steps, never inside one, and after the step whose cost crossed the bound: + // what a step spent is not known until it has been taken. + if (outOfBudget(tokens, turnSeconds(), recorded?.tokens)) { + rt.warn?.("turn stopped where it was: it is out of budget", { + sessionID, + steps: taken, + turn: { tokens, seconds: Math.round(turnSeconds()) }, + session: { + tokens: recorded?.tokens ?? spent.tokens + tokens, + seconds: Math.round(spent.seconds + turnSeconds()), + recorded: recorded !== undefined, + }, + budget: BUDGET, + }) + break + } const r: StepDrainResult = await rt.runTurnStep({ sessionID, step, promotion, first, force }) + tokens += r.spent?.tokens ?? 0 + recorded = r.session ?? recorded // Inside the loop as well, because one drain is a whole turn: a long one outgrows the // history without ever reaching the next drain's check. if (rt.historyWantsRollover?.()) rolloverPending = true @@ -124,6 +226,9 @@ export const makeSupervisor = (rt: SupervisorRuntime, options?: WorkflowOptions) promotion = r.promotion first = false } + // Whatever ended the turn, the provider billed for what it ran. + spent.tokens += tokens + spent.seconds += turnSeconds() }) .finally(() => { inFlight = null @@ -171,7 +276,7 @@ export const makeSupervisor = (rt: SupervisorRuntime, options?: WorkflowOptions) // continue-as-new: only from the main method, only when quiescent. Carry the pending wake so // queued work survives the boundary (a pure rollover carries false -> no spurious drain). if (rt.continueAsNew && rolloverPending && quiescent()) { - await rt.continueAsNew(sessionID, pendingWake) + await rt.continueAsNew(sessionID, pendingWake, spent) } if (pendingWake) { pendingWake = false diff --git a/packages/temporal/src/workflow.ts b/packages/temporal/src/workflow.ts index 7cfcdc26f10e..f14eec127461 100644 --- a/packages/temporal/src/workflow.ts +++ b/packages/temporal/src/workflow.ts @@ -28,7 +28,7 @@ import { WorkflowExecutionAlreadyStartedError } from "@temporalio/common" import type { StepActivities, SteppedTurnActivities } from "./activities" import { isHaltFailure, makeSteppedTurn } from "./stepped-turn" import { SIGNALS, RESUME_UPDATE, WORKFLOW_ID_PREFIX } from "./protocol" -import { makeSupervisor, type SupervisorRuntime } from "./supervisor" +import { makeSupervisor, type SupervisorRuntime, type WorkflowOptions } from "./supervisor" const activityOptions = { // The heartbeat is the liveness bound (it stops within seconds of a worker death and Temporal @@ -133,6 +133,9 @@ const runtime: SupervisorRuntime = { // The server's own read of whether this run has grown enough to roll over. The drain count alone // misses it: a stepped turn is thousands of events, so a handful of drains can cross the limit. historyWantsRollover: () => workflowInfo().continueAsNewSuggested, + // The workflow's clock, not the machine's: it reads the same on a replay, and so does its timer. + now: () => Date.now(), + sleep: (ms) => sleep(ms), warn: (message, attributes) => log.warn(message, attributes), } @@ -174,6 +177,11 @@ export interface SessionTurnOptions { * on two hosts each publish a tree without the other's work. The client decides, because only it * can read whether the store is shared. */ readonly serialTools?: boolean + /** What one turn may spend before the supervisor stops driving it, on tokens or on wall clock. + * Off unless set. See `WorkflowOptions` in the supervisor for what it does and does not bound. */ + readonly budget?: WorkflowOptions["budget"] + /** What the session had spent when this run started. Set by a rollover, not by a client. */ + readonly spent?: WorkflowOptions["spent"] } export async function sessionTurn(sessionID: string, options?: SessionTurnOptions): Promise { @@ -182,16 +190,31 @@ export async function sessionTurn(sessionID: string, options?: SessionTurnOption const idleTimeout = options?.idleTimeout const stepped = options?.stepped === true const serialTools = options?.serialTools === true - if (!idleTimeout && !stepped) return workflows.sessionTurn(sessionID, startWithWake) + const budget = options?.budget + if (!idleTimeout && !stepped && !budget) return workflows.sessionTurn(sessionID, startWithWake) return makeSupervisor( { ...(stepped ? steppedRuntime(serialTools) : runtime), // The mode has to survive the boundary, or a long session silently reverts to whole-step - // activities the first time it rolls over. - continueAsNew: (id, wake) => - continueAsNew(id, { startWithWake: wake, idleTimeout, stepped, serialTools }), + // activities the first time it rolls over. The budget travels for the same reason: a turn + // that starts after a rollover is a turn, and it is the operator's bound either way. + continueAsNew: (id, wake, spent) => + continueAsNew(id, { + startWithWake: wake, + idleTimeout, + stepped, + serialTools, + budget, + spent, + }), }, - idleTimeout ? { idleTimeout } : undefined, + idleTimeout || budget + ? { + ...(idleTimeout ? { idleTimeout } : {}), + ...(budget ? { budget } : {}), + ...(options?.spent ? { spent: options.spent } : {}), + } + : undefined, ).sessionTurn(sessionID, startWithWake) } diff --git a/packages/temporal/test/config.test.ts b/packages/temporal/test/config.test.ts new file mode 100644 index 000000000000..6a3a9efe71c9 --- /dev/null +++ b/packages/temporal/test/config.test.ts @@ -0,0 +1,42 @@ +// The budget reaches the supervisor through the workflow input, and the workflow input is built from +// what this module read. A bound that is set and read as nothing is the one failure an operator +// cannot see, so what is set but unreadable refuses to start rather than reading as none. +import { afterEach, expect, it } from "bun:test" +import { TemporalConfig } from "../src/config" + +const VARS = [ + "OPENCODE_TEMPORAL_BUDGET_TOKENS", + "OPENCODE_TEMPORAL_BUDGET_SECONDS", + "OPENCODE_TEMPORAL_BUDGET_HARD_SECONDS", + "OPENCODE_TEMPORAL_BUDGET_SESSION_TOKENS", + "OPENCODE_TEMPORAL_BUDGET_SESSION_SECONDS", +] + +afterEach(() => { + for (const name of VARS) delete process.env[name] +}) + +it("reads no budget when nothing is set", () => { + expect(TemporalConfig.fromEnv().budget).toBeUndefined() +}) + +it("reads the bounds that are set and leaves the rest off", () => { + process.env.OPENCODE_TEMPORAL_BUDGET_TOKENS = "250000" + process.env.OPENCODE_TEMPORAL_BUDGET_SESSION_SECONDS = "3600" + + const config = TemporalConfig.fromEnv() + + // Off means absent, not zero: the supervisor reads a missing bound as no bound. + expect(config.budget).toEqual({ tokens: 250000, sessionSeconds: 3600 }) + expect(TemporalConfig.preflight(config)).toEqual([]) + expect(TemporalConfig.describe(config).budget).toBe("tokens=250000 sessionSeconds=3600") +}) + +it("refuses a bound it cannot read as a positive number", () => { + process.env.OPENCODE_TEMPORAL_BUDGET_SECONDS = "ten minutes" + + const config = TemporalConfig.fromEnv() + + expect(config.budget).toBeUndefined() + expect(TemporalConfig.preflight(config).some((p) => p.includes("OPENCODE_TEMPORAL_BUDGET_SECONDS"))).toBe(true) +}) diff --git a/packages/temporal/test/drain-spend.test.ts b/packages/temporal/test/drain-spend.test.ts new file mode 100644 index 000000000000..16643009ebdc --- /dev/null +++ b/packages/temporal/test/drain-spend.test.ts @@ -0,0 +1,113 @@ +// What a step reports it was billed for, which is what the turn's budget adds up. +// +// Two paths report it and they have to report the same number: the split step takes it off the +// settlement the attempt hands back, and the whole-step drain takes it off the result the runner +// returns. Before that second one existed, a deployment that never split a step had no bound on +// what a turn spends, only on how many steps it took. + +import { expect, it } from "bun:test" +import { Context, Effect, Layer } from "effect" +import { billed } from "@opencode-ai/core/session/runner/publish-llm-event" +import { SessionRunner } from "@opencode-ai/core/session/runner" +import { makeDrains, type DrainDeps } from "../src/drain" +import { makeSteppedTurn, type SteppedActivities } from "../src/stepped-turn" + +// One settlement, and both paths have to make the same number of it. The workflow file writes this +// arithmetic out again because it is bundled into the sandbox and may not import core, so nothing +// but this holds the two copies together. +const SETTLEMENT = { + finish: "stop", + tokens: { input: 100, output: 20, reasoning: 5, cache: { read: 1000, write: 30 } }, +} +const BILLED = { tokens: 1155 } + +it("counts everything the provider bills for, cache included", () => { + // A long context is mostly cache reads. A bound that drops them is one it walks straight through. + expect(billed(SETTLEMENT)).toEqual(BILLED) +}) + +it("makes the same number of it in the workflow, which cannot import that one", async () => { + const activities: SteppedActivities = { + runModelCall: async () => ({ + kind: "called", + step: 1, + calls: [], + settlement: SETTLEMENT, + owner: "run:1:1", + }), + runToolCall: async () => ({ outcome: "settled" }), + sealStep: async () => ({ continue: false, step: 2, promotion: null }), + } + + const result = await makeSteppedTurn({ + activities, + isCancellation: () => false, + isHalt: () => false, + })({ sessionID: "ses_spend", step: 1, promotion: null, first: true, force: false }) + + expect(result.spent).toEqual(BILLED) +}) + +it("says nothing, rather than nothing spent, when the attempt reported no counts", () => { + expect(billed(undefined)).toBeUndefined() +}) + +// The session row carries its own running totals, which is where a session's bound reads from. +const SESSION_ROW = { + id: "ses_spend", + location: { directory: "/project" }, + tokens: { input: 40, output: 5, reasoning: 1, cache: { read: 200, write: 4 } }, +} + +const drainsOver = (runStep: SessionRunner.Interface["runStep"]) => + makeDrains({ + store: { get: () => Effect.succeed(SESSION_ROW) } as unknown as DrainDeps["store"], + locations: { + get: () => Layer.succeed(SessionRunner.Service, { runStep } as unknown as SessionRunner.Interface), + } as unknown as DrainDeps["locations"], + ctx: Context.empty() as DrainDeps["ctx"], + events: { claim: () => Effect.void } as unknown as DrainDeps["events"], + worktrees: { + ensure: () => Effect.void, + beginWrite: () => Effect.void, + endWrite: () => Effect.void, + } as unknown as DrainDeps["worktrees"], + }) + +it("carries what a whole step spent out to the supervisor", async () => { + const { stepDrain } = drainsOver((() => + Effect.succeed({ + continue: true, + step: 2, + promotion: undefined, + spent: { tokens: 1155 }, + })) as unknown as SessionRunner.Interface["runStep"]) + + const result = await stepDrain( + { sessionID: "ses_spend", step: 1, promotion: null, first: true, force: false }, + new AbortController().signal, + ) + + expect(result.spent).toEqual({ tokens: 1155 }) + // And what the session has been billed in total, off its own row: 40 in, 5 out, 1 reasoning, + // 200 cache read, 4 cache write. That number outlives this run, where the one above does not. + expect(result.session).toEqual({ tokens: 250 }) +}) + +it("carries nothing out when the step had nothing to say", async () => { + const { stepDrain } = drainsOver((() => + Effect.succeed({ + continue: false, + step: 2, + promotion: undefined, + })) as unknown as SessionRunner.Interface["runStep"]) + + const result = await stepDrain( + { sessionID: "ses_spend", step: 1, promotion: null, first: true, force: false }, + new AbortController().signal, + ) + + expect(result.spent).toBeUndefined() + // The session's own total is not about this step, so it is reported either way. + expect(result.session).toEqual({ tokens: 250 }) +}) diff --git a/packages/temporal/test/session-supervisor-ceiling.test.ts b/packages/temporal/test/session-supervisor-ceiling.test.ts index 6a31cb260ca7..8e79de4f26a0 100644 --- a/packages/temporal/test/session-supervisor-ceiling.test.ts +++ b/packages/temporal/test/session-supervisor-ceiling.test.ts @@ -2,6 +2,9 @@ // tool forever, or a step that keeps handing itself back because its host keeps dying. Every one of // those steps is a model call somebody pays for, and each of them succeeds, so nothing below the // supervisor can see it. Driven by a fake runtime whose steps always ask for another. +// +// The ceiling is a runaway guard with a fixed number. The budget below it is policy: what an +// operator decided this turn may spend, which is the thing the workflow having the loop is for. import { it, expect } from "bun:test" import { makeSupervisor, type SupervisorRuntime } from "../src/supervisor" import type { StepDrainResult } from "../src/activities" @@ -30,10 +33,34 @@ class LoopingRuntime implements SupervisorRuntime { setUpdateHandler = () => {} // Always another step, up to a bound of its own: without one a ceiling that failed to hold would // hang this test rather than fail it. + // As expensive and as slow as the case wants, and never finished unless it says so. + tokensPerStep = 0 + secondsPerStep = 0 + clock = 0 + now = () => this.clock finishAfter = 50 - runTurnStep = async () => (++this.steps < this.finishAfter ? MORE : DONE) + // What the record says the session has been billed, which the drains read off the session's own + // row. Unset means a host that cannot say, and then the supervisor's own count is what is used. + recorded: number | undefined = undefined + runTurnStep = async () => { + if (this.realTimePerStep) await new Promise((r) => setTimeout(r, this.realTimePerStep)) + if (this.cancelled) return DONE + this.clock += this.secondsPerStep * 1000 + const session = this.recorded === undefined ? {} : { session: { tokens: this.recorded } } + return ++this.steps < this.finishAfter + ? { ...MORE, spent: { tokens: this.tokensPerStep }, ...session } + : { ...DONE, ...session } + } + // Real time rather than the workflow clock, for the one bound that is a timer. + realTimePerStep = 0 runInDrainScope = (fn: () => Promise) => fn() - cancelCurrentScope = () => {} + // What a cancellation does here: the step loop is a plain loop, so the fake ends it by refusing + // to take another step, which is what the scope's cancellation does to the real one. + cancelled = false + cancelCurrentScope = () => { + this.cancelled = true + } + sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms)) isCancellation = () => false isRootCancelled = () => false warn = (message: string) => { @@ -50,11 +77,88 @@ it("stops driving a turn that keeps asking for another step", async () => { expect(rt.warnings).toEqual(["turn hit the step ceiling and was left where it stopped"]) }) -it("leaves a turn that finishes on its own alone", async () => { +it("stops a turn that has spent the tokens the operator allowed it", async () => { + const rt = new LoopingRuntime() + rt.tokensPerStep = 100 + // Three steps at 100 is 300, which is over 250. The bound is read after a step, because what a + // step spent is not known until it has been taken. + await makeSupervisor(rt, { budget: { tokens: 250 }, idleTimeout: "1 millisecond" }).sessionTurn("ses_tokens") + + expect(rt.steps).toBe(3) + expect(rt.warnings).toEqual(["turn stopped where it was: it is out of budget"]) +}) + +it("stops a turn that has run for the time the operator allowed it", async () => { const rt = new LoopingRuntime() - rt.finishAfter = 3 - await makeSupervisor(rt, { idleTimeout: "1 millisecond" }).sessionTurn("ses_finishes") + rt.secondsPerStep = 1 + await makeSupervisor(rt, { budget: { seconds: 2 }, idleTimeout: "1 millisecond" }).sessionTurn("ses_seconds") expect(rt.steps).toBe(3) + expect(rt.warnings).toEqual(["turn stopped where it was: it is out of budget"]) +}) + +it("leaves a turn inside its bound alone", async () => { + const inside = new LoopingRuntime() + inside.tokensPerStep = 1 + await makeSupervisor(inside, { budget: { tokens: 250 }, idleTimeout: "1 millisecond" }).sessionTurn("ses_inside") + expect(inside.steps).toBe(50) + expect(inside.warnings).toEqual([]) +}) + +it("bounds what a whole session spends, not only what one turn does", async () => { + // A per-turn bound says what one answer may cost. This is the number somebody is billed for, and + // a session of cheap turns can pass every per-turn bound and still run all night. + const rt = new LoopingRuntime() + rt.tokensPerStep = 100 + rt.finishAfter = 2 + const supervisor = makeSupervisor(rt, { + budget: { tokens: 10_000, sessionTokens: 250 }, + idleTimeout: "1 millisecond", + }) + await supervisor.sessionTurn("ses_session_budget") + + // Two steps at 100 is inside the turn's own bound and inside the session's, so the turn answers. + expect(rt.steps).toBe(2) expect(rt.warnings).toEqual([]) }) + +it("carries what a session spent across a rollover", async () => { + // Without the carry a busy session gets its allowance back every time it outgrows a run's + // history, which is the one thing a long session is sure to do. + const rt = new LoopingRuntime() + rt.tokensPerStep = 100 + await makeSupervisor(rt, { + budget: { sessionTokens: 250 }, + spent: { tokens: 200, seconds: 0 }, + idleTimeout: "1 millisecond", + }).sessionTurn("ses_carried") + + // 200 already spent, so the first step's 100 crosses 250 and the second is not scheduled. + expect(rt.steps).toBe(1) + expect(rt.warnings).toEqual(["turn stopped where it was: it is out of budget"]) +}) + +it("measures a session against what the record says it spent, not against this run's count", async () => { + // A session that went idle and was woken again is a fresh run with an empty count of its own. + // What it has been billed is in the session's own row, and the drains report it. + const rt = new LoopingRuntime() + rt.tokensPerStep = 10 + rt.recorded = 400 + await makeSupervisor(rt, { budget: { sessionTokens: 250 }, idleTimeout: "1 millisecond" }).sessionTurn("ses_recorded") + + // One step, whose ten tokens are nowhere near the bound: the record is what crosses it. + expect(rt.steps).toBe(1) + expect(rt.warnings).toEqual(["turn stopped where it was: it is out of budget"]) +}) + +it("stops a turn at its deadline, inside the step that was running", async () => { + // Every other bound is read between two steps, so a step that has started is left to finish. A + // deadline is the one that fires while one is running, which is what a user pressing stop does. + const rt = new LoopingRuntime() + rt.realTimePerStep = 40 + await makeSupervisor(rt, { budget: { hardSeconds: 0.1 }, idleTimeout: "1 millisecond" }).sessionTurn("ses_deadline") + + expect(rt.cancelled).toBe(true) + expect(rt.steps).toBeLessThan(50) + expect(rt.warnings).toEqual(["turn stopped where it was: its deadline passed"]) +}) diff --git a/packages/temporal/test/stepped-turn.test.ts b/packages/temporal/test/stepped-turn.test.ts index 36fee9799c4c..0a1dfa3403c0 100644 --- a/packages/temporal/test/stepped-turn.test.ts +++ b/packages/temporal/test/stepped-turn.test.ts @@ -55,6 +55,16 @@ const fakes = ( } describe("stepped turn", () => { + it("reports nothing for a step whose attempt said nothing about tokens", async () => { + const { activities } = fakes({ kind: "called", step: 2, calls: [call("call_a")], owner: "run-1:model-1:1" }) + + const result = await makeSteppedTurn({ activities, isCancellation, isHalt })(INPUT) + + // Absent rather than zero: a budget that reads a missing count as nothing spent is one that + // never stops a turn whose attempts do not report. + expect(result.spent).toBeUndefined() + }) + it("dispatches nothing and does not seal when the step is already settled", async () => { const settled: StepDrainResult = { continue: false, step: 2, promotion: null } const { activities, tools, seals } = fakes({ kind: "settled", result: settled }) @@ -87,7 +97,8 @@ describe("stepped turn", () => { expect(seals[0]?.owner).toBe("run-1:model-1:1") // The finish reason lives only in the attempt's memory, so the seal has to be handed it. expect(seals[0]?.settlement?.finish).toBe("tool-calls") - expect(result).toEqual(SEALED) + // And what the step cost rides back with the seal's own result, for the turn's budget. + expect(result).toEqual({ ...SEALED, spent: { tokens: 3 } }) }) it("reports the calls it did not settle, and says nothing when they all settled", async () => {