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/test/suite/agent-session-model-extension.test.ts
+33 -9
View File
@@ -1,5 +1,5 @@
import type { AgentTool, ThinkingLevel } from "@earendil-works/pi-agent-core";
import { fauxAssistantMessage, fauxToolCall, type Model } from "@earendil-works/pi-ai";
import { fauxAssistantMessage, fauxToolCall, type Model, type Usage } from "@earendil-works/pi-ai";
import { Type } from "typebox";
import { afterEach, describe, expect, it } from "vitest";
import type { BuildSystemPromptOptions, ExtensionAPI } from "../../src/index.ts";
@@ -156,6 +156,23 @@ describe("AgentSession model and extension characterization", () => {
});
it("allows extension tool_result handlers to modify tool results", async () => {
const toolUsage: Usage = {
input: 1,
output: 2,
cacheRead: 3,
cacheWrite: 4,
totalTokens: 10,
cost: { input: 0.1, output: 0.2, cacheRead: 0.3, cacheWrite: 0.4, total: 1 },
};
const patchedToolUsage: Usage = {
input: 5,
output: 6,
cacheRead: 7,
cacheWrite: 8,
totalTokens: 26,
cost: { input: 0.5, output: 0.6, cacheRead: 0.7, cacheWrite: 0.8, total: 2.6 },
};
let observedToolUsage: Usage | undefined;
const echoTool: AgentTool = {
name: "echo",
label: "Echo",
@@ -163,17 +180,21 @@ describe("AgentSession model and extension characterization", () => {
parameters: Type.Object({ text: Type.String() }),
execute: async (_toolCallId, params) => {
const text = typeof params === "object" && params !== null && "text" in params ? String(params.text) : "";
return { content: [{ type: "text", text }], details: { text } };
return { content: [{ type: "text", text }], details: { text }, usage: toolUsage };
},
};
const harness = await createHarness({
tools: [echoTool],
extensionFactories: [
(pi) => {
pi.on("tool_result", async () => ({
content: [{ type: "text", text: "patched result" }],
details: { patched: true },
}));
pi.on("tool_result", async (event) => {
observedToolUsage = event.usage;
return {
content: [{ type: "text", text: "patched result" }],
details: { patched: true },
usage: patchedToolUsage,
};
});
},
],
});
@@ -196,9 +217,12 @@ describe("AgentSession model and extension characterization", () => {
await harness.session.prompt("hi");
expect(getAssistantTexts(harness)).toContain("patched result");
expect(
harness.session.messages.find((message) => message.role === "toolResult" && message.details?.patched === true),
).toBeDefined();
const toolResult = harness.session.messages.find(
(message) => message.role === "toolResult" && message.details?.patched === true,
);
expect(observedToolUsage).toEqual(toolUsage);
expect(toolResult).toBeDefined();
expect(toolResult?.role === "toolResult" ? toolResult.usage : undefined).toEqual(patchedToolUsage);
});
it("allows extension context handlers to modify messages before the LLM call", async () => {