feat(pi-local): fix bugs, add RPC mode, improve cost tracking and output handling

Major improvements to the Pi local adapter:

Bug Fixes (Greptile-identified):
- Fix string interpolation in models.ts error message (was showing literal ${detail})
- Fix tool matching in parse.ts to use toolCallId instead of toolName for correct
  multi-call handling and result assignment
- Fix dead code in execute.ts by tracking instructionsReadFailed flag

Feature Improvements:
- Switch from print mode (-p) to RPC mode (--mode rpc) to prevent agent from
  exiting prematurely and ensure proper lifecycle completion
- Add stdin command sending via JSON-RPC format for prompt delivery
- Add line buffering in execute.ts to handle partial JSON chunks correctly
- Filter RPC protocol messages (response, extension_ui_request, etc.) from transcript

Cost Tracking:
- Extract cost and usage data from turn_end assistant messages
- Support both Pi format (input/output/cacheRead/cost.total) and generic format
- Add tests for cost extraction and accumulation across multiple turns

All tests pass (12/12), typecheck clean, server builds successfully.
This commit is contained in:
Richard Anaya
2026-03-07 07:23:44 -08:00
parent 6077ae6064
commit a6b5f12daf
5 changed files with 206 additions and 20 deletions

View File

@@ -100,6 +100,115 @@ describe("parsePiJsonl", () => {
expect(parsed.toolCalls[0].isError).toBe(true);
expect(parsed.toolCalls[0].result).toBe("File not found");
});
it("extracts usage and cost from turn_end events", () => {
const stdout = [
JSON.stringify({
type: "turn_end",
message: {
role: "assistant",
content: "Response with usage",
usage: {
input: 100,
output: 50,
cacheRead: 20,
totalTokens: 170,
cost: {
input: 0.001,
output: 0.0015,
cacheRead: 0.0001,
cacheWrite: 0,
total: 0.0026,
},
},
},
toolResults: [],
}),
].join("\n");
const parsed = parsePiJsonl(stdout);
expect(parsed.usage.inputTokens).toBe(100);
expect(parsed.usage.outputTokens).toBe(50);
expect(parsed.usage.cachedInputTokens).toBe(20);
expect(parsed.usage.costUsd).toBeCloseTo(0.0026, 4);
});
it("accumulates usage from multiple turns", () => {
const stdout = [
JSON.stringify({
type: "turn_end",
message: {
role: "assistant",
content: "First response",
usage: {
input: 50,
output: 25,
cacheRead: 0,
cost: { total: 0.001 },
},
},
}),
JSON.stringify({
type: "turn_end",
message: {
role: "assistant",
content: "Second response",
usage: {
input: 30,
output: 20,
cacheRead: 10,
cost: { total: 0.0015 },
},
},
}),
].join("\n");
const parsed = parsePiJsonl(stdout);
expect(parsed.usage.inputTokens).toBe(80);
expect(parsed.usage.outputTokens).toBe(45);
expect(parsed.usage.cachedInputTokens).toBe(10);
expect(parsed.usage.costUsd).toBeCloseTo(0.0025, 4);
});
it("handles standalone usage events with Pi format", () => {
const stdout = [
JSON.stringify({
type: "usage",
usage: {
input: 200,
output: 100,
cacheRead: 50,
cost: { total: 0.005 },
},
}),
].join("\n");
const parsed = parsePiJsonl(stdout);
expect(parsed.usage.inputTokens).toBe(200);
expect(parsed.usage.outputTokens).toBe(100);
expect(parsed.usage.cachedInputTokens).toBe(50);
expect(parsed.usage.costUsd).toBe(0.005);
});
it("handles standalone usage events with generic format", () => {
const stdout = [
JSON.stringify({
type: "usage",
usage: {
inputTokens: 150,
outputTokens: 75,
cachedInputTokens: 25,
costUsd: 0.003,
},
}),
].join("\n");
const parsed = parsePiJsonl(stdout);
expect(parsed.usage.inputTokens).toBe(150);
expect(parsed.usage.outputTokens).toBe(75);
expect(parsed.usage.cachedInputTokens).toBe(25);
expect(parsed.usage.costUsd).toBe(0.003);
});
});
describe("isPiUnknownSessionError", () => {