diff --git a/dist/server.js b/dist/server.js index cefe93e..26d8606 100644 --- a/dist/server.js +++ b/dist/server.js @@ -152,12 +152,20 @@ var PendingAttemptSchema = Schema.Struct({ armNoProgress: Schema.Boolean, previousLastContinuationAt: Schema.NullOr(Schema.Number) }); +var UsageTrackerSchema = Schema.Struct({ + baseline: Schema.optionalWith(Schema.Unknown, { default: () => null }), + lastObserved: Schema.optionalWith(Schema.Unknown, { default: () => null }), + baseTokens: Schema.optionalWith(Schema.Unknown, { default: () => null }), + pendingBaseline: Schema.optionalWith(Schema.Unknown, { default: () => null }), + pendingBaseTokens: Schema.optionalWith(Schema.Unknown, { default: () => null }) +}); var GoalSchema = Schema.Struct({ sessionID: Schema.String, objective: Schema.String, status: Schema.Literal("active", "paused", "budgetLimited", "usageLimited", "complete", "unmet"), tokenBudget: NullableNumber, tokensUsed: Schema.Number, + usageTrackers: Schema.optionalWith(Schema.Record({ key: Schema.String, value: UsageTrackerSchema }), { default: () => ({}) }), timeUsedSeconds: Schema.Number, createdAt: Schema.Number, updatedAt: Schema.Number, @@ -315,12 +323,34 @@ function normalizeGoal(goal) { goal.maxAutoTurns = positiveIntegerOrNull(goal.maxAutoTurns); goal.maxDurationSeconds = positiveIntegerOrNull(goal.maxDurationSeconds); goal.tokenBudget = positiveIntegerOrNull(goal.tokenBudget); + goal.usageTrackers = normalizeUsageTrackers(goal.usageTrackers); goal.noProgressTokenThreshold = positiveIntegerOrNull(goal.noProgressTokenThreshold) ?? DEFAULT_NO_PROGRESS_TOKEN_THRESHOLD; goal.maxNoProgressTurns = positiveIntegerOrNull(goal.maxNoProgressTurns) ?? DEFAULT_MAX_NO_PROGRESS_TURNS; goal.budgetWrapupSent = goal.budgetWrapupSent === true; goal.stopReason ??= null; return goal; } +function normalizeUsageTrackers(trackers) { + const normalized = {}; + for (const [source, rawTracker] of Object.entries(trackers ?? {})) { + const tracker = rawTracker; + const baseline = nonNegativeIntegerOrNull(tracker?.baseline); + const lastObserved = nonNegativeIntegerOrNull(tracker?.lastObserved); + const baseTokens = nonNegativeIntegerOrNull(tracker?.baseTokens); + if (source && baseline != null && lastObserved != null && baseTokens != null && lastObserved >= baseline) { + const pendingBaseline = nonNegativeIntegerOrNull(tracker.pendingBaseline); + const pendingBaseTokens = nonNegativeIntegerOrNull(tracker.pendingBaseTokens); + normalized[source] = { + baseline, + lastObserved, + baseTokens, + pendingBaseline, + pendingBaseTokens: pendingBaseline == null ? null : pendingBaseTokens + }; + } + } + return normalized; +} function normalizePendingAttempt(attempt) { if (!attempt || typeof attempt !== "object") return null; @@ -365,6 +395,9 @@ function positiveIntegerOrNull(value) { function nonNegativeInteger(value, fallback) { return typeof value === "number" && Number.isSafeInteger(value) && value >= 0 ? value : fallback; } +function nonNegativeIntegerOrNull(value) { + return typeof value === "number" && Number.isSafeInteger(value) && value >= 0 ? value : null; +} function isClosed(status) { return status === "complete" || status === "unmet"; } @@ -444,6 +477,7 @@ async function createGoal(sessionID, objective, options) { status: normalizedOptions.initialStatus, tokenBudget: normalizedOptions.tokenBudget, tokensUsed: 0, + usageTrackers: {}, timeUsedSeconds: 0, createdAt: now, updatedAt: now, @@ -605,14 +639,66 @@ async function clearGoal(sessionID) { return existed; }); } -async function accountUsage(sessionID, tokensUsed) { +async function accountUsage(sessionID, tokensUsed, options) { return mutate((state) => { const goal = state.goals[sessionID]; if (!goal) return null; accountWallClock(goal); if (typeof tokensUsed === "number" && Number.isFinite(tokensUsed)) { - goal.tokensUsed = Math.max(goal.tokensUsed, Math.max(0, Math.ceil(tokensUsed))); + const observed = Math.max(0, Math.ceil(tokensUsed)); + if (options?.cumulative === true) { + const source = options.source?.trim() || "default"; + let tracker = goal.usageTrackers[source]; + if (!tracker) { + const initialBaseline = nonNegativeIntegerOrNull(options.initialBaseline); + tracker = initialBaseline != null && initialBaseline <= observed ? { + baseline: initialBaseline, + lastObserved: observed, + baseTokens: goal.tokensUsed, + pendingBaseline: null, + pendingBaseTokens: null + } : { + baseline: observed, + lastObserved: observed, + baseTokens: goal.tokensUsed, + pendingBaseline: null, + pendingBaseTokens: null + }; + goal.usageTrackers[source] = tracker; + } else if (observed < tracker.lastObserved) { + const initialBaseline = nonNegativeIntegerOrNull(options.initialBaseline); + if (initialBaseline != null && initialBaseline <= observed) { + tracker = { + baseline: initialBaseline, + lastObserved: observed, + baseTokens: goal.tokensUsed, + pendingBaseline: null, + pendingBaseTokens: null + }; + goal.usageTrackers[source] = tracker; + } else if (tracker.pendingBaseline == null || observed < tracker.pendingBaseline) { + tracker.pendingBaseline = observed; + tracker.pendingBaseTokens = goal.tokensUsed; + } else { + tracker = { + baseline: tracker.pendingBaseline, + lastObserved: observed, + baseTokens: tracker.pendingBaseTokens ?? goal.tokensUsed, + pendingBaseline: null, + pendingBaseTokens: null + }; + goal.usageTrackers[source] = tracker; + } + } else { + tracker.lastObserved = observed; + tracker.pendingBaseline = null; + tracker.pendingBaseTokens = null; + } + goal.tokensUsed = Math.max(goal.tokensUsed, tracker.baseTokens + observed - tracker.baseline); + } else { + goal.tokensUsed = Math.max(goal.tokensUsed, observed); + } } maybeStopForBudget(goal); goal.updatedAt = nowSeconds(); @@ -1097,7 +1183,7 @@ function commandNameFromOptions(options) { function positiveIntegerOrNull2(value) { return typeof value === "number" && Number.isSafeInteger(value) && value > 0 ? value : null; } -function nonNegativeIntegerOrNull(value) { +function nonNegativeIntegerOrNull2(value) { return typeof value === "number" && Number.isSafeInteger(value) && value >= 0 ? value : null; } function timeoutMillisecondsFromSeconds(value) { @@ -1190,9 +1276,9 @@ function outputTokensFromMessage(message) { return outputTokensFromRecord(message.info.tokens); return; } -function tokensFromMessages(messages) { +function usageFromMessages(messages) { const exactTotal = messages.reduce((sum, message) => sum + (exactTokensFromMessage(message) ?? 0), 0); - return exactTotal > 0 ? exactTotal : estimateMessages(messages); + return exactTotal > 0 ? { tokens: exactTotal, source: "v1.messages.exact" } : { tokens: estimateMessages(messages), source: "v1.messages.estimated" }; } function taskHeader(output) { const resultIndex = output.search(//); @@ -1696,6 +1782,7 @@ async function createGoalFromTool(input, context, services) { agent: typeof context.agent === "string" ? context.agent : null, initialStatus: planningOnly ? "paused" : "active" }); + await services.initializeUsage?.(context.sessionID); return JSON.stringify(planningOnly ? { goal, plan_mode_notice: PLAN_MODE_CREATE_NOTICE } : { goal }, null, 2); } async function updateGoalObjectiveFromTool(input, context, services) { @@ -1758,7 +1845,7 @@ var server = async ({ client }, options) => { const autoContinue = options?.auto_continue ?? true; const deferWhileTasksActive = options?.defer_while_tasks_active ?? true; const maxAutoTurns = positiveIntegerOrNull2(options?.max_auto_turns) ?? DEFAULT_MAX_AUTO_TURNS; - const minInterval = nonNegativeIntegerOrNull(options?.min_continue_interval_seconds) ?? DEFAULT_CONTINUE_INTERVAL_SECONDS; + const minInterval = nonNegativeIntegerOrNull2(options?.min_continue_interval_seconds) ?? DEFAULT_CONTINUE_INTERVAL_SECONDS; const maxTurnTimeMs = timeoutMillisecondsFromSeconds(options?.max_turn_time); const maxPromptFailures = positiveIntegerOrNull2(options?.max_prompt_failures) ?? DEFAULT_MAX_PROMPT_FAILURES; const registerCommand = options?.register_command ?? true; @@ -2167,7 +2254,8 @@ var server = async ({ client }, options) => { const sessionID = "sessionID" in input && typeof input.sessionID === "string" ? input.sessionID : output.messages.find((message) => typeof message.info.sessionID === "string")?.info.sessionID; if (!sessionID) return; - await accountUsage(sessionID, tokensFromMessages(output.messages)); + const usage = usageFromMessages(output.messages); + await accountUsage(sessionID, usage.tokens, { cumulative: true, source: usage.source }); const observed = await recordAssistantMessage(sessionID, latestAssistantMessage(output.messages), options ?? {}); await reconcileLocalMarkerAfterProgress(locallyDeliveredPendingSessions, sessionID, observed.goal); const scheduled = scheduledContinuations.get(sessionID); @@ -2298,7 +2386,7 @@ async function setupV2(context) { const autoContinue = options.auto_continue ?? true; const deferWhileTasksActive = options.defer_while_tasks_active ?? true; const maxAutoTurns = positiveIntegerOrNull2(options.max_auto_turns) ?? DEFAULT_MAX_AUTO_TURNS; - const minInterval = nonNegativeIntegerOrNull(options.min_continue_interval_seconds) ?? DEFAULT_CONTINUE_INTERVAL_SECONDS; + const minInterval = nonNegativeIntegerOrNull2(options.min_continue_interval_seconds) ?? DEFAULT_CONTINUE_INTERVAL_SECONDS; const maxTurnTimeMs = timeoutMillisecondsFromSeconds(options.max_turn_time); const maxPromptFailures = positiveIntegerOrNull2(options.max_prompt_failures) ?? DEFAULT_MAX_PROMPT_FAILURES; const registerCommand = options.register_command ?? true; @@ -2314,11 +2402,21 @@ async function setupV2(context) { const toolAttempts = new Map; const planAgents = restrictedAgentSet(options); const isPlanAgent = (agent) => typeof agent === "string" && planAgents.has(agent.trim().toLowerCase()); - const goalServices = { options, isPlanAgent }; const activeContinuationsV2 = new Set; const latestStepBySession = new Map; const stepTextBuffers = new Map; const stepTokenSums = new Map; + const goalServices = { + options, + isPlanAgent, + initializeUsage: async (sessionID) => { + try { + await accountUsage(sessionID, stepTokenSums.get(sessionID) ?? 0, { cumulative: true, source: "v2.steps" }); + } catch (error) { + v2ErrorLog("Failed to initialize goal usage accounting", error); + } + } + }; const registrations = []; let disposed = false; function stepKey(sessionID, messageID2) { @@ -2719,7 +2817,11 @@ async function setupV2(context) { if (typeof tokens === "number") { const sum = (stepTokenSums.get(sessionID) ?? 0) + tokens; stepTokenSums.set(sessionID, sum); - await accountUsage(sessionID, sum); + await accountUsage(sessionID, sum, { + cumulative: true, + source: "v2.steps", + initialBaseline: Math.ceil(sum - tokens) + }); } const text = stepTextBuffers.get(stepKey(sessionID, messageID2)) ?? ""; stepTextBuffers.delete(stepKey(sessionID, messageID2)); @@ -2755,7 +2857,11 @@ async function setupV2(context) { if (typeof tokens === "number") { const sum = (stepTokenSums.get(sessionID) ?? 0) + tokens; stepTokenSums.set(sessionID, sum); - await accountUsage(sessionID, sum); + await accountUsage(sessionID, sum, { + cumulative: true, + source: "v2.steps", + initialBaseline: Math.ceil(sum - tokens) + }); } const text = stepTextBuffers.get(stepKey(sessionID, messageID2)) ?? ""; stepTextBuffers.delete(stepKey(sessionID, messageID2)); @@ -2788,7 +2894,7 @@ async function setupV2(context) { return; const tokens = tokensFromRecord(data.tokens); if (typeof tokens === "number") - await accountUsage(sessionID, tokens); + await accountUsage(sessionID, tokens, { cumulative: true, source: "v2.session" }); return; } } diff --git a/src/server.ts b/src/server.ts index 73eb501..979f578 100644 --- a/src/server.ts +++ b/src/server.ts @@ -247,9 +247,11 @@ function outputTokensFromMessage(message: { info?: unknown; parts?: unknown[] }) return undefined } -function tokensFromMessages(messages: { info?: unknown; parts?: unknown[] }[]) { +function usageFromMessages(messages: { info?: unknown; parts?: unknown[] }[]) { const exactTotal = messages.reduce((sum, message) => sum + (exactTokensFromMessage(message) ?? 0), 0) - return exactTotal > 0 ? exactTotal : estimateMessages(messages) + return exactTotal > 0 + ? { tokens: exactTotal, source: "v1.messages.exact" } + : { tokens: estimateMessages(messages), source: "v1.messages.estimated" } } function taskHeader(output: string) { @@ -767,6 +769,7 @@ type ToolExecContext = { type GoalServices = { options: Options isPlanAgent: (agent: unknown) => boolean + initializeUsage?: (sessionID: string) => Promise } async function createGoalFromTool(input: CreateGoalArgs, context: ToolExecContext, services: GoalServices) { @@ -780,6 +783,7 @@ async function createGoalFromTool(input: CreateGoalArgs, context: ToolExecContex agent: typeof context.agent === "string" ? context.agent : null, initialStatus: planningOnly ? "paused" : "active", }) + await services.initializeUsage?.(context.sessionID) return JSON.stringify(planningOnly ? { goal, plan_mode_notice: PLAN_MODE_CREATE_NOTICE } : { goal }, null, 2) } @@ -1348,7 +1352,8 @@ const server: Plugin = async ({ client }, options?: Options) => { ? input.sessionID : output.messages.find((message) => typeof message.info.sessionID === "string")?.info.sessionID if (!sessionID) return - await accountUsage(sessionID, tokensFromMessages(output.messages)) + const usage = usageFromMessages(output.messages) + await accountUsage(sessionID, usage.tokens, { cumulative: true, source: usage.source }) const observed = await recordAssistantMessage(sessionID, latestAssistantMessage(output.messages), options ?? {}) await reconcileLocalMarkerAfterProgress(locallyDeliveredPendingSessions, sessionID, observed.goal) const scheduled = scheduledContinuations.get(sessionID) @@ -1511,11 +1516,21 @@ async function setupV2(context: PluginV2.Plugin.Context): Promise() const planAgents = restrictedAgentSet(options) const isPlanAgent = (agent: unknown) => typeof agent === "string" && planAgents.has(agent.trim().toLowerCase()) - const goalServices: GoalServices = { options, isPlanAgent } const activeContinuationsV2 = new Set() const latestStepBySession = new Map() const stepTextBuffers = new Map() const stepTokenSums = new Map() + const goalServices: GoalServices = { + options, + isPlanAgent, + initializeUsage: async (sessionID) => { + try { + await accountUsage(sessionID, stepTokenSums.get(sessionID) ?? 0, { cumulative: true, source: "v2.steps" }) + } catch (error) { + v2ErrorLog("Failed to initialize goal usage accounting", error) + } + }, + } const registrations: Array<{ dispose(): Promise }> = [] let disposed = false @@ -1940,7 +1955,11 @@ async function setupV2(context: PluginV2.Plugin.Context): Promise timeUsedSeconds: number createdAt: number updatedAt: number @@ -112,6 +113,14 @@ export type Goal = { continuationBaselineSummary: string } +type UsageTracker = { + baseline: number + lastObserved: number + baseTokens: number + pendingBaseline: number | null + pendingBaseTokens: number | null +} + type State = { version: 1 goals: Record @@ -169,12 +178,20 @@ const PendingAttemptSchema = Schema.Struct({ armNoProgress: Schema.Boolean, previousLastContinuationAt: Schema.NullOr(Schema.Number), }) +const UsageTrackerSchema = Schema.Struct({ + baseline: Schema.optionalWith(Schema.Unknown, { default: () => null }), + lastObserved: Schema.optionalWith(Schema.Unknown, { default: () => null }), + baseTokens: Schema.optionalWith(Schema.Unknown, { default: () => null }), + pendingBaseline: Schema.optionalWith(Schema.Unknown, { default: () => null }), + pendingBaseTokens: Schema.optionalWith(Schema.Unknown, { default: () => null }), +}) const GoalSchema = Schema.Struct({ sessionID: Schema.String, objective: Schema.String, status: Schema.Literal("active", "paused", "budgetLimited", "usageLimited", "complete", "unmet"), tokenBudget: NullableNumber, tokensUsed: Schema.Number, + usageTrackers: Schema.optionalWith(Schema.Record({ key: Schema.String, value: UsageTrackerSchema }), { default: () => ({}) }), timeUsedSeconds: Schema.Number, createdAt: Schema.Number, updatedAt: Schema.Number, @@ -214,7 +231,7 @@ const StateSchema = Schema.Struct({ // getGoalInternal / the internal snapshot type instead. export type GoalSnapshot = Omit< Goal, - "lastAccountedAt" | "autoTurns" | "lastContinuationAt" | "pendingAttempt" + "lastAccountedAt" | "autoTurns" | "lastContinuationAt" | "pendingAttempt" | "usageTrackers" > & { remainingTokens: number | null sampledAt: number @@ -405,6 +422,7 @@ function normalizeGoal(goal: Goal) { goal.maxAutoTurns = positiveIntegerOrNull(goal.maxAutoTurns) goal.maxDurationSeconds = positiveIntegerOrNull(goal.maxDurationSeconds) goal.tokenBudget = positiveIntegerOrNull(goal.tokenBudget) + goal.usageTrackers = normalizeUsageTrackers(goal.usageTrackers) goal.noProgressTokenThreshold = positiveIntegerOrNull(goal.noProgressTokenThreshold) ?? DEFAULT_NO_PROGRESS_TOKEN_THRESHOLD goal.maxNoProgressTurns = positiveIntegerOrNull(goal.maxNoProgressTurns) ?? DEFAULT_MAX_NO_PROGRESS_TURNS goal.budgetWrapupSent = goal.budgetWrapupSent === true @@ -412,6 +430,28 @@ function normalizeGoal(goal: Goal) { return goal } +function normalizeUsageTrackers(trackers: Record | undefined) { + const normalized: Record = {} + for (const [source, rawTracker] of Object.entries(trackers ?? {})) { + const tracker = rawTracker as Partial + const baseline = nonNegativeIntegerOrNull(tracker?.baseline) + const lastObserved = nonNegativeIntegerOrNull(tracker?.lastObserved) + const baseTokens = nonNegativeIntegerOrNull(tracker?.baseTokens) + if (source && baseline != null && lastObserved != null && baseTokens != null && lastObserved >= baseline) { + const pendingBaseline = nonNegativeIntegerOrNull(tracker.pendingBaseline) + const pendingBaseTokens = nonNegativeIntegerOrNull(tracker.pendingBaseTokens) + normalized[source] = { + baseline, + lastObserved, + baseTokens, + pendingBaseline, + pendingBaseTokens: pendingBaseline == null ? null : pendingBaseTokens, + } + } + } + return normalized +} + function normalizePendingAttempt(attempt: PendingAttempt | null | undefined): PendingAttempt | null { if (!attempt || typeof attempt !== "object") return null return { @@ -464,6 +504,10 @@ function nonNegativeInteger(value: unknown, fallback: number) { return typeof value === "number" && Number.isSafeInteger(value) && value >= 0 ? value : fallback } +function nonNegativeIntegerOrNull(value: unknown) { + return typeof value === "number" && Number.isSafeInteger(value) && value >= 0 ? value : null +} + function isClosed(status: GoalStatus) { return status === "complete" || status === "unmet" } @@ -557,6 +601,7 @@ export async function createGoal(sessionID: string, objective: string, options?: status: normalizedOptions.initialStatus, tokenBudget: normalizedOptions.tokenBudget, tokensUsed: 0, + usageTrackers: {}, timeUsedSeconds: 0, createdAt: now, updatedAt: now, @@ -736,13 +781,74 @@ export async function clearGoal(sessionID: string) { }) } -export async function accountUsage(sessionID: string, tokensUsed?: number) { +export async function accountUsage( + sessionID: string, + tokensUsed?: number, + options?: { cumulative?: boolean; source?: string; initialBaseline?: number }, +) { return mutate((state) => { const goal = state.goals[sessionID] if (!goal) return null accountWallClock(goal) if (typeof tokensUsed === "number" && Number.isFinite(tokensUsed)) { - goal.tokensUsed = Math.max(goal.tokensUsed, Math.max(0, Math.ceil(tokensUsed))) + const observed = Math.max(0, Math.ceil(tokensUsed)) + if (options?.cumulative === true) { + const source = options.source?.trim() || "default" + let tracker = goal.usageTrackers[source] + if (!tracker) { + const initialBaseline = nonNegativeIntegerOrNull(options.initialBaseline) + tracker = + initialBaseline != null && initialBaseline <= observed + ? { + baseline: initialBaseline, + lastObserved: observed, + baseTokens: goal.tokensUsed, + pendingBaseline: null, + pendingBaseTokens: null, + } + : { + baseline: observed, + lastObserved: observed, + baseTokens: goal.tokensUsed, + pendingBaseline: null, + pendingBaseTokens: null, + } + goal.usageTrackers[source] = tracker + } else if (observed < tracker.lastObserved) { + const initialBaseline = nonNegativeIntegerOrNull(options.initialBaseline) + if (initialBaseline != null && initialBaseline <= observed) { + tracker = { + baseline: initialBaseline, + lastObserved: observed, + baseTokens: goal.tokensUsed, + pendingBaseline: null, + pendingBaseTokens: null, + } + goal.usageTrackers[source] = tracker + } else if (tracker.pendingBaseline == null || observed < tracker.pendingBaseline) { + // Require a second consistent low observation before treating an + // un-signaled decrease as compaction rather than a partial sample. + tracker.pendingBaseline = observed + tracker.pendingBaseTokens = goal.tokensUsed + } else { + tracker = { + baseline: tracker.pendingBaseline, + lastObserved: observed, + baseTokens: tracker.pendingBaseTokens ?? goal.tokensUsed, + pendingBaseline: null, + pendingBaseTokens: null, + } + goal.usageTrackers[source] = tracker + } + } else { + tracker.lastObserved = observed + tracker.pendingBaseline = null + tracker.pendingBaseTokens = null + } + goal.tokensUsed = Math.max(goal.tokensUsed, tracker.baseTokens + observed - tracker.baseline) + } else { + goal.tokensUsed = Math.max(goal.tokensUsed, observed) + } } maybeStopForBudget(goal) goal.updatedAt = nowSeconds() diff --git a/test/server-v2.test.ts b/test/server-v2.test.ts index 87b0b3f..f8a81da 100644 --- a/test/server-v2.test.ts +++ b/test/server-v2.test.ts @@ -1,5 +1,5 @@ import { afterEach, beforeEach, expect, test } from "bun:test" -import { mkdtemp, rm, writeFile } from "node:fs/promises" +import { mkdtemp, readFile, rm, writeFile } from "node:fs/promises" import { join } from "node:path" import { tmpdir } from "node:os" import plugin from "../src/server" @@ -342,22 +342,104 @@ test("V2 events account usage and checkpoints from step/usage events", async () expect(contentOf(readAfterStep)).toContain('"tokensUsed": 70') expect(contentOf(readAfterStep)).toContain("IMPLEMENTED_THE_FEATURE") - // Cumulative usage.updated is authoritative and raises the accounted total. + // The first cumulative observation establishes a baseline without counting + // the session's pre-goal usage or replacing step-derived goal usage. mock.stream.push({ type: "session.usage.updated", created: Date.now(), data: { sessionID: "ses_v2", tokens: { input: 200, output: 50, reasoning: 0, cache: { read: 10, write: 0 } } }, }) - await waitFor(async () => contentOf(await goalTool(mock, "get_goal").execute({}, toolContext())).includes('"tokensUsed": 260')) + await waitFor(async () => { + const state = JSON.parse(await readFile(process.env.OPENCODE_GOAL_STATE_PATH!, "utf8")) as { + goals: Record }> + } + return JSON.stringify(state.goals.ses_v2?.usageTrackers?.["v2.session"]) === + JSON.stringify({ baseline: 260, lastObserved: 260, baseTokens: 70, pendingBaseline: null, pendingBaseTokens: null }) + }) + + mock.stream.push({ + type: "session.usage.updated", + created: Date.now(), + data: { sessionID: "ses_v2", tokens: { input: 215, output: 50, reasoning: 0, cache: { read: 10, write: 0 } } }, + }) + await waitFor(async () => contentOf(await goalTool(mock, "get_goal").execute({}, toolContext())).includes('"tokensUsed": 85')) const read = await goalTool(mock, "get_goal").execute({}, toolContext()) - expect(contentOf(read)).toContain('"tokensUsed": 260') + expect(contentOf(read)).toContain('"tokensUsed": 85') expect(contentOf(read)).toContain("IMPLEMENTED_THE_FEATURE") mock.stream.end() await cleanup() }) +test("V2 step accounting excludes steps observed before goal creation", async () => { + const mock = makeMockContext({ auto_continue: false }) + const cleanup = await plugin.setup(mock as never) + + mock.stream.push({ + type: "session.step.ended", + created: 1, + data: { + sessionID: "ses_v2", + assistantMessageID: "msg_before_goal", + tokens: { input: 50_000, output: 0, reasoning: 0, cache: { read: 0, write: 0 } }, + }, + }) + await new Promise((resolve) => setTimeout(resolve, 20)) + await goalTool(mock, "create_goal").execute( + { objective: "measure only goal work", token_budget: 1_000 }, + toolContext(), + ) + + mock.stream.push({ + type: "session.step.ended", + created: 2, + data: { + sessionID: "ses_v2", + assistantMessageID: "msg_goal_work", + tokens: { input: 200, output: 100, reasoning: 0, cache: { read: 0, write: 0 } }, + }, + }) + + await waitFor(async () => (await getGoal("ses_v2"))?.tokensUsed === 300) + expect(await getGoal("ses_v2")).toMatchObject({ status: "active", tokensUsed: 300 }) + + mock.stream.end() + await cleanup() +}) + +test("V2 step and session sources do not double-count when session usage arrives first", async () => { + const mock = makeMockContext({ auto_continue: false }) + const cleanup = await plugin.setup(mock as never) + await createGoalViaV2Tool(mock, "reconcile usage sources") + + mock.stream.push({ + type: "session.usage.updated", + created: 1, + data: { sessionID: "ses_v2", tokens: { input: 500_000, output: 0, reasoning: 0, cache: { read: 0, write: 0 } } }, + }) + mock.stream.push({ + type: "session.usage.updated", + created: 2, + data: { sessionID: "ses_v2", tokens: { input: 500_250, output: 0, reasoning: 0, cache: { read: 0, write: 0 } } }, + }) + mock.stream.push({ + type: "session.step.ended", + created: 3, + data: { + sessionID: "ses_v2", + assistantMessageID: "msg_goal_work", + tokens: { input: 200, output: 100, reasoning: 0, cache: { read: 0, write: 0 } }, + }, + }) + + await waitFor(async () => (await getGoal("ses_v2"))?.tokensUsed === 300) + expect((await getGoal("ses_v2"))?.tokensUsed).toBe(300) + + mock.stream.end() + await cleanup() +}) + test("V2 failed steps account usage and replace stale assistant progress", async () => { const mock = makeMockContext({ auto_continue: false }) const cleanup = await plugin.setup(mock as never) @@ -575,6 +657,19 @@ test("V2 successful tool progress cancels no-pending transport recovery", async const cleanup = await plugin.setup(mock as never) await createGoalViaV2Tool(mock, "cancel recovery via tool progress") + mock.stream.push({ + type: "session.usage.updated", + created: 0, + data: { sessionID: "ses_v2", tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } } }, + }) + await waitFor(async () => { + const state = JSON.parse(await readFile(process.env.OPENCODE_GOAL_STATE_PATH!, "utf8")) as { + goals: Record }> + } + return JSON.stringify(state.goals.ses_v2?.usageTrackers?.["v2.session"]) === + JSON.stringify({ baseline: 0, lastObserved: 0, baseTokens: 0, pendingBaseline: null, pendingBaseTokens: null }) + }) + mock.stream.push({ type: "session.error", created: 1, diff --git a/test/server.test.ts b/test/server.test.ts index 0743380..10f431b 100644 --- a/test/server.test.ts +++ b/test/server.test.ts @@ -222,7 +222,7 @@ OpenCode goal mode policy: ) const read = await requireTool(tools.get_goal, "get_goal").execute({}, context) expect(String(read)).toContain('"objective": "OBJECTIVE_SHOULD_NOT_LEAK_7f31"') - expect(String(read)).toContain('"tokensUsed": 460') + expect(String(read)).toContain('"tokensUsed": 0') expect(String(read)).toContain('"timeUsedSeconds": 5') expect(String(read)).toContain("CHECKPOINT_SHOULD_NOT_LEAK_4b72") await transform("ses_lifecycle") @@ -278,6 +278,17 @@ OpenCode goal mode policy: ], } as never, ) + await hooks["experimental.chat.messages.transform"]!( + {}, + { + messages: [ + { + info: { id: "msg_budget_2", role: "assistant", sessionID: "ses_budget" }, + parts: [{ type: "step-finish", tokens: { input: 17, output: 5 } }], + }, + ], + } as never, + ) const budgetLimited = await requireTool(tools.get_goal, "get_goal").execute({}, budgetContext) expect(String(budgetLimited)).toContain('"status": "budgetLimited"') expect(String(budgetLimited)).toContain("Do not start or continue substantive work") @@ -483,6 +494,12 @@ test("message transform prefers exact step token usage", async () => { const context = { sessionID: "ses_1" } as never await requireTool(tools.create_goal, "create_goal").execute({ objective: "finish" }, context) + await hooks["experimental.chat.messages.transform"]!( + { sessionID: "ses_1" } as never, + { + messages: [{ info: { sessionID: "ses_1" }, parts: [{ type: "step-finish", tokens: { input: 1, output: 0 } }] }], + } as never, + ) await hooks["experimental.chat.messages.transform"]!( {}, { @@ -492,7 +509,7 @@ test("message transform prefers exact step token usage", async () => { parts: [ { type: "step-finish", - tokens: { input: 10, output: 5, reasoning: 2, cache: { read: 3, write: 4 } }, + tokens: { input: 11, output: 5, reasoning: 2, cache: { read: 3, write: 4 } }, }, ], }, @@ -504,6 +521,40 @@ test("message transform prefers exact step token usage", async () => { expect(String(read)).toContain('"tokensUsed": 24') }) +test("message transform excludes session usage observed before goal work", async () => { + const hooks = await plugin.server( + { client: { session: { promptAsync: async () => {} } } } as never, + { auto_continue: false }, + ) + const tools = hooks.tool! + const context = { sessionID: "ses_1" } as never + await requireTool(tools.create_goal, "create_goal").execute({ objective: "finish", token_budget: 10 }, context) + + const transform = (total: number) => + hooks["experimental.chat.messages.transform"]!( + {}, + { + messages: [ + { + info: { sessionID: "ses_1" }, + parts: [{ type: "step-finish", tokens: { input: total, output: 0 } }], + }, + ], + } as never, + ) + + await transform(1_000) + expect(await getGoal("ses_1")).toMatchObject({ status: "active", tokensUsed: 0 }) + await hooks["experimental.chat.messages.transform"]!( + { sessionID: "ses_1" } as never, + { messages: [] } as never, + ) + await transform(1_005) + expect(await getGoal("ses_1")).toMatchObject({ status: "active", tokensUsed: 5 }) + const read = await requireTool(tools.get_goal, "get_goal").execute({}, context) + expect(String(read)).not.toContain("usageTrackers") +}) + test("per-prompt chat hook recovers from an empty state file", async () => { await writeFile(process.env.OPENCODE_GOAL_STATE_PATH!, "", "utf8") const hooks = await plugin.server( diff --git a/test/state.test.ts b/test/state.test.ts index 9083396..bf5e5e5 100644 --- a/test/state.test.ts +++ b/test/state.test.ts @@ -81,6 +81,152 @@ test("token usage marks goals budget limited", async () => { expect(updated?.stopReason).toContain("token budget reached") }) +test("cumulative usage establishes a private tracker and grows by its delta across state reloads", async () => { + await createGoal("ses_1", "measure goal usage", null) + await accountUsage("ses_1", 20) + + const first = await accountUsage("ses_1", 100, { cumulative: true, source: "messages" }) + expect(first?.tokensUsed).toBe(20) + const persistedAfterFirst = JSON.parse(await readFile(process.env.OPENCODE_GOAL_STATE_PATH!, "utf8")) as { + goals: Record }> + } + expect(persistedAfterFirst.goals.ses_1?.usageTrackers?.messages).toEqual({ + baseline: 100, + lastObserved: 100, + baseTokens: 20, + pendingBaseline: null, + pendingBaseTokens: null, + }) + + const grown = await accountUsage("ses_1", 105, { cumulative: true, source: "messages" }) + expect(grown?.tokensUsed).toBe(25) + expect((await getGoal("ses_1"))?.tokensUsed).toBe(25) +}) + +test("cumulative usage rebases after a session counter reset without decreasing usage", async () => { + await createGoal("ses_1", "survive compaction", null) + await accountUsage("ses_1", 100, { cumulative: true, source: "messages" }) + await accountUsage("ses_1", 110, { cumulative: true, source: "messages" }) + expect((await getGoal("ses_1"))?.tokensUsed).toBe(10) + + const reset = await accountUsage("ses_1", 20, { cumulative: true, source: "messages" }) + expect(reset?.tokensUsed).toBe(10) + const resumed = await accountUsage("ses_1", 25, { cumulative: true, source: "messages" }) + expect(resumed?.tokensUsed).toBe(15) +}) + +test("a transient cumulative dip does not inflate usage when the source recovers", async () => { + await createGoal("ses_1", "ignore partial observations", null) + await accountUsage("ses_1", 100, { cumulative: true, source: "messages" }) + await accountUsage("ses_1", 110, { cumulative: true, source: "messages" }) + + expect((await accountUsage("ses_1", 0, { cumulative: true, source: "messages" }))?.tokensUsed).toBe(10) + expect((await accountUsage("ses_1", 115, { cumulative: true, source: "messages" }))?.tokensUsed).toBe(15) +}) + +test("independent cumulative sources do not add overlapping usage", async () => { + await createGoal("ses_1", "compare usage sources", null) + await accountUsage("ses_1", 100, { cumulative: true, source: "messages" }) + await accountUsage("ses_1", 110, { cumulative: true, source: "messages" }) + await accountUsage("ses_1", 1_000, { cumulative: true, source: "events" }) + await accountUsage("ses_1", 1_005, { cumulative: true, source: "events" }) + expect((await getGoal("ses_1"))?.tokensUsed).toBe(15) + + await accountUsage("ses_1", 115, { cumulative: true, source: "messages" }) + expect((await getGoal("ses_1"))?.tokensUsed).toBe(15) +}) + +test("an explicit initial baseline counts the first cumulative observation delta", async () => { + await createGoal("ses_1", "count first step", null) + + await accountUsage("ses_1", 1_030, { + cumulative: true, + source: "steps", + initialBaseline: 1_000, + }) + const observed = await accountUsage("ses_1", 1_040, { cumulative: true, source: "steps", initialBaseline: 1_030 }) + + expect(observed?.tokensUsed).toBe(40) + + const afterRestart = await accountUsage("ses_1", 20, { cumulative: true, source: "steps", initialBaseline: 0 }) + expect(afterRestart?.tokensUsed).toBe(60) +}) + +test("an explicit baseline preserves usage for legacy goals without a tracker", async () => { + await createGoal("ses_1", "continue after upgrade", null) + await accountUsage("ses_1", 50) + + const observed = await accountUsage("ses_1", 2, { cumulative: true, source: "steps", initialBaseline: 0 }) + + expect(observed?.tokensUsed).toBe(52) +}) + +test("old persisted goals without usage trackers default to an empty record", async () => { + await createGoal("ses_1", "read old state", null) + const persisted = JSON.parse(await readFile(process.env.OPENCODE_GOAL_STATE_PATH!, "utf8")) as { + goals: Record> + } + delete persisted.goals.ses_1?.usageTrackers + await writeFile(process.env.OPENCODE_GOAL_STATE_PATH!, JSON.stringify(persisted), "utf8") + + expect((await getGoal("ses_1"))?.tokensUsed).toBe(0) + await accountUsage("ses_1", 40, { cumulative: true, source: "messages" }) + const rewritten = JSON.parse(await readFile(process.env.OPENCODE_GOAL_STATE_PATH!, "utf8")) as { + goals: Record }> + } + expect(rewritten.goals.ses_1?.usageTrackers?.messages).toEqual({ + baseline: 40, + lastObserved: 40, + baseTokens: 0, + pendingBaseline: null, + pendingBaseTokens: null, + }) +}) + +test("invalid persisted usage trackers are discarded", async () => { + await createGoal("ses_1", "normalize accounting state", null) + const persisted = JSON.parse(await readFile(process.env.OPENCODE_GOAL_STATE_PATH!, "utf8")) as { + goals: Record> + } + persisted.goals.ses_1!.usageTrackers = { + valid: { baseline: 10, lastObserved: 20, baseTokens: 5 }, + fractional: { baseline: 1.5, lastObserved: 20, baseTokens: 0 }, + backwards: { baseline: 20, lastObserved: 10, baseTokens: 0 }, + missing: { lastObserved: 20 }, + text: { baseline: "10", lastObserved: 20, baseTokens: 0 }, + } + await writeFile(process.env.OPENCODE_GOAL_STATE_PATH!, JSON.stringify(persisted), "utf8") + + await accountUsage("ses_1") + const rewritten = JSON.parse(await readFile(process.env.OPENCODE_GOAL_STATE_PATH!, "utf8")) as { + goals: Record }> + } + expect(rewritten.goals.ses_1?.usageTrackers).toEqual({ + valid: { baseline: 10, lastObserved: 20, baseTokens: 5, pendingBaseline: null, pendingBaseTokens: null }, + }) +}) + +test("usage trackers are not exposed by public or internal snapshots", async () => { + const created = await createGoal("ses_1", "hide accounting internals", null) + await accountUsage("ses_1", 100, { cumulative: true, source: "messages" }) + + expect("usageTrackers" in created).toBe(false) + expect("usageTrackers" in (await getGoal("ses_1"))!).toBe(false) + expect("usageTrackers" in (await getGoalInternal("ses_1"))!).toBe(false) +}) + +test("direct usage accounting remains the default and does not establish a tracker", async () => { + await createGoal("ses_1", "preserve direct accounting", null) + await accountUsage("ses_1", 12) + const unchanged = await accountUsage("ses_1", 8) + expect(unchanged?.tokensUsed).toBe(12) + + const persisted = JSON.parse(await readFile(process.env.OPENCODE_GOAL_STATE_PATH!, "utf8")) as { + goals: Record }> + } + expect(persisted.goals.ses_1?.usageTrackers).toEqual({}) +}) + test("reserves continuation until max auto turns is reached", async () => { await createGoal("ses_1", "continue", null) expect(await reserveContinuation("ses_1", 1, 0)).not.toBeNull()