fix: account completion calling turn usage

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.

src/goal-accounting.ts

@@ -1,11 +1,13 @@
-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";
 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,
   };

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);
   };
 

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"

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);

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);
       }