Keyboard shortcuts

Press or to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

Build: Run Spans and /stats Token Accounting

Maps to: Phase 5 (telemetry + /stats). Kind: Build.

Objective

Make the coordinator observable and accountable. You will add five things, all small, all leaning on the same principle: the telemetry is not new bookkeeping bolted on the side — it reads the tokens the workers already report and the run ids the domain already carries.

  1. telemetry::init — a tracing subscriber, filtered by RUST_LOG, safe to call from every test.
  2. TraceLayer::new_for_http() on the router — one line, and every request is a span.
  3. #[tracing::instrument] on process_run — the run id follows every job through dispatch.
  4. Store::usage_by_model() returning Vec<ModelUsage> — token totals per model, aggregated from the stored outcomes.
  5. GET /stats — turns those totals into tokens and an estimated cost per model, plus a grand total.

By the end, cargo run prints run-stamped logs you can filter, and GET /stats answers "how many tokens has each model burned, and what would it cost?"

Scaffold

Edit crates/control-store/Cargo.toml — nothing new is required; serde is already a dependency (ModelUsage derives Serialize).

Edit crates/panoptes-control/Cargo.toml — add, all { workspace = true }:

  • tracing — the #[instrument] macro and the info!/warn! macros.
  • tracing-subscriber (features ["env-filter"]) — the fmt subscriber and EnvFilter. The feature is what turns RUST_LOG into a filter; without it, EnvFilter does not exist.
  • tower-http (features ["trace"]) — supplies TraceLayer. The trace feature is why the layer is in scope.

tracing and tracing-subscriber are already available to control-store if you choose to log there, but the aggregation itself needs no logging.

Files you touch:

  • crates/control-store/src/lib.rs — add ModelUsage and Store::usage_by_model.
  • crates/panoptes-control/src/lib.rs — add the telemetry module, the TraceLayer on app, est_cost_usd, and the get_stats handler wired to GET /stats.
  • crates/panoptes-control/src/scheduler.rs — add the #[instrument] attribute to process_run.

Expected result:

  • cargo test -p control-store → the store suite grows by 1: usage_by_model_sums_across_jobs.
  • cargo test -p panoptes-control → the API suite grows by 1: stats_reports_tokens_and_cost_per_model.

The spec (givens)

telemetry::init

A module with one function that installs the global subscriber:

pub mod telemetry {
    pub fn init() { /* you write this */ }
}

The contract — three properties the rest of the system depends on:

  • Filtered by RUST_LOG, defaulting to "info". Build an EnvFilter from the default env; if the variable is unset, fall back to EnvFilter::new("info"). This is the exact try_from_default_env().unwrap_or_else(...) shape from the concept chapter.
  • Idempotent. Use try_init, not init. It returns a Result instead of panicking when a subscriber is already installed, so every test can call telemetry::init() in its setup and only the first call wins. Discard the Result (let _ = ...).
  • fmt subscriber. A plain tracing_subscriber::fmt() with the filter attached is all this course needs.

main.rs calls telemetry::init() once at startup, before serve.

→ Answer key

TraceLayer on the router

app gains exactly one line — a layer, applied after the routes and before .with_state:

.layer(TraceLayer::new_for_http())

Every request now opens a span carrying method and path and logs its status and latency on completion. No handler changes. (Layer order matters in general; for this single layer, placing it after the routes is correct.)

→ Answer key

The run span on process_run

Add this attribute — and nothing else — to process_run in scheduler.rs:

#[tracing::instrument(skip(self, run), fields(run_id = %run.id))]
async fn process_run(&self, run: Run) -> Result<(), ControlError> { ... }

skip(self, run) keeps the two large arguments out of the span's fields; fields(run_id = %run.id) records just the id, via Display (that is the %). From here every info!/warn! the scheduler emits inside the run — including the existing "planned jobs for run" line — is stamped with run_id. The body does not change.

→ Answer key

Store::usage_by_model and ModelUsage

The aggregate type, in control-store:

/// Aggregate token usage for one model, across every recorded job.
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
pub struct ModelUsage {
    pub model: String,
    pub input_tokens: u64,
    pub output_tokens: u64,
    pub calls: u64,
}

impl Store {
    pub async fn usage_by_model(&self) -> Result<Vec<ModelUsage>, ControlError>;
}

The contract:

  • Read every job that has an outcome (the same rows run_results reads), deserialize each JobOutcome, and walk its records.
  • Group by record.model. For each record, add usage.input_tokens to the model's input_tokens, usage.output_tokens to its output_tokens, and 1 to callscalls counts records (one model call each), not jobs.
  • Widen the per-record u32 token counts to u64 as you sum, so a long run cannot overflow the total.
  • Return one ModelUsage per distinct model. A BTreeMap<String, ModelUsage> keyed by model name gives you grouping and a stable, sorted order for free — worth it so the output and the tests are deterministic.

→ Answer key

GET /stats, est_cost_usd, and the pricing

The cost helper is an illustrative flat model — $3 per 1,000,000 input tokens, $15 per 1,000,000 output tokens:

fn est_cost_usd(u: &ModelUsage) -> f64; // input/1e6 * 3.0 + output/1e6 * 15.0

get_stats calls store.usage_by_model(), maps each ModelUsage to a JSON object, sums the per-model costs into a grand total, and answers this exact shape:

// GET /stats  → 200
{
  "models": [
    {
      "model": "claude",
      "input_tokens": 1000000,
      "output_tokens": 1000000,
      "calls": 1,
      "est_cost_usd": 18.0
    }
  ],
  "total_est_cost_usd": 18.0
}

The worked number pins the arithmetic: 1,000,000 input @ $3/M = $3, plus 1,000,000 output @ $15/M = $15, is $18.00. Wire it as .route("/stats", get(get_stats)); the handler returns Result<Json<serde_json::Value>, ApiError> and uses ? on the store call, exactly like the Part V handlers.

→ Answer key

Concepts exercised

  • #[tracing::instrument] with skip/fields to carry a run_id through an async call and its nested work (from Concept: tracing).
  • EnvFilter + try_init for a RUST_LOG-controlled, idempotent subscriber.
  • TraceLayer as a single tower layer that traces every request uniformly.
  • Aggregating stored outcomes in Rust with a BTreeMap for grouped, deterministic results.
  • Computing a derived value (cost) from data the system already records, and shaping it into a JSON response with the Part V handler pattern.

The build loop (you drive)

Test 1 — usage_by_model_sums_across_jobs (in control-store/src/lib.rs, #[tokio::test])

  1. Write the failing test. Insert a run with job_count = 2 and two jobs. Record an outcome for job 1 with one record usage { input: 100, output: 10 }, and an outcome for job 2 with two records, each usage { input: 50, output: 5 } — all model: "claude". Call store.usage_by_model() and assert: stats.len() == 1, stats[0].model == "claude", stats[0].input_tokens == 200, stats[0].output_tokens == 20, stats[0].calls == 3.
  2. Predict: three records land across two jobs. Before you implement the loop — is calls going to be 2 or 3? Say which, and what that tells you about whether calls counts jobs or model calls.
  3. Run — it fails to compile (no usage_by_model yet).
  4. Implement ModelUsage and usage_by_model per the contract above. What it computes is specified; how you group is yours (a BTreeMap is the clean way).
  5. Run green, commit.

Test 2 — stats_reports_tokens_and_cost_per_model (in panoptes-control/src/lib.rs, #[tokio::test])

  1. Write the failing test. Spawn the app on a random port with a shared in-memory Store (the existing spawn() helper returns both). Seed a Running run, insert one job for it, and record_job_outcome with a single record whose usage is { input: 1_000_000, output: 1_000_000 }, model: "claude". Then GET /stats, parse the JSON, and assert: stats["models"][0]["model"] == "claude", stats["models"][0]["input_tokens"] == 1_000_000, and stats["total_est_cost_usd"] == 18.0.
  2. Predict: with input and output both at 1,000,000, work the pricing by hand — $3/M in, $15/M out. What is est_cost_usd for this model, and what is total_est_cost_usd? (If you get $18.0, your mental model of the formula matches the test.)
  3. Run — it fails (no /stats route, no get_stats, no est_cost_usd).
  4. Implement est_cost_usd, get_stats, and the .route("/stats", get(get_stats)) wiring. Add the TraceLayer line and telemetry::init while you are in this file.
  5. Run green, commit.
Predict before you run The /stats test never starts a scheduler and never sends a model a real prompt — it seeds an outcome directly into the store and asks the API to report on it. Why is that a complete test of the stats path? (Hint: get_stats reads usage_by_model, which reads stored outcomes. The number on the wire depends only on what is in the store, not on how it got there — so seeding the store is exactly seeding the input.)
Why cost is derived, never stored Notice what you did not add: a cost column, a running total the workers update, a second source of truth. est_cost_usd is a pure function of the token counts the workers already report, computed at read time in /stats. Change the pricing and every historical run reprices correctly, because nothing was frozen at write time. That is the same discipline as the JSONL contract and the IntoResponse error mapping — compute the derived thing in one place, from the data you already trust.

Done when

cargo test -p control-store and cargo test -p panoptes-control are both green with the two new tests passing; usage_by_model sums to 200 / 20 / 3 across the two jobs; GET /stats reports input_tokens: 1000000 and total_est_cost_usd: 18.0 for the seeded run; cargo run emits run_id-stamped log lines that RUST_LOG=warn can quiet and RUST_LOG=debug can widen.