naisys
Version:
NAISYS - Autonomous AI agent runner with built-in context management and cost tracking
265 lines • 10.6 kB
JavaScript
import { calculatePeriodBoundaries } from "@naisys/common";
import { CostControlSchema, HubEvents, } from "@naisys/hub-protocol";
export const SPEND_LIMIT_TIMEOUT_SECONDS = 60;
export class SpendLimitError extends Error {
constructor(message) {
super(message);
this.name = "SpendLimitError";
}
}
export function createCostTracker({ globalConfig }, { agentConfig }, modelService, runService, hubClient, hubCostBuffer, localUserId, promptNotification,
/** When set, this tracker is for an ephemeral subagent: spend-limit checks
* defer to the parent and costs roll into the parent's accounting via
* addSubagentCost. */
parentCostTracker) {
// In-memory per-model aggregated costs (always maintained, both modes)
const modelCosts = new Map();
// Running total cost for this agent
let totalCost = 0;
// Period tracking for local mode spend limits
let periodCost = 0;
let currentPeriodEnd = 0;
// Hub mode: budget tracking from COST_WRITE responses
let hubBudgetLeft = null;
let costSinceLastSnapshot = 0;
// Children that need to be woken when this tracker's suspension lifts.
const subagentSubscriptions = new Set();
// Register for budget updates from the shared buffer
if (hubCostBuffer) {
hubCostBuffer.registerBudgetCallback(localUserId, (budgetLeft) => {
hubBudgetLeft = budgetLeft;
costSinceLastSnapshot = 0;
});
}
// Hub mode: receive cost control messages from hub
let hubCostControlReason;
if (hubClient) {
hubClient.registerEvent(HubEvents.COST_CONTROL, (data) => {
const parsed = CostControlSchema.parse(data);
if (parsed.userId !== localUserId)
return;
if (parsed.enabled) {
hubCostControlReason = undefined;
promptNotification.notify({
wake: "always",
userId: localUserId,
commentOutput: ["Cost control suspension lifted, resuming"],
});
// Wake any subagents that were paused on the parent's suspension
for (const subagentUserId of subagentSubscriptions) {
promptNotification.notify({
wake: "always",
userId: subagentUserId,
commentOutput: ["Cost control suspension lifted, resuming"],
});
}
}
else {
hubCostControlReason = parsed.reason;
}
});
}
if (parentCostTracker) {
parentCostTracker.subscribeSubagent(localUserId);
}
function updateInMemory(modelKey, cost, inputTokens, outputTokens, cacheWriteTokens, cacheReadTokens) {
const existing = modelCosts.get(modelKey);
if (existing) {
existing.cost += cost;
existing.inputTokens += inputTokens;
existing.outputTokens += outputTokens;
existing.cacheWriteTokens += cacheWriteTokens;
existing.cacheReadTokens += cacheReadTokens;
}
else {
modelCosts.set(modelKey, {
cost,
inputTokens,
outputTokens,
cacheWriteTokens,
cacheReadTokens,
});
}
totalCost += cost;
costSinceLastSnapshot += cost;
// Subagents still update periodCost for the ns-cost period display; spend-limit/budget checks defer to the parent.
addCostToPeriod(cost);
}
function pushToBuffer(source, modelKey, cost, inputTokens, outputTokens, cacheWriteTokens, cacheReadTokens) {
hubCostBuffer.pushEntry({
userId: localUserId,
runId: runService.getRunId(),
sessionId: runService.getSessionId(),
source,
model: modelKey,
cost,
inputTokens,
outputTokens,
cacheWriteTokens,
cacheReadTokens,
});
}
// Record token usage for LLM calls
function recordTokens(source, modelKey, inputTokens = 0, outputTokens = 0, cacheWriteTokens = 0, cacheReadTokens = 0) {
const model = modelService.getLlmModel(modelKey);
const tokenUsage = {
inputTokens,
outputTokens,
cacheWriteTokens,
cacheReadTokens,
};
const cost = calculateCostFromTokens(tokenUsage, model);
updateInMemory(modelKey, cost, inputTokens, outputTokens, cacheWriteTokens, cacheReadTokens);
if (hubCostBuffer) {
pushToBuffer(source, modelKey, cost, inputTokens, outputTokens, cacheWriteTokens, cacheReadTokens);
}
parentCostTracker?.addSubagentCost(cost);
}
// Record fixed cost for non-token services like image generation
function recordCost(cost, source, modelKey) {
updateInMemory(modelKey, cost, 0, 0, 0, 0);
if (hubCostBuffer) {
pushToBuffer(source, modelKey, cost, 0, 0, 0, 0);
}
parentCostTracker?.addSubagentCost(cost);
}
/** Roll a subagent's cost into this tracker's running totals. The per-model
* breakdown stays scoped to costs the agent spent directly. */
function addSubagentCost(cost) {
totalCost += cost;
costSinceLastSnapshot += cost;
addCostToPeriod(cost);
}
// Common function to calculate cost from token usage
function calculateCostFromTokens(tokenUsage, model) {
const inputCost = (tokenUsage.inputTokens * model.inputCost) / 1_000_000;
const outputCost = (tokenUsage.outputTokens * model.outputCost) / 1_000_000;
const cacheWriteCost = (tokenUsage.cacheWriteTokens * (model.cacheWriteCost || 0)) / 1_000_000;
const cacheReadCost = (tokenUsage.cacheReadTokens * (model.cacheReadCost || 0)) / 1_000_000;
return inputCost + outputCost + cacheWriteCost + cacheReadCost;
}
function addCostToPeriod(cost) {
const now = Date.now();
if (now >= currentPeriodEnd) {
// Period rolled over, reset
const spendLimitHours = agentConfig().spendLimitHours || globalConfig().spendLimitHours;
if (spendLimitHours !== undefined) {
const { periodEnd } = calculatePeriodBoundaries(spendLimitHours);
currentPeriodEnd = periodEnd.getTime();
}
periodCost = 0;
}
periodCost += cost;
}
// Check if the current spend limit has been reached and throw an error if so
// In hub mode, checks the cost control state received from the hub
function checkSpendLimit() {
if (parentCostTracker) {
parentCostTracker.checkSpendLimit();
return;
}
if (hubClient) {
if (!hubClient.isConnected()) {
throw new SpendLimitError("LLM Spend limit check failed: not connected to hub");
}
else if (hubCostControlReason) {
throw new SpendLimitError(`LLM ${hubCostControlReason}`);
}
return;
}
// Else we check local in-memory spend limits
const spendLimitHours = agentConfig().spendLimitHours || globalConfig().spendLimitHours;
const spendLimit = agentConfig().spendLimitDollars || globalConfig().spendLimitDollars || -1;
let currentCost;
let periodDescription;
if (spendLimitHours !== undefined) {
const { periodStart, periodEnd } = calculatePeriodBoundaries(spendLimitHours);
// Ensure period is current
const now = Date.now();
if (now >= currentPeriodEnd) {
currentPeriodEnd = periodEnd.getTime();
periodCost = 0;
}
currentCost = periodCost;
const formatTime = (date) => {
return date.toLocaleTimeString("en-US", {
hour: "2-digit",
minute: "2-digit",
hour12: true,
});
};
periodDescription = `per ${spendLimitHours} hour${spendLimitHours !== 1 ? "s" : ""} (current period: ${formatTime(periodStart)} - ${formatTime(periodEnd)})`;
}
else {
currentCost = totalCost;
periodDescription = "total";
}
if (spendLimit < currentCost) {
const userDescription = agentConfig().spendLimitDollars
? `${agentConfig().username}`
: "all users";
throw new SpendLimitError(`LLM Spend limit of $${spendLimit} ${periodDescription} reached for ${userDescription}, current cost $${currentCost.toFixed(2)}`);
}
}
function cleanup() {
if (hubCostBuffer) {
hubCostBuffer.unregisterBudgetCallback(localUserId);
}
if (parentCostTracker) {
parentCostTracker.unsubscribeSubagent(localUserId);
}
}
function getCostControlReason() {
return hubCostControlReason;
}
function subscribeSubagent(userId) {
subagentSubscriptions.add(userId);
}
function unsubscribeSubagent(userId) {
subagentSubscriptions.delete(userId);
}
// Exposed for costDisplayService
function getModelCosts() {
return modelCosts;
}
function getTotalCost() {
return totalCost;
}
function getPeriodInfo() {
const spendLimitHours = agentConfig().spendLimitHours || globalConfig().spendLimitHours;
if (spendLimitHours === undefined)
return null;
const { periodStart, periodEnd } = calculatePeriodBoundaries(spendLimitHours);
return { periodCost, periodStart, periodEnd };
}
function resetCosts() {
modelCosts.clear();
totalCost = 0;
periodCost = 0;
}
/** Returns the estimated remaining budget, or null if no per-agent limit */
function getBudgetLeft() {
if (parentCostTracker)
return parentCostTracker.getBudgetLeft();
if (hubBudgetLeft === null)
return null;
return Math.max(0, hubBudgetLeft - costSinceLastSnapshot);
}
return {
recordTokens,
recordCost,
calculateCostFromTokens,
checkSpendLimit,
cleanup,
getModelCosts,
getTotalCost,
getPeriodInfo,
getBudgetLeft,
resetCosts,
addSubagentCost,
getCostControlReason,
subscribeSubagent,
unsubscribeSubagent,
};
}
//# sourceMappingURL=costTracker.js.map