diff --git a/CHANGELOG.md b/CHANGELOG.md index dd2026d..08ce838 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,9 @@ # Changelog +## Unreleased + +- Include the assistant response calling `update_goal` in the final token report, without charging a replacement goal for a response generated for its predecessor. + ## 0.3.1 - 2026-09-22 - Align the native host test fixtures with Pi 0.87's registered tool definitions and session projection, and qualify official and fork hosts in one CI job. Runtime behavior is unchanged from 0.3.0. diff --git a/src/goal-accounting.ts b/src/goal-accounting.ts index 4d4cee4..3fb2885 100644 --- a/src/goal-accounting.ts +++ b/src/goal-accounting.ts @@ -1,4 +1,4 @@ -import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent"; +import type { ExtensionAPI, ExtensionContext, SessionEntry } from "@earendil-works/pi-coding-agent"; import { budgetLimitPrompt } from "./prompts.js"; import { applyUsage } from "./state.js"; @@ -6,6 +6,8 @@ import { CUSTOM_ENTRY_TYPE, type ThreadGoal } from "./types.js"; export interface AccountingState { activeGoalId: string | null; + /** Response owner stays fixed through pause/resume and replacements; cleared once turn_end accounts it. */ + turnGoalId: string | null; lastAccountedAt: number | null; budgetWarningSentFor: string | null; } @@ -24,6 +26,7 @@ export interface AssistantTurnMessage { export function createAccountingState(): AccountingState { return { activeGoalId: null, + turnGoalId: null, lastAccountedAt: null, budgetWarningSentFor: null, }; @@ -43,6 +46,20 @@ export function assistantTurnTokens(message: AssistantTurnMessage): number { return usageChannelTokens(message.usage.input) + usageChannelTokens(message.usage.output); } +/** The current response is persisted before its tools execute, but counted at turn_end. */ +export function assistantTurnTokensForToolCall(entries: SessionEntry[], toolCallId: string): number { + for (let i = entries.length - 1; i >= 0; i -= 1) { + const entry = entries[i]; + if (entry?.type !== "message" || entry.message.role !== "assistant") { + continue; + } + return entry.message.content.some((part) => part.type === "toolCall" && part.id === toolCallId) + ? assistantTurnTokens(entry.message) + : 0; + } + return 0; +} + export function isAbortedAssistantMessage(message: AssistantTurnMessage): boolean { return message.role === "assistant" && message.stopReason === "aborted"; } @@ -59,15 +76,12 @@ interface GoalAccountingDeps { } export function createGoalAccounting(deps: GoalAccountingDeps) { - const clearActiveAccounting = (): void => { - const accounting = deps.getAccounting(); - accounting.activeGoalId = null; - accounting.lastAccountedAt = null; - }; - - const beginAccounting = (): void => { + const beginAccounting = (newTurn = true): void => { const goal = deps.getGoal(); const accounting = deps.getAccounting(); + if (newTurn) { + accounting.turnGoalId = goal?.status === "active" ? goal.goalId : null; + } if (!goal || goal.status !== "active") { accounting.activeGoalId = null; accounting.lastAccountedAt = null; @@ -87,16 +101,24 @@ export function createGoalAccounting(deps: GoalAccountingDeps) { const goal = deps.getGoal(); const accounting = deps.getAccounting(); const canAccount = goal?.status === "active" || (accountBudgetLimited && goal?.status === "budgetLimited"); - if (!goal || accounting.activeGoalId !== goal.goalId || !canAccount) { - beginAccounting(); + if (!goal || !canAccount) { + beginAccounting(false); return; } + if (accounting.activeGoalId !== goal.goalId) { + // Re-arm elapsed time after resume/replacement without changing the response's owner. + beginAccounting(false); + if (accounting.activeGoalId !== goal.goalId) { + return; + } + } const now = Date.now(); const elapsed = accounting.lastAccountedAt === null ? 0 : Math.floor((now - accounting.lastAccountedAt) / 1000); accounting.lastAccountedAt = now; - const result = applyUsage(goal, completedTurnTokens, elapsed, { + const tokens = accounting.turnGoalId === goal.goalId ? completedTurnTokens : 0; + const result = applyUsage(goal, tokens, elapsed, { expectedGoalId: accounting.activeGoalId, accountBudgetLimited, }); @@ -121,7 +143,6 @@ export function createGoalAccounting(deps: GoalAccountingDeps) { }; return { - clearActiveAccounting, beginAccounting, accountProgress, }; diff --git a/src/goal-runtime-controller.ts b/src/goal-runtime-controller.ts index d1e6aee..f5be2a8 100644 --- a/src/goal-runtime-controller.ts +++ b/src/goal-runtime-controller.ts @@ -2,7 +2,7 @@ import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-a import { registerGoalCommand } from "./commands.js"; import { createContinuationScheduler } from "./continuation-scheduler.js"; -import { createGoalAccounting } from "./goal-accounting.js"; +import { assistantTurnTokensForToolCall, createGoalAccounting } from "./goal-accounting.js"; import { createGoalPersistence } from "./goal-persistence.js"; import { createGoalRuntimeEventHandlers, @@ -31,7 +31,7 @@ export interface GoalRuntimeController extends GoalRuntimeEventHandlers { getGoalStartTurnStrategy(): GoalStartTurnStrategy; setGoal(goal: ThreadGoal, source: GoalEntrySource, ctx: ExtensionContext): void; clearGoal(source: GoalEntrySource, ctx: ExtensionContext): void; - completeGoal(source: GoalEntrySource, ctx: ExtensionContext): GoalResult; + completeGoal(source: GoalEntrySource, ctx: ExtensionContext, toolCallId: string): GoalResult; cancelProviderLimitAutoResume(goalId: string, ctx: StatusContext): void; resumeGoalWithContinuation(goalId: string, source: GoalEntrySource, ctx: StatusContext): GoalResult; } @@ -163,9 +163,14 @@ export function createGoalRuntimeController(pi: ExtensionAPI): GoalRuntimeContro resumeGoalWithContinuation, }); - const completeGoal = (source: GoalEntrySource, ctx: ExtensionContext): GoalResult => { + const completeGoal = ( + source: GoalEntrySource, + ctx: ExtensionContext, + toolCallId: string, + ): GoalResult => { providerLimitAutoResume.clear(); - goalAccounting.accountProgress(ctx, false, 0, true); + const completedTurnTokens = assistantTurnTokensForToolCall(ctx.sessionManager.getBranch(), toolCallId); + goalAccounting.accountProgress(ctx, false, completedTurnTokens, true); return stateController.completeGoal(source, ctx); }; diff --git a/src/goal-runtime-event-handler-types.ts b/src/goal-runtime-event-handler-types.ts index f8a22d2..c56d049 100644 --- a/src/goal-runtime-event-handler-types.ts +++ b/src/goal-runtime-event-handler-types.ts @@ -121,6 +121,7 @@ export interface GoalRuntimeInputContextHandlerContext extends StaleQueuedWorkEf export interface GoalRuntimeTurnHandlerContext extends StaleQueuedWorkEffectContext { runtimeState: Pick< GoalRuntimeState, + | "accounting" | "agentRunFromContinuation" | "agentRunToolNames" | "currentTurnIndex" diff --git a/src/goal-runtime-turn-handlers.ts b/src/goal-runtime-turn-handlers.ts index e32ef28..a71df1e 100644 --- a/src/goal-runtime-turn-handlers.ts +++ b/src/goal-runtime-turn-handlers.ts @@ -47,6 +47,7 @@ export function createTurnEventHandlers(deps: GoalRuntimeTurnHandlerContext) { const completedTurnTokens = assistantTurnTokens(event.message); goalAccounting.accountProgress(ctx, true, completedTurnTokens); + runtimeState.accounting.turnGoalId = null; stateController.flushGoalPersistence("runtime"); if (isAbortedAssistantMessage(event.message)) { stateController.pauseForAbort(ctx); diff --git a/src/tools.ts b/src/tools.ts index c0af258..e670673 100644 --- a/src/tools.ts +++ b/src/tools.ts @@ -36,7 +36,7 @@ const UpdateGoalParams = Type.Object({ export interface ToolHost { getGoal(): ThreadGoal | null; setGoal(goal: ThreadGoal, source: GoalEntrySource, ctx: ExtensionContext): void; - completeGoal(source: GoalEntrySource, ctx: ExtensionContext): GoalResult; + completeGoal(source: GoalEntrySource, ctx: ExtensionContext, toolCallId: string): GoalResult; } function textResult( @@ -100,8 +100,8 @@ export function registerGoalTools(pi: ExtensionAPI, host: ToolHost): void { promptGuidelines: TOOL_PROMPT_GUIDELINES, parameters: UpdateGoalParams, executionMode: "sequential", - async execute(_toolCallId, _params, _signal, _onUpdate, ctx) { - const result = host.completeGoal("tool", ctx); + async execute(toolCallId, _params, _signal, _onUpdate, ctx) { + const result = host.completeGoal("tool", ctx, toolCallId); if (!result.ok || !result.goal) { throwToolError(result.message); } diff --git a/test/continuation.test.ts b/test/continuation.test.ts index 9ec539a..affda64 100644 --- a/test/continuation.test.ts +++ b/test/continuation.test.ts @@ -1,6 +1,8 @@ import assert from "node:assert/strict"; import { mock, test } from "node:test"; +import type { AssistantMessage } from "@earendil-works/pi-ai"; + import { formatFooterStatus } from "../src/format.js"; import { isGoalCustomEntry, setEntry } from "../src/state.js"; import { CUSTOM_ENTRY_TYPE } from "../src/types.js"; @@ -8,6 +10,7 @@ import { assistantMessage, createRuntimeHarness, emitPersistentAssistantError, + emitToolExecutionEnd, fireProviderLimitAutoResume, flushContinuationScheduler, goalUserContextMessage, @@ -16,6 +19,22 @@ import { sessionShutdownEvent, } from "./support/runtime-harness.js"; +function assistantToolUseMessage( + input: number, + output: number, + toolCalls: ReadonlyArray<{ id: string; name: string }>, +): AssistantMessage { + return { + ...assistantMessage("toolUse", { input, output }), + content: toolCalls.map((toolCall) => ({ + type: "toolCall" as const, + id: toolCall.id, + name: toolCall.name, + arguments: {}, + })), + }; +} + test("aborted turns pause goals and do not queue continuation", async () => { const harness = createRuntimeHarness(); await harness.runCommand("ship it"); @@ -41,6 +60,19 @@ test("aborted turns pause goals and do not queue continuation", async () => { assert.equal(harness.sentMessages.length, 0); }); +test("resuming between aborted turn_end and agent_end does not count the response twice", async () => { + const harness = createRuntimeHarness(); + await harness.runTool("create_goal", { objective: "ship it" }); + await harness.emit("turn_start", { type: "turn_start", turnIndex: 0, timestamp: 1 }); + const message = assistantMessage("aborted", { input: 40, output: 2 }); + await harness.emit("turn_end", { type: "turn_end", turnIndex: 0, message, toolResults: [] }); + assert.equal(harness.snapshot().goal?.usage.tokensUsed, 42); + await harness.runCommand("resume"); + await harness.emit("agent_end", { type: "agent_end", messages: [message] }); + + assert.equal(harness.snapshot().goal?.usage.tokensUsed, 42); +}); + test("a new user-driven agent start leaves a paused goal paused", async () => { const harness = createRuntimeHarness(); await harness.runCommand("ship it"); @@ -176,6 +208,176 @@ test("tool-use turn ends do not queue continuation before tool execution finishe assert.equal(harness.sentMessages.length, 0); }); +test("update_goal accounts its calling turn", async () => { + const harness = createRuntimeHarness(); + await harness.runTool("create_goal", { objective: "ship it" }); + + const message = assistantToolUseMessage(100, 20, [{ id: "update-call", name: "update_goal" }]); + await harness.emit("turn_start", { type: "turn_start", turnIndex: 0, timestamp: 1 }); + harness.appendMessage(message); + await harness.runTool("update_goal", { status: "complete" }, "update-call"); + await harness.emit("turn_end", { + type: "turn_end", + turnIndex: 0, + message, + toolResults: [], + }); + + assert.equal(harness.snapshot().goal?.status, "complete"); + assert.equal(harness.snapshot().goal?.usage.tokensUsed, 120); + assert.equal(harness.sentMessages.length, 0); +}); + +test("update_goal charges the current assistant message when a historical call reuses its ID", async () => { + const harness = createRuntimeHarness(); + await harness.runTool("create_goal", { objective: "ship it" }); + + harness.appendMessage(assistantToolUseMessage(7, 3, [{ id: "", name: "update_goal" }])); + const message = assistantToolUseMessage(100, 20, [{ id: "", name: "update_goal" }]); + await harness.emit("turn_start", { type: "turn_start", turnIndex: 0, timestamp: 1 }); + harness.appendMessage(message); + await harness.runTool("update_goal", { status: "complete" }, ""); + await harness.emit("turn_end", { + type: "turn_end", + turnIndex: 0, + message, + toolResults: [], + }); + + assert.equal(harness.snapshot().goal?.status, "complete"); + assert.equal(harness.snapshot().goal?.usage.tokensUsed, 120); + assert.equal(harness.sentMessages.length, 0); +}); + +test("direct update_goal execution does not charge a historical call with the same ID", async () => { + const harness = createRuntimeHarness(); + await harness.runTool("create_goal", { objective: "ship it" }); + + harness.appendMessage(assistantToolUseMessage(7, 3, [{ id: "", name: "update_goal" }])); + harness.appendMessage(assistantToolUseMessage(100, 20, [{ id: "bash-call", name: "bash" }])); + await harness.emit("turn_start", { type: "turn_start", turnIndex: 0, timestamp: 1 }); + await harness.runTool("update_goal", { status: "complete" }, ""); + + assert.equal(harness.snapshot().goal?.status, "complete"); + assert.equal(harness.snapshot().goal?.usage.tokensUsed, 0); + assert.equal(harness.sentMessages.length, 0); +}); + +test("duplicate update_goal completion does not double count its calling turn", async () => { + const harness = createRuntimeHarness(); + await harness.runTool("create_goal", { objective: "ship it" }); + + const message = assistantToolUseMessage(100, 20, [{ id: "update-call", name: "update_goal" }]); + await harness.emit("turn_start", { type: "turn_start", turnIndex: 0, timestamp: 1 }); + harness.appendMessage(message); + await harness.runTool("update_goal", { status: "complete" }, "update-call"); + await harness.runTool("update_goal", { status: "complete" }, "update-call"); + await harness.emit("turn_end", { + type: "turn_end", + turnIndex: 0, + message, + toolResults: [], + }); + + const goal = harness.snapshot().goal; + assert.equal(goal?.status, "complete"); + assert.equal(goal?.usage.tokensUsed, 120); + assert.equal( + harness.entries.filter( + (entry) => + entry.type === "custom" && + entry.customType === CUSTOM_ENTRY_TYPE && + isGoalCustomEntry(entry.data) && + entry.data.kind === "set" && + entry.data.goal.status === "complete", + ).length, + 1, + ); + assert.equal(harness.sentMessages.length, 0); +}); + +test("update_goal counts a multi-tool assistant message once", async () => { + const harness = createRuntimeHarness(); + await harness.runTool("create_goal", { objective: "ship it" }); + + const message = assistantToolUseMessage(100, 20, [ + { id: "bash-call", name: "bash" }, + { id: "update-call", name: "update_goal" }, + ]); + await harness.emit("turn_start", { type: "turn_start", turnIndex: 0, timestamp: 1 }); + harness.appendMessage(message); + await harness.emit("tool_execution_end", { + type: "tool_execution_end", + toolCallId: "bash-call", + toolName: "bash", + args: {}, + result: {}, + isError: false, + }); + harness.appendMessage({ + role: "toolResult", + toolCallId: "bash-call", + toolName: "bash", + content: [{ type: "text", text: "done" }], + isError: false, + timestamp: 1, + }); + await harness.runTool("update_goal", { status: "complete" }, "update-call"); + await harness.emit("turn_end", { + type: "turn_end", + turnIndex: 0, + message, + toolResults: [], + }); + + assert.equal(harness.snapshot().goal?.status, "complete"); + assert.equal(harness.snapshot().goal?.usage.tokensUsed, 120); + assert.equal(harness.sentMessages.length, 0); +}); + +test("update_goal preserves completion when its calling turn crosses the budget", async () => { + const harness = createRuntimeHarness(); + await harness.runTool("create_goal", { objective: "ship it", token_budget: 500_000 }); + + const message = assistantToolUseMessage(499_999, 2, [{ id: "update-call", name: "update_goal" }]); + await harness.emit("turn_start", { type: "turn_start", turnIndex: 0, timestamp: 1 }); + harness.appendMessage(message); + const result = await harness.runTool("update_goal", { status: "complete" }, "update-call"); + assert.match(JSON.stringify(result), /tokens used: 500,001 of 500,000/); + await harness.emit("turn_end", { + type: "turn_end", + turnIndex: 0, + message, + toolResults: [], + }); + + const goal = harness.snapshot().goal; + assert.equal(goal?.status, "complete"); + assert.equal(goal?.usage.tokensUsed, 500_001); + assert.equal(harness.sentMessages.length, 0); +}); + +test("direct update_goal execution without a matching assistant entry contributes zero", async () => { + const harness = createRuntimeHarness(); + await harness.runTool("create_goal", { objective: "ship it" }); + + const unrelatedMessage = assistantToolUseMessage(100, 20, [{ id: "bash-call", name: "bash" }]); + const callingMessage = assistantToolUseMessage(100, 20, [{ id: "update-call", name: "update_goal" }]); + await harness.emit("turn_start", { type: "turn_start", turnIndex: 0, timestamp: 1 }); + harness.appendMessage(unrelatedMessage); + await harness.runTool("update_goal", { status: "complete" }, "update-call"); + await harness.emit("turn_end", { + type: "turn_end", + turnIndex: 0, + message: callingMessage, + toolResults: [], + }); + + assert.equal(harness.snapshot().goal?.status, "complete"); + assert.equal(harness.snapshot().goal?.usage.tokensUsed, 0); + assert.equal(harness.sentMessages.length, 0); +}); + test("successful budget-crossing turn clears stale recovery footer attention", async () => { const harness = createRuntimeHarness(); await harness.runTool("create_goal", { objective: "ship it", token_budget: 500_000 }); @@ -257,6 +459,69 @@ test("replacement during an in-flight turn does not charge old tokens to the new assert.equal(harness.sentMessages.length, 1); }); +for (const [complete, precedingTool] of [[false, false], [false, true], [true, false], [true, true]]) { + test(`pause/resume preserves response tokens (complete: ${complete}, preceding tool: ${precedingTool})`, async () => { + const harness = createRuntimeHarness(); + await harness.runTool("create_goal", { objective: "ship it" }); + await harness.emit("turn_start", { type: "turn_start", turnIndex: 0, timestamp: 1 }); + const message = complete || precedingTool + ? assistantToolUseMessage(100, 20, [ + ...(precedingTool ? [{ id: "bash-call", name: "bash" }] : []), + ...(complete ? [{ id: "update-call", name: "update_goal" }] : []), + ]) + : assistantMessage("stop", { input: 100, output: 20 }); + harness.appendMessage(message); + await harness.runCommand("pause"); + await harness.runCommand("resume"); + if (precedingTool) { + await emitToolExecutionEnd(harness); + } + if (complete) { + await harness.runTool("update_goal", { status: "complete" }, "update-call"); + } + await harness.emit("turn_end", { type: "turn_end", turnIndex: 0, message, toolResults: [] }); + + const goal = harness.snapshot().goal; + assert.equal(goal?.status, complete ? "complete" : "active"); + assert.equal(goal?.usage.tokensUsed, 120); + }); +} + +for (const replacementSource of ["command", "tool"] as const) { + test(`${replacementSource} replacement counts tool time without claiming the old response's tokens`, async () => { + mock.timers.enable({ apis: ["Date"], now: 1_000_000 }); + try { + const harness = createRuntimeHarness(); + await harness.runTool("create_goal", { objective: "old goal" }); + await harness.emit("turn_start", { type: "turn_start", turnIndex: 0, timestamp: 1 }); + const message = assistantToolUseMessage(100, 20, [ + { id: "create-call", name: "create_goal" }, + { id: "update-call", name: "update_goal" }, + ]); + if (replacementSource === "command") { + await harness.runCommand("new goal"); + harness.appendMessage(message); + } else { + harness.appendMessage(message); + await harness.runTool("create_goal", { objective: "new goal", replace_existing: true }, "create-call"); + } + await emitToolExecutionEnd(harness); + mock.timers.tick(5_000); + await emitToolExecutionEnd(harness); + await harness.runTool("update_goal", { status: "complete" }, "update-call"); + await harness.emit("turn_end", { type: "turn_end", turnIndex: 0, message, toolResults: [] }); + + const goal = harness.snapshot().goal; + assert.equal(goal?.objective, "new goal"); + assert.equal(goal?.status, "complete"); + assert.equal(goal?.usage.tokensUsed, 0); + assert.equal(goal?.usage.activeSeconds, 5); + } finally { + mock.timers.reset(); + } + }); +} + test("goal tools return Codex-shaped response details", async () => { const harness = createRuntimeHarness(); const created = (await harness.runTool("create_goal", { diff --git a/test/sdk-runtime-smoke.test.ts b/test/sdk-runtime-smoke.test.ts index 6e5dba0..891be32 100644 --- a/test/sdk-runtime-smoke.test.ts +++ b/test/sdk-runtime-smoke.test.ts @@ -17,17 +17,20 @@ import { } from "@earendil-works/pi-coding-agent"; import goalExtension, { __testHooks } from "../src/index.js"; +import { reconstructGoal } from "../src/state.js"; import { CUSTOM_ENTRY_TYPE } from "../src/types.js"; function assistantResponse( model: Parameters[0], contextTokens: number, - text: string, + content: string | AssistantMessage["content"], ): ReturnType { const stream = createAssistantMessageEventStream(); + const parts: AssistantMessage["content"] = typeof content === "string" ? [{ type: "text", text: content }] : content; + const stopReason = parts.some((part) => part.type === "toolCall") ? "toolUse" : "stop"; const message: AssistantMessage = { role: "assistant", - content: [{ type: "text", text }], + content: parts, api: model.api, provider: model.provider, model: model.id, @@ -39,17 +42,88 @@ function assistantResponse( totalTokens: contextTokens, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, - stopReason: "stop", + stopReason, timestamp: Date.now(), }; queueMicrotask(() => { stream.push({ type: "start", partial: message }); - stream.push({ type: "done", reason: "stop", message }); + stream.push({ type: "done", reason: stopReason, message }); stream.end(); }); return stream; } +test("SDK completion report includes its calling response exactly once", async () => { + const modelRuntime = await ModelRuntime.create({ + allowModelNetwork: false, + credentials: new InMemoryCredentialStore(), + modelsPath: null, + }); + const loader = new DefaultResourceLoader({ + cwd: process.cwd(), + agentDir: getAgentDir(), + noContextFiles: true, + noExtensions: true, + extensionFactories: [goalExtension], + }); + await loader.reload(); + modelRuntime.registerProvider("sdk-smoke", { apiKey: "test" }); + const { session } = await createAgentSession({ + cwd: process.cwd(), + agentDir: getAgentDir(), + model: { + provider: "sdk-smoke", + id: "completion", + name: "SDK Completion Smoke", + api: "openai-completions", + baseUrl: "http://localhost", + reasoning: false, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 1_000_000, + maxTokens: 1_000, + }, + modelRuntime, + noTools: "builtin", + resourceLoader: loader, + sessionManager: SessionManager.inMemory(process.cwd()), + settingsManager: SettingsManager.inMemory({ compaction: { enabled: false } }), + }); + let requests = 0; + session.agent.streamFunction = (model) => { + requests += 1; + if (requests > 2) { + assert.equal(requests, 3, "completed goals must not queue another continuation"); + return assistantResponse(model, 7, "Done"); + } + return assistantResponse(model, requests === 1 ? 410_000 : 100_000, [{ + type: "toolCall", + id: `call-${requests}`, + name: requests === 1 ? "get_goal" : "update_goal", + arguments: requests === 1 ? {} : { status: "complete" }, + }]); + }; + try { + const runner = session.extensionRunner; + const createGoal = runner.getToolDefinition("create_goal"); + assert.ok(createGoal); + await createGoal.execute("create", { objective: "ship it", token_budget: 500_000 }, + undefined, undefined, runner.createContext()); + await session.prompt("Complete the goal"); + + const entries = session.sessionManager.getBranch(); + const goal = reconstructGoal(entries).goal; + assert.equal(goal?.status, "complete"); + assert.equal(goal?.usage.tokensUsed, 510_000); + const result = entries.find((entry) => + entry.type === "message" && entry.message.role === "toolResult" && entry.message.toolName === "update_goal"); + assert.match(JSON.stringify(result), /tokens used: 510,000 of 500,000/); + assert.equal(requests, 3); + } finally { + session.dispose(); + } +}); + function goalIdFromToolResult(result: unknown): string { assert.ok(result && typeof result === "object"); const details = (result as { details?: unknown }).details; diff --git a/test/support/runtime-harness.ts b/test/support/runtime-harness.ts index 7bb1fbb..2ffa115 100644 --- a/test/support/runtime-harness.ts +++ b/test/support/runtime-harness.ts @@ -1,7 +1,13 @@ import assert from "node:assert/strict"; import { mock } from "node:test"; -import type { ExtensionAPI, ExtensionCommandContext, ExtensionContext } from "@earendil-works/pi-coding-agent"; +import type { AssistantMessage } from "@earendil-works/pi-ai"; +import type { + ExtensionAPI, + ExtensionCommandContext, + ExtensionContext, + SessionMessageEntry, +} from "@earendil-works/pi-coding-agent"; import goalExtension, { __testHooks } from "../../src/index.js"; import { isContextOverflowError } from "../../src/recovery.js"; @@ -89,7 +95,7 @@ export function createRuntimeHarness(options: { const handlers = new Map(); const sentMessages: SentMessage[] = []; const sentUserMessages: SentUserMessage[] = []; - const tools = new Map) => Promise>(); + const tools = new Map) => Promise>(); const compactCalls: Array<{ customInstructions?: string; onComplete?: (result: { @@ -189,9 +195,9 @@ export function createRuntimeHarness(options: { }, registerShortcut() {}, registerTool(tool) { - tools.set(tool.name, (params) => + tools.set(tool.name, (toolCallId, params) => tool.execute( - "tool-call", + toolCallId, params as Parameters[1], undefined, undefined, @@ -400,13 +406,24 @@ export function createRuntimeHarness(options: { return results; } - async function runTool(name: string, params: Record) { + async function runTool(name: string, params: Record, toolCallId = "tool-call") { const tool = tools.get(name); assert.ok(tool, `Expected tool ${name} to be registered.`); - return tool(params); + return tool(toolCallId, params); + } + + function appendMessage(message: SessionMessageEntry["message"]): void { + entries.push({ + type: "message", + id: `entry-${++entryIndex}`, + parentId: null, + timestamp: new Date(0).toISOString(), + message, + }); } return { + appendMessage, compactCalls, footerStatuses, emit, @@ -611,7 +628,7 @@ export function assistantMessage( stopReason: "stop" | "aborted" | "length" | "toolUse" | "error", usage: TestAssistantUsage, errorMessage?: string, -) { +): AssistantMessage { const cacheRead = usage.cacheRead ?? 0; const cacheWrite = usage.cacheWrite ?? 0;