diff --git a/changelog.d/fix-67-end-turn-accounting.fixed.md b/changelog.d/fix-67-end-turn-accounting.fixed.md new file mode 100644 index 000000000..dcf24db7e --- /dev/null +++ b/changelog.d/fix-67-end-turn-accounting.fixed.md @@ -0,0 +1,30 @@ +- **A turn ended by a tool result (`endTurn`, such as `skip_reply`) is recorded as a completed + turn** (#67). It emits `inference:completed` with the stream's cumulative usage, then + `inference:turn_ended`, writes a successful inference-log entry (`stopReason: 'turn_ended'`), + and adds its usage once to the session totals (`usage:updated`, `getSessionUsage()`). Like a + completion, it resets the consecutive-failure streak and the poison-history rewind budget, ends + a refusal-rewind episode (a forced `/unstick` reports that the model responded), and records a + provider-cooldown recovery. In hosts such as connectome-host, consumers that end a turn's display + on `inference:completed` (a typing indicator, notices that close on success, a voice relay) now + end it on these turns too, and a consumer that totals `inference:completed`'s `tokenUsage` now + counts them too; one that also adds each `inference:usage` sample counts them twice. +- **Every stream's billed rounds reach the session totals once, however the stream ends.** A + context-budget or physical-window restart is counted and logged where it is decided + (`stopReason: 'context_budget'` or `'physical_window'`), without `inference:completed`. A stream + that errors, aborts or throws after finishing rounds counts them, and its failure entry in the + inference log carries them. A turn abandoned by an operator quiesce, a stream superseded by a + newer one, one whose agent was disposed, or one that ended with no final response and no other + end, which wrote no entry before, now writes a failure entry naming what ended it and carrying + its rounds, and `inferenceCount` counts it, when it had finished a round; one that hadn't has + nothing to carry and writes none, as before. If the store fails while a stream's end is + recorded, the error is reported on stderr and what comes after it still runs: the rest of a + tool-ended turn's end, a restart's next stream, and the stream's teardown. Exactly one + entry of each stream carries its usage, so totals summed from the inference log count each + stream once, as the session's do. A health snapshot's token totals read the log's newest 1,000 + entries, so with these new entries that window spans fewer turns. +- An inference-log entry for a stream the framework ended at a tool boundary (a tool-ended turn or + a restart) records a note in place of the compiled request, so these frequent entries don't each + add a whole-context blob to the store. They are still new entries: the inference log keeps no + retention, so it grows by one small entry for every tool-ended turn and restart, and by one + entry, with its request, for each abandoned, superseded or disposed stream that had finished a + round. diff --git a/src/framework.ts b/src/framework.ts index 0f36f28f9..f45248623 100644 --- a/src/framework.ts +++ b/src/framework.ts @@ -4,7 +4,7 @@ import { INLINE_WITHHELD_TEXT, REFERENCE_LIMITS, classifyBlock, isInlineContradi import { ReferenceFetcher, DEFAULT_FETCH_MAX_BYTES, EAGER_FETCH_TIMEOUT_MS } from './mcpl/reference-fetcher.js'; import { appendFileSync, existsSync, mkdirSync, readFileSync, renameSync, writeFileSync } from 'node:fs'; import { JsStore } from '@animalabs/chronicle'; -import type { Membrane, ContentBlock, NormalizedRequest, YieldingStream, ToolResult as MembraneToolResult, ToolResultContentBlock } from '@animalabs/membrane'; +import type { Membrane, ContentBlock, NormalizedRequest, NormalizedResponse, YieldingStream, DetailedUsage, ToolResult as MembraneToolResult, ToolResultContentBlock } from '@animalabs/membrane'; import { MembraneError } from '@animalabs/membrane'; import { ContextManager, PassthroughStrategy, WindowedPassthroughStrategy, OverBudgetError, UncoveredDropError, stripUnpairedToolBlocks } from '@animalabs/context-manager'; import type { CacheWireReceipt } from './kv-unified-wire.js'; @@ -669,6 +669,18 @@ const PROCESS_LOG_ID = 'framework/process-log'; const TURN_CHECKPOINTS_ID = 'framework/turn-checkpoints'; // legacy single-map layout, read-only fallback const TURN_CHECKPOINTS_TREE_ID = 'framework/turn-checkpoints/tree'; +/** Why the framework ended a stream at a tool boundary: an endTurn tool + * result, or a context-budget or physical-window restart (the + * `inference:stream_restarted` reasons). It is the stream's inference-log + * stop reason. */ +type StreamBoundaryEnd = 'turn_ended' | 'context_budget' | 'physical_window'; + +/** The request a boundary-ended stream's inference-log entry records in + * place of its body. A body is the whole compiled context, one blob per + * entry; these ends happen every tool-ended turn and every context-budget + * or physical-window restart. */ +const BOUNDARY_END_REQUEST = { note: 'request body not kept for a stream ended at a tool boundary' }; + /** Maximum number of turn checkpoints to keep per agent. */ const MAX_TURN_CHECKPOINTS = 20; const DEFAULT_DISCORD_AWARENESS_DEADLINE_MS = 10_000; @@ -1715,6 +1727,14 @@ export class AgentFramework { private processLoggingPersist: boolean; private processLoggingBroadcast: boolean; private activeStreams: Map> = new Map(); + // The framework ends a stream at a tool boundary from handleProcessEvent: + // an endTurn result, or a context-budget or physical-window restart. Each + // live stream registers here how to record that end (its usage, its log + // entry and, for endTurn, the completed turn), so the record is made where + // the decision is taken, before the cancellation's asynchronous teardown + // or an ephemeral's disposal. Keyed by the physical stream, so a + // successor's record is never touched. + private streamBoundaryEnds = new WeakMap void>(); /** A tool boundary injected fresh CONVERSATIONAL input (a real message — * not a reaction or a system marker) into the live stream. Tells @@ -1795,7 +1815,8 @@ export class AgentFramework { * The name is single-generation, so these frames still own their terminal * typing/outgoing close even after runEphemeralToCompletion deregisters them. */ private disposedEphemeralFrames: WeakSet = new WeakSet(); - /** Per-agent count of consecutive exhausted inferences (reset on any success). + /** Per-agent count of consecutive exhausted inferences, reset when a turn + * completes (inference:completed), not when a call answers mid-turn. * Drives hard-down escalation — see noteInferenceExhausted. */ private consecutiveInferenceFailures: Map = new Map(); /** N consecutive failed inferences ⇒ the agent is treated as hard-down and @@ -1819,7 +1840,8 @@ export class AgentFramework { * rewind loop — see refusalHandling + rewindTriggeringTurn. */ private refusalRewinds: Map = new Map(); /** Per-agent count of poison-history rewinds performed by the hard-down - * breaker (noteInferenceExhausted). Reset on any successful inference. + * breaker (noteInferenceExhausted). Reset when a turn completes + * (inference:completed). * Bounds the automatic quarantine loop the same way refusalRewinds bounds * the refusal loop, so the breaker can never shed the whole history. */ private exhaustionRewinds: Map = new Map(); @@ -6920,6 +6942,23 @@ export class AgentFramework { } } + /** + * A turn that ended WITHOUT a refusal (a natural completion that wasn't + * one, or a tool-ended turn) ends the rewind episode: the model responded. + * Leave the consolidated marker in place as the durable record; just + * clear the per-episode counters, and report a forced `/unstick` done. + */ + private endRefusalEpisode(agentName: string): void { + if (this.forcedRewind.has(agentName)) { + this.finishUnstick(agentName, true); + } + if (this.refusalRewinds.get(agentName)) { + this.refusalRewinds.set(agentName, 0); + } + this.rewindEpisode.delete(agentName); + this.refusalStreak.delete(agentName); + } + /** * Conclude an active `/unstick` session: post the outcome (what was shed and * whether the model stopped refusing) to the channel the command came from, @@ -9267,6 +9306,11 @@ export class AgentFramework { // driveStream's finally is the backstop.) this.eventGate?.onInferenceEnded(agent.name); this.settleAgent(agent.name, { stopReason: 'turn_ended', speech: '' }, this.turnRouting.get(agent.name)); + // A successful turn end, recorded as a natural completion's is: + // inference:completed with the rounds' usage, then turn_ended. + if (currentState.stream) { + this.streamBoundaryEnds.get(currentState.stream)?.('turn_ended'); + } this.emitTrace({ type: 'inference:turn_ended', agentName: agent.name }); } else if (overBudget || overPhysical) { // Context budget exceeded: break the stream, let compile() compress. @@ -9278,6 +9322,9 @@ export class AgentFramework { // endTurn above). if (currentState.stream) { this.frameworkCancelledStreams.set(`${agent.name}:${agent.streamId}`, 'budget_restart'); + // Its rounds were billed; the restart opens a fresh stream + // with its own usage. + this.streamBoundaryEnds.get(currentState.stream)?.(overBudget ? 'context_budget' : 'physical_window'); } agent.cancelStream(); this.emitTrace({ @@ -12357,6 +12404,116 @@ export class AgentFramework { const startTime = Date.now(); const requestId = `${agent.name}-${startTime}-${Math.random().toString(36).slice(2, 8)}`; const myStreamId = agent.streamId; + // Membrane's usage events are cumulative across the stream's tool loop, + // so the latest is everything the stream has been billed for so far: + // the measure of a stream that ends without a final response. + let latestUsage: DetailedUsage | undefined; + const tokenUsageOf = (usage: DetailedUsage | undefined): InferenceLogEntry['tokenUsage'] => usage + ? { + input: usage.inputTokens, + output: usage.outputTokens, + cacheCreation: usage.cacheCreationTokens, + cacheRead: usage.cacheReadTokens, + } + : undefined; + // The stream's billed usage reaches the session totals once, however the + // stream ends. A completion counts its final response's detailed usage, + // and the guard's abandoned round its detailed usage, else the latest + // sample, else the response's own. Every other end counts the rounds it + // finished: a final response without detailed usage, the boundary ends + // below, and `finally` for the rest. Each caller passes the one value it + // also reports. + let usageCounted = false; + let countedUsage: DetailedUsage | undefined; + const countUsage = (usage: DetailedUsage | undefined): void => { + if (usageCounted) return; + usageCounted = true; + countedUsage = usage; + if (!usage) return; + this.usageTracker.onInferenceCompleted(agent.name, usage, usage.estimatedCost + ? { total: usage.estimatedCost.total, currency: usage.estimatedCost.currency } + : undefined); + this.persistUsageState(); + }; + // ...and is carried by exactly one of its inference-log entries, so + // totals summed from the log count each stream once, as the session's + // do. The first entry that has usage to carry carries it (the value the + // entry names, else the rounds finished so far): an entry before it with + // none leaves it to the next, and no entry after it carries it again. An + // end that writes no entry of its own is logged in `finally`, when it has + // rounds to carry. + let usageLogged = false; + const logStream = (entry: Omit): void => { + const tokenUsage = usageLogged ? undefined : (entry.tokenUsage ?? tokenUsageOf(latestUsage)); + if (tokenUsage) usageLogged = true; + this.logInference({ + timestamp: startTime, + agentName: agent.name, + requestId, + ...entry, + tokenUsage, + }); + }; + // What ended a stream whose end writes no log entry of its own, for the + // entry `finally` writes for it. + let unloggedEnd: string | undefined; + // A turn that ended successfully: with its final response (natural + // completion), or without one, when a tool result with endTurn ended it. + const completeTurn = (usage: DetailedUsage | undefined, durationMs: number, response?: NormalizedResponse): void => { + // The trace, the session totals and the log entry all say the same. + const used = usage ?? latestUsage; + const tokenUsage = tokenUsageOf(used); + this.emitTrace({ + type: 'inference:completed', + agentName: agent.name, + durationMs, + tokenUsage, + }); + countUsage(used); + logStream(response + ? { + success: true, + request: compiledRequest ?? { note: 'streaming request' }, + response: response.raw ?? { note: 'streaming response' }, + durationMs, + tokenUsage, + stopReason: response.stopReason, + } + : { + success: true, + request: BOUNDARY_END_REQUEST, + response: { note: 'stream ended by tool result' }, + durationMs, + tokenUsage, + stopReason: 'turn_ended', + }); + }; + this.streamBoundaryEnds.set(stream, (end) => { + const durationMs = Date.now() - startTime; + const turnEnded = end === 'turn_ended'; + if (turnEnded) this.recordProviderAccelerationRecovery(agent, compiledRequest, end); + // The accounting writes to the store, and a write that fails must not + // skip what the boundary does after this call: the rest of a turn's + // end, or a restart's cancel and its request, and then the event's + // dispatch to modules. + try { + if (turnEnded) { + completeTurn(latestUsage, durationMs); + } else { + countUsage(latestUsage); + logStream({ + success: true, + request: BOUNDARY_END_REQUEST, + response: { note: 'stream restarted to compress its context' }, + durationMs, + stopReason: end, + }); + } + } catch (error) { + console.error(`[usage] agent=${agent.name} could not record a stream's usage at its end: ${describeFailureWithFrames(error)}`); + } + if (turnEnded) this.endRefusalEpisode(agent.name); + }); // Membrane usage events are cumulative across the native/XML tool loop. // Keep the previous cumulative sample so consumers that operate at the // physical provider-call boundary (estimator calibration and the @@ -12888,29 +13045,15 @@ export class AgentFramework { proseStream?.reset(); if (turnProseRouting === 'explicit') proseStream?.retarget(streamRouting.destination); if (withheld) { - const usage = response.details?.usage ?? response.usage; - const tokenUsage = usage ? { - input: usage.inputTokens, output: usage.outputTokens, - cacheCreation: usage.cacheCreationTokens, cacheRead: usage.cacheReadTokens, - } : undefined; - this.noteRefusal(agent.name, category, tokenUsage); // The abandoned stream was billed (membrane usage here is // cumulative across this physical tool loop): count it before - // the retry opens a fresh stream with its own usage. - const abandonedUsage = response.details?.usage; - if (abandonedUsage) { - this.usageTracker.onInferenceCompleted(agent.name, { - inputTokens: abandonedUsage.inputTokens, - outputTokens: abandonedUsage.outputTokens, - cacheCreationTokens: abandonedUsage.cacheCreationTokens, - cacheReadTokens: abandonedUsage.cacheReadTokens, - }, abandonedUsage.estimatedCost - ? { total: abandonedUsage.estimatedCost.total, currency: abandonedUsage.estimatedCost.currency } - : undefined); - this.persistUsageState(); - } - this.logInference({ - timestamp: startTime, agentName: agent.name, requestId, + // the retry opens a fresh stream with its own usage. One value + // for the count, the log entry and the refusal note. + const abandoned: DetailedUsage | undefined = response.details?.usage ?? latestUsage ?? response.usage; + const tokenUsage = tokenUsageOf(abandoned); + this.noteRefusal(agent.name, category, tokenUsage); + countUsage(abandoned); + logStream({ success: false, error: 'Tool output withheld by the guard', request: compiledRequest ?? {}, response, durationMs, tokenUsage, stopReason: 'refusal', }); @@ -13084,14 +13227,8 @@ export class AgentFramework { const admittedSegments = admitFinal ? finalSegments.map((seg) => this.admitProse(seg)) : []; const du = response.details?.usage; - const tokenUsage = du - ? { - input: du.inputTokens, - output: du.outputTokens, - cacheCreation: du.cacheCreationTokens, - cacheRead: du.cacheReadTokens, - } - : undefined; + // A refusal's note carries what completeTurn counts. + const tokenUsage = tokenUsageOf(du ?? latestUsage); // Async completion work may overlap a replacement stream. Recheck // physical ownership before any terminal state/settlement side effect. @@ -13119,35 +13256,7 @@ export class AgentFramework { .join('\n'), }, streamRouting); - this.emitTrace({ - type: 'inference:completed', - agentName: agent.name, - durationMs, - tokenUsage, - }); - - if (du) { - this.usageTracker.onInferenceCompleted(agent.name, { - inputTokens: du.inputTokens, - outputTokens: du.outputTokens, - cacheCreationTokens: du.cacheCreationTokens, - cacheReadTokens: du.cacheReadTokens, - }, du.estimatedCost ? { total: du.estimatedCost.total, currency: du.estimatedCost.currency } : undefined); - this.persistUsageState(); - } - - // Log inference - this.logInference({ - timestamp: startTime, - agentName: agent.name, - requestId, - success: true, - request: compiledRequest ?? { note: 'streaming request' }, - response: response.raw ?? { note: 'streaming response' }, - durationMs, - tokenUsage, - stopReason: response.stopReason, - }); + completeTurn(du, durationMs, response); // Surface refusals instead of going silently mute: stderr line // (headless inference failures are otherwise under-logged) + an @@ -13255,17 +13364,7 @@ export class AgentFramework { void this.reactToRefusal(agent.name, category); } } else { - // A turn that completed WITHOUT a refusal ends the rewind episode: - // the model responded. Leave the consolidated marker in place as - // the durable record; just clear the per-episode counters. - if (this.forcedRewind.has(agent.name)) { - this.finishUnstick(agent.name, true); - } - if (this.refusalRewinds.get(agent.name)) { - this.refusalRewinds.set(agent.name, 0); - } - this.rewindEpisode.delete(agent.name); - this.refusalStreak.delete(agent.name); + this.endRefusalEpisode(agent.name); } // Dispatch speech (and thoughts if any) @@ -13474,10 +13573,7 @@ export class AgentFramework { stack: err.stack, }); - this.logInference({ - timestamp: startTime, - agentName: agent.name, - requestId, + logStream({ success: false, error: err.message, request: compiledRequest ?? { note: 'streaming request failed' }, @@ -13593,6 +13689,7 @@ export class AgentFramework { this.eventGate?.onInferenceEnded(agent.name); } if (cancelKind === 'quiesce_abandoned') { + unloggedEnd = 'Turn abandoned by operator quiesce'; // Operator-initiated cancel during a quiesce drain. Settle // the turn honestly — deliberately NO inference:exhausted: // that trace centrally drives the consecutive-failure @@ -13662,10 +13759,7 @@ export class AgentFramework { // state. Without this, abort-terminated inferences are // invisible to forensic queries (only request-side telemetry // via llm-calls.jsonl shows them, and only by absence). - this.logInference({ - timestamp: startTime, - agentName: agent.name, - requestId, + logStream({ success: false, error: `Stream aborted: ${reason}`, request: compiledRequest ?? { note: 'streaming request aborted' }, @@ -13677,6 +13771,8 @@ export class AgentFramework { } case 'usage': { + // Cumulative across the stream's rounds: replace, never sum. + latestUsage = event.usage; // The stream's accumulated fresh input, which is what the // maxStreamTokens budget checks. agent.lastStreamInputTokens = event.usage.inputTokens; @@ -13800,7 +13896,7 @@ export class AgentFramework { const durationMs = Date.now() - startTime; if (ownsProviderGate && this.holdProviderAcceleration(agent, err, trigger)) { this.emitTrace({ type: 'inference:failed', agentName: agent.name, error: err.message, stack: err.stack }); - this.logInference({ timestamp: startTime, agentName: agent.name, requestId, success: false, + logStream({ success: false, error: `Provider acceleration cooldown: ${err.message}`, request: compiledRequest ?? { note: 'streaming request rate-limited' }, durationMs }); this.abortAgentScript(agent.name, 'provider acceleration cooldown'); @@ -13814,7 +13910,7 @@ export class AgentFramework { const statedWait = ownsProviderGate ? undefined : this.nonOwningProviderWait(agent, err, trigger); if (statedWait === 'held') { this.emitTrace({ type: 'inference:failed', agentName: agent.name, error: err.message, stack: err.stack }); - this.logInference({ timestamp: startTime, agentName: agent.name, requestId, success: false, + logStream({ success: false, error: `Provider wait: ${err.message}`, request: compiledRequest ?? { note: 'streaming request rate-limited' }, durationMs }); this.abortAgentScript(agent.name, 'provider wait'); @@ -13841,10 +13937,7 @@ export class AgentFramework { // Postmortem 2026-05-28 P2 #7: catch-path failures (stream itself // threw) were previously only visible via in-memory trace listeners. // Persist to inference log for forensic attribution. - this.logInference({ - timestamp: startTime, - agentName: agent.name, - requestId, + logStream({ success: false, error: `Stream threw: ${err.message}`, request: compiledRequest ?? { note: 'streaming request threw' }, @@ -13855,6 +13948,35 @@ export class AgentFramework { agent.reset(); if (this.agents.get(agent.name) === agent && agent.streamId === myStreamId) this.eventGate?.onInferenceEnded(agent.name); } finally { + // The accounting writes to the store, and a write that fails must not + // skip the teardown after it. + try { + // Every end the cases above didn't count (a stream error or abort, a + // throw, a superseded generation) still billed the rounds it finished. + countUsage(latestUsage); + // An end that wrote no entry carrying them (a quiesce-abandoned turn, + // a stream superseded by a newer one, a disposed agent's stream) writes + // one now, as a failure naming what ended it, so the log carries what + // the totals count. With no finished round there is nothing to carry, + // and no entry is written, as before. + if (countedUsage && !usageLogged) { + const replaced = this.agents.get(agent.name) !== agent; + const superseded = generationLost || agent.streamId !== myStreamId; + logStream({ + success: false, + error: unloggedEnd ?? (replaced + ? 'Stream abandoned: its agent was disposed or replaced' + : superseded + ? 'Stream superseded: a newer stream of the agent took over' + : 'Stream ended without a final response'), + request: compiledRequest ?? { note: 'streaming request' }, + durationMs: Date.now() - startTime, + tokenUsage: tokenUsageOf(countedUsage), + }); + } + } catch (error) { + console.error(`[usage] agent=${agent.name} could not record a stream's usage at its end: ${describeFailureWithFrames(error)}`); + } const unsettledKvSubmissions = drainKvSubmissionIds?.() ?? []; if (unsettledKvSubmissions.length > 0) { try { @@ -15269,6 +15391,10 @@ export class AgentFramework { return clone; } + /** Writes one inference-log entry. A stream's entries go through its + * driveStream's `logStream`, which gives the stream's usage to one entry + * only; an end that calls this directly for a stream writes a second + * entry for it, and its usage can be counted twice in the log. */ private logInference(entry: InferenceLogEntry): void { // Store large request/response as blobs const entryToStore = { ...entry }; @@ -15885,9 +16011,14 @@ export class AgentFramework { }, ); } else if (event.type === 'inference:completed') { - // A successful response — even mid-turn between tool calls — proves the - // agent isn't hard-down; clear its consecutive-failure streak (and the - // poison-history breaker's rewind budget). + // A completed turn, a tool-ended one included, shows the agent can + // complete turns: clear its consecutive-failure streak and the + // poison-history breaker's rewind budget. A provider call answered + // mid-turn (inference:usage) doesn't. The streak exists to catch an + // agent that can't complete a turn, and an agent whose every turn has a + // first call that answers and a later one that fails is exactly that: + // a reset per answered call would keep it below the hard-down + // threshold, so neither the alert nor the breaker would come. const name = event.agentName as string | undefined; if (name && this.consecutiveInferenceFailures.get(name)) { this.consecutiveInferenceFailures.set(name, 0); @@ -16403,7 +16534,7 @@ export class AgentFramework { * push event to wake the agent onto the same context. * * Bounded: at most `refusalHandling.maxRewinds` (default 3, hard cap 10) - * sheds per failure episode; the budget resets on any successful inference. + * sheds per failure episode; the budget resets when a turn completes. * At the cap (or with nothing left to shed) it stops — the agent stays up * and hard-down logging continues, but no further history is consumed. */ diff --git a/src/types/framework.ts b/src/types/framework.ts index 5725aee5d..003ed8d92 100644 --- a/src/types/framework.ts +++ b/src/types/framework.ts @@ -413,18 +413,30 @@ export interface InferenceLogEntry { success: boolean; /** Error message if failed */ error?: string; - /** Request data or blob ID if large */ + /** Request data or blob ID if large. A stream the framework ended at a + * tool boundary (stopReason 'turn_ended', 'context_budget' or + * 'physical_window') records a note in place of the body. */ request: unknown | { blobId: string }; /** Response data or blob ID if large (only if successful) */ response?: unknown | { blobId: string }; durationMs: number; + /** The stream's usage: its final response's, or, for a stream that ended + * without one, the provider rounds it finished before it ended. Exactly + * one entry of a stream carries it, the one its end wrote; a stream whose + * end writes none (a quiesce-abandoned turn, a superseded stream, a + * disposed agent's stream) gets a failure entry that carries it, when it + * had finished a round. One that hadn't has nothing to carry and gets no + * entry. */ tokenUsage?: { input: number; output: number; cacheCreation?: number; cacheRead?: number; }; - /** Stop reason from the model */ + /** Stop reason from the model, or why the framework ended the stream at a + * tool boundary: 'turn_ended' (a tool result with endTurn), or + * 'context_budget' / 'physical_window' (a context-budget or + * physical-window restart). */ stopReason?: string; } diff --git a/src/types/module.ts b/src/types/module.ts index ca9cb3816..b8f6794e4 100644 --- a/src/types/module.ts +++ b/src/types/module.ts @@ -327,7 +327,8 @@ export interface ModuleContext { * Modules use this to react to agent lifecycle without polling — e.g. an * adapter that wants to show a "typing" indicator can bracket it between * `inference:started` and the terminal events: `inference:completed` (the - * normal end), `inference:turn_ended` (a tool result ended the turn), + * model answered and the turn ended, with a final response or with a tool + * result that ended it, which `inference:turn_ended` then follows), * `inference:aborted`, `inference:failed` and `inference:exhausted`, all of * which carry `agentName`. A turn's tool-call cycles stay inside its * stream, so one bracket spans them. A context-budget restart emits diff --git a/src/types/trace.ts b/src/types/trace.ts index 343fa10cb..e678a46c9 100644 --- a/src/types/trace.ts +++ b/src/types/trace.ts @@ -55,6 +55,12 @@ export type TraceEvent = silent?: boolean; }) | (TraceEventBase & { + /** + * The model answered and the turn ended: with the stream's final + * response, or with a tool result carrying `endTurn`, which + * `inference:turn_ended` then follows. `tokenUsage` is the stream's + * usage, cumulative across its rounds. + */ type: 'inference:completed'; agentName: string; durationMs: number; diff --git a/test/end-turn-usage.test.ts b/test/end-turn-usage.test.ts new file mode 100644 index 000000000..9efbb5c9e --- /dev/null +++ b/test/end-turn-usage.test.ts @@ -0,0 +1,1050 @@ +import { describe, it } from 'node:test'; +import assert from 'node:assert/strict'; +import { mkdtempSync, rmSync } from 'node:fs'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; +import { MembraneError } from '@animalabs/membrane'; +import type { DetailedUsage, NormalizedRequest, NormalizedResponse, StreamEvent, YieldingStream } from '@animalabs/membrane'; +import { AgentFramework } from '../src/index.js'; +import { HealthModule } from '../src/modules/health/index.js'; +import type { EventResponse, Module, ProcessEvent, ToolCall, TraceEvent } from '../src/index.js'; +import { createMockResponse, MockMembrane, MockYieldingStream } from './helpers/mock-membrane.js'; + +class TurnTools implements Module { + readonly name = 'test'; + async start() {} + async stop() {} + async onProcess() { return {}; } + getTools() { + return ['echo', 'finish'].map((name) => ({ + name, + description: name, + inputSchema: { type: 'object' as const, properties: {} }, + })); + } + async handleToolCall(call: ToolCall) { + return { success: true, endTurn: call.name.endsWith('finish'), data: {} }; + } +} + +const firstUsage: DetailedUsage = { + inputTokens: 100, outputTokens: 20, cacheCreationTokens: 30, cacheReadTokens: 40, + estimatedCost: { input: 0.01, output: 0.02, total: 0.03, currency: 'USD' }, +}; +const totalUsage: DetailedUsage = { + inputTokens: 250, outputTokens: 55, cacheCreationTokens: 45, cacheReadTokens: 90, + estimatedCost: { input: 0.03, output: 0.05, total: 0.08, currency: 'USD' }, +}; + +function toolResponse(name: string, id: string, usage?: DetailedUsage): NormalizedResponse { + const response = createMockResponse([{ type: 'tool_use', name: 'test--' + name, id, input: {} }], 'tool_use'); + // Membrane emits cumulative usage across all rounds of a yielding stream. + response.usage = usage!; + return response; +} + +function expectedTokens(usage: DetailedUsage) { + return { + input: usage.inputTokens, output: usage.outputTokens, + cacheCreation: usage.cacheCreationTokens, cacheRead: usage.cacheReadTokens, + }; +} + +function expectedTotals(usage: DetailedUsage) { + return { + inputTokens: usage.inputTokens, outputTokens: usage.outputTokens, + cacheCreationTokens: usage.cacheCreationTokens ?? 0, + cacheReadTokens: usage.cacheReadTokens ?? 0, + estimatedCost: usage.estimatedCost + ? { total: usage.estimatedCost.total, currency: usage.estimatedCost.currency } + : undefined, + }; +} + +async function fixture() { + const dir = mkdtempSync(join(tmpdir(), 'af-end-turn-usage-')); + const storePath = join(dir, 'store.chronicle'); + const membrane = new MockMembrane(); + const framework = await AgentFramework.create({ + storePath, membrane: membrane.asMembrane(), agents: [], + modules: [new TurnTools()], syncIntervalMs: 0, + }); + const traces: TraceEvent[] = []; + const completionStates: Array = []; + // What the session totals already hold when the turn's terminal traces fire. + const countedAt: Array<{ type: string; inferenceCount: number }> = []; + framework.onTrace((event) => { + traces.push(event); + if (event.type === 'inference:completed') { + completionStates.push(framework.getAgent(event.agentName)?.state.status); + } + if (event.type === 'inference:turn_ended' || event.type === 'inference:stream_restarted') { + countedAt.push({ type: event.type, inferenceCount: framework.getSessionUsage().inferenceCount }); + } + }); + framework.start(); + return { dir, storePath, membrane, framework, traces, completionStates, countedAt }; +} + +/** Every stream the framework has driven so far, by the promise driveStream + * returns, so a test can wait until each one's teardown has run. */ +const driven = new WeakMap[]>(); +function streamsDriven(framework: AgentFramework): () => Promise { + let handles = driven.get(framework); + if (!handles) { + const all: Promise[] = []; + handles = all; + driven.set(framework, all); + const streams = (framework as unknown as { activeStreams: Map> }).activeStreams; + const set = streams.set.bind(streams); + streams.set = (name, handle) => { all.push(handle); return set(name, handle); }; + } + const watched = handles; + return async () => { await Promise.allSettled(watched); }; +} + +async function run(framework: AgentFramework, name: string) { + const tornDown = streamsDriven(framework); + const { agent, contextManager } = await framework.createEphemeralAgent({ + name, model: 'test-model', systemPrompt: 'Do the task.', allowedTools: 'all', + proseRouting: 'disabled', + }); + contextManager.addMessage('user', [{ type: 'text', text: 'Run once.' }]); + const result = await framework.runEphemeralToCompletion(agent, contextManager); + // A cancelled stream's teardown runs after the run settles: wait for it + // before checking counts. + await tornDown(); + return result; +} + +describe('tool-ended turn usage accounting', () => { + for (const rounds of [1, 2]) { + it('accounts once for ' + rounds + ' tool round(s), with cumulative cache usage and cost', async () => { + const f = await fixture(); + try { + if (rounds === 2) f.membrane.pushResponse(toolResponse('echo', 'echo-1', firstUsage)); + f.membrane.pushResponse(toolResponse('finish', 'finish-1', totalUsage)); + const result = await run(f.framework, 'worker'); + assert.equal(result.toolCallsCount, rounds); + assert.equal(f.membrane.calls.length, 1, 'endTurn must not request another inference'); + assert.equal(f.membrane.lastStream!.receivedToolResults.length, rounds - 1); + assert.deepEqual(f.completionStates, ['idle']); + const completed = f.traces.filter((e) => e.type === 'inference:completed'); + assert.equal(completed.length, 1); + assert.equal(completed[0].agentName, 'worker'); + assert.deepEqual(completed[0].tokenUsage, expectedTokens(totalUsage)); + assert.ok(completed[0].durationMs >= 0); + const logs = f.framework.queryInferenceLogs({ agentName: 'worker' }); + assert.equal(logs.total, 1); + assert.equal(logs.entries[0].entry.success, true); + assert.equal(logs.entries[0].entry.stopReason, 'turn_ended'); + assert.deepEqual(logs.entries[0].entry.tokenUsage, expectedTokens(totalUsage)); + assert.deepEqual(logs.entries[0].entry.response, { note: 'stream ended by tool result' }); + assert.deepEqual(logs.entries[0].entry.request, { note: 'request body not kept for a stream ended at a tool boundary' }, + 'no compiled-context blob per tool-ended turn'); + assert.deepEqual(f.countedAt, [{ type: 'inference:turn_ended', inferenceCount: 1 }], + 'counted before turn_ended fires, as a completion is before its caller resumes'); + assert.equal(logs.entries[0].entry.durationMs, completed[0].durationMs); + assert.equal(f.traces.filter((e) => e.type === 'inference:turn_ended').length, 1); + assert.equal(f.traces.filter((e) => e.type === 'usage:updated').length, 1); + assert.equal(f.traces.filter((e) => e.type === 'inference:exhausted').length, 0); + const snapshot = f.framework.getSessionUsage(); + assert.equal(snapshot.inferenceCount, 1); + assert.deepEqual(snapshot.totals, expectedTotals(totalUsage)); + assert.deepEqual(snapshot.byAgent, [{ + agentName: 'worker', usage: expectedTotals(totalUsage), inferenceCount: 1, + }]); + } finally { + await f.framework.stop(); + rmSync(f.dir, { recursive: true, force: true }); + } + }); + } + + it('keeps natural completion authoritative and counts neither cumulative samples nor final usage twice', async () => { + const f = await fixture(); + try { + f.membrane.pushResponse(toolResponse('echo', 'echo-1', firstUsage)); + const response = createMockResponse([{ type: 'text', text: 'Done.' }]); + response.usage = { ...totalUsage, outputTokens: 50 }; // final details remain authoritative + response.details!.usage = totalUsage; + f.membrane.pushResponse(response); + await run(f.framework, 'ordinary'); + assert.deepEqual(f.completionStates, ['idle']); + const completed = f.traces.filter((e) => e.type === 'inference:completed'); + assert.equal(completed.length, 1); + assert.deepEqual(completed[0].tokenUsage, expectedTokens(totalUsage)); + assert.equal(f.traces.filter((e) => e.type === 'usage:updated').length, 1); + assert.equal(f.traces.filter((e) => e.type === 'inference:turn_ended').length, 0); + assert.equal(f.framework.getSessionUsage().inferenceCount, 1); + assert.deepEqual(f.framework.getSessionUsage().totals, expectedTotals(totalUsage)); + const logs = f.framework.queryInferenceLogs({ agentName: 'ordinary' }); + assert.equal(logs.total, 1, 'natural completion logs exactly once'); + assert.equal(logs.entries[0].entry.stopReason, 'end_turn'); + assert.deepEqual(logs.entries[0].entry.tokenUsage, expectedTokens(totalUsage)); + } finally { + await f.framework.stop(); + rmSync(f.dir, { recursive: true, force: true }); + } + }); + + it('counts the rounds\' samples when a final response carries no usage', async () => { + const f = await fixture(); + try { + f.membrane.pushResponse(toolResponse('echo', 'echo-1', firstUsage)); + const response = createMockResponse([{ type: 'text', text: 'Done.' }]); + response.usage = totalUsage; // the last round's cumulative sample; no details.usage + f.membrane.pushResponse(response); + await run(f.framework, 'ordinary'); + const completed = f.traces.filter((e) => e.type === 'inference:completed'); + assert.equal(completed.length, 1); + assert.deepEqual(completed[0].tokenUsage, expectedTokens(totalUsage)); + assert.equal(f.traces.filter((e) => e.type === 'usage:updated').length, 1); + assert.deepEqual(f.framework.getSessionUsage().totals, expectedTotals(totalUsage)); + const logs = f.framework.queryInferenceLogs({ agentName: 'ordinary' }); + assert.deepEqual(logs.entries[0].entry.tokenUsage, expectedTokens(totalUsage), 'the trace, the totals and the log agree'); + } finally { + await f.framework.stop(); + rmSync(f.dir, { recursive: true, force: true }); + } + }); + + it('attributes successive streams separately and persists their totals through store reopen', async () => { + const f = await fixture(); + let stopped = false; + try { + f.membrane.pushResponse(toolResponse('finish', 'finish-a', firstUsage)); + await run(f.framework, 'alice'); + f.membrane.pushResponse(toolResponse('finish', 'finish-b', totalUsage)); + await run(f.framework, 'bob'); + const snapshot = f.framework.getSessionUsage(); + assert.equal(snapshot.inferenceCount, 2); + assert.deepEqual(snapshot.byAgent, [ + { agentName: 'alice', usage: expectedTotals(firstUsage), inferenceCount: 1 }, + { agentName: 'bob', usage: expectedTotals(totalUsage), inferenceCount: 1 }, + ]); + assert.deepEqual(f.traces.filter((e) => e.type === 'usage:updated').map((e) => e.agentName), ['alice', 'bob']); + const logs = f.framework.queryInferenceLogs(); + assert.equal(logs.total, 2); + assert.equal(new Set(logs.entries.map((e) => e.entry.requestId)).size, 2); + assert.deepEqual(logs.entries.map((e) => e.entry.agentName).sort(), ['alice', 'bob']); + await f.framework.stop(); + stopped = true; + const reopened = await AgentFramework.create({ + storePath: f.storePath, membrane: f.membrane.asMembrane(), agents: [], modules: [], syncIntervalMs: 0, + }); + try { + assert.deepEqual(reopened.getSessionUsage(), snapshot); + assert.deepEqual(reopened.queryInferenceLogs(), logs); + } finally { + await reopened.stop(); + } + } finally { + if (!stopped) await f.framework.stop(); + rmSync(f.dir, { recursive: true, force: true }); + } + }); + + it('includes tool-ended turns in the log-backed health snapshot', async () => { + const f = await fixture(); + try { + f.membrane.pushResponse(toolResponse('echo', 'echo-1', firstUsage)); + f.membrane.pushResponse(toolResponse('finish', 'finish-1', totalUsage)); + await run(f.framework, 'worker'); + const health = new HealthModule({ timeZone: 'UTC' }); + health.bind(f.framework, f.framework.getStore()); + const result = await health.handleToolCall({ + id: 'health-snapshot', name: 'snapshot', input: { includeSubagents: false }, + }); + assert.equal(result.success, true); + const data = result.data as { + inferences: { successCount: number; errorCount: number }; + tokenTotalsByAgent: Record; + }; + assert.equal(data.inferences.successCount, 1); + assert.equal(data.inferences.errorCount, 0); + assert.deepEqual(data.tokenTotalsByAgent.worker, { + input: totalUsage.inputTokens, output: totalUsage.outputTokens, + cacheRead: totalUsage.cacheReadTokens, inferences: 1, + }); + } finally { + await f.framework.stop(); + rmSync(f.dir, { recursive: true, force: true }); + } + }); + + it('reports completion without inventing usage when no sample is available', async () => { + const f = await fixture(); + try { + f.membrane.pushResponse(toolResponse('finish', 'finish-1')); + await run(f.framework, 'unknown-usage'); + const completed = f.traces.filter((e) => e.type === 'inference:completed'); + assert.equal(completed.length, 1); + assert.equal(completed[0].tokenUsage, undefined); + assert.equal(f.traces.filter((e) => e.type === 'usage:updated').length, 0); + assert.equal(f.framework.getSessionUsage().inferenceCount, 0); + const logs = f.framework.queryInferenceLogs({ agentName: 'unknown-usage' }); + assert.equal(logs.total, 1); + assert.equal(logs.entries[0].entry.success, true); + assert.equal(logs.entries[0].entry.tokenUsage, undefined); + } finally { + await f.framework.stop(); + rmSync(f.dir, { recursive: true, force: true }); + } + }); +}); + +/** One queued response per stream: a restart opens a fresh stream. */ +class SequentialStreamMembrane extends MockMembrane { + private next = 0; + override streamYielding(request: NormalizedRequest): YieldingStream { + this.calls.push(request); + const stream = new MockYieldingStream(this.responses.slice(this.next, ++this.next)); + this.lastStream = stream; + return stream; + } +} + +/** A stream that finishes one tool round with `usage`, then ends as `ending` + * says: with that event, by throwing it, or, for 'close', by just closing. */ +class RoundThenEnd implements YieldingStream { + isWaitingForTools = false; + pendingToolCallIds: string[] = []; + toolDepth = 0; + isCancelled = false; + private resumed = false; + private resume: (() => void) | null = null; + constructor(private readonly usage: DetailedUsage, private readonly ending: StreamEvent | Error | 'close') {} + provideToolResults(): void { + this.isWaitingForTools = false; + this.resumed = true; + this.resume?.(); + } + cancel(): void { + this.isCancelled = true; + this.resumed = true; + this.resume?.(); + } + async *[Symbol.asyncIterator](): AsyncIterator { + yield { type: 'usage', usage: this.usage } as StreamEvent; + const call = { type: 'tool_use' as const, id: 'echo-1', name: 'test--echo', input: {} }; + this.isWaitingForTools = true; + this.pendingToolCallIds = [call.id]; + yield { + type: 'tool-calls', + calls: [{ id: call.id, name: call.name, input: {} }], + context: { rawText: '', preamble: '', depth: 0, previousResults: [], accumulated: '', roundContent: [call] }, + } as StreamEvent; + if (!this.resumed) await new Promise((r) => { this.resume = r; }); + if (this.isCancelled) { + yield { type: 'aborted', reason: 'user' } as StreamEvent; + return; + } + if (this.ending === 'close') return; + if (this.ending instanceof Error) throw this.ending; + yield this.ending; + } +} + +class StreamsMembrane extends MockMembrane { + constructor(private readonly streams: Array<() => YieldingStream>) { super(); } + override streamYielding(request: NormalizedRequest): YieldingStream { + this.calls.push(request); + const make = this.streams.shift(); + return make ? make() : new MockYieldingStream([createMockResponse([{ type: 'text', text: 'done' }])]); + } +} + +async function pollUntil(cond: () => boolean, timeoutMs = 5_000): Promise { + const deadline = Date.now() + timeoutMs; + while (!cond() && Date.now() < deadline) await new Promise((r) => setTimeout(r, 10)); + return cond(); +} + +function finalResponse(text: string, usage: DetailedUsage): NormalizedResponse { + const response = createMockResponse([{ type: 'text', text }]); + response.usage = usage; + response.details!.usage = usage; + return response; +} + +const secondUsage: DetailedUsage = { + inputTokens: 300, outputTokens: 7, cacheCreationTokens: 0, cacheReadTokens: 200, + estimatedCost: { input: 0.02, output: 0.01, total: 0.03, currency: 'USD' }, +}; + +function summed(a: DetailedUsage, b: DetailedUsage) { + return { + inputTokens: a.inputTokens + b.inputTokens, outputTokens: a.outputTokens + b.outputTokens, + cacheCreationTokens: (a.cacheCreationTokens ?? 0) + (b.cacheCreationTokens ?? 0), + cacheReadTokens: (a.cacheReadTokens ?? 0) + (b.cacheReadTokens ?? 0), + estimatedCost: { total: a.estimatedCost!.total + b.estimatedCost!.total, currency: 'USD' }, + }; +} + +describe('a stream counts its usage once, however it ends', () => { + const nearCap: DetailedUsage = { + inputTokens: 5_000, outputTokens: 5, cacheCreationTokens: 0, cacheReadTokens: 195_000, + estimatedCost: { input: 0.05, output: 0.01, total: 0.06, currency: 'USD' }, + }; + for (const [reason, limits, usage] of [ + ['context_budget', { maxStreamTokens: 50 }, firstUsage], + ['physical_window', { physicalWindowTokens: 200_000, maxTokens: 4_000 }, nearCap], + ] as const) { + it(`a ${reason} restart counts and logs the stream it cancels, without completing the turn`, async () => { + const dir = mkdtempSync(join(tmpdir(), 'af-end-turn-usage-')); + const membrane = new SequentialStreamMembrane(); + membrane.pushResponse(toolResponse('echo', 'echo-1', usage)); + membrane.pushResponse(finalResponse('Recovered.', secondUsage)); + const framework = await AgentFramework.create({ + storePath: join(dir, 'store.chronicle'), membrane: membrane.asMembrane(), agents: [], + modules: [new TurnTools()], syncIntervalMs: 0, + }); + const traces: TraceEvent[] = []; + const countedAtRestart: number[] = []; + framework.onTrace((event) => { + traces.push(event); + if (event.type === 'inference:stream_restarted') countedAtRestart.push(framework.getSessionUsage().inferenceCount); + }); + framework.start(); + const tornDown = streamsDriven(framework); + try { + const { agent, contextManager } = await framework.createEphemeralAgent({ + name: 'worker', model: 'test-model', systemPrompt: 'Do the task.', allowedTools: 'all', + proseRouting: 'disabled', ...limits, + }); + contextManager.addMessage('user', [{ type: 'text', text: 'Run once.' }]); + await framework.runEphemeralToCompletion(agent, contextManager); + // The cancelled stream's teardown has run, so whatever it would log is in. + await tornDown(); + + assert.equal(membrane.calls.length, 2, 'the restart opened a second stream'); + assert.equal((traces.find((e) => e.type === 'inference:stream_restarted') as { reason?: string })?.reason, reason); + assert.deepEqual(countedAtRestart, [1], 'the cancelled stream is counted where the restart is decided'); + assert.equal(traces.filter((e) => e.type === 'inference:completed').length, 1, 'only the turn\'s end completes'); + assert.equal(traces.filter((e) => e.type === 'usage:updated').length, 2); + const snapshot = framework.getSessionUsage(); + assert.equal(snapshot.inferenceCount, 2); + assert.deepEqual(snapshot.totals, summed(usage, secondUsage)); + const logs = framework.queryInferenceLogs({ agentName: 'worker' }).entries.map((e) => e.entry).reverse(); + assert.equal(logs.length, 2); + assert.deepEqual( + { success: logs[0].success, stopReason: logs[0].stopReason, tokenUsage: logs[0].tokenUsage, request: logs[0].request, response: logs[0].response }, + { + success: true, stopReason: reason, tokenUsage: expectedTokens(usage), + request: { note: 'request body not kept for a stream ended at a tool boundary' }, + response: { note: 'stream restarted to compress its context' }, + }, + ); + assert.equal(logs[1].stopReason, 'end_turn'); + assert.deepEqual(logs[1].tokenUsage, expectedTokens(secondUsage)); + } finally { + await framework.stop(); + rmSync(dir, { recursive: true, force: true }); + } + }); + } + + for (const [name, ending] of [ + ['a stream error', { type: 'error', error: new MembraneError({ type: 'auth', retryable: false, message: 'key revoked', rawError: { status: 401 } }) } as StreamEvent], + ['an abort', { type: 'aborted', reason: 'timeout' } as StreamEvent], + ['a thrown stream', new Error('socket hang up')], + ] as const) { + it(`${name} after a finished round counts that round once and logs it with the failure`, async () => { + const dir = mkdtempSync(join(tmpdir(), 'af-end-turn-usage-')); + const membrane = new StreamsMembrane([() => new RoundThenEnd(firstUsage, ending)]); + const framework = await AgentFramework.create({ + storePath: join(dir, 'store.chronicle'), membrane: membrane.asMembrane(), agents: [], + modules: [new TurnTools()], syncIntervalMs: 0, + }); + const traces: TraceEvent[] = []; + framework.onTrace((event) => traces.push(event)); + framework.start(); + try { + const { agent, contextManager } = await framework.createEphemeralAgent({ + name: 'worker', model: 'test-model', systemPrompt: 'Do the task.', allowedTools: 'all', proseRouting: 'disabled', + }); + contextManager.addMessage('user', [{ type: 'text', text: 'Run once.' }]); + // The run settles once the failure path has decided: it rejects with + // the stream's failure, or, had a retry been scheduled, it would wait + // for it and resolve on the default response. + let bound: ReturnType | undefined; + const outcome = await Promise.race([ + framework.runEphemeralToCompletion(agent, contextManager).then(() => 'resolved', (error: unknown) => error), + new Promise((resolve) => { bound = setTimeout(() => resolve('timed out'), 10_000); }), + ]).finally(() => clearTimeout(bound)); + assert.ok(outcome instanceof Error, `the run failed with the stream: ${String(outcome)}`); + assert.equal(membrane.calls.length, 1, 'no retry was scheduled by the time the run failed'); + assert.ok(await pollUntil(() => traces.some((e) => e.type === 'usage:updated')), 'the finished round was counted'); + + assert.equal(traces.filter((e) => e.type === 'inference:completed').length, 0); + assert.equal(traces.filter((e) => e.type === 'usage:updated').length, 1); + assert.deepEqual(framework.getSessionUsage().totals, expectedTotals(firstUsage)); + const logs = framework.queryInferenceLogs({ agentName: 'worker' }).entries.map((e) => e.entry); + assert.equal(logs.length, 1); + assert.equal(logs[0].success, false); + assert.deepEqual(logs[0].tokenUsage, expectedTokens(firstUsage), 'the failure entry carries the finished round'); + } finally { + await framework.stop(); + rmSync(dir, { recursive: true, force: true }); + } + }); + } +}); + +class InputModule implements Module { + readonly name = 'input'; + async start() {} + async stop() {} + getTools() { return []; } + async handleToolCall() { return { success: false, isError: true, error: 'none' }; } + async onProcess(event: ProcessEvent): Promise { + if (event.type !== 'external-message') return {}; + return { addMessages: [{ participant: 'User', content: [{ type: 'text', text: String(event.content) }] }], requestInference: true }; + } +} + +function accelerationLimit(request: unknown): MembraneError { + return new MembraneError({ + type: 'rate_limit', retryable: true, httpStatus: 429, + message: "This request would exceed your organization's maximum usage increase rate for input tokens per minute", + rawError: { status: 429 }, rawRequest: request, + }); +} + +class ThrowingStream implements YieldingStream { + isWaitingForTools = false; + pendingToolCallIds: string[] = []; + toolDepth = 0; + isCancelled = false; + constructor(private readonly error: Error) {} + provideToolResults(): void { throw new Error('not waiting'); } + cancel(): void { this.isCancelled = true; } + async *[Symbol.asyncIterator](): AsyncIterator { + throw this.error; + } +} + +/** A resident whose streams come from `streams`, one per inference, with a 1 s provider cooldown. */ +async function resident(streams: Array<(request: NormalizedRequest) => YieldingStream>) { + const dir = mkdtempSync(join(tmpdir(), 'af-end-turn-usage-')); + const membrane = new MockMembrane(); + membrane.streamYielding = (request: NormalizedRequest): YieldingStream => { + membrane.calls.push(request); + const make = streams.shift(); + return make ? make(request) : new MockYieldingStream([createMockResponse([{ type: 'text', text: 'done' }])]); + }; + const framework = await AgentFramework.create({ + storePath: join(dir, 'store.chronicle'), membrane: membrane.asMembrane(), + agents: [{ name: 'resident', model: 'test-model', systemPrompt: 'system', allowedTools: 'all' }], + modules: [new InputModule(), new TurnTools()], syncIntervalMs: 0, maintenanceIntervalMs: 0, + }); + const internal = framework as unknown as { + providerAccelerationDefaultCooldownMs: number; + providerAccelerationJitterMs: number; + providerAccelerationLastRecovery: Map; + }; + internal.providerAccelerationDefaultCooldownMs = 1_000; + internal.providerAccelerationJitterMs = 0; + const traces: TraceEvent[] = []; + framework.onTrace((event) => traces.push(event)); + return { dir, membrane, framework, internal, traces }; +} + +describe('a tool-ended turn ends what a successful turn ends', () => { + for (const [ending, response] of [ + ['a tool-ended turn', () => toolResponse('finish', 'finish-1', totalUsage)], + ['a natural completion', () => createMockResponse([{ type: 'text', text: 'Done.' }])], + ] as const) { + it(`the failure streak and the refusal-rewind episode, a forced /unstick included: ${ending}`, async () => { + const f = await fixture(); + try { + const internal = f.framework as unknown as Record< + 'consecutiveInferenceFailures' | 'exhaustionRewinds' | 'refusalRewinds' | 'refusalStreak' | 'rewindEpisode' | 'forcedRewind', + Map + >; + for (const key of ['consecutiveInferenceFailures', 'exhaustionRewinds', 'refusalRewinds', 'refusalStreak'] as const) { + internal[key].set('worker', 2); + } + internal.rewindEpisode.set('worker', { markerId: 'marker', count: 1, category: 'cyber' }); + internal.forcedRewind.set('worker', { remaining: 1, removed: [], serverId: 'ops', channelId: 'ops-channel' }); + const posted: Array<{ channelId: string; content: string }> = []; + (f.framework as unknown as { postToChannel: (s: string, c: string, content: string) => Promise }) + .postToChannel = async (_server, channelId, content) => { posted.push({ channelId, content }); }; + f.membrane.pushResponse(response()); + await run(f.framework, 'worker'); + assert.equal(posted.length, 1, 'the /unstick outcome was posted'); + assert.equal(posted[0].channelId, 'ops-channel'); + assert.match(posted[0].content, /the model responded/); + assert.equal(internal.consecutiveInferenceFailures.get('worker'), 0, 'the hard-down streak'); + assert.equal(internal.exhaustionRewinds.get('worker'), 0, 'the poison-history rewind budget'); + assert.equal(internal.refusalRewinds.get('worker'), 0); + assert.equal(internal.refusalStreak.has('worker'), false); + assert.equal(internal.rewindEpisode.has('worker'), false); + assert.equal(internal.forcedRewind.has('worker'), false, 'the /unstick session reports the model responded'); + } finally { + await f.framework.stop(); + rmSync(f.dir, { recursive: true, force: true }); + } + }); + } + + it('a provider cooldown, which records its recovery', async () => { + const r = await resident([ + (request) => new ThrowingStream(accelerationLimit(request)), + () => new MockYieldingStream([toolResponse('finish', 'finish-1', totalUsage)]), + ]); + r.framework.start(); + try { + r.framework.pushEvent({ type: 'external-message', source: 'test', content: 'first', metadata: {} } as unknown as ProcessEvent); + await r.framework.runUntilIdle(); + assert.equal(r.membrane.calls.length, 1, 'held for the cooldown'); + assert.ok(await pollUntil(() => r.membrane.calls.length === 2, 10_000), 'the cooldown ended and the call was retried'); + assert.ok(await pollUntil(() => r.traces.some((e) => e.type === 'inference:turn_ended')), 'the retried turn ended by tool'); + assert.equal(r.membrane.calls.length, 2); + assert.equal(r.internal.providerAccelerationLastRecovery.get('resident')?.stopReason, 'turn_ended'); + } finally { + await r.framework.stop(); + rmSync(r.dir, { recursive: true, force: true }); + } + }); + + it('a cooldown that cuts a turn after a finished round logs and counts that round', async () => { + const r = await resident([(request) => new RoundThenEnd(firstUsage, accelerationLimit(request))]); + r.framework.start(); + try { + r.framework.pushEvent({ type: 'external-message', source: 'test', content: 'first', metadata: {} } as unknown as ProcessEvent); + await r.framework.runUntilIdle(); + assert.ok(await pollUntil(() => r.traces.some((e) => e.type === 'usage:updated')), 'the finished round was counted'); + const logs = r.framework.queryInferenceLogs({ agentName: 'resident' }).entries.map((e) => e.entry); + assert.equal(logs.length, 1); + assert.match(String(logs[0].error), /^Provider acceleration cooldown: /); + assert.deepEqual(logs[0].tokenUsage, expectedTokens(firstUsage)); + assert.deepEqual(r.framework.getSessionUsage().totals, expectedTotals(firstUsage)); + assert.equal(r.traces.filter((e) => e.type === 'usage:updated').length, 1); + } finally { + await r.framework.stop(); + rmSync(r.dir, { recursive: true, force: true }); + } + }); + + it("a conversation agent's stated provider wait that cuts a turn after a finished round logs and counts that round once", async () => { + const statedWait = (request: unknown) => new MembraneError({ + type: 'rate_limit', retryable: true, httpStatus: 429, retryAfterMs: 60_000, + message: 'zz rate limit reached for requests', rawError: { status: 429 }, rawRequest: request, + }); + const r = await resident([(request) => new RoundThenEnd(firstUsage, statedWait(request))]); + // A conversation agent doesn't own provider admission, so a stated wait + // holds its turn (the "Provider wait" end) instead of the organization's + // cooldown. Its turns aren't an external message's default target, so + // the turn is started directly. + const internal = r.framework as unknown as { + conversationAgentHomes: Map; + startAgentStream(agent: unknown, trigger: unknown): Promise; + }; + r.framework.start(); + try { + internal.conversationAgentHomes.set('resident', 'world:zz'); + r.framework.getAgent('resident')!.getContextManager().addMessage('User', [{ type: 'text', text: 'first' }]); + await internal.startAgentStream(r.framework.getAgent('resident'), + { agentName: 'resident', reason: 'conversation', source: 'test', timestamp: Date.now() }); + assert.ok(await pollUntil(() => r.traces.some((e) => e.type === 'usage:updated') + && r.framework.queryInferenceLogs({ agentName: 'resident' }).total > 0), 'the finished round was counted and logged'); + const logs = r.framework.queryInferenceLogs({ agentName: 'resident' }).entries.map((e) => e.entry); + assert.equal(logs.length, 1, 'the wait writes the stream\'s one entry'); + assert.match(String(logs[0].error), /^Provider wait: /); + assert.deepEqual(logs[0].tokenUsage, expectedTokens(firstUsage), 'and it carries the finished round'); + assert.deepEqual(r.framework.getSessionUsage().totals, expectedTotals(firstUsage)); + assert.equal(r.traces.filter((e) => e.type === 'usage:updated').length, 1); + } finally { + await r.framework.stop(); + rmSync(r.dir, { recursive: true, force: true }); + } + }); + + it("a refusal completion's note carries what its count counted", async () => { + // The refused response carries no detailed usage, so the stream is + // counted from its latest usage sample (10 input, 5 output). + const r = await resident([() => new MockYieldingStream([createMockResponse([{ type: 'text', text: 'zz-refused' }], 'refusal')])]); + const notes: Array<{ tokens: unknown; totals: unknown }> = []; + const internal = r.framework as unknown as { logFailure: (entry: Record) => void }; + const logFailure = internal.logFailure.bind(r.framework); + // What the session had counted when the refusal was noted: this stream + // alone, before any retry a rewind might queue. + internal.logFailure = (entry) => { + if (entry.kind === 'refusal') notes.push({ tokens: entry.tokens, totals: structuredClone(r.framework.getSessionUsage().totals) }); + logFailure(entry); + }; + r.framework.start(); + try { + r.framework.pushEvent({ type: 'external-message', source: 'test', content: 'first', metadata: {} } as unknown as ProcessEvent); + await r.framework.runUntilIdle(); + assert.ok(await pollUntil(() => notes.length > 0), 'the refusal was noted'); + const counted: DetailedUsage = { inputTokens: 10, outputTokens: 5 }; + assert.deepEqual(notes[0].totals, expectedTotals(counted), 'counted from the latest sample'); + assert.deepEqual(notes[0].tokens, expectedTokens(counted), 'and the note carries the same'); + } finally { + await r.framework.stop(); + rmSync(r.dir, { recursive: true, force: true }); + } + }); + + it('turns that each fail after an answered call reach the hard-down threshold and run the breaker: an answered call is not a completed turn', async () => { + // Each turn's first call answers (its round's usage arrives) and the + // call after its tool result is refused as the provider can't take it, + // which shedMayHelp matches. + const refused = (): StreamEvent => ({ + type: 'error', + error: new MembraneError({ + type: 'invalid_request', retryable: false, httpStatus: 400, + message: 'zz the provider could not process the tool result', rawError: { status: 400 }, + }), + }) as StreamEvent; + const r = await resident([1, 2, 3].map(() => () => new RoundThenEnd(firstUsage, refused()))); + const internal = r.framework as unknown as { consecutiveInferenceFailures: Map }; + const streaks: number[] = []; + const answeredBefore: number[] = []; + let answered = 0; + // emitTrace counts the failure before listeners hear it. + r.framework.onTrace((event) => { + if (event.type === 'inference:usage') answered++; + if (event.type === 'inference:exhausted') { + streaks.push(internal.consecutiveInferenceFailures.get('resident') ?? 0); + answeredBefore.push(answered); + } + }); + const lines: string[] = []; + const original = console.error; + console.error = (...args: unknown[]) => { lines.push(args.map(String).join(' ')); }; + r.framework.start(); + try { + for (const [turn, text] of ['first', 'second', 'third'].entries()) { + r.framework.pushEvent({ type: 'external-message', source: 'test', content: text, metadata: {} } as unknown as ProcessEvent); + await r.framework.runUntilIdle(); + assert.ok(await pollUntil(() => streaks.length === turn + 1), `turn ${turn + 1} failed`); + } + assert.deepEqual(answeredBefore, [1, 2, 3], 'each failed turn had a call that answered'); + assert.deepEqual(streaks, [1, 2, 3], 'an answered call leaves the streak standing'); + const alerts = r.traces.filter((e) => e.type === 'ops:alert' && (e as { kind?: string }).kind === 'hard-down'); + assert.equal(alerts.length, 1, 'the third failed turn raises the hard-down alert'); + assert.ok(lines.some((line) => /^\[inference-rewind\] agent=resident auto-quarantined /.test(line)), + 'and the poison-history breaker sheds the newest exchange'); + } finally { + console.error = original; + await r.framework.stop(); + rmSync(r.dir, { recursive: true, force: true }); + } + }); +}); + +/** A completion the stream would yield next; these tests end it before it's handled. */ +const unseenCompletion = (): StreamEvent => + ({ type: 'complete', response: createMockResponse([{ type: 'text', text: 'unseen' }]) }) as StreamEvent; + +/** A stream whose first round never returns: it waits until it's cancelled. */ +class FirstRoundHeld implements YieldingStream { + isWaitingForTools = false; + pendingToolCallIds: string[] = []; + toolDepth = 0; + isCancelled = false; + started = false; + private wake: (() => void) | null = null; + provideToolResults(): void { throw new Error('not waiting'); } + cancel(): void { + this.isCancelled = true; + this.wake?.(); + } + async *[Symbol.asyncIterator](): AsyncIterator { + this.started = true; + if (!this.isCancelled) await new Promise((resolve) => { this.wake = resolve; }); + yield { type: 'aborted', reason: 'user' } as StreamEvent; + } +} + +/** A tool that returns only once released, so the stream waits on it. */ +class HeldTools implements Module { + readonly name = 'test'; + release: () => void = () => {}; + private readonly held = new Promise((resolve) => { this.release = resolve; }); + calledWith: string[] = []; + async start() {} + async stop() {} + async onProcess() { return {}; } + getTools() { + return [{ name: 'echo', description: 'echo', inputSchema: { type: 'object' as const, properties: {} } }]; + } + async handleToolCall(call: ToolCall) { + this.calledWith.push(call.id); + await this.held; + return { success: true, data: {} }; + } +} + +describe('a stream that ends without writing its own entry is logged once, in finally', () => { + async function residentWith(stream: YieldingStream, tools: Module) { + const dir = mkdtempSync(join(tmpdir(), 'af-end-turn-usage-')); + const membrane = new MockMembrane(); + let made = false; + membrane.streamYielding = (request: NormalizedRequest): YieldingStream => { + membrane.calls.push(request); + if (made) return new MockYieldingStream([createMockResponse([{ type: 'text', text: 'done' }])]); + made = true; + return stream; + }; + const framework = await AgentFramework.create({ + storePath: join(dir, 'store.chronicle'), membrane: membrane.asMembrane(), + agents: [{ name: 'resident', model: 'test-model', systemPrompt: 'system', allowedTools: 'all' }], + modules: [new InputModule(), tools], syncIntervalMs: 0, maintenanceIntervalMs: 0, + }); + const traces: TraceEvent[] = []; + framework.onTrace((event) => traces.push(event)); + framework.start(); + framework.pushEvent({ type: 'external-message', source: 'test', content: 'go', metadata: {} } as unknown as ProcessEvent); + return { dir, membrane, framework, traces }; + } + + function loggedOnce(framework: AgentFramework, traces: TraceEvent[], error: RegExp) { + assert.equal(traces.filter((e) => e.type === 'usage:updated').length, 1, 'counted once'); + assert.deepEqual(framework.getSessionUsage().totals, expectedTotals(firstUsage)); + const logs = framework.queryInferenceLogs({ agentName: 'resident' }).entries.map((e) => e.entry); + assert.equal(logs.length, 1, 'one entry'); + assert.equal(logs[0].success, false); + assert.match(String(logs[0].error), error); + assert.deepEqual(logs[0].tokenUsage, expectedTokens(firstUsage), 'the entry carries what was counted'); + } + + it('a turn abandoned by an operator quiesce after a finished round', async () => { + const tools = new HeldTools(); + const r = await residentWith(new RoundThenEnd(firstUsage, unseenCompletion()), tools); + try { + assert.ok(await pollUntil(() => tools.calledWith.length === 1), 'the round finished and its tool is running'); + await r.framework.quiesce({ timeoutMs: 50, abandon: true }); + assert.ok(await pollUntil(() => r.framework.queryInferenceLogs({ agentName: 'resident' }).total > 0), 'logged'); + loggedOnce(r.framework, r.traces, /^Turn abandoned by operator quiesce$/); + } finally { + // Let the held tool finish while the framework's queue is still open. + tools.release(); + await pollUntil(() => r.traces.some((e) => + (e.type === 'tool:completed' || e.type === 'tool:failed') && (e as { callId?: string }).callId === 'echo-1')); + await r.framework.stop(); + rmSync(r.dir, { recursive: true, force: true }); + } + }); + + it('a failure after a completion is logged without carrying the usage again', async () => { + const r = await resident([() => new MockYieldingStream([finalResponse('Done.', totalUsage)])]); + // Something after the completion's entry throws: the stream's catch + // writes a failure entry too. (A resident: an ephemeral is deregistered + // once its completion settles it, so its catch treats the throw as a + // superseded stream's.) + const internal = r.framework as unknown as { endRefusalEpisode: (name: string) => void }; + const end = internal.endRefusalEpisode.bind(r.framework); + internal.endRefusalEpisode = () => { + internal.endRefusalEpisode = end; + throw new Error('after the completion'); + }; + r.framework.start(); + try { + r.framework.pushEvent({ type: 'external-message', source: 'test', content: 'go', metadata: {} } as unknown as ProcessEvent); + assert.ok(await pollUntil(() => r.framework.queryInferenceLogs({ agentName: 'resident' }).total === 2), 'both entries'); + // Both entries carry the stream's start time, so find each by outcome. + const logs = r.framework.queryInferenceLogs({ agentName: 'resident' }).entries.map((e) => e.entry); + const completed = logs.find((e) => e.success); + const failed = logs.find((e) => !e.success); + assert.deepEqual(completed?.tokenUsage, expectedTokens(totalUsage), 'the completion carries the usage'); + assert.equal(failed?.error, 'Stream threw: after the completion'); + assert.equal(failed?.tokenUsage, undefined, 'the later failure does not carry it again'); + assert.deepEqual(r.framework.getSessionUsage().totals, expectedTotals(totalUsage)); + } finally { + await r.framework.stop(); + rmSync(r.dir, { recursive: true, force: true }); + } + }); + + it('a stream whose agent was replaced under its name after a finished round', async () => { + const stream = new RoundThenEnd(firstUsage, unseenCompletion()); + let framework: AgentFramework | undefined; + let original: object | undefined; + const agentsOf = () => (framework as unknown as { agents: Map }).agents; + const provide = stream.provideToolResults.bind(stream); + // The agent object is replaced under its name while this stream waits on its tool. + stream.provideToolResults = () => { + original = agentsOf().get('resident'); + agentsOf().set('resident', Object.create(original!)); + provide(); + }; + const r = await residentWith(stream, new TurnTools()); + framework = r.framework; + try { + assert.ok(await pollUntil(() => r.framework.queryInferenceLogs({ agentName: 'resident' }).total > 0), 'logged'); + loggedOnce(r.framework, r.traces, /^Stream abandoned: its agent was disposed or replaced$/); + } finally { + if (original) agentsOf().set('resident', original); + await r.framework.stop(); + rmSync(r.dir, { recursive: true, force: true }); + } + }); + + it('a stream that just closes after a finished round', async () => { + const r = await residentWith(new RoundThenEnd(firstUsage, 'close'), new TurnTools()); + try { + assert.ok(await pollUntil(() => r.framework.queryInferenceLogs({ agentName: 'resident' }).total > 0), 'logged'); + loggedOnce(r.framework, r.traces, /^Stream ended without a final response$/); + } finally { + await r.framework.stop(); + rmSync(r.dir, { recursive: true, force: true }); + } + }); + + it('a turn abandoned by an operator quiesce before its first round returns has nothing to carry, and writes no entry', async () => { + const stream = new FirstRoundHeld(); + const r = await residentWith(stream, new TurnTools()); + const internal = r.framework as unknown as { activeTurnTokens: Map; activeStreams: Map }; + try { + assert.ok(await pollUntil(() => stream.started), 'the stream is running'); + await r.framework.quiesce({ timeoutMs: 50, abandon: true }); + assert.ok(await pollUntil(() => stream.isCancelled), 'the abandon cancelled it'); + // The stream and its turn are released at the end of finally, after its accounting. + assert.ok(await pollUntil(() => !internal.activeTurnTokens.has('resident') && !internal.activeStreams.has('resident')), + 'its teardown ran'); + assert.equal(r.framework.queryInferenceLogs({ agentName: 'resident' }).total, 0, 'no entry'); + assert.equal(r.traces.filter((e) => e.type === 'usage:updated').length, 0, 'nothing counted'); + assert.equal(r.framework.getSessionUsage().inferenceCount, 0); + } finally { + await r.framework.stop(); + rmSync(r.dir, { recursive: true, force: true }); + } + }); + + it("a write that fails while finally logs the stream doesn't skip the teardown after it", async () => { + const tools = new HeldTools(); + const r = await residentWith(new RoundThenEnd(firstUsage, unseenCompletion()), tools); + const internal = r.framework as unknown as { + logInference: (entry: { error?: string }) => void; + activeTurnTokens: Map; + activeStreams: Map; + }; + const write = internal.logInference.bind(r.framework); + internal.logInference = (entry) => { + if (entry.error === 'Turn abandoned by operator quiesce') throw new Error('zz store refused the write'); + write(entry); + }; + const lines: string[] = []; + const original = console.error; + console.error = (...args: unknown[]) => { lines.push(args.map(String).join(' ')); }; + try { + assert.ok(await pollUntil(() => tools.calledWith.length === 1), 'the round finished and its tool is running'); + await r.framework.quiesce({ timeoutMs: 50, abandon: true }); + assert.ok(await pollUntil(() => lines.some((line) => line.startsWith('[usage] agent=resident could not record'))), + 'the failed write is said'); + assert.ok(await pollUntil(() => !internal.activeTurnTokens.has('resident') && !internal.activeStreams.has('resident')), + 'and the stream and its turn are still released'); + } finally { + console.error = original; + tools.release(); + await pollUntil(() => r.traces.some((e) => + (e.type === 'tool:completed' || e.type === 'tool:failed') && (e as { callId?: string }).callId === 'echo-1')); + await r.framework.stop(); + rmSync(r.dir, { recursive: true, force: true }); + } + }); + + it('a stream superseded by a newer one after a finished round', async () => { + const stream = new RoundThenEnd(firstUsage, unseenCompletion()); + let framework: AgentFramework | undefined; + const provide = stream.provideToolResults.bind(stream); + // A newer stream of the agent takes over while this one waits on its tool. + stream.provideToolResults = () => { + (framework!.getAgent('resident') as unknown as { _streamId: number })._streamId++; + provide(); + }; + const r = await residentWith(stream, new TurnTools()); + framework = r.framework; + try { + assert.ok(await pollUntil(() => r.framework.queryInferenceLogs({ agentName: 'resident' }).total > 0), 'logged'); + loggedOnce(r.framework, r.traces, /^Stream superseded/); + } finally { + await r.framework.stop(); + rmSync(r.dir, { recursive: true, force: true }); + } + }); +}); + +describe("a write that fails at a tool boundary doesn't skip what the boundary does after it", () => { + /** Makes the framework's inference-log write throw for the entries + * `refuse` picks, as a store that refuses the write would, and collects + * stderr so a test can see the failure reported. */ + function refuseWrites(framework: AgentFramework, refuse: (entry: { stopReason?: string }) => boolean) { + const internal = framework as unknown as { logInference: (entry: { stopReason?: string }) => void }; + const write = internal.logInference.bind(framework); + internal.logInference = (entry) => { + if (refuse(entry)) throw new Error('zz store refused the write'); + write(entry); + }; + const lines: string[] = []; + const original = console.error; + console.error = (...args: unknown[]) => { lines.push(args.map(String).join(' ')); }; + return { + reported: (agent: string) => lines.some((line) => + line.startsWith(`[usage] agent=${agent} could not record a stream's usage at its end`) + && line.includes('zz store refused the write')), + restore: () => { console.error = original; }, + }; + } + + it('a restart still opens its next stream', async () => { + const dir = mkdtempSync(join(tmpdir(), 'af-end-turn-usage-')); + const membrane = new SequentialStreamMembrane(); + membrane.pushResponse(toolResponse('echo', 'echo-1', firstUsage)); + membrane.pushResponse(finalResponse('Recovered.', secondUsage)); + const framework = await AgentFramework.create({ + storePath: join(dir, 'store.chronicle'), membrane: membrane.asMembrane(), agents: [], + modules: [new TurnTools()], syncIntervalMs: 0, + }); + const traces: TraceEvent[] = []; + framework.onTrace((event) => traces.push(event)); + const stderr = refuseWrites(framework, (entry) => entry.stopReason === 'context_budget'); + framework.start(); + try { + const { agent, contextManager } = await framework.createEphemeralAgent({ + name: 'worker', model: 'test-model', systemPrompt: 'Do the task.', allowedTools: 'all', + proseRouting: 'disabled', maxStreamTokens: 50, + }); + contextManager.addMessage('user', [{ type: 'text', text: 'Run once.' }]); + const run = framework.runEphemeralToCompletion(agent, contextManager); + // Awaited below once the restart is in. Without the restart it never + // settles, and stop() mustn't turn it into an unhandled rejection. + run.catch(() => {}); + assert.ok(await pollUntil(() => membrane.calls.length === 2), 'the restart opened its next stream'); + await run; + assert.equal(traces.filter((e) => e.type === 'inference:stream_restarted').length, 1, 'the restart is announced'); + assert.equal(traces.filter((e) => e.type === 'inference:completed').length, 1, 'the turn ends on its next stream'); + assert.ok(stderr.reported('worker'), 'the failed write is said'); + } finally { + stderr.restore(); + await framework.stop(); + rmSync(dir, { recursive: true, force: true }); + } + }); + + it("a tool-ended turn still ends its refusal episode, says it ended, and hands the tool's result to the modules", async () => { + const f = await fixture(); + const stderr = refuseWrites(f.framework, (entry) => entry.stopReason === 'turn_ended'); + try { + const refusalStreak = (f.framework as unknown as { refusalStreak: Map }).refusalStreak; + refusalStreak.set('worker', 2); + f.membrane.pushResponse(toolResponse('finish', 'finish-1', totalUsage)); + await run(f.framework, 'worker'); + assert.equal(refusalStreak.has('worker'), false, 'the refusal episode ends'); + assert.equal(f.traces.filter((e) => e.type === 'inference:turn_ended').length, 1, 'inference:turn_ended'); + assert.ok(await pollUntil(() => f.traces.some((e) => e.type === 'process:completed' + && (e as { processEvent?: { type?: string } }).processEvent?.type === 'tool-result')), + "the tool result's event reaches the modules"); + assert.ok(stderr.reported('worker'), 'the failed write is said'); + } finally { + stderr.restore(); + await f.framework.stop(); + rmSync(f.dir, { recursive: true, force: true }); + } + }); +}); diff --git a/test/tool-result-guard.test.ts b/test/tool-result-guard.test.ts index f5d6b02b5..1ca41a243 100644 --- a/test/tool-result-guard.test.ts +++ b/test/tool-result-guard.test.ts @@ -1252,7 +1252,50 @@ test('usage of an abandoned guarded round is counted in session totals', async ( try { await h.run(); const totals = h.framework.getSessionUsage().totals as unknown as Record; - assert.ok(totals.inputTokens >= 1_007, `refused round usage missing: ${JSON.stringify(totals)}`); + assert.equal(totals.inputTokens, 1_007, `the refused round is counted, once: ${JSON.stringify(totals)}`); + } finally { await h.framework.stop(); } +}); + +test('an abandoned guarded round whose response carries no usage still counts its rounds', async () => { + // Membrane's usage samples are cumulative across a stream's rounds: the + // withheld stream's two rounds emit 10, then 20 input tokens, and the + // retry's one round 10. No final response carries details.usage, so each + // stream counts its latest sample. + const sampled = (response: NormalizedResponse, inputTokens: number, outputTokens: number) => + ({ ...response, usage: { inputTokens, outputTokens } }) as NormalizedResponse; + const h = await harness([[sampled(calls('one'), 10, 5), sampled(refused(), 20, 10)], [sampled(answer(), 10, 5)]], { toolResultGuard: true }); + try { + await h.run(); + const totals = h.framework.getSessionUsage().totals as unknown as Record; + assert.equal(totals.inputTokens, 30, `each stream is counted once, from its latest sample: ${JSON.stringify(totals)}`); + // The log says what the totals count: the withheld stream, then the retry. + const logs = h.framework.queryInferenceLogs({ agentName: 'assistant' }).entries.map((e) => e.entry).reverse(); + assert.deepEqual(logs.map((e) => [e.success, e.tokenUsage]), + [[false, { input: 20, output: 10 }], [true, { input: 10, output: 5 }]]); + } finally { await h.framework.stop(); } +}); + +test("an abandoned guarded round with no usage sample counts and notes its response's own usage", async () => { + const h = await harness([[calls('one'), refused()], [answer()]], { toolResultGuard: true }); + // A stream that emits no usage events: only the refused response's own + // usage (10 input tokens, from the shared mock) says what it was billed. + const make = h.membrane.streamYielding.bind(h.membrane); + h.membrane.streamYielding = (request, options) => { + const stream = make(request, options); + const events = stream[Symbol.asyncIterator].bind(stream); + stream[Symbol.asyncIterator] = async function* () { + const it = events(); + for (let r = await it.next(); !r.done; r = await it.next()) if (r.value.type !== 'usage') yield r.value; + }; + return stream; + }; + try { + await h.run(); + const totals = h.framework.getSessionUsage().totals as unknown as Record; + assert.equal(totals.inputTokens, 10, `the refused round is counted from its response: ${JSON.stringify(totals)}`); + const withheld = h.framework.queryInferenceLogs({ agentName: 'assistant' }).entries.map((e) => e.entry) + .find((e) => e.stopReason === 'refusal'); + assert.deepEqual(withheld?.tokenUsage, { input: 10, output: 5 }, "the withheld stream's entry carries the response's usage"); } finally { await h.framework.stop(); } });