toad-eye 2.2.0 → 2.4.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (47) hide show
  1. package/README.md +166 -152
  2. package/dist/agent.d.ts +31 -28
  3. package/dist/agent.js +116 -35
  4. package/dist/alerts/channels.js +24 -6
  5. package/dist/alerts/index.js +46 -2
  6. package/dist/alerts/manager.d.ts +1 -0
  7. package/dist/alerts/manager.js +36 -21
  8. package/dist/budget/tracker.d.ts +16 -6
  9. package/dist/budget/tracker.js +38 -17
  10. package/dist/budget/types.d.ts +2 -0
  11. package/dist/cli.d.ts +8 -1
  12. package/dist/cli.js +106 -54
  13. package/dist/core/metrics.d.ts +4 -0
  14. package/dist/core/metrics.js +20 -1
  15. package/dist/core/pricing.d.ts +2 -0
  16. package/dist/core/pricing.js +4 -0
  17. package/dist/core/spans.d.ts +5 -3
  18. package/dist/core/spans.js +80 -21
  19. package/dist/core/tracer.d.ts +6 -3
  20. package/dist/core/tracer.js +49 -7
  21. package/dist/drift/baseline.js +7 -1
  22. package/dist/drift/monitor.js +6 -6
  23. package/dist/index.d.ts +7 -2
  24. package/dist/index.js +6 -1
  25. package/dist/instrumentations/anthropic.d.ts +1 -1
  26. package/dist/instrumentations/anthropic.js +2 -3
  27. package/dist/instrumentations/create.js +80 -19
  28. package/dist/instrumentations/gemini.d.ts +1 -1
  29. package/dist/instrumentations/gemini.js +6 -7
  30. package/dist/instrumentations/openai.d.ts +1 -1
  31. package/dist/instrumentations/openai.js +2 -3
  32. package/dist/instrumentations/registry.d.ts +3 -1
  33. package/dist/instrumentations/registry.js +28 -5
  34. package/dist/instrumentations/types.d.ts +4 -2
  35. package/dist/types/attributes.d.ts +20 -0
  36. package/dist/types/attributes.js +23 -1
  37. package/dist/types/config.d.ts +22 -2
  38. package/dist/types/index.d.ts +3 -3
  39. package/dist/types/metrics.d.ts +2 -0
  40. package/dist/types/metrics.js +3 -0
  41. package/dist/types/providers.d.ts +6 -1
  42. package/dist/types/spans.d.ts +21 -2
  43. package/dist/vercel.d.ts +53 -0
  44. package/dist/vercel.js +106 -0
  45. package/package.json +29 -3
  46. package/templates/grafana/dashboards/agent-workflow.json +207 -0
  47. package/templates/otel-collector.yml +25 -1
package/dist/agent.js CHANGED
@@ -1,30 +1,60 @@
1
1
  import { trace, SpanStatusCode } from "@opentelemetry/api";
2
2
  import { GEN_AI_ATTRS, INSTRUMENTATION_NAME } from "./types/index.js";
3
3
  import { recordAgentSteps, recordAgentToolUsage } from "./core/metrics.js";
4
- import { getConfig } from "./core/tracer.js";
4
+ import { getConfig, shouldEmitDeprecatedAttrs } from "./core/tracer.js";
5
5
  const tracer = trace.getTracer(INSTRUMENTATION_NAME);
6
+ const DEFAULT_MAX_STEPS = 25;
6
7
  /**
7
- * Record a single agent step as a child span (ReAct pattern: think → act → observe → answer).
8
- *
9
- * Creates a short-lived span under the current active context. Call this inside
10
- * {@link traceAgentQuery} (preferred) or any other active span (e.g. traceLLMCall).
11
- *
12
- * - Respects `recordContent` config — content is omitted when recording is disabled.
13
- * - For `act` steps with a `toolName`, automatically increments the `agent.tool_usage` metric.
8
+ * Record a single agent step as a child span (ReAct pattern).
9
+ * Supports: think, act, observe, answer, handoff.
14
10
  */
15
11
  function traceAgentStep(input) {
16
- const span = tracer.startSpan(`agent.step.${input.type}`);
12
+ const spanName = input.type === "act" && input.toolName !== undefined
13
+ ? `execute_tool ${input.toolName}`
14
+ : `gen_ai.agent.step.${input.type}`;
15
+ const span = tracer.startSpan(spanName);
17
16
  const config = getConfig();
18
17
  const recordContent = config?.recordContent !== false;
18
+ const emitDeprecated = shouldEmitDeprecatedAttrs();
19
19
  span.setAttributes({
20
- [GEN_AI_ATTRS.AGENT_STEP_TYPE]: input.type,
21
- [GEN_AI_ATTRS.AGENT_STEP_NUMBER]: input.stepNumber,
20
+ // Canonical toad_eye namespace
21
+ [GEN_AI_ATTRS.TOAD_AGENT_STEP_TYPE]: input.type,
22
+ [GEN_AI_ATTRS.TOAD_AGENT_STEP_NUMBER]: input.stepNumber,
23
+ // Deprecated aliases (skipped when OTEL_SEMCONV_STABILITY_OPT_IN=gen_ai_latest_experimental)
24
+ ...(emitDeprecated && {
25
+ [GEN_AI_ATTRS.AGENT_STEP_TYPE]: input.type,
26
+ [GEN_AI_ATTRS.AGENT_STEP_NUMBER]: input.stepNumber,
27
+ }),
22
28
  ...(input.toolName !== undefined && {
23
- [GEN_AI_ATTRS.AGENT_TOOL_NAME]: input.toolName,
29
+ // OTel spec attribute
30
+ [GEN_AI_ATTRS.TOOL_NAME]: input.toolName,
31
+ // Deprecated alias
32
+ ...(emitDeprecated && {
33
+ [GEN_AI_ATTRS.AGENT_TOOL_NAME]: input.toolName,
34
+ }),
35
+ }),
36
+ ...(input.toolType !== undefined && {
37
+ [GEN_AI_ATTRS.TOOL_TYPE]: input.toolType,
24
38
  }),
25
39
  ...(recordContent &&
26
40
  input.content !== undefined && {
27
- [GEN_AI_ATTRS.AGENT_STEP_CONTENT]: input.content,
41
+ [GEN_AI_ATTRS.TOAD_AGENT_STEP_CONTENT]: input.content,
42
+ ...(emitDeprecated && {
43
+ [GEN_AI_ATTRS.AGENT_STEP_CONTENT]: input.content,
44
+ }),
45
+ }),
46
+ // Handoff attributes
47
+ ...(input.toAgent !== undefined && {
48
+ [GEN_AI_ATTRS.TOAD_AGENT_HANDOFF_TO]: input.toAgent,
49
+ ...(emitDeprecated && {
50
+ [GEN_AI_ATTRS.AGENT_HANDOFF_TO]: input.toAgent,
51
+ }),
52
+ }),
53
+ ...(input.handoffReason !== undefined && {
54
+ [GEN_AI_ATTRS.TOAD_AGENT_HANDOFF_REASON]: input.handoffReason,
55
+ ...(emitDeprecated && {
56
+ [GEN_AI_ATTRS.AGENT_HANDOFF_REASON]: input.handoffReason,
57
+ }),
28
58
  }),
29
59
  });
30
60
  if (input.type === "act" && input.toolName !== undefined) {
@@ -34,51 +64,102 @@ function traceAgentStep(input) {
34
64
  span.end();
35
65
  }
36
66
  /**
37
- * Wrap an entire agent query in a parent span, producing structured ReAct traces.
67
+ * Wrap an entire agent query in a parent span with ReAct tracing.
38
68
  *
39
- * The callback receives a `step` function to record individual steps.
40
- * When the callback completes, the wrapper automatically records the
41
- * `agent.steps_per_query` histogram metric with the total step count.
69
+ * Accepts either a string query (backward-compatible) or an object with
70
+ * agent metadata per OTel GenAI spec (agentName, agentId).
42
71
  *
43
- * Span hierarchy in Jaeger:
44
- * ```
45
- * [parent] agent.query "Is anything dangerous near Earth?"
46
- * ├── agent.step.think "I need to check asteroids"
47
- * ├── agent.step.act tool=near-earth-asteroids
48
- * ├── agent.step.observe "7 asteroids found"
49
- * └── agent.step.answer "7 asteroids passing safely..."
50
- * ```
72
+ * Supports multi-agent via nesting — child traceAgentQuery calls
73
+ * automatically become child spans in Jaeger.
74
+ *
75
+ * Loop detection: counts observe→think transitions. Records loop count
76
+ * as span attribute. Emits warning when maxSteps exceeded.
51
77
  *
52
78
  * @example
53
79
  * ```ts
54
- * const result = await traceAgentQuery("Is anything dangerous near Earth?", async (step) => {
55
- * step({ type: "think", stepNumber: 1, content: "I need to check asteroids" });
56
- * const data = await callTool("near-earth-asteroids");
57
- * step({ type: "act", stepNumber: 2, toolName: "near-earth-asteroids" });
58
- * step({ type: "observe", stepNumber: 3, content: `${data.length} asteroids found` });
59
- * step({ type: "answer", stepNumber: 4, content: "7 asteroids passing safely..." });
60
- * return { answer: "7 asteroids passing safely..." };
80
+ * // Simple form (backward-compatible)
81
+ * await traceAgentQuery('Is anything dangerous?', async (step) => { ... });
82
+ *
83
+ * // Object form with agent metadata (OTel GenAI spec)
84
+ * await traceAgentQuery(
85
+ * { query: 'Is anything dangerous?', agentName: 'space-monitor', agentId: 'agent-001' },
86
+ * async (step) => { ... }
87
+ * );
88
+ *
89
+ * // Multi-agent: orchestrator delegates to specialist
90
+ * await traceAgentQuery({ query: 'Analyze data', agentName: 'orchestrator' }, async (step) => {
91
+ * step({ type: 'think', stepNumber: 1, content: 'Need domain expert' });
92
+ * step({ type: 'handoff', stepNumber: 2, toAgent: 'specialist', handoffReason: 'domain expertise' });
93
+ *
94
+ * const result = await traceAgentQuery({ query: 'Analyze', agentName: 'specialist' }, async (step) => {
95
+ * step({ type: 'act', stepNumber: 1, toolName: 'analyze', toolType: 'function' });
96
+ * step({ type: 'answer', stepNumber: 2, content: 'analysis complete' });
97
+ * return { answer: 'done' };
98
+ * });
61
99
  * });
62
100
  * ```
63
101
  */
64
- export async function traceAgentQuery(query, fn) {
65
- return tracer.startActiveSpan(`agent.query`, async (span) => {
102
+ export async function traceAgentQuery(queryOrInput, fn, options) {
103
+ const resolved = typeof queryOrInput === "string"
104
+ ? { query: queryOrInput, agentName: undefined, agentId: undefined }
105
+ : queryOrInput;
106
+ const spanName = resolved.agentName
107
+ ? `invoke_agent ${resolved.agentName}`
108
+ : `invoke_agent`;
109
+ return tracer.startActiveSpan(spanName, async (span) => {
66
110
  const config = getConfig();
67
111
  const recordContent = config?.recordContent !== false;
112
+ const emitDeprecated = shouldEmitDeprecatedAttrs();
113
+ const maxSteps = options?.maxSteps ?? DEFAULT_MAX_STEPS;
114
+ // OTel GenAI agent attributes
115
+ if (resolved.agentName !== undefined) {
116
+ span.setAttribute(GEN_AI_ATTRS.AGENT_NAME, resolved.agentName);
117
+ }
118
+ if (resolved.agentId !== undefined) {
119
+ span.setAttribute(GEN_AI_ATTRS.AGENT_ID, resolved.agentId);
120
+ }
68
121
  if (recordContent) {
69
- span.setAttribute(GEN_AI_ATTRS.AGENT_STEP_CONTENT, query);
122
+ span.setAttribute(GEN_AI_ATTRS.TOAD_AGENT_STEP_CONTENT, resolved.query);
123
+ if (emitDeprecated) {
124
+ span.setAttribute(GEN_AI_ATTRS.AGENT_STEP_CONTENT, resolved.query);
125
+ }
70
126
  }
71
127
  let stepCount = 0;
128
+ let loopCount = 0;
129
+ let lastStepType;
130
+ let maxStepsWarned = false;
72
131
  try {
73
132
  const result = await fn((input) => {
74
133
  stepCount++;
134
+ // Loop detection: observe → think = one loop iteration
135
+ if (input.type === "think" && lastStepType === "observe") {
136
+ loopCount++;
137
+ }
138
+ lastStepType = input.type;
139
+ // Max steps guard
140
+ if (stepCount > maxSteps && !maxStepsWarned) {
141
+ maxStepsWarned = true;
142
+ console.warn(`toad-eye: agent exceeded maxSteps (${maxSteps}). Current: ${stepCount}`);
143
+ span.addEvent("agent.max_steps_exceeded", {
144
+ "agent.max_steps": maxSteps,
145
+ "agent.current_steps": stepCount,
146
+ });
147
+ }
75
148
  traceAgentStep(input);
76
149
  });
150
+ span.setAttribute(GEN_AI_ATTRS.TOAD_AGENT_LOOP_COUNT, loopCount);
151
+ if (emitDeprecated) {
152
+ span.setAttribute(GEN_AI_ATTRS.AGENT_LOOP_COUNT, loopCount);
153
+ }
77
154
  span.setStatus({ code: SpanStatusCode.OK });
78
155
  return result;
79
156
  }
80
157
  catch (error) {
81
158
  const message = error instanceof Error ? error.message : String(error);
159
+ span.setAttribute(GEN_AI_ATTRS.TOAD_AGENT_LOOP_COUNT, loopCount);
160
+ if (emitDeprecated) {
161
+ span.setAttribute(GEN_AI_ATTRS.AGENT_LOOP_COUNT, loopCount);
162
+ }
82
163
  span.setStatus({ code: SpanStatusCode.ERROR, message });
83
164
  throw error;
84
165
  }
@@ -1,4 +1,10 @@
1
- import nodemailer from "nodemailer";
1
+ // Cache nodemailer transporters by connection identity to avoid creating a new
2
+ // SMTP connection for every alert (frequent alerts would exhaust connection pools).
3
+ // eslint-disable-next-line @typescript-eslint/no-explicit-any
4
+ const smtpTransporterCache = new Map();
5
+ function smtpCacheKey(config) {
6
+ return `${config.host}:${config.port}:${config.user}`;
7
+ }
2
8
  function formatMessage(alert) {
3
9
  const topModelsText = alert.topModels.length > 0
4
10
  ? `\nTop models by spend:\n${alert.topModels.map((m) => ` • ${m.model}: $${m.value.toFixed(4)}`).join("\n")}`
@@ -52,11 +58,23 @@ async function sendWebhook(config, alert) {
52
58
  throw new Error(`Webhook delivery failed: ${res.status}`);
53
59
  }
54
60
  async function sendEmail(config, alert) {
55
- const transporter = nodemailer.createTransport({
56
- host: config.host,
57
- port: config.port,
58
- auth: { user: config.user, pass: config.password },
59
- });
61
+ const key = smtpCacheKey(config);
62
+ let transporter = smtpTransporterCache.get(key);
63
+ if (!transporter) {
64
+ let createTransport;
65
+ try {
66
+ ({ createTransport } = await import("nodemailer"));
67
+ }
68
+ catch {
69
+ throw new Error("toad-eye: email alerts require nodemailer — install it: npm install nodemailer");
70
+ }
71
+ transporter = createTransport({
72
+ host: config.host,
73
+ port: config.port,
74
+ auth: { user: config.user, pass: config.password },
75
+ });
76
+ smtpTransporterCache.set(key, transporter);
77
+ }
60
78
  const to = typeof config.to === "string"
61
79
  ? config.to
62
80
  : Array.from(config.to).join(", ");
@@ -3,10 +3,54 @@ export { parseCondition } from "./conditions.js";
3
3
  import { readFileSync } from "node:fs";
4
4
  import yaml from "yaml";
5
5
  import { AlertManager } from "./manager.js";
6
+ /** Resolve ${ENV_VAR} references in string values throughout the config. */
7
+ function resolveEnvVars(value) {
8
+ return value.replace(/\$\{([^}]+)\}/g, (_, name) => {
9
+ return process.env[name] ?? `\${${name}}`;
10
+ });
11
+ }
12
+ function interpolateConfig(obj) {
13
+ if (typeof obj === "string")
14
+ return resolveEnvVars(obj);
15
+ if (Array.isArray(obj))
16
+ return obj.map(interpolateConfig);
17
+ if (obj !== null && typeof obj === "object") {
18
+ return Object.fromEntries(Object.entries(obj).map(([k, v]) => [
19
+ k,
20
+ interpolateConfig(v),
21
+ ]));
22
+ }
23
+ return obj;
24
+ }
25
+ function validateConfig(config) {
26
+ if (typeof config !== "object" || config === null) {
27
+ throw new Error("toad-eye: alerts config must be an object");
28
+ }
29
+ const c = config;
30
+ if (!Array.isArray(c["alerts"])) {
31
+ throw new Error('toad-eye: alerts config missing required "alerts" array');
32
+ }
33
+ for (const rule of c["alerts"]) {
34
+ if (typeof rule !== "object" || rule === null) {
35
+ throw new Error("toad-eye: each alert rule must be an object");
36
+ }
37
+ const r = rule;
38
+ for (const field of ["name", "metric", "condition"]) {
39
+ if (typeof r[field] !== "string" || !r[field]) {
40
+ throw new Error(`toad-eye: alert rule missing required string field "${field}"`);
41
+ }
42
+ }
43
+ if (!Array.isArray(r["channels"]) || r["channels"].length === 0) {
44
+ throw new Error(`toad-eye: alert rule "${String(r["name"])}" must have at least one channel`);
45
+ }
46
+ }
47
+ }
6
48
  export function startAlertsFromFile(configPath) {
7
49
  const raw = readFileSync(configPath, "utf-8");
8
- const config = yaml.parse(raw);
9
- const manager = new AlertManager(config);
50
+ const parsed = yaml.parse(raw);
51
+ const interpolated = interpolateConfig(parsed);
52
+ validateConfig(interpolated);
53
+ const manager = new AlertManager(interpolated);
10
54
  manager.start();
11
55
  return manager;
12
56
  }
@@ -8,6 +8,7 @@ export declare class AlertManager {
8
8
  stop(): void;
9
9
  private isInCooldown;
10
10
  private evaluate;
11
+ private evaluateRule;
11
12
  private fireChannels;
12
13
  private postAnnotation;
13
14
  }
@@ -5,6 +5,7 @@ const DEFAULT_PROMETHEUS_URL = "http://localhost:9090";
5
5
  const DEFAULT_GRAFANA_URL = "http://localhost:3100";
6
6
  const DEFAULT_EVAL_INTERVAL_SECONDS = 60;
7
7
  const DEFAULT_COOLDOWN_MINUTES = 30;
8
+ const MIN_EVAL_INTERVAL_SECONDS = 10;
8
9
  export class AlertManager {
9
10
  config;
10
11
  cooldowns = new Map();
@@ -13,10 +14,15 @@ export class AlertManager {
13
14
  this.config = config;
14
15
  }
15
16
  start() {
16
- const intervalMs = (this.config.evalIntervalSeconds ?? DEFAULT_EVAL_INTERVAL_SECONDS) * 1000;
17
+ const configured = this.config.evalIntervalSeconds ?? DEFAULT_EVAL_INTERVAL_SECONDS;
18
+ const clampedSeconds = Math.max(configured, MIN_EVAL_INTERVAL_SECONDS);
19
+ if (configured < MIN_EVAL_INTERVAL_SECONDS) {
20
+ console.warn(`[toad-eye alerts] evalIntervalSeconds ${configured} is below minimum ${MIN_EVAL_INTERVAL_SECONDS}s — using ${clampedSeconds}s`);
21
+ }
22
+ const intervalMs = clampedSeconds * 1000;
17
23
  this.intervalId = setInterval(() => void this.evaluate(), intervalMs);
18
24
  void this.evaluate();
19
- console.log(`[toad-eye alerts] Started — ${this.config.alerts.length} rule(s), eval every ${intervalMs / 1000}s`);
25
+ console.log(`[toad-eye alerts] Started — ${this.config.alerts.length} rule(s), eval every ${clampedSeconds}s`);
20
26
  }
21
27
  stop() {
22
28
  if (this.intervalId !== undefined) {
@@ -38,26 +44,35 @@ export class AlertManager {
38
44
  async evaluate() {
39
45
  const prometheusUrl = this.config.prometheusUrl ?? DEFAULT_PROMETHEUS_URL;
40
46
  const grafanaUrl = this.config.grafanaUrl ?? DEFAULT_GRAFANA_URL;
41
- for (const rule of this.config.alerts) {
42
- if (this.isInCooldown(rule))
43
- continue;
44
- const result = await evaluateCondition(prometheusUrl, rule);
45
- if (!result?.triggered)
46
- continue;
47
- this.cooldowns.set(rule.name, Date.now());
48
- const firedAlert = {
49
- rule,
50
- value: result.value,
51
- threshold: result.threshold,
52
- topModels: result.topModels,
53
- firedAt: new Date(),
54
- };
55
- console.log(`[toad-eye alerts] 🚨 "${rule.name}" fired — value=${result.value.toFixed(4)}, threshold=${result.threshold}`);
56
- await Promise.allSettled([
57
- this.fireChannels(firedAlert),
58
- this.postAnnotation(grafanaUrl, firedAlert.rule.name, firedAlert.value, firedAlert.threshold),
59
- ]);
47
+ // Evaluate all rules in parallel — each rule is independent
48
+ await Promise.allSettled(this.config.alerts.map((rule) => this.evaluateRule(prometheusUrl, grafanaUrl, rule)));
49
+ }
50
+ async evaluateRule(prometheusUrl, grafanaUrl, rule) {
51
+ if (this.isInCooldown(rule))
52
+ return;
53
+ let result;
54
+ try {
55
+ result = await evaluateCondition(prometheusUrl, rule);
56
+ }
57
+ catch (err) {
58
+ console.error(`[toad-eye alerts] Failed to evaluate rule "${rule.name}":`, err);
59
+ return;
60
60
  }
61
+ if (!result?.triggered)
62
+ return;
63
+ this.cooldowns.set(rule.name, Date.now());
64
+ const firedAlert = {
65
+ rule,
66
+ value: result.value,
67
+ threshold: result.threshold,
68
+ topModels: result.topModels,
69
+ firedAt: new Date(),
70
+ };
71
+ console.log(`[toad-eye alerts] 🚨 "${rule.name}" fired — value=${result.value.toFixed(4)}, threshold=${result.threshold}`);
72
+ await Promise.allSettled([
73
+ this.fireChannels(firedAlert),
74
+ this.postAnnotation(grafanaUrl, firedAlert.rule.name, firedAlert.value, firedAlert.threshold),
75
+ ]);
61
76
  }
62
77
  async fireChannels(alert) {
63
78
  const channels = this.config.channels ?? {};
@@ -9,18 +9,28 @@ export declare class BudgetTracker {
9
9
  private resetIfNewDay;
10
10
  /**
11
11
  * Check budget BEFORE making an LLM call.
12
- * Returns modified provider/model if downgrade mode triggers,
13
- * throws ToadBudgetExceededError if block mode triggers,
14
- * returns null if no action needed.
12
+ * - block mode: throws ToadBudgetExceededError if budget is exceeded.
13
+ * - downgrade mode: returns modified provider/model + budget type that triggered it.
14
+ * - warn mode: returns null (warning emitted after call via recordCost).
15
+ *
16
+ * When estimatedCost is provided, it is added to reservedCost so concurrent
17
+ * in-flight requests are counted against the budget before their actual cost is known.
18
+ * Call recordCost with the same estimatedCost to release the reservation.
15
19
  */
16
- checkBefore(provider: string, model: string, userId?: string): {
20
+ checkBefore(provider: string, model: string, userId?: string, estimatedCost?: number): {
17
21
  provider: string;
18
22
  model: string;
23
+ budget: BudgetExceededInfo["budget"];
19
24
  } | null;
20
- /** Record cost AFTER a successful LLM call. Returns exceeded info if budget was just crossed. */
21
- recordCost(cost: number, model: string, userId?: string): BudgetExceededInfo | null;
25
+ /**
26
+ * Record cost AFTER a successful LLM call. Returns exceeded info if budget was just crossed.
27
+ * Pass the same estimatedCost used in checkBefore() to release the reservation atomically.
28
+ */
29
+ recordCost(cost: number, model: string, userId?: string, reservedAmount?: number): BudgetExceededInfo | null;
22
30
  /** Find which budget (if any) is exceeded. */
23
31
  private findExceeded;
32
+ /** Release a cost reservation when an LLM call fails (no cost was incurred). */
33
+ releaseReservation(reservedAmount: number): void;
24
34
  /** Get current budget usage as percentage (0-100+). */
25
35
  getUsagePercent(): number;
26
36
  /** Get current state snapshot (for persistence or debugging). */
@@ -14,6 +14,7 @@ export class BudgetTracker {
14
14
  this.state = {
15
15
  date: todayUTC(),
16
16
  totalCost: 0,
17
+ reservedCost: 0,
17
18
  perUser: new Map(),
18
19
  perModel: new Map(),
19
20
  };
@@ -25,6 +26,7 @@ export class BudgetTracker {
25
26
  this.state = {
26
27
  date: today,
27
28
  totalCost: 0,
29
+ reservedCost: 0,
28
30
  perUser: new Map(),
29
31
  perModel: new Map(),
30
32
  };
@@ -32,27 +34,40 @@ export class BudgetTracker {
32
34
  }
33
35
  /**
34
36
  * Check budget BEFORE making an LLM call.
35
- * Returns modified provider/model if downgrade mode triggers,
36
- * throws ToadBudgetExceededError if block mode triggers,
37
- * returns null if no action needed.
37
+ * - block mode: throws ToadBudgetExceededError if budget is exceeded.
38
+ * - downgrade mode: returns modified provider/model + budget type that triggered it.
39
+ * - warn mode: returns null (warning emitted after call via recordCost).
40
+ *
41
+ * When estimatedCost is provided, it is added to reservedCost so concurrent
42
+ * in-flight requests are counted against the budget before their actual cost is known.
43
+ * Call recordCost with the same estimatedCost to release the reservation.
38
44
  */
39
- checkBefore(provider, model, userId) {
45
+ checkBefore(provider, model, userId, estimatedCost = 0) {
40
46
  this.resetIfNewDay();
41
- const exceeded = this.findExceeded(model, userId);
42
- if (!exceeded)
43
- return null;
44
- if (this.mode === "block") {
45
- throw new ToadBudgetExceededError(exceeded);
46
- }
47
- if (this.mode === "downgrade" && this.downgrade) {
48
- return this.downgrade({ provider, model });
47
+ const exceeded = this.findExceeded(model, userId, estimatedCost);
48
+ if (exceeded) {
49
+ if (this.mode === "block") {
50
+ throw new ToadBudgetExceededError(exceeded);
51
+ }
52
+ if (this.mode === "downgrade" && this.downgrade) {
53
+ this.state.reservedCost += estimatedCost;
54
+ return {
55
+ ...this.downgrade({ provider, model }),
56
+ budget: exceeded.budget,
57
+ };
58
+ }
49
59
  }
50
- // 'warn' mode — continue, warning emitted after call
60
+ // warn mode or no budget exceeded: reserve estimated cost and proceed
61
+ this.state.reservedCost += estimatedCost;
51
62
  return null;
52
63
  }
53
- /** Record cost AFTER a successful LLM call. Returns exceeded info if budget was just crossed. */
54
- recordCost(cost, model, userId) {
64
+ /**
65
+ * Record cost AFTER a successful LLM call. Returns exceeded info if budget was just crossed.
66
+ * Pass the same estimatedCost used in checkBefore() to release the reservation atomically.
67
+ */
68
+ recordCost(cost, model, userId, reservedAmount = 0) {
55
69
  this.resetIfNewDay();
70
+ this.state.reservedCost = Math.max(0, this.state.reservedCost - reservedAmount);
56
71
  this.state.totalCost += cost;
57
72
  this.state.perModel.set(model, (this.state.perModel.get(model) ?? 0) + cost);
58
73
  if (userId) {
@@ -61,9 +76,10 @@ export class BudgetTracker {
61
76
  return this.findExceeded(model, userId);
62
77
  }
63
78
  /** Find which budget (if any) is exceeded. */
64
- findExceeded(model, userId) {
79
+ findExceeded(model, userId, additionalCost = 0) {
65
80
  if (this.config.daily !== undefined &&
66
- this.state.totalCost >= this.config.daily) {
81
+ this.state.totalCost + this.state.reservedCost + additionalCost >=
82
+ this.config.daily) {
67
83
  return {
68
84
  budget: "daily",
69
85
  limit: this.config.daily,
@@ -95,6 +111,10 @@ export class BudgetTracker {
95
111
  }
96
112
  return null;
97
113
  }
114
+ /** Release a cost reservation when an LLM call fails (no cost was incurred). */
115
+ releaseReservation(reservedAmount) {
116
+ this.state.reservedCost = Math.max(0, this.state.reservedCost - reservedAmount);
117
+ }
98
118
  /** Get current budget usage as percentage (0-100+). */
99
119
  getUsagePercent() {
100
120
  if (this.config.daily === undefined || this.config.daily === 0)
@@ -114,6 +134,7 @@ export class BudgetTracker {
114
134
  this.state = {
115
135
  date: saved.date,
116
136
  totalCost: saved.totalCost,
137
+ reservedCost: 0,
117
138
  perUser: new Map(Object.entries(saved.perUser)),
118
139
  perModel: new Map(Object.entries(saved.perModel)),
119
140
  };
@@ -21,6 +21,8 @@ export type DowngradeCallback = (original: {
21
21
  export interface BudgetState {
22
22
  date: string;
23
23
  totalCost: number;
24
+ /** Sum of estimated costs for in-flight requests (released when recordCost is called). */
25
+ reservedCost: number;
24
26
  perUser: Map<string, number>;
25
27
  perModel: Map<string, number>;
26
28
  }
package/dist/cli.d.ts CHANGED
@@ -1,2 +1,9 @@
1
1
  #!/usr/bin/env node
2
- export {};
2
+ /** Find the docker container matching a service by exact service name. Exported for testing. */
3
+ export declare function findContainerByService(containers: ReadonlyArray<{
4
+ Service: string;
5
+ State: string;
6
+ }>, serviceKey: string): {
7
+ Service: string;
8
+ State: string;
9
+ } | undefined;