126 lines
4.3 KiB
TypeScript
126 lines
4.3 KiB
TypeScript
import type { AgentRole, TaskStore } from "@fusion/core";
|
|
import type { AgentSession } from "@mariozechner/pi-coding-agent";
|
|
import { createLogger } from "./logger.js";
|
|
|
|
const log = createLogger("session-token-usage");
|
|
const cacheMetricsLog = createLogger("token-cache-metrics");
|
|
|
|
interface SessionBaseline {
|
|
input: number;
|
|
output: number;
|
|
cached: number;
|
|
cacheWrite: number;
|
|
}
|
|
|
|
// Per-session cumulative-token baselines so repeated calls only persist deltas.
|
|
// The session object is keyed weakly so disposed sessions get garbage-collected.
|
|
const sessionBaselines = new WeakMap<AgentSession, SessionBaseline>();
|
|
|
|
interface SessionStatsLike {
|
|
tokens?: {
|
|
input?: number;
|
|
output?: number;
|
|
cacheRead?: number;
|
|
cacheWrite?: number;
|
|
};
|
|
}
|
|
|
|
function readSessionStats(session: AgentSession): SessionStatsLike | undefined {
|
|
const accessor = (session as unknown as { getSessionStats?: () => SessionStatsLike }).getSessionStats;
|
|
if (typeof accessor !== "function") return undefined;
|
|
try {
|
|
return accessor.call(session);
|
|
} catch {
|
|
return undefined;
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Capture the session's cumulative token usage and accumulate any *new* deltas
|
|
* onto `task.tokenUsage`. Safe to call repeatedly on the same session — each
|
|
* call only persists what's been added since the previous call (per-session
|
|
* baseline tracking). Failures are logged and swallowed so token bookkeeping
|
|
* never blocks the task pipeline.
|
|
*/
|
|
export async function accumulateSessionTokenUsage(
|
|
store: TaskStore,
|
|
taskId: string,
|
|
session: AgentSession,
|
|
options?: { agentId?: string; role?: AgentRole },
|
|
): Promise<void> {
|
|
try {
|
|
const stats = readSessionStats(session);
|
|
const tokens = stats?.tokens;
|
|
if (!tokens) return;
|
|
|
|
const currentInput = tokens.input ?? 0;
|
|
const currentOutput = tokens.output ?? 0;
|
|
const currentCached = tokens.cacheRead ?? 0;
|
|
const currentCacheWrite = tokens.cacheWrite ?? 0;
|
|
|
|
const baseline = sessionBaselines.get(session) ?? { input: 0, output: 0, cached: 0, cacheWrite: 0 };
|
|
const inputDelta = Math.max(0, currentInput - baseline.input);
|
|
const outputDelta = Math.max(0, currentOutput - baseline.output);
|
|
const cachedDelta = Math.max(0, currentCached - baseline.cached);
|
|
const cacheWriteDelta = Math.max(0, currentCacheWrite - baseline.cacheWrite);
|
|
|
|
sessionBaselines.set(session, {
|
|
input: currentInput,
|
|
output: currentOutput,
|
|
cached: currentCached,
|
|
cacheWrite: currentCacheWrite,
|
|
});
|
|
|
|
if (inputDelta === 0 && outputDelta === 0 && cachedDelta === 0 && cacheWriteDelta === 0) return;
|
|
|
|
const task = await store.getTask(taskId);
|
|
const now = new Date().toISOString();
|
|
const newInput = (task.tokenUsage?.inputTokens ?? 0) + inputDelta;
|
|
const newOutput = (task.tokenUsage?.outputTokens ?? 0) + outputDelta;
|
|
const newCached = (task.tokenUsage?.cachedTokens ?? 0) + cachedDelta;
|
|
const newCacheWrite = (task.tokenUsage?.cacheWriteTokens ?? 0) + cacheWriteDelta;
|
|
|
|
const role = options?.role ?? "executor";
|
|
const tokenUsage = {
|
|
inputTokens: newInput,
|
|
outputTokens: newOutput,
|
|
cachedTokens: newCached,
|
|
cacheWriteTokens: newCacheWrite,
|
|
totalTokens: newInput + newOutput + newCached + newCacheWrite,
|
|
firstUsedAt: task.tokenUsage?.firstUsedAt ?? now,
|
|
lastUsedAt: now,
|
|
};
|
|
|
|
cacheMetricsLog.log(JSON.stringify({
|
|
taskId,
|
|
agentId: options?.agentId,
|
|
role,
|
|
inputTokens: tokenUsage.inputTokens,
|
|
cachedTokens: tokenUsage.cachedTokens,
|
|
cacheWriteTokens: tokenUsage.cacheWriteTokens,
|
|
hitRatio: computeCacheHitRatio(tokenUsage.inputTokens, tokenUsage.cachedTokens),
|
|
}));
|
|
|
|
await store.updateTask(taskId, { tokenUsage });
|
|
} catch (err) {
|
|
const message = err instanceof Error ? err.message : String(err);
|
|
log.warn(`${taskId}: session token usage accumulate failed: ${message}`);
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Compute the cache hit ratio: `cachedTokens / (inputTokens + cachedTokens)`.
|
|
* Returns a number in [0, 1], or 0 when both arguments are 0.
|
|
*
|
|
* Compatible with canonical stored `task.tokenUsage` fields: pass raw
|
|
* `inputTokens` and cache-read `cachedTokens`.
|
|
*/
|
|
export function computeCacheHitRatio(
|
|
inputTokens: number,
|
|
cachedTokens: number,
|
|
): number {
|
|
const total = inputTokens + cachedTokens;
|
|
if (total === 0) return 0;
|
|
return cachedTokens / total;
|
|
}
|