fix: persist usage from rpc

This commit is contained in:
Peter Steinberger
2025-12-10 01:37:57 +00:00
parent 49e70746f0
commit c4b02645f5
3 changed files with 26 additions and 12 deletions

View File

@@ -57,4 +57,15 @@ describe("pi agent helpers", () => {
expect(tool?.toolName).toBe("bash"); expect(tool?.toolName).toBe("bash");
expect(tool?.meta).toBe("ls -la"); expect(tool?.meta).toBe("ls -la");
}); });
it("keeps usage meta even when assistant message has no text", () => {
const stdout = [
'{"type":"message_start","message":{"role":"assistant"}}',
'{"type":"message_end","message":{"role":"assistant","content":[{"type":"thinking","thinking":"hmm"}],"usage":{"input":10,"output":5},"model":"pi-1","provider":"inflection","stopReason":"end"}}',
].join("\n");
const parsed = piSpec.parseOutput(stdout);
expect(parsed.texts?.length ?? 0).toBe(0);
expect((parsed.meta?.usage as { input?: number })?.input).toBe(10);
expect(parsed.meta?.model).toBe("pi-1");
});
}); });

View File

@@ -95,10 +95,11 @@ function parsePiJson(raw: string): AgentParseResult {
.trim(); .trim();
if (isAssistantMessage) { if (isAssistantMessage) {
// Always remember the last assistant message so we keep usage/model meta even when no text.
lastAssistant = msg;
if (msgText && msgText !== lastPushed) { if (msgText && msgText !== lastPushed) {
texts.push(msgText); texts.push(msgText);
lastPushed = msgText; lastPushed = msgText;
lastAssistant = msg;
} }
} else if (isToolResult && msg.content) { } else if (isToolResult && msg.content) {
const toolText = msg.content const toolText = msg.content
@@ -119,15 +120,14 @@ function parsePiJson(raw: string): AgentParseResult {
} }
} }
const meta: AgentMeta | undefined = const meta: AgentMeta | undefined = lastAssistant
lastAssistant && texts.length ? {
? { model: lastAssistant.model,
model: lastAssistant.model, provider: lastAssistant.provider,
provider: lastAssistant.provider, stopReason: lastAssistant.stopReason,
stopReason: lastAssistant.stopReason, usage: lastAssistant.usage,
usage: lastAssistant.usage, }
} : undefined;
: undefined;
return { return {
texts, texts,

View File

@@ -59,12 +59,15 @@ function stripRpcNoise(raw: string): string {
typeof role === "string" && role.toLowerCase().includes("tool"); typeof role === "string" && role.toLowerCase().includes("tool");
if (!isAssistant && !isToolRole) continue; if (!isAssistant && !isToolRole) continue;
// Ignore assistant messages that have no text content (pure toolcall scaffolding). // Ignore assistant messages that have no text content unless they carry usage (final message_end often does).
if (msg?.role === "assistant" && Array.isArray(msg?.content)) { if (msg?.role === "assistant" && Array.isArray(msg?.content)) {
const hasText = msg.content.some( const hasText = msg.content.some(
(c: unknown) => (c as { type?: string })?.type === "text", (c: unknown) => (c as { type?: string })?.type === "text",
); );
if (!hasText) continue; const hasUsage =
typeof msg?.usage === "object" &&
(msg.usage?.input != null || msg.usage?.output != null);
if (!hasText && !hasUsage) continue;
} }
} catch { } catch {
// not JSON; keep as-is // not JSON; keep as-is