Skip to content
Open
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
5 changes: 5 additions & 0 deletions src/electron/main/ipc/register-main-ipc-handlers.ts
Original file line number Diff line number Diff line change
Expand Up @@ -205,6 +205,11 @@ export function registerMainIpcHandlers(options: RegisterMainIpcHandlersOptions)
ipcChannels.storageKeys,
async (_event, store: string) => options.resolveStore(store).keys(),
);
ipcMain.handle(
ipcChannels.getUsageSummary,
async (_event, params: { since?: number }) =>
withRuntime((runtimeController) => runtimeController.getUsageSummary(params)),
);

// Shell, dialog, and app handlers.
ipcMain.handle(ipcChannels.copyText, (_event, text: string) => {
Expand Down
135 changes: 135 additions & 0 deletions src/electron/main/runtime/agent-runtime/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -302,6 +302,40 @@ interface AgentSessionTokenTotals {

type AgentTokenUsageSnapshot = AgentMessageUsage;

interface AgentLiteActionHttp {
mintContainerToken(
groupFolder: string,
isMain: boolean,
): { url: string; token: string } | null;
}

type AgentLiteAgentWithActions = AgentLiteAgent & {
actionsHttp?: AgentLiteActionHttp;
};

interface UsageByModelRow {
model: string;
input_tokens: number;
output_tokens: number;
cost_usd: number | null;
}

interface UsageBySessionRow {
session_id: string;
agent_name: string | null;
input_tokens: number;
output_tokens: number;
cost_usd: number | null;
}

interface UsageSummaryResult {
total_tokens: number;
total_cost_usd: number | null;
request_count: number;
by_model: UsageByModelRow[];
by_session: UsageBySessionRow[];
}

function asNonNegativeInteger(value: unknown): number | null {
if (typeof value !== 'number' || !Number.isFinite(value) || value < 0) {
return null;
Expand All @@ -318,6 +352,62 @@ function asFiniteNumber(value: unknown): number | null {
return value;
}

function isUsageSummaryResult(value: unknown): value is UsageSummaryResult {
if (!value || typeof value !== 'object') {
return false;
}

const summary = value as UsageSummaryResult;
return typeof summary.total_tokens === 'number'
&& typeof summary.request_count === 'number'
&& Array.isArray(summary.by_model)
&& Array.isArray(summary.by_session);
}

function addNullableCost(left: number | null, right: number | null): number | null {
if (left === null && right === null) {
return null;
}

return (left ?? 0) + (right ?? 0);
}

function mergeUsageSummaries(summaries: UsageSummaryResult[]): UsageSummaryResult {
const byModel = new Map<string, UsageByModelRow>();
const bySession: UsageBySessionRow[] = [];
let totalTokens = 0;
let totalCostUsd: number | null = null;
let requestCount = 0;

for (const summary of summaries) {
totalTokens += summary.total_tokens;
totalCostUsd = addNullableCost(totalCostUsd, summary.total_cost_usd);
requestCount += summary.request_count;
bySession.push(...summary.by_session);

for (const row of summary.by_model) {
const existing = byModel.get(row.model);

byModel.set(row.model, existing
? {
cost_usd: addNullableCost(existing.cost_usd, row.cost_usd),
input_tokens: existing.input_tokens + row.input_tokens,
model: row.model,
output_tokens: existing.output_tokens + row.output_tokens,
}
: { ...row });
}
}

return {
by_model: [...byModel.values()],
by_session: bySession,
request_count: requestCount,
total_cost_usd: totalCostUsd,
total_tokens: totalTokens,
};
}

function captureAgentTokenUsageSnapshot(
message: unknown,
totals: AgentSessionTokenTotals,
Expand Down Expand Up @@ -731,6 +821,22 @@ export class AgentRuntime implements AgentRuntimeContract {
await this.telegram.refreshRuntimeState({ forceReconnect: true });
}

/** Returns token usage summary from running AgentLite agents. */
async getUsageSummary(params: { since?: number }): Promise<UsageSummaryResult | null> {
const summaries = await Promise.all(
[...this.lifecycle.allRuntimes()].map(([, duneAgent]) =>
this.callAgentLiteAction(duneAgent, 'usage_get_summary', params)),
);
const usableSummaries = summaries.filter((summary): summary is UsageSummaryResult =>
isUsageSummaryResult(summary));

if (usableSummaries.length === 0) {
return null;
}

return mergeUsageSummaries(usableSummaries);
}

/** Resets agent. */
reset() {
this.messageStream.clear();
Expand Down Expand Up @@ -758,6 +864,35 @@ export class AgentRuntime implements AgentRuntimeContract {
void this.telegram.disconnectAll();
}

private async callAgentLiteAction(
duneAgent: DuneAgent,
name: string,
payload: Record<string, unknown>,
): Promise<unknown> {
const actionHost = duneAgent.agentLiteAgent as AgentLiteAgentWithActions;
const binding = actionHost.actionsHttp?.mintContainerToken('main', true);

if (!binding) {
return null;
}

const response = await fetch(`${binding.url}/call`, {
body: JSON.stringify({ name, payload }),
headers: {
Authorization: `Bearer ${binding.token}`,
'Content-Type': 'application/json',
},
method: 'POST',
});

if (!response.ok) {
return null;
}

const body = await response.json() as { result?: unknown };
return body.result ?? null;
}

// -------------------------------------------------------------------------
// Agent CRUD
// -------------------------------------------------------------------------
Expand Down
10 changes: 10 additions & 0 deletions src/electron/main/runtime/desktop-runtime-controller.ts
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ import {

/** Active runtime shape. */
type ActiveRuntime = AgentRuntimeContract & {
getUsageSummary?: (params: { since?: number }) => Promise<unknown>;
reloadExternalChannels?: () => Promise<void>;
shutdown?: () => Promise<void>;
};
Expand Down Expand Up @@ -194,6 +195,15 @@ export class DesktopRuntimeController {
this.activeRuntime.service.selectAgent(agentId);
}

/** Returns token usage summary, or null when AgentLite tracking is unavailable. */
async getUsageSummary(params: { since?: number }): Promise<unknown> {
try {
return await this.activeRuntime.getUsageSummary?.(params) ?? null;
} catch {
return null;
}
}

/** Resets desktop runtime. */
async reset() {
if (typeof this.activeRuntime.reset === 'function') {
Expand Down
1 change: 1 addition & 0 deletions src/electron/preload.ts
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,7 @@ const bridge: DesktopBridge = {
storageKeys: (store) => ipcRenderer.invoke(ipcChannels.storageKeys, store),
storageSet: (store, key, value) => ipcRenderer.invoke(ipcChannels.storageSet, store, key, value),
selectProjectDirectory: () => ipcRenderer.invoke(ipcChannels.selectProjectDirectory),
getUsageSummary: (params) => ipcRenderer.invoke(ipcChannels.getUsageSummary, params),
subscribe: (listener) => {
/** Handles snapshot. */
const handleSnapshot = (
Expand Down
Loading