diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index bfe57eab..40229340 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -117,7 +117,7 @@ jobs:
# repo means find_repo_root() returns None, the local-build short-
# circuit is skipped, and the CLI tries to download from a GitHub
# release that does not yet exist for the in-flight version on
- # release-bump PRs (e.g. 0.4.23 before publish-binaries has run).
+ # release-bump PRs (e.g. 0.4.29 before publish-binaries has run).
- name: Seed ~/.plano cache for zero-config test
run: |
VERSION=$(sed -nE 's/^__version__ = "(.*)"$/\1/p' cli/planoai/__init__.py)
@@ -183,13 +183,13 @@ jobs:
load: true
tags: |
${{ env.PLANO_DOCKER_IMAGE }}
- ${{ env.DOCKER_IMAGE }}:0.4.23
+ ${{ env.DOCKER_IMAGE }}:0.4.29
${{ env.DOCKER_IMAGE }}:latest
cache-from: type=gha
cache-to: type=gha,mode=max
- name: Save image as artifact
- run: docker save ${{ env.PLANO_DOCKER_IMAGE }} ${{ env.DOCKER_IMAGE }}:0.4.23 ${{ env.DOCKER_IMAGE }}:latest -o /tmp/plano-image.tar
+ run: docker save ${{ env.PLANO_DOCKER_IMAGE }} ${{ env.DOCKER_IMAGE }}:0.4.29 ${{ env.DOCKER_IMAGE }}:latest -o /tmp/plano-image.tar
- name: Upload image artifact
uses: actions/upload-artifact@v6
diff --git a/.github/workflows/update-providers.yml b/.github/workflows/update-providers.yml
index 0add8a98..affc089b 100644
--- a/.github/workflows/update-providers.yml
+++ b/.github/workflows/update-providers.yml
@@ -57,6 +57,7 @@ jobs:
GOOGLE_API_KEY: ${{ secrets.GOOGLE_API_KEY }}
OPENROUTER_API_KEY: ${{ secrets.OPENROUTER_API_KEY }}
AI_GATEWAY_API_KEY: ${{ secrets.AI_GATEWAY_API_KEY }}
+ META_MODELS_API_KEY: ${{ secrets.META_MODELS_API_KEY }}
run: cargo run --bin fetch_models --features model-fetch
- name: Create pull request
diff --git a/Dockerfile b/Dockerfile
index ad0ca707..c273e57a 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -47,10 +47,13 @@ FROM docker.io/envoyproxy/envoy:${ENVOY_VERSION} AS envoy
FROM python:3.14-slim AS arch
+# Install runtime deps first, then upgrade — so security patches also land on
+# packages pulled in transitively by the install, not just those already present
+# in the base image.
RUN set -eux; \
apt-get update; \
+ apt-get install -y --no-install-recommends gettext-base procps; \
apt-get upgrade -y; \
- apt-get install -y --no-install-recommends gettext-base curl procps; \
apt-get clean; rm -rf /var/lib/apt/lists/*
RUN pip install --no-cache-dir supervisor
diff --git a/apps/www/src/components/Hero.tsx b/apps/www/src/components/Hero.tsx
index b9d5b170..f4849b04 100644
--- a/apps/www/src/components/Hero.tsx
+++ b/apps/www/src/components/Hero.tsx
@@ -24,7 +24,7 @@ export function Hero() {
>
- v0.4.23
+ v0.4.29
—
diff --git a/build_filter_image.sh b/build_filter_image.sh
index 624955c2..5640f680 100644
--- a/build_filter_image.sh
+++ b/build_filter_image.sh
@@ -1 +1 @@
-docker build -f Dockerfile . -t katanemo/plano -t katanemo/plano:0.4.23
+docker build -f Dockerfile . -t katanemo/plano -t katanemo/plano:0.4.29
diff --git a/cli/planoai/__init__.py b/cli/planoai/__init__.py
index dc0c543a..c28f402c 100644
--- a/cli/planoai/__init__.py
+++ b/cli/planoai/__init__.py
@@ -1,3 +1,3 @@
"""Plano CLI - Intelligent Prompt Gateway."""
-__version__ = "0.4.23"
+__version__ = "0.4.29"
diff --git a/cli/planoai/consts.py b/cli/planoai/consts.py
index 8f13ba50..8377f583 100644
--- a/cli/planoai/consts.py
+++ b/cli/planoai/consts.py
@@ -5,7 +5,7 @@ PLANO_COLOR = "#969FF4"
SERVICE_NAME_ARCHGW = "plano"
PLANO_DOCKER_NAME = "plano"
-PLANO_DOCKER_IMAGE = os.getenv("PLANO_DOCKER_IMAGE", "katanemo/plano:0.4.23")
+PLANO_DOCKER_IMAGE = os.getenv("PLANO_DOCKER_IMAGE", "katanemo/plano:0.4.29")
DEFAULT_OTEL_TRACING_GRPC_ENDPOINT = "http://localhost:4317"
# Native mode constants
diff --git a/cli/planoai/obs/pricing.py b/cli/planoai/obs/pricing.py
index 6f2ce5b4..0a8b7321 100644
--- a/cli/planoai/obs/pricing.py
+++ b/cli/planoai/obs/pricing.py
@@ -1,7 +1,8 @@
-"""DigitalOcean Gradient pricing catalog for the obs console.
+"""Model pricing catalog for the obs console.
-Ported loosely from ``crates/brightstaff/src/router/model_metrics.rs::fetch_do_pricing``.
-Single-source: one fetch at startup, cached for the life of the process.
+Mirrors ``crates/brightstaff/src/router/model_metrics.rs``. The source is
+configurable: ``digitalocean`` (DO GenAI catalog) or ``models.dev``. A single
+fetch at startup is cached for the life of the process.
"""
from __future__ import annotations
@@ -14,7 +15,18 @@ from typing import Any
import requests
-DEFAULT_PRICING_URL = "https://api.digitalocean.com/v2/gen-ai/models/catalog"
+DO_PRICING_URL = "https://api.digitalocean.com/v2/gen-ai/models/catalog"
+MODELS_DEV_URL = "https://models.dev/api.json"
+
+# Backwards-compatible default (DigitalOcean) used when no provider is given.
+DEFAULT_PRICING_URL = DO_PRICING_URL
+DEFAULT_PRICING_PROVIDER = "digitalocean"
+
+_DEFAULT_URLS = {
+ "digitalocean": DO_PRICING_URL,
+ "models.dev": MODELS_DEV_URL,
+}
+
FETCH_TIMEOUT_SECS = 5.0
@@ -51,36 +63,52 @@ class PricingCatalog:
return list(self._prices.keys())[:n]
@classmethod
- def fetch(cls, url: str = DEFAULT_PRICING_URL) -> "PricingCatalog":
- """Fetch pricing from DO's catalog endpoint. On failure, returns an
+ def fetch(
+ cls,
+ provider: str = DEFAULT_PRICING_PROVIDER,
+ url: str | None = None,
+ ) -> "PricingCatalog":
+ """Fetch pricing from the configured catalog. On failure, returns an
empty catalog (cost column will be blank).
- The catalog endpoint is public — no auth required, no signup — so
- ``planoai obs`` gets cost data on first run out of the box.
+ ``provider`` selects the parser/default URL: ``digitalocean`` or
+ ``models.dev``. Both catalog endpoints are public — no auth required —
+ so ``planoai obs`` gets cost data on first run out of the box.
"""
+ provider = (provider or DEFAULT_PRICING_PROVIDER).strip().lower()
+ resolved_url = url or _DEFAULT_URLS.get(provider, DO_PRICING_URL)
try:
- resp = requests.get(url, timeout=FETCH_TIMEOUT_SECS)
+ resp = requests.get(resolved_url, timeout=FETCH_TIMEOUT_SECS)
resp.raise_for_status()
data = resp.json()
except Exception as exc: # noqa: BLE001 — best-effort; never fatal
logger.warning(
- "DO pricing fetch failed: %s; cost column will be blank.",
+ "%s pricing fetch failed: %s; cost column will be blank.",
+ provider,
exc,
)
return cls()
- prices = _parse_do_pricing(data)
+ if provider == "models.dev":
+ prices = _parse_models_dev_pricing(data)
+ else:
+ prices = _parse_do_pricing(data)
+
if not prices:
- # Dump the first entry's raw shape so we can see which fields DO
- # actually returned — helps when the catalog adds new fields or
- # the response doesn't match our parser.
+ # Dump a sample of the raw shape so we can see which fields the
+ # catalog returned — helps when it adds new fields or the response
+ # doesn't match our parser.
import json as _json
- sample_items = _coerce_items(data)
- sample = sample_items[0] if sample_items else data
+ if provider == "models.dev" and isinstance(data, dict):
+ sample = next(iter(data.values()), data)
+ else:
+ sample_items = _coerce_items(data)
+ sample = sample_items[0] if sample_items else data
logger.warning(
- "DO pricing response had no parseable entries; cost column "
+ "%s pricing response had no parseable entries; cost column "
"will be blank. Sample entry: %s",
+ provider,
_json.dumps(sample, default=str)[:400],
)
return cls(prices)
@@ -278,6 +306,75 @@ def _parse_do_pricing(data: Any) -> dict[str, ModelPrice]:
return prices
+def _parse_models_dev_pricing(data: Any) -> dict[str, ModelPrice]:
+ """Parse a models.dev ``api.json`` response into a ModelPrice map.
+
+ models.dev shape (top-level object keyed by provider id)::
+
+ {
+ "anthropic": {
+ "models": {
+ "claude-opus-4-5": {
+ "cost": {"input": 5, "output": 25, "cache_read": 0.5}
+ }
+ }
+ },
+ ...
+ }
+
+ ``cost.*`` values are USD per *million* tokens, so we divide by 1e6 to get a
+ per-token rate. First-party providers use bare model keys, so we register
+ both ``provider/model`` (matching Plano's routing names) and the bare model
+ id as a fallback.
+ """
+ prices: dict[str, ModelPrice] = {}
+ if not isinstance(data, dict):
+ return prices
+
+ for provider_id, provider in data.items():
+ if not isinstance(provider, dict):
+ continue
+ models = provider.get("models")
+ if not isinstance(models, dict):
+ continue
+ for model_key, model in models.items():
+ if not isinstance(model, dict):
+ continue
+ cost = model.get("cost")
+ if not isinstance(cost, dict):
+ continue
+ input_pm = _as_float(cost.get("input"))
+ output_pm = _as_float(cost.get("output"))
+ if input_pm is None or output_pm is None:
+ continue
+ # Skip 0-rate entries so cost falls back to `—` rather than $0.0000.
+ if input_pm == 0 and output_pm == 0:
+ continue
+ cached_pm = _as_float(cost.get("cache_read"))
+ price = ModelPrice(
+ input_per_token_usd=input_pm / 1_000_000,
+ output_per_token_usd=output_pm / 1_000_000,
+ cached_input_per_token_usd=(
+ cached_pm / 1_000_000 if cached_pm is not None else None
+ ),
+ )
+ composite = f"{provider_id}/{model_key}"
+ prices[composite] = price
+ prices.setdefault(composite.lower(), price)
+ prices.setdefault(str(model_key), price)
+ prices.setdefault(str(model_key).lower(), price)
+ return prices
+
+
+def _as_float(value: Any) -> float | None:
+ if value is None:
+ return None
+ try:
+ return float(value)
+ except (TypeError, ValueError):
+ return None
+
+
def _coerce_items(data: Any) -> list[dict]:
if isinstance(data, list):
return [x for x in data if isinstance(x, dict)]
diff --git a/cli/planoai/obs_cmd.py b/cli/planoai/obs_cmd.py
index 6249df30..a0867b6e 100644
--- a/cli/planoai/obs_cmd.py
+++ b/cli/planoai/obs_cmd.py
@@ -2,9 +2,12 @@
from __future__ import annotations
+import logging
+import os
import time
import rich_click as click
+import yaml
from rich.console import Console
from rich.live import Live
@@ -15,8 +18,50 @@ from planoai.obs.collector import (
LLMCallStore,
ObsCollector,
)
-from planoai.obs.pricing import PricingCatalog
+from planoai.obs.pricing import DEFAULT_PRICING_PROVIDER, PricingCatalog
from planoai.obs.render import render
+from planoai.utils import find_config_file
+
+logger = logging.getLogger(__name__)
+
+
+def _resolve_pricing_source(
+ config_file: str | None,
+ provider_override: str | None,
+ url_override: str | None,
+) -> tuple[str, str | None]:
+ """Pick the cost pricing source.
+
+ Precedence: explicit CLI overrides > the first ``type: cost`` entry in
+ ``model_metrics_sources`` from the Plano config > the DigitalOcean default.
+ """
+ provider = DEFAULT_PRICING_PROVIDER
+ url: str | None = None
+
+ config_path = find_config_file(file=config_file)
+ if config_path and os.path.exists(config_path):
+ try:
+ with open(config_path, "r") as f:
+ config = yaml.safe_load(f) or {}
+ sources = config.get("model_metrics_sources") or []
+ for source in sources:
+ if isinstance(source, dict) and source.get("type") == "cost":
+ if source.get("provider"):
+ provider = str(source["provider"])
+ if source.get("url"):
+ url = str(source["url"])
+ break
+ except Exception as exc: # noqa: BLE001 — config is optional for obs
+ logger.warning(
+ "could not read pricing source from %s: %s", config_path, exc
+ )
+
+ if provider_override:
+ provider = provider_override
+ if url_override:
+ url = url_override
+
+ return provider, url
@click.command(name="obs", help="Live observability console for Plano LLM traffic.")
@@ -48,13 +93,42 @@ from planoai.obs.render import render
show_default=True,
help="TUI refresh interval.",
)
-def obs(port: int, host: str, capacity: int, refresh_ms: int) -> None:
+@click.option(
+ "--config",
+ "config_file",
+ type=str,
+ default=None,
+ help="Path to the Plano config to read the pricing source from "
+ "(defaults to ./config.yaml or ./plano_config.yaml).",
+)
+@click.option(
+ "--pricing-provider",
+ type=click.Choice(["digitalocean", "models.dev"]),
+ default=None,
+ help="Override the cost pricing provider (otherwise read from config).",
+)
+@click.option(
+ "--pricing-url",
+ type=str,
+ default=None,
+ help="Override the pricing catalog URL (otherwise read from config / provider default).",
+)
+def obs(
+ port: int,
+ host: str,
+ capacity: int,
+ refresh_ms: int,
+ config_file: str | None,
+ pricing_provider: str | None,
+ pricing_url: str | None,
+) -> None:
console = Console()
+ provider, url = _resolve_pricing_source(config_file, pricing_provider, pricing_url)
console.print(
- f"[bold {PLANO_COLOR}]planoai obs[/] — loading DO pricing catalog...",
+ f"[bold {PLANO_COLOR}]planoai obs[/] — loading {provider} pricing catalog...",
end="",
)
- pricing = PricingCatalog.fetch()
+ pricing = PricingCatalog.fetch(provider=provider, url=url)
if len(pricing):
sample = ", ".join(pricing.sample_models(3))
console.print(
@@ -63,7 +137,7 @@ def obs(port: int, host: str, capacity: int, refresh_ms: int) -> None:
else:
console.print(
" [yellow]no pricing loaded[/] — "
- "[dim]cost column will be blank (DO catalog unreachable)[/]"
+ f"[dim]cost column will be blank ({provider} catalog unreachable)[/]"
)
store = LLMCallStore(capacity=capacity)
diff --git a/cli/planoai/templates/conversational_state_v1_responses.yaml b/cli/planoai/templates/conversational_state_v1_responses.yaml
index 403278a9..11fb7477 100644
--- a/cli/planoai/templates/conversational_state_v1_responses.yaml
+++ b/cli/planoai/templates/conversational_state_v1_responses.yaml
@@ -11,7 +11,7 @@ model_providers:
default: true
# Anthropic Models
- - model: anthropic/claude-sonnet-4-20250514
+ - model: anthropic/claude-sonnet-4-6
access_key: $ANTHROPIC_API_KEY
listeners:
diff --git a/cli/planoai/templates/preference_aware_routing.yaml b/cli/planoai/templates/preference_aware_routing.yaml
index e38b3881..1fcb6bf4 100644
--- a/cli/planoai/templates/preference_aware_routing.yaml
+++ b/cli/planoai/templates/preference_aware_routing.yaml
@@ -12,7 +12,7 @@ model_providers:
- name: code understanding
description: understand and explain existing code snippets, functions, or libraries
- - model: anthropic/claude-sonnet-4-20250514
+ - model: anthropic/claude-sonnet-4-6
access_key: $ANTHROPIC_API_KEY
routing_preferences:
- name: code generation
diff --git a/cli/pyproject.toml b/cli/pyproject.toml
index e0a90e11..97f37678 100644
--- a/cli/pyproject.toml
+++ b/cli/pyproject.toml
@@ -1,6 +1,6 @@
[project]
name = "planoai"
-version = "0.4.23"
+version = "0.4.29"
description = "Python-based CLI tool to manage Plano."
authors = [{name = "Katanemo Labs, Inc."}]
readme = "README.md"
diff --git a/cli/test/test_config_generator.py b/cli/test/test_config_generator.py
index 9aade29e..2f9834ca 100644
--- a/cli/test/test_config_generator.py
+++ b/cli/test/test_config_generator.py
@@ -386,6 +386,33 @@ model_providers:
access_key: $OPENAI_API_KEY
default: true
+""",
+ },
+ {
+ "id": "valid_tracing_posthog_exporter",
+ "expected_error": None,
+ "plano_config": """
+version: v0.4.0
+
+listeners:
+ - name: llm
+ type: model
+ port: 12000
+
+model_providers:
+ - model: openai/gpt-4o-mini
+ access_key: $OPENAI_API_KEY
+ default: true
+
+tracing:
+ random_sampling: 100
+ exporters:
+ - type: posthog
+ url: https://us.i.posthog.com
+ api_key: $POSTHOG_API_KEY
+ distinct_id_header: x-user-id
+ capture_messages: false
+
""",
},
]
@@ -584,7 +611,7 @@ model_providers:
- name: code understanding
description: understand and explain existing code snippets, functions, or libraries
- - model: anthropic/claude-sonnet-4-20250514
+ - model: anthropic/claude-sonnet-4-6
access_key: $ANTHROPIC_API_KEY
routing_preferences:
- name: code generation
@@ -601,9 +628,7 @@ model_providers:
by_name = {entry["name"]: entry for entry in top_level}
assert set(by_name) == {"code understanding", "code generation"}
assert by_name["code understanding"]["models"] == ["openai/gpt-4o"]
- assert by_name["code generation"]["models"] == [
- "anthropic/claude-sonnet-4-20250514"
- ]
+ assert by_name["code generation"]["models"] == ["anthropic/claude-sonnet-4-6"]
assert (
by_name["code understanding"]["description"]
== "understand and explain existing code snippets, functions, or libraries"
@@ -626,7 +651,7 @@ model_providers:
- name: code generation
description: generating new code snippets, functions, or boilerplate based on user prompts or requirements
- - model: anthropic/claude-sonnet-4-20250514
+ - model: anthropic/claude-sonnet-4-6
access_key: $ANTHROPIC_API_KEY
routing_preferences:
- name: code generation
@@ -641,7 +666,7 @@ model_providers:
assert entry["name"] == "code generation"
assert entry["models"] == [
"openai/gpt-4o",
- "anthropic/claude-sonnet-4-20250514",
+ "anthropic/claude-sonnet-4-6",
]
assert config_yaml["version"] == "v0.4.0"
@@ -658,7 +683,7 @@ listeners:
model_providers:
- model: openai/gpt-4o
access_key: $OPENAI_API_KEY
- - model: anthropic/claude-sonnet-4-20250514
+ - model: anthropic/claude-sonnet-4-6
access_key: $ANTHROPIC_API_KEY
routing_preferences:
@@ -666,7 +691,7 @@ routing_preferences:
description: generating new code snippets or boilerplate
models:
- openai/gpt-4o
- - anthropic/claude-sonnet-4-20250514
+ - anthropic/claude-sonnet-4-6
"""
config_yaml = yaml.safe_load(plano_config)
before = yaml.safe_dump(config_yaml, sort_keys=True)
diff --git a/cli/test/test_obs_pricing.py b/cli/test/test_obs_pricing.py
index 02247d3d..322607a9 100644
--- a/cli/test/test_obs_pricing.py
+++ b/cli/test/test_obs_pricing.py
@@ -144,3 +144,68 @@ def test_parse_do_catalog_divides_large_values_as_per_million():
prices = _parse_do_pricing(sample)
assert prices["mystery-model"].input_per_token_usd == 5.0 / 1_000_000
assert prices["mystery-model"].output_per_token_usd == 15.0 / 1_000_000
+
+
+_MODELS_DEV_SAMPLE = {
+ "anthropic": {
+ "id": "anthropic",
+ "models": {
+ "claude-opus-4-5": {
+ "id": "claude-opus-4-5",
+ "cost": {"input": 5, "output": 25, "cache_read": 0.5},
+ }
+ },
+ },
+ "groq": {
+ "id": "groq",
+ "models": {
+ "llama-3.3-70b-versatile": {
+ "id": "llama-3.3-70b-versatile",
+ "cost": {"input": 0.59, "output": 0.79},
+ },
+ # No cost block → skipped.
+ "whisper-large-v3-turbo": {"id": "whisper-large-v3-turbo"},
+ },
+ },
+}
+
+
+def test_parse_models_dev_composes_provider_keys_and_per_token_rates():
+ from planoai.obs.pricing import _parse_models_dev_pricing
+
+ prices = _parse_models_dev_pricing(_MODELS_DEV_SAMPLE)
+
+ # models.dev cost values are per-million → divided by 1e6.
+ opus = prices["anthropic/claude-opus-4-5"]
+ assert opus.input_per_token_usd == 5 / 1_000_000
+ assert opus.output_per_token_usd == 25 / 1_000_000
+ assert opus.cached_input_per_token_usd == 0.5 / 1_000_000
+
+ # Composite provider/model keys match Plano's routing names.
+ assert "groq/llama-3.3-70b-versatile" in prices
+ # Bare model id registered as a fallback.
+ assert "llama-3.3-70b-versatile" in prices
+ # Models without a cost block are skipped.
+ assert "groq/whisper-large-v3-turbo" not in prices
+
+
+def test_models_dev_catalog_cost_computation():
+ from planoai.obs.pricing import PricingCatalog, _parse_models_dev_pricing
+
+ catalog = PricingCatalog(_parse_models_dev_pricing(_MODELS_DEV_SAMPLE))
+ # 1000 input @ 5e-6 = 0.005; 500 output @ 25e-6 = 0.0125
+ cost = catalog.cost_for_call(_call("anthropic/claude-opus-4-5", 1000, 500))
+ assert cost == round(0.005 + 0.0125, 6)
+
+
+def test_models_dev_skips_zero_rate_entries():
+ from planoai.obs.pricing import _parse_models_dev_pricing
+
+ sample = {
+ "free": {
+ "models": {
+ "promo-model": {"cost": {"input": 0, "output": 0}},
+ }
+ }
+ }
+ assert _parse_models_dev_pricing(sample) == {}
diff --git a/cli/uv.lock b/cli/uv.lock
index 98d50481..fafbaabb 100644
--- a/cli/uv.lock
+++ b/cli/uv.lock
@@ -110,7 +110,7 @@ name = "exceptiongroup"
version = "1.3.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
- { name = "typing-extensions", marker = "python_full_version < '3.13'" },
+ { name = "typing-extensions" },
]
sdist = { url = "https://files.pythonhosted.org/packages/0b/9f/a65090624ecf468cdca03533906e7c69ed7588582240cfe7cc9e770b50eb/exceptiongroup-1.3.0.tar.gz", hash = "sha256:b241f5885f560bc56a59ee63ca4c6a8bfa46ae4ad651af316d4e81817bb9fd88", size = 29749, upload-time = "2025-05-10T17:42:51.123Z" }
wheels = [
@@ -337,7 +337,7 @@ wheels = [
[[package]]
name = "planoai"
-version = "0.4.23"
+version = "0.4.29"
source = { editable = "." }
dependencies = [
{ name = "click" },
diff --git a/config/envoy.template.yaml b/config/envoy.template.yaml
index b2b9fb1f..22b1a535 100644
--- a/config/envoy.template.yaml
+++ b/config/envoy.template.yaml
@@ -196,7 +196,11 @@ static_resources:
name: decompress
typed_config:
"@type": "type.googleapis.com/envoy.extensions.compression.gzip.decompressor.v3.Gzip"
- window_bits: 9
+ # Must be >= the window_bits of whatever compressed the response
+ # (upstreams use 15; our compressor uses 10). A smaller inflate window
+ # makes zlib silently stop emitting data mid-stream, truncating SSE
+ # responses after the first few KB. 15 handles any compliant gzip.
+ window_bits: 15
chunk_size: 8192
# If this ratio is set too low, then body data will not be decompressed completely.
max_inflate_ratio: 1000
@@ -269,6 +273,13 @@ static_resources:
auto_host_rewrite: true
cluster: {{ cluster_name }}
timeout: 300s
+ {% if cluster.prefix_affinity %}
+ # Feeds the cluster's RING_HASH lb_policy: requests with the
+ # same prompt-prefix hash stick to the same replica.
+ hash_policy:
+ - header:
+ header_name: x-plano-prefix-hash
+ {% endif %}
{% endfor %}
http_filters:
- name: envoy.filters.http.router
@@ -361,7 +372,11 @@ static_resources:
name: decompress
typed_config:
"@type": "type.googleapis.com/envoy.extensions.compression.gzip.decompressor.v3.Gzip"
- window_bits: 9
+ # Must be >= the window_bits of whatever compressed the response
+ # (upstreams use 15; our compressor uses 10). A smaller inflate window
+ # makes zlib silently stop emitting data mid-stream, truncating SSE
+ # responses after the first few KB. 15 handles any compliant gzip.
+ window_bits: 15
chunk_size: 8192
# If this ratio is set too low, then body data will not be decompressed completely.
max_inflate_ratio: 1000
@@ -442,7 +457,11 @@ static_resources:
name: decompress
typed_config:
"@type": "type.googleapis.com/envoy.extensions.compression.gzip.decompressor.v3.Gzip"
- window_bits: 9
+ # Must be >= the window_bits of whatever compressed the response
+ # (upstreams use 15; our compressor uses 10). A smaller inflate window
+ # makes zlib silently stop emitting data mid-stream, truncating SSE
+ # responses after the first few KB. 15 handles any compliant gzip.
+ window_bits: 15
chunk_size: 8192
# If this ratio is set too low, then body data will not be decompressed completely.
max_inflate_ratio: 1000
@@ -577,6 +596,12 @@ static_resources:
name: decompress
typed_config:
"@type": "type.googleapis.com/envoy.extensions.compression.gzip.decompressor.v3.Gzip"
+ # Must be >= the window_bits of whatever compressed the response
+ # (upstreams use 15; our compressor uses 10; Envoy's default here is
+ # only 12). A smaller inflate window makes zlib silently stop emitting
+ # data mid-stream, truncating SSE responses after the first few KB.
+ # 15 handles any compliant gzip.
+ window_bits: 15
chunk_size: 8192
# If this ratio is set too low, then body data will not be decompressed completely.
max_inflate_ratio: 1000
@@ -977,9 +1002,18 @@ static_resources:
{% else -%}
connect_timeout: {{ upstream_connect_timeout | default('5s') }}
{% endif -%}
+ {% if cluster.prefix_affinity -%}
+ # KV-aware replica stickiness: resolve every replica address and consistent-hash
+ # requests (by x-plano-prefix-hash, see the route-level hash_policy) so the same
+ # prompt prefix lands on the replica holding its warm KV cache.
+ type: STRICT_DNS
+ dns_lookup_family: V4_ONLY
+ lb_policy: RING_HASH
+ {% else -%}
type: LOGICAL_DNS
dns_lookup_family: V4_ONLY
lb_policy: ROUND_ROBIN
+ {% endif -%}
load_assignment:
cluster_name: {{ cluster_name }}
endpoints:
diff --git a/config/plano_config_schema.yaml b/config/plano_config_schema.yaml
index 2ecf3892..0b92c777 100644
--- a/config/plano_config_schema.yaml
+++ b/config/plano_config_schema.yaml
@@ -155,6 +155,13 @@ properties:
- https
http_host:
type: string
+ prefix_affinity:
+ type: boolean
+ description: >
+ For self-hosted multi-replica backends (e.g. vLLM): consistent-hash requests
+ to replicas by the x-plano-prefix-hash header so the same prompt prefix lands
+ on the same replica and reuses its KV cache. Uses ring-hash load balancing
+ across all resolved endpoint addresses. Default false.
additionalProperties: false
required:
- endpoint
@@ -316,6 +323,33 @@ properties:
orchestrator_model_context_length:
type: integer
description: "Maximum token length for the orchestrator/routing model context window. Default is 8192."
+ prompt_caching:
+ type: object
+ description: >
+ Automatic provider prompt caching, configured once for the whole Plano instance.
+ Disabled by default; set enabled: true to opt in. Prompt caching never changes
+ which model routing selects — it only keeps a conversation on the same
+ model/provider so the upstream prompt cache stays warm across turns, and
+ auto-injects provider cache-control markers where supported.
+ properties:
+ enabled:
+ type: boolean
+ description: "Master switch. Default false (opt-in). Applies across the entire instance."
+ session_affinity:
+ type: boolean
+ description: "Auto-derive a session key from the prompt prefix (system + tools + first user message) when X-Model-Affinity is absent, so follow-up turns reuse the same warm cache. Default true when enabled."
+ inject_cache_control:
+ type: boolean
+ description: "Auto-insert ephemeral cache_control breakpoints for providers that need explicit markers (Anthropic). Default true when enabled."
+ min_prefix_tokens:
+ type: integer
+ minimum: 1
+ description: "Skip breakpoint injection when the estimated stable prefix is below this many tokens. Default 1024."
+ session_ttl_seconds:
+ type: integer
+ minimum: 1
+ description: "Pin lifetime for implicit/explicit sessions; align with the provider cache window (e.g. 3600 for Anthropic 1h caching). Defaults to routing.session_ttl_seconds."
+ additionalProperties: false
system_prompt:
type: string
prompt_targets:
@@ -447,6 +481,28 @@ properties:
additionalProperties:
type: string
additionalProperties: false
+ exporters:
+ type: array
+ items:
+ oneOf:
+ - type: object
+ properties:
+ type:
+ type: string
+ const: posthog
+ url:
+ type: string
+ api_key:
+ type: string
+ distinct_id_header:
+ type: string
+ capture_messages:
+ type: boolean
+ additionalProperties: false
+ required:
+ - type
+ - url
+ - api_key
additionalProperties: false
mode:
type: string
@@ -488,6 +544,53 @@ properties:
Optional HTTP header name whose value is used as a tenant prefix in the cache key.
When set, keys are scoped as plano:affinity:{tenant_id}:{session_id}.
additionalProperties: false
+ routing_budget:
+ type: object
+ description: >
+ Per-session cost gate on model switching. Self-sufficient and independent of
+ prompt_caching — it applies whenever configured (presence of this block turns
+ it on), derives implicit sessions on its own, and always prices warm anchors
+ at their cached input rate (provider caches are assumed real). The default
+ posture is to stick to the model a session is warm on. When routing proposes a
+ different model while that session's provider cache is plausibly still warm
+ (inferred from the idle gap vs. the provider's cache window), the actual
+ input-token cost of abandoning the cache —
+ context_tokens x (candidate_uncached_input_rate - anchor_cached_input_rate),
+ output cost deliberately excluded — accrues into the session's cumulative
+ switch spend. A paid switch is allowed only while that spend stays within
+ max_overhead_pct percent of the session's running never-switch baseline (what
+ staying on the anchor would have cost). An outright-cheaper switch is free but
+ never reduces the spend. Requires a cost source in model_metrics_sources.
+ properties:
+ max_overhead_pct:
+ type: number
+ minimum: 0
+ description: >
+ Cap on cumulative switching overhead, as a percentage of what the session
+ would have cost by never switching (a whole number: 20 = 20%). The promise
+ is "this conversation bills at most max_overhead_pct% above never-switching."
+ 0 means never pay to switch (only outright-cheaper switches are allowed);
+ larger values buy more quality-driven switches. Typical range 10-30.
+ replenish_on_rebind:
+ type: boolean
+ description: "Reset the running baseline/spend totals when a cold session re-binds. Default true."
+ cache_read_discount:
+ type: number
+ minimum: 0
+ maximum: 1
+ description: >
+ Fallback used to estimate a model's cached input rate when the pricing feed
+ doesn't publish one (cached_rate = input_rate x discount). A pricing detail,
+ not a cost policy. Default 0.1.
+ record_counterfactual:
+ type: boolean
+ description: >
+ When true, a vetoed switch records the route the gate would have taken had
+ the switch been allowed, as the plano.switch.counterfactual_route span
+ attribute. Telemetry only — the counterfactual model is never dispatched.
+ Useful for evals/benchmarks. Default false.
+ required: ["max_overhead_pct"]
+ additionalProperties: false
additionalProperties: false
state_storage:
type: object
@@ -582,13 +685,17 @@ properties:
type: string
enum:
- digitalocean
+ - models.dev
+ url:
+ type: string
+ description: "Optional override for the pricing catalog endpoint. Defaults per provider (digitalocean: DO GenAI catalog; models.dev: https://models.dev/api.json)."
refresh_interval:
type: integer
minimum: 1
description: "Refresh interval in seconds"
model_aliases:
type: object
- description: "Map DO catalog keys (lowercase(creator)/model_id) to Plano model names used in routing_preferences. Example: 'openai/openai-gpt-oss-120b: openai/gpt-4o'"
+ description: "Map catalog keys to Plano model names used in routing_preferences. DigitalOcean keys are 'lowercase(creator)/model_id'; models.dev keys are 'creator/model_id'. Example: 'openai/openai-gpt-oss-120b: openai/gpt-4o'"
additionalProperties:
type: string
required:
diff --git a/crates/brightstaff/src/affinity.rs b/crates/brightstaff/src/affinity.rs
new file mode 100644
index 00000000..bf05d956
--- /dev/null
+++ b/crates/brightstaff/src/affinity.rs
@@ -0,0 +1,181 @@
+//! Implicit session-affinity key derivation.
+//!
+//! When a request carries no explicit `X-Model-Affinity` header, Plano derives a
+//! stable session key from the parts of the prompt that repeat verbatim at the head
+//! of every turn — the same bytes the provider's prompt cache is keyed on:
+//!
+//! ```text
+//! session_key = hash(system + tools + first_user_message)
+//! prefix_hash = hash(system + tools)
+//! ```
+//!
+//! The session key is constant for the life of a conversation (history grows at the
+//! tail, not the head), so turns 2+ reuse the same pin without any client changes.
+//! The prefix hash covers only the fully-stable segment and is stored with the pin
+//! for drift detection: if it changes, the provider cache is already lost and
+//! re-routing fresh is safe.
+//!
+//! Only salted hashes are ever stored — never prompt content.
+
+use hermesllm::apis::openai::{Message, Role};
+
+/// Salt folded into every hash so stored keys can't be trivially correlated with
+/// prompt content across systems. Deterministic across processes/replicas so a
+/// shared Redis session cache keys consistently.
+const HASH_SALT: &str = "plano-affinity-v1";
+
+/// Prefix distinguishing derived keys from client-supplied `X-Model-Affinity` ids.
+const IMPLICIT_KEY_PREFIX: &str = "implicit:";
+
+/// Derived affinity identifiers for one request.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct ImplicitAffinity {
+ /// Session-cache key: `implicit:{hex}` over system + tools + first user message.
+ pub session_key: String,
+ /// Hash of the stable prefix only (system + tools), for drift detection.
+ pub prefix_hash: u64,
+}
+
+/// FNV-1a 64-bit — stable across processes and Rust versions (unlike `DefaultHasher`),
+/// dependency-free, and plenty for cache keying (collisions merely over-pin, which is
+/// cache-friendly).
+fn fnv1a64(chunks: &[&str]) -> u64 {
+ const OFFSET_BASIS: u64 = 0xcbf2_9ce4_8422_2325;
+ const PRIME: u64 = 0x0000_0100_0000_01b3;
+ let mut hash = OFFSET_BASIS;
+ let mut feed = |bytes: &[u8]| {
+ for &b in bytes {
+ hash ^= b as u64;
+ hash = hash.wrapping_mul(PRIME);
+ }
+ // Field separator so ("ab","c") and ("a","bc") hash differently.
+ hash ^= 0x1f;
+ hash = hash.wrapping_mul(PRIME);
+ };
+ feed(HASH_SALT.as_bytes());
+ for chunk in chunks {
+ feed(chunk.as_bytes());
+ }
+ hash
+}
+
+/// Derive the implicit affinity key from parsed request messages and tool names.
+///
+/// Returns `None` when there is no user message to anchor on (nothing distinguishes
+/// the conversation, so pinning would be meaningless).
+pub fn derive_implicit_affinity(
+ messages: &[Message],
+ tool_names: Option<&[String]>,
+ tenant_id: Option<&str>,
+) -> Option {
+ let system_text: String = messages
+ .iter()
+ .filter(|m| matches!(m.role, Role::System | Role::Developer))
+ .filter_map(|m| m.content.as_ref().map(|c| c.to_string()))
+ .collect::>()
+ .join("\n");
+
+ let first_user = messages
+ .iter()
+ .find(|m| matches!(m.role, Role::User))
+ .and_then(|m| m.content.as_ref().map(|c| c.to_string()))?;
+
+ let tools_text = tool_names.map(|names| names.join(",")).unwrap_or_default();
+ let tenant = tenant_id.unwrap_or_default();
+
+ let prefix_hash = fnv1a64(&[tenant, &system_text, &tools_text]);
+ let session_hash = fnv1a64(&[tenant, &system_text, &tools_text, &first_user]);
+
+ Some(ImplicitAffinity {
+ session_key: format!("{IMPLICIT_KEY_PREFIX}{session_hash:016x}"),
+ prefix_hash,
+ })
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use hermesllm::apis::openai::MessageContent;
+
+ fn msg(role: Role, text: &str) -> Message {
+ Message {
+ role,
+ content: Some(MessageContent::Text(text.to_string())),
+ name: None,
+ tool_call_id: None,
+ tool_calls: None,
+ }
+ }
+
+ #[test]
+ fn key_is_stable_as_history_grows() {
+ let turn1 = vec![
+ msg(Role::System, "You are a coding agent."),
+ msg(Role::User, "Fix the bug in main.rs"),
+ ];
+ let turn5 = vec![
+ msg(Role::System, "You are a coding agent."),
+ msg(Role::User, "Fix the bug in main.rs"),
+ msg(Role::Assistant, "Done. Anything else?"),
+ msg(Role::User, "Now add tests"),
+ msg(Role::Assistant, "Added."),
+ ];
+
+ let a1 = derive_implicit_affinity(&turn1, None, None).unwrap();
+ let a5 = derive_implicit_affinity(&turn5, None, None).unwrap();
+ assert_eq!(a1.session_key, a5.session_key);
+ assert_eq!(a1.prefix_hash, a5.prefix_hash);
+ assert!(a1.session_key.starts_with("implicit:"));
+ }
+
+ #[test]
+ fn different_first_user_message_yields_different_session() {
+ let base = msg(Role::System, "You are a coding agent.");
+ let a = derive_implicit_affinity(
+ &[base.clone(), msg(Role::User, "conversation A")],
+ None,
+ None,
+ )
+ .unwrap();
+ let b = derive_implicit_affinity(&[base, msg(Role::User, "conversation B")], None, None)
+ .unwrap();
+ // Same stable prefix, different conversations.
+ assert_eq!(a.prefix_hash, b.prefix_hash);
+ assert_ne!(a.session_key, b.session_key);
+ }
+
+ #[test]
+ fn changed_system_prompt_changes_prefix_hash() {
+ let a = derive_implicit_affinity(
+ &[msg(Role::System, "v1 prompt"), msg(Role::User, "hi")],
+ None,
+ None,
+ )
+ .unwrap();
+ let b = derive_implicit_affinity(
+ &[msg(Role::System, "v2 prompt"), msg(Role::User, "hi")],
+ None,
+ None,
+ )
+ .unwrap();
+ assert_ne!(a.prefix_hash, b.prefix_hash);
+ assert_ne!(a.session_key, b.session_key);
+ }
+
+ #[test]
+ fn tools_and_tenant_are_part_of_the_key() {
+ let messages = [msg(Role::System, "s"), msg(Role::User, "u")];
+ let plain = derive_implicit_affinity(&messages, None, None).unwrap();
+ let with_tools =
+ derive_implicit_affinity(&messages, Some(&["get_weather".to_string()]), None).unwrap();
+ let with_tenant = derive_implicit_affinity(&messages, None, Some("acme")).unwrap();
+ assert_ne!(plain.session_key, with_tools.session_key);
+ assert_ne!(plain.session_key, with_tenant.session_key);
+ }
+
+ #[test]
+ fn no_user_message_yields_none() {
+ assert!(derive_implicit_affinity(&[msg(Role::System, "s")], None, None).is_none());
+ assert!(derive_implicit_affinity(&[], None, None).is_none());
+ }
+}
diff --git a/crates/brightstaff/src/app_state.rs b/crates/brightstaff/src/app_state.rs
index 1d534e89..da8d15b9 100644
--- a/crates/brightstaff/src/app_state.rs
+++ b/crates/brightstaff/src/app_state.rs
@@ -1,7 +1,10 @@
use std::collections::HashMap;
use std::sync::Arc;
-use common::configuration::{Agent, FilterPipeline, Listener, ModelAlias, SpanAttributes};
+use common::configuration::{
+ Agent, EffectivePromptCaching, EffectiveRoutingBudget, FilterPipeline, Listener, ModelAlias,
+ SpanAttributes,
+};
use common::llm_providers::LlmProviders;
use tokio::sync::RwLock;
@@ -21,10 +24,20 @@ pub struct AppState {
pub state_storage: Option>,
pub llm_provider_url: String,
pub span_attributes: Option,
+ /// Request header whose value populates the observability `distinct_id`
+ /// (e.g. PostHog). Sourced from `tracing.exporters[].distinct_id_header`.
+ /// `None` means LLM events are captured anonymously.
+ pub distinct_id_header: Option,
/// Shared HTTP client for upstream LLM requests (connection pooling / keep-alive).
pub http_client: reqwest::Client,
pub filter_pipeline: Arc,
/// When false, agentic signal analysis is skipped on LLM responses to save CPU.
/// Controlled by `overrides.disable_signals` in plano config.
pub signals_enabled: bool,
+ /// Instance-wide automatic prompt-caching settings, resolved once from the
+ /// top-level `prompt_caching` config. Disabled by default (opt-in).
+ pub prompt_caching: EffectivePromptCaching,
+ /// Per-session model-switch cost gate, resolved from `routing.routing_budget`.
+ /// Independent of prompt caching; `None` when not configured (off by default).
+ pub routing_budget: Option,
}
diff --git a/crates/brightstaff/src/handlers/function_calling.rs b/crates/brightstaff/src/handlers/function_calling.rs
index 3e2543bc..24ca5c52 100644
--- a/crates/brightstaff/src/handlers/function_calling.rs
+++ b/crates/brightstaff/src/handlers/function_calling.rs
@@ -1239,6 +1239,7 @@ impl ArchFunctionHandler {
total_tokens: 0,
prompt_tokens_details: None,
completion_tokens_details: None,
+ ..Default::default()
},
system_fingerprint: None,
service_tier: None,
diff --git a/crates/brightstaff/src/handlers/llm/mod.rs b/crates/brightstaff/src/handlers/llm/mod.rs
index 3336209f..5cccf930 100644
--- a/crates/brightstaff/src/handlers/llm/mod.rs
+++ b/crates/brightstaff/src/handlers/llm/mod.rs
@@ -1,6 +1,9 @@
use bytes::Bytes;
use common::configuration::{FilterPipeline, ModelAlias};
-use common::consts::{ARCH_IS_STREAMING_HEADER, ARCH_PROVIDER_HINT_HEADER, MODEL_AFFINITY_HEADER};
+use common::consts::{
+ ARCH_IS_STREAMING_HEADER, ARCH_PROVIDER_HINT_HEADER, MODEL_AFFINITY_HEADER, PLANO_CACHE_HEADER,
+ PLANO_PREFIX_HASH_HEADER,
+};
use common::llm_providers::LlmProviders;
use hermesllm::apis::openai::Message;
use hermesllm::apis::openai_responses::InputParam;
@@ -19,6 +22,8 @@ use tokio::sync::RwLock;
use tracing::{debug, info, info_span, warn, Instrument};
pub(crate) mod model_selection;
+pub(crate) mod prompt_caching;
+pub(crate) mod session_router;
use crate::app_state::AppState;
use crate::handlers::agents::pipeline::PipelineProcessor;
@@ -31,7 +36,7 @@ use crate::state::{
};
use crate::streaming::{
create_streaming_response, create_streaming_response_with_output_filter, truncate_message,
- LlmMetricsCtx, ObservableStreamProcessor, StreamProcessor,
+ LlmMetricsCtx, ObservableStreamProcessor, SessionUpdateCtx, StreamProcessor,
};
use crate::tracing::{
collect_custom_trace_attributes, llm as tracing_llm, operation_component,
@@ -93,8 +98,26 @@ async fn llm_chat_inner(
}
});
- // Session pinning: extract session ID and check cache before routing
- let session_id: Option = request_headers
+ // Stamp the caller identity for downstream exporters (e.g. PostHog
+ // `distinct_id`). Sourced from the configured `distinct_id_header`; when the
+ // header is absent the event is exported anonymously.
+ if let Some(header_name) = state.distinct_id_header.as_deref() {
+ if let Some(distinct_id) = request_headers
+ .get(header_name)
+ .and_then(|v| v.to_str().ok())
+ .map(str::trim)
+ .filter(|s| !s.is_empty())
+ {
+ get_active_span(|span| {
+ span.set_attribute(opentelemetry::KeyValue::new(
+ tracing_plano::DISTINCT_ID,
+ distinct_id.to_string(),
+ ));
+ });
+ }
+ }
+
+ let explicit_session_id: Option = request_headers
.get(MODEL_AFFINITY_HEADER)
.and_then(|h| h.to_str().ok())
.map(|s| s.to_string());
@@ -104,40 +127,17 @@ async fn llm_chat_inner(
.and_then(|hdr| request_headers.get(hdr))
.and_then(|v| v.to_str().ok())
.map(|s| s.to_string());
- let cached_route = if let Some(ref sid) = session_id {
- state
- .orchestrator_service
- .get_cached_route(sid, tenant_id.as_deref())
- .await
- } else {
- None
- };
- let (pinned_model, pinned_route_name): (Option, Option) = match cached_route {
- Some(c) => (Some(c.model_name), c.route_name),
- None => (None, None),
- };
-
- // Record session id on the LLM span for the observability console.
- if let Some(ref sid) = session_id {
- get_active_span(|span| {
- span.set_attribute(opentelemetry::KeyValue::new(
- tracing_plano::SESSION_ID,
- sid.clone(),
- ));
- });
- }
- if let Some(ref route_name) = pinned_route_name {
- get_active_span(|span| {
- span.set_attribute(opentelemetry::KeyValue::new(
- tracing_plano::ROUTE_NAME,
- route_name.clone(),
- ));
- });
- }
+ // `X-Plano-Cache: off` disables implicit pinning + marker injection for one call.
+ let cache_off_for_request = request_headers
+ .get(PLANO_CACHE_HEADER)
+ .and_then(|h| h.to_str().ok())
+ .is_some_and(|v| v.eq_ignore_ascii_case("off"));
let full_qualified_llm_provider_url = format!("{}{}", state.llm_provider_url, request_path);
// --- Phase 1: Parse and validate the incoming request ---
+ // (Parsing happens before session-key derivation: the implicit affinity key is
+ // computed from the request body's stable prompt prefix.)
let parsed = match parse_and_validate_request(
request,
&request_path,
@@ -168,6 +168,14 @@ async fn llm_chat_inner(
provider_id,
} = parsed;
+ // Prompt caching is configured once for the whole instance and never influences
+ // which model routing selects — it only keeps a conversation on the same
+ // model/provider so the upstream prompt cache stays warm across turns. Note:
+ // session affinity and pin lookup are derived later (Phase 2a), *after* input
+ // filters and Responses-API state merge, so the prefix hash reflects the exact
+ // stable prefix sent upstream rather than the pre-filter request body.
+ let prompt_caching = state.prompt_caching;
+
// Record LLM-specific span attributes
let span = tracing::Span::current();
if let Some(temp) = temperature {
@@ -264,6 +272,18 @@ async fn llm_chat_inner(
*bad_request.status_mut() = StatusCode::BAD_REQUEST;
return Ok(bad_request);
}
+
+ // Auto-inject prompt-cache markers (grouped in `prompt_caching`). No-op unless
+ // caching is enabled for this request and the model needs explicit markers.
+ prompt_caching::inject_cache_markers(
+ &mut client_request,
+ provider_id,
+ &model_name_only,
+ &upstream_api,
+ &prompt_caching,
+ cache_off_for_request,
+ &alias_resolved_model,
+ );
}
// --- Phase 2: Resolve conversation state (v1/responses API) ---
@@ -282,6 +302,40 @@ async fn llm_chat_inner(
Err(response) => return Ok(response),
};
+ // --- Phase 2a: Session affinity (see `session_router`) ---
+ // Derived here, after input filters and Responses-API state merge, so the
+ // implicit session key and prefix hash are computed over the exact stable
+ // prefix (system + tools + first user message) that is actually sent upstream —
+ // the same bytes the provider's prompt cache is keyed on. The prefix hash is
+ // derived even when caching is off, so the `x-plano-prefix-hash` RING_HASH
+ // replica-stickiness header still works.
+ let routing_budget = state.routing_budget;
+ // Derive the implicit session key when either prompt-caching affinity or the
+ // routing budget is active — the budget needs a session anchor even with caching off.
+ let implicit_affinity_enabled = prompt_caching.session_affinity || routing_budget.is_some();
+ let request_messages = client_request.get_messages();
+ let session_router::SessionResolution {
+ request_prefix_hash,
+ session_id,
+ } = session_router::resolve_session(
+ explicit_session_id,
+ &request_messages,
+ tool_names.as_deref(),
+ tenant_id.as_deref(),
+ implicit_affinity_enabled,
+ cache_off_for_request,
+ );
+
+ // Record session id on the LLM span for the observability console.
+ if let Some(ref sid) = session_id {
+ get_active_span(|span| {
+ span.set_attribute(opentelemetry::KeyValue::new(
+ tracing_plano::SESSION_ID,
+ sid.clone(),
+ ));
+ });
+ }
+
// Serialize request for upstream BEFORE router consumes it
let client_request_bytes_for_upstream: Bytes =
match ProviderRequestType::to_bytes(&client_request) {
@@ -294,78 +348,151 @@ async fn llm_chat_inner(
}
};
- // --- Phase 3: Route the request (or use pinned model from session cache) ---
- let resolved_model = if let Some(cached_model) = pinned_model {
- info!(
- session_id = %session_id.as_deref().unwrap_or(""),
- model = %cached_model,
- "using pinned routing decision from cache"
- );
- cached_model
+ // Context size for the switch-cost estimate, computed before the router consumes
+ // the request. Only needed when stickiness could act on a session.
+ let context_tokens: u64 = if session_id.is_some() && routing_budget.is_some() {
+ session_router::actual_context_tokens(&request_messages, &model_name_only)
} else {
- let routing_span = info_span!(
- "routing",
- component = "routing",
- http.method = "POST",
- http.target = %request_path,
- model.requested = %model_from_request,
- model.alias_resolved = %alias_resolved_model,
- route.selected_model = tracing::field::Empty,
- routing.determination_ms = tracing::field::Empty,
- );
- let routing_result = match async {
- set_service_name(operation_component::ROUTING);
- router_chat_get_upstream_model(
- Arc::clone(&state.orchestrator_service),
- client_request,
- &request_path,
- &request_id,
- inline_routing_preferences,
- )
- .await
- }
- .instrument(routing_span)
- .await
- {
- Ok(result) => result,
- Err(err) => {
- let mut internal_error = Response::new(full(err.message));
- *internal_error.status_mut() = err.status_code;
- return Ok(internal_error);
- }
- };
-
- let (router_selected_model, route_name) =
- (routing_result.model_name, routing_result.route_name);
- let model = if router_selected_model != "none" {
- router_selected_model
- } else {
- alias_resolved_model.clone()
- };
-
- // Record route name on the LLM span (only when the orchestrator produced one).
- if let Some(ref rn) = route_name {
- if !rn.is_empty() && rn != "none" {
- get_active_span(|span| {
- span.set_attribute(opentelemetry::KeyValue::new(
- tracing_plano::ROUTE_NAME,
- rn.clone(),
- ));
- });
- }
- }
-
- if let Some(ref sid) = session_id {
- state
- .orchestrator_service
- .cache_route(sid.clone(), tenant_id.as_deref(), model.clone(), route_name)
- .await;
- }
-
- model
+ 0
};
+
+ // --- Phase 3: Route the request (quality), then apply the session-cache decision ---
+ // Routing stays cache-blind: the quality router always picks a candidate. The
+ // session router then honors it or sticks to the warm anchor (see `session_router`).
+ let routing_span = info_span!(
+ "routing",
+ component = "routing",
+ http.method = "POST",
+ http.target = %request_path,
+ model.requested = %model_from_request,
+ model.alias_resolved = %alias_resolved_model,
+ route.selected_model = tracing::field::Empty,
+ routing.determination_ms = tracing::field::Empty,
+ );
+ let routing_result = match async {
+ set_service_name(operation_component::ROUTING);
+ router_chat_get_upstream_model(
+ Arc::clone(&state.orchestrator_service),
+ client_request,
+ &request_path,
+ &request_id,
+ inline_routing_preferences,
+ )
+ .await
+ }
+ .instrument(routing_span)
+ .await
+ {
+ Ok(result) => result,
+ Err(err) => {
+ let mut internal_error = Response::new(full(err.message));
+ *internal_error.status_mut() = err.status_code;
+ return Ok(internal_error);
+ }
+ };
+
+ let candidate_model = if routing_result.model_name != "none" {
+ routing_result.model_name
+ } else {
+ alias_resolved_model.clone()
+ };
+ let candidate_route = routing_result.route_name;
+
+ let decision = session_router::route(
+ &state.orchestrator_service,
+ routing_budget.as_ref(),
+ session_router::RouteFacts {
+ session_id: session_id.as_deref(),
+ tenant_id: tenant_id.as_deref(),
+ prefix_hash: request_prefix_hash,
+ context_tokens,
+ candidate_model: &candidate_model,
+ candidate_route: candidate_route.as_deref(),
+ },
+ )
+ .await;
+
+ let resolved_model = decision.model;
+ let resolved_route_name = decision.route_name;
+
+ // Record route name on the LLM span (only when a real route was produced).
+ if let Some(ref rn) = resolved_route_name {
+ if !rn.is_empty() && rn != "none" {
+ get_active_span(|span| {
+ span.set_attribute(opentelemetry::KeyValue::new(
+ tracing_plano::ROUTE_NAME,
+ rn.clone(),
+ ));
+ });
+ }
+ }
+
tracing::Span::current().record(tracing_llm::MODEL_NAME, resolved_model.as_str());
+ // Record the provider (derived from the `provider/model` prefix) so
+ // observability exporters can populate provider fields (e.g. PostHog
+ // `$ai_provider`).
+ let (resolved_provider, _) = bs_metrics::split_provider_model(&resolved_model);
+ if resolved_provider != "unknown" {
+ get_active_span(|span| {
+ span.set_attribute(opentelemetry::KeyValue::new(
+ tracing_llm::PROVIDER,
+ resolved_provider.to_string(),
+ ));
+ });
+ }
+
+ // Resolve the dispatched model's catalog rates now (this side is async; the
+ // response path that prices the turn is synchronous). `None` when no cost feed is
+ // configured → no per-request/session cost is computed.
+ let cost_rates = if session_id.is_some() {
+ state
+ .orchestrator_service
+ .model_rates(&resolved_model)
+ .await
+ } else {
+ None
+ };
+ let cache_read_discount = state
+ .routing_budget
+ .as_ref()
+ .map(|b| b.cache_read_discount)
+ .unwrap_or(common::configuration::DEFAULT_CACHE_READ_DISCOUNT);
+
+ // Response-side refresh: update `last_used` + the context-size estimate from the
+ // real response. The routing decision itself was already persisted by `route()`.
+ let session_update_ctx: Option =
+ session_id.as_ref().map(|sid| SessionUpdateCtx {
+ orchestrator: Arc::clone(&state.orchestrator_service),
+ session_id: sid.clone(),
+ tenant_id: tenant_id.clone(),
+ anchor_model: resolved_model.clone(),
+ default_model: decision.default_model.clone(),
+ route_name: resolved_route_name.clone(),
+ prefix_hash: request_prefix_hash,
+ baseline_usd: decision.baseline_usd,
+ switch_spend_usd: decision.switch_spend_usd,
+ switches: decision.switches,
+ history: decision.history.clone(),
+ session_cost_usd: decision.session_cost_usd,
+ cost_rates,
+ cache_read_discount,
+ context_tokens: decision.cached_tokens,
+ gc_ttl: decision.gc_ttl,
+ });
+
+ // Forward the prefix hash so self-hosted multi-replica backends can do
+ // KV-aware replica stickiness (consistent-hash on this header at the
+ // cluster layer).
+ if let Some(hash) = request_prefix_hash {
+ if let Ok(val) = header::HeaderValue::from_str(&format!("{hash:016x}")) {
+ request_headers.insert(
+ header::HeaderName::from_static(PLANO_PREFIX_HASH_HEADER),
+ val,
+ );
+ }
+ }
+
// --- Phase 4: Forward to upstream and stream back ---
send_upstream(
&state.http_client,
@@ -383,6 +510,7 @@ async fn llm_chat_inner(
state.state_storage.clone(),
request_id,
&state.filter_pipeline,
+ session_update_ctx,
)
.await
}
@@ -660,6 +788,7 @@ async fn send_upstream(
state_storage: Option>,
request_id: String,
filter_pipeline: &Arc,
+ session_update_ctx: Option,
) -> Result>, hyper::Error> {
let span_name = if model_from_request == resolved_model {
format!("POST {} {}", request_path, resolved_model)
@@ -775,7 +904,7 @@ async fn send_upstream(
let byte_stream = llm_response.bytes_stream();
// Create base processor for metrics and tracing
- let base_processor = ObservableStreamProcessor::new(
+ let mut base_processor = ObservableStreamProcessor::new(
operation_component::LLM,
span_name,
request_start_time,
@@ -786,6 +915,13 @@ async fn send_upstream(
model: metric_model.clone(),
upstream_status: upstream_status.as_u16(),
});
+ // Only refresh the session binding for successful responses — errors carry no
+ // usage block and shouldn't reset the warmth clock.
+ if upstream_status.is_success() {
+ if let Some(update_ctx) = session_update_ctx {
+ base_processor = base_processor.with_session_update(update_ctx);
+ }
+ }
let output_filter_request_headers = if filter_pipeline.has_output_filters() {
Some(request_headers.clone())
diff --git a/crates/brightstaff/src/handlers/llm/prompt_caching.rs b/crates/brightstaff/src/handlers/llm/prompt_caching.rs
new file mode 100644
index 00000000..2d0b5a08
--- /dev/null
+++ b/crates/brightstaff/src/handlers/llm/prompt_caching.rs
@@ -0,0 +1,70 @@
+//! Prompt-caching request handling for the LLM path.
+//!
+//! This module owns the "make the upstream provider's prompt cache work" concern:
+//! resolving the correct cache-marking strategy for the `(gateway × model family ×
+//! upstream API)` combination and injecting the markers into the outbound request.
+//! It never influences routing — see [`super::session_router`] for the session-cache
+//! lookup and the routing-budget switch-cost concern.
+
+use common::configuration::EffectivePromptCaching;
+use hermesllm::clients::SupportedUpstreamAPIs;
+use hermesllm::{cache_marker_strategy, CacheMarkerStrategy, ProviderId, ProviderRequestType};
+use tracing::debug;
+
+/// Auto-inject prompt-cache markers into `client_request`.
+///
+/// The strategy is resolved from `(gateway × model family × upstream API)` so that
+/// Anthropic-family models cache whether they arrive over the native Messages API or
+/// an OpenAI-compatible gateway (DigitalOcean, OpenRouter), while OpenAI-family models
+/// (which cache automatically) and unimplemented backends are left untouched.
+///
+/// A no-op when caching is disabled, `inject_cache_control` is off, the request opted
+/// out (`X-Plano-Cache: off`), or the provider caches automatically. Injection is
+/// idempotent (client-supplied markers are respected) and threshold-guarded against
+/// the provider's minimum cacheable prefix.
+pub fn inject_cache_markers(
+ client_request: &mut ProviderRequestType,
+ provider_id: ProviderId,
+ model_name_only: &str,
+ upstream_api: &SupportedUpstreamAPIs,
+ prompt_caching: &EffectivePromptCaching,
+ cache_off_for_request: bool,
+ alias_resolved_model: &str,
+) {
+ if !(prompt_caching.enabled && prompt_caching.inject_cache_control && !cache_off_for_request) {
+ return;
+ }
+
+ match cache_marker_strategy(provider_id, model_name_only, upstream_api) {
+ CacheMarkerStrategy::AnthropicMessagesBreakpoints {
+ min_prefix_tokens, ..
+ } => {
+ if let ProviderRequestType::MessagesRequest(req) = client_request {
+ let threshold = prompt_caching.min_prefix_tokens.max(min_prefix_tokens);
+ if req.inject_cache_breakpoints(threshold) {
+ debug!(
+ model = %alias_resolved_model,
+ min_prefix_tokens = threshold,
+ "injected anthropic ephemeral cache breakpoints"
+ );
+ }
+ }
+ }
+ CacheMarkerStrategy::OpenAiContentPartCacheControl {
+ min_prefix_tokens,
+ ttl,
+ } => {
+ if let ProviderRequestType::ChatCompletionsRequest(req) = client_request {
+ let threshold = prompt_caching.min_prefix_tokens.max(min_prefix_tokens);
+ if req.inject_cache_control(ttl, threshold) {
+ debug!(
+ model = %alias_resolved_model,
+ min_prefix_tokens = threshold,
+ "injected openai content-part cache_control"
+ );
+ }
+ }
+ }
+ CacheMarkerStrategy::Automatic | CacheMarkerStrategy::None => {}
+ }
+}
diff --git a/crates/brightstaff/src/handlers/llm/session_router.rs b/crates/brightstaff/src/handlers/llm/session_router.rs
new file mode 100644
index 00000000..651b517d
--- /dev/null
+++ b/crates/brightstaff/src/handlers/llm/session_router.rs
@@ -0,0 +1,1010 @@
+//! Session-cache-aware routing.
+//!
+//! Routing itself stays cache-blind: the `llm_router` (quality) still picks a
+//! candidate model for every request. This module then decides whether to *honor*
+//! that candidate or stick to the session's warm anchor, based on:
+//!
+//! * **Cache warmth** — inferred structurally from how long ago the session was last
+//! used vs. the provider's cache window ([`hermesllm::provider_cache_capability`]),
+//! so it works on the decision path with no provider response in hand.
+//! * **A cumulative per-session overhead cap** — a paid switch (the candidate must
+//! re-ingest the context at its uncached rate) is allowed only while total switch
+//! spend stays within `max_overhead_pct`% of the session's running *never-switch*
+//! baseline (what staying on the session's `default_model` would have cost — priced
+//! independently of the current anchor, which drifts as switches happen). An
+//! outright-cheaper switch is free but never reduces that spend. The promise: this
+//! conversation bills at most `max_overhead_pct`% above never-switching. A warm
+//! anchor is priced at its *cached* rate and the switch pays the cache-loss delta:
+//! provider caches are assumed real whenever the budget is active (most providers
+//! cache automatically), independent of `prompt_caching.enabled`, which only
+//! controls Plano's own marker injection and caching-without-budget affinity.
+//!
+//! The default posture is to stick. Quality and cost stay separate: the router decides
+//! whether a switch *improves quality*; the overhead cap decides whether it is *affordable*.
+//!
+//! Prompt-cache *marker injection* is a separate concern — see [`super::prompt_caching`].
+
+use std::time::{Duration, SystemTime};
+
+use common::configuration::EffectiveRoutingBudget;
+use hermesllm::apis::openai::Message;
+use hermesllm::{provider_cache_capability, ProviderCacheCapability, ProviderId};
+use opentelemetry::trace::get_active_span;
+use opentelemetry::KeyValue;
+use tracing::{debug, info};
+
+use crate::affinity::derive_implicit_affinity;
+use crate::metrics as bs_metrics;
+use crate::metrics::labels as metric_labels;
+use crate::router::orchestrator::OrchestratorService;
+use crate::session_cache::{record_route_visit, RouteVisit, SessionBinding};
+use crate::tracing::plano as tracing_plano;
+
+/// Resolved session identity for one request.
+pub struct SessionResolution {
+ /// Stable prefix hash (system + tools + first user message), independent of
+ /// `prompt_caching.enabled` so it can still drive the `x-plano-prefix-hash`
+ /// RING_HASH replica-stickiness header. `None` when the request opted out or has
+ /// no anchorable prompt.
+ pub request_prefix_hash: Option,
+ /// Session key: the explicit `X-Model-Affinity` value, or the implicit prefix-hash
+ /// key when implicit affinity is active. `None` when there is nothing to anchor to.
+ pub session_id: Option,
+}
+
+/// Resolve the session key and prefix hash from the (already filtered / state-merged)
+/// request. An explicit affinity header always anchors; the implicit key is derived
+/// when `implicit_affinity_enabled` is set — true when either prompt caching's
+/// `session_affinity` or the routing budget is active, so stickiness works whether or
+/// not prompt caching is enabled. The prefix hash is derived regardless (only
+/// `X-Plano-Cache: off` or an unanchorable prompt suppresses it) so the
+/// `x-plano-prefix-hash` RING_HASH replica-stickiness header still works.
+pub fn resolve_session(
+ explicit_session_id: Option,
+ messages: &[Message],
+ tool_names: Option<&[String]>,
+ tenant_id: Option<&str>,
+ implicit_affinity_enabled: bool,
+ cache_off_for_request: bool,
+) -> SessionResolution {
+ let implicit_affinity = if cache_off_for_request {
+ None
+ } else {
+ derive_implicit_affinity(messages, tool_names, tenant_id)
+ };
+ let request_prefix_hash = implicit_affinity.as_ref().map(|a| a.prefix_hash);
+
+ let session_id = match explicit_session_id {
+ Some(sid) => Some(sid),
+ None if implicit_affinity_enabled && !cache_off_for_request => {
+ implicit_affinity.as_ref().map(|a| a.session_key.clone())
+ }
+ None => None,
+ };
+
+ SessionResolution {
+ request_prefix_hash,
+ session_id,
+ }
+}
+
+/// Extra memory retention beyond the warmth window, so a still-warm binding is never
+/// GC'd out from under the router before it could plausibly go cold.
+const GC_SLACK: Duration = Duration::from_secs(60);
+
+/// Stable request facts the router reasons about. Independent of the transport (full
+/// proxy vs. decision endpoint) so both paths route identically.
+pub struct RouteFacts<'a> {
+ /// Session key (explicit `X-Model-Affinity` or the implicit prefix key). `None`
+ /// disables stickiness for this request (nothing to anchor to).
+ pub session_id: Option<&'a str>,
+ pub tenant_id: Option<&'a str>,
+ /// Stable prompt-prefix hash; a mismatch vs. the stored binding means the provider
+ /// cache is already lost, so a switch is free.
+ pub prefix_hash: Option,
+ /// Context size in tokens (the tokens a switch would re-ingest). The request-side
+ /// count of the real messages; the binding's usage-refined count is preferred when
+ /// warm (see [`actual_context_tokens`]).
+ pub context_tokens: u64,
+ /// The model the quality router picked for this request.
+ pub candidate_model: &'a str,
+ pub candidate_route: Option<&'a str>,
+}
+
+/// The routing decision plus the session state to carry into the response side.
+pub struct RouteDecision {
+ /// The model to actually dispatch to (the anchor when a switch was vetoed).
+ pub model: String,
+ pub route_name: Option,
+ /// The session's never-switch model for this episode — carried to the response side
+ /// so the usage-refresh preserves it on the binding.
+ pub default_model: String,
+ /// Whether the session's cache was inferred warm at decision time.
+ pub warm: bool,
+ /// Cumulative never-switch baseline (USD) after this decision.
+ pub baseline_usd: f64,
+ /// Cumulative switch spend (USD) after this decision.
+ pub switch_spend_usd: f64,
+ /// Cumulative actual conversation cost (USD) so far — carried to the response side,
+ /// which adds this turn's real cost and re-persists it.
+ pub session_cost_usd: f64,
+ /// Cumulative switches taken this session (after this decision).
+ pub switches: u32,
+ /// Bounded per-model route history after this decision — carried to the response
+ /// side so the usage-refresh preserves it (and refines the anchor's token count).
+ pub history: Vec,
+ /// Context-token estimate persisted with the binding (refined later from usage).
+ pub cached_tokens: u64,
+ /// GC bound the binding was stored with (reused when the response side refreshes).
+ pub gc_ttl: Duration,
+}
+
+/// Count the request's context size in tokens from the real message content, using the
+/// tiktoken-based counter when available and falling back to the chars/4 heuristic. This
+/// is the request-side figure; on the full-proxy path the binding is later refined with
+/// the provider's own reported prompt-token count (see [`SessionBinding::cached_tokens`]),
+/// which the router prefers when present.
+pub fn actual_context_tokens(messages: &[Message], model: &str) -> u64 {
+ let text: String = messages
+ .iter()
+ .filter_map(|m| m.content.as_ref().map(|c| c.to_string()))
+ .collect::>()
+ .join("\n");
+ match common::tokenizer::token_count(model, &text) {
+ Ok(count) => count as u64,
+ Err(_) => (text.len() / 4) as u64,
+ }
+}
+
+/// Resolve a provider-qualified model id (e.g. `openai/gpt-4o`) to its cache window.
+/// Unknown providers fall back to the conservative default.
+fn capability_for_model(model: &str) -> ProviderCacheCapability {
+ let provider_part = model.split_once('/').map(|(p, _)| p).unwrap_or(model);
+ ProviderId::try_from(provider_part)
+ .map(provider_cache_capability)
+ .unwrap_or_default()
+}
+
+/// How long a binding on this model can sit idle before its cache is certainly cold.
+fn warmth_window(cap: &ProviderCacheCapability) -> Duration {
+ if cap.extended_retention {
+ cap.extended_ttl
+ } else {
+ cap.idle_ttl.min(cap.hard_ttl)
+ }
+}
+
+/// Whether the session's provider cache is plausibly still warm given how long ago it
+/// was last used. Returns the warmth verdict and the measured idle gap.
+fn warmth(
+ binding: &SessionBinding,
+ cap: &ProviderCacheCapability,
+ now: SystemTime,
+) -> (bool, Duration) {
+ let idle = now
+ .duration_since(binding.last_used)
+ .unwrap_or(Duration::ZERO);
+ let warm = if cap.extended_retention {
+ idle <= cap.extended_ttl
+ } else {
+ idle <= cap.idle_ttl && idle <= cap.hard_ttl
+ };
+ (warm, idle)
+}
+
+/// Decide the final model for this request and persist the updated session binding.
+///
+/// Never overrides the router on a *cold* session — it only protects a warm cache. The
+/// returned [`RouteDecision`] carries the model to dispatch plus the session state the
+/// response side reuses when it refreshes the binding from real usage.
+pub async fn route(
+ orchestrator: &OrchestratorService,
+ routing_budget: Option<&EffectiveRoutingBudget>,
+ facts: RouteFacts<'_>,
+) -> RouteDecision {
+ let now = SystemTime::now();
+ let candidate_gc_ttl = warmth_window(&capability_for_model(facts.candidate_model)) + GC_SLACK;
+
+ // No session to anchor to: honor the candidate, persist nothing.
+ let Some(session_id) = facts.session_id else {
+ return RouteDecision {
+ model: facts.candidate_model.to_string(),
+ route_name: facts.candidate_route.map(str::to_string),
+ default_model: facts.candidate_model.to_string(),
+ warm: false,
+ baseline_usd: 0.0,
+ switch_spend_usd: 0.0,
+ session_cost_usd: 0.0,
+ switches: 0,
+ history: Vec::new(),
+ cached_tokens: facts.context_tokens,
+ gc_ttl: candidate_gc_ttl,
+ };
+ };
+
+ let existing = orchestrator.get_binding(session_id, facts.tenant_id).await;
+
+ // Warmth + prefix drift. A drifted prefix means the cache is already cold.
+ let (warm, idle) = match &existing {
+ Some(b) => warmth(b, &capability_for_model(&b.anchor_model), now),
+ None => (false, Duration::ZERO),
+ };
+ let drifted = match (
+ existing.as_ref().and_then(|b| b.prefix_hash),
+ facts.prefix_hash,
+ ) {
+ (Some(stored), Some(current)) => stored != current,
+ _ => false,
+ };
+ let effective_warm = warm && !drifted;
+
+ // Cumulative actual conversation cost so far (through prior turns). Conversation-
+ // level: preserved across warm/cold re-binds; the response side adds this turn.
+ let session_cost_usd = existing.as_ref().map(|b| b.session_cost_usd).unwrap_or(0.0);
+
+ // Resolve the final model, cumulative baseline/spend, switch count, and telemetry.
+ let mut model = facts.candidate_model.to_string();
+ let mut route_name = facts.candidate_route.map(str::to_string);
+ // The session's never-switch model for this episode — priced into the baseline.
+ let default_model;
+ let baseline_usd;
+ let mut switch_spend_usd;
+ let mut switches;
+ let mut cost_opt: Option = None;
+ let mut ceiling_opt: Option = None;
+ let mut candidate_warm_tokens: u64 = 0;
+ let mut counterfactual: Option = None;
+ let decision_label: &'static str;
+ let reason: &'static str;
+
+ match existing.as_ref() {
+ Some(b) if effective_warm => {
+ switches = b.switches;
+ // The model the session would have stayed on had it never switched. Older
+ // bindings (persisted before this field existed) fall back to the anchor.
+ let session_default = if b.default_model.is_empty() {
+ b.anchor_model.clone()
+ } else {
+ b.default_model.clone()
+ };
+ // Prefer the provider's real prompt-token count from the prior turn over the
+ // request-side estimate — it's the actual context the session carries.
+ let context_tokens = if b.cached_tokens > 0 {
+ b.cached_tokens
+ } else {
+ facts.context_tokens
+ };
+ // Grow the never-switch baseline by this turn's read cost on the *default*
+ // model — the money the session would spend by never switching. This is the
+ // denominator the overhead cap is measured against. Missing pricing → no
+ // growth this turn.
+ let turn_baseline = match routing_budget {
+ Some(cfg) => orchestrator
+ .context_read_cost_in_usd(
+ context_tokens,
+ &session_default,
+ cfg.cache_read_discount,
+ )
+ .await
+ .unwrap_or(0.0),
+ None => 0.0,
+ };
+ baseline_usd = b.baseline_usd + turn_baseline;
+ switch_spend_usd = b.switch_spend_usd;
+ default_model = session_default;
+
+ if facts.candidate_model == b.anchor_model {
+ // Router agrees with the anchor — stick, no cost.
+ decision_label = metric_labels::SWITCH_DECISION_ALLOWED;
+ reason = metric_labels::SWITCH_REASON_SAME_ANCHOR;
+ } else if let Some(cfg) = routing_budget {
+ // Ceiling: at most `max_overhead_pct`% of the cumulative baseline may be
+ // spent on switching over this warm episode.
+ let ceiling = (cfg.max_overhead_pct / 100.0) * baseline_usd;
+ ceiling_opt = Some(ceiling);
+ // Credit any context the candidate still has cached from an earlier visit
+ // this session: a return to a still-warm model re-reads only the tokens
+ // appended since, not the whole context (the A→B→A case).
+ candidate_warm_tokens = b
+ .history
+ .iter()
+ .find(|v| v.model == facts.candidate_model)
+ .filter(|v| {
+ now.duration_since(v.last_used).unwrap_or(Duration::MAX)
+ <= warmth_window(&capability_for_model(facts.candidate_model))
+ })
+ .map(|v| v.cached_tokens.min(context_tokens))
+ .unwrap_or(0);
+ match orchestrator
+ .estimate_switch_cost_in_usd(
+ context_tokens,
+ &b.anchor_model,
+ facts.candidate_model,
+ candidate_warm_tokens,
+ cfg.cache_read_discount,
+ )
+ .await
+ {
+ // No pricing for one side — fail open (switch freely) rather than
+ // veto the router on guesswork.
+ None => {
+ switches += 1;
+ decision_label = metric_labels::SWITCH_DECISION_ALLOWED;
+ reason = metric_labels::SWITCH_REASON_NO_PRICING;
+ debug!(
+ anchor = %b.anchor_model,
+ candidate = %facts.candidate_model,
+ "switch allowed — missing pricing data, cannot gate"
+ );
+ }
+ Some(cost) => {
+ cost_opt = Some(cost);
+ if cost <= 0.0 {
+ // Outright cheaper: allowed for free. Does NOT reduce spend —
+ // the "saving" is vs a path we didn't take, not real money.
+ switches += 1;
+ decision_label = metric_labels::SWITCH_DECISION_ALLOWED;
+ reason = metric_labels::SWITCH_REASON_FREE;
+ info!(
+ anchor = %b.anchor_model,
+ candidate = %facts.candidate_model,
+ switch_cost_in_usd = cost,
+ "switch allowed — candidate is no more expensive than staying"
+ );
+ } else if switch_spend_usd + cost <= ceiling {
+ switch_spend_usd += cost;
+ switches += 1;
+ decision_label = metric_labels::SWITCH_DECISION_ALLOWED;
+ reason = metric_labels::SWITCH_REASON_WITHIN_CAP;
+ info!(
+ anchor = %b.anchor_model,
+ candidate = %facts.candidate_model,
+ switch_cost_in_usd = cost,
+ switch_spend_in_usd = switch_spend_usd,
+ overhead_ceiling_in_usd = ceiling,
+ "switch allowed — within session overhead cap"
+ );
+ } else {
+ // Unaffordable: retain the warm anchor.
+ if cfg.record_counterfactual {
+ counterfactual = Some(match route_name.as_deref() {
+ Some(rn) if !rn.is_empty() && rn != "none" => {
+ format!("{} ({rn})", facts.candidate_model)
+ }
+ _ => facts.candidate_model.to_string(),
+ });
+ }
+ model = b.anchor_model.clone();
+ route_name = b.route_name.clone();
+ decision_label = metric_labels::SWITCH_DECISION_RETAINED;
+ reason = metric_labels::SWITCH_REASON_OVER_CAP;
+ info!(
+ anchor = %b.anchor_model,
+ candidate = %facts.candidate_model,
+ switch_cost_in_usd = cost,
+ switch_spend_in_usd = switch_spend_usd,
+ overhead_ceiling_in_usd = ceiling,
+ "switch vetoed — would exceed session overhead cap, retaining anchor"
+ );
+ }
+ }
+ }
+ } else {
+ // Warm but no budget configured — follow the router freely.
+ switches += 1;
+ decision_label = metric_labels::SWITCH_DECISION_ALLOWED;
+ reason = metric_labels::SWITCH_REASON_FREE;
+ }
+ bs_metrics::record_session_switch_decision(decision_label, reason);
+ }
+ _ => {
+ // Cold (or no binding, or drifted): honor the candidate and (re)start a
+ // fresh warm episode. Switches reset — this is a new cache lifetime. On
+ // rebind we reset the running totals unless replenish_on_rebind is off, in
+ // which case the prior episode's baseline/spend carry over.
+ let (base, spend) = match (routing_budget, existing.as_ref()) {
+ (Some(cfg), Some(b)) if !cfg.replenish_on_rebind => {
+ (b.baseline_usd, b.switch_spend_usd)
+ }
+ _ => (0.0, 0.0),
+ };
+ baseline_usd = base;
+ switch_spend_usd = spend;
+ switches = 0;
+ // Fresh episode anchors on the model we're about to dispatch. When totals are
+ // carried across a rebind (replenish off), keep the prior default so the
+ // baseline lineage stays consistent.
+ default_model = match (routing_budget, existing.as_ref()) {
+ (Some(cfg), Some(b)) if !cfg.replenish_on_rebind && !b.default_model.is_empty() => {
+ b.default_model.clone()
+ }
+ _ => model.clone(),
+ };
+ }
+ }
+
+ // Context count persisted with the binding (refined later from real usage).
+ let cached_tokens = if facts.context_tokens > 0 {
+ facts.context_tokens
+ } else {
+ existing.as_ref().map(|b| b.cached_tokens).unwrap_or(0)
+ };
+ let gc_ttl = warmth_window(&capability_for_model(&model)) + GC_SLACK;
+
+ // Route history: a drifted prefix invalidates every model's cache, so start fresh;
+ // otherwise carry it forward. Record this turn's dispatched model (refined with the
+ // real token count on the response side). Stale entries decay via the warmth check.
+ let mut history = if drifted {
+ Vec::new()
+ } else {
+ existing
+ .as_ref()
+ .map(|b| b.history.clone())
+ .unwrap_or_default()
+ };
+ record_route_visit(&mut history, &model, now, cached_tokens);
+
+ // Observability: cache warmth + budget/switch state on the current span.
+ get_active_span(|span| {
+ span.set_attribute(KeyValue::new(tracing_plano::CACHE_WARM, effective_warm));
+ span.set_attribute(KeyValue::new(
+ tracing_plano::CACHE_IDLE_MS,
+ idle.as_millis() as i64,
+ ));
+ if routing_budget.is_some() {
+ // Consumed overhead as a percentage of the never-switch baseline — directly
+ // comparable to the configured max_overhead_pct. Zero before any baseline.
+ let overhead_pct = if baseline_usd > 0.0 {
+ 100.0 * switch_spend_usd / baseline_usd
+ } else {
+ 0.0
+ };
+ span.set_attribute(KeyValue::new(
+ tracing_plano::SESSION_OVERHEAD_PCT,
+ overhead_pct,
+ ));
+ span.set_attribute(KeyValue::new(
+ tracing_plano::SESSION_SWITCH_SPEND_IN_USD,
+ switch_spend_usd,
+ ));
+ span.set_attribute(KeyValue::new(
+ tracing_plano::SESSION_BASELINE_IN_USD,
+ baseline_usd,
+ ));
+ span.set_attribute(KeyValue::new(
+ tracing_plano::SESSION_SWITCHES,
+ switches as i64,
+ ));
+ }
+ // Cumulative actual conversation cost (through prior turns) — emitted for every
+ // session, independent of the routing budget.
+ span.set_attribute(KeyValue::new(
+ tracing_plano::SESSION_TOTAL_COST_IN_USD,
+ session_cost_usd,
+ ));
+ if let Some(cost) = cost_opt {
+ span.set_attribute(KeyValue::new(tracing_plano::SWITCH_COST_IN_USD, cost));
+ span.set_attribute(KeyValue::new(
+ tracing_plano::SWITCH_CANDIDATE_WARM_TOKENS,
+ candidate_warm_tokens as i64,
+ ));
+ if let Some(ceiling) = ceiling_opt {
+ span.set_attribute(KeyValue::new(
+ tracing_plano::SWITCH_OVERHEAD_CEILING_IN_USD,
+ ceiling,
+ ));
+ }
+ span.set_attribute(KeyValue::new(
+ tracing_plano::SWITCH_DECISION,
+ if model == facts.candidate_model {
+ metric_labels::SWITCH_DECISION_ALLOWED
+ } else {
+ metric_labels::SWITCH_DECISION_RETAINED
+ },
+ ));
+ }
+ if let Some(ref cf) = counterfactual {
+ span.set_attribute(KeyValue::new(
+ tracing_plano::SWITCH_COUNTERFACTUAL_ROUTE,
+ cf.clone(),
+ ));
+ }
+ });
+
+ orchestrator
+ .store_binding(
+ session_id,
+ facts.tenant_id,
+ SessionBinding {
+ anchor_model: model.clone(),
+ default_model: default_model.clone(),
+ route_name: route_name.clone(),
+ prefix_hash: facts.prefix_hash,
+ last_used: now,
+ cached_tokens,
+ baseline_usd,
+ switch_spend_usd,
+ switches,
+ session_cost_usd,
+ history: history.clone(),
+ },
+ Some(gc_ttl),
+ )
+ .await;
+
+ RouteDecision {
+ model,
+ route_name,
+ default_model,
+ warm: effective_warm,
+ baseline_usd,
+ switch_spend_usd,
+ session_cost_usd,
+ switches,
+ history,
+ cached_tokens,
+ gc_ttl,
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ fn cap_5m_1h() -> ProviderCacheCapability {
+ ProviderCacheCapability {
+ idle_ttl: Duration::from_secs(300),
+ hard_ttl: Duration::from_secs(3600),
+ extended_retention: false,
+ extended_ttl: Duration::from_secs(3600),
+ }
+ }
+
+ fn binding_used_ago(secs: u64) -> SessionBinding {
+ SessionBinding {
+ anchor_model: "anthropic/claude-sonnet-4-5".to_string(),
+ default_model: "anthropic/claude-sonnet-4-5".to_string(),
+ route_name: None,
+ prefix_hash: Some(1),
+ last_used: SystemTime::now() - Duration::from_secs(secs),
+ cached_tokens: 100_000,
+ baseline_usd: 1.0,
+ switch_spend_usd: 0.0,
+ switches: 0,
+ session_cost_usd: 0.0,
+ history: Vec::new(),
+ }
+ }
+
+ #[test]
+ fn warm_within_idle_window() {
+ let (warm, _) = warmth(&binding_used_ago(60), &cap_5m_1h(), SystemTime::now());
+ assert!(warm);
+ }
+
+ #[test]
+ fn cold_past_idle_window() {
+ let (warm, _) = warmth(&binding_used_ago(600), &cap_5m_1h(), SystemTime::now());
+ assert!(!warm);
+ }
+
+ #[test]
+ fn extended_retention_keeps_warm_past_idle() {
+ let cap = ProviderCacheCapability {
+ extended_retention: true,
+ ..cap_5m_1h()
+ };
+ // 10 minutes idle: cold under 5m, warm under the 1h extended window.
+ let (warm, _) = warmth(&binding_used_ago(600), &cap, SystemTime::now());
+ assert!(warm);
+ }
+
+ #[test]
+ fn capability_resolves_from_model_prefix() {
+ // Known provider prefix resolves; unknown falls back to the default.
+ let anthropic = capability_for_model("anthropic/claude-sonnet-4-5");
+ assert_eq!(anthropic, ProviderCacheCapability::default());
+ let unknown = capability_for_model("madeup/model-x");
+ assert_eq!(unknown, ProviderCacheCapability::default());
+ }
+
+ // ---- route() budget behavior ----
+
+ use crate::router::model_metrics::{ModelMetricsService, ModelRates};
+ use crate::session_cache::memory::MemorySessionCache;
+ use std::collections::HashMap;
+ use std::sync::Arc;
+
+ // Anchor `expensive` cached rate 0.3, candidate `pricey` input 5.0, candidate `cheap`
+ // input 0.1. With a 100k-token context the paid switch to `pricey` costs
+ // 0.1M * (5.0 - 0.3) = $0.47; the `cheap` switch is 0.1M * (0.1 - 0.3) = -$0.02 (free).
+ // Each warm turn grows the never-switch baseline by 0.1M * 0.3 = $0.03.
+ fn orch_with_rates() -> OrchestratorService {
+ let mut rates = HashMap::new();
+ rates.insert(
+ "anthropic/expensive".to_string(),
+ ModelRates {
+ input_per_million: 3.0,
+ output_per_million: 15.0,
+ cache_read_per_million: Some(0.3),
+ },
+ );
+ rates.insert(
+ "openai/pricey".to_string(),
+ ModelRates {
+ input_per_million: 5.0,
+ output_per_million: 15.0,
+ cache_read_per_million: Some(0.5),
+ },
+ );
+ rates.insert(
+ "google/cheap".to_string(),
+ ModelRates {
+ input_per_million: 0.1,
+ output_per_million: 0.4,
+ cache_read_per_million: Some(0.01),
+ },
+ );
+ let metrics = Arc::new(ModelMetricsService::from_rates_for_test(rates));
+ let cache = Arc::new(MemorySessionCache::new(100));
+ OrchestratorService::with_routing(
+ "http://localhost/v1/chat/completions".to_string(),
+ "m".to_string(),
+ "p".to_string(),
+ None,
+ Some(metrics),
+ Some(600),
+ cache,
+ None,
+ 8192,
+ )
+ }
+
+ fn routing_budget(pct: f64) -> EffectiveRoutingBudget {
+ EffectiveRoutingBudget {
+ max_overhead_pct: pct,
+ replenish_on_rebind: true,
+ cache_read_discount: 0.1,
+ record_counterfactual: false,
+ }
+ }
+
+ /// Seed a warm binding on the `expensive` anchor with a pre-accumulated never-switch
+ /// baseline (`baseline_usd`) and switch spend (`switch_spend_usd`), simulating a
+ /// session that has already run for some turns.
+ async fn seed_warm_binding(
+ orch: &OrchestratorService,
+ baseline_usd: f64,
+ switch_spend_usd: f64,
+ idle_secs: u64,
+ ) {
+ orch.store_binding(
+ "s1",
+ None,
+ SessionBinding {
+ anchor_model: "anthropic/expensive".to_string(),
+ default_model: "anthropic/expensive".to_string(),
+ route_name: None,
+ prefix_hash: Some(1),
+ last_used: SystemTime::now() - Duration::from_secs(idle_secs),
+ cached_tokens: 100_000,
+ baseline_usd,
+ switch_spend_usd,
+ switches: 0,
+ session_cost_usd: 0.0,
+ history: Vec::new(),
+ },
+ Some(Duration::from_secs(3600)),
+ )
+ .await;
+ }
+
+ fn facts_for<'a>(candidate: &'a str) -> RouteFacts<'a> {
+ RouteFacts {
+ session_id: Some("s1"),
+ tenant_id: None,
+ prefix_hash: Some(1),
+ context_tokens: 0,
+ candidate_model: candidate,
+ candidate_route: None,
+ }
+ }
+
+ #[tokio::test]
+ async fn paid_switch_within_cap_accrues_spend() {
+ let orch = orch_with_rates();
+ // Baseline $2.00 already accrued; this turn adds $0.03 -> $2.03. At 25% the
+ // ceiling is $0.5075, which covers the $0.47 switch to `pricey`.
+ seed_warm_binding(&orch, 2.0, 0.0, 30).await;
+ let st = routing_budget(25.0);
+ let d = route(&orch, Some(&st), facts_for("openai/pricey")).await;
+
+ assert_eq!(d.model, "openai/pricey");
+ assert!(d.warm);
+ assert_eq!(d.switches, 1);
+ assert!(
+ (d.switch_spend_usd - 0.47).abs() < 1e-6,
+ "spend {} != 0.47",
+ d.switch_spend_usd
+ );
+ assert!(
+ (d.baseline_usd - 2.03).abs() < 1e-6,
+ "baseline {} != 2.03",
+ d.baseline_usd
+ );
+ }
+
+ #[tokio::test]
+ async fn paid_switch_over_cap_retains_anchor() {
+ let orch = orch_with_rates();
+ // Baseline $1.00 (+$0.03 this turn). At 25% the ceiling is ~$0.2575 < $0.47.
+ seed_warm_binding(&orch, 1.0, 0.0, 30).await;
+ let st = routing_budget(25.0);
+ let d = route(&orch, Some(&st), facts_for("openai/pricey")).await;
+
+ assert_eq!(d.model, "anthropic/expensive");
+ assert!(d.warm);
+ assert_eq!(d.switches, 0);
+ // Vetoed switch spends nothing.
+ assert!((d.switch_spend_usd - 0.0).abs() < 1e-6);
+ }
+
+ #[tokio::test]
+ async fn cheaper_switch_is_free() {
+ let orch = orch_with_rates();
+ seed_warm_binding(&orch, 1.0, 0.10, 30).await;
+ let st = routing_budget(25.0);
+ let d = route(&orch, Some(&st), facts_for("google/cheap")).await;
+
+ assert_eq!(d.model, "google/cheap");
+ assert!(d.warm);
+ assert_eq!(d.switches, 1);
+ // Free switches never touch the running spend — it stays at 0.10.
+ assert!(
+ (d.switch_spend_usd - 0.10).abs() < 1e-6,
+ "spend {} != 0.10",
+ d.switch_spend_usd
+ );
+ }
+
+ #[tokio::test]
+ async fn cold_session_resets_and_follows_router() {
+ let orch = orch_with_rates();
+ // 10 minutes idle: past Anthropic's 5m idle window -> cold. Prior episode spent
+ // its overhead; the fresh episode resets baseline and spend to zero.
+ seed_warm_binding(&orch, 5.0, 3.0, 600).await;
+ let st = routing_budget(25.0);
+ let d = route(&orch, Some(&st), facts_for("openai/pricey")).await;
+
+ assert_eq!(d.model, "openai/pricey");
+ assert!(!d.warm);
+ assert_eq!(d.switches, 0);
+ assert!((d.baseline_usd - 0.0).abs() < 1e-6, "baseline reset");
+ assert!((d.switch_spend_usd - 0.0).abs() < 1e-6, "spend reset");
+ }
+
+ #[tokio::test]
+ async fn no_session_honors_candidate() {
+ let orch = orch_with_rates();
+ let st = routing_budget(1.0);
+ let facts = RouteFacts {
+ session_id: None,
+ tenant_id: None,
+ prefix_hash: Some(1),
+ context_tokens: 0,
+ candidate_model: "openai/pricey",
+ candidate_route: None,
+ };
+ let d = route(&orch, Some(&st), facts).await;
+ assert_eq!(d.model, "openai/pricey");
+ assert!(!d.warm);
+ }
+
+ #[tokio::test]
+ async fn baseline_priced_on_default_not_anchor() {
+ // A session that started on `expensive` (default) but has since switched to
+ // `cheap` (current anchor). The never-switch baseline must keep growing at the
+ // *default* model's cached rate (0.3/M), not the cheap anchor's (0.01/M).
+ let orch = orch_with_rates();
+ orch.store_binding(
+ "s1",
+ None,
+ SessionBinding {
+ anchor_model: "google/cheap".to_string(),
+ default_model: "anthropic/expensive".to_string(),
+ route_name: None,
+ prefix_hash: Some(1),
+ last_used: SystemTime::now(),
+ cached_tokens: 100_000,
+ baseline_usd: 0.0,
+ switch_spend_usd: 0.0,
+ switches: 1,
+ session_cost_usd: 0.0,
+ history: Vec::new(),
+ },
+ Some(Duration::from_secs(3600)),
+ )
+ .await;
+ let st = routing_budget(25.0);
+ // Router agrees with the current anchor -> same-anchor, no switch.
+ let d = route(&orch, Some(&st), facts_for("google/cheap")).await;
+
+ assert_eq!(d.model, "google/cheap");
+ assert!(d.warm);
+ assert_eq!(d.switches, 1);
+ // 100_000 tokens x $0.30/M (default `expensive`) = $0.03 — not $0.001 (cheap).
+ assert!(
+ (d.baseline_usd - 0.03).abs() < 1e-6,
+ "baseline {} != 0.03 (should price the default model, not the anchor)",
+ d.baseline_usd
+ );
+ assert_eq!(d.default_model, "anthropic/expensive");
+ }
+
+ #[tokio::test]
+ async fn warm_return_is_affordable() {
+ // Warm on `expensive`, but `pricey` was used recently and still holds the whole
+ // 100k context. Baseline $0.40 (+$0.03 this turn = $0.43); at 25% the ceiling is
+ // ~$0.1075. Returning to `pricey` re-reads the 100k at its *cached* rate (0.5),
+ // so the switch costs 100k x (0.5 - 0.3)/1M = $0.02 — under the ceiling, allowed.
+ // A cold switch would re-read at 5.0 -> $0.47 and be vetoed.
+ let orch = orch_with_rates();
+ orch.store_binding(
+ "s1",
+ None,
+ SessionBinding {
+ anchor_model: "anthropic/expensive".to_string(),
+ default_model: "anthropic/expensive".to_string(),
+ route_name: None,
+ prefix_hash: Some(1),
+ last_used: SystemTime::now() - Duration::from_secs(30),
+ cached_tokens: 100_000,
+ baseline_usd: 0.40,
+ switch_spend_usd: 0.0,
+ switches: 1,
+ session_cost_usd: 0.0,
+ history: vec![RouteVisit {
+ model: "openai/pricey".to_string(),
+ last_used: SystemTime::now() - Duration::from_secs(30),
+ cached_tokens: 100_000,
+ }],
+ },
+ Some(Duration::from_secs(3600)),
+ )
+ .await;
+ let st = routing_budget(25.0);
+ let d = route(&orch, Some(&st), facts_for("openai/pricey")).await;
+
+ assert_eq!(d.model, "openai/pricey", "warm return should be allowed");
+ assert_eq!(d.switches, 2);
+ assert!(
+ (d.switch_spend_usd - 0.02).abs() < 1e-6,
+ "spend {} != 0.02 (warm return should charge only the cached-rate delta)",
+ d.switch_spend_usd
+ );
+ // The candidate's own history entry is refreshed to the current model.
+ assert!(d.history.iter().any(|v| v.model == "openai/pricey"));
+ }
+
+ /// End-to-end cost validation over a full session lifetime: drive `route()` turn by
+ /// turn through the real session cache and pricing math, and check the feature's core
+ /// promise at every turn — cumulative switch spend never exceeds `max_overhead_pct`%
+ /// of the never-switch baseline.
+ ///
+ /// With a 100k context, each warm turn grows the baseline by 100k x $0.30/M = $0.03
+ /// (the `expensive` default's cached rate) and a switch to `pricey` costs
+ /// 100k x (5.0 - 0.3)/M = $0.47. At a 25% cap the ceiling is 0.25 x $0.03 x k after
+ /// k warm turns, so the switch must be vetoed through warm turn 62
+ /// (ceiling $0.465 < $0.47) and allowed exactly at warm turn 63 (ceiling $0.4725).
+ #[tokio::test]
+ async fn multi_turn_session_cost_stays_within_overhead_cap() {
+ let orch = orch_with_rates();
+ let st = routing_budget(25.0);
+ let cap_fraction = 25.0 / 100.0;
+ let facts_with_context = |candidate: &'static str| RouteFacts {
+ session_id: Some("s1"),
+ tenant_id: None,
+ prefix_hash: Some(1),
+ context_tokens: 100_000,
+ candidate_model: candidate,
+ candidate_route: None,
+ };
+
+ // Turn 1 — cold start: no binding yet, the candidate is honored and becomes both
+ // the anchor and the session's never-switch default.
+ let d = route(&orch, Some(&st), facts_with_context("anthropic/expensive")).await;
+ assert_eq!(d.model, "anthropic/expensive");
+ assert!(!d.warm);
+ assert_eq!(d.default_model, "anthropic/expensive");
+
+ // Warm turns: the router keeps proposing the pricier model every turn. The gate
+ // must veto until the accrued baseline makes the switch affordable.
+ let mut first_switch_warm_turn: Option = None;
+ let mut warm_turns = 0u32;
+ let mut last = None;
+ for turn in 1..=70u32 {
+ let d = route(&orch, Some(&st), facts_with_context("openai/pricey")).await;
+ assert!(d.warm, "turn {turn} should be warm (used moments ago)");
+ warm_turns = turn;
+
+ // The invariant under validation: spend never exceeds the cap.
+ assert!(
+ d.switch_spend_usd <= cap_fraction * d.baseline_usd + 1e-9,
+ "turn {turn}: spend {} exceeds {}% of baseline {}",
+ d.switch_spend_usd,
+ st.max_overhead_pct,
+ d.baseline_usd
+ );
+ // Baseline must track the independently computed never-switch cost:
+ // $0.03 per warm turn on the default model's cached rate.
+ let expected_baseline = 0.03 * turn as f64;
+ assert!(
+ (d.baseline_usd - expected_baseline).abs() < 1e-6,
+ "turn {turn}: baseline {} != expected never-switch cost {expected_baseline}",
+ d.baseline_usd
+ );
+
+ if d.model == "openai/pricey" {
+ first_switch_warm_turn = Some(turn);
+ assert_eq!(d.switches, 1);
+ assert!(
+ (d.switch_spend_usd - 0.47).abs() < 1e-6,
+ "allowed switch should charge the full cache-loss delta"
+ );
+ last = Some(d);
+ break;
+ }
+ // Vetoed turns retain the anchor and spend nothing.
+ assert_eq!(d.model, "anthropic/expensive");
+ assert_eq!(d.switches, 0);
+ assert!((d.switch_spend_usd - 0.0).abs() < 1e-9);
+ }
+
+ // ceiling(k) = 0.25 x 0.03k first covers $0.47 at k = 63.
+ assert_eq!(
+ first_switch_warm_turn,
+ Some(63),
+ "switch should flip from veto to allow exactly when the ceiling covers its cost"
+ );
+ let after_switch = last.unwrap();
+
+ // Same-anchor turn after the switch: no further spend accrues.
+ let d = route(&orch, Some(&st), facts_with_context("openai/pricey")).await;
+ warm_turns += 1;
+ assert_eq!(d.model, "openai/pricey");
+ assert_eq!(d.switches, 1);
+ assert!((d.switch_spend_usd - after_switch.switch_spend_usd).abs() < 1e-9);
+
+ // A -> B -> A return: `expensive` was dispatched last turn-but-one, so its cache
+ // is still warm — the return re-reads at cached rates and is free (never vetoed,
+ // never accrues spend).
+ let d = route(&orch, Some(&st), facts_with_context("anthropic/expensive")).await;
+ warm_turns += 1;
+ assert_eq!(
+ d.model, "anthropic/expensive",
+ "warm return must be allowed"
+ );
+ assert_eq!(d.switches, 2);
+ assert!(
+ (d.switch_spend_usd - after_switch.switch_spend_usd).abs() < 1e-9,
+ "free return must not change the running spend"
+ );
+
+ // Final end-to-end check of the promise: total switch overhead across the whole
+ // session is within max_overhead_pct% of the independently computed
+ // never-switch baseline.
+ let never_switch_cost = 0.03 * warm_turns as f64;
+ assert!(
+ (d.baseline_usd - never_switch_cost).abs() < 1e-6,
+ "final baseline {} != independently computed never-switch cost {never_switch_cost}",
+ d.baseline_usd
+ );
+ assert!(
+ d.switch_spend_usd <= cap_fraction * d.baseline_usd + 1e-9,
+ "session overhead {} exceeds the promised {}% of {}",
+ d.switch_spend_usd,
+ st.max_overhead_pct,
+ d.baseline_usd
+ );
+ }
+}
diff --git a/crates/brightstaff/src/handlers/routing_service.rs b/crates/brightstaff/src/handlers/routing_service.rs
index b93b1422..0fc759f1 100644
--- a/crates/brightstaff/src/handlers/routing_service.rs
+++ b/crates/brightstaff/src/handlers/routing_service.rs
@@ -1,9 +1,11 @@
use bytes::Bytes;
-use common::configuration::{SpanAttributes, TopLevelRoutingPreference};
+use common::configuration::{
+ EffectivePromptCaching, EffectiveRoutingBudget, SpanAttributes, TopLevelRoutingPreference,
+};
use common::consts::{MODEL_AFFINITY_HEADER, REQUEST_ID_HEADER};
use common::errors::BrightStaffError;
use hermesllm::clients::SupportedAPIsFromClient;
-use hermesllm::ProviderRequestType;
+use hermesllm::{ProviderRequest, ProviderRequestType};
use http_body_util::combinators::BoxBody;
use http_body_util::{BodyExt, Full};
use hyper::{Request, Response, StatusCode};
@@ -12,6 +14,7 @@ use tracing::{debug, info, info_span, warn, Instrument};
use super::extract_or_generate_traceparent;
use crate::handlers::llm::model_selection::router_chat_get_upstream_model;
+use crate::handlers::llm::session_router;
use crate::metrics as bs_metrics;
use crate::metrics::labels as metric_labels;
use crate::router::orchestrator::OrchestratorService;
@@ -60,11 +63,14 @@ struct RoutingDecisionResponse {
pinned: bool,
}
+#[allow(clippy::too_many_arguments)]
pub async fn routing_decision(
request: Request,
orchestrator_service: Arc,
request_path: String,
span_attributes: &Option,
+ prompt_caching: EffectivePromptCaching,
+ routing_budget: Option,
) -> Result>, hyper::Error> {
let request_headers = request.headers().clone();
let request_id: String = request_headers
@@ -73,7 +79,7 @@ pub async fn routing_decision(
.map(|s| s.to_string())
.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
- let session_id: Option = request_headers
+ let explicit_session_id: Option = request_headers
.get(MODEL_AFFINITY_HEADER)
.and_then(|h| h.to_str().ok())
.map(|s| s.to_string());
@@ -101,8 +107,10 @@ pub async fn routing_decision(
request_path,
request_headers,
custom_attrs,
- session_id,
+ explicit_session_id,
tenant_id,
+ prompt_caching,
+ routing_budget,
)
.instrument(request_span)
.await
@@ -116,8 +124,10 @@ async fn routing_decision_inner(
request_path: String,
request_headers: hyper::HeaderMap,
custom_attrs: std::collections::HashMap,
- session_id: Option,
+ explicit_session_id: Option,
tenant_id: Option,
+ prompt_caching: EffectivePromptCaching,
+ routing_budget: Option,
) -> Result>, hyper::Error> {
set_service_name(operation_component::ROUTING);
opentelemetry::trace::get_active_span(|span| {
@@ -135,37 +145,9 @@ async fn routing_decision_inner(
.unwrap_or("unknown")
.to_string();
- if let Some(ref sid) = session_id {
- if let Some(cached) = orchestrator_service
- .get_cached_route(sid, tenant_id.as_deref())
- .await
- {
- info!(
- session_id = %sid,
- model = %cached.model_name,
- route = ?cached.route_name,
- "returning pinned routing decision from cache"
- );
- let response = RoutingDecisionResponse {
- models: vec![cached.model_name],
- route: cached.route_name,
- trace_id,
- session_id: Some(sid.clone()),
- pinned: true,
- };
- let json = serde_json::to_string(&response).unwrap();
- let body = Full::new(Bytes::from(json))
- .map_err(|never| match never {})
- .boxed();
- return Ok(Response::builder()
- .status(StatusCode::OK)
- .header("Content-Type", "application/json")
- .body(body)
- .unwrap());
- }
- }
-
- // Parse request body
+ // Parse the request body up front so a pin can be validated against prefix drift
+ // and a fresh pin can be stored with its prefix hash. This endpoint shares the
+ // session cache with the LLM handler, so both must key drift the same way.
let raw_bytes = request.collect().await?.to_bytes();
debug!(
@@ -202,6 +184,39 @@ async fn routing_decision_inner(
}
};
+ // `X-Plano-Cache: off` opts this request out of implicit affinity (same sentinel
+ // the LLM handler honors), so callers can bypass stickiness per request.
+ let cache_off_for_request = request_headers
+ .get(common::consts::PLANO_CACHE_HEADER)
+ .and_then(|h| h.to_str().ok())
+ .is_some_and(|v| v.eq_ignore_ascii_case("off"));
+
+ let request_messages = client_request.get_messages();
+ let tool_names = client_request.get_tool_names();
+
+ // Session key + prefix hash resolved identically to the LLM handler so pins
+ // interoperate across the full-proxy and decision paths.
+ // Derive the implicit session key when either prompt-caching affinity or the routing
+ // budget is active, so the budget works the same way with caching off.
+ let implicit_affinity_enabled = prompt_caching.session_affinity || routing_budget.is_some();
+ let session_router::SessionResolution {
+ request_prefix_hash,
+ session_id,
+ } = session_router::resolve_session(
+ explicit_session_id,
+ &request_messages,
+ tool_names.as_deref(),
+ tenant_id.as_deref(),
+ implicit_affinity_enabled,
+ cache_off_for_request,
+ );
+
+ let context_tokens: u64 = if session_id.is_some() && routing_budget.is_some() {
+ session_router::actual_context_tokens(&request_messages, client_request.model())
+ } else {
+ 0
+ };
+
let routing_result = router_chat_get_upstream_model(
Arc::clone(&orchestrator_service),
client_request,
@@ -213,23 +228,37 @@ async fn routing_decision_inner(
match routing_result {
Ok(result) => {
- if let Some(ref sid) = session_id {
- orchestrator_service
- .cache_route(
- sid.clone(),
- tenant_id.as_deref(),
- result.model_name.clone(),
- result.route_name.clone(),
- )
- .await;
+ let candidate_model = result.model_name.clone();
+ let decision = session_router::route(
+ &orchestrator_service,
+ routing_budget.as_ref(),
+ session_router::RouteFacts {
+ session_id: session_id.as_deref(),
+ tenant_id: tenant_id.as_deref(),
+ prefix_hash: request_prefix_hash,
+ context_tokens,
+ candidate_model: &candidate_model,
+ candidate_route: result.route_name.as_deref(),
+ },
+ )
+ .await;
+
+ // Front the ranked fallback list with the decided model (the anchor, when a
+ // switch was vetoed), so 429/5xx fallbacks still work.
+ let mut models = result.models;
+ if models.first() != Some(&decision.model) {
+ models.retain(|m| m != &decision.model);
+ models.insert(0, decision.model.clone());
}
let response = RoutingDecisionResponse {
- models: result.models,
- route: result.route_name,
+ models,
+ route: decision.route_name,
trace_id,
session_id,
- pinned: false,
+ // `pinned` signals a warm, stuck session — safe for callers to treat as
+ // "keep this provider's cache warm".
+ pinned: decision.warm,
};
// Distinguish "decision served" (a concrete model picked) from
@@ -247,6 +276,7 @@ async fn routing_decision_inner(
primary_model = %response.models.first().map(|s| s.as_str()).unwrap_or("none"),
total_models = response.models.len(),
route = ?response.route,
+ pinned = response.pinned,
"routing decision completed"
);
diff --git a/crates/brightstaff/src/lib.rs b/crates/brightstaff/src/lib.rs
index 66c6eadf..f0b6b13d 100644
--- a/crates/brightstaff/src/lib.rs
+++ b/crates/brightstaff/src/lib.rs
@@ -1,3 +1,4 @@
+pub mod affinity;
pub mod app_state;
pub mod handlers;
pub mod metrics;
diff --git a/crates/brightstaff/src/main.rs b/crates/brightstaff/src/main.rs
index 90ed84c3..38af28b5 100644
--- a/crates/brightstaff/src/main.rs
+++ b/crates/brightstaff/src/main.rs
@@ -216,7 +216,14 @@ async fn init_app_state(
if latency_count > 1 {
return Err("model_metrics_sources: only one latency metrics source is allowed".into());
}
- let svc = ModelMetricsService::new(sources, reqwest::Client::new()).await;
+ // The initial pricing fetch is awaited before listeners come up, so bound it —
+ // an unresponsive feed (models.dev / DO catalog) must not hang startup. The
+ // same client backs the periodic refresh loop.
+ let metrics_client = reqwest::Client::builder()
+ .connect_timeout(std::time::Duration::from_secs(5))
+ .timeout(std::time::Duration::from_secs(30))
+ .build()?;
+ let svc = ModelMetricsService::new(sources, metrics_client).await;
Some(Arc::new(svc))
} else {
None
@@ -327,8 +334,71 @@ async fn init_app_state(
.as_ref()
.and_then(|tracing| tracing.span_attributes.clone());
+ // Resolve the distinct_id header from the first PostHog exporter that
+ // declares one, so the LLM handler can stamp `plano.distinct_id` on spans.
+ let distinct_id_header = config
+ .tracing
+ .as_ref()
+ .and_then(|tracing| tracing.exporters.as_ref())
+ .and_then(|exporters| {
+ exporters.iter().find_map(|exporter| match exporter {
+ common::configuration::Exporter::Posthog(posthog) => {
+ posthog.distinct_id_header.clone()
+ }
+ })
+ });
+
let signals_enabled = !overrides.disable_signals.unwrap_or(false);
+ let prompt_caching =
+ common::configuration::EffectivePromptCaching::from_config(config.prompt_caching.as_ref())?;
+
+ // Routing budget lives under `routing` and is independent of prompt caching.
+ let routing_budget = common::configuration::EffectiveRoutingBudget::from_config(
+ config
+ .routing
+ .as_ref()
+ .and_then(|r| r.routing_budget.as_ref()),
+ )?;
+
+ // The routing-budget cost gate needs per-model pricing to compute switch cost.
+ if routing_budget.is_some() {
+ use common::configuration::MetricsSource;
+ let has_cost_source = config
+ .model_metrics_sources
+ .as_deref()
+ .unwrap_or_default()
+ .iter()
+ .any(|s| matches!(s, MetricsSource::Cost(_)));
+ if !has_cost_source {
+ return Err(
+ "routing.routing_budget is configured but no cost metrics source is \
+ configured — add a cost source (e.g. models.dev) to model_metrics_sources so \
+ per-model input/cached rates are available for the switch-cost calculation"
+ .into(),
+ );
+ }
+ }
+
+ // Session bindings (anchor model, budget spend) are keyed by tenant + session.
+ // Without a tenant header, sessions from different customers share one keyspace,
+ // so identical prompts from different tenants would collide on the same binding.
+ let session_state_in_use = prompt_caching.session_affinity || routing_budget.is_some();
+ if session_state_in_use
+ && config
+ .routing
+ .as_ref()
+ .and_then(|r| r.session_cache.as_ref())
+ .and_then(|c| c.tenant_header.as_ref())
+ .is_none()
+ {
+ warn!(
+ "session affinity / routing budget is enabled but routing.session_cache.tenant_header \
+ is not set — all requests share one session keyspace; set tenant_header to isolate \
+ sessions per customer in multi-tenant deployments"
+ );
+ }
+
Ok(AppState {
orchestrator_service,
model_aliases: config.model_aliases.clone(),
@@ -338,9 +408,12 @@ async fn init_app_state(
state_storage,
llm_provider_url,
span_attributes,
+ distinct_id_header,
http_client: reqwest::Client::new(),
filter_pipeline,
signals_enabled,
+ prompt_caching,
+ routing_budget,
})
}
@@ -512,6 +585,8 @@ async fn dispatch(
Arc::clone(&state.orchestrator_service),
stripped,
&state.span_attributes,
+ state.prompt_caching,
+ state.routing_budget,
)
.with_context(parent_cx)
.await;
diff --git a/crates/brightstaff/src/metrics/labels.rs b/crates/brightstaff/src/metrics/labels.rs
index 4eaf3e59..a3865cef 100644
--- a/crates/brightstaff/src/metrics/labels.rs
+++ b/crates/brightstaff/src/metrics/labels.rs
@@ -18,6 +18,10 @@ pub const ROUTE_LLM: &str = "llm";
// Token kind for brightstaff_llm_tokens_total.
pub const TOKEN_KIND_PROMPT: &str = "prompt";
pub const TOKEN_KIND_COMPLETION: &str = "completion";
+/// Input tokens served from the provider's prompt cache (billed at the cached rate).
+pub const TOKEN_KIND_CACHE_READ: &str = "cache_read";
+/// Input tokens written into the provider's prompt cache (cache-creation surcharge).
+pub const TOKEN_KIND_CACHE_WRITE: &str = "cache_write";
// LLM error_class values (match docstring in metrics/mod.rs).
pub const LLM_ERR_NONE: &str = "none";
@@ -36,3 +40,31 @@ pub const ROUTING_SVC_POLICY_ERROR: &str = "policy_error";
pub const SESSION_CACHE_HIT: &str = "hit";
pub const SESSION_CACHE_MISS: &str = "miss";
pub const SESSION_CACHE_STORE: &str = "store";
+
+// Prompt cache outcome values (brightstaff_prompt_cache_requests_total).
+pub const PROMPT_CACHE_HIT: &str = "hit";
+pub const PROMPT_CACHE_MISS: &str = "miss";
+
+// Session binding lifecycle events (brightstaff_session_binding_events_total).
+/// An existing binding was refreshed from observed usage (TTL extended, token counts
+/// and running cost updated) after a turn completed.
+pub const BINDING_EVENT_REFRESH: &str = "refresh";
+
+// Session-stickiness decisions (brightstaff_session_switch_decisions_total).
+// `decision` label — the coarse outcome:
+/// The proposed switch was honored (free, within the overhead cap, or unpriced fail-open).
+pub const SWITCH_DECISION_ALLOWED: &str = "allowed";
+/// The switch would have exceeded the session's overhead cap — the warm anchor was retained.
+pub const SWITCH_DECISION_RETAINED: &str = "retained";
+
+// `reason` label — why the decision was made:
+/// The router agreed with the warm anchor; no switch was needed.
+pub const SWITCH_REASON_SAME_ANCHOR: &str = "same_anchor";
+/// The candidate was outright cheaper (negative cost) — a free switch.
+pub const SWITCH_REASON_FREE: &str = "free";
+/// A paid switch that kept cumulative spend within the session's overhead cap.
+pub const SWITCH_REASON_WITHIN_CAP: &str = "within_cap";
+/// A paid switch that would have pushed cumulative spend over the overhead cap — retained.
+pub const SWITCH_REASON_OVER_CAP: &str = "over_cap";
+/// Pricing was missing for one side, so the switch was allowed without a cost gate.
+pub const SWITCH_REASON_NO_PRICING: &str = "no_pricing";
diff --git a/crates/brightstaff/src/metrics/mod.rs b/crates/brightstaff/src/metrics/mod.rs
index 34679cca..1f5ce909 100644
--- a/crates/brightstaff/src/metrics/mod.rs
+++ b/crates/brightstaff/src/metrics/mod.rs
@@ -18,6 +18,9 @@
//! `brightstaff_router_decision_duration_seconds`,
//! `brightstaff_routing_service_requests_total`,
//! `brightstaff_session_cache_events_total`.
+//! - Prompt caching: `brightstaff_prompt_cache_requests_total`,
+//! `brightstaff_session_binding_events_total` (cache read/write tokens are emitted as
+//! `kind=cache_read|cache_write` on `brightstaff_llm_tokens_total`).
//! - Process: via `metrics-process`.
//! - Build: `brightstaff_build_info`.
@@ -172,6 +175,16 @@ fn describe_all() {
"brightstaff_session_cache_events_total",
"Session affinity cache lookups and stores, by outcome."
);
+ describe_counter!(
+ "brightstaff_prompt_cache_requests_total",
+ "LLM responses with usage data, by provider, model and prompt-cache outcome \
+ (hit = cache read/creation tokens reported, miss = none). The miss rate is \
+ the silent cache-miss baseline."
+ );
+ describe_counter!(
+ "brightstaff_session_binding_events_total",
+ "Session binding lifecycle events, by event (currently: refresh)."
+ );
describe_gauge!(
"brightstaff_build_info",
@@ -375,3 +388,37 @@ pub fn record_session_cache_event(outcome: &'static str) {
)
.increment(1);
}
+
+/// Record whether a completed LLM response reported prompt-cache activity.
+/// `hit` means the provider billed cache-read or cache-creation tokens.
+pub fn record_prompt_cache_outcome(provider: &str, model: &str, outcome: &'static str) {
+ counter!(
+ "brightstaff_prompt_cache_requests_total",
+ "provider" => provider.to_string(),
+ "model" => model.to_string(),
+ "outcome" => outcome,
+ )
+ .increment(1);
+}
+
+/// Record a session binding lifecycle event (see `metrics::labels::BINDING_EVENT_*`).
+pub fn record_session_binding_event(event: &'static str) {
+ counter!(
+ "brightstaff_session_binding_events_total",
+ "event" => event,
+ )
+ .increment(1);
+}
+
+/// Record a session-stickiness decision on a proposed model switch. `decision` is the
+/// coarse outcome (`allowed`/`retained`, see `metrics::labels::SWITCH_DECISION_*`) and
+/// `reason` explains why (`same_anchor`/`free`/`within_cap`/`over_cap`/`no_pricing`,
+/// see `metrics::labels::SWITCH_REASON_*`).
+pub fn record_session_switch_decision(decision: &'static str, reason: &'static str) {
+ counter!(
+ "brightstaff_session_switch_decisions_total",
+ "decision" => decision,
+ "reason" => reason,
+ )
+ .increment(1);
+}
diff --git a/crates/brightstaff/src/router/model_metrics.rs b/crates/brightstaff/src/router/model_metrics.rs
index 1adb408d..2990a046 100644
--- a/crates/brightstaff/src/router/model_metrics.rs
+++ b/crates/brightstaff/src/router/model_metrics.rs
@@ -9,41 +9,124 @@ use tokio::sync::RwLock;
use tracing::{debug, info, warn};
const DO_PRICING_URL: &str = "https://api.digitalocean.com/v2/gen-ai/models/catalog";
+const MODELS_DEV_URL: &str = "https://models.dev/api.json";
+
+/// DigitalOcean publishes prices per token; scale to the per-million convention
+/// that `ModelRates` and the absolute-USD consumers (switch-cost gate) expect.
+const TOKENS_PER_MILLION: f64 = 1_000_000.0;
+
+/// Structured per-million-token USD rates for one model, as published by the cost
+/// feed. Kept alongside the blended ranking metric so cost-sensitive features (e.g.
+/// the session-stickiness switch-cost gate) can reason about input vs cached-input
+/// pricing without touching `rank_models`.
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub struct ModelRates {
+ pub input_per_million: f64,
+ pub output_per_million: f64,
+ /// Cached (prompt-cache read) input rate. Present when the feed publishes it
+ /// (models.dev `cost.cache_read`, DO catalog `cache_read_input_price_per_million`
+ /// / `input_cache_read`); absent for models the feed prices without a cache SKU.
+ pub cache_read_per_million: Option,
+}
+
+impl ModelRates {
+ /// Blended input+output metric used for `prefer: cheapest` ranking.
+ fn blended(&self) -> f64 {
+ self.input_per_million + self.output_per_million
+ }
+
+ /// Cached input rate, falling back to `input * cache_read_discount` when the
+ /// feed doesn't publish a cached rate.
+ pub fn cached_input_rate(&self, cache_read_discount: f64) -> f64 {
+ self.cache_read_per_million
+ .unwrap_or(self.input_per_million * cache_read_discount)
+ }
+
+ /// Actual USD cost of one request from observed token usage, split into
+ /// `(input_cost, output_cost)`. Cache *creation* tokens are priced at the plain
+ /// input rate (no write premium). Cache *read* tokens use the cached rate.
+ ///
+ /// `prompt_includes_cached` reflects the provider convention: OpenAI-shape usage
+ /// folds cached tokens into `prompt_tokens` (uncached = prompt - cached), while
+ /// Anthropic-shape reports uncached input separately (uncached = prompt_tokens).
+ pub fn request_cost_usd(
+ &self,
+ prompt_tokens: u64,
+ cached_input_tokens: u64,
+ cache_creation_tokens: u64,
+ completion_tokens: u64,
+ prompt_includes_cached: bool,
+ cache_read_discount: f64,
+ ) -> (f64, f64) {
+ let read_rate = self.cached_input_rate(cache_read_discount);
+ let uncached = if prompt_includes_cached {
+ prompt_tokens.saturating_sub(cached_input_tokens)
+ } else {
+ prompt_tokens
+ };
+ let input = (uncached as f64 * self.input_per_million
+ + cached_input_tokens as f64 * read_rate
+ + cache_creation_tokens as f64 * self.input_per_million)
+ / TOKENS_PER_MILLION;
+ let output = completion_tokens as f64 * self.output_per_million / TOKENS_PER_MILLION;
+ (input, output)
+ }
+}
+
+/// Derive the blended ranking map from the structured rates map.
+fn blended_costs(rates: &HashMap) -> HashMap {
+ rates
+ .iter()
+ .map(|(k, r)| (k.clone(), r.blended()))
+ .collect()
+}
pub struct ModelMetricsService {
cost: Arc>>,
+ rates: Arc>>,
latency: Arc>>,
}
impl ModelMetricsService {
pub async fn new(sources: &[MetricsSource], client: reqwest::Client) -> Self {
let cost_data = Arc::new(RwLock::new(HashMap::new()));
+ let rates_data = Arc::new(RwLock::new(HashMap::new()));
let latency_data = Arc::new(RwLock::new(HashMap::new()));
for source in sources {
match source {
- MetricsSource::Cost(cfg) => match cfg.provider {
- CostProvider::Digitalocean => {
- let aliases = cfg.model_aliases.clone().unwrap_or_default();
- let data = fetch_do_pricing(&client, &aliases).await;
- info!(models = data.len(), "fetched digitalocean pricing");
- *cost_data.write().await = data;
+ MetricsSource::Cost(cfg) => {
+ let provider = cfg.provider.clone();
+ let url = cfg
+ .url
+ .clone()
+ .unwrap_or_else(|| default_cost_url(&provider).to_string());
+ let aliases = cfg.model_aliases.clone().unwrap_or_default();
+ let provider_name = cost_provider_name(&provider);
- if let Some(interval_secs) = cfg.refresh_interval {
- let cost_clone = Arc::clone(&cost_data);
- let client_clone = client.clone();
- let interval = Duration::from_secs(interval_secs);
- tokio::spawn(async move {
- loop {
- tokio::time::sleep(interval).await;
- let data = fetch_do_pricing(&client_clone, &aliases).await;
- info!(models = data.len(), "refreshed digitalocean pricing");
- *cost_clone.write().await = data;
- }
- });
- }
+ let data = fetch_cost_pricing(&provider, &url, &client, &aliases).await;
+ info!(models = data.len(), provider = provider_name, url = %url, "fetched cost pricing");
+ *cost_data.write().await = blended_costs(&data);
+ *rates_data.write().await = data;
+
+ if let Some(interval_secs) = cfg.refresh_interval {
+ let cost_clone = Arc::clone(&cost_data);
+ let rates_clone = Arc::clone(&rates_data);
+ let client_clone = client.clone();
+ let interval = Duration::from_secs(interval_secs);
+ tokio::spawn(async move {
+ loop {
+ tokio::time::sleep(interval).await;
+ let data =
+ fetch_cost_pricing(&provider, &url, &client_clone, &aliases)
+ .await;
+ info!(models = data.len(), provider = provider_name, url = %url, "refreshed cost pricing");
+ *cost_clone.write().await = blended_costs(&data);
+ *rates_clone.write().await = data;
+ }
+ });
}
- },
+ }
MetricsSource::Latency(cfg) => match cfg.provider {
LatencyProvider::Prometheus => {
let data = fetch_prometheus_metrics(&cfg.url, &cfg.query, &client).await;
@@ -73,10 +156,34 @@ impl ModelMetricsService {
ModelMetricsService {
cost: cost_data,
+ rates: rates_data,
latency: latency_data,
}
}
+ /// Build a service directly from a rates map (no network). Test-only: lets
+ /// handler/router tests exercise switch-cost math with deterministic pricing.
+ #[cfg(test)]
+ pub fn from_rates_for_test(rates: HashMap) -> Self {
+ ModelMetricsService {
+ cost: Arc::new(RwLock::new(blended_costs(&rates))),
+ rates: Arc::new(RwLock::new(rates)),
+ latency: Arc::new(RwLock::new(HashMap::new())),
+ }
+ }
+
+ /// Structured per-million rates for a model, if the cost feed published them.
+ /// Falls back to the bare model id (without `provider/` prefix) like the
+ /// blended cost map does.
+ pub async fn model_rates(&self, model: &str) -> Option {
+ let rates = self.rates.read().await;
+ rates.get(model).copied().or_else(|| {
+ model
+ .split_once('/')
+ .and_then(|(_, bare)| rates.get(bare).copied())
+ })
+ }
+
/// Rank `models` by `policy`, returning them in preference order.
/// Models with no metric data are appended at the end in their original order.
pub async fn rank_models(&self, models: &[String], policy: &SelectionPolicy) -> Vec {
@@ -159,42 +266,172 @@ struct DoModel {
pricing: Option,
}
+/// DigitalOcean catalog pricing. Despite the `_per_million` field names, the DO
+/// API returns these as USD **per token** (e.g. gpt-4o input `2.5e-6`), so they
+/// are scaled to per-million at ingestion to match `ModelRates`' convention.
+///
+/// The catalog publishes the cached-read rate under two names (a migration in
+/// flight): `cache_read_input_price_per_million` and `input_cache_read`. Entries
+/// may carry either or both; both are parsed and the explicit name wins.
#[derive(serde::Deserialize)]
struct DoPricing {
input_price_per_million: Option,
output_price_per_million: Option,
+ cache_read_input_price_per_million: Option,
+ input_cache_read: Option,
+}
+
+impl DoPricing {
+ fn cache_read(&self) -> Option {
+ self.cache_read_input_price_per_million
+ .or(self.input_cache_read)
+ }
+}
+
+#[derive(serde::Deserialize)]
+struct ModelsDevProvider {
+ #[serde(default)]
+ models: HashMap,
+}
+
+#[derive(serde::Deserialize)]
+struct ModelsDevModel {
+ cost: Option,
+}
+
+#[derive(serde::Deserialize)]
+struct ModelsDevCost {
+ input: Option,
+ output: Option,
+ cache_read: Option,
+}
+
+fn default_cost_url(provider: &CostProvider) -> &'static str {
+ match provider {
+ CostProvider::Digitalocean => DO_PRICING_URL,
+ CostProvider::ModelsDev => MODELS_DEV_URL,
+ }
+}
+
+fn cost_provider_name(provider: &CostProvider) -> &'static str {
+ match provider {
+ CostProvider::Digitalocean => "digitalocean",
+ CostProvider::ModelsDev => "models.dev",
+ }
+}
+
+async fn fetch_cost_pricing(
+ provider: &CostProvider,
+ url: &str,
+ client: &reqwest::Client,
+ aliases: &HashMap,
+) -> HashMap {
+ match provider {
+ CostProvider::Digitalocean => fetch_do_pricing(url, client, aliases).await,
+ CostProvider::ModelsDev => fetch_models_dev_pricing(url, client, aliases).await,
+ }
}
async fn fetch_do_pricing(
+ url: &str,
client: &reqwest::Client,
aliases: &HashMap,
-) -> HashMap {
- match client.get(DO_PRICING_URL).send().await {
+) -> HashMap {
+ match client.get(url).send().await {
Ok(resp) => match resp.json::().await {
- Ok(list) => list
- .data
- .into_iter()
- .filter_map(|m| {
- let pricing = m.pricing?;
- let raw_key = m.model_id.clone();
- let key = aliases.get(&raw_key).cloned().unwrap_or(raw_key);
- let cost = pricing.input_price_per_million.unwrap_or(0.0)
- + pricing.output_price_per_million.unwrap_or(0.0);
- Some((key, cost))
- })
- .collect(),
+ Ok(list) => parse_do_pricing(list, aliases),
Err(err) => {
- warn!(error = %err, url = DO_PRICING_URL, "failed to parse digitalocean pricing response");
+ warn!(error = %err, url = %url, "failed to parse digitalocean pricing response");
HashMap::new()
}
},
Err(err) => {
- warn!(error = %err, url = DO_PRICING_URL, "failed to fetch digitalocean pricing");
+ warn!(error = %err, url = %url, "failed to fetch digitalocean pricing");
HashMap::new()
}
}
}
+/// Map the DO catalog into `ModelRates`, scaling DO's per-token prices into the
+/// per-million convention the rest of the router uses.
+fn parse_do_pricing(
+ list: DoModelList,
+ aliases: &HashMap,
+) -> HashMap {
+ list.data
+ .into_iter()
+ .filter_map(|m| {
+ let pricing = m.pricing?;
+ let raw_key = m.model_id.clone();
+ let key = aliases.get(&raw_key).cloned().unwrap_or(raw_key);
+ // DO reports prices per token; scale to per-million so the
+ // absolute-USD consumers (the switch-cost gate) are correct.
+ // Relative ranking via `blended()` is unaffected either way.
+ let rates = ModelRates {
+ input_per_million: pricing.input_price_per_million.unwrap_or(0.0)
+ * TOKENS_PER_MILLION,
+ output_per_million: pricing.output_price_per_million.unwrap_or(0.0)
+ * TOKENS_PER_MILLION,
+ cache_read_per_million: pricing.cache_read().map(|r| r * TOKENS_PER_MILLION),
+ };
+ Some((key, rates))
+ })
+ .collect()
+}
+
+/// models.dev publishes a top-level object keyed by provider id; each provider
+/// carries a `models` map whose keys are `creator/model` ids and whose `cost`
+/// block holds per-million USD rates. We keep the structured rates (including
+/// `cache_read` when published) and key the result by `creator/model_id` so it
+/// lines up with Plano's `provider/model` routing names.
+async fn fetch_models_dev_pricing(
+ url: &str,
+ client: &reqwest::Client,
+ aliases: &HashMap,
+) -> HashMap {
+ match client.get(url).send().await {
+ Ok(resp) => match resp.json::>().await {
+ Ok(providers) => parse_models_dev_pricing(providers, aliases),
+ Err(err) => {
+ warn!(error = %err, url = %url, "failed to parse models.dev pricing response");
+ HashMap::new()
+ }
+ },
+ Err(err) => {
+ warn!(error = %err, url = %url, "failed to fetch models.dev pricing");
+ HashMap::new()
+ }
+ }
+}
+
+fn parse_models_dev_pricing(
+ providers: HashMap,
+ aliases: &HashMap,
+) -> HashMap {
+ let mut out = HashMap::new();
+ for (provider_id, provider) in providers {
+ for (model_key, model) in provider.models {
+ let Some(cost) = model.cost else { continue };
+ let (Some(input), Some(output)) = (cost.input, cost.output) else {
+ continue;
+ };
+ // First-party providers use bare model keys (`claude-opus-4-5`),
+ // so compose `provider/model` to line up with Plano routing names.
+ let raw_key = format!("{provider_id}/{model_key}");
+ let rates = ModelRates {
+ input_per_million: input,
+ output_per_million: output,
+ cache_read_per_million: cost.cache_read,
+ };
+ let key = aliases.get(&raw_key).cloned().unwrap_or(raw_key);
+ out.insert(key, rates);
+ // Also register the bare model id as a fallback lookup.
+ out.entry(model_key).or_insert(rates);
+ }
+ }
+ out
+}
+
#[derive(serde::Deserialize)]
struct PrometheusResponse {
data: PrometheusData,
@@ -255,6 +492,46 @@ mod tests {
SelectionPolicy { prefer }
}
+ fn service_with(
+ cost: HashMap,
+ latency: HashMap,
+ ) -> ModelMetricsService {
+ ModelMetricsService {
+ cost: Arc::new(RwLock::new(cost)),
+ rates: Arc::new(RwLock::new(HashMap::new())),
+ latency: Arc::new(RwLock::new(latency)),
+ }
+ }
+
+ #[test]
+ fn request_cost_openai_convention_subtracts_cached_from_prompt() {
+ // input $3/M, output $15/M, cached read $0.30/M. prompt_tokens (=1000) INCLUDES
+ // the 400 cached tokens → uncached = 600.
+ let rates = ModelRates {
+ input_per_million: 3.0,
+ output_per_million: 15.0,
+ cache_read_per_million: Some(0.30),
+ };
+ let (input, output) = rates.request_cost_usd(1000, 400, 0, 200, true, 0.1);
+ // 600*3 + 400*0.30 = 1800 + 120 = 1920 micro-$/M → /1e6
+ assert!((input - (1920.0 / 1_000_000.0)).abs() < 1e-12);
+ assert!((output - (200.0 * 15.0 / 1_000_000.0)).abs() < 1e-12);
+ }
+
+ #[test]
+ fn request_cost_anthropic_convention_prices_uncached_and_creation() {
+ // Anthropic: prompt_tokens (=input_tokens=100) EXCLUDES cached read (30) and
+ // creation (20). No published cached rate → falls back to input*discount.
+ let rates = ModelRates {
+ input_per_million: 3.0,
+ output_per_million: 15.0,
+ cache_read_per_million: None,
+ };
+ let (input, _output) = rates.request_cost_usd(100, 30, 20, 50, false, 0.1);
+ // uncached 100*3 + cached 30*(3*0.1=0.3) + creation 20*3 = 300 + 9 + 60 = 369
+ assert!((input - (369.0 / 1_000_000.0)).abs() < 1e-12);
+ }
+
#[test]
fn test_rank_by_ascending_metric_picks_lowest_first() {
let models = vec!["a".to_string(), "b".to_string(), "c".to_string()];
@@ -285,15 +562,15 @@ mod tests {
#[tokio::test]
async fn test_rank_models_cheapest() {
- let service = ModelMetricsService {
- cost: Arc::new(RwLock::new({
+ let service = service_with(
+ {
let mut m = HashMap::new();
m.insert("gpt-4o".to_string(), 0.005);
m.insert("gpt-4o-mini".to_string(), 0.0001);
m
- })),
- latency: Arc::new(RwLock::new(HashMap::new())),
- };
+ },
+ HashMap::new(),
+ );
let models = vec!["gpt-4o".to_string(), "gpt-4o-mini".to_string()];
let result = service
.rank_models(&models, &make_policy(SelectionPreference::Cheapest))
@@ -303,15 +580,12 @@ mod tests {
#[tokio::test]
async fn test_rank_models_fastest() {
- let service = ModelMetricsService {
- cost: Arc::new(RwLock::new(HashMap::new())),
- latency: Arc::new(RwLock::new({
- let mut m = HashMap::new();
- m.insert("gpt-4o".to_string(), 200.0);
- m.insert("claude-sonnet".to_string(), 120.0);
- m
- })),
- };
+ let service = service_with(HashMap::new(), {
+ let mut m = HashMap::new();
+ m.insert("gpt-4o".to_string(), 200.0);
+ m.insert("claude-sonnet".to_string(), 120.0);
+ m
+ });
let models = vec!["gpt-4o".to_string(), "claude-sonnet".to_string()];
let result = service
.rank_models(&models, &make_policy(SelectionPreference::Fastest))
@@ -321,10 +595,7 @@ mod tests {
#[tokio::test]
async fn test_rank_models_fallback_no_metrics() {
- let service = ModelMetricsService {
- cost: Arc::new(RwLock::new(HashMap::new())),
- latency: Arc::new(RwLock::new(HashMap::new())),
- };
+ let service = service_with(HashMap::new(), HashMap::new());
let models = vec!["model-a".to_string(), "model-b".to_string()];
let result = service
.rank_models(&models, &make_policy(SelectionPreference::Cheapest))
@@ -334,14 +605,14 @@ mod tests {
#[tokio::test]
async fn test_rank_models_partial_data_appended_last() {
- let service = ModelMetricsService {
- cost: Arc::new(RwLock::new({
+ let service = service_with(
+ {
let mut m = HashMap::new();
m.insert("gpt-4o".to_string(), 0.005);
m
- })),
- latency: Arc::new(RwLock::new(HashMap::new())),
- };
+ },
+ HashMap::new(),
+ );
let models = vec!["gpt-4o-mini".to_string(), "gpt-4o".to_string()];
let result = service
.rank_models(&models, &make_policy(SelectionPreference::Cheapest))
@@ -351,15 +622,15 @@ mod tests {
#[tokio::test]
async fn test_rank_models_none_preserves_order() {
- let service = ModelMetricsService {
- cost: Arc::new(RwLock::new({
+ let service = service_with(
+ {
let mut m = HashMap::new();
m.insert("gpt-4o-mini".to_string(), 0.0001);
m.insert("gpt-4o".to_string(), 0.005);
m
- })),
- latency: Arc::new(RwLock::new(HashMap::new())),
- };
+ },
+ HashMap::new(),
+ );
let models = vec!["gpt-4o".to_string(), "gpt-4o-mini".to_string()];
let result = service
.rank_models(&models, &make_policy(SelectionPreference::None))
@@ -368,6 +639,201 @@ mod tests {
assert_eq!(result, vec!["gpt-4o", "gpt-4o-mini"]);
}
+ #[test]
+ fn test_parse_do_pricing_scales_per_token_to_per_million() {
+ // DO's catalog returns USD *per token* despite the `_per_million` names.
+ // These are the real values from the live catalog.
+ let json = r#"{
+ "data": [
+ {
+ "model_id": "digitalocean/anthropic-claude-4.6-sonnet",
+ "pricing": {
+ "input_price_per_million": 3e-6,
+ "output_price_per_million": 15e-6,
+ "cache_read_input_price_per_million": 3e-7
+ }
+ },
+ {
+ "model_id": "digitalocean/openai-gpt-4o",
+ "pricing": {
+ "input_price_per_million": 2.5e-6,
+ "output_price_per_million": 10e-6
+ }
+ }
+ ]
+ }"#;
+ let list: DoModelList = serde_json::from_str(json).unwrap();
+ let rates = parse_do_pricing(list, &HashMap::new());
+
+ // Scaled to per-million so the switch-cost gate compares real dollars.
+ let sonnet = rates
+ .get("digitalocean/anthropic-claude-4.6-sonnet")
+ .unwrap();
+ assert!((sonnet.input_per_million - 3.0).abs() < 1e-9);
+ assert!((sonnet.output_per_million - 15.0).abs() < 1e-9);
+ // DO *does* publish a cached-read rate — it must be captured, not dropped.
+ assert_eq!(sonnet.cache_read_per_million, Some(0.3));
+
+ let gpt4o = rates.get("digitalocean/openai-gpt-4o").unwrap();
+ assert!((gpt4o.input_per_million - 2.5).abs() < 1e-9);
+ assert_eq!(gpt4o.cache_read_per_million, None);
+
+ // Regression: a ~5k-token switch off warm sonnet to cold gpt-4o must now
+ // cost real dollars, not ~1e-8. This is the bug the live gate test caught.
+ let switch_cost =
+ 5_000.0 / 1_000_000.0 * (gpt4o.input_per_million - sonnet.cached_input_rate(0.1));
+ assert!(
+ switch_cost > 0.01,
+ "expected ~$0.011 of switch cost, got {switch_cost}"
+ );
+ }
+
+ #[test]
+ fn test_parse_do_pricing_reads_cache_rate_from_either_field_name() {
+ // Real catalog shape (Jul 2026): DO publishes the cached-read rate under
+ // both `cache_read_input_price_per_million` and `input_cache_read` for
+ // some models, and only one of the two for others. All variants must
+ // resolve; the explicit name wins when both are present.
+ let json = r#"{
+ "data": [
+ {
+ "model_id": "kimi-k2.5",
+ "pricing": {
+ "input_price_per_million": 3.75e-7,
+ "output_price_per_million": 2.025e-6,
+ "cache_read_input_price_per_million": 2.03e-7,
+ "input_cache_read": 2.03e-7
+ }
+ },
+ {
+ "model_id": "glm-5",
+ "pricing": {
+ "input_price_per_million": 6e-7,
+ "output_price_per_million": 2.2e-6,
+ "input_cache_read": 1.1e-7
+ }
+ },
+ {
+ "model_id": "legacy-model",
+ "pricing": {
+ "input_price_per_million": 5e-7,
+ "output_price_per_million": 1e-6
+ }
+ }
+ ]
+ }"#;
+ let list: DoModelList = serde_json::from_str(json).unwrap();
+ let rates = parse_do_pricing(list, &HashMap::new());
+
+ // Both names present -> parsed (and scaled per-token -> per-million).
+ let kimi = rates.get("kimi-k2.5").unwrap();
+ assert!((kimi.input_per_million - 0.375).abs() < 1e-9);
+ assert!((kimi.cache_read_per_million.unwrap() - 0.203).abs() < 1e-9);
+
+ // Only the new `input_cache_read` name -> still parsed.
+ let glm = rates.get("glm-5").unwrap();
+ assert!((glm.cache_read_per_million.unwrap() - 0.11).abs() < 1e-9);
+
+ // Neither -> None, so the budget falls back to input x cache_read_discount.
+ let legacy = rates.get("legacy-model").unwrap();
+ assert_eq!(legacy.cache_read_per_million, None);
+ assert!((legacy.cached_input_rate(0.1) - 0.05).abs() < 1e-9);
+ }
+
+ #[test]
+ fn test_parse_models_dev_pricing_composes_provider_keys() {
+ let json = r#"{
+ "anthropic": {
+ "models": {
+ "claude-opus-4-5": {"cost": {"input": 5.0, "output": 25.0, "cache_read": 0.5}}
+ }
+ },
+ "groq": {
+ "models": {
+ "llama-3.3-70b-versatile": {"cost": {"input": 0.59, "output": 0.79}},
+ "whisper-large-v3-turbo": {"cost": null}
+ }
+ }
+ }"#;
+ let providers: HashMap = serde_json::from_str(json).unwrap();
+ let aliases = HashMap::new();
+ let rates = parse_models_dev_pricing(providers, &aliases);
+
+ let opus = rates.get("anthropic/claude-opus-4-5").unwrap();
+ assert_eq!(opus.blended(), 30.0);
+ assert_eq!(opus.cache_read_per_million, Some(0.5));
+ let llama = rates.get("groq/llama-3.3-70b-versatile").unwrap();
+ assert_eq!(llama.blended(), 1.38);
+ assert_eq!(llama.cache_read_per_million, None);
+ // bare fallback also registered
+ assert_eq!(
+ rates.get("claude-opus-4-5").map(|r| r.blended()),
+ Some(30.0)
+ );
+ // models with no cost block are skipped
+ assert!(!rates.contains_key("groq/whisper-large-v3-turbo"));
+ }
+
+ #[test]
+ fn test_parse_models_dev_pricing_applies_aliases() {
+ let json = r#"{
+ "openai": {"models": {"gpt-oss-120b": {"cost": {"input": 1.0, "output": 2.0}}}}
+ }"#;
+ let providers: HashMap = serde_json::from_str(json).unwrap();
+ let mut aliases = HashMap::new();
+ aliases.insert(
+ "openai/gpt-oss-120b".to_string(),
+ "openai/gpt-4o".to_string(),
+ );
+ let rates = parse_models_dev_pricing(providers, &aliases);
+
+ assert_eq!(rates.get("openai/gpt-4o").map(|r| r.blended()), Some(3.0));
+ assert!(!rates.contains_key("openai/gpt-oss-120b"));
+ }
+
+ #[test]
+ fn test_cached_input_rate_prefers_published_rate() {
+ let with_feed = ModelRates {
+ input_per_million: 3.0,
+ output_per_million: 15.0,
+ cache_read_per_million: Some(0.3),
+ };
+ // Published cache_read wins; the discount is ignored.
+ assert_eq!(with_feed.cached_input_rate(0.5), 0.3);
+
+ let without_feed = ModelRates {
+ input_per_million: 3.0,
+ output_per_million: 15.0,
+ cache_read_per_million: None,
+ };
+ // Falls back to input * discount.
+ assert!((without_feed.cached_input_rate(0.1) - 0.3).abs() < 1e-9);
+ }
+
+ #[tokio::test]
+ async fn test_model_rates_falls_back_to_bare_model_id() {
+ let service = ModelMetricsService {
+ cost: Arc::new(RwLock::new(HashMap::new())),
+ rates: Arc::new(RwLock::new({
+ let mut m = HashMap::new();
+ m.insert(
+ "claude-sonnet-4-5".to_string(),
+ ModelRates {
+ input_per_million: 3.0,
+ output_per_million: 15.0,
+ cache_read_per_million: Some(0.3),
+ },
+ );
+ m
+ })),
+ latency: Arc::new(RwLock::new(HashMap::new())),
+ };
+ // provider-prefixed lookup falls back to the bare id
+ let rates = service.model_rates("vercel/claude-sonnet-4-5").await;
+ assert_eq!(rates.map(|r| r.input_per_million), Some(3.0));
+ assert!(service.model_rates("unknown/model").await.is_none());
+ }
+
#[test]
fn test_rank_by_ascending_metric_nan_treated_as_missing() {
let models = vec![
diff --git a/crates/brightstaff/src/router/orchestrator.rs b/crates/brightstaff/src/router/orchestrator.rs
index 2d7b25de..050797d3 100644
--- a/crates/brightstaff/src/router/orchestrator.rs
+++ b/crates/brightstaff/src/router/orchestrator.rs
@@ -20,9 +20,41 @@ use crate::metrics::labels as metric_labels;
use crate::router::orchestrator_model_v1;
use crate::session_cache::SessionCache;
-pub use crate::session_cache::CachedRoute;
+pub use crate::session_cache::SessionBinding;
const DEFAULT_SESSION_TTL_SECONDS: u64 = 600;
+const TOKENS_PER_MILLION: f64 = 1_000_000.0;
+
+/// Input-cost of moving a session off its anchor onto `candidate`.
+///
+/// Staying re-reads the whole context at the anchor's read rate (`anchor_read_rate` — the
+/// cached rate when prompt caching keeps the anchor warm, the plain uncached rate when it
+/// doesn't). The candidate re-reads `candidate_warm_tokens` (whatever it still has cached
+/// from an earlier visit this session) at *its* cached rate, and the remaining,
+/// freshly-appended tokens at its uncached rate. When the candidate is cold
+/// (`candidate_warm_tokens == 0`) this reduces to the whole context at the uncached rate —
+/// i.e. a first-time switch. Crediting warm tokens is what makes an A→B→A return cheap:
+/// on the way back A still holds most of the context, so only the delta is re-ingested.
+///
+/// Deliberately input-only: output-token savings are unknowable before the response is
+/// generated (reasoning models can emit 5-10x the tokens for the same task), so they are
+/// never credited. Rates are USD per million tokens. Negative when the switch is outright
+/// cheaper than staying; the caller draws a positive cost down from the overhead cap.
+pub fn switch_cost_in_usd(
+ context_tokens: u64,
+ candidate_warm_tokens: u64,
+ anchor_read_rate: f64,
+ candidate_uncached_rate: f64,
+ candidate_cached_rate: f64,
+) -> f64 {
+ let warm = candidate_warm_tokens.min(context_tokens);
+ let fresh = context_tokens - warm;
+ let candidate_cost = (fresh as f64 * candidate_uncached_rate
+ + warm as f64 * candidate_cached_rate)
+ / TOKENS_PER_MILLION;
+ let anchor_cost = context_tokens as f64 * anchor_read_rate / TOKENS_PER_MILLION;
+ candidate_cost - anchor_cost
+}
pub struct OrchestratorService {
orchestrator_url: String,
@@ -126,43 +158,109 @@ impl OrchestratorService {
}
}
- pub async fn get_cached_route(
+ /// Look up a session binding. Warmth is the caller's concern (time since
+ /// `last_used`); this only reports whether a binding exists.
+ pub async fn get_binding(
&self,
session_id: &str,
tenant_id: Option<&str>,
- ) -> Option {
+ ) -> Option {
let cache = self.session_cache.as_ref()?;
let result = cache.get(&Self::session_key(tenant_id, session_id)).await;
- bs_metrics::record_session_cache_event(if result.is_some() {
- metric_labels::SESSION_CACHE_HIT
- } else {
- metric_labels::SESSION_CACHE_MISS
+ bs_metrics::record_session_cache_event(match result {
+ Some(_) => metric_labels::SESSION_CACHE_HIT,
+ None => metric_labels::SESSION_CACHE_MISS,
});
result
}
- pub async fn cache_route(
+ /// The GC bound for a session binding: the per-scope override when provided,
+ /// otherwise the global `routing.session_ttl_seconds`. This only governs when an
+ /// idle binding is reclaimed from memory — not whether its cache is warm.
+ pub fn effective_session_ttl(&self, ttl_override_seconds: Option) -> Duration {
+ ttl_override_seconds
+ .map(Duration::from_secs)
+ .unwrap_or(self.session_ttl)
+ }
+
+ /// Persist a session binding with a GC bound (defaults to `routing.session_ttl_seconds`
+ /// when `gc_ttl` is `None`).
+ pub async fn store_binding(
&self,
- session_id: String,
+ session_id: &str,
tenant_id: Option<&str>,
- model_name: String,
- route_name: Option,
+ binding: SessionBinding,
+ gc_ttl: Option,
) {
if let Some(ref cache) = self.session_cache {
cache
.put(
- &Self::session_key(tenant_id, &session_id),
- CachedRoute {
- model_name,
- route_name,
- },
- self.session_ttl,
+ &Self::session_key(tenant_id, session_id),
+ binding,
+ gc_ttl.unwrap_or(self.session_ttl),
)
.await;
bs_metrics::record_session_cache_event(metric_labels::SESSION_CACHE_STORE);
}
}
+ /// Structured per-million pricing for a model, from the configured cost feed.
+ /// `None` when no cost source is configured or the model is unknown to the feed.
+ pub async fn model_rates(&self, model: &str) -> Option {
+ self.metrics_service.as_ref()?.model_rates(model).await
+ }
+
+ /// Estimate the input-cost (USD) of switching a session from `anchor_model`
+ /// (the model that handled the latest request) to `candidate_model`.
+ ///
+ /// The anchor is warm (this is only called for warm sessions), so staying re-reads the
+ /// context at its *cached* rate, while `candidate_warm_tokens` (context the candidate
+ /// still holds from an earlier visit this session) re-read at the candidate's cached
+ /// rate. Provider caches are assumed real whenever the routing budget is active —
+ /// most providers cache automatically, with or without Plano's marker injection
+ /// (`prompt_caching.enabled` is not consulted here).
+ ///
+ /// Fetches per-model rates from the configured cost feed; returns `None` when pricing
+ /// is missing for either side so the caller can fail open (switch freely) rather than
+ /// veto the router on guesswork. `cache_read_discount` estimates a model's cached-read
+ /// rate when the feed doesn't publish one. Negative when the switch is outright cheaper.
+ pub async fn estimate_switch_cost_in_usd(
+ &self,
+ context_tokens: u64,
+ anchor_model: &str,
+ candidate_model: &str,
+ candidate_warm_tokens: u64,
+ cache_read_discount: f64,
+ ) -> Option {
+ let anchor = self.model_rates(anchor_model).await?;
+ let candidate = self.model_rates(candidate_model).await?;
+ Some(switch_cost_in_usd(
+ context_tokens,
+ candidate_warm_tokens,
+ anchor.cached_input_rate(cache_read_discount),
+ candidate.input_per_million,
+ candidate.cached_input_rate(cache_read_discount),
+ ))
+ }
+
+ /// This turn's contribution to the session's *never-switch* baseline: the USD cost of
+ /// reading `context_tokens` on `model` — the session's `default_model`, i.e. what it
+ /// would have paid by never switching (not the possibly-drifted current anchor). Priced
+ /// at the cached input rate: the never-switch path stays warm, and provider caches are
+ /// assumed real whenever the routing budget is active. Summed across turns, this is
+ /// the denominator the percentage overhead cap is measured against. `None` when the
+ /// model has no pricing (the caller then can't grow the baseline this turn).
+ pub async fn context_read_cost_in_usd(
+ &self,
+ context_tokens: u64,
+ model: &str,
+ cache_read_discount: f64,
+ ) -> Option {
+ let rates = self.model_rates(model).await?;
+ let context_millions = context_tokens as f64 / TOKENS_PER_MILLION;
+ Some(context_millions * rates.cached_input_rate(cache_read_discount))
+ }
+
// ---- LLM routing ----
pub async fn determine_route(
@@ -348,86 +446,190 @@ mod tests {
)
}
- #[tokio::test]
- async fn test_cache_miss_returns_none() {
- let svc = make_orchestrator_service(600, 100);
- assert!(svc
- .get_cached_route("unknown-session", None)
- .await
- .is_none());
+ fn binding(model: &str, route_name: Option<&str>) -> SessionBinding {
+ SessionBinding {
+ anchor_model: model.to_string(),
+ default_model: model.to_string(),
+ route_name: route_name.map(|r| r.to_string()),
+ prefix_hash: None,
+ last_used: std::time::SystemTime::now(),
+ cached_tokens: 0,
+ baseline_usd: 0.0,
+ switch_spend_usd: 0.0,
+ switches: 0,
+ session_cost_usd: 0.0,
+ history: Vec::new(),
+ }
}
#[tokio::test]
- async fn test_cache_hit_returns_cached_route() {
+ async fn test_cache_miss_returns_none() {
let svc = make_orchestrator_service(600, 100);
- svc.cache_route(
- "s1".to_string(),
- None,
- "gpt-4o".to_string(),
- Some("code".to_string()),
- )
- .await;
+ assert!(svc.get_binding("unknown-session", None).await.is_none());
+ }
- let cached = svc.get_cached_route("s1", None).await.unwrap();
- assert_eq!(cached.model_name, "gpt-4o");
+ #[tokio::test]
+ async fn test_cache_hit_returns_binding() {
+ let svc = make_orchestrator_service(600, 100);
+ svc.store_binding("s1", None, binding("gpt-4o", Some("code")), None)
+ .await;
+
+ let cached = svc.get_binding("s1", None).await.unwrap();
+ assert_eq!(cached.anchor_model, "gpt-4o");
assert_eq!(cached.route_name, Some("code".to_string()));
}
#[tokio::test]
async fn test_cache_expired_entry_returns_none() {
let svc = make_orchestrator_service(0, 100);
- svc.cache_route("s1".to_string(), None, "gpt-4o".to_string(), None)
+ svc.store_binding("s1", None, binding("gpt-4o", None), None)
.await;
- assert!(svc.get_cached_route("s1", None).await.is_none());
+ assert!(svc.get_binding("s1", None).await.is_none());
}
#[tokio::test]
async fn test_expired_entries_not_returned() {
let svc = make_orchestrator_service(0, 100);
- svc.cache_route("s1".to_string(), None, "gpt-4o".to_string(), None)
+ svc.store_binding("s1", None, binding("gpt-4o", None), None)
.await;
- svc.cache_route("s2".to_string(), None, "claude".to_string(), None)
+ svc.store_binding("s2", None, binding("claude", None), None)
.await;
- assert!(svc.get_cached_route("s1", None).await.is_none());
- assert!(svc.get_cached_route("s2", None).await.is_none());
+ assert!(svc.get_binding("s1", None).await.is_none());
+ assert!(svc.get_binding("s2", None).await.is_none());
}
#[tokio::test]
async fn test_cache_evicts_oldest_when_full() {
let svc = make_orchestrator_service(600, 2);
- svc.cache_route("s1".to_string(), None, "model-a".to_string(), None)
+ svc.store_binding("s1", None, binding("model-a", None), None)
.await;
tokio::time::sleep(Duration::from_millis(10)).await;
- svc.cache_route("s2".to_string(), None, "model-b".to_string(), None)
+ svc.store_binding("s2", None, binding("model-b", None), None)
.await;
- svc.cache_route("s3".to_string(), None, "model-c".to_string(), None)
+ svc.store_binding("s3", None, binding("model-c", None), None)
.await;
- assert!(svc.get_cached_route("s1", None).await.is_none());
- assert!(svc.get_cached_route("s2", None).await.is_some());
- assert!(svc.get_cached_route("s3", None).await.is_some());
+ assert!(svc.get_binding("s1", None).await.is_none());
+ assert!(svc.get_binding("s2", None).await.is_some());
+ assert!(svc.get_binding("s3", None).await.is_some());
}
#[tokio::test]
async fn test_cache_update_existing_session_does_not_evict() {
let svc = make_orchestrator_service(600, 2);
- svc.cache_route("s1".to_string(), None, "model-a".to_string(), None)
+ svc.store_binding("s1", None, binding("model-a", None), None)
.await;
- svc.cache_route("s2".to_string(), None, "model-b".to_string(), None)
+ svc.store_binding("s2", None, binding("model-b", None), None)
.await;
- svc.cache_route(
- "s1".to_string(),
+ svc.store_binding("s1", None, binding("model-a-updated", Some("route")), None)
+ .await;
+
+ let s1 = svc.get_binding("s1", None).await.unwrap();
+ assert_eq!(s1.anchor_model, "model-a-updated");
+ assert!(svc.get_binding("s2", None).await.is_some());
+ }
+
+ #[tokio::test]
+ async fn test_gc_ttl_override_extends_binding_lifetime() {
+ // Global GC bound of 0 would reclaim immediately; the per-call override keeps it.
+ let svc = make_orchestrator_service(0, 100);
+ svc.store_binding(
+ "s1",
None,
- "model-a-updated".to_string(),
- Some("route".to_string()),
+ binding("gpt-4o", None),
+ Some(Duration::from_secs(600)),
)
.await;
+ let cached = svc.get_binding("s1", None).await.unwrap();
+ assert_eq!(cached.anchor_model, "gpt-4o");
+ }
- let s1 = svc.get_cached_route("s1", None).await.unwrap();
- assert_eq!(s1.model_name, "model-a-updated");
- assert!(svc.get_cached_route("s2", None).await.is_some());
+ #[tokio::test]
+ async fn test_binding_fields_round_trip_through_cache() {
+ let svc = make_orchestrator_service(600, 100);
+ let mut b = binding("gpt-4o", None);
+ b.prefix_hash = Some(0xdead_beef);
+ b.cached_tokens = 12_345;
+ b.baseline_usd = 1.5;
+ b.switch_spend_usd = 0.42;
+ b.switches = 3;
+ b.session_cost_usd = 2.75;
+ svc.store_binding("s1", None, b, None).await;
+
+ let cached = svc.get_binding("s1", None).await.unwrap();
+ assert_eq!(cached.prefix_hash, Some(0xdead_beef));
+ assert_eq!(cached.cached_tokens, 12_345);
+ assert!((cached.baseline_usd - 1.5).abs() < 1e-9);
+ assert!((cached.switch_spend_usd - 0.42).abs() < 1e-9);
+ assert_eq!(cached.switches, 3);
+ assert!((cached.session_cost_usd - 2.75).abs() < 1e-9);
+ }
+
+ // ---- switch-cost math ----
+ //
+ // Real models.dev rates (USD per million input tokens):
+ // claude-opus-4-1: input 15, cache_read 1.5
+ // claude-sonnet-4-5: input 3, cache_read 0.3
+ // claude-haiku-4-5: input 1, cache_read 0.1
+ // gpt-4.1: input 2, cache_read 0.5
+
+ #[test]
+ fn negative_cost_when_candidate_undercuts_cached_rate() {
+ // Anchor opus (cached 1.5) -> haiku (uncached 1.0) over 100k context, cold
+ // candidate: cost = 0.1M x (1.0 - 1.5) = -$0.05 — cheaper even after re-reading.
+ let cost = switch_cost_in_usd(100_000, 0, 1.5, 1.0, 0.1);
+ assert!((cost - (-0.05)).abs() < 1e-9);
+ }
+
+ #[test]
+ fn positive_cost_when_candidate_pricier_than_cached_rate() {
+ // Anchor opus (cached 1.5) -> gpt-4.1 (uncached 2.0) over 100k, cold candidate:
+ // cost = 0.1M x (2.0 - 1.5) = +$0.05.
+ let cost = switch_cost_in_usd(100_000, 0, 1.5, 2.0, 0.5);
+ assert!((cost - 0.05).abs() < 1e-9);
+ }
+
+ #[test]
+ fn large_context_amplifies_cost() {
+ // Anchor sonnet (cached 0.3) -> gpt-5.5-class (uncached 5.0) over 150k, cold:
+ // cost = 0.15M x (5.0 - 0.3) = +$0.705.
+ let cost = switch_cost_in_usd(150_000, 0, 0.3, 5.0, 0.5);
+ assert!((cost - 0.705).abs() < 1e-9);
+ }
+
+ #[test]
+ fn cost_scales_linearly_with_context() {
+ let small = switch_cost_in_usd(10_000, 0, 0.3, 0.8, 0.1);
+ let large = switch_cost_in_usd(1_000_000, 0, 0.3, 0.8, 0.1);
+ assert!((large / small - 100.0).abs() < 1e-6);
+ }
+
+ #[test]
+ fn tiny_context_cost_is_negligible() {
+ // 2k-token chat: even an expensive candidate costs ~$0.009.
+ let cost = switch_cost_in_usd(2_000, 0, 0.3, 5.0, 0.5);
+ assert!(cost < 0.01);
+ }
+
+ #[test]
+ fn warm_return_charges_only_the_delta() {
+ // Anchor sonnet (cached 0.3). Candidate gpt-4.1 (uncached 2.0, cached 0.5) is
+ // still warm from an earlier visit holding 90k of the 100k context. Only the
+ // 10k fresh tokens re-read at 2.0; the 90k warm tokens re-read at 0.5:
+ // candidate = (10k x 2.0 + 90k x 0.5)/1M = 0.020 + 0.045 = $0.065
+ // anchor = 100k x 0.3 /1M = $0.030
+ // switch = 0.065 - 0.030 = +$0.035
+ let warm = switch_cost_in_usd(100_000, 90_000, 0.3, 2.0, 0.5);
+ assert!(
+ (warm - 0.035).abs() < 1e-9,
+ "warm-return cost {warm} != 0.035"
+ );
+ // Cold, same switch re-reads the whole 100k at 2.0: (100k x 2.0)/1M - 0.03 = $0.17.
+ let cold = switch_cost_in_usd(100_000, 0, 0.3, 2.0, 0.5);
+ assert!((cold - 0.17).abs() < 1e-9, "cold cost {cold} != 0.17");
+ assert!(warm < cold, "returning to a warm model must be cheaper");
}
}
diff --git a/crates/brightstaff/src/session_cache/memory.rs b/crates/brightstaff/src/session_cache/memory.rs
index 54185738..6d67a8b4 100644
--- a/crates/brightstaff/src/session_cache/memory.rs
+++ b/crates/brightstaff/src/session_cache/memory.rs
@@ -9,9 +9,9 @@ use lru::LruCache;
use tokio::sync::Mutex;
use tracing::info;
-use super::{CachedRoute, SessionCache};
+use super::{SessionBinding, SessionCache};
-type CacheStore = Mutex>;
+type CacheStore = Mutex>;
pub struct MemorySessionCache {
store: Arc,
@@ -38,6 +38,8 @@ impl MemorySessionCache {
async fn evict_expired(store: &CacheStore) {
let mut cache = store.lock().await;
+ // The TTL is only a GC bound: drop bindings once they've outlived the window
+ // in which they could plausibly still be warm.
let expired: Vec = cache
.iter()
.filter(|(_, (_, inserted_at, ttl))| inserted_at.elapsed() >= *ttl)
@@ -59,21 +61,21 @@ impl MemorySessionCache {
#[async_trait]
impl SessionCache for MemorySessionCache {
- async fn get(&self, key: &str) -> Option {
+ async fn get(&self, key: &str) -> Option {
let mut cache = self.store.lock().await;
- if let Some((route, inserted_at, ttl)) = cache.get(key) {
+ if let Some((binding, inserted_at, ttl)) = cache.get(key) {
if inserted_at.elapsed() < *ttl {
- return Some(route.clone());
+ return Some(binding.clone());
}
}
None
}
- async fn put(&self, key: &str, route: CachedRoute, ttl: Duration) {
+ async fn put(&self, key: &str, binding: SessionBinding, ttl: Duration) {
self.store
.lock()
.await
- .put(key.to_string(), (route, Instant::now(), ttl));
+ .put(key.to_string(), (binding, Instant::now(), ttl));
}
async fn remove(&self, key: &str) {
diff --git a/crates/brightstaff/src/session_cache/mod.rs b/crates/brightstaff/src/session_cache/mod.rs
index 4cf0dda5..8a5e2e07 100644
--- a/crates/brightstaff/src/session_cache/mod.rs
+++ b/crates/brightstaff/src/session_cache/mod.rs
@@ -1,4 +1,5 @@
use std::sync::Arc;
+use std::time::{SystemTime, UNIX_EPOCH};
use async_trait::async_trait;
use common::configuration::Configuration;
@@ -8,21 +9,155 @@ use tracing::{debug, info};
pub mod memory;
pub mod redis;
+/// A conversation's binding to a model, plus the state the session router needs to
+/// reason about cache warmth and switch affordability across turns.
+///
+/// Warmth is no longer derived from the cache's own expiry — the entry is kept alive
+/// as a plain KV value (subject only to a GC bound) and the router decides warmth from
+/// [`SessionBinding::last_used`] against the provider's cache window. This is what lets
+/// the decision path reason about warmth without ever seeing a provider response.
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
-pub struct CachedRoute {
- pub model_name: String,
+pub struct SessionBinding {
+ /// Provider-qualified model that handled the latest request (e.g. `openai/gpt-4o`).
+ /// This is what the session is currently *warm on*; a proposed switch's cost is
+ /// measured against reading the context at this model's cached rate. It tracks the
+ /// last dispatched model and changes whenever a switch is honored.
+ pub anchor_model: String,
+ /// Provider-qualified model the session started on this warm episode — the model it
+ /// would have stayed on had it *never switched*. The never-switch baseline is priced
+ /// against this (not `anchor_model`, which drifts as switches happen). Set when a warm
+ /// episode begins and preserved across its turns.
+ #[serde(default)]
+ pub default_model: String,
+ #[serde(default, skip_serializing_if = "Option::is_none")]
pub route_name: Option,
+ /// Hash of the stable prompt prefix (system + tools) observed when the binding was
+ /// stored. Used to detect prefix drift: if a later request's prefix hash differs,
+ /// the provider cache is already lost so a switch is free.
+ #[serde(default, skip_serializing_if = "Option::is_none")]
+ pub prefix_hash: Option,
+ /// When this session was last dispatched. Warmth = `now - last_used` compared
+ /// against the provider's idle/hard cache window.
+ #[serde(default = "SystemTime::now", with = "epoch_secs")]
+ pub last_used: SystemTime,
+ /// Best estimate of the cacheable context size (input tokens) — the tokens a switch
+ /// would have to re-ingest at the uncached rate. Refined from real usage on the
+ /// full-proxy path; the tokenizer estimate on the decision path.
+ #[serde(default)]
+ pub cached_tokens: u64,
+ /// Cumulative *never-switch* baseline (USD) for this warm episode: the running cost
+ /// the session would have paid by staying on its `default_model`. Grows each warm
+ /// turn. This is the denominator the percentage overhead cap is measured against.
+ #[serde(default)]
+ pub baseline_usd: f64,
+ /// Cumulative overhead (USD) actually spent on paid switches this warm episode.
+ /// Monotonic: paid switches add to it, free/cheaper switches never subtract. A paid
+ /// switch is allowed only while `switch_spend_usd + cost <= pct * baseline_usd`.
+ #[serde(default)]
+ pub switch_spend_usd: f64,
+ /// Number of model switches taken during this warm session (observability).
+ #[serde(default)]
+ pub switches: u32,
+ /// Bounded most-recent-first-ish record of the models this session has been
+ /// dispatched to, each with when it was last used and how large the context was
+ /// then. Lets the switch-cost estimate credit a *return* to a still-warm model (it
+ /// re-reads only the tokens appended since, not the whole context) and gives future
+ /// routing policies per-model recency to reason about. Capped at
+ /// [`MAX_ROUTE_HISTORY`] distinct models (LRU-evicted).
+ #[serde(default, skip_serializing_if = "Vec::is_empty")]
+ pub history: Vec,
+ /// Cumulative *actual* cost (USD, input + output) of the whole conversation, priced
+ /// from the configured catalog rates and refined from real usage each turn on the
+ /// full-proxy path. Conversation-level (not per warm episode): it persists across
+ /// cold re-binds and is best-effort (resets only if the binding is evicted).
+ #[serde(default)]
+ pub session_cost_usd: f64,
+}
+
+/// One model this session has been dispatched to, with the recency and context size
+/// needed to estimate whether its provider cache is still warm on a later return.
+#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
+pub struct RouteVisit {
+ /// Provider-qualified model id (e.g. `openai/gpt-4o`).
+ pub model: String,
+ /// When this model last handled a turn in this session.
+ #[serde(with = "epoch_secs")]
+ pub last_used: SystemTime,
+ /// Context size (input tokens) that model saw on its last turn — the prefix it may
+ /// still have cached if it's revisited before its cache window elapses.
+ #[serde(default)]
+ pub cached_tokens: u64,
+}
+
+/// Max distinct models retained in [`SessionBinding::history`]. Small: real sessions
+/// touch a handful of models, and the entry has to stay compact on the Redis wire.
+pub const MAX_ROUTE_HISTORY: usize = 8;
+
+/// Record (or refresh) a model visit in `history`, then LRU-evict down to
+/// [`MAX_ROUTE_HISTORY`]. Refreshing an existing model updates its recency and context
+/// size in place rather than appending a duplicate.
+pub fn record_route_visit(
+ history: &mut Vec,
+ model: &str,
+ last_used: SystemTime,
+ cached_tokens: u64,
+) {
+ if let Some(entry) = history.iter_mut().find(|e| e.model == model) {
+ entry.last_used = last_used;
+ entry.cached_tokens = cached_tokens;
+ } else {
+ history.push(RouteVisit {
+ model: model.to_string(),
+ last_used,
+ cached_tokens,
+ });
+ }
+ while history.len() > MAX_ROUTE_HISTORY {
+ if let Some(idx) = history
+ .iter()
+ .enumerate()
+ .min_by_key(|(_, e)| e.last_used)
+ .map(|(i, _)| i)
+ {
+ history.remove(idx);
+ } else {
+ break;
+ }
+ }
+}
+
+/// Serde helper: persist `SystemTime` as whole epoch seconds so the Redis wire format
+/// is stable and compact (the default `SystemTime` representation is version-fragile).
+mod epoch_secs {
+ use super::{Duration, SystemTime, UNIX_EPOCH};
+ use serde::{Deserialize, Deserializer, Serializer};
+
+ pub fn serialize(t: &SystemTime, s: S) -> Result {
+ let secs = t
+ .duration_since(UNIX_EPOCH)
+ .map(|d| d.as_secs())
+ .unwrap_or(0);
+ s.serialize_u64(secs)
+ }
+
+ pub fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result {
+ let secs = u64::deserialize(d)?;
+ Ok(UNIX_EPOCH + Duration::from_secs(secs))
+ }
}
#[async_trait]
pub trait SessionCache: Send + Sync {
- /// Look up a cached routing decision by key.
- async fn get(&self, key: &str) -> Option;
+ /// Look up a session binding by key. `None` when absent or GC-evicted. Warmth is
+ /// the caller's concern (time since `last_used`), not the cache's.
+ async fn get(&self, key: &str) -> Option;
- /// Store a routing decision in the session cache with the given TTL.
- async fn put(&self, key: &str, route: CachedRoute, ttl: Duration);
+ /// Store a session binding with the given GC TTL. The TTL is only a memory bound
+ /// (keep the entry around at least as long as it could plausibly be warm); it does
+ /// not define warmth.
+ async fn put(&self, key: &str, binding: SessionBinding, ttl: Duration);
- /// Remove a cached routing decision by key.
+ /// Remove a session binding by key.
async fn remove(&self, key: &str);
}
diff --git a/crates/brightstaff/src/session_cache/redis.rs b/crates/brightstaff/src/session_cache/redis.rs
index 29630acc..d905965d 100644
--- a/crates/brightstaff/src/session_cache/redis.rs
+++ b/crates/brightstaff/src/session_cache/redis.rs
@@ -4,7 +4,7 @@ use async_trait::async_trait;
use redis::aio::MultiplexedConnection;
use redis::AsyncCommands;
-use super::{CachedRoute, SessionCache};
+use super::{SessionBinding, SessionCache};
const KEY_PREFIX: &str = "plano:affinity:";
@@ -26,18 +26,20 @@ impl RedisSessionCache {
#[async_trait]
impl SessionCache for RedisSessionCache {
- async fn get(&self, key: &str) -> Option {
+ async fn get(&self, key: &str) -> Option {
let mut conn = self.conn.clone();
let value: Option = conn.get(Self::make_key(key)).await.ok()?;
value.and_then(|v| serde_json::from_str(&v).ok())
}
- async fn put(&self, key: &str, route: CachedRoute, ttl: Duration) {
+ async fn put(&self, key: &str, binding: SessionBinding, ttl: Duration) {
let mut conn = self.conn.clone();
- let Ok(json) = serde_json::to_string(&route) else {
+ // The Redis TTL is only a GC bound; warmth is decided by the router from
+ // `binding.last_used`, not by expiry here.
+ let ttl_secs = ttl.as_secs().max(1);
+ let Ok(json) = serde_json::to_string(&binding) else {
return;
};
- let ttl_secs = ttl.as_secs().max(1);
let _: Result<(), _> = conn.set_ex(Self::make_key(key), json, ttl_secs).await;
}
diff --git a/crates/brightstaff/src/signals/otel.rs b/crates/brightstaff/src/signals/otel.rs
index deb3c1b5..176875cc 100644
--- a/crates/brightstaff/src/signals/otel.rs
+++ b/crates/brightstaff/src/signals/otel.rs
@@ -1,27 +1,21 @@
//! Helpers for emitting `SignalReport` data to OpenTelemetry spans.
//!
-//! Two sets of attributes are emitted:
-//!
-//! - **Legacy** keys under `signals.*` (e.g. `signals.frustration.count`),
-//! computed from the new layered counts. Preserved for one release for
-//! backward compatibility with existing dashboards.
-//! - **New** layered keys (e.g. `signals.interaction.misalignment.count`),
-//! one set of `count`/`severity` attributes per category, plus per-instance
-//! span events named `signal.`.
+//! Layered keys (e.g. `signals.interaction.misalignment.count`) are emitted,
+//! one set of `count`/`severity` attributes per category, plus per-instance
+//! span events named `signal.`.
use opentelemetry::trace::SpanRef;
use opentelemetry::KeyValue;
-use crate::signals::schemas::{SignalGroup, SignalReport, SignalType};
+use crate::signals::schemas::{SignalGroup, SignalReport};
-/// Emit both legacy and layered OTel attributes/events for a `SignalReport`.
+/// Emit layered OTel attributes/events for a `SignalReport`.
///
/// Returns `true` if any "concerning" signal was found, mirroring the previous
/// behavior used to flag the span operation name.
pub fn emit_signals_to_span(span: &SpanRef<'_>, report: &SignalReport) -> bool {
emit_overall(span, report);
emit_layered_attributes(span, report);
- emit_legacy_attributes(span, report);
emit_signal_events(span, report);
is_concerning(report)
@@ -90,69 +84,6 @@ fn emit_layered_attributes(span: &SpanRef<'_>, report: &SignalReport) {
);
}
-fn count_of(report: &SignalReport, t: SignalType) -> usize {
- report.iter_signals().filter(|s| s.signal_type == t).count()
-}
-
-/// Emit the legacy attribute keys consumed by existing dashboards. These are
-/// derived from the new `SignalReport` so no detector contract is broken.
-fn emit_legacy_attributes(span: &SpanRef<'_>, report: &SignalReport) {
- use crate::tracing::signals as legacy;
-
- // signals.follow_up.repair.{count,ratio} - misalignment proxies repairs.
- let repair_count = report.interaction.misalignment.count;
- let user_turns = report.turn_metrics.user_turns.max(1) as f32;
- if repair_count > 0 {
- span.set_attribute(KeyValue::new(legacy::REPAIR_COUNT, repair_count as i64));
- let ratio = repair_count as f32 / user_turns;
- span.set_attribute(KeyValue::new(legacy::REPAIR_RATIO, format!("{:.3}", ratio)));
- }
-
- // signals.frustration.{count,severity} - disengagement.negative_stance is
- // the closest legacy analog of "frustration".
- let frustration_count = count_of(report, SignalType::DisengagementNegativeStance);
- if frustration_count > 0 {
- span.set_attribute(KeyValue::new(
- legacy::FRUSTRATION_COUNT,
- frustration_count as i64,
- ));
- let severity = match frustration_count {
- 0 => 0,
- 1..=2 => 1,
- 3..=4 => 2,
- _ => 3,
- };
- span.set_attribute(KeyValue::new(legacy::FRUSTRATION_SEVERITY, severity as i64));
- }
-
- // signals.repetition.count - stagnation (repetition + dragging).
- if report.interaction.stagnation.count > 0 {
- span.set_attribute(KeyValue::new(
- legacy::REPETITION_COUNT,
- report.interaction.stagnation.count as i64,
- ));
- }
-
- // signals.escalation.requested - any escalation/quit signal.
- let escalated = report.interaction.disengagement.signals.iter().any(|s| {
- matches!(
- s.signal_type,
- SignalType::DisengagementEscalation | SignalType::DisengagementQuit
- )
- });
- if escalated {
- span.set_attribute(KeyValue::new(legacy::ESCALATION_REQUESTED, true));
- }
-
- // signals.positive_feedback.count - satisfaction signals.
- if report.interaction.satisfaction.count > 0 {
- span.set_attribute(KeyValue::new(
- legacy::POSITIVE_FEEDBACK_COUNT,
- report.interaction.satisfaction.count as i64,
- ));
- }
-}
-
fn emit_signal_events(span: &SpanRef<'_>, report: &SignalReport) {
for sig in report.iter_signals() {
let event_name = format!("signal.{}", sig.signal_type.as_str());
@@ -231,11 +162,4 @@ mod tests {
let r = report_with_escalation();
assert!(is_concerning(&r));
}
-
- #[test]
- fn count_of_returns_per_type_count() {
- let r = report_with_escalation();
- assert_eq!(count_of(&r, SignalType::DisengagementEscalation), 1);
- assert_eq!(count_of(&r, SignalType::DisengagementNegativeStance), 0);
- }
}
diff --git a/crates/brightstaff/src/streaming.rs b/crates/brightstaff/src/streaming.rs
index 26af8672..6cbf25c0 100644
--- a/crates/brightstaff/src/streaming.rs
+++ b/crates/brightstaff/src/streaming.rs
@@ -22,10 +22,15 @@ const STREAM_BUFFER_SIZE: usize = 16;
const USAGE_BUFFER_MAX: usize = 2 * 1024 * 1024;
use crate::metrics as bs_metrics;
use crate::metrics::labels as metric_labels;
+use crate::router::model_metrics::ModelRates;
+use crate::router::orchestrator::OrchestratorService;
+use crate::session_cache::{record_route_visit, RouteVisit, SessionBinding};
use crate::signals::otel::emit_signals_to_span;
use crate::signals::{SignalAnalyzer, FLAG_MARKER};
use crate::tracing::{llm, set_service_name};
use hermesllm::apis::openai::Message;
+use std::sync::Arc;
+use std::time::{Duration, SystemTime};
/// Parsed usage + resolved-model details from a provider response.
#[derive(Debug, Default, Clone)]
@@ -39,6 +44,11 @@ struct ExtractedUsage {
/// The model the upstream actually used. For router aliases (e.g.
/// `router:software-engineering`), this differs from the request model.
resolved_model: Option,
+ /// Provider convention for `prompt_tokens`: OpenAI-shape usage folds cached tokens
+ /// *into* `prompt_tokens` (so uncached = prompt - cached), while Anthropic-shape
+ /// reports uncached `input_tokens` separately from `cache_read`/`cache_creation`.
+ /// Set at parse time from which field `prompt_tokens` was sourced. Drives cost math.
+ prompt_includes_cached: bool,
}
impl ExtractedUsage {
@@ -57,8 +67,9 @@ impl ExtractedUsage {
}
}
if let Some(u) = value.get("usage") {
- // OpenAI-shape usage
+ // OpenAI-shape usage: `prompt_tokens` includes the cached subset.
out.prompt_tokens = u.get("prompt_tokens").and_then(|v| v.as_i64());
+ out.prompt_includes_cached = out.prompt_tokens.is_some();
out.completion_tokens = u.get("completion_tokens").and_then(|v| v.as_i64());
out.total_tokens = u.get("total_tokens").and_then(|v| v.as_i64());
out.cached_input_tokens = u
@@ -187,6 +198,44 @@ pub struct LlmMetricsCtx {
pub upstream_status: u16,
}
+/// Response-side session-update context: refreshes the session binding once the real
+/// response is in hand, so `last_used` and the context-size estimate (`cached_tokens`)
+/// reflect the turn that just completed. The routing decision (model, budget, switches)
+/// was already made and persisted on the request side by [`super::handlers::llm::session_router`];
+/// this only refines the fields that need the response.
+pub struct SessionUpdateCtx {
+ pub orchestrator: Arc,
+ pub session_id: String,
+ pub tenant_id: Option,
+ /// Provider-qualified model this request actually ran on (the router's final pick).
+ pub anchor_model: String,
+ /// The session's never-switch model for this episode — preserved across the refresh.
+ pub default_model: String,
+ pub route_name: Option,
+ pub prefix_hash: Option,
+ /// Cumulative never-switch baseline from the decision — preserved across the refresh.
+ pub baseline_usd: f64,
+ /// Cumulative switch spend from the decision — preserved across the refresh.
+ pub switch_spend_usd: f64,
+ /// Cumulative switch count from the routing decision — preserved across the refresh.
+ pub switches: u32,
+ /// Per-model route history from the decision — preserved across the refresh, with
+ /// the anchor's entry refined to the real prompt-token count.
+ pub history: Vec,
+ /// Cumulative actual conversation cost (USD) through prior turns. This turn's real
+ /// cost is added on top once usage is known.
+ pub session_cost_usd: f64,
+ /// Catalog rates for the dispatched model, resolved request-side (the response path
+ /// is synchronous). `None` when no cost source is configured → no cost is computed.
+ pub cost_rates: Option,
+ /// Cached-read discount used to price cached input when the feed omits a cached rate.
+ pub cache_read_discount: f64,
+ /// Context-size count to fall back to when the response carries no usage block.
+ pub context_tokens: u64,
+ /// GC bound to store the refreshed binding with.
+ pub gc_ttl: Duration,
+}
+
/// A processor that tracks streaming metrics
pub struct ObservableStreamProcessor {
service_name: String,
@@ -201,6 +250,7 @@ pub struct ObservableStreamProcessor {
/// from the buffer (they still pass through to the client).
response_buffer: Vec,
llm_metrics: Option,
+ session_update: Option,
metrics_recorded: bool,
}
@@ -237,6 +287,7 @@ impl ObservableStreamProcessor {
messages,
response_buffer: Vec::new(),
llm_metrics: None,
+ session_update: None,
metrics_recorded: false,
}
}
@@ -247,6 +298,103 @@ impl ObservableStreamProcessor {
self.llm_metrics = Some(ctx);
self
}
+
+ /// Attach session-update context so the processor refreshes `last_used` and the
+ /// context-size estimate from the real response once usage is known.
+ pub fn with_session_update(mut self, ctx: SessionUpdateCtx) -> Self {
+ self.session_update = Some(ctx);
+ self
+ }
+
+ /// Refresh the session binding from the response so warmth (`last_used`) and the
+ /// context-size estimate (`cached_tokens`) reflect the completed turn.
+ fn handle_session_update(&mut self, usage: &ExtractedUsage) {
+ let Some(update) = self.session_update.take() else {
+ return;
+ };
+ let SessionUpdateCtx {
+ orchestrator,
+ session_id,
+ tenant_id,
+ anchor_model,
+ default_model,
+ route_name,
+ prefix_hash,
+ baseline_usd,
+ switch_spend_usd,
+ switches,
+ mut history,
+ session_cost_usd,
+ cost_rates,
+ cache_read_discount,
+ context_tokens,
+ gc_ttl,
+ } = update;
+
+ // Prefer the real prompt-token count (the tokens a future switch would re-read)
+ // over the request-side estimate; fall back to the estimate when absent.
+ let cached_tokens = usage
+ .prompt_tokens
+ .filter(|&p| p > 0)
+ .map(|p| p as u64)
+ .unwrap_or(context_tokens);
+
+ // Refine the anchor's route-history entry with the real context size, so a later
+ // return to this model prices its warm portion off actual usage.
+ record_route_visit(
+ &mut history,
+ &anchor_model,
+ SystemTime::now(),
+ cached_tokens,
+ );
+
+ // Price this turn from the catalog rates and roll it into the conversation
+ // total. Emit per-request cost on the (llm) span and carry the running total
+ // into the binding so the next turn's routing span can surface it.
+ let session_cost_usd = if let Some(rates) = cost_rates {
+ let (input_cost, output_cost) = rates.request_cost_usd(
+ usage.prompt_tokens.unwrap_or(0).max(0) as u64,
+ usage.cached_input_tokens.unwrap_or(0).max(0) as u64,
+ usage.cache_creation_tokens.unwrap_or(0).max(0) as u64,
+ usage.completion_tokens.unwrap_or(0).max(0) as u64,
+ usage.prompt_includes_cached,
+ cache_read_discount,
+ );
+ let span = tracing::Span::current();
+ let otel_span = span.context();
+ let otel_span = otel_span.span();
+ otel_span.set_attribute(KeyValue::new(llm::INPUT_COST_IN_USD, input_cost));
+ otel_span.set_attribute(KeyValue::new(llm::OUTPUT_COST_IN_USD, output_cost));
+ otel_span.set_attribute(KeyValue::new(
+ llm::TOTAL_COST_IN_USD,
+ input_cost + output_cost,
+ ));
+ session_cost_usd + input_cost + output_cost
+ } else {
+ session_cost_usd
+ };
+
+ bs_metrics::record_session_binding_event(metric_labels::BINDING_EVENT_REFRESH);
+ let binding = SessionBinding {
+ anchor_model,
+ default_model,
+ route_name,
+ prefix_hash,
+ last_used: SystemTime::now(),
+ cached_tokens,
+ baseline_usd,
+ switch_spend_usd,
+ switches,
+ session_cost_usd,
+ history,
+ };
+ // Fire-and-forget: binding bookkeeping must not delay stream completion.
+ tokio::spawn(async move {
+ orchestrator
+ .store_binding(&session_id, tenant_id.as_deref(), binding, Some(gc_ttl))
+ .await;
+ });
+ }
}
impl StreamProcessor for ObservableStreamProcessor {
@@ -357,19 +505,53 @@ impl StreamProcessor for ObservableStreamProcessor {
v.max(0) as u64,
);
}
+ // Prompt-cache token counters + hit/miss baseline (cache-blindness is
+ // invisible without these: a reroute that burns a warm cache shows up
+ // here as a miss with zero cache_read tokens).
+ if let Some(v) = usage.cached_input_tokens {
+ bs_metrics::record_llm_tokens(
+ &ctx.provider,
+ &ctx.model,
+ metric_labels::TOKEN_KIND_CACHE_READ,
+ v.max(0) as u64,
+ );
+ }
+ if let Some(v) = usage.cache_creation_tokens {
+ bs_metrics::record_llm_tokens(
+ &ctx.provider,
+ &ctx.model,
+ metric_labels::TOKEN_KIND_CACHE_WRITE,
+ v.max(0) as u64,
+ );
+ }
+ if usage.prompt_tokens.is_some() {
+ let cache_active = usage.cached_input_tokens.unwrap_or(0) > 0
+ || usage.cache_creation_tokens.unwrap_or(0) > 0;
+ bs_metrics::record_prompt_cache_outcome(
+ &ctx.provider,
+ &ctx.model,
+ if cache_active {
+ metric_labels::PROMPT_CACHE_HIT
+ } else {
+ metric_labels::PROMPT_CACHE_MISS
+ },
+ );
+ }
if usage.prompt_tokens.is_none() && usage.completion_tokens.is_none() {
bs_metrics::record_llm_tokens_usage_missing(&ctx.provider, &ctx.model);
}
self.metrics_recorded = true;
}
+
+ // Session-binding refresh: update `last_used` + the context-size estimate from
+ // the completed turn (the routing decision itself was made request-side).
+ self.handle_session_update(&usage);
// Release the buffered bytes early; nothing downstream needs them.
self.response_buffer.clear();
self.response_buffer.shrink_to_fit();
// Analyze signals if messages are available and record as span
- // attributes + per-signal events. We dual-emit legacy aggregate keys
- // and the new layered taxonomy so existing dashboards keep working
- // while new consumers can opt into the richer hierarchy.
+ // attributes + per-signal events using the layered signal taxonomy.
if let Some(ref messages) = self.messages {
let analyzer = SignalAnalyzer::default();
let report = analyzer.analyze_openai(messages);
@@ -614,6 +796,8 @@ mod usage_extraction_tests {
assert_eq!(u.total_tokens, Some(46));
assert_eq!(u.cached_input_tokens, Some(5));
assert_eq!(u.reasoning_tokens, None);
+ // OpenAI folds cached tokens into `prompt_tokens`.
+ assert!(u.prompt_includes_cached);
}
#[test]
@@ -625,6 +809,8 @@ mod usage_extraction_tests {
assert_eq!(u.total_tokens, Some(150));
assert_eq!(u.cached_input_tokens, Some(30));
assert_eq!(u.cache_creation_tokens, Some(20));
+ // Anthropic reports uncached `input_tokens` separately from cached reads.
+ assert!(!u.prompt_includes_cached);
}
#[test]
diff --git a/crates/brightstaff/src/tracing/constants.rs b/crates/brightstaff/src/tracing/constants.rs
index 79a40401..f5a3a0c0 100644
--- a/crates/brightstaff/src/tracing/constants.rs
+++ b/crates/brightstaff/src/tracing/constants.rs
@@ -92,6 +92,18 @@ pub mod llm {
/// (OpenAI `completion_tokens_details.reasoning_tokens`, Google `thoughts_token_count`)
pub const REASONING_TOKENS: &str = "llm.usage.reasoning_tokens";
+ /// This request's input-token cost (USD), priced from the catalog rates: uncached
+ /// input at the input rate, cached reads at the cached rate, cache creation at the
+ /// plain input rate. Present only when a cost source is configured.
+ pub const INPUT_COST_IN_USD: &str = "llm.usage.input_cost_usd";
+
+ /// This request's output-token cost (USD) = completion tokens x output rate.
+ pub const OUTPUT_COST_IN_USD: &str = "llm.usage.output_cost_usd";
+
+ /// This request's total cost (USD) = input + output. Sum across a session's turns
+ /// (group by `plano.session_id`) for the conversation total.
+ pub const TOTAL_COST_IN_USD: &str = "llm.usage.total_cost_usd";
+
/// Temperature parameter used
pub const TEMPERATURE: &str = "llm.temperature";
@@ -145,6 +157,66 @@ pub mod plano {
/// "software-engineering"). Absent when the client routed directly
/// to a concrete model.
pub const ROUTE_NAME: &str = "plano.route.name";
+
+ /// Caller identity used to populate downstream observability `distinct_id`
+ /// fields (e.g. PostHog). Sourced from the configured
+ /// `tracing.exporters[].distinct_id_header`. Absent for anonymous calls.
+ pub const DISTINCT_ID: &str = "plano.distinct_id";
+
+ /// Whether the session's provider cache was inferred warm at decision time
+ /// (from the idle gap vs. the provider's cache window).
+ pub const CACHE_WARM: &str = "plano.cache.warm";
+
+ /// How long (ms) since the session was last used — the idle gap warmth is measured
+ /// against.
+ pub const CACHE_IDLE_MS: &str = "plano.cache.idle_ms";
+
+ /// Cumulative switching overhead consumed this session, as a percentage of the
+ /// never-switch baseline (`100 * switch_spend / baseline`). Directly comparable to
+ /// the configured `routing.routing_budget.max_overhead_pct`.
+ pub const SESSION_OVERHEAD_PCT: &str = "plano.session.overhead_pct";
+
+ /// Cumulative overhead (USD) actually spent on paid switches this session — the
+ /// numerator behind `plano.session.overhead_pct`.
+ pub const SESSION_SWITCH_SPEND_IN_USD: &str = "plano.session.switch_spend_in_usd";
+
+ /// Cumulative never-switch baseline (USD) — what staying on the anchor would have
+ /// cost so far. The denominator behind `plano.session.overhead_pct`.
+ pub const SESSION_BASELINE_IN_USD: &str = "plano.session.baseline_in_usd";
+
+ /// Cumulative number of model switches taken during this warm session.
+ pub const SESSION_SWITCHES: &str = "plano.session.switches";
+
+ /// Cumulative *actual* cost (USD, input + output) of the whole conversation, priced
+ /// from the configured catalog rates and refined from real usage each turn. Emitted
+ /// on the routing span; reflects cost through the previous turn (this turn isn't
+ /// billed yet at decision time). Full-proxy path only.
+ pub const SESSION_TOTAL_COST_IN_USD: &str = "plano.session.total_cost_in_usd";
+
+ /// Actual input-cost (USD) of the proposed model switch — computed from input-token
+ /// pricing only, output-token cost deliberately excluded. Negative when the candidate
+ /// is outright cheaper than staying on the warm anchor.
+ pub const SWITCH_COST_IN_USD: &str = "plano.switch.cost_in_usd";
+
+ /// Tokens the switch candidate still has cached from an earlier visit this session
+ /// (a return to a still-warm model). These re-read at the candidate's cached rate
+ /// instead of its uncached rate, which is why the switch cost can be far below a
+ /// full re-ingest. Zero for a first-time (cold) switch.
+ pub const SWITCH_CANDIDATE_WARM_TOKENS: &str = "plano.switch.candidate_warm_tokens";
+
+ /// The overhead ceiling (USD) available when the switch was evaluated —
+ /// `max_overhead_pct% * baseline`. A paid switch is allowed while cumulative spend
+ /// plus this switch's cost stays under it. Directly comparable to `cost_in_usd`.
+ pub const SWITCH_OVERHEAD_CEILING_IN_USD: &str = "plano.switch.overhead_ceiling_in_usd";
+
+ /// Switch outcome: "allowed" or "retained".
+ pub const SWITCH_DECISION: &str = "plano.switch.decision";
+
+ /// The route (`provider/model`, plus route name when routed) the routing-budget gate
+ /// *would* have selected had the switch been allowed. Recorded only on a `retained`
+ /// decision when `routing.routing_budget.record_counterfactual` is enabled.
+ /// Telemetry only — the counterfactual model is never dispatched.
+ pub const SWITCH_COUNTERFACTUAL_ROUTE: &str = "plano.switch.counterfactual_route";
}
// =============================================================================
@@ -183,27 +255,6 @@ pub mod signals {
/// Efficiency score (0.0-1.0)
pub const EFFICIENCY_SCORE: &str = "signals.efficiency_score";
-
- /// Number of repair attempts detected
- pub const REPAIR_COUNT: &str = "signals.follow_up.repair.count";
-
- /// Ratio of repairs to user turns
- pub const REPAIR_RATIO: &str = "signals.follow_up.repair.ratio";
-
- /// Number of frustration indicators detected
- pub const FRUSTRATION_COUNT: &str = "signals.frustration.count";
-
- /// Frustration severity level (0-3)
- pub const FRUSTRATION_SEVERITY: &str = "signals.frustration.severity";
-
- /// Number of repetition instances detected
- pub const REPETITION_COUNT: &str = "signals.repetition.count";
-
- /// Whether escalation was requested (user asked for human help)
- pub const ESCALATION_REQUESTED: &str = "signals.escalation.requested";
-
- /// Number of positive feedback indicators detected
- pub const POSITIVE_FEEDBACK_COUNT: &str = "signals.positive_feedback.count";
}
// =============================================================================
diff --git a/crates/brightstaff/src/tracing/init.rs b/crates/brightstaff/src/tracing/init.rs
index ed351148..b9560423 100644
--- a/crates/brightstaff/src/tracing/init.rs
+++ b/crates/brightstaff/src/tracing/init.rs
@@ -11,8 +11,8 @@ use tracing_subscriber::registry::LookupSpan;
use tracing_subscriber::util::SubscriberInitExt;
use tracing_subscriber::EnvFilter;
-use super::ServiceNameOverrideExporter;
-use common::configuration::Tracing;
+use super::{PostHogExporter, ServiceNameOverrideExporter};
+use common::configuration::{Exporter, PosthogExporter, Tracing};
struct BracketedTime;
@@ -90,26 +90,53 @@ pub fn init_tracer(tracing_config: Option<&Tracing>) -> &'static SdkTracerProvid
let random_sampling = tracing_config.and_then(|t| t.random_sampling).unwrap_or(0);
- let tracing_enabled = random_sampling > 0 && otel_endpoint.is_some();
+ // Collect PostHog export destinations from `tracing.exporters`.
+ let posthog_exporters: Vec = tracing_config
+ .and_then(|t| t.exporters.as_ref())
+ .map(|exporters| {
+ exporters
+ .iter()
+ .map(|Exporter::Posthog(posthog)| posthog.clone())
+ .collect()
+ })
+ .unwrap_or_default();
+
+ // Tracing is enabled when sampling is on and there is at least one
+ // destination — an OTLP collector and/or a configured exporter.
+ let has_destination = otel_endpoint.is_some() || !posthog_exporters.is_empty();
+ let tracing_enabled = random_sampling > 0 && has_destination;
eprintln!(
- "initializing tracing: tracing_enabled={}, otel_endpoint={:?}, random_sampling={}",
- tracing_enabled, otel_endpoint, random_sampling
+ "initializing tracing: tracing_enabled={}, otel_endpoint={:?}, random_sampling={}, posthog_exporters={}",
+ tracing_enabled, otel_endpoint, random_sampling, posthog_exporters.len()
);
- // Create OTLP exporter to send spans to collector.
- // Use `if let` to destructure the endpoint, avoiding an unwrap.
- if let Some(endpoint) = otel_endpoint.as_deref().filter(|_| tracing_enabled) {
+ if tracing_enabled {
if std::env::var("OTEL_SERVICE_NAME").is_err() {
std::env::set_var("OTEL_SERVICE_NAME", "plano");
}
+
+ // Compose the tracer provider from all configured destinations. Each
+ // `with_batch_exporter` registers an independent span processor, so
+ // every span fans out to the OTLP collector and every exporter.
+ let mut builder = SdkTracerProvider::builder();
+
// Create ServiceNameOverrideExporter to support per-span service names
// This allows spans to have different service names (e.g., plano(orchestrator),
// plano(filter), plano(llm)) by setting the "service.name.override" attribute
- let exporter = ServiceNameOverrideExporter::new(endpoint);
+ if let Some(endpoint) = otel_endpoint.as_deref() {
+ builder = builder.with_batch_exporter(ServiceNameOverrideExporter::new(endpoint));
+ }
- let provider = SdkTracerProvider::builder()
- .with_batch_exporter(exporter)
- .build();
+ // PostHog exporters translate LLM spans into `$ai_generation` events.
+ for posthog in &posthog_exporters {
+ builder = builder.with_batch_exporter(PostHogExporter::new(
+ &posthog.url,
+ &posthog.api_key,
+ posthog.capture_messages.unwrap_or(false),
+ ));
+ }
+
+ let provider = builder.build();
global::set_tracer_provider(provider.clone());
diff --git a/crates/brightstaff/src/tracing/mod.rs b/crates/brightstaff/src/tracing/mod.rs
index 8e09a21c..dac26232 100644
--- a/crates/brightstaff/src/tracing/mod.rs
+++ b/crates/brightstaff/src/tracing/mod.rs
@@ -1,6 +1,7 @@
mod constants;
mod custom_attributes;
mod init;
+mod posthog_exporter;
mod service_name_exporter;
pub use constants::{
@@ -8,6 +9,7 @@ pub use constants::{
};
pub use custom_attributes::collect_custom_trace_attributes;
pub use init::init_tracer;
+pub use posthog_exporter::PostHogExporter;
pub use service_name_exporter::{ServiceNameOverrideExporter, SERVICE_NAME_OVERRIDE_KEY};
use opentelemetry::trace::get_active_span;
diff --git a/crates/brightstaff/src/tracing/posthog_exporter.rs b/crates/brightstaff/src/tracing/posthog_exporter.rs
new file mode 100644
index 00000000..53d3ccef
--- /dev/null
+++ b/crates/brightstaff/src/tracing/posthog_exporter.rs
@@ -0,0 +1,402 @@
+//! PostHog Span Exporter
+//!
+//! A custom [`SpanExporter`] that translates Plano's LLM spans into PostHog
+//! [`$ai_generation`](https://posthog.com/docs/ai-observability/generations)
+//! events and POSTs them to PostHog's capture API (`{url}/batch/`).
+//!
+//! This makes PostHog a first-class, provider-agnostic export target: a user
+//! only points `tracing.exporters` at their PostHog URL + project token and
+//! every LLM call is captured — mirroring LiteLLM's `posthog` callback.
+//!
+//! # Behaviour
+//!
+//! - Receives every span in the provider (like all batch exporters do) and
+//! keeps only LLM generation spans, identified by the presence of the
+//! [`llm::MODEL_NAME`] (`llm.model`) attribute.
+//! - Maps span attributes onto `$ai_*` PostHog properties (model, provider,
+//! latency, tokens, http status, ...).
+//! - `distinct_id` is read from the [`plano::DISTINCT_ID`] span attribute (set
+//! by the LLM handler from the configured `distinct_id_header`). When absent
+//! the event is captured anonymously (`$process_person_profile = false`).
+//! - Network failures are logged and dropped — telemetry export never blocks or
+//! fails request processing.
+
+use std::time::Duration;
+
+use opentelemetry::{Array, Value};
+use opentelemetry_sdk::error::OTelSdkResult;
+use opentelemetry_sdk::trace::{SpanData, SpanExporter};
+use opentelemetry_sdk::Resource;
+use serde_json::{json, Map, Value as JsonValue};
+use time::format_description::well_known::Rfc3339;
+use time::OffsetDateTime;
+
+use super::{http, llm, plano};
+
+/// PostHog event name for an individual LLM call.
+const AI_GENERATION_EVENT: &str = "$ai_generation";
+
+/// PostHog capture path appended to the configured host.
+const CAPTURE_PATH: &str = "batch/";
+
+/// A [`SpanExporter`] that ships LLM spans to PostHog as `$ai_generation` events.
+pub struct PostHogExporter {
+ client: reqwest::Client,
+ /// Fully-qualified capture endpoint, e.g. `https://us.i.posthog.com/batch/`.
+ endpoint: String,
+ /// PostHog project API key (token).
+ api_key: String,
+ /// Whether to attach the truncated user message preview as `$ai_input`.
+ capture_messages: bool,
+}
+
+impl std::fmt::Debug for PostHogExporter {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ f.debug_struct("PostHogExporter")
+ .field("endpoint", &self.endpoint)
+ .field("capture_messages", &self.capture_messages)
+ .finish()
+ }
+}
+
+impl PostHogExporter {
+ /// Create a new PostHog exporter.
+ ///
+ /// # Arguments
+ /// * `url` – PostHog host (e.g. `https://us.i.posthog.com`). The `/batch/`
+ /// capture path is appended automatically.
+ /// * `api_key` – PostHog project API key (token).
+ /// * `capture_messages` – when true, send the user message preview as
+ /// `$ai_input`.
+ pub fn new(url: &str, api_key: &str, capture_messages: bool) -> Self {
+ let endpoint = format!("{}/{}", url.trim_end_matches('/'), CAPTURE_PATH);
+ let client = reqwest::Client::builder()
+ .timeout(Duration::from_secs(10))
+ .build()
+ .unwrap_or_default();
+ Self {
+ client,
+ endpoint,
+ api_key: api_key.to_string(),
+ capture_messages,
+ }
+ }
+
+ /// Build the PostHog `batch` payload from a batch of spans, keeping only LLM
+ /// generation spans. Returns `None` when no LLM spans are present.
+ fn build_payload(&self, batch: &[SpanData]) -> Option {
+ let events: Vec = batch
+ .iter()
+ .filter_map(|span| self.build_generation_event(span))
+ .collect();
+
+ if events.is_empty() {
+ return None;
+ }
+
+ Some(json!({
+ "api_key": self.api_key,
+ "batch": events,
+ }))
+ }
+
+ /// Translate a single span into a PostHog `$ai_generation` event, or `None`
+ /// if the span is not an LLM generation span.
+ fn build_generation_event(&self, span: &SpanData) -> Option {
+ // Only LLM generation spans carry `llm.model`.
+ let model = find_attr(span, llm::MODEL_NAME)?;
+
+ let mut props = Map::new();
+ props.insert("$ai_model".to_string(), otel_value_to_json(model));
+ props.insert(
+ "$ai_trace_id".to_string(),
+ json!(span.span_context.trace_id().to_string()),
+ );
+ if span.parent_span_id != opentelemetry::trace::SpanId::INVALID {
+ props.insert(
+ "$ai_parent_id".to_string(),
+ json!(span.parent_span_id.to_string()),
+ );
+ }
+
+ if let Some(provider) = find_attr(span, llm::PROVIDER) {
+ props.insert("$ai_provider".to_string(), otel_value_to_json(provider));
+ }
+
+ // Latency / TTFT are stored in milliseconds; PostHog wants seconds.
+ if let Some(ms) = find_i64(span, llm::DURATION_MS) {
+ props.insert("$ai_latency".to_string(), json!(ms as f64 / 1000.0));
+ }
+ if let Some(ms) = find_i64(span, llm::TIME_TO_FIRST_TOKEN_MS) {
+ props.insert(
+ "$ai_time_to_first_token".to_string(),
+ json!(ms as f64 / 1000.0),
+ );
+ props.insert("$ai_stream".to_string(), json!(true));
+ }
+
+ if let Some(tokens) = find_i64(span, llm::PROMPT_TOKENS) {
+ props.insert("$ai_input_tokens".to_string(), json!(tokens));
+ }
+ if let Some(tokens) = find_i64(span, llm::COMPLETION_TOKENS) {
+ props.insert("$ai_output_tokens".to_string(), json!(tokens));
+ }
+
+ if let Some(status) = find_i64(span, http::STATUS_CODE) {
+ props.insert("$ai_http_status".to_string(), json!(status));
+ if status >= 400 {
+ props.insert("$ai_is_error".to_string(), json!(true));
+ }
+ }
+
+ if self.capture_messages {
+ if let Some(preview) = find_attr(span, llm::USER_MESSAGE_PREVIEW) {
+ props.insert(
+ "$ai_input".to_string(),
+ json!([{ "role": "user", "content": value_to_string(preview) }]),
+ );
+ }
+ }
+
+ // distinct_id: identified when the configured header was present,
+ // otherwise anonymous (do not create/update a person profile).
+ match find_attr(span, plano::DISTINCT_ID) {
+ Some(id) => {
+ props.insert("distinct_id".to_string(), otel_value_to_json(id));
+ }
+ None => {
+ props.insert(
+ "distinct_id".to_string(),
+ json!(span.span_context.trace_id().to_string()),
+ );
+ props.insert("$process_person_profile".to_string(), json!(false));
+ }
+ }
+
+ // Pass through any other non-reserved attributes (custom span attributes
+ // such as static tags or header-derived tenant ids) as plain properties.
+ for kv in span.attributes.iter() {
+ let key = kv.key.as_str();
+ if is_reserved_attr(key) {
+ continue;
+ }
+ props
+ .entry(key.to_string())
+ .or_insert_with(|| otel_value_to_json(&kv.value));
+ }
+
+ let mut event = Map::new();
+ event.insert("event".to_string(), json!(AI_GENERATION_EVENT));
+ event.insert("properties".to_string(), JsonValue::Object(props));
+ if let Ok(ts) = OffsetDateTime::from(span.end_time).format(&Rfc3339) {
+ event.insert("timestamp".to_string(), json!(ts));
+ }
+
+ Some(JsonValue::Object(event))
+ }
+}
+
+impl SpanExporter for PostHogExporter {
+ fn export(
+ &self,
+ batch: Vec,
+ ) -> impl std::future::Future