summaryrefslogtreecommitdiffhomepage
path: root/packages/kernel/src
diff options
context:
space:
mode:
Diffstat (limited to 'packages/kernel/src')
-rw-r--r--packages/kernel/src/contracts/conversation.ts2
-rw-r--r--packages/kernel/src/contracts/index.ts2
-rw-r--r--packages/kernel/src/runtime/run-turn.test.ts151
-rw-r--r--packages/kernel/src/runtime/run-turn.ts20
4 files changed, 171 insertions, 4 deletions
diff --git a/packages/kernel/src/contracts/conversation.ts b/packages/kernel/src/contracts/conversation.ts
index 0080964..f4a342d 100644
--- a/packages/kernel/src/contracts/conversation.ts
+++ b/packages/kernel/src/contracts/conversation.ts
@@ -11,6 +11,7 @@ export type {
ErrorChunk,
Role,
StepId,
+ StepMetrics,
StoredChunk,
SystemChunk,
TextChunk,
@@ -18,4 +19,5 @@ export type {
ToolCallChunk,
ToolResultChunk,
TurnId,
+ TurnMetrics,
} from "@dispatch/wire";
diff --git a/packages/kernel/src/contracts/index.ts b/packages/kernel/src/contracts/index.ts
index 38f1442..b5802f3 100644
--- a/packages/kernel/src/contracts/index.ts
+++ b/packages/kernel/src/contracts/index.ts
@@ -18,6 +18,7 @@ export type {
ErrorChunk,
Role,
StepId,
+ StepMetrics,
StoredChunk,
SystemChunk,
TextChunk,
@@ -25,6 +26,7 @@ export type {
ToolCallChunk,
ToolResultChunk,
TurnId,
+ TurnMetrics,
} from "./conversation.js";
export type { ToolDispatchPolicy } from "./dispatch.js";
export type {
diff --git a/packages/kernel/src/runtime/run-turn.test.ts b/packages/kernel/src/runtime/run-turn.test.ts
index ce654d5..6db6645 100644
--- a/packages/kernel/src/runtime/run-turn.test.ts
+++ b/packages/kernel/src/runtime/run-turn.test.ts
@@ -1307,6 +1307,157 @@ describe("runTurn", () => {
// Only one decode span (for the second step)
expect(decodeOpens).toHaveLength(1);
});
+
+ it("turn span close stamps usage.inputTokens / usage.outputTokens (dotted)", async () => {
+ const provider = createFakeProvider([
+ [
+ { type: "text-delta", delta: "hi" },
+ { type: "usage", usage: { inputTokens: 10, outputTokens: 5 } },
+ { type: "finish", reason: "stop" },
+ ],
+ ]);
+
+ const { logger, sink } = createTestLogger();
+
+ await runTurn({
+ provider,
+ messages: [userMessage],
+ tools: [],
+ dispatch: { maxConcurrent: 1, eager: false },
+ conversationId: "conv-1",
+ turnId: "turn-1",
+ emit: () => {},
+ logger,
+ });
+
+ const turnClose = sink.records.find((r) => r.kind === "span-close" && r.name === "turn");
+ expect(turnClose).toBeDefined();
+ if (turnClose?.kind === "span-close") {
+ expect(turnClose.attributes?.["usage.inputTokens"]).toBe(10);
+ expect(turnClose.attributes?.["usage.outputTokens"]).toBe(5);
+ expect(turnClose.attributes?.usage_inputTokens).toBeUndefined();
+ expect(turnClose.attributes?.usage_outputTokens).toBeUndefined();
+ }
+ });
+
+ it("step span close stamps usage.inputTokens / usage.outputTokens (dotted)", async () => {
+ const provider = createFakeProvider([
+ [
+ { type: "text-delta", delta: "hi" },
+ { type: "usage", usage: { inputTokens: 7, outputTokens: 3 } },
+ { type: "finish", reason: "stop" },
+ ],
+ ]);
+
+ const { logger, sink } = createTestLogger();
+
+ await runTurn({
+ provider,
+ messages: [userMessage],
+ tools: [],
+ dispatch: { maxConcurrent: 1, eager: false },
+ conversationId: "conv-1",
+ turnId: "turn-1",
+ emit: () => {},
+ logger,
+ });
+
+ const stepClose = sink.records.find((r) => r.kind === "span-close" && r.name === "step");
+ expect(stepClose).toBeDefined();
+ if (stepClose?.kind === "span-close") {
+ expect(stepClose.attributes?.["usage.inputTokens"]).toBe(7);
+ expect(stepClose.attributes?.["usage.outputTokens"]).toBe(3);
+ expect(stepClose.attributes?.usage_inputTokens).toBeUndefined();
+ expect(stepClose.attributes?.usage_outputTokens).toBeUndefined();
+ }
+ });
+
+ it("turn + step spans stamp usage.cacheReadTokens / usage.cacheWriteTokens when the provider Usage carries them", async () => {
+ const provider = createFakeProvider([
+ [
+ { type: "text-delta", delta: "hi" },
+ {
+ type: "usage",
+ usage: { inputTokens: 10, outputTokens: 5, cacheReadTokens: 3, cacheWriteTokens: 2 },
+ },
+ { type: "finish", reason: "stop" },
+ ],
+ ]);
+
+ const { logger, sink } = createTestLogger();
+
+ await runTurn({
+ provider,
+ messages: [userMessage],
+ tools: [],
+ dispatch: { maxConcurrent: 1, eager: false },
+ conversationId: "conv-1",
+ turnId: "turn-1",
+ emit: () => {},
+ logger,
+ });
+
+ const turnClose = sink.records.find((r) => r.kind === "span-close" && r.name === "turn");
+ const stepClose = sink.records.find((r) => r.kind === "span-close" && r.name === "step");
+
+ expect(turnClose).toBeDefined();
+ if (turnClose?.kind === "span-close") {
+ expect(turnClose.attributes?.["usage.inputTokens"]).toBe(10);
+ expect(turnClose.attributes?.["usage.outputTokens"]).toBe(5);
+ expect(turnClose.attributes?.["usage.cacheReadTokens"]).toBe(3);
+ expect(turnClose.attributes?.["usage.cacheWriteTokens"]).toBe(2);
+ }
+
+ expect(stepClose).toBeDefined();
+ if (stepClose?.kind === "span-close") {
+ expect(stepClose.attributes?.["usage.inputTokens"]).toBe(10);
+ expect(stepClose.attributes?.["usage.outputTokens"]).toBe(5);
+ expect(stepClose.attributes?.["usage.cacheReadTokens"]).toBe(3);
+ expect(stepClose.attributes?.["usage.cacheWriteTokens"]).toBe(2);
+ }
+ });
+
+ it("turn + step spans OMIT the cache-token attrs when the provider Usage lacks them", async () => {
+ const provider = createFakeProvider([
+ [
+ { type: "text-delta", delta: "hi" },
+ { type: "usage", usage: { inputTokens: 10, outputTokens: 5 } },
+ { type: "finish", reason: "stop" },
+ ],
+ ]);
+
+ const { logger, sink } = createTestLogger();
+
+ await runTurn({
+ provider,
+ messages: [userMessage],
+ tools: [],
+ dispatch: { maxConcurrent: 1, eager: false },
+ conversationId: "conv-1",
+ turnId: "turn-1",
+ emit: () => {},
+ logger,
+ });
+
+ const turnClose = sink.records.find((r) => r.kind === "span-close" && r.name === "turn");
+ const stepClose = sink.records.find((r) => r.kind === "span-close" && r.name === "step");
+
+ expect(turnClose).toBeDefined();
+ if (turnClose?.kind === "span-close") {
+ expect(turnClose.attributes?.["usage.inputTokens"]).toBe(10);
+ expect(turnClose.attributes?.["usage.outputTokens"]).toBe(5);
+ expect(turnClose.attributes?.["usage.cacheReadTokens"]).toBeUndefined();
+ expect(turnClose.attributes?.["usage.cacheWriteTokens"]).toBeUndefined();
+ }
+
+ expect(stepClose).toBeDefined();
+ if (stepClose?.kind === "span-close") {
+ expect(stepClose.attributes?.["usage.inputTokens"]).toBe(10);
+ expect(stepClose.attributes?.["usage.outputTokens"]).toBe(5);
+ expect(stepClose.attributes?.["usage.cacheReadTokens"]).toBeUndefined();
+ expect(stepClose.attributes?.["usage.cacheWriteTokens"]).toBeUndefined();
+ }
+ });
});
describe("provider logger threading", () => {
diff --git a/packages/kernel/src/runtime/run-turn.ts b/packages/kernel/src/runtime/run-turn.ts
index 06069a2..d5db1bf 100644
--- a/packages/kernel/src/runtime/run-turn.ts
+++ b/packages/kernel/src/runtime/run-turn.ts
@@ -50,6 +50,20 @@ function addUsage(a: Usage, b: Usage): Usage {
return { inputTokens, outputTokens };
}
+function usageAttrs(usage: Usage): Record<string, string | number | boolean | null> {
+ const attrs: Record<string, string | number | boolean | null> = {
+ "usage.inputTokens": usage.inputTokens,
+ "usage.outputTokens": usage.outputTokens,
+ };
+ if (usage.cacheReadTokens !== undefined) {
+ attrs["usage.cacheReadTokens"] = usage.cacheReadTokens;
+ }
+ if (usage.cacheWriteTokens !== undefined) {
+ attrs["usage.cacheWriteTokens"] = usage.cacheWriteTokens;
+ }
+ return attrs;
+}
+
function appendTextDelta(chunks: Chunk[], delta: string): void {
const lastIdx = chunks.length - 1;
const last = chunks[lastIdx];
@@ -409,8 +423,7 @@ async function executeStep(ctx: StepContext): Promise<StepResult> {
stepSpan.end({
attrs: {
finishReason,
- usage_inputTokens: stepUsage.inputTokens,
- usage_outputTokens: stepUsage.outputTokens,
+ ...usageAttrs(stepUsage),
},
});
} catch {
@@ -533,8 +546,7 @@ export async function runTurn(input: RunTurnInput): Promise<RunTurnResult> {
turnSpan.end({
attrs: {
finishReason,
- usage_inputTokens: totalUsage.inputTokens,
- usage_outputTokens: totalUsage.outputTokens,
+ ...usageAttrs(totalUsage),
},
});
} catch {
generated by cgit v1.2.3 (git 2.39.1) at 2026-06-29 10:06:17 +0000