Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion apps/extension/entrypoints/sidepanel/agent-chat-atoms.ts
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,10 @@ export const contextUsageAtomFamily = atomFamily((_conversationId: string) =>
atom<ContextUsage | undefined>()
);

// Per-conversation running session spend (USD) from streamed usage.cost (in-memory only).
// Same lifecycle as drafts/usage: kept across close and compaction, evicted on delete and sign-out.
export const sessionCostAtomFamily = atomFamily((_conversationId: string) => atom(0));

// Event id of the assistant message currently streaming content for a conversation.
// Cleared when that message's stream ends (even if later tool rounds continue).
export const streamingMessageIdAtomFamily = atomFamily((_conversationId: string) =>
Expand All @@ -30,10 +34,11 @@ export const streamingMessageIdAtomFamily = atomFamily((_conversationId: string)
export const runningConversationIdsAtom = atom<readonly string[]>([]);
export const compactingConversationIdsAtom = atom<readonly string[]>([]);

// Evict a single conversation's in-memory atoms (draft + context usage + streaming id).
// Evict a single conversation's in-memory atoms (draft + context usage + session cost + streaming id).
export const evictConversationAtoms = (conversationId: string): void => {
draftAtomFamily.remove(conversationId);
contextUsageAtomFamily.remove(conversationId);
sessionCostAtomFamily.remove(conversationId);
streamingMessageIdAtomFamily.remove(conversationId);
};

Expand All @@ -47,6 +52,7 @@ export const clearPerConversationAtoms = (): void => {
const ids = new Set([
...draftAtomFamily.getParams(),
...contextUsageAtomFamily.getParams(),
...sessionCostAtomFamily.getParams(),
...streamingMessageIdAtomFamily.getParams(),
]);
for (const id of ids) {
Expand Down
13 changes: 12 additions & 1 deletion apps/extension/entrypoints/sidepanel/agent-chat-panel.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import {
evictConversationAtoms,
remoteMcpStoreAtom,
runningConversationIdsAtom,
sessionCostAtomFamily,
streamingMessageIdAtomFamily,
} from './agent-chat-atoms';
import {
Expand All @@ -24,6 +25,7 @@ import {
compactConversationEvents,
hasCompactableHistory,
} from '@/src/shared/agent-context-compaction';
import type { TurnUsage } from '@/src/shared/agent-llm-turn-runner-core';
import { defaultMode } from '@/src/shared/agent-chat-placeholder';
import { getKiloApiBaseUrl } from '@/src/shared/auth';
import type { StoredAuth } from '@/src/shared/auth';
Expand All @@ -46,6 +48,7 @@ import { AgentFooterControls } from './agent-footer-controls';
import { ContextDonut } from './context-donut';
import { runDangerousLlmTurn, runSafeLlmTurn } from './agent-turn-runners';
import { AUTO_COMPACT_RATIO, getContextRatio } from '@/src/shared/context-usage';
import { addSessionCost } from '@/src/shared/session-cost';
import { useTabDebugger } from './use-tab-debugger';
import { ConversationList } from './conversation-list';
import { ConversationTabs } from './conversation-tabs';
Expand Down Expand Up @@ -162,6 +165,7 @@ export const AgentChatPanel = ({
const isCompacting = compactingConversationIds.includes(activeConversationId);
const activeUsage = useAtomValue(contextUsageAtomFamily(activeConversationId));
const activePromptTokens = activeUsage?.promptTokens ?? 0;
const activeSessionCostUsd = useAtomValue(sessionCostAtomFamily(activeConversationId));
const streamingMessageId = useAtomValue(streamingMessageIdAtomFamily(activeConversationId));
const contextLength = selectedModel?.contextLength;

Expand Down Expand Up @@ -246,10 +250,12 @@ export const AgentChatPanel = ({
void compactActiveConversation();
}}
promptTokens={activePromptTokens}
sessionCostUsd={activeSessionCostUsd}
/>
),
[
activePromptTokens,
activeSessionCostUsd,
canCompactActive,
compactActiveConversation,
contextLength,
Expand Down Expand Up @@ -471,10 +477,15 @@ export const AgentChatPanel = ({
}
};
let currentRunHasUsage = false;
const updateRunUsage = (usage: { promptTokens: number }): void => {
const updateRunUsage = (usage: TurnUsage): void => {
if (isCurrentRun()) {
currentRunHasUsage = true;
store.set(contextUsageAtomFamily(conversationId), { promptTokens: usage.promptTokens });
const previousCost = store.get(sessionCostAtomFamily(conversationId));
store.set(
sessionCostAtomFamily(conversationId),
addSessionCost(previousCost, usage.costUsd)
);
}
};

Expand Down
5 changes: 5 additions & 0 deletions apps/extension/entrypoints/sidepanel/context-donut.tsx
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import { useRef } from 'react';
import type { JSX } from 'react';
import { formatContextSummary, getContextRatio, getContextTone } from '@/src/shared/context-usage';
import { formatSessionCost } from '@/src/shared/session-cost';

const toneStroke: Record<'danger' | 'safe' | 'warn', string> = {
danger: '#f87171',
Expand All @@ -16,18 +17,21 @@ export const ContextDonut = ({
contextLength,
onCompact,
promptTokens,
sessionCostUsd,
}: {
canCompact: boolean;
contextLength: number | undefined;
onCompact: () => void;
promptTokens: number;
sessionCostUsd: number;
}): JSX.Element => {
const detailsRef = useRef<HTMLDetailsElement>(null);
const ratio = getContextRatio(promptTokens, contextLength);
const stroke = ratio === undefined ? '#52525b' : toneStroke[getContextTone(ratio)];
const dash = ratio === undefined ? 0 : ratio * CIRCUMFERENCE;
const summary = formatContextSummary(promptTokens, contextLength);
const label = `Context usage: ${summary}`;
const sessionCostLabel = formatSessionCost(sessionCostUsd);

return (
<details className="relative shrink-0" ref={detailsRef}>
Expand All @@ -54,6 +58,7 @@ export const ContextDonut = ({
<div className="absolute bottom-10 right-0 z-20 w-56 rounded-md border border-zinc-800 bg-zinc-950 p-3 text-xs text-zinc-300 shadow-lg shadow-zinc-950/60">
<p className="font-medium text-zinc-100">Context</p>
<p className="mt-1 text-zinc-400">{summary}</p>
<p className="mt-1 text-zinc-400">Session cost {sessionCostLabel}</p>
<button
className="mt-3 h-7 w-full rounded-md bg-[#EDFF00] px-2 text-xs font-semibold text-zinc-950 outline-none transition hover:bg-[#d9ea00] focus:ring-2 focus:ring-[#EDFF00]/50 disabled:cursor-not-allowed disabled:bg-zinc-800 disabled:text-zinc-500"
disabled={!canCompact}
Expand Down
20 changes: 16 additions & 4 deletions apps/extension/src/shared/agent-chat-atoms.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,37 +9,49 @@ import {
draftAtomFamily,
evictConversationAtoms,
runningConversationIdsAtom,
sessionCostAtomFamily,
streamingMessageIdAtomFamily,
} from '@/entrypoints/sidepanel/agent-chat-atoms';

describe('per-conversation atom eviction', () => {
it('evictConversationAtoms resets draft, usage, and streaming id for a conversation id', () => {
it('evictConversationAtoms resets draft, usage, session cost, and streaming id for a conversation id', () => {
const store = getDefaultStore();
store.set(draftAtomFamily('conversation-1'), 'hello');
store.set(contextUsageAtomFamily('conversation-1'), { promptTokens: 42 });
store.set(sessionCostAtomFamily('conversation-1'), 0.0123);
store.set(streamingMessageIdAtomFamily('conversation-1'), 'msg-streaming');

evictConversationAtoms('conversation-1');

// A fresh atom (post-remove) starts from its initial value.
expect(store.get(draftAtomFamily('conversation-1'))).toBe('');
expect(store.get(contextUsageAtomFamily('conversation-1'))).toBeUndefined();
expect(store.get(sessionCostAtomFamily('conversation-1'))).toBe(0);
expect(store.get(streamingMessageIdAtomFamily('conversation-1'))).toBeUndefined();
});

it('clearPerConversationAtoms wipes all drafts, usage, streaming ids, and run-state on sign-out', () => {
it('clearPerConversationAtoms wipes all drafts, usage, session cost, and streaming ids on sign-out', () => {
const store = getDefaultStore();
store.set(draftAtomFamily('conversation-1'), 'prev account draft');
store.set(contextUsageAtomFamily('conversation-2'), { promptTokens: 999 });
store.set(sessionCostAtomFamily('conversation-3'), 1.5);
store.set(streamingMessageIdAtomFamily('conversation-3'), 'msg-only-streaming');
store.set(runningConversationIdsAtom, ['conversation-1']);
store.set(compactingConversationIdsAtom, ['conversation-2']);

clearPerConversationAtoms();

expect(store.get(draftAtomFamily('conversation-1'))).toBe('');
expect(store.get(contextUsageAtomFamily('conversation-2'))).toBeUndefined();
expect(store.get(sessionCostAtomFamily('conversation-3'))).toBe(0);
expect(store.get(streamingMessageIdAtomFamily('conversation-3'))).toBeUndefined();
});

it('clearPerConversationAtoms clears run-state on sign-out', () => {
const store = getDefaultStore();
store.set(runningConversationIdsAtom, ['conversation-1']);
store.set(compactingConversationIdsAtom, ['conversation-2']);

clearPerConversationAtoms();

expect(store.get(runningConversationIdsAtom)).toStrictEqual([]);
expect(store.get(compactingConversationIdsAtom)).toStrictEqual([]);
});
Expand Down
11 changes: 10 additions & 1 deletion apps/extension/src/shared/agent-llm-turn-runner-core.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -14,10 +14,12 @@ function* createGatewayResponses(): Generator<Response, Response> {
yield streamResponse([
'data: {"choices":[{"delta":{"content":"Reading"}}]}\n\n',
'data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_snapshot","type":"function","function":{"name":"get_page_snapshot","arguments":"{}"}}]}}]}\n\n',
'data: {"choices":[],"usage":{"prompt_tokens":100,"completion_tokens":5,"total_tokens":105,"cost":0.0007}}\n\n',
'data: [DONE]\n\n',
]);
yield streamResponse([
'data: {"choices":[{"delta":{"content":"Done."}}]}\n\n',
'data: {"choices":[],"usage":{"prompt_tokens":200,"completion_tokens":3,"total_tokens":203,"cost":0.001}}\n\n',
'data: [DONE]\n\n',
]);
return streamResponse(['data: [DONE]\n\n']);
Expand Down Expand Up @@ -63,7 +65,7 @@ describe('agent LLM turn runner core', () => {
const fetch: FetchLike = () =>
streamResponse([
'data: {"choices":[{"delta":{"content":"Done."}}]}\n\n',
'data: {"choices":[],"usage":{"completion_tokens":5,"prompt_tokens":999,"total_tokens":1004}}\n\n',
'data: {"choices":[],"usage":{"completion_tokens":5,"prompt_tokens":999,"total_tokens":1004,"cost":0.0123}}\n\n',
'data: [DONE]\n\n',
]);

Expand All @@ -88,13 +90,15 @@ describe('agent LLM turn runner core', () => {
});

expect(usageCalls).toContainEqual({
costUsd: 0.0123,
promptTokens: 999,
});
});

it('streams, runs tools, and continues with tool results', async () => {
const appendedEvents: AgentConversationEvent[] = [];
const updatedMessages: string[] = [];
const usageCalls: unknown[] = [];
const fetchCalls: unknown[] = [];
const responses = createGatewayResponses();
const fetch: FetchLike = (_input, init) => {
Expand All @@ -115,6 +119,7 @@ describe('agent LLM turn runner core', () => {
maxToolRounds: 4,
model: 'anthropic/claude-sonnet-4',
noResponseMessage: 'The model did not return a response.',
onUsage: usage => usageCalls.push(usage),
signal: undefined,
toToolCallEvents: (toolCalls: KiloGatewayToolCallRequest[]) =>
toolCalls.map(toolCall =>
Expand Down Expand Up @@ -152,6 +157,10 @@ describe('agent LLM turn runner core', () => {
{ role: 'assistant', text: 'Done.', type: 'message' },
]);
expect(fetchCalls).toHaveLength(2);
expect(usageCalls).toStrictEqual([
{ costUsd: 0.0007, promptTokens: 100 },
{ costUsd: 0.001, promptTokens: 200 },
]);
});

it('allows twenty tool rounds before asking the user to continue', async () => {
Expand Down
1 change: 1 addition & 0 deletions apps/extension/src/shared/agent-llm-turn-runner-core.ts
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import type { EvalTabResult } from './tab-debugger';
type ToolCallEvent = Extract<AgentConversationEvent, { readonly type: 'tool-call' }>;

export interface TurnUsage {
readonly costUsd?: number;
readonly promptTokens: number;
}

Expand Down
1 change: 1 addition & 0 deletions apps/extension/src/shared/kilo-gateway-chat-client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@ export interface KiloGatewayChatCompletion {
readonly reasoningDetails?: readonly unknown[];
readonly toolCalls: KiloGatewayToolCallRequest[];
readonly usage?: {
readonly costUsd?: number;
readonly promptTokens: number;
};
}
45 changes: 45 additions & 0 deletions apps/extension/src/shared/kilo-gateway-chat-stream-client.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -583,4 +583,49 @@ describe('kilo gateway chat stream client', () => {
promptTokens: 1200,
});
});

it('carries costUsd when the usage chunk includes cost', () => {
const sse = [
'data: {"choices":[{"delta":{"content":"hi"}}]}\n\n',
'data: {"choices":[],"usage":{"prompt_tokens":1200,"completion_tokens":34,"total_tokens":1234,"cost":0.0123}}\n\n',
'data: [DONE]\n\n',
].join('');

const completion = parseKiloGatewayChatCompletionStream(sse, () => {});

expect(completion.usage).toStrictEqual({
costUsd: 0.0123,
promptTokens: 1200,
});
});

it('omits costUsd when the usage chunk has no cost field', () => {
const sse = [
'data: {"choices":[{"delta":{"content":"hi"}}]}\n\n',
'data: {"choices":[],"usage":{"prompt_tokens":800,"completion_tokens":10,"total_tokens":810}}\n\n',
'data: [DONE]\n\n',
].join('');

const completion = parseKiloGatewayChatCompletionStream(sse, () => {});

expect(completion.usage).toStrictEqual({
promptTokens: 800,
});
expect(completion.usage).not.toHaveProperty('costUsd');
});

it('parses prompt_tokens when cost is null and omits costUsd', () => {
const sse = [
'data: {"choices":[{"delta":{"content":"hi"}}]}\n\n',
'data: {"choices":[],"usage":{"prompt_tokens":500,"completion_tokens":2,"total_tokens":502,"cost":null}}\n\n',
'data: [DONE]\n\n',
].join('');

const completion = parseKiloGatewayChatCompletionStream(sse, () => {});

expect(completion.usage).toStrictEqual({
promptTokens: 500,
});
expect(completion.usage).not.toHaveProperty('costUsd');
});
});
9 changes: 7 additions & 2 deletions apps/extension/src/shared/kilo-gateway-chat-stream-client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -87,8 +87,9 @@ const streamingToolCallDeltaSchema = z.object({
id: z.string().optional(),
index: z.number(),
});
// Only prompt_tokens is consumed (the context-usage ratio); other usage fields are ignored.
// Prompt_tokens drives the context-usage ratio; cost is optional session spend (USD).
const usageSchema = z.object({
cost: z.number().nullish(),
prompt_tokens: z.number(),
});
const streamDataSchema = z.object({
Expand Down Expand Up @@ -261,7 +262,11 @@ const applyStreamingData = (
}

if (parsed.data.usage !== undefined && parsed.data.usage !== null) {
accumulator.usage = { promptTokens: parsed.data.usage.prompt_tokens };
const { cost, prompt_tokens: promptTokens } = parsed.data.usage;
accumulator.usage = {
promptTokens,
...(typeof cost === 'number' ? { costUsd: cost } : {}),
};
}

const choice = parsed.data.choices.at(0);
Expand Down
31 changes: 31 additions & 0 deletions apps/extension/src/shared/session-cost.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
import { describe, expect, it } from 'vitest';
import { addSessionCost, formatSessionCost } from './session-cost';

describe('session cost helpers', () => {
it('formats zero and non-positive values as $0.0000', () => {
expect(formatSessionCost(0)).toBe('$0.0000');
expect(formatSessionCost(-0.01)).toBe('$0.0000');
expect(formatSessionCost(Number.NaN)).toBe('$0.0000');
expect(formatSessionCost(Number.POSITIVE_INFINITY)).toBe('$0.0000');
});

it('formats positive costs to four decimal places', () => {
expect(formatSessionCost(0.0123)).toBe('$0.0123');
expect(formatSessionCost(1.234)).toBe('$1.2340');
expect(formatSessionCost(5e-5)).toBe('$0.0001');
});

it('adds finite non-negative costs', () => {
expect(addSessionCost(0.0123, 0.0007)).toBeCloseTo(0.013);
expect(addSessionCost(0.013, 0.001)).toBeCloseTo(0.014);
expect(addSessionCost(0, 0)).toBe(0);
});

it('leaves the previous total unchanged for missing, non-finite, or negative cost', () => {
const usageWithoutCost: { costUsd?: number } = {};
expect(addSessionCost(0.5, usageWithoutCost.costUsd)).toBe(0.5);
expect(addSessionCost(0.5, Number.NaN)).toBe(0.5);
expect(addSessionCost(0.5, Number.POSITIVE_INFINITY)).toBe(0.5);
expect(addSessionCost(0.5, -0.01)).toBe(0.5);
});
});
20 changes: 20 additions & 0 deletions apps/extension/src/shared/session-cost.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
/** Format session spend like mobile `formatCost`: always `$X.XXXX`. */
export const formatSessionCost = (cost: number): string => {
if (!Number.isFinite(cost) || cost <= 0) {
return '$0.0000';
}

return `$${cost.toFixed(4)}`;
};

/**
* Add a completion's USD cost to a running session total.
* Missing, non-finite, or negative values leave the previous total unchanged.
*/
export const addSessionCost = (previous: number, costUsd: number | undefined): number => {
if (costUsd === undefined || !Number.isFinite(costUsd) || costUsd < 0) {
return previous;
}

return previous + costUsd;
};
Loading