# DD-24 — Lab Specification: Build a Cost-Aware Agent Loop

**Duration**: 60-75 minutes · **Language**: TypeScript (Node 20+) · **External dependencies**: None (pure TypeScript, mock provider)

This lab builds a cost-aware execution loop that wraps a mock agent loop with budget enforcement. You will implement the `TokenMeter` class, the `costAwareLoop` function, and a runaway-loop demo that triggers the circuit breaker. This is the layer Pi lacks.

---

## Phase 0 — Setup (5 min)

1. Ensure Node 20+ is installed: `node --version`
2. Create a working directory:
```bash
mkdir cost-aware-agent && cd cost-aware-agent
npm init -y
npm install -D tsx typescript @types/node
```
3. Create three files: `token-meter.ts`, `cost-aware-loop.ts`, `demo.ts`.

**Checkpoint**: `npx tsx --version` prints a version.

---

## Phase 1 — Implement the TokenMeter (15 min)

Write `token-meter.ts` implementing the interceptor from the teaching document (section DD-24.3):

```typescript
export interface TokenUsage {
  promptTokens: number;
  completionTokens: number;
  cumulative: number;
}

export interface BudgetConfig {
  taskBudget: number;          // max tokens for this task
  throttleAt: number;          // 0.8 = throttle at 80%
  rollbackAt: number;          // 1.0 = rollback at 100%
  retryStormThreshold: number; // max consecutive retries before circuit-break
}

export type BeforeCallDecision = "ok" | "throttle" | "rollback";

export class RetryStormError extends Error {
  constructor(message: string) {
    super(message);
    this.name = "RetryStormError";
  }
}
```

Implement the `TokenMeter` class with:
- `constructor(config: BudgetConfig)`
- `beforeCall(promptTokens: number): BeforeCallDecision` — checks projected usage against budget, returns "ok" / "throttle" / "rollback"
- `afterCall(usage: { promptTokens: number; completionTokens: number }, isRetry: boolean): void` — records usage, tracks consecutive retries, throws `RetryStormError` on storm
- `getUsage(): TokenUsage`
- `isThrottled(): boolean`

Key details:
- `beforeCall` computes `projected = cumulative + promptTokens`, then `ratio = projected / taskBudget`
- `afterCall` increments counters. On `isRetry=true`, increments `consecutiveRetries`; if >= threshold, throw `RetryStormError`. On `isRetry=false`, reset `consecutiveRetries = 0`.

**Reference**: The teaching document (01-teaching-document.md, section DD-24.3) has the full implementation. Use it as your guide.

**Checkpoint**: Create a `TokenMeter` with budget 10000, throttle 0.8, rollback 1.0, retryStormThreshold 3. Call `beforeCall(5000)` → "ok". Call `afterCall({promptTokens: 5000, completionTokens: 1000}, false)`. Call `beforeCall(7000)` → "throttle". Call `beforeCall(15000)` → "rollback".

---

## Phase 2 — Implement a Mock Provider (10 min)

Write a mock provider that simulates an LLM API call:

```typescript
export interface ProviderResponse {
  content: string;
  usage: { promptTokens: number; completionTokens: number };
  isError?: boolean;  // if true, the loop will likely retry
}

export type Provider = (model: string, messages: any[]) => Promise<ProviderResponse>;
```

Create a mock provider factory that can be configured to:
- Succeed normally (return content, accurate usage, no error)
- Fail N times then succeed (simulate a transient error + retry)
- Fail always (simulate a permanently broken tool / retry storm)
- Return inflated token counts (simulate the multiplier stack)

```typescript
export function makeMockProvider(opts: {
  failFirstN?: number;      // fail the first N calls, then succeed
  failAlways?: boolean;     // always fail (for storm demo)
  completionTokens?: number;// tokens per response
}): Provider {
  let callCount = 0;
  return async (model: string, messages: any[]): Promise<ProviderResponse> => {
    callCount++;
    const promptTokens = Math.round(JSON.stringify(messages).length / 4);

    if (opts.failAlways || (opts.failFirstN && callCount <= opts.failFirstN)) {
      return {
        content: "Error: tool execution failed",
        usage: { promptTokens, completionTokens: 50 },
        isError: true,
      };
    }

    return {
      content: `[DONE] Task completed using ${model}.`,
      usage: { promptTokens, completionTokens: opts.completionTokens ?? 200 },
    };
  };
}
```

**Checkpoint**: Create a mock provider with `failFirstN: 2`. Call it 3 times. The first two return `isError: true`, the third succeeds.

---

## Phase 3 — Implement the Cost-Aware Loop (20 min)

Write `cost-aware-loop.ts` implementing the loop from the teaching document (section DD-24.7):

```typescript
import { TokenMeter, RetryStormError, BudgetConfig } from "./token-meter";
import { Provider, ProviderResponse } from "./mock-provider";

export interface AgentConfig {
  budget: BudgetConfig;
  strongModel: string;    // e.g., "gpt-4o"
  cheapModel: string;     // e.g., "gpt-4o-mini"
  hitlThreshold: number;  // e.g., 0.8
  maxIterations: number;  // hard cap (fail-safe)
}

export interface LoopResult {
  result: string;
  totalCost: number;      // cumulative tokens
  iterations: number;
  throttled: boolean;
  hitlTriggered: boolean;
  aborted: boolean;
  abortReason?: string;
}
```

Implement `costAwareLoop(task, config, provider)` that:

1. Creates a `TokenMeter` with the budget config.
2. Runs the loop up to `maxIterations` times.
3. Before each provider call: estimates prompt tokens, calls `meter.beforeCall()`. If "rollback", abort with reason "budget exceeded".
4. Checks HITL threshold: if usage ratio >= `hitlThreshold` and not already triggered, log a warning and set `hitlTriggered = true`.
5. Selects model: if `meter.isThrottled()`, use `cheapModel`; else use `strongModel`.
6. Calls the provider. On exception, record as retry.
7. After each call: calls `meter.afterCall(usage, isRetry)`. Catches `RetryStormError` and aborts with reason "circuit breaker: retry storm".
8. Appends the response to messages. If response contains "[DONE]" or is not an error, break (task complete).
9. Returns a `LoopResult`.

**Reference**: The teaching document (01-teaching-document.md, section DD-24.7) has the full implementation.

**Checkpoint**: Create a loop with budget 50000, throttle 0.8, rollback 1.0, retryStormThreshold 5, maxIterations 50. Run with a normal provider. Verify it completes and returns `aborted: false`, `throttled: false`.

---

## Phase 4 — The Runaway Loop Demo (10 min)

Write `demo.ts` that demonstrates the circuit breaker:

```typescript
// Demo 1: Normal operation
const normalProvider = makeMockProvider({ completionTokens: 200 });
const result1 = await costAwareLoop(
  "Summarize this document",
  { budget: { taskBudget: 50000, throttleAt: 0.8, rollbackAt: 1.0, retryStormThreshold: 5 },
    strongModel: "gpt-4o", cheapModel: "gpt-4o-mini", hitlThreshold: 0.8, maxIterations: 50 },
  normalProvider,
);
console.log("Demo 1 (normal):", result1);

// Demo 2: Retry storm — circuit breaker trips
const stormProvider = makeMockProvider({ failAlways: true });
const result2 = await costAwareLoop(
  "Task that always fails",
  { budget: { taskBudget: 50000, throttleAt: 0.8, rollbackAt: 1.0, retryStormThreshold: 3 },
    strongModel: "gpt-4o", cheapModel: "gpt-4o-mini", hitlThreshold: 0.8, maxIterations: 50 },
  stormProvider,
);
console.log("Demo 2 (storm):", result2);
// Expected: aborted=true, abortReason contains "circuit breaker", iterations <= retryStormThreshold

// Demo 3: Budget exceeded — rollback
const expensiveProvider = makeMockProvider({ completionTokens: 20000 });
const result3 = await costAwareLoop(
  "Expensive task",
  { budget: { taskBudget: 10000, throttleAt: 0.8, rollbackAt: 1.0, retryStormThreshold: 5 },
    strongModel: "gpt-4o", cheapModel: "gpt-4o-mini", hitlThreshold: 0.8, maxIterations: 50 },
  expensiveProvider,
);
console.log("Demo 3 (budget exceeded):", result3);
// Expected: aborted=true, abortReason contains "budget exceeded"

// Demo 4: Throttle triggers — cheap model used
const throttleProvider = makeMockProvider({ completionTokens: 5000 });
const result4 = await costAwareLoop(
  "Task that hits throttle",
  { budget: { taskBudget: 10000, throttleAt: 0.5, rollbackAt: 1.0, retryStormThreshold: 5 },
    strongModel: "gpt-4o", cheapModel: "gpt-4o-mini", hitlThreshold: 0.9, maxIterations: 50 },
  throttleProvider,
);
console.log("Demo 4 (throttle):", result4);
// Expected: throttled=true, the loop used cheapModel for later iterations
```

**Checkpoint**: Run `npx tsx demo.ts`. All four demos produce the expected results:
- Demo 1: completes normally, aborted=false
- Demo 2: circuit breaker trips, aborted=true, abortReason mentions "retry storm"
- Demo 3: budget exceeded, aborted=true, abortReason mentions "budget"
- Demo 4: throttle triggered, throttled=true

---

## Phase 5 — The Re-sent Context Visualization (10 min)

Add instrumentation to the loop that visualizes the multiplier stack:

```typescript
// After each iteration, log the cumulative token usage
// and the percentage that is re-sent context (the same tokens, billed again)
function logMultiplierStack(meter: TokenMeter, iteration: number): void {
  const usage = meter.getUsage();
  const reSent = usage.promptTokens - (usage.promptTokens / iteration);
  // ^ rough estimate: promptTokens accumulate; only the first call's tokens are "new"
  console.log(
    `  iter ${iteration}: cumulative=${usage.cumulative}, ` +
    `re-sent context ≈ ${Math.round((reSent / usage.cumulative) * 100)}% of total`
  );
}
```

Call this after each `meter.afterCall()`. Run Demo 1 again and observe how the re-sent-context percentage grows over iterations — approaching 62% (the Stanford figure) as the session extends.

**Deliverable**: A short note documenting that re-sent context dominates cumulative token cost as sessions extend. This is the invisible 62%.

---

## Phase 6 — Stretch Goals

1. **Add HITL simulation.** Implement a `requestHumanApproval` function that returns a Promise (in the lab, resolves immediately; in production, awaits a UI interaction). Trigger it at the HITL threshold. Have the loop pause and resume based on approval.

2. **Add model-routing logic.** Instead of switching models only at throttle, route based on task phase: the first 3 iterations (planning) use the strong model; subsequent iterations (execution) use the cheap model. Log which model was used per iteration.

3. **Add prompt-caching simulation.** Add a `cacheablePrefixTokens` field to the provider config. If the first N tokens of the prompt are identical to the previous call, apply a 90% discount to those tokens. Log the savings. Then inject a dynamic timestamp into the prefix and observe the savings drop to zero.

4. **Add cost-aware tool selection.** Define tools with cost annotations (`read_file: 0`, `web_search: 500`, `paid_api: 5000`). At throttle, skip tools above a cost threshold. Log which tools were skipped.

5. **Add the fan-out simulation.** Extend the loop to support sub-agent calls. Track fan-out depth. Circuit-break at depth 3 (3^3 = 27 leaf calls). Verify the breaker prevents the exponential explosion.

---

## Deliverables Checklist

- [ ] `token-meter.ts` — TokenMeter class with beforeCall/afterCall, throttle/rollback/circuit-break enforcement
- [ ] `mock-provider.ts` (or inline) — configurable mock provider (normal, fail-N-then-succeed, fail-always, expensive)
- [ ] `cost-aware-loop.ts` — costAwareLoop function wrapping the provider with budget enforcement, HITL, model selection
- [ ] `demo.ts` — four demos: normal, retry storm (breaker), budget exceeded (rollback), throttle
- [ ] Re-sent context visualization showing the percentage approaching 62%
- [ ] (stretch) at least one of: HITL simulation, model routing, prompt caching, cost-aware tool selection, fan-out simulation

---

## Solution Key

The teaching document (01-teaching-document.md, sections DD-24.3 and DD-24.7) contains the full reference implementations of both `TokenMeter` and `costAwareLoop`. The key implementation details to verify:

1. **`beforeCall` checks PROJECTED usage** (cumulative + estimated prompt tokens), not just current cumulative. Prevention before the call.
2. **`afterCall` throws `RetryStormError`** when consecutive retries hit the threshold. The loop catches this and aborts.
3. **Throttle sets a flag** (`isThrottled`) that persists for the rest of the session. Once throttled, the loop uses the cheap model.
4. **HITL triggers exactly once** — the `hitlTriggered` flag prevents re-triggering on subsequent iterations.
5. **There is a hard `maxIterations` cap** as a fail-safe. Even if every other safeguard fails, the loop cannot run forever.

---

## What This Lab Teaches

This lab is the difference between understanding cost-awareness as a concept and having built the enforcement layer. After completing it, you will have:

- Implemented the `@token_meter` interceptor pattern with all three enforcement actions.
- Seen the circuit breaker trip on a retry storm and prevent budget combustion.
- Seen the rollback action fire when budget is exceeded.
- Seen the throttle action switch models mid-session.
- Visualized the re-sent-context multiplier approaching 62%.
- Internalized the prime directive (throttle must be paired with quality monitoring — a stretch goal in this lab).

The loop you build here is the foundation for Course 4 Module E03 (Cost-Aware Agent Architectures), which extends it to fleet scale: per-tenant budgets, chargeback, cost dashboards, and the instrumentation that makes the hidden 80% visible.