From cccea0e7b2aca94d2fcc207123a7e2b204adf7ee Mon Sep 17 00:00:00 2001 From: Drewry Pope Date: Wed, 28 Oct 2026 15:45:44 -0500 Subject: [PATCH 1/3] feat(core): stitch continued route attempts into one turn Response chunks produced by an output-token-limit continuation were logged as unrelated route_attempt rows. Record the attempt that started the response on every continuation and roll attempts up per turn so route_usage reports one token and cost figure for a multi-chunk response. - route_attempt gains continuation_of and continuation_index plus an index on the link - the runner passes the root assistant message and chunk index; the writer resolves the root row id at persist time, so a retried or hedged first chunk still links to the attempt that answered - ModelUsage.turns groups a root attempt, its retries and hedge legs and its continuations; route_usage lists continued turns - tests cover linking, failed chunks and turn totals --- packages/core/schema.json | 38 +++- packages/core/src/database/migration.gen.ts | 2 + ...261028203055_route_attempt_continuation.ts | 15 ++ packages/core/src/database/schema.gen.ts | 3 + packages/core/src/model-route-log.ts | 65 ++++++- packages/core/src/model-route-log/sql.ts | 7 + packages/core/src/model-usage.ts | 80 ++++++++ packages/core/src/session/runner/llm.ts | 13 ++ packages/core/src/session/runner/step.ts | 9 + packages/core/src/tool/plugin/route-usage.ts | 23 ++- packages/core/test/model-route-log.test.ts | 184 ++++++++++++++++++ packages/core/test/model-usage.test.ts | 139 +++++++++++++ packages/core/test/session-runner.test.ts | 36 ++++ 13 files changed, 606 insertions(+), 8 deletions(-) create mode 100644 packages/core/src/database/migration/20261028203055_route_attempt_continuation.ts diff --git a/packages/core/schema.json b/packages/core/schema.json index 31264785f644..48c2dfc3b872 100644 --- a/packages/core/schema.json +++ b/packages/core/schema.json @@ -1,8 +1,8 @@ { "version": "7", "dialect": "sqlite", - "id": "3b8437ee-1a59-400d-9a69-ad74f9d1e8f5", - "prevIds": ["c9735d25-ef18-4dc1-ae6d-00504c68d035"], + "id": "2d44615a-d376-47d0-bd31-b05ed305aeb1", + "prevIds": ["3b8437ee-1a59-400d-9a69-ad74f9d1e8f5"], "ddl": [ { "name": "account_state", @@ -1384,6 +1384,26 @@ "entityType": "columns", "table": "route_attempt" }, + { + "type": "integer", + "notNull": false, + "autoincrement": false, + "default": null, + "generated": null, + "name": "continuation_of", + "entityType": "columns", + "table": "route_attempt" + }, + { + "type": "integer", + "notNull": false, + "autoincrement": false, + "default": "0", + "generated": null, + "name": "continuation_index", + "entityType": "columns", + "table": "route_attempt" + }, { "type": "text", "notNull": false, @@ -3389,6 +3409,20 @@ "entityType": "indexes", "table": "route_attempt" }, + { + "columns": [ + { + "value": "continuation_of", + "isExpression": false + } + ], + "isUnique": false, + "where": null, + "origin": "manual", + "name": "route_attempt_continuation_idx", + "entityType": "indexes", + "table": "route_attempt" + }, { "columns": [ { diff --git a/packages/core/src/database/migration.gen.ts b/packages/core/src/database/migration.gen.ts index 3db5cafd249e..6e89c693cc03 100644 --- a/packages/core/src/database/migration.gen.ts +++ b/packages/core/src/database/migration.gen.ts @@ -56,6 +56,7 @@ import m53 from "./migration/20261010180456_project_time_archived.js" import m54 from "./migration/20261014094820_route_attempt_cost_usd.js" import m55 from "./migration/20261016221426_messaging.js" import m56 from "./migration/20261025175023_session_subagent_policy.js" +import m57 from "./migration/20261028203055_route_attempt_continuation.js" export const migrations = [ m00, @@ -115,4 +116,5 @@ export const migrations = [ m54, m55, m56, + m57, ] satisfies DatabaseMigration.Migration[] diff --git a/packages/core/src/database/migration/20261028203055_route_attempt_continuation.ts b/packages/core/src/database/migration/20261028203055_route_attempt_continuation.ts new file mode 100644 index 000000000000..d8f825f0e1a4 --- /dev/null +++ b/packages/core/src/database/migration/20261028203055_route_attempt_continuation.ts @@ -0,0 +1,15 @@ +import { Effect } from "effect" +import type { DatabaseMigration } from "../migration.js" + +const migration: DatabaseMigration.Migration = { + id: "20261028203055_route_attempt_continuation", + up(tx) { + return Effect.gen(function* () { + yield* tx.run(`ALTER TABLE \`route_attempt\` ADD \`continuation_of\` integer;`) + yield* tx.run(`ALTER TABLE \`route_attempt\` ADD \`continuation_index\` integer DEFAULT 0;`) + yield* tx.run(`CREATE INDEX \`route_attempt_continuation_idx\` ON \`route_attempt\` (\`continuation_of\`);`) + }) + }, +} + +export default migration diff --git a/packages/core/src/database/schema.gen.ts b/packages/core/src/database/schema.gen.ts index 66069e7d87b6..5a19854f389d 100644 --- a/packages/core/src/database/schema.gen.ts +++ b/packages/core/src/database/schema.gen.ts @@ -200,6 +200,8 @@ const schema: Omit = { \`tokens_cache_read\` integer, \`tokens_cache_write\` integer, \`cost_usd\` real, + \`continuation_of\` integer, + \`continuation_index\` integer DEFAULT 0, \`quota\` text ); `) @@ -445,6 +447,7 @@ const schema: Omit = { yield* tx.run( `CREATE INDEX \`route_attempt_session_idx\` ON \`route_attempt\` (\`session_id\`,\`time_started\`);`, ) + yield* tx.run(`CREATE INDEX \`route_attempt_continuation_idx\` ON \`route_attempt\` (\`continuation_of\`);`) yield* tx.run(`CREATE INDEX \`route_decision_route_time_idx\` ON \`route_decision\` (\`route_id\`,\`time\`);`) yield* tx.run(`CREATE INDEX \`route_decision_session_idx\` ON \`route_decision\` (\`session_id\`,\`time\`);`) yield* tx.run( diff --git a/packages/core/src/model-route-log.ts b/packages/core/src/model-route-log.ts index 91f8cafc80fc..37b881ad824b 100644 --- a/packages/core/src/model-route-log.ts +++ b/packages/core/src/model-route-log.ts @@ -16,7 +16,13 @@ export type Note = typeof RouteNoteTable.$inferInsert export type Event = | { readonly kind: "decision"; readonly row: Decision } - | { readonly kind: "attempt"; readonly row: Attempt } + | { + readonly kind: "attempt" + readonly row: Attempt + /** The assistant message of the attempt this request continues. The writer resolves it to that attempt's row id + * when it persists the row, because the id only exists after the root was inserted. */ + readonly continuationOfMessageID?: string + } | { readonly kind: "health"; readonly row: Health } | { readonly kind: "note"; readonly row: Note } @@ -272,15 +278,37 @@ export const failureFields = (reason: unknown, options?: { readonly captureBody? const persist = (db: Database.Interface["db"], events: readonly Event[]) => Effect.gen(function* () { const decisions = events.flatMap((event) => (event.kind === "decision" ? [event.row] : [])) - const attempts = events.flatMap((event) => (event.kind === "attempt" ? [event.row] : [])) + const attempts = events.flatMap((event) => (event.kind === "attempt" ? [event] : [])) const health = events.flatMap((event) => (event.kind === "health" ? [event.row] : [])) if (decisions.length) yield* db.insert(RouteDecisionTable).values(decisions).run() - if (attempts.length) yield* db.insert(RouteAttemptTable).values(attempts).run() + // One statement per attempt: a continuation resolves its root from the row inserted just before it. + for (const attempt of attempts) yield* persistAttempt(db, attempt) if (health.length) yield* db.insert(RouteHealthTable).values(health).run() const notes = events.flatMap((event) => (event.kind === "note" ? [event.row] : [])) if (notes.length) yield* db.insert(RouteNoteTable).values(notes).run() }) +/** Inserts one attempt and, for a continuation, points `continuation_of` at the attempt that started the response. + * The root is read back by its assistant message, so the link survives the root having been persisted in an earlier + * batch and stays correct across a restart. Retries and hedge losers share that message, so the newest standalone row + * is the attempt that actually answered — the one whose response hit the output token limit. */ +const persistAttempt = (db: Database.Interface["db"], event: Extract) => + Effect.gen(function* () { + const inserted = yield* db + .insert(RouteAttemptTable) + .values(event.row) + .returning({ id: RouteAttemptTable.id }) + .get() + if (event.continuationOfMessageID === undefined) return + yield* db + .update(RouteAttemptTable) + .set({ + continuation_of: sql`(select root.id from route_attempt root where root.session_id = ${event.row.session_id} and root.assistant_message_id = ${event.continuationOfMessageID} and root.continuation_index = 0 order by root.id desc limit 1)`, + }) + .where(eq(RouteAttemptTable.id, inserted.id)) + .run() + }) + /** Receives the still-active adjustments stored by earlier runs. Set by the routing module, which * this one cannot import without a cycle. */ let restore: ((notes: readonly Record[]) => void) | undefined @@ -627,6 +655,37 @@ export const attemptRows = (db: Database.Interface["db"], since: number, routeID .limit(ROW_LIMIT) .all() +/** Attempt rows in the window with their continuation link, so a response split by output-token-limit + * continuations can be rolled up as one turn. Newest first; `ModelUsage.turns` does the arithmetic. */ +export const turnAttempts = (db: Database.Interface["db"], since: number, sessionID?: string) => + db + .select({ + id: RouteAttemptTable.id, + continuationOf: RouteAttemptTable.continuation_of, + continuationIndex: RouteAttemptTable.continuation_index, + assistantMessageID: RouteAttemptTable.assistant_message_id, + sessionID: RouteAttemptTable.session_id, + timeStarted: RouteAttemptTable.time_started, + timeEnded: RouteAttemptTable.time_ended, + outcome: RouteAttemptTable.outcome, + input: RouteAttemptTable.tokens_input, + output: RouteAttemptTable.tokens_output, + reasoning: RouteAttemptTable.tokens_reasoning, + cacheRead: RouteAttemptTable.tokens_cache_read, + cacheWrite: RouteAttemptTable.tokens_cache_write, + costUsd: RouteAttemptTable.cost_usd, + }) + .from(RouteAttemptTable) + .where( + and( + gte(RouteAttemptTable.time_started, since), + sessionID === undefined ? undefined : eq(RouteAttemptTable.session_id, sessionID), + ), + ) + .orderBy(desc(RouteAttemptTable.time_started)) + .limit(ROW_LIMIT) + .all() + /** Selections in the window, newest first: who was eligible, the weights and exclusions behind each, and who was chosen. */ export const decisionRows = (db: Database.Interface["db"], since: number, routeID?: string) => db diff --git a/packages/core/src/model-route-log/sql.ts b/packages/core/src/model-route-log/sql.ts index c416207935a7..d2fcd23d737b 100644 --- a/packages/core/src/model-route-log/sql.ts +++ b/packages/core/src/model-route-log/sql.ts @@ -67,6 +67,12 @@ export const RouteAttemptTable = sqliteTable( /** What the request cost in USD, from the model's catalog prices and the tokens it used. Null when the request * did not finish with usage. Stored with the attempt so the cost outlives the session it belongs to. */ cost_usd: real(), + /** The attempt that started the response this request continued, when a response hit the output token limit and + * the session runner asked for the rest. Null for a request that stands on its own. A turn is a root attempt plus + * every attempt that names it here, so its tokens and cost can be read as one figure. */ + continuation_of: integer(), + /** 0 for a standalone request, then 1, 2, … for each chunk that continues the root attempt. */ + continuation_index: integer().default(0), /** Provider rate-limit and quota headers or fields, when the error carried them. */ quota: text({ mode: "json" }).$type>(), }, @@ -74,6 +80,7 @@ export const RouteAttemptTable = sqliteTable( index("route_attempt_target_time_idx").on(table.provider_id, table.model_id, table.time_started), index("route_attempt_route_time_idx").on(table.route_id, table.time_started), index("route_attempt_session_idx").on(table.session_id, table.time_started), + index("route_attempt_continuation_idx").on(table.continuation_of), ], ) diff --git a/packages/core/src/model-usage.ts b/packages/core/src/model-usage.ts index e855afa6ac00..705f02de2119 100644 --- a/packages/core/src/model-usage.ts +++ b/packages/core/src/model-usage.ts @@ -239,6 +239,86 @@ export interface DayReport { > } +/** One request as far as turn stitching needs it: the link to the attempt it continues, plus what it cost. */ +export interface TurnAttemptRow { + readonly id: number + readonly continuationOf: number | null + readonly continuationIndex: number | null + readonly assistantMessageID: string | null + readonly sessionID: string | null + readonly timeStarted: number + readonly timeEnded: number + readonly outcome: string + readonly input: number | null + readonly output: number | null + readonly reasoning: number | null + readonly cacheRead: number | null + readonly cacheWrite: number | null + readonly costUsd: number | null +} + +export interface TurnRow extends Totals { + /** Id of the attempt that started the turn; for a partial turn whose root fell outside the window, the id it names. */ + readonly rootID: number + readonly sessionID: string | null + /** Response segments: the root attempt plus each output-token-limit continuation, however many requests each took. */ + readonly chunks: number + readonly first: number + readonly last: number + readonly durationMs: number +} + +/** Rolls attempts up into turns: a root attempt plus every attempt whose `continuation_of` names it. Retries and hedge + * legs share the assistant message of the chunk they belong to, so the newest standalone attempt of a message is that + * chunk's root and the others join it; `requests` therefore counts every billed attempt, while `chunks` counts response + * segments, so a response split by output-token-limit continuations reads as one cost instead of several rows. */ +export const turns = (rows: readonly TurnAttemptRow[]): TurnRow[] => { + // Each chunk's root is the newest standalone attempt of its assistant message; retries and hedge legs of the same + // chunk share the message and join it instead of forming turns of their own. + const roots = new Map() + for (const row of rows) { + if ((row.continuationIndex ?? 0) !== 0 || row.assistantMessageID === null) continue + roots.set(row.assistantMessageID, Math.max(roots.get(row.assistantMessageID) ?? 0, row.id)) + } + const rootOf = (row: TurnAttemptRow) => + row.continuationOf ?? (row.assistantMessageID === null ? row.id : (roots.get(row.assistantMessageID) ?? row.id)) + const grouped = new Map() + for (const row of rows) { + const rootID = rootOf(row) + const group = grouped.get(rootID) + if (group) group.push(row) + else grouped.set(rootID, [row]) + } + return [...grouped] + .map(([rootID, attempts]) => { + const root = attempts.find((attempt) => attempt.id === rootID) + const totals = sum( + attempts.map((attempt) => ({ + requests: 1, + failures: attempt.outcome === "failure" || attempt.outcome === "timeout" ? 1 : 0, + input: attempt.input ?? 0, + output: attempt.output ?? 0, + reasoning: attempt.reasoning ?? 0, + cacheRead: attempt.cacheRead ?? 0, + cacheWrite: attempt.cacheWrite ?? 0, + costUsd: attempt.costUsd ?? 0, + })), + ) + const first = Math.min(...attempts.map((attempt) => attempt.timeStarted)) + const last = Math.max(...attempts.map((attempt) => attempt.timeEnded)) + return { + ...totals, + rootID, + sessionID: root?.sessionID ?? attempts[0]?.sessionID ?? null, + chunks: 1 + Math.max(...attempts.map((attempt) => attempt.continuationIndex ?? 0)), + first, + last, + durationMs: last - first, + } + }) + .sort((a, b) => b.last - a.last || b.rootID - a.rootID) +} + export const dayReports = (rows: readonly DayRow[]): DayReport[] => { const days = new Map() for (const row of rows) days.set(row.day, [...(days.get(row.day) ?? []), row]) diff --git a/packages/core/src/session/runner/llm.ts b/packages/core/src/session/runner/llm.ts index 00157e719663..87f8005d1f3b 100644 --- a/packages/core/src/session/runner/llm.ts +++ b/packages/core/src/session/runner/llm.ts @@ -209,6 +209,10 @@ const layer = Layer.effect( const runStep = Effect.fn("SessionRunner.runStep")(function* (first: SessionContext.Loaded, step: number) { const sessionID = first.session.id let assistantMessageID = SessionMessage.ID.create() + // Set once this step's response has been continued after an output token limit. Every later chunk names the + // attempt that started the response, so a turn's tokens and cost can be read as one figure. + let continuationOf: SessionMessage.ID | undefined + let continuationIndex = 0 const retry = yield* SessionRunnerRetry.make(bus, sessionID) let initial: SessionContext.Loaded | undefined = first let recoverOverflow = true @@ -221,6 +225,8 @@ const layer = Layer.effect( if (compacted.status === "failed") return yield* new StepFailedError({ error: compacted.error }) if (compacted.status === "completed") { assistantMessageID = SessionMessage.ID.create() + continuationOf = undefined + continuationIndex = 0 continue } const stepLimitReached = loaded.agent.info.steps !== undefined && step >= loaded.agent.info.steps @@ -254,6 +260,7 @@ const layer = Layer.effect( isLocationClosed: lifecycle.isClosed, sessionID, assistantMessageID, + ...(continuationOf === undefined ? {} : { continuation: { of: continuationOf, index: continuationIndex } }), agent: loaded.agent.id, model: loaded.model, prepared, @@ -297,6 +304,8 @@ const layer = Layer.effect( }) yield* bus.publish(SessionEvent.Synthetic, { sessionID, text: CONTINUE_AFTER_INCOMPLETE_STREAM }) assistantMessageID = SessionMessage.ID.create() + continuationOf = undefined + continuationIndex = 0 }), OutputLimit: Effect.fn("SessionRunner.continueOutputLimit")(function* (outcome) { // The current response is included. User input or any other finish breaks the streak. @@ -326,11 +335,15 @@ const layer = Layer.effect( }) if (outcome.needsContinuation) return true yield* bus.publish(SessionEvent.Synthetic, { sessionID, text: CONTINUE_AFTER_OUTPUT_LIMIT }) + continuationOf ??= assistantMessageID + continuationIndex++ assistantMessageID = SessionMessage.ID.create() }), Compacted: Effect.fnUntraced(function* () { recoverOverflow = false assistantMessageID = SessionMessage.ID.create() + continuationOf = undefined + continuationIndex = 0 }), RecoverFull: Effect.fnUntraced(function* () { recoverContinuation = false diff --git a/packages/core/src/session/runner/step.ts b/packages/core/src/session/runner/step.ts index 0555ca67f0f8..ebebde049794 100644 --- a/packages/core/src/session/runner/step.ts +++ b/packages/core/src/session/runner/step.ts @@ -54,6 +54,9 @@ interface Input { readonly isLocationClosed: () => boolean readonly sessionID: SessionSchema.ID readonly assistantMessageID: SessionMessage.ID + /** Set when this request continues a response that hit the output token limit: the assistant message of the + * attempt that started the response, and this chunk's index (1 for the first continuation). */ + readonly continuation?: { readonly of: SessionMessage.ID; readonly index: number } readonly agent: Agent.ID readonly model: SessionRunnerModel.Resolved /** Prompt-size estimate for budget accounting when a hedge is cancelled before provider usage arrives. */ @@ -345,7 +348,9 @@ export const make = Effect.gen(function* () { retryable: false, output_started: false, response_ms: streamedAt - requestStarted, + continuation_index: input.continuation?.index, }, + continuationOfMessageID: input.continuation?.of, }) return Outcome.Compacted() } @@ -417,7 +422,9 @@ export const make = Effect.gen(function* () { tokens_cache_read: tokens?.cache.read, tokens_cache_write: tokens?.cache.write, cost_usd: tokens ? SessionUsage.calculateCost(input.model.cost, tokens) : undefined, + continuation_index: input.continuation?.index, }, + continuationOfMessageID: input.continuation?.of, }) } if ( @@ -664,7 +671,9 @@ export const make = Effect.gen(function* () { response_ms: raced.loserElapsedMs, tokens_input: input.estimatedInputTokens, tokens_estimated: true, + continuation_index: input.continuation?.index, }, + continuationOfMessageID: input.continuation?.of, }) } const winner = raced.winner === "hedge" ? raced.value : undefined diff --git a/packages/core/src/tool/plugin/route-usage.ts b/packages/core/src/tool/plugin/route-usage.ts index 99dfd44faf06..361b310c6be8 100644 --- a/packages/core/src/tool/plugin/route-usage.ts +++ b/packages/core/src/tool/plugin/route-usage.ts @@ -10,7 +10,7 @@ import { ModelUsage } from "../../model-usage.js" export const name = "route_usage" -export const description = `Read what routed work has used: totals per provider and model, per session (requests, tokens, cost, the pools and models it used), per day, and each plan's quota windows with the share used in its last window. +export const description = `Read what routed work has used: totals per provider and model, per session (requests, tokens, cost, the pools and models it used), per day, per continued turn (the chunks of one response joined back together), and each plan's quota windows with the share used in its last window. Use it for questions like "which sessions used the most", "did any pool use this provider", "how much of my Claude window is gone", or "what did this week cost". Set \`hours\` for the window (default 168; up to a year). Set \`provider\` or \`model\` to narrow the totals, \`route\` to one pool, or \`session\` to one session. Every target lists each pool it belongs to: used, offered but never chosen, or configured but never reached. The history is kept indefinitely, so a longer window is always available.` @@ -28,6 +28,7 @@ export const Output = Schema.Struct({ targets: Schema.Array(Schema.Unknown), sessions: Schema.Array(Schema.Unknown), days: Schema.Array(Schema.Unknown), + turns: Schema.Array(Schema.Unknown), plans: Schema.Array(Schema.Unknown), noDatabase: Schema.Boolean, }) @@ -91,6 +92,7 @@ export const collect = (input: Input, source: Source = NO_SOURCE, options: Optio } const targets = yield* ModelRouteLog.targetUsage(db, since) const days = yield* ModelRouteLog.dailyUsage(db, since) + const turns = yield* ModelRouteLog.turnAttempts(db, since, input.session) const offers = (yield* ModelRouteLog.offerUsage(db, since, input.route)) as ModelUsage.OfferRow[] const windowTotals: { plan: string; window: string; totals: ModelUsage.Totals }[] = [] for (const { plan, window } of windows) { @@ -102,7 +104,7 @@ export const collect = (input: Input, source: Source = NO_SOURCE, options: Optio totals: ModelUsage.sum(rows.filter((row) => row.providerID === plan.providerID)), }) } - return { sessions, sessionTargets, targets, days, offers, windowTotals, windowSessionTargets } + return { sessions, sessionTargets, targets, days, turns, offers, windowTotals, windowSessionTargets } }).pipe(Effect.orDie), ) const history = Option.getOrUndefined(read) @@ -142,6 +144,11 @@ export const collect = (input: Input, source: Source = NO_SOURCE, options: Optio session.targets.some((item) => matches(ModelUsage.splitTarget(item.target))), ) const days = ModelUsage.dayReports((history?.days ?? []).filter(matches)) + // A turn is a root attempt plus every chunk that continued it. Only continued turns are worth listing here: + // a single-request turn is already a row of the target, session and day totals. + const turns = ModelUsage.turns(history?.turns ?? []) + .filter((turn) => turn.chunks > 1) + .slice(0, input.limit ?? 20) const planReports = windows.map(({ plan, window }) => { const found = history?.windowTotals.find((item) => item.plan === plan.id && item.window === window.name) const report = ModelUsage.windowReport(plan, window, found?.totals ?? ModelUsage.zero) @@ -159,6 +166,7 @@ export const collect = (input: Input, source: Source = NO_SOURCE, options: Optio targets, sessions, days, + turns, plans: planReports, } }) @@ -170,7 +178,9 @@ const describe = (report: Effect.Success>) => { if (report.noDatabase) lines.push("Routing history is unavailable: no database is attached.") if (report.planError) lines.push(`Plan records could not be read: ${report.planError}`) lines.push(`Window: the last ${report.hours} hours, from ${report.since}.`) - lines.push(`${report.targets.length} targets, ${report.sessions.length} sessions listed, ${report.days.length} days.`) + lines.push( + `${report.targets.length} targets, ${report.sessions.length} sessions listed, ${report.turns.length} continued turns, ${report.days.length} days.`, + ) for (const target of report.targets) { const pools = target.pools.map( (pool) => `${pool.routeID} ${pool.state} (offered ${pool.offered}, chosen ${pool.chosen})`, @@ -190,6 +200,12 @@ const describe = (report: Effect.Success>) => { for (const day of report.days.slice(-14)) { lines.push(`${day.day}: ${day.requests} requests, ${day.metered} metered tokens, ${money(day.costUsd)}`) } + // Continued responses: the tokens and cost of every chunk of the turn, not one line per chunk. + for (const turn of report.turns.slice(0, 10)) { + lines.push( + `turn ${turn.rootID}${turn.sessionID === null ? "" : ` (session ${turn.sessionID})`}: ${turn.chunks} chunks, ${turn.requests} requests, ${ModelUsage.metered(turn)} metered tokens, ${money(turn.costUsd)}, ${(turn.durationMs / 1_000).toFixed(1)}s${turn.failures > 0 ? `, ${turn.failures} failed` : ""}`, + ) + } for (const plan of report.plans) { const limit = plan.limit === null ? "no limit recorded" : `of ${plan.limit} ${plan.unit}` lines.push( @@ -214,6 +230,7 @@ export const add = (editor: ToolEditor, source: () => Source = () => NO_SOURCE, targets: report.targets, sessions: report.sessions, days: report.days, + turns: report.turns, plans: report.plans, noDatabase: report.noDatabase, }, diff --git a/packages/core/test/model-route-log.test.ts b/packages/core/test/model-route-log.test.ts index a431d41a51a1..e949343d95aa 100644 --- a/packages/core/test/model-route-log.test.ts +++ b/packages/core/test/model-route-log.test.ts @@ -5,11 +5,14 @@ import { ModelRoute } from "@opencode/core/model-route" import { Permission } from "@opencode/core/permission" import { ModelRouteLog } from "@opencode/core/model-route-log" import { ModelRouteReview } from "@opencode/core/model-route-review" +import { ModelUsage } from "@opencode/core/model-usage" import { AIError, AuthenticationError, QuotaExceededError, RateLimitError } from "@opencode/ai" import { RouteStatsTool } from "@opencode/core/tool/plugin/route-stats" +import { RouteUsageTool } from "@opencode/core/tool/plugin/route-usage" import { RouteAttemptTable, RouteDecisionTable, RouteHealthTable } from "@opencode/core/model-route-log/sql" import { LayerNode } from "@opencode/util/effect/layer-node" import { testEffect } from "./lib/effect" +import { withTempDir } from "./fixture/tmpdir" const it = testEffect(LayerNode.compile(LayerNode.group([ModelRouteLog.node, Database.node]))) @@ -232,6 +235,187 @@ describe("usage history", () => { ) }) +describe("continuation attempts", () => { + it.live("links each chunk of a continued response to the attempt that started it", () => + Effect.gen(function* () { + const db = (yield* Database.Service).db + ModelRouteLog.record({ + kind: "attempt", + row: attempt({ + session_id: "ses_continued", + assistant_message_id: "msg_root", + time_started: 1_000, + time_ended: 2_000, + tokens_input: 1_000, + tokens_output: 100, + cost_usd: 0.1, + }), + }) + ModelRouteLog.record({ + kind: "attempt", + continuationOfMessageID: "msg_root", + row: attempt({ + session_id: "ses_continued", + assistant_message_id: "msg_chunk_1", + continuation_index: 1, + time_started: 3_000, + time_ended: 4_000, + tokens_input: 1_100, + tokens_output: 80, + cost_usd: 0.08, + }), + }) + // A continuation that fails keeps its own outcome and reason, and still belongs to the turn. + ModelRouteLog.record({ + kind: "attempt", + continuationOfMessageID: "msg_root", + row: attempt({ + session_id: "ses_continued", + assistant_message_id: "msg_chunk_2", + continuation_index: 2, + outcome: "failure", + output_started: false, + time_started: 5_000, + time_ended: 6_000, + tokens_input: 1_200, + cost_usd: 0.02, + ...ModelRouteLog.failureFields({ _tag: "QuotaExceeded", message: "You exceeded your current quota" }), + }), + }) + yield* Effect.sleep("100 millis") + + const rows = yield* db.select().from(RouteAttemptTable).orderBy(RouteAttemptTable.id).all() + const root = rows.find((row) => row.assistant_message_id === "msg_root") + expect(root).toMatchObject({ continuation_of: null, continuation_index: 0 }) + expect(rows.find((row) => row.assistant_message_id === "msg_chunk_1")).toMatchObject({ + continuation_of: root?.id, + continuation_index: 1, + cost_usd: 0.08, + }) + expect(rows.find((row) => row.assistant_message_id === "msg_chunk_2")).toMatchObject({ + continuation_of: root?.id, + continuation_index: 2, + outcome: "failure", + error_tag: "QuotaExceeded", + }) + + // The per-session totals already count every chunk; the turn view gives the same usage as one response. + const session = (yield* ModelRouteLog.sessionUsage(db, 0, 50)).find((row) => row.sessionID === "ses_continued") + expect(session).toMatchObject({ requests: 3, failures: 1, input: 3_300, output: 180 }) + const [turn] = ModelUsage.turns(yield* ModelRouteLog.turnAttempts(db, 0, "ses_continued")) + expect(turn).toMatchObject({ + rootID: root?.id, + chunks: 3, + requests: 3, + failures: 1, + input: 3_300, + output: 180, + first: 1_000, + last: 6_000, + durationMs: 5_000, + }) + expect(turn?.costUsd).toBeCloseTo(0.2) + }), + ) + + it.live("links to the attempt that answered when the first chunk was retried", () => + Effect.gen(function* () { + const db = (yield* Database.Service).db + ModelRouteLog.record({ + kind: "attempt", + row: attempt({ + session_id: "ses_retried", + assistant_message_id: "msg_retried", + outcome: "failure", + output_started: false, + time_started: 1_000, + time_ended: 2_000, + }), + }) + ModelRouteLog.record({ + kind: "attempt", + row: attempt({ + session_id: "ses_retried", + assistant_message_id: "msg_retried", + time_started: 2_100, + time_ended: 3_000, + tokens_input: 500, + tokens_output: 40, + cost_usd: 0.05, + }), + }) + ModelRouteLog.record({ + kind: "attempt", + continuationOfMessageID: "msg_retried", + row: attempt({ + session_id: "ses_retried", + assistant_message_id: "msg_retried_chunk", + continuation_index: 1, + time_started: 4_000, + time_ended: 5_000, + tokens_input: 600, + tokens_output: 30, + cost_usd: 0.03, + }), + }) + yield* Effect.sleep("100 millis") + + const rows = yield* db.select().from(RouteAttemptTable).orderBy(RouteAttemptTable.id).all() + const answered = rows.find((row) => row.time_started === 2_100) + const chunk = rows.find((row) => row.assistant_message_id === "msg_retried_chunk") + expect(chunk?.continuation_of).toBe(answered?.id) + expect(chunk?.continuation_of).not.toBe(rows.find((row) => row.time_started === 1_000)?.id) + // The retry that produced no response still belongs to the turn it was part of. + const turns = ModelUsage.turns(yield* ModelRouteLog.turnAttempts(db, 0, "ses_retried")) + expect(turns).toHaveLength(1) + expect(turns[0]).toMatchObject({ rootID: answered?.id, chunks: 2, requests: 3, failures: 1 }) + }), + ) + + it.live("route_usage lists a continued response as one turn", () => + withTempDir((tmp) => + Effect.gen(function* () { + const now = Date.now() + ModelRouteLog.record({ + kind: "attempt", + row: attempt({ + session_id: "ses_turn_report", + assistant_message_id: "msg_report_root", + time_started: now - 3_000, + time_ended: now - 2_000, + tokens_input: 100, + tokens_output: 40, + cost_usd: 0.04, + }), + }) + ModelRouteLog.record({ + kind: "attempt", + continuationOfMessageID: "msg_report_root", + row: attempt({ + session_id: "ses_turn_report", + assistant_message_id: "msg_report_chunk", + continuation_index: 1, + time_started: now - 1_000, + time_ended: now, + tokens_input: 200, + tokens_output: 20, + cost_usd: 0.02, + }), + }) + yield* Effect.sleep("100 millis") + + const report = yield* RouteUsageTool.collect({ hours: 1, session: "ses_turn_report" }, undefined, { + dataDir: tmp.path, + now: () => now, + }) + expect(report.turns).toHaveLength(1) + expect(report.turns[0]).toMatchObject({ chunks: 2, requests: 2, input: 300, output: 60, durationMs: 3_000 }) + expect(report.sessions[0]).toMatchObject({ requests: 2, input: 300, output: 60 }) + }), + ), + ) +}) + describe("route note restore", () => { it.live("restores active adjustments, drops expired and removed ones", () => Effect.gen(function* () { diff --git a/packages/core/test/model-usage.test.ts b/packages/core/test/model-usage.test.ts index a183ffd59438..b808ef58a160 100644 --- a/packages/core/test/model-usage.test.ts +++ b/packages/core/test/model-usage.test.ts @@ -14,6 +14,24 @@ const totals = (overrides: Partial = {}): ModelUsage.Totals = ...overrides, }) +const attemptRow = (overrides: Partial): ModelUsage.TurnAttemptRow => ({ + id: 1, + continuationOf: null, + continuationIndex: 0, + assistantMessageID: "msg_root", + sessionID: "ses_turn", + timeStarted: 1_000, + timeEnded: 2_000, + outcome: "success", + input: 100, + output: 50, + reasoning: 0, + cacheRead: 0, + cacheWrite: 0, + costUsd: 0.01, + ...overrides, +}) + describe("ModelUsage", () => { test("metered tokens are input, output and reasoning, with cache kept apart", () => { const row = totals({ input: 100, output: 50, reasoning: 10, cacheRead: 999, cacheWrite: 5 }) @@ -216,6 +234,127 @@ describe("ModelUsage", () => { expect(serialised["constructor"]).toEqual({ requests: 5, metered: 150, costUsd: 2 }) }) + test("turns stitch a response split by output token limit continuations into one cost", () => { + const turns = ModelUsage.turns([ + attemptRow({ id: 1, timeStarted: 1_000, timeEnded: 2_000, input: 1_000, output: 100, costUsd: 0.1 }), + attemptRow({ + id: 2, + continuationOf: 1, + continuationIndex: 1, + assistantMessageID: "msg_chunk_1", + timeStarted: 3_000, + timeEnded: 4_000, + input: 1_100, + output: 80, + costUsd: 0.08, + }), + // A failed chunk is still part of the turn: its own failure stays on its row, and its cost counts. + attemptRow({ + id: 3, + continuationOf: 1, + continuationIndex: 2, + assistantMessageID: "msg_chunk_2", + outcome: "failure", + timeStarted: 5_000, + timeEnded: 6_000, + input: 1_200, + output: 0, + costUsd: 0.02, + }), + attemptRow({ + id: 4, + assistantMessageID: "msg_other", + sessionID: "ses_other", + timeStarted: 7_000, + timeEnded: 8_000, + }), + ]) + + expect(turns).toHaveLength(2) + const continued = turns.find((turn) => turn.rootID === 1) + expect(continued).toMatchObject({ + rootID: 1, + sessionID: "ses_turn", + chunks: 3, + requests: 3, + failures: 1, + input: 3_300, + output: 180, + first: 1_000, + last: 6_000, + durationMs: 5_000, + }) + expect(continued?.costUsd).toBeCloseTo(0.2) + expect(turns.find((turn) => turn.rootID === 4)).toMatchObject({ chunks: 1, requests: 1, failures: 0 }) + }) + + test("a continuation whose root fell outside the window is still grouped by the id it names", () => { + const [turn] = ModelUsage.turns([ + { + id: 5, + continuationOf: 9, + continuationIndex: 1, + assistantMessageID: "msg_partial", + sessionID: "ses_turn", + timeStarted: 1_000, + timeEnded: 2_000, + outcome: "success", + input: 100, + output: 50, + reasoning: 0, + cacheRead: 0, + cacheWrite: 0, + costUsd: 0.01, + }, + ]) + expect(turn).toMatchObject({ rootID: 9, chunks: 2, requests: 1, costUsd: 0.01 }) + }) + + test("retries and hedge legs of the first chunk join its turn instead of splitting it", () => { + const turns = ModelUsage.turns([ + attemptRow({ + id: 1, + assistantMessageID: "msg_retried", + outcome: "failure", + timeStarted: 1_000, + timeEnded: 1_500, + input: 400, + output: 0, + costUsd: 0, + }), + attemptRow({ + id: 2, + assistantMessageID: "msg_retried", + timeStarted: 2_000, + timeEnded: 3_000, + input: 500, + output: 40, + costUsd: 0.05, + }), + attemptRow({ + id: 3, + assistantMessageID: "msg_retried_chunk", + continuationOf: 2, + continuationIndex: 1, + timeStarted: 4_000, + timeEnded: 5_000, + input: 600, + output: 30, + costUsd: 0.03, + }), + ]) + expect(turns).toHaveLength(1) + expect(turns[0]).toMatchObject({ + rootID: 2, + chunks: 2, + requests: 3, + failures: 1, + input: 1_500, + output: 70, + durationMs: 4_000, + }) + }) + test("a plan window reports what it has used, its percent and what remains", () => { const plan = { id: "anthropic-max", diff --git a/packages/core/test/session-runner.test.ts b/packages/core/test/session-runner.test.ts index 164fa3647544..f6c7c61ad24d 100644 --- a/packages/core/test/session-runner.test.ts +++ b/packages/core/test/session-runner.test.ts @@ -76,6 +76,7 @@ import { McpInstructions } from "@opencode/core/mcp/instructions" import { SessionSystemPrompt } from "@opencode/core/session/system-prompt" import { ID, Model } from "@opencode/core/model" import { ModelRoute } from "@opencode/core/model-route" +import { ModelRouteLog } from "@opencode/core/model-route-log" import { Location } from "@opencode/core/location" import { Provider } from "@opencode/core/provider" import { Cause, Context, DateTime, Deferred, Effect, Exit, Fiber, Layer, Queue, Schema, Scope, Stream } from "effect" @@ -5773,6 +5774,41 @@ describe("SessionRunnerLLM", () => { }) } + scenario("links output token limit continuations to the attempt that started the response", function* (s) { + const recorded: ModelRouteLog.Event[] = [] + ModelRouteLog.setSink((event) => recorded.push(event)) + yield* Effect.addFinalizer(() => Effect.sync(() => ModelRouteLog.setSink(undefined))) + const truncated = (id: string, text: string) => + TestLLM.complete( + { reason: { normalized: "length" } }, + LLMEvent.textStart({ id }), + LLMEvent.textDelta({ id, text }), + LLMEvent.textEnd({ id }), + ) + yield* s.llm.push( + truncated("partial", "Partial"), + truncated("more", "More"), + TestLLM.text("Finished", "finished"), + ) + + yield* s.runPrompt("Complete the response") + + const attempts = recorded.flatMap((event) => (event.kind === "attempt" ? [event] : [])) + expect(attempts).toHaveLength(3) + const [root, first, second] = attempts + const rootMessageID = root?.row.assistant_message_id ?? undefined + expect(root?.continuationOfMessageID).toBeUndefined() + expect(root?.row.continuation_index).toBeUndefined() + expect(rootMessageID).toBeDefined() + expect(first?.continuationOfMessageID).toBe(rootMessageID) + expect(first?.row.continuation_index).toBe(1) + expect(first?.row.assistant_message_id).not.toBe(rootMessageID) + // The second chunk names the response's first attempt, not the chunk before it. + expect(second?.continuationOfMessageID).toBe(rootMessageID) + expect(second?.row.continuation_index).toBe(2) + expect(second?.row.assistant_message_id).not.toBe(first?.row.assistant_message_id) + }) + for (const tools of ["hosted", "unfinished"] as const) { scenario(`caps output token limit continuations and resets for new input (${tools})`, function* (s) { const toolEvents = From 062ac897e83966ad67323ae75bea8f785e7c3ad7 Mon Sep 17 00:00:00 2001 From: Drewry Pope Date: Wed, 28 Oct 2026 16:30:42 -0500 Subject: [PATCH 2/3] fix(core): keep the continuation link across a tool-driven step A truncated response that also called a local tool ends its step so the tool can settle, and the response continues in the next step. Carry the response's root attempt and chunk index into that step, and drop it when newly promoted input starts a separate response. Also narrow the route_usage turn list by the provider and model filters the other reports apply. --- packages/core/src/model-route-log.ts | 2 ++ packages/core/src/model-usage.ts | 20 +++++++---- packages/core/src/session/runner/llm.ts | 35 +++++++++++++++----- packages/core/src/session/runner/step.ts | 5 ++- packages/core/src/tool/plugin/route-usage.ts | 2 +- packages/core/test/model-route-log.test.ts | 7 ++++ packages/core/test/model-usage.test.ts | 4 +++ packages/core/test/session-runner.test.ts | 30 +++++++++++++++++ 8 files changed, 89 insertions(+), 16 deletions(-) diff --git a/packages/core/src/model-route-log.ts b/packages/core/src/model-route-log.ts index 37b881ad824b..e0e5517defef 100644 --- a/packages/core/src/model-route-log.ts +++ b/packages/core/src/model-route-log.ts @@ -665,6 +665,8 @@ export const turnAttempts = (db: Database.Interface["db"], since: number, sessio continuationIndex: RouteAttemptTable.continuation_index, assistantMessageID: RouteAttemptTable.assistant_message_id, sessionID: RouteAttemptTable.session_id, + providerID: RouteAttemptTable.provider_id, + modelID: RouteAttemptTable.model_id, timeStarted: RouteAttemptTable.time_started, timeEnded: RouteAttemptTable.time_ended, outcome: RouteAttemptTable.outcome, diff --git a/packages/core/src/model-usage.ts b/packages/core/src/model-usage.ts index 705f02de2119..ca8bb8e7ea50 100644 --- a/packages/core/src/model-usage.ts +++ b/packages/core/src/model-usage.ts @@ -246,6 +246,8 @@ export interface TurnAttemptRow { readonly continuationIndex: number | null readonly assistantMessageID: string | null readonly sessionID: string | null + readonly providerID: string + readonly modelID: string readonly timeStarted: number readonly timeEnded: number readonly outcome: string @@ -261,6 +263,9 @@ export interface TurnRow extends Totals { /** Id of the attempt that started the turn; for a partial turn whose root fell outside the window, the id it names. */ readonly rootID: number readonly sessionID: string | null + /** The target the response started on; a failover mid-turn does not move the turn to another model. */ + readonly providerID: string + readonly modelID: string /** Response segments: the root attempt plus each output-token-limit continuation, however many requests each took. */ readonly chunks: number readonly first: number @@ -282,16 +287,17 @@ export const turns = (rows: readonly TurnAttemptRow[]): TurnRow[] => { } const rootOf = (row: TurnAttemptRow) => row.continuationOf ?? (row.assistantMessageID === null ? row.id : (roots.get(row.assistantMessageID) ?? row.id)) - const grouped = new Map() + const grouped = new Map() for (const row of rows) { const rootID = rootOf(row) const group = grouped.get(rootID) - if (group) group.push(row) - else grouped.set(rootID, [row]) + if (group) group.attempts.push(row) + else grouped.set(rootID, { lead: row, attempts: [row] }) } return [...grouped] - .map(([rootID, attempts]) => { - const root = attempts.find((attempt) => attempt.id === rootID) + .map(([rootID, { lead, attempts }]) => { + // The lead row stands in when the root fell outside the window, so the turn keeps a target to report. + const root = attempts.find((attempt) => attempt.id === rootID) ?? lead const totals = sum( attempts.map((attempt) => ({ requests: 1, @@ -309,7 +315,9 @@ export const turns = (rows: readonly TurnAttemptRow[]): TurnRow[] => { return { ...totals, rootID, - sessionID: root?.sessionID ?? attempts[0]?.sessionID ?? null, + sessionID: root.sessionID, + providerID: root.providerID, + modelID: root.modelID, chunks: 1 + Math.max(...attempts.map((attempt) => attempt.continuationIndex ?? 0)), first, last, diff --git a/packages/core/src/session/runner/llm.ts b/packages/core/src/session/runner/llm.ts index 87f8005d1f3b..b4db7d9ff40e 100644 --- a/packages/core/src/session/runner/llm.ts +++ b/packages/core/src/session/runner/llm.ts @@ -38,6 +38,9 @@ const CONTINUE_AFTER_INCOMPLETE_STREAM = const CONTINUE_AFTER_OUTPUT_LIMIT = "Your last response hit the output token limit (stop reason: length). Do not apologize, recap, or repeat yourself. Break the remaining work into smaller pieces." +/** How a step ended: whether the turn needs another step, and the link that step's response continues. */ +type StepResult = { readonly needsContinuation: boolean; readonly continuation?: SessionStep.ContinuationLink } + const layer = Layer.effect( Service, Effect.gen(function* () { @@ -59,6 +62,9 @@ const layer = Layer.effect( let force = input.force let continuing = input.continuation !== undefined let step = input.continuation?.step ?? 1 + // A step that hit the output token limit with tools left to run ends and continues in the next step; the + // response's root and chunk index carry over until new input is promoted. + let stepContinuation: SessionStep.ContinuationLink | undefined let entering = true const promotable = input.promotable ?? "input" if (!force && !continuing) { @@ -179,7 +185,7 @@ const layer = Layer.effect( onlyIfMissing: true, }) if (promoted > 0) step = 1 - return { _tag: "Ready" as const, context: yield* context.load(selected) } + return { _tag: "Ready" as const, context: yield* context.load(selected), promoted } }), ) if (ready) return ready @@ -191,7 +197,11 @@ const layer = Layer.effect( while (true) { const next = yield* advanceToStep() if (next._tag !== "Ready") return next - continuing = yield* runStep(next.context, step) + // Promoted input starts a new response; only a tool-driven step continuation keeps the link. + if (next.promoted > 0) stepContinuation = undefined + const completed = yield* runStep(next.context, step, stepContinuation) + continuing = completed.needsContinuation + stepContinuation = completed.continuation step++ force = false entering = false @@ -206,13 +216,18 @@ const layer = Layer.effect( }) /** Owns logical Step policy; each attempt owns its streaming, tools, and durable settlement. */ - const runStep = Effect.fn("SessionRunner.runStep")(function* (first: SessionContext.Loaded, step: number) { + const runStep = Effect.fn("SessionRunner.runStep")(function* ( + first: SessionContext.Loaded, + step: number, + continuation?: SessionStep.ContinuationLink, + ) { const sessionID = first.session.id let assistantMessageID = SessionMessage.ID.create() // Set once this step's response has been continued after an output token limit. Every later chunk names the - // attempt that started the response, so a turn's tokens and cost can be read as one figure. - let continuationOf: SessionMessage.ID | undefined - let continuationIndex = 0 + // attempt that started the response, so a turn's tokens and cost can be read as one figure. A step that + // continues a previous step's truncated response starts from that link. + let continuationOf: SessionMessage.ID | undefined = continuation?.of + let continuationIndex = continuation?.index ?? 0 const retry = yield* SessionRunnerRetry.make(bus, sessionID) let initial: SessionContext.Loaded | undefined = first let recoverOverflow = true @@ -285,7 +300,7 @@ const layer = Layer.effect( ), }) const completed = yield* SessionStep.Outcome.$match(outcome, { - Completed: (outcome) => Effect.succeed(outcome.needsContinuation), + Completed: (outcome) => Effect.succeed({ needsContinuation: outcome.needsContinuation }), Retry: (outcome) => retry.wait({ decision: outcome.decision, @@ -333,7 +348,11 @@ const layer = Layer.effect( message: "Response still truncated after two output token limit continuations", }, }) - if (outcome.needsContinuation) return true + if (outcome.needsContinuation) + return { + needsContinuation: true, + continuation: { of: continuationOf ?? assistantMessageID, index: continuationIndex + 1 }, + } yield* bus.publish(SessionEvent.Synthetic, { sessionID, text: CONTINUE_AFTER_OUTPUT_LIMIT }) continuationOf ??= assistantMessageID continuationIndex++ diff --git a/packages/core/src/session/runner/step.ts b/packages/core/src/session/runner/step.ts index ebebde049794..cb2f25cb9b6d 100644 --- a/packages/core/src/session/runner/step.ts +++ b/packages/core/src/session/runner/step.ts @@ -50,13 +50,16 @@ export type Outcome = Data.TaggedEnum<{ }> export const Outcome = Data.taggedEnum() +/** The attempt that started a response continued after an output token limit, and which chunk this request is. */ +export type ContinuationLink = { readonly of: SessionMessage.ID; readonly index: number } + interface Input { readonly isLocationClosed: () => boolean readonly sessionID: SessionSchema.ID readonly assistantMessageID: SessionMessage.ID /** Set when this request continues a response that hit the output token limit: the assistant message of the * attempt that started the response, and this chunk's index (1 for the first continuation). */ - readonly continuation?: { readonly of: SessionMessage.ID; readonly index: number } + readonly continuation?: ContinuationLink readonly agent: Agent.ID readonly model: SessionRunnerModel.Resolved /** Prompt-size estimate for budget accounting when a hedge is cancelled before provider usage arrives. */ diff --git a/packages/core/src/tool/plugin/route-usage.ts b/packages/core/src/tool/plugin/route-usage.ts index 361b310c6be8..39c099909f99 100644 --- a/packages/core/src/tool/plugin/route-usage.ts +++ b/packages/core/src/tool/plugin/route-usage.ts @@ -147,7 +147,7 @@ export const collect = (input: Input, source: Source = NO_SOURCE, options: Optio // A turn is a root attempt plus every chunk that continued it. Only continued turns are worth listing here: // a single-request turn is already a row of the target, session and day totals. const turns = ModelUsage.turns(history?.turns ?? []) - .filter((turn) => turn.chunks > 1) + .filter((turn) => turn.chunks > 1 && matches(turn)) .slice(0, input.limit ?? 20) const planReports = windows.map(({ plan, window }) => { const found = history?.windowTotals.find((item) => item.plan === plan.id && item.window === window.name) diff --git a/packages/core/test/model-route-log.test.ts b/packages/core/test/model-route-log.test.ts index e949343d95aa..23d06c10e751 100644 --- a/packages/core/test/model-route-log.test.ts +++ b/packages/core/test/model-route-log.test.ts @@ -411,6 +411,13 @@ describe("continuation attempts", () => { expect(report.turns).toHaveLength(1) expect(report.turns[0]).toMatchObject({ chunks: 2, requests: 2, input: 300, output: 60, durationMs: 3_000 }) expect(report.sessions[0]).toMatchObject({ requests: 2, input: 300, output: 60 }) + + // A provider filter narrows the turn list like it narrows the target and day reports. + const filtered = yield* RouteUsageTool.collect({ hours: 1, provider: "azure" }, undefined, { + dataDir: tmp.path, + now: () => now, + }) + expect(filtered.turns).toEqual([]) }), ), ) diff --git a/packages/core/test/model-usage.test.ts b/packages/core/test/model-usage.test.ts index b808ef58a160..d92c69c37dc3 100644 --- a/packages/core/test/model-usage.test.ts +++ b/packages/core/test/model-usage.test.ts @@ -20,6 +20,8 @@ const attemptRow = (overrides: Partial): ModelUsage.T continuationIndex: 0, assistantMessageID: "msg_root", sessionID: "ses_turn", + providerID: "openai", + modelID: "gpt-6-luna", timeStarted: 1_000, timeEnded: 2_000, outcome: "success", @@ -296,6 +298,8 @@ describe("ModelUsage", () => { continuationIndex: 1, assistantMessageID: "msg_partial", sessionID: "ses_turn", + providerID: "openai", + modelID: "gpt-6-luna", timeStarted: 1_000, timeEnded: 2_000, outcome: "success", diff --git a/packages/core/test/session-runner.test.ts b/packages/core/test/session-runner.test.ts index f6c7c61ad24d..67d0df20063f 100644 --- a/packages/core/test/session-runner.test.ts +++ b/packages/core/test/session-runner.test.ts @@ -5809,6 +5809,36 @@ describe("SessionRunnerLLM", () => { expect(second?.row.assistant_message_id).not.toBe(first?.row.assistant_message_id) }) + scenario("keeps the continuation link when a truncated response runs a local tool first", function* (s) { + const recorded: ModelRouteLog.Event[] = [] + ModelRouteLog.setSink((event) => recorded.push(event)) + yield* Effect.addFinalizer(() => Effect.sync(() => ModelRouteLog.setSink(undefined))) + // A local tool that never settles forces the truncated step to end so the tool can be settled, and the + // response continues in the next step. + yield* s.llm.push( + TestLLM.complete( + { reason: { normalized: "length" } }, + LLMEvent.toolInputStart({ id: "call-incomplete", name: "echo" }), + LLMEvent.toolInputDelta({ id: "call-incomplete", name: "echo", text: '{"text":"partial' }), + LLMEvent.textStart({ id: "partial" }), + LLMEvent.textDelta({ id: "partial", text: "Partial" }), + LLMEvent.textEnd({ id: "partial" }), + ), + TestLLM.text("Finished", "finished"), + ) + + yield* s.runPrompt("Complete the response") + + const attempts = recorded.flatMap((event) => (event.kind === "attempt" ? [event] : [])) + expect(attempts).toHaveLength(2) + const [root, continuation] = attempts + const rootMessageID = root?.row.assistant_message_id ?? undefined + expect(root?.row.continuation_index).toBeUndefined() + expect(rootMessageID).toBeDefined() + expect(continuation?.continuationOfMessageID).toBe(rootMessageID) + expect(continuation?.row.continuation_index).toBe(1) + }) + for (const tools of ["hosted", "unfinished"] as const) { scenario(`caps output token limit continuations and resets for new input (${tools})`, function* (s) { const toolEvents = From 79b5de760a448cb20515cad3875d721e62c2c029 Mon Sep 17 00:00:00 2001 From: Drewry Pope Date: Thu, 29 Oct 2026 04:44:09 -0500 Subject: [PATCH 3/3] fix(core): carry the response continuation link across a Location handoff When a truncated response leaves a local tool to settle, the step ends with the link to the attempt that started the response and the next step continues it. If a move is delivered before that next step, drain returned `Moved` (or `Reloaded`) with `{ step }` only, so the resumed chunk was written without `continuation_of` and started a new turn, splitting the response's tokens, cost and duration across two turns in `ModelUsage.turns` and `route_usage`. `SessionRunner.Continuation` now carries the step's `SessionStep.ContinuationLink` and the drain restores it from its input, so both exits keep the link. The reset on promoted input is unchanged. --- packages/core/src/session/runner/index.ts | 7 +++- packages/core/src/session/runner/llm.ts | 17 +++++++-- packages/core/test/session-runner.test.ts | 45 +++++++++++++++++++++++ 3 files changed, 64 insertions(+), 5 deletions(-) diff --git a/packages/core/src/session/runner/index.ts b/packages/core/src/session/runner/index.ts index 493713fa6768..a98a3f5d5976 100644 --- a/packages/core/src/session/runner/index.ts +++ b/packages/core/src/session/runner/index.ts @@ -7,6 +7,7 @@ import { SessionSchema } from "../schema.js" import type { Promotable } from "../inbox.js" import type { AgentNotFoundError, MessageDecodeError, StepFailedError } from "../error.js" import { SessionRunnerModel } from "./model.js" +import type { SessionStep } from "./step.js" import type { Instructions } from "../../instructions/index.js" export type RunError = @@ -18,7 +19,11 @@ export type RunError = | Instructions.InitializationBlocked | FileSystem.DirectoryNotFoundError -export type Continuation = { readonly step: number } +export type Continuation = { + readonly step: number + /** The response link this step's continuation carries, restored when the drain resumes at a new Location. */ + readonly link?: SessionStep.ContinuationLink +} export type DrainResult = Data.TaggedEnum<{ Complete: {} diff --git a/packages/core/src/session/runner/llm.ts b/packages/core/src/session/runner/llm.ts index b4db7d9ff40e..ba4aa381ea5c 100644 --- a/packages/core/src/session/runner/llm.ts +++ b/packages/core/src/session/runner/llm.ts @@ -63,8 +63,8 @@ const layer = Layer.effect( let continuing = input.continuation !== undefined let step = input.continuation?.step ?? 1 // A step that hit the output token limit with tools left to run ends and continues in the next step; the - // response's root and chunk index carry over until new input is promoted. - let stepContinuation: SessionStep.ContinuationLink | undefined + // response's root and chunk index carry over until new input is promoted, including across a Location handoff. + let stepContinuation: SessionStep.ContinuationLink | undefined = input.continuation?.link let entering = true const promotable = input.promotable ?? "input" if (!force && !continuing) { @@ -83,7 +83,12 @@ const layer = Layer.effect( while (true) { if (lifecycle.isClosed()) { yield* restore(modelTransport.close(sessionID)) - return DrainResult.Reloaded({ force, continuation: continuing ? { step } : undefined }) + return DrainResult.Reloaded({ + force, + continuation: continuing + ? { step, ...(stepContinuation ? { link: stepContinuation } : {}) } + : undefined, + }) } // Location entry and idle boundaries allow queued controls, not necessarily queued prompts. const pending = yield* SessionInbox.serialized( @@ -117,7 +122,11 @@ const layer = Layer.effect( step = 1 } if (pending?.type === "move") - return DrainResult.Moved({ continuation: continuing ? { step } : undefined }) + return DrainResult.Moved({ + continuation: continuing + ? { step, ...(stepContinuation ? { link: stepContinuation } : {}) } + : undefined, + }) if (pending?.type === "compaction") { const compacted = yield* restore( Effect.gen(function* () { diff --git a/packages/core/test/session-runner.test.ts b/packages/core/test/session-runner.test.ts index 67d0df20063f..d83ca02aa1c7 100644 --- a/packages/core/test/session-runner.test.ts +++ b/packages/core/test/session-runner.test.ts @@ -5839,6 +5839,51 @@ describe("SessionRunnerLLM", () => { expect(continuation?.row.continuation_index).toBe(1) }) + scenario("keeps the continuation link when a move hands off before the continued step", function* (s) { + const recorded: ModelRouteLog.Event[] = [] + ModelRouteLog.setSink((event) => recorded.push(event)) + yield* Effect.addFinalizer(() => Effect.sync(() => ModelRouteLog.setSink(undefined))) + // The truncated step ends with the link and the next step continues the response, but a move delivered in + // between hands the drain off to a new Location first, so the link has to travel in the drain result. + yield* s.admit("Echo before moving") + yield* s.llm.push( + TestLLM.complete( + { reason: { normalized: "length" } }, + LLMEvent.toolCall({ id: "call-move", name: "echo", input: { text: "moving" } }), + LLMEvent.textStart({ id: "partial" }), + LLMEvent.textDelta({ id: "partial", text: "Partial" }), + LLMEvent.textEnd({ id: "partial" }), + ), + TestLLM.text("Continued", "continued-after-move"), + ) + const tools = yield* s.blockTools() + const run = yield* s.resume.pipe(Effect.forkChild) + yield* tools.started + yield* s.sessionInbox.admit({ + id: SessionMessage.ID.create(), + sessionID, + item: { + type: "move", + payload: { + location: Location.Ref.make({ directory: AbsolutePath.make("/project") }), + projectID: Project.ID.global, + }, + delivery: "steer", + }, + }) + yield* tools.release + yield* Fiber.join(run) + + const attempts = recorded.flatMap((event) => (event.kind === "attempt" ? [event] : [])) + expect(attempts).toHaveLength(2) + const [root, continuation] = attempts + const rootMessageID = root?.row.assistant_message_id ?? undefined + expect(root?.row.continuation_index).toBeUndefined() + expect(rootMessageID).toBeDefined() + expect(continuation?.continuationOfMessageID).toBe(rootMessageID) + expect(continuation?.row.continuation_index).toBe(1) + }) + for (const tools of ["hosted", "unfinished"] as const) { scenario(`caps output token limit continuations and resets for new input (${tools})`, function* (s) { const toolEvents =