ton/pi_harness
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add usage info to branch summary, compaction and tool result entries (#6671)

Browse Source 
* add usage info to branch summary entries

* add usage to compaction entries

* allow custom tools to report llm usage in tool results

* allow observing and patching usage in tool_result hooks

* agent-harness: save usage in entries

for compaction, branch summaries and tool results
This commit is contained in:
David Brailovsky
2026-07-20 16:41:43 +02:00
committed by GitHub
parent c179395218
commit 2fd3868401
29 changed files with 844 additions and 122 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/coding-agent/src/core/agent-session.ts
+33 -20
View File
@@ -32,6 +32,7 @@ import type {
Model,
ProviderHeaders,
TextContent,
Usage,
} from "@earendil-works/pi-ai/compat";
import {
clampThinkingLevel,
@@ -104,6 +105,7 @@ import { type BuildSystemPromptOptions, buildSystemPrompt } from "./system-promp
import { type BashOperations, createLocalBashOperations } from "./tools/bash.ts";
import { createAllToolDefinitions } from "./tools/index.ts";
import { createToolDefinitionFromAgentTool } from "./tools/tool-definition-wrapper.ts";
import { addUsageToTotals, createUsageTotals } from "./usage-totals.ts";
// ============================================================================
// Skill Block Parsing
@@ -482,6 +484,7 @@ export class AgentSession {
content: result.content,
details: result.details,
isError,
usage: result.usage,
});
if (!hookResult) {
@@ -492,6 +495,7 @@ export class AgentSession {
content: hookResult.content,
details: hookResult.details,
isError: hookResult.isError ?? isError,
usage: hookResult.usage,
};
};
}
@@ -1812,6 +1816,7 @@ export class AgentSession {
let summary: string;
let firstKeptEntryId: string;
let tokensBefore: number;
let usage: Usage | undefined;
let details: unknown;
if (extensionCompaction) {
@@ -1819,6 +1824,7 @@ export class AgentSession {
summary = extensionCompaction.summary;
firstKeptEntryId = extensionCompaction.firstKeptEntryId;
tokensBefore = extensionCompaction.tokensBefore;
usage = extensionCompaction.usage;
details = extensionCompaction.details;
} else {
// Generate compaction result
@@ -1836,6 +1842,7 @@ export class AgentSession {
summary = result.summary;
firstKeptEntryId = result.firstKeptEntryId;
tokensBefore = result.tokensBefore;
usage = result.usage;
details = result.details;
}
@@ -1843,7 +1850,7 @@ export class AgentSession {
throw new Error("Compaction cancelled");
}
this.sessionManager.appendCompaction(summary, firstKeptEntryId, tokensBefore, details, fromExtension);
this.sessionManager.appendCompaction(summary, firstKeptEntryId, tokensBefore, details, fromExtension, usage);
const newEntries = this.sessionManager.getEntries();
const sessionContext = this.sessionManager.buildSessionContext();
this.agent.state.messages = sessionContext.messages;
@@ -1869,6 +1876,7 @@ export class AgentSession {
firstKeptEntryId,
tokensBefore,
estimatedTokensAfter,
usage,
details,
};
this._emit({
@@ -2084,6 +2092,7 @@ export class AgentSession {
let summary: string;
let firstKeptEntryId: string;
let tokensBefore: number;
let usage: Usage | undefined;
let details: unknown;
if (extensionCompaction) {
@@ -2091,6 +2100,7 @@ export class AgentSession {
summary = extensionCompaction.summary;
firstKeptEntryId = extensionCompaction.firstKeptEntryId;
tokensBefore = extensionCompaction.tokensBefore;
usage = extensionCompaction.usage;
details = extensionCompaction.details;
} else {
// Generate compaction result
@@ -2108,6 +2118,7 @@ export class AgentSession {
summary = compactResult.summary;
firstKeptEntryId = compactResult.firstKeptEntryId;
tokensBefore = compactResult.tokensBefore;
usage = compactResult.usage;
details = compactResult.details;
}
@@ -2122,7 +2133,7 @@ export class AgentSession {
return false;
}
this.sessionManager.appendCompaction(summary, firstKeptEntryId, tokensBefore, details, fromExtension);
this.sessionManager.appendCompaction(summary, firstKeptEntryId, tokensBefore, details, fromExtension, usage);
const newEntries = this.sessionManager.getEntries();
const sessionContext = this.sessionManager.buildSessionContext();
this.agent.state.messages = sessionContext.messages;
@@ -2148,6 +2159,7 @@ export class AgentSession {
firstKeptEntryId,
tokensBefore,
estimatedTokensAfter,
usage,
details,
};
this._emit({ type: "compaction_end", reason, result, aborted: false, willRetry });
@@ -2872,7 +2884,7 @@ export class AgentSession {
this._branchSummaryAbortController = new AbortController();
try {
let extensionSummary: { summary: string; details?: unknown } | undefined;
let extensionSummary: { summary: string; details?: unknown; usage?: Usage } | undefined;
let fromExtension = false;
// Emit session_before_tree event
@@ -2907,6 +2919,7 @@ export class AgentSession {
// Run default summarizer if needed
let summaryText: string | undefined;
let summaryDetails: unknown;
let summaryUsage: Usage | undefined;
if (options.summarize && entriesToSummarize.length > 0 && !extensionSummary) {
const model = this.model!;
const { apiKey, headers, env } = await this._getSummarizationRequestAuth(model);
@@ -2929,6 +2942,7 @@ export class AgentSession {
throw new Error(result.error);
}
summaryText = result.summary;
summaryUsage = result.usage;
summaryDetails = {
readFiles: result.readFiles || [],
modifiedFiles: result.modifiedFiles || [],
@@ -2936,6 +2950,7 @@ export class AgentSession {
} else if (extensionSummary) {
summaryText = extensionSummary.summary;
summaryDetails = extensionSummary.details;
summaryUsage = extensionSummary.usage;
}
// Determine the new leaf position based on target type
@@ -2965,6 +2980,7 @@ export class AgentSession {
summaryText,
summaryDetails,
fromExtension,
summaryUsage,
);
summaryEntry = this.sessionManager.getEntry(summaryId) as BranchSummaryEntry;
@@ -3037,13 +3053,12 @@ export class AgentSession {
let toolResults = 0;
let totalMessages = 0;
let toolCalls = 0;
let totalInput = 0;
let totalOutput = 0;
let totalCacheRead = 0;
let totalCacheWrite = 0;
let totalCost = 0;
const usageTotals = createUsageTotals();
for (const entry of this.sessionManager.getEntries()) {
if ((entry.type === "branch_summary" || entry.type === "compaction") && entry.usage) {
addUsageToTotals(usageTotals, entry.usage);
}
if (entry.type !== "message") continue;
totalMessages++;
const message = entry.message;
@@ -3051,18 +3066,16 @@ export class AgentSession {
userMessages++;
} else if (message.role === "toolResult") {
toolResults++;
if (message.usage) {
addUsageToTotals(usageTotals, message.usage);
}
} else if (message.role === "assistant") {
assistantMessages++;
const assistantMsg = message as AssistantMessage;
if (Array.isArray(assistantMsg.content)) {
toolCalls += assistantMsg.content.filter((c) => c.type === "toolCall").length;
}
const usage = assistantMsg.usage;
totalInput += usage.input;
totalOutput += usage.output;
totalCacheRead += usage.cacheRead;
totalCacheWrite += usage.cacheWrite;
totalCost += usage.cost.total;
addUsageToTotals(usageTotals, assistantMsg.usage);
}
}
@@ -3075,13 +3088,13 @@ export class AgentSession {
toolResults,
totalMessages,
tokens: {
input: totalInput,
output: totalOutput,
cacheRead: totalCacheRead,
cacheWrite: totalCacheWrite,
total: totalInput + totalOutput + totalCacheRead + totalCacheWrite,
input: usageTotals.input,
output: usageTotals.output,
cacheRead: usageTotals.cacheRead,
cacheWrite: usageTotals.cacheWrite,
total: usageTotals.input + usageTotals.output + usageTotals.cacheRead + usageTotals.cacheWrite,
},
cost: totalCost,
cost: usageTotals.cost,
contextUsage: this.getContextUsage(),
};
}