-# Smart AI model cascading for cost optimization
+# Agent Runtime Intelligence Layer
[](https://pypi.org/project/cascadeflow/)
[](https://www.npmjs.com/package/@cascadeflow/core)
@@ -17,10 +17,11 @@
[](https://pepy.tech/project/cascadeflow)
[](https://www.npmjs.com/search?q=%40cascadeflow)
[](https://github.com/lemony-ai/cascadeflow/actions/workflows/test.yml)
+[](https://docs.cascadeflow.dev)
[](./docs/)
[](./docs/)
[](https://x.com/saschabuehrle)
-[](https://github.com/lemony-ai/cascadeflow)
+[](https://github.com/lemony-ai/cascadeflow/stargazers)
Python](#-python) • [
TypeScript](#-typescript) • [
n8n](#-n8n-integration) • [
OpenClaw](https://clawhub.ai/saschabuehrle/cascadeflow) • [📖 Docs](./docs/) • [💡 Examples](#examples)**
+**[ Python](#-python) • [ TypeScript](#-typescript) • [ n8n](#-n8n-integration) • [ OpenClaw](https://clawhub.ai/saschabuehrle/cascadeflow) • [Full Docs](https://docs.cascadeflow.dev) • [📖 Docs](./docs/) • [💡 Examples](#examples)**
---
-**Stop Bleeding Money on AI Calls. Cut Costs 30-65% in 3 Lines of Code.**
+**The in-process intelligence layer for AI agents.** Optimize cost, latency, quality, budget, compliance, and energy — inside the execution loop, not at the HTTP boundary.
-40-70% of text prompts and 20-60% of agent calls don't need expensive flagship models. You're overpaying every single day.
-
-*cascadeflow fixes this with intelligent model cascading, available in Python and TypeScript.*
+cascadeflow works where external proxies can't: per-step model decisions based on agent state, per-tool-call budget gating, runtime stop/continue/escalate actions, and business KPI injection during agent loops. Sub-1ms overhead. Works with LangChain, OpenAI Agents SDK, CrewAI, Google ADK, n8n, and Vercel AI SDK.
```python
pip install cascadeflow
@@ -52,6 +51,17 @@ npm install @cascadeflow/core
## Why cascadeflow?
+### Proxy vs In-Process Harness
+
+| Dimension | External Proxy | cascadeflow Harness |
+|---|---|---|
+| **Scope** | HTTP request boundary | Inside agent execution loop |
+| **Dimensions** | Cost only | Cost + quality + latency + budget + compliance + energy |
+| **Latency overhead** | 10-50ms network RTT | <1ms in-process |
+| **Business logic** | None | KPI weights and targets |
+| **Enforcement** | None (observe only) | stop, deny_tool, switch_model |
+| **Auditability** | Request logs | Per-step decision traces |
+
cascadeflow is an intelligent AI model cascading library that dynamically selects the optimal model for each query or tool call through speculative execution. It's based on the research that 40-70% of queries don't require slow, expensive flagship models, and domain-specific smaller models often outperform large general-purpose models on specialized tasks. For the remaining queries that need advanced reasoning, cascadeflow automatically escalates to flagship models if needed.
### Use Cases
@@ -140,6 +150,34 @@ In practice, 60-70% of queries are handled by small, efficient models (8-20x cos
---
+## Harness API
+
+Three tiers of integration — zero-change observability to full policy control:
+
+**Tier 1: Zero-change observability**
+```python
+import cascadeflow
+cascadeflow.init(mode="observe")
+# All OpenAI/Anthropic SDK calls are now tracked. No code changes needed.
+```
+
+**Tier 2: Scoped runs with budget**
+```python
+with cascadeflow.run(budget=0.50, max_tool_calls=10) as session:
+ result = await agent.run("Analyze this dataset")
+ print(session.summary()) # cost, latency, energy, steps, tool calls
+ print(session.trace()) # full decision audit trail
+```
+
+**Tier 3: Decorated agents with policy**
+```python
+@cascadeflow.agent(budget=0.20, compliance="gdpr", kpi_weights={"quality": 0.6, "cost": 0.3, "latency": 0.1})
+async def my_agent(query: str):
+ return await llm.complete(query)
+```
+
+---
+
## Quick Start
### Drop-In Gateway (Existing Apps)
@@ -724,6 +762,12 @@ console.log(`Warnings: ${validation.warnings}`);
| 📋 **Message & Tool Call Lists** | Full conversation history with tool_calls and tool_call_id preservation across turns |
| 🪝 **Hooks & Callbacks** | Telemetry callbacks, cost events, and streaming hooks for observability |
| 🏭 **Production Ready** | Streaming, batch processing, tool handling, reasoning model support, caching, error recovery, anomaly detection |
+| 💳 **Budget Enforcement** | Per-run and per-user budget caps with automatic stop actions when limits are exceeded |
+| 🔒 **Compliance Gating** | GDPR, HIPAA, PCI, and strict model allowlists — block non-compliant models before execution |
+| 📊 **KPI-Weighted Routing** | Inject business priorities (quality, cost, latency, energy) as weights into every model decision |
+| 🌱 **Energy Tracking** | Deterministic compute-intensity coefficients for carbon-aware AI operations |
+| 🔍 **Decision Traces** | Full per-step audit trail: action, reason, model, cost, budget state, enforcement status |
+| ⚙️ **Harness Modes** | off / observe / enforce — roll out safely with observe, then switch to enforce when ready |
---
@@ -774,7 +818,7 @@ If you use cascadeflow in your research or project, please cite:
```bibtex
@software{cascadeflow2025,
author = {Lemony Inc., Sascha Buehrle and Contributors},
- title = {cascadeflow: Smart AI model cascading for cost optimization},
+ title = {cascadeflow: Agent runtime intelligence layer for AI agent workflows},
year = {2025},
publisher = {GitHub},
url = {https://github.com/lemony-ai/cascadeflow}
diff --git a/cascadeflow/__init__.py b/cascadeflow/__init__.py
index 1b61a9f3..af4c429a 100644
--- a/cascadeflow/__init__.py
+++ b/cascadeflow/__init__.py
@@ -1,30 +1,23 @@
"""
-cascadeflow - Smart AI model cascading for cost optimization.
-
-Route queries intelligently across multiple AI models from tiny SLMs
-to frontier LLMs based on complexity, domain, and budget.
-
-Features:
-- 🚀 Speculative cascades (2-3x faster)
-- 💰 60-95% cost savings
-- 🎯 Per-prompt domain detection
-- 🎨 2.0x domain boost for specialists
-- 🔍 Multi-factor optimization
-- 🆓 Free tier (Ollama + Groq)
-- ⚡ 3 lines of code
-
-Example:
- >>> from cascadeflow import CascadeAgent, CascadePresets
- >>>
- >>> # Auto-detect available models
- >>> models = CascadePresets.auto_detect_models()
- >>>
- >>> # Create agent with intelligence layer
- >>> agent = CascadeAgent(models, enable_caching=True)
- >>>
- >>> # Run query (automatically optimized!)
- >>> result = await agent.run("Fix this Python bug")
- >>> print(f"Used {result.model_used} - Cost: ${result.cost:.6f}")
+cascadeflow - Agent runtime intelligence layer.
+
+In-process harness that optimizes cost, latency, quality, budget, compliance,
+and energy across AI agent workflows. Works inside agent execution loops with
+full state awareness -- not an external proxy.
+
+Quick start:
+ import cascadeflow
+ cascadeflow.init(mode="observe")
+ # All OpenAI/Anthropic SDK calls are now tracked and traced.
+
+Key APIs:
+ cascadeflow.init(mode) -- activate harness (off | observe | enforce)
+ cascadeflow.run(budget) -- scoped run with budget/trace
+ @cascadeflow.agent(budget) -- policy metadata on agent functions
+ session.summary() -- structured metrics
+ session.trace() -- full decision audit trail
+
+Integrations: LangChain, OpenAI Agents SDK, CrewAI, Google ADK, n8n, Vercel AI SDK
"""
__version__ = "1.0.0"
@@ -240,6 +233,10 @@
)
# NEW: Harness API scaffold (V2 core branch)
+# NOTE: harness.agent is NOT re-exported here — it would shadow the
+# cascadeflow.agent *module* and break dotted-path resolution
+# (e.g. patch("cascadeflow.agent.PROVIDER_REGISTRY")).
+# Use ``from cascadeflow.harness import agent`` instead.
from .harness import (
HarnessConfig,
HarnessInitReport,
@@ -247,7 +244,6 @@
init,
reset,
run,
- agent as harness_agent,
get_harness_config,
get_current_run,
)
@@ -401,7 +397,6 @@
"init",
"reset",
"run",
- "harness_agent",
"get_harness_config",
"get_current_run",
# ===== PROVIDERS =====
diff --git a/cascadeflow/harness/__init__.py b/cascadeflow/harness/__init__.py
index 43a03662..74c07219 100644
--- a/cascadeflow/harness/__init__.py
+++ b/cascadeflow/harness/__init__.py
@@ -14,11 +14,13 @@
HarnessInitReport,
HarnessRunContext,
agent,
+ get_harness_callback_manager,
get_current_run,
get_harness_config,
init,
reset,
run,
+ set_harness_callback_manager,
)
__all__ = [
@@ -29,6 +31,8 @@
"run",
"agent",
"get_current_run",
+ "get_harness_callback_manager",
"get_harness_config",
+ "set_harness_callback_manager",
"reset",
]
diff --git a/cascadeflow/harness/api.py b/cascadeflow/harness/api.py
index a71d5f5a..95ff4245 100644
--- a/cascadeflow/harness/api.py
+++ b/cascadeflow/harness/api.py
@@ -4,8 +4,10 @@
import json
import logging
import os
+import time
from contextvars import ContextVar, Token
from dataclasses import dataclass, field
+from functools import wraps
from importlib.util import find_spec
from pathlib import Path
from typing import Any, Callable, Literal, Optional, TypeVar, cast
@@ -39,7 +41,21 @@ class HarnessInitReport:
@dataclass
class HarnessRunContext:
+ """Scoped run context for tracking harness metrics across LLM calls.
+
+ Thread safety: the context is stored in a ``ContextVar`` and is safe for
+ asyncio (each task gets its own copy of the token). However, the context
+ object itself uses plain attribute mutation (``+=``) for counters. If
+ multiple OS threads share the *same* ``HarnessRunContext`` instance,
+ concurrent updates may race. Each ``with run(...)`` scope should be
+ confined to a single thread or asyncio task.
+ """
+
run_id: str = field(default_factory=lambda: uuid4().hex[:12])
+ _started_monotonic: float = field(default_factory=time.monotonic, init=False, repr=False)
+ started_at_ms: float = field(default_factory=lambda: time.time() * 1000)
+ ended_at_ms: Optional[float] = None
+ duration_ms: Optional[float] = None
mode: HarnessMode = "off"
budget_max: Optional[float] = None
tool_calls_max: Optional[int] = None
@@ -73,6 +89,9 @@ def __enter__(self) -> HarnessRunContext:
return self
def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
+ self.ended_at_ms = time.time() * 1000
+ self.duration_ms = max(0.0, (time.monotonic() - self._started_monotonic) * 1000.0)
+ self._log_summary()
if self._token is not None:
_current_run.reset(self._token)
self._token = None
@@ -86,6 +105,44 @@ async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
def trace(self) -> list[dict[str, Any]]:
return list(self._trace)
+ def summary(self) -> dict[str, Any]:
+ return {
+ "run_id": self.run_id,
+ "mode": self.mode,
+ "step_count": self.step_count,
+ "tool_calls": self.tool_calls,
+ "cost": self.cost,
+ "savings": self.savings,
+ "latency_used_ms": self.latency_used_ms,
+ "energy_used": self.energy_used,
+ "budget_max": self.budget_max,
+ "budget_remaining": self.budget_remaining,
+ "last_action": self.last_action,
+ "model_used": self.model_used,
+ "duration_ms": self.duration_ms,
+ }
+
+ def _log_summary(self) -> None:
+ if self.mode == "off" or self.step_count <= 0:
+ return
+ logger.info(
+ (
+ "harness run summary run_id=%s mode=%s steps=%d tool_calls=%d "
+ "cost=%.6f latency_ms=%.2f energy=%.4f last_action=%s model=%s "
+ "budget_remaining=%s"
+ ),
+ self.run_id,
+ self.mode,
+ self.step_count,
+ self.tool_calls,
+ self.cost,
+ self.latency_used_ms,
+ self.energy_used,
+ self.last_action,
+ self.model_used,
+ self.budget_remaining,
+ )
+
def record(
self,
action: str,
@@ -95,19 +152,42 @@ def record(
applied: Optional[bool] = None,
decision_mode: Optional[str] = None,
) -> None:
- self.last_action = action
- self.model_used = model
+ safe_action = _sanitize_trace_value(action, max_length=_MAX_ACTION_LEN)
+ if not safe_action:
+ logger.warning("record() called with empty action, defaulting to 'allow'")
+ safe_action = "allow"
+ safe_reason = _sanitize_trace_value(reason, max_length=_MAX_REASON_LEN) or "unspecified"
+ safe_model = (
+ _sanitize_trace_value(model, max_length=_MAX_MODEL_LEN) if model is not None else None
+ )
+
+ self.last_action = safe_action
+ self.model_used = safe_model
entry: dict[str, Any] = {
- "action": action,
- "reason": reason,
- "model": model,
+ "action": safe_action,
+ "reason": safe_reason,
+ "model": safe_model,
"run_id": self.run_id,
+ "mode": self.mode,
+ "step": self.step_count,
+ "timestamp_ms": time.time() * 1000,
+ "tool_calls_total": self.tool_calls,
+ "cost_total": self.cost,
+ "latency_used_ms": self.latency_used_ms,
+ "energy_used": self.energy_used,
+ "budget_state": {
+ "max": self.budget_max,
+ "remaining": self.budget_remaining,
+ },
}
if applied is not None:
entry["applied"] = applied
if decision_mode is not None:
entry["decision_mode"] = decision_mode
self._trace.append(entry)
+ if len(self._trace) > _MAX_TRACE_ENTRIES:
+ self._trace = self._trace[-_MAX_TRACE_ENTRIES:]
+ _emit_harness_decision(entry)
_harness_config: HarnessConfig = HarnessConfig()
@@ -115,6 +195,7 @@ def record(
"cascadeflow_harness_run", default=None
)
_is_instrumented: bool = False
+_harness_callback_manager: Any = None
_UNSET = object()
@@ -124,6 +205,32 @@ def _validate_mode(mode: str) -> HarnessMode:
return cast(HarnessMode, mode)
+_VALID_COMPLIANCE_VALUES = {"gdpr", "hipaa", "pci", "strict"}
+
+
+def _validate_harness_params(
+ *,
+ budget: Optional[float],
+ max_tool_calls: Optional[int],
+ max_latency_ms: Optional[float],
+ max_energy: Optional[float],
+ compliance: Optional[str],
+) -> None:
+ """Validate harness parameters, raising ValueError for invalid inputs."""
+ if budget is not None and budget < 0:
+ raise ValueError(f"budget must be non-negative, got {budget}")
+ if max_tool_calls is not None and max_tool_calls < 0:
+ raise ValueError(f"max_tool_calls must be non-negative, got {max_tool_calls}")
+ if max_latency_ms is not None and max_latency_ms < 0:
+ raise ValueError(f"max_latency_ms must be non-negative, got {max_latency_ms}")
+ if max_energy is not None and max_energy < 0:
+ raise ValueError(f"max_energy must be non-negative, got {max_energy}")
+ if compliance is not None and compliance.strip().lower() not in _VALID_COMPLIANCE_VALUES:
+ raise ValueError(
+ f"compliance must be one of {sorted(_VALID_COMPLIANCE_VALUES)}, got {compliance!r}"
+ )
+
+
def _detect_sdks() -> dict[str, bool]:
return {
"openai": find_spec("openai") is not None,
@@ -139,6 +246,15 @@ def get_current_run() -> Optional[HarnessRunContext]:
return _current_run.get()
+def get_harness_callback_manager() -> Any:
+ return _harness_callback_manager
+
+
+def set_harness_callback_manager(callback_manager: Any) -> None:
+ global _harness_callback_manager
+ _harness_callback_manager = callback_manager
+
+
def reset() -> None:
"""
Reset harness global state and unpatch instrumented clients.
@@ -148,15 +264,72 @@ def reset() -> None:
global _harness_config
global _is_instrumented
+ global _harness_callback_manager
+ global _cached_cascade_decision_event
- from cascadeflow.harness.instrument import unpatch_openai
+ from cascadeflow.harness.instrument import unpatch_anthropic, unpatch_openai
unpatch_openai()
+ unpatch_anthropic()
_harness_config = HarnessConfig()
_is_instrumented = False
+ _harness_callback_manager = None
+ _cached_cascade_decision_event = None
_current_run.set(None)
+_MAX_ACTION_LEN = 64
+_MAX_REASON_LEN = 160
+_MAX_MODEL_LEN = 128
+_MAX_ENV_JSON_LEN = 4096
+_MAX_TRACE_ENTRIES = 1000
+
+
+def _sanitize_trace_value(value: Any, *, max_length: int) -> Optional[str]:
+ if value is None:
+ return None
+ text = str(value).replace("\n", " ").replace("\r", " ").strip()
+ text = "".join(c for c in text if c.isprintable())
+ if len(text) > max_length:
+ text = text[: max_length - 3] + "..."
+ return text or None
+
+
+_cached_cascade_decision_event: Any = None
+
+
+def _emit_harness_decision(entry: dict[str, Any]) -> None:
+ global _cached_cascade_decision_event
+
+ manager = get_harness_callback_manager()
+ if manager is None:
+ return
+
+ trigger = getattr(manager, "trigger", None)
+ if not callable(trigger):
+ logger.debug("harness callback manager has no trigger() method")
+ return
+
+ if _cached_cascade_decision_event is None:
+ try:
+ from cascadeflow.telemetry.callbacks import CallbackEvent
+
+ _cached_cascade_decision_event = CallbackEvent.CASCADE_DECISION
+ except Exception:
+ logger.debug("telemetry callbacks unavailable for harness decision emit", exc_info=True)
+ return
+
+ try:
+ trigger(
+ _cached_cascade_decision_event,
+ query="[harness]",
+ data=dict(entry),
+ workflow="harness",
+ )
+ except Exception:
+ logger.debug("failed to emit harness decision callback", exc_info=True)
+
+
def _parse_bool(raw: str) -> bool:
normalized = raw.strip().lower()
return normalized in {"1", "true", "yes", "on"}
@@ -171,6 +344,8 @@ def _parse_int(raw: str) -> int:
def _parse_json_dict(raw: str) -> dict[str, float]:
+ if len(raw) > _MAX_ENV_JSON_LEN:
+ raise ValueError(f"JSON config exceeds {_MAX_ENV_JSON_LEN} characters for harness env var")
value = json.loads(raw)
if not isinstance(value, dict):
raise ValueError("expected JSON object")
@@ -305,9 +480,12 @@ def init(
kpi_targets: Optional[dict[str, float]] | object = _UNSET,
kpi_weights: Optional[dict[str, float]] | object = _UNSET,
compliance: Optional[str] | object = _UNSET,
+ callback_manager: Any | object = _UNSET,
) -> HarnessInitReport:
"""
- Initialize global harness settings and instrument detected SDK clients.
+ Initialize global harness settings.
+
+ This is a scaffold API for V2 work and intentionally performs no request patching yet.
"""
global _harness_config
@@ -338,8 +516,18 @@ def init(
resolved_compliance = _resolve_value(
"compliance", compliance, env_config, file_config, None, sources
)
+ if callback_manager is not _UNSET:
+ set_harness_callback_manager(callback_manager)
+ sources["callback_manager"] = "code"
validated_mode = _validate_mode(str(resolved_mode))
+ _validate_harness_params(
+ budget=cast(Optional[float], resolved_budget),
+ max_tool_calls=cast(Optional[int], resolved_max_tool_calls),
+ max_latency_ms=cast(Optional[float], resolved_max_latency_ms),
+ max_energy=cast(Optional[float], resolved_max_energy),
+ compliance=cast(Optional[str], resolved_compliance),
+ )
_harness_config = HarnessConfig(
mode=validated_mode,
verbose=bool(resolved_verbose),
@@ -361,13 +549,29 @@ def init(
if patch_openai():
instrumented.append("openai")
- elif validated_mode == "off":
- from cascadeflow.harness.instrument import is_patched, unpatch_openai
+ else:
+ detected_but_not_instrumented.append("openai")
+
+ if validated_mode != "off" and sdk_presence["anthropic"]:
+ from cascadeflow.harness.instrument import patch_anthropic
- if is_patched():
+ if patch_anthropic():
+ instrumented.append("anthropic")
+ else:
+ detected_but_not_instrumented.append("anthropic")
+
+ if validated_mode == "off":
+ from cascadeflow.harness.instrument import (
+ is_anthropic_patched,
+ is_openai_patched,
+ unpatch_anthropic,
+ unpatch_openai,
+ )
+
+ if is_openai_patched():
unpatch_openai()
- if sdk_presence["anthropic"]:
- detected_but_not_instrumented.append("anthropic")
+ if is_anthropic_patched():
+ unpatch_anthropic()
if _is_instrumented:
logger.debug("harness init called again; instrumentation remains idempotent")
@@ -415,6 +619,14 @@ def run(
resolved_kpi_weights = kpi_weights if kpi_weights is not None else config.kpi_weights
resolved_compliance = compliance if compliance is not None else config.compliance
+ _validate_harness_params(
+ budget=resolved_budget,
+ max_tool_calls=resolved_tool_calls,
+ max_latency_ms=resolved_latency,
+ max_energy=resolved_energy,
+ compliance=resolved_compliance,
+ )
+
return HarnessRunContext(
mode=config.mode,
budget_max=resolved_budget,
@@ -453,18 +665,18 @@ def decorator(func: F) -> F:
if inspect.iscoroutinefunction(func):
+ @wraps(func)
async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
return await func(*args, **kwargs)
async_wrapper.__cascadeflow_agent_policy__ = metadata # type: ignore[attr-defined]
- async_wrapper.__name__ = getattr(func, "__name__", "wrapped_agent")
return cast(F, async_wrapper)
+ @wraps(func)
def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
return func(*args, **kwargs)
sync_wrapper.__cascadeflow_agent_policy__ = metadata # type: ignore[attr-defined]
- sync_wrapper.__name__ = getattr(func, "__name__", "wrapped_agent")
return cast(F, sync_wrapper)
return decorator
diff --git a/cascadeflow/harness/instrument.py b/cascadeflow/harness/instrument.py
index c2fbd7ab..4b08b9f6 100644
--- a/cascadeflow/harness/instrument.py
+++ b/cascadeflow/harness/instrument.py
@@ -1,11 +1,10 @@
-"""OpenAI Python client auto-instrumentation for cascadeflow harness.
+"""Python SDK auto-instrumentation for cascadeflow harness.
-Patches ``openai.resources.chat.completions.Completions.create`` (sync) and
-``AsyncCompletions.create`` (async) to intercept LLM calls for observe/enforce
-modes.
+Patches OpenAI and Anthropic SDK request methods to intercept LLM calls for
+observe/enforce modes.
-This module is called internally by ``cascadeflow.harness.init()``. Users
-should not call ``patch_openai`` / ``unpatch_openai`` directly.
+This module is called internally by ``cascadeflow.harness.init()``. Users
+should not call patch/unpatch helpers directly.
Implementation notes:
- Patching is class-level (all current and future client instances).
@@ -51,6 +50,9 @@
_openai_patched: bool = False
_original_sync_create: Any = None
_original_async_create: Any = None
+_anthropic_patched: bool = False
+_original_anthropic_sync_create: Any = None
+_original_anthropic_async_create: Any = None
_MODEL_TOTAL_COSTS: dict[str, float] = {
name: _model_total_price_shared(name) for name in _PRICING_MODELS
@@ -140,7 +142,7 @@ def _estimate_energy(model: str, prompt_tokens: int, completion_tokens: int) ->
return _estimate_energy_shared(model, prompt_tokens, completion_tokens)
-def _count_tool_calls_in_response(response: Any) -> int:
+def _count_tool_calls_in_openai_response(response: Any) -> int:
"""Count tool calls in a non-streaming ChatCompletion response."""
choices = getattr(response, "choices", None)
if not choices:
@@ -154,7 +156,7 @@ def _count_tool_calls_in_response(response: Any) -> int:
return len(tool_calls)
-def _extract_usage(response: Any) -> tuple[int, int]:
+def _extract_openai_usage(response: Any) -> tuple[int, int]:
"""Extract (prompt_tokens, completion_tokens) from a response."""
usage = getattr(response, "usage", None)
if usage is None:
@@ -165,6 +167,29 @@ def _extract_usage(response: Any) -> tuple[int, int]:
)
+def _extract_anthropic_usage(response: Any) -> tuple[int, int]:
+ """Extract (input_tokens, output_tokens) from an Anthropic response."""
+ usage = getattr(response, "usage", None)
+ if usage is None:
+ return 0, 0
+ return (
+ getattr(usage, "input_tokens", 0) or 0,
+ getattr(usage, "output_tokens", 0) or 0,
+ )
+
+
+def _count_tool_calls_in_anthropic_response(response: Any) -> int:
+ """Count Anthropic ``tool_use`` blocks in a non-streaming response."""
+ content = getattr(response, "content", None)
+ if not content:
+ return 0
+ count = 0
+ for block in content:
+ if getattr(block, "type", None) == "tool_use":
+ count += 1
+ return count
+
+
def _model_total_cost(model: str) -> float:
return _MODEL_TOTAL_COSTS.get(model, _model_total_price_shared(model))
@@ -596,6 +621,9 @@ def __next__(self) -> Any:
except StopIteration:
self._finalize()
raise
+ except Exception:
+ self._finalize()
+ raise
def __enter__(self) -> _InstrumentedStream:
if hasattr(self._stream, "__enter__"):
@@ -625,6 +653,9 @@ async def __anext__(self) -> Any:
except StopAsyncIteration:
self._finalize()
raise
+ except Exception:
+ self._finalize()
+ raise
async def __aenter__(self) -> _InstrumentedAsyncStream:
if hasattr(self._stream, "__aenter__"):
@@ -638,6 +669,174 @@ async def __aexit__(self, *args: Any) -> bool:
return False
+class _InstrumentedAnthropicStreamBase:
+ """Shared stream-wrapper logic for sync and async Anthropic streams."""
+
+ __slots__ = (
+ "_stream",
+ "_ctx",
+ "_model",
+ "_start_time",
+ "_pre_action",
+ "_pre_reason",
+ "_pre_model",
+ "_pre_applied",
+ "_decision_mode",
+ "_input_tokens",
+ "_output_tokens",
+ "_tool_call_count",
+ "_finalized",
+ )
+
+ def __init__(
+ self,
+ stream: Any,
+ ctx: Any,
+ model: str,
+ start_time: float,
+ pre_action: str = "allow",
+ pre_reason: str = "observe",
+ pre_model: str | None = None,
+ pre_applied: bool = True,
+ decision_mode: str = "observe",
+ ) -> None:
+ self._stream = stream
+ self._ctx = ctx
+ self._model = model
+ self._start_time = start_time
+ self._pre_action = pre_action
+ self._pre_reason = pre_reason
+ self._pre_model = pre_model or model
+ self._pre_applied = pre_applied
+ self._decision_mode = decision_mode
+ self._input_tokens: int = 0
+ self._output_tokens: int = 0
+ self._tool_call_count: int = 0
+ self._finalized: bool = False
+
+ def close(self) -> None:
+ self._finalize()
+ if hasattr(self._stream, "close"):
+ self._stream.close()
+
+ def _inspect_event(self, event: Any) -> None:
+ event_type = getattr(event, "type", None)
+
+ if event_type == "message_start":
+ message = getattr(event, "message", None)
+ usage = getattr(message, "usage", None)
+ if usage is not None:
+ input_tokens = getattr(usage, "input_tokens", None)
+ output_tokens = getattr(usage, "output_tokens", None)
+ if isinstance(input_tokens, (int, float)):
+ self._input_tokens = int(input_tokens) if input_tokens > 0 else 0
+ if isinstance(output_tokens, (int, float)):
+ self._output_tokens = int(output_tokens) if output_tokens > 0 else 0
+ return
+
+ usage = getattr(event, "usage", None)
+ if usage is not None:
+ input_tokens = getattr(usage, "input_tokens", None)
+ output_tokens = getattr(usage, "output_tokens", None)
+ if isinstance(input_tokens, (int, float)) and input_tokens > 0:
+ self._input_tokens = int(input_tokens)
+ if isinstance(output_tokens, (int, float)):
+ self._output_tokens = int(output_tokens) if output_tokens > 0 else 0
+
+ if event_type == "content_block_start":
+ content_block = getattr(event, "content_block", None)
+ block_type = getattr(content_block, "type", None)
+ if block_type in {"tool_use", "server_tool_use"}:
+ self._tool_call_count += 1
+
+ def _finalize(self) -> None:
+ if self._finalized:
+ return
+ self._finalized = True
+
+ if self._ctx is None:
+ return
+
+ elapsed_ms = (time.monotonic() - self._start_time) * 1000
+ _update_context(
+ self._ctx,
+ self._model,
+ self._input_tokens,
+ self._output_tokens,
+ self._tool_call_count,
+ elapsed_ms,
+ action=self._pre_action,
+ action_reason=self._pre_reason,
+ action_model=self._pre_model,
+ applied=self._pre_applied,
+ decision_mode=self._decision_mode,
+ )
+
+
+class _InstrumentedAnthropicStream(_InstrumentedAnthropicStreamBase):
+ """Wraps an Anthropic sync stream and tracks usage at stream end."""
+
+ __slots__ = ()
+
+ def __iter__(self) -> _InstrumentedAnthropicStream:
+ return self
+
+ def __next__(self) -> Any:
+ try:
+ event = next(self._stream)
+ self._inspect_event(event)
+ return event
+ except StopIteration:
+ self._finalize()
+ raise
+ except Exception:
+ self._finalize()
+ raise
+
+ def __enter__(self) -> _InstrumentedAnthropicStream:
+ if hasattr(self._stream, "__enter__"):
+ self._stream.__enter__()
+ return self
+
+ def __exit__(self, *args: Any) -> bool:
+ self._finalize()
+ if hasattr(self._stream, "__exit__"):
+ return self._stream.__exit__(*args) # type: ignore[no-any-return]
+ return False
+
+
+class _InstrumentedAnthropicAsyncStream(_InstrumentedAnthropicStreamBase):
+ """Wraps an Anthropic async stream and tracks usage at stream end."""
+
+ __slots__ = ()
+
+ def __aiter__(self) -> _InstrumentedAnthropicAsyncStream:
+ return self
+
+ async def __anext__(self) -> Any:
+ try:
+ event = await self._stream.__anext__()
+ self._inspect_event(event)
+ return event
+ except StopAsyncIteration:
+ self._finalize()
+ raise
+ except Exception:
+ self._finalize()
+ raise
+
+ async def __aenter__(self) -> _InstrumentedAnthropicAsyncStream:
+ if hasattr(self._stream, "__aenter__"):
+ await self._stream.__aenter__()
+ return self
+
+ async def __aexit__(self, *args: Any) -> bool:
+ self._finalize()
+ if hasattr(self._stream, "__aexit__"):
+ return await self._stream.__aexit__(*args) # type: ignore[no-any-return]
+ return False
+
+
# ---------------------------------------------------------------------------
# Wrapper factories
# ---------------------------------------------------------------------------
@@ -713,8 +912,8 @@ def _finalize_interception(
if (not state.is_stream) and ctx:
elapsed_ms = (time.monotonic() - state.start_time) * 1000
- prompt_tokens, completion_tokens = _extract_usage(response)
- tool_call_count = _count_tool_calls_in_response(response)
+ prompt_tokens, completion_tokens = _extract_openai_usage(response)
+ tool_call_count = _count_tool_calls_in_openai_response(response)
_update_context(
ctx,
state.model,
@@ -810,6 +1009,158 @@ async def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
return wrapper
+def _make_patched_anthropic_create(original_fn: Any) -> Any:
+ """Create a patched version of ``anthropic.Messages.create``."""
+
+ @functools.wraps(original_fn)
+ def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
+ from cascadeflow.harness.api import get_current_run, get_harness_config
+
+ config = get_harness_config()
+ ctx = get_current_run()
+ mode = ctx.mode if ctx else config.mode
+
+ if mode == "off":
+ return original_fn(self, *args, **kwargs)
+
+ model: str = kwargs.get("model", "unknown")
+ pre_action = "allow"
+ pre_reason = mode
+ pre_model = model
+ pre_applied = True
+
+ if ctx:
+ kwargs, model, pre_action, pre_reason, pre_model, pre_applied = (
+ _resolve_pre_call_decision(
+ ctx,
+ mode,
+ model,
+ kwargs,
+ )
+ )
+
+ is_stream = bool(kwargs.get("stream", False))
+ start_time = time.monotonic()
+ response = original_fn(self, *args, **kwargs)
+
+ if not ctx:
+ logger.debug(
+ "harness %s (anthropic): model=%s (no active run scope, metrics not tracked)",
+ mode,
+ model,
+ )
+ return response
+
+ if is_stream:
+ return _InstrumentedAnthropicStream(
+ response,
+ ctx,
+ model,
+ start_time,
+ pre_action,
+ pre_reason,
+ pre_model,
+ pre_applied,
+ mode,
+ )
+
+ elapsed_ms = (time.monotonic() - start_time) * 1000
+ input_tokens, output_tokens = _extract_anthropic_usage(response)
+ tool_call_count = _count_tool_calls_in_anthropic_response(response)
+ _update_context(
+ ctx,
+ model,
+ input_tokens,
+ output_tokens,
+ tool_call_count,
+ elapsed_ms,
+ action=pre_action,
+ action_reason=pre_reason,
+ action_model=pre_model,
+ applied=pre_applied,
+ decision_mode=mode,
+ )
+ return response
+
+ return wrapper
+
+
+def _make_patched_anthropic_async_create(original_fn: Any) -> Any:
+ """Create a patched version of ``anthropic.AsyncMessages.create``."""
+
+ @functools.wraps(original_fn)
+ async def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
+ from cascadeflow.harness.api import get_current_run, get_harness_config
+
+ config = get_harness_config()
+ ctx = get_current_run()
+ mode = ctx.mode if ctx else config.mode
+
+ if mode == "off":
+ return await original_fn(self, *args, **kwargs)
+
+ model: str = kwargs.get("model", "unknown")
+ pre_action = "allow"
+ pre_reason = mode
+ pre_model = model
+ pre_applied = True
+
+ if ctx:
+ kwargs, model, pre_action, pre_reason, pre_model, pre_applied = (
+ _resolve_pre_call_decision(
+ ctx,
+ mode,
+ model,
+ kwargs,
+ )
+ )
+
+ is_stream = bool(kwargs.get("stream", False))
+ start_time = time.monotonic()
+ response = await original_fn(self, *args, **kwargs)
+
+ if not ctx:
+ logger.debug(
+ "harness %s async (anthropic): model=%s (no active run scope, metrics not tracked)",
+ mode,
+ model,
+ )
+ return response
+
+ if is_stream:
+ return _InstrumentedAnthropicAsyncStream(
+ response,
+ ctx,
+ model,
+ start_time,
+ pre_action,
+ pre_reason,
+ pre_model,
+ pre_applied,
+ mode,
+ )
+
+ elapsed_ms = (time.monotonic() - start_time) * 1000
+ input_tokens, output_tokens = _extract_anthropic_usage(response)
+ tool_call_count = _count_tool_calls_in_anthropic_response(response)
+ _update_context(
+ ctx,
+ model,
+ input_tokens,
+ output_tokens,
+ tool_call_count,
+ elapsed_ms,
+ action=pre_action,
+ action_reason=pre_reason,
+ action_model=pre_model,
+ applied=pre_applied,
+ decision_mode=mode,
+ )
+ return response
+
+ return wrapper
+
+
# ---------------------------------------------------------------------------
# Public API (called by cascadeflow.harness.api)
# ---------------------------------------------------------------------------
@@ -846,6 +1197,37 @@ def patch_openai() -> bool:
return True
+def patch_anthropic() -> bool:
+ """Patch the Anthropic Python client for harness instrumentation.
+
+ Returns ``True`` if patching succeeded, ``False`` if anthropic is not
+ installed. Idempotent: safe to call multiple times.
+ """
+ global _anthropic_patched, _original_anthropic_sync_create, _original_anthropic_async_create
+
+ if _anthropic_patched:
+ logger.debug("anthropic already patched, skipping")
+ return True
+
+ try:
+ from anthropic.resources.messages import AsyncMessages, Messages
+ except ImportError:
+ logger.debug("anthropic package not available, skipping instrumentation")
+ return False
+
+ _original_anthropic_sync_create = Messages.create
+ _original_anthropic_async_create = AsyncMessages.create
+
+ Messages.create = _make_patched_anthropic_create(_original_anthropic_sync_create) # type: ignore[assignment]
+ AsyncMessages.create = _make_patched_anthropic_async_create( # type: ignore[assignment]
+ _original_anthropic_async_create,
+ )
+
+ _anthropic_patched = True
+ logger.info("anthropic client instrumented (sync + async)")
+ return True
+
+
def unpatch_openai() -> None:
"""Restore original OpenAI client methods.
@@ -873,6 +1255,43 @@ def unpatch_openai() -> None:
logger.info("openai client unpatched")
-def is_patched() -> bool:
+def unpatch_anthropic() -> None:
+ """Restore original Anthropic client methods.
+
+ Safe to call even if not patched. Used by ``reset()`` and tests.
+ """
+ global _anthropic_patched, _original_anthropic_sync_create, _original_anthropic_async_create
+
+ if not _anthropic_patched:
+ return
+
+ try:
+ from anthropic.resources.messages import AsyncMessages, Messages
+ except ImportError:
+ _anthropic_patched = False
+ return
+
+ if _original_anthropic_sync_create is not None:
+ Messages.create = _original_anthropic_sync_create # type: ignore[assignment]
+ if _original_anthropic_async_create is not None:
+ AsyncMessages.create = _original_anthropic_async_create # type: ignore[assignment]
+
+ _original_anthropic_sync_create = None
+ _original_anthropic_async_create = None
+ _anthropic_patched = False
+ logger.info("anthropic client unpatched")
+
+
+def is_openai_patched() -> bool:
"""Return whether the OpenAI client is currently patched."""
return _openai_patched
+
+
+def is_anthropic_patched() -> bool:
+ """Return whether the Anthropic client is currently patched."""
+ return _anthropic_patched
+
+
+def is_patched() -> bool:
+ """Return whether any supported Python SDK is currently patched."""
+ return _openai_patched or _anthropic_patched
diff --git a/cascadeflow/harness/pricing.py b/cascadeflow/harness/pricing.py
index bd86323e..81a1de06 100644
--- a/cascadeflow/harness/pricing.py
+++ b/cascadeflow/harness/pricing.py
@@ -1,11 +1,17 @@
"""Shared harness pricing and energy profiles.
This module centralizes model-cost and energy-estimation defaults used by
-harness integrations (OpenAI auto-instrumentation, OpenAI Agents SDK, CrewAI).
+harness integrations (OpenAI auto-instrumentation, OpenAI Agents SDK, CrewAI,
+Google ADK).
+
+A future pricing registry will consolidate with ``cascadeflow.pricing``
+and LiteLLM live data. Until then this module is the canonical source
+for harness-level cost/energy estimation.
"""
from __future__ import annotations
+import re as _re
from typing import Final
# USD per 1M tokens (input, output).
@@ -21,15 +27,22 @@
"o1": (15.00, 60.00),
"o1-mini": (3.00, 12.00),
"o3-mini": (1.10, 4.40),
- # Anthropic aliases used by CrewAI model names.
+ # Anthropic
"claude-sonnet-4": (3.00, 15.00),
"claude-haiku-3.5": (1.00, 5.00),
"claude-opus-4.5": (5.00, 25.00),
+ # Google Gemini
+ "gemini-2.5-flash": (0.15, 0.60),
+ "gemini-2.5-pro": (1.25, 10.00),
+ "gemini-2.0-flash": (0.10, 0.40),
+ "gemini-1.5-flash": (0.075, 0.30),
+ "gemini-1.5-pro": (1.25, 5.00),
}
DEFAULT_PRICING_USD_PER_M: Final[tuple[float, float]] = (2.50, 10.00)
# Deterministic proxy coefficients for energy tracking.
ENERGY_COEFFICIENTS: Final[dict[str, float]] = {
+ # OpenAI
"gpt-4o": 1.0,
"gpt-4o-mini": 0.3,
"gpt-5": 1.2,
@@ -40,6 +53,16 @@
"o1": 2.0,
"o1-mini": 0.8,
"o3-mini": 0.5,
+ # Anthropic
+ "claude-sonnet-4": 1.0,
+ "claude-haiku-3.5": 0.3,
+ "claude-opus-4.5": 1.8,
+ # Google Gemini
+ "gemini-2.5-flash": 0.3,
+ "gemini-2.5-pro": 1.2,
+ "gemini-2.0-flash": 0.25,
+ "gemini-1.5-flash": 0.2,
+ "gemini-1.5-pro": 1.0,
}
DEFAULT_ENERGY_COEFFICIENT: Final[float] = 1.0
ENERGY_OUTPUT_WEIGHT: Final[float] = 1.5
@@ -60,19 +83,85 @@
)
+# ---------------------------------------------------------------------------
+# Fuzzy model-name resolution
+# ---------------------------------------------------------------------------
+
+# Pre-compiled pattern for stripping version/preview/date suffixes.
+# Matches: -preview, -preview-05-20, -20250120, -latest, -exp-0827, etc.
+_VERSION_SUFFIX_RE = _re.compile(
+ r"(-preview(?:-\d{2,4}-\d{2})?|-\d{8,}|-latest|-exp(?:-\d+)?|-it)$"
+)
+
+# Cache for resolved model → pricing key lookups.
+_pricing_key_cache: dict[str, str | None] = {}
+
+
+def _resolve_pricing_key(model: str) -> str | None:
+ """Resolve a model name to a known pricing table key.
+
+ Tries exact match first, then strips version/preview/date suffixes,
+ then tries longest-prefix match against known model names.
+ Returns ``None`` when no match is found (caller should use defaults).
+ """
+ if model in _pricing_key_cache:
+ return _pricing_key_cache[model]
+
+ # Exact match
+ if model in PRICING_USD_PER_M:
+ _pricing_key_cache[model] = model
+ return model
+
+ # Strip version suffixes and retry
+ stripped = _VERSION_SUFFIX_RE.sub("", model)
+ if stripped != model and stripped in PRICING_USD_PER_M:
+ _pricing_key_cache[model] = stripped
+ return stripped
+
+ # Longest-prefix match (e.g. "gemini-2.5-flash-8b" → "gemini-2.5-flash")
+ best: str | None = None
+ best_len = 0
+ for known in PRICING_USD_PER_M:
+ if model.startswith(known) and len(known) > best_len:
+ best = known
+ best_len = len(known)
+ if best is not None:
+ _pricing_key_cache[model] = best
+ return best
+
+ _pricing_key_cache[model] = None
+ return None
+
+
+# ---------------------------------------------------------------------------
+# Public estimation helpers
+# ---------------------------------------------------------------------------
+
+
def estimate_cost(model: str, input_tokens: int, output_tokens: int) -> float:
"""Estimate USD cost from token usage."""
- in_price, out_price = PRICING_USD_PER_M.get(model, DEFAULT_PRICING_USD_PER_M)
+ key = _resolve_pricing_key(model)
+ in_price, out_price = (
+ PRICING_USD_PER_M.get(key, DEFAULT_PRICING_USD_PER_M) if key else DEFAULT_PRICING_USD_PER_M
+ )
return (input_tokens / 1_000_000.0) * in_price + (output_tokens / 1_000_000.0) * out_price
def estimate_energy(model: str, input_tokens: int, output_tokens: int) -> float:
"""Estimate deterministic proxy energy units."""
- coefficient = ENERGY_COEFFICIENTS.get(model, DEFAULT_ENERGY_COEFFICIENT)
- return coefficient * (input_tokens + (output_tokens * ENERGY_OUTPUT_WEIGHT))
+ key = _resolve_pricing_key(model)
+ coeff = (
+ ENERGY_COEFFICIENTS.get(key, DEFAULT_ENERGY_COEFFICIENT)
+ if key
+ else DEFAULT_ENERGY_COEFFICIENT
+ )
+ return coeff * (input_tokens + (output_tokens * ENERGY_OUTPUT_WEIGHT))
def model_total_price(model: str) -> float:
"""Return total (input + output) price per 1M tokens."""
- in_price, out_price = PRICING_USD_PER_M.get(model, DEFAULT_PRICING_USD_PER_M)
+ key = _resolve_pricing_key(model)
+ in_price, out_price = (
+ PRICING_USD_PER_M.get(key, DEFAULT_PRICING_USD_PER_M) if key else DEFAULT_PRICING_USD_PER_M
+ )
return in_price + out_price
diff --git a/cascadeflow/integrations/__init__.py b/cascadeflow/integrations/__init__.py
index 33552773..61c3ebbd 100644
--- a/cascadeflow/integrations/__init__.py
+++ b/cascadeflow/integrations/__init__.py
@@ -185,6 +185,28 @@
crewai_is_enabled = None
crewai_get_config = None
+# Try to import Google ADK integration
+try:
+ from .google_adk import (
+ GOOGLE_ADK_AVAILABLE,
+ GoogleADKHarnessConfig,
+ CascadeFlowADKPlugin,
+ enable as google_adk_enable,
+ disable as google_adk_disable,
+ is_available as google_adk_is_available,
+ is_enabled as google_adk_is_enabled,
+ get_config as google_adk_get_config,
+ )
+except ImportError:
+ GOOGLE_ADK_AVAILABLE = False
+ GoogleADKHarnessConfig = None
+ CascadeFlowADKPlugin = None
+ google_adk_enable = None
+ google_adk_disable = None
+ google_adk_is_available = None
+ google_adk_is_enabled = None
+ google_adk_get_config = None
+
__all__ = []
if LITELLM_AVAILABLE:
@@ -285,6 +307,20 @@
]
)
+if GOOGLE_ADK_AVAILABLE:
+ __all__.extend(
+ [
+ "GOOGLE_ADK_AVAILABLE",
+ "GoogleADKHarnessConfig",
+ "CascadeFlowADKPlugin",
+ "google_adk_enable",
+ "google_adk_disable",
+ "google_adk_is_available",
+ "google_adk_is_enabled",
+ "google_adk_get_config",
+ ]
+ )
+
# Integration capabilities
INTEGRATION_CAPABILITIES = {
"litellm": LITELLM_AVAILABLE,
@@ -294,6 +330,7 @@
"openclaw": OPENCLAW_AVAILABLE,
"paygentic": PAYGENTIC_AVAILABLE,
"crewai": CREWAI_AVAILABLE,
+ "google_adk": GOOGLE_ADK_AVAILABLE,
}
@@ -319,4 +356,5 @@ def get_integration_info():
"openclaw_available": OPENCLAW_AVAILABLE,
"paygentic_available": PAYGENTIC_AVAILABLE,
"crewai_available": CREWAI_AVAILABLE,
+ "google_adk_available": GOOGLE_ADK_AVAILABLE,
}
diff --git a/cascadeflow/integrations/google_adk.py b/cascadeflow/integrations/google_adk.py
new file mode 100644
index 00000000..325d21b2
--- /dev/null
+++ b/cascadeflow/integrations/google_adk.py
@@ -0,0 +1,486 @@
+"""Google ADK (Agent Development Kit) harness integration for cascadeflow.
+
+Uses ADK's ``BasePlugin`` system to intercept all LLM calls across all agents
+in a Runner, feeding metrics into ``cascadeflow.harness`` run contexts.
+
+This module is optional — ``pip install cascadeflow[google-adk]`` pulls in the
+google-adk dependency. When google-adk is not installed the public helpers
+return gracefully and ``GOOGLE_ADK_AVAILABLE`` is ``False``.
+
+Integration surface:
+ - ``enable()``: create and return a plugin instance
+ - ``disable()``: deactivate the plugin and clean up
+ - ``CascadeFlowADKPlugin``: BasePlugin subclass for Runner(plugins=[...])
+
+Unlike CrewAI (global hooks), ADK plugins are registered per-Runner.
+``enable()`` returns the plugin instance; the user passes it to
+``Runner(plugins=[plugin])``.
+
+Design note — no tool gating:
+ ADK's ``tools_dict`` is part of agent definition, not per-call.
+ Budget gate via ``before_model_callback`` provides sufficient cost control.
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from dataclasses import dataclass
+from importlib.util import find_spec
+from typing import Any, Optional
+
+from cascadeflow.harness.api import get_current_run
+from cascadeflow.harness.pricing import estimate_cost, estimate_energy
+
+logger = logging.getLogger("cascadeflow.integrations.google_adk")
+
+GOOGLE_ADK_AVAILABLE = find_spec("google.adk") is not None
+
+# Resolve the base class: use ADK's BasePlugin when available, else object.
+_ADKBasePlugin: type
+if GOOGLE_ADK_AVAILABLE:
+ try:
+ from google.adk.plugins import BasePlugin as _ADKBasePlugin # type: ignore[assignment]
+ except ImportError:
+ _ADKBasePlugin = object # type: ignore[assignment,misc]
+ GOOGLE_ADK_AVAILABLE = False
+else:
+ _ADKBasePlugin = object # type: ignore[assignment,misc]
+
+
+# ---------------------------------------------------------------------------
+# Configuration
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class GoogleADKHarnessConfig:
+ """Runtime configuration for the Google ADK harness integration.
+
+ fail_open:
+ If ``True`` (default), errors inside callbacks never break ADK
+ execution — they are logged and swallowed.
+ enable_budget_gate:
+ If ``True`` (default), ``before_model_callback`` blocks calls when
+ the harness run budget is exhausted (enforce mode only).
+ """
+
+ fail_open: bool = True
+ enable_budget_gate: bool = True
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _normalize_model_name(model: str) -> str:
+ """Strip LiteLlm-style provider prefix (``openai/gpt-4o`` → ``gpt-4o``).
+
+ Also handles ``models/gemini-2.5-flash`` → ``gemini-2.5-flash``.
+ """
+ if "/" in model:
+ return model.rsplit("/", 1)[-1]
+ return model
+
+
+def _count_function_calls(content: Any) -> int:
+ """Count ``function_call`` parts in an ADK LlmResponse content."""
+ if content is None:
+ return 0
+ parts = getattr(content, "parts", None)
+ if not parts:
+ return 0
+ count = 0
+ for part in parts:
+ if getattr(part, "function_call", None) is not None:
+ count += 1
+ return count
+
+
+# ---------------------------------------------------------------------------
+# Plugin
+# ---------------------------------------------------------------------------
+
+
+class CascadeFlowADKPlugin(_ADKBasePlugin): # type: ignore[misc]
+ """Google ADK BasePlugin with cascadeflow harness awareness.
+
+ Intercepts every LLM call across all agents in a Runner to provide:
+ - Budget enforcement (enforce mode: short-circuits with error response)
+ - Cost, latency, and energy tracking
+ - Tool call counting
+ - Full trace recording into HarnessRunContext
+ """
+
+ def __init__(self, config: Optional[GoogleADKHarnessConfig] = None) -> None:
+ # google-adk BasePlugin requires a stable plugin name.
+ try:
+ super().__init__(name="cascadeflow_harness")
+ except TypeError:
+ # Fallback for local test environments where BasePlugin is ``object``.
+ super().__init__()
+ self.name = "cascadeflow_harness"
+ self._config = config or GoogleADKHarnessConfig()
+ self._active = True
+ self._call_seq: int = 0
+ # Track call metadata between before/after callbacks.
+ # Keyed by id(callback_context) to guarantee uniqueness even when
+ # two concurrent calls share (invocation_id, agent_name).
+ self._call_start_times: dict[int, float] = {}
+ self._call_models: dict[int, str] = {}
+ # Fallback mapping for runtimes that provide distinct callback_context
+ # objects between before/after callbacks.
+ self._call_fallback_keys: dict[tuple[str, str], list[int]] = {}
+
+ @staticmethod
+ def _callback_key(callback_context: Any) -> int:
+ """Return a unique key for a callback_context object.
+
+ Uses ``id()`` which is guaranteed unique for the lifetime of the
+ object — ADK keeps the same CallbackContext alive across the
+ before/after/error callback sequence for a single LLM call.
+ """
+ return id(callback_context)
+
+ @staticmethod
+ def _fallback_key(callback_context: Any) -> tuple[str, str]:
+ """Return a stable fallback key for correlation across callbacks."""
+ invocation_id = str(getattr(callback_context, "invocation_id", "") or "")
+ agent_name = str(getattr(callback_context, "agent_name", "") or "")
+ return (invocation_id, agent_name)
+
+ def _track_call_key(self, callback_context: Any, key: int) -> None:
+ """Register key in fallback queue for cross-object callback matching."""
+ fallback_key = self._fallback_key(callback_context)
+ if not fallback_key[0] and not fallback_key[1]:
+ return
+ self._call_fallback_keys.setdefault(fallback_key, []).append(key)
+
+ def _resolve_call_key(self, callback_context: Any) -> int | None:
+ """Resolve stored key for callback context across runtime variants."""
+ key = self._callback_key(callback_context)
+ if key in self._call_models or key in self._call_start_times:
+ return key
+
+ fallback_key = self._fallback_key(callback_context)
+ keys = self._call_fallback_keys.get(fallback_key)
+ if not keys:
+ return None
+
+ resolved = keys.pop(0)
+ if not keys:
+ self._call_fallback_keys.pop(fallback_key, None)
+ return resolved
+
+ async def before_model_callback(
+ self,
+ callback_context: Any,
+ llm_request: Any,
+ ) -> Any:
+ """Budget gate and timing setup.
+
+ Returns ``None`` to proceed normally, or an ``LlmResponse`` with
+ an error to short-circuit the call when budget is exhausted.
+ """
+ if not self._active:
+ return None
+
+ try:
+ ctx = get_current_run()
+ if ctx is None:
+ return None
+ if ctx.mode == "off":
+ return None
+
+ # Extract model name from request
+ model_raw = getattr(llm_request, "model", None) or "unknown"
+ model = _normalize_model_name(str(model_raw))
+
+ key = self._callback_key(callback_context)
+
+ # Budget gate in enforce mode
+ if (
+ self._config.enable_budget_gate
+ and ctx.mode == "enforce"
+ and ctx.budget_max is not None
+ and ctx.cost >= ctx.budget_max
+ ):
+ logger.warning(
+ "google-adk: blocking LLM call — budget exhausted "
+ "(spent $%.4f of $%.4f max)",
+ ctx.cost,
+ ctx.budget_max,
+ )
+ ctx.record(action="stop", reason="budget_exhausted", model=model)
+ return self._make_budget_error_response(ctx)
+
+ # Record start time and model for after_model_callback
+ self._call_start_times[key] = time.monotonic()
+ self._call_models[key] = model
+ self._track_call_key(callback_context, key)
+
+ return None
+ except Exception:
+ if self._config.fail_open:
+ logger.debug("google-adk before_model_callback error (fail_open)", exc_info=True)
+ return None
+ raise
+
+ async def after_model_callback(
+ self,
+ callback_context: Any,
+ llm_response: Any,
+ ) -> Any:
+ """Extract tokens, count tool calls, estimate cost/energy, update run context."""
+ if not self._active:
+ return None
+
+ try:
+ ctx = get_current_run()
+ if ctx is None:
+ return None
+ if ctx.mode == "off":
+ return None
+
+ key = self._resolve_call_key(callback_context)
+
+ # Recover model name stored during before_model_callback
+ model = self._call_models.pop(key, "unknown") if key is not None else "unknown"
+
+ # Extract token counts from usage_metadata
+ input_tokens, output_tokens = self._extract_tokens(llm_response)
+
+ # Count function_call parts in response content
+ content = getattr(llm_response, "content", None)
+ tool_calls = _count_function_calls(content)
+
+ # Cost and energy estimation
+ cost = estimate_cost(model, input_tokens, output_tokens)
+ energy = estimate_energy(model, input_tokens, output_tokens)
+
+ # Latency
+ start_time = self._call_start_times.pop(key, None) if key is not None else None
+ elapsed_ms = (time.monotonic() - start_time) * 1000 if start_time else 0.0
+
+ # Update run context
+ ctx.cost += cost
+ ctx.step_count += 1
+ ctx.latency_used_ms += elapsed_ms
+ ctx.energy_used += energy
+ ctx.tool_calls += tool_calls
+
+ if ctx.budget_max is not None:
+ ctx.budget_remaining = ctx.budget_max - ctx.cost
+
+ ctx.model_used = model
+ ctx.record(action="allow", reason=ctx.mode, model=model)
+
+ logger.debug(
+ "google-adk: tracked call model=%s cost=$%.6f latency=%.0fms tools=%d",
+ model,
+ cost,
+ elapsed_ms,
+ tool_calls,
+ )
+
+ return None
+ except Exception:
+ if self._config.fail_open:
+ logger.debug("google-adk after_model_callback error (fail_open)", exc_info=True)
+ return None
+ raise
+
+ async def on_model_error_callback(
+ self,
+ callback_context: Any,
+ llm_request: Any = None,
+ error: Exception | None = None,
+ ) -> Any:
+ """Record error in trace and clean up timing state."""
+ if not self._active:
+ return None
+
+ try:
+ # Backward-compatible calling form used in existing tests:
+ # on_model_error_callback(callback_context, error)
+ if error is None and isinstance(llm_request, Exception):
+ error = llm_request
+
+ key = self._resolve_call_key(callback_context)
+ model = self._call_models.pop(key, "unknown") if key is not None else "unknown"
+ if key is not None:
+ self._call_start_times.pop(key, None)
+
+ ctx = get_current_run()
+ if ctx is not None and error is not None:
+ error_type = type(error).__name__
+ ctx.record(
+ action="error",
+ reason=f"model_error:{error_type}",
+ model=model,
+ )
+
+ return None
+ except Exception:
+ if self._config.fail_open:
+ logger.debug("google-adk on_model_error_callback error (fail_open)", exc_info=True)
+ return None
+ raise
+
+ def deactivate(self) -> None:
+ """Make all callbacks no-ops without unregistering from Runner."""
+ self._active = False
+ self._call_seq = 0
+ self._call_start_times.clear()
+ self._call_models.clear()
+ self._call_fallback_keys.clear()
+
+ @staticmethod
+ def _extract_tokens(llm_response: Any) -> tuple[int, int]:
+ """Extract input/output token counts from an ADK LlmResponse.
+
+ ADK responses carry ``usage_metadata`` with ``prompt_token_count``
+ and ``candidates_token_count``. Falls back to estimating from
+ content text (4 chars ≈ 1 token).
+ """
+ usage = getattr(llm_response, "usage_metadata", None)
+ if usage is not None:
+ input_tokens = getattr(usage, "prompt_token_count", 0) or 0
+ output_tokens = getattr(usage, "candidates_token_count", 0) or 0
+ if input_tokens > 0 or output_tokens > 0:
+ return int(input_tokens), int(output_tokens)
+
+ # Fallback: estimate from content text
+ content = getattr(llm_response, "content", None)
+ if content is not None:
+ parts = getattr(content, "parts", None)
+ if parts:
+ text_chars = sum(len(getattr(p, "text", "") or "") for p in parts)
+ return 0, max(text_chars // 4, 1)
+
+ return 0, 0
+
+ @staticmethod
+ def _make_budget_error_response(ctx: Any) -> Any:
+ """Build an LlmResponse that short-circuits the LLM call.
+
+ When ADK is available we return a real ``LlmResponse``. When not
+ (shouldn't happen in practice), we return a sentinel dict.
+
+ The user-facing message is intentionally generic to avoid leaking
+ internal spend/limit numbers. Exact figures are logged separately.
+ """
+ # Generic message safe for end-user exposure.
+ msg = "cascadeflow harness budget exceeded"
+ # Detailed figures for operators only.
+ logger.warning(
+ "google-adk: budget exceeded — spent $%.4f of $%.4f max",
+ ctx.cost,
+ ctx.budget_max,
+ )
+ if GOOGLE_ADK_AVAILABLE:
+ try:
+ from google.adk.models import LlmResponse # type: ignore[import-untyped]
+ from google.genai.types import Content, Part # type: ignore[import-untyped]
+
+ return LlmResponse(
+ content=Content(parts=[Part(text=msg)]),
+ error_code="BUDGET_EXCEEDED",
+ error_message=msg,
+ )
+ except ImportError:
+ pass
+
+ return {"error_code": "BUDGET_EXCEEDED", "error_message": msg}
+
+
+# ---------------------------------------------------------------------------
+# Module-level state
+# ---------------------------------------------------------------------------
+
+_config: GoogleADKHarnessConfig = GoogleADKHarnessConfig()
+_plugin_instance: Optional[CascadeFlowADKPlugin] = None
+_enabled: bool = False
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+
+def is_available() -> bool:
+ """Return whether the google-adk package is installed."""
+ return GOOGLE_ADK_AVAILABLE
+
+
+def is_enabled() -> bool:
+ """Return whether a plugin instance has been created via ``enable()``."""
+ return _enabled
+
+
+def get_config() -> GoogleADKHarnessConfig:
+ """Return a copy of the current configuration."""
+ return GoogleADKHarnessConfig(
+ fail_open=_config.fail_open,
+ enable_budget_gate=_config.enable_budget_gate,
+ )
+
+
+def enable(
+ config: Optional[GoogleADKHarnessConfig] = None,
+) -> CascadeFlowADKPlugin:
+ """Create a cascadeflow-instrumented ADK plugin instance.
+
+ Unlike CrewAI (global hooks), ADK plugins are per-Runner. Pass the
+ returned plugin to ``Runner(plugins=[plugin])``.
+
+ Idempotent: returns the same instance on repeated calls unless
+ ``disable()`` was called in between.
+
+ Args:
+ config: Optional configuration overrides.
+
+ Returns:
+ ``CascadeFlowADKPlugin`` instance ready for ``Runner(plugins=[...])``.
+ """
+ global _config, _plugin_instance, _enabled
+
+ if _enabled and _plugin_instance is not None:
+ logger.debug("google-adk plugin already enabled; returning existing instance")
+ return _plugin_instance
+
+ if config is not None:
+ _config = config
+
+ _plugin_instance = CascadeFlowADKPlugin(config=_config)
+ _enabled = True
+ logger.info("google-adk harness plugin created")
+ return _plugin_instance
+
+
+def disable() -> None:
+ """Deactivate the plugin and clear module state.
+
+ Safe to call even if not enabled.
+ """
+ global _plugin_instance, _enabled
+
+ if _plugin_instance is not None:
+ _plugin_instance.deactivate()
+
+ _plugin_instance = None
+ _enabled = False
+ logger.info("google-adk harness plugin disabled")
+
+
+__all__ = [
+ "GOOGLE_ADK_AVAILABLE",
+ "GoogleADKHarnessConfig",
+ "CascadeFlowADKPlugin",
+ "enable",
+ "disable",
+ "is_available",
+ "is_enabled",
+ "get_config",
+]
diff --git a/cascadeflow/integrations/langchain/__init__.py b/cascadeflow/integrations/langchain/__init__.py
index 45c6ea2f..7b3f9551 100644
--- a/cascadeflow/integrations/langchain/__init__.py
+++ b/cascadeflow/integrations/langchain/__init__.py
@@ -54,6 +54,14 @@
CascadeFlowCallbackHandler,
get_cascade_callback,
)
+from .harness_callback import (
+ HarnessAwareCascadeFlowCallbackHandler,
+ get_harness_callback,
+)
+from .harness_state import (
+ apply_langgraph_state,
+ extract_langgraph_state,
+)
__all__ = [
# Main classes
@@ -93,4 +101,8 @@
# LangChain callback handlers
"CascadeFlowCallbackHandler",
"get_cascade_callback",
+ "HarnessAwareCascadeFlowCallbackHandler",
+ "get_harness_callback",
+ "extract_langgraph_state",
+ "apply_langgraph_state",
]
diff --git a/cascadeflow/integrations/langchain/harness_callback.py b/cascadeflow/integrations/langchain/harness_callback.py
new file mode 100644
index 00000000..01f08d8c
--- /dev/null
+++ b/cascadeflow/integrations/langchain/harness_callback.py
@@ -0,0 +1,248 @@
+"""Harness-aware callbacks for LangChain/LangGraph integration.
+
+Enforce-mode limitations (LangChain callback architecture):
+ - ``stop`` (budget/latency/energy exceeded): fully enforced — raises
+ BudgetExceededError or HarnessStopError from ``on_llm_start``.
+ - ``deny_tool`` (tool-call cap): fully enforced at the tool level via
+ ``on_tool_start`` — raises HarnessStopError before tool execution.
+ - ``switch_model``: **observe-only** — LangChain dispatches the LLM call
+ before ``on_llm_start`` returns, so the callback cannot redirect to a
+ different model. The decision is recorded with ``applied=False``.
+ - ``deny_tool`` at LLM level (pre-call decision): **observe-only** — the
+ callback cannot strip tools from an already-dispatched LLM request.
+ The decision is recorded with ``applied=False``.
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from contextlib import contextmanager
+from typing import Any, Optional
+
+from cascadeflow.harness import get_current_run
+from cascadeflow.harness.pricing import estimate_cost, estimate_energy
+from cascadeflow.schema.exceptions import HarnessStopError
+
+from .harness_state import apply_langgraph_state, extract_langgraph_state
+from .langchain_callbacks import CascadeFlowCallbackHandler
+from .utils import extract_token_usage
+
+logger = logging.getLogger("cascadeflow.harness.langchain")
+
+
+class HarnessAwareCascadeFlowCallbackHandler(CascadeFlowCallbackHandler):
+ """LangChain callback that bridges native lifecycle events into HarnessRunContext.
+
+ See module docstring for enforce-mode limitations on ``switch_model``
+ and LLM-level ``deny_tool``.
+ """
+
+ def __init__(self, *, fail_open: bool = True):
+ super().__init__()
+ self.fail_open = fail_open
+ self._llm_started_at: Optional[float] = None
+ self._pre_action: str = "allow"
+ self._pre_reason: str = "allow"
+ self._pre_model: Optional[str] = None
+ self._pre_recorded: bool = False
+
+ def _handle_harness_error(self, error: Exception) -> None:
+ if self.fail_open:
+ logger.exception("langchain harness callback failed (fail-open)", exc_info=error)
+ return
+ raise error
+
+ def _sync_state(self, payload: dict[str, Any]) -> None:
+ run_ctx = get_current_run()
+ if run_ctx is None:
+ return
+ state = extract_langgraph_state(payload)
+ if state:
+ apply_langgraph_state(run_ctx, state)
+
+ def on_llm_start(self, serialized: dict[str, Any], prompts: list[str], **kwargs: Any) -> None:
+ super().on_llm_start(serialized=serialized, prompts=prompts, **kwargs)
+ self._llm_started_at = time.monotonic()
+ self._pre_action = "allow"
+ self._pre_reason = "allow"
+ self._pre_model = self.current_model
+ self._pre_recorded = False
+
+ try:
+ self._sync_state(kwargs)
+
+ run_ctx = get_current_run()
+ if run_ctx is None:
+ return
+
+ model_name = self.current_model or "unknown"
+ invocation_params = kwargs.get("invocation_params")
+ has_tools = False
+ if isinstance(invocation_params, dict):
+ has_tools = bool(invocation_params.get("tools"))
+ if not has_tools:
+ has_tools = bool(kwargs.get("tools"))
+
+ from cascadeflow.harness.instrument import (
+ _evaluate_pre_call_decision,
+ _raise_stop_error,
+ ) # noqa: I001
+
+ decision = _evaluate_pre_call_decision(run_ctx, model_name, has_tools=has_tools)
+ self._pre_action = decision.action
+ self._pre_reason = decision.reason
+ self._pre_model = decision.target_model
+
+ if run_ctx.mode == "observe":
+ if decision.action != "allow":
+ run_ctx.record(
+ action=decision.action,
+ reason=decision.reason,
+ model=decision.target_model,
+ applied=False,
+ decision_mode="observe",
+ )
+ self._pre_recorded = True
+ return
+
+ if run_ctx.mode != "enforce":
+ return
+
+ if decision.action == "stop":
+ run_ctx.record(
+ action="stop",
+ reason=decision.reason,
+ model=model_name,
+ applied=True,
+ decision_mode="enforce",
+ )
+ self._pre_recorded = True
+ _raise_stop_error(run_ctx, decision.reason)
+
+ if decision.action == "switch_model":
+ run_ctx.record(
+ action="switch_model",
+ reason=decision.reason,
+ model=decision.target_model,
+ applied=False,
+ decision_mode="enforce",
+ )
+ self._pre_recorded = True
+
+ if decision.action == "deny_tool" and has_tools:
+ run_ctx.record(
+ action="deny_tool",
+ reason=decision.reason,
+ model=model_name,
+ applied=False,
+ decision_mode="enforce",
+ )
+ self._pre_recorded = True
+
+ except Exception as exc:
+ self._handle_harness_error(exc)
+
+ def on_llm_end(self, response: Any, **kwargs: Any) -> None:
+ super().on_llm_end(response=response, **kwargs)
+
+ try:
+ self._sync_state(kwargs)
+ run_ctx = get_current_run()
+ if run_ctx is None:
+ return
+
+ model_name = self.current_model
+ if not model_name and getattr(response, "llm_output", None):
+ model_name = response.llm_output.get("model_name")
+ model_name = model_name or "unknown"
+
+ token_usage = extract_token_usage(response)
+ prompt_tokens = int(token_usage["input"])
+ completion_tokens = int(token_usage["output"])
+ elapsed_ms = 0.0
+ if self._llm_started_at is not None:
+ elapsed_ms = (time.monotonic() - self._llm_started_at) * 1000.0
+
+ run_ctx.step_count += 1
+ run_ctx.cost += estimate_cost(model_name, prompt_tokens, completion_tokens)
+ run_ctx.energy_used += estimate_energy(model_name, prompt_tokens, completion_tokens)
+ run_ctx.latency_used_ms += elapsed_ms
+
+ if run_ctx.budget_max is not None:
+ run_ctx.budget_remaining = run_ctx.budget_max - run_ctx.cost
+
+ if self._pre_action == "allow":
+ run_ctx.record(
+ action="allow",
+ reason="langchain_step",
+ model=model_name,
+ applied=True,
+ decision_mode=run_ctx.mode,
+ )
+ elif not self._pre_recorded:
+ run_ctx.record(
+ action=self._pre_action,
+ reason=self._pre_reason,
+ model=self._pre_model or model_name,
+ applied=False,
+ decision_mode=run_ctx.mode,
+ )
+
+ except Exception as exc:
+ self._handle_harness_error(exc)
+ finally:
+ self._llm_started_at = None
+ self._pre_action = "allow"
+ self._pre_reason = "allow"
+ self._pre_model = None
+ self._pre_recorded = False
+
+ def on_tool_start(self, serialized: dict[str, Any], input_str: str, **kwargs: Any) -> Any:
+ try:
+ self._sync_state(kwargs)
+ run_ctx = get_current_run()
+ if run_ctx is None:
+ return None
+ if run_ctx.tool_calls_max is None:
+ return None
+
+ if run_ctx.tool_calls >= run_ctx.tool_calls_max:
+ if run_ctx.mode == "observe":
+ run_ctx.record(
+ action="deny_tool",
+ reason="max_tool_calls_reached",
+ model=self.current_model,
+ applied=False,
+ decision_mode="observe",
+ )
+ return None
+ if run_ctx.mode == "enforce":
+ run_ctx.record(
+ action="deny_tool",
+ reason="max_tool_calls_reached",
+ model=self.current_model,
+ applied=True,
+ decision_mode="enforce",
+ )
+ raise HarnessStopError(
+ "cascadeflow harness deny_tool: max tool calls reached",
+ reason="max_tool_calls_reached",
+ )
+
+ # Track executed tools (not predicted tool calls in LLM output).
+ run_ctx.tool_calls += 1
+ return None
+ except Exception as exc:
+ self._handle_harness_error(exc)
+ return None
+
+
+@contextmanager
+def get_harness_callback(*, fail_open: bool = True):
+ """Context manager that yields a harness-aware LangChain callback handler."""
+ callback = HarnessAwareCascadeFlowCallbackHandler(fail_open=fail_open)
+ yield callback
+
+
+__all__ = ["HarnessAwareCascadeFlowCallbackHandler", "get_harness_callback"]
diff --git a/cascadeflow/integrations/langchain/harness_state.py b/cascadeflow/integrations/langchain/harness_state.py
new file mode 100644
index 00000000..b4b40da5
--- /dev/null
+++ b/cascadeflow/integrations/langchain/harness_state.py
@@ -0,0 +1,124 @@
+"""LangGraph/LangChain state extraction helpers for harness integration."""
+
+from __future__ import annotations
+
+from typing import Any, Mapping, Optional
+
+
+def _as_int(value: Any) -> Optional[int]:
+ try:
+ if value is None:
+ return None
+ return int(value)
+ except (TypeError, ValueError):
+ return None
+
+
+def _as_float(value: Any) -> Optional[float]:
+ try:
+ if value is None:
+ return None
+ return float(value)
+ except (TypeError, ValueError):
+ return None
+
+
+def _extract_candidate_state(source: Any) -> Optional[Mapping[str, Any]]:
+ """Extract a named state container from a mapping.
+
+ Only returns state from explicitly named keys (langgraph_state, graph_state,
+ state). Returns None when no named key matches — avoids treating arbitrary
+ kwargs as harness state.
+ """
+ if not isinstance(source, Mapping):
+ return None
+
+ for key in ("langgraph_state", "graph_state", "state"):
+ candidate = source.get(key)
+ if isinstance(candidate, Mapping):
+ return candidate
+
+ return None
+
+
+def extract_langgraph_state(payload: Any) -> dict[str, Any]:
+ """Extract normalized harness-relevant fields from LangGraph-style state payloads."""
+
+ candidates: list[Mapping[str, Any]] = []
+ root = _extract_candidate_state(payload)
+ if root is not None:
+ candidates.append(root)
+
+ if isinstance(payload, Mapping):
+ metadata = payload.get("metadata")
+ if isinstance(metadata, Mapping):
+ state_from_metadata = _extract_candidate_state(metadata)
+ if state_from_metadata is not None:
+ candidates.append(state_from_metadata)
+
+ configurable = payload.get("configurable")
+ if isinstance(configurable, Mapping):
+ state_from_configurable = _extract_candidate_state(configurable)
+ if state_from_configurable is not None:
+ candidates.append(state_from_configurable)
+
+ merged: dict[str, Any] = {}
+ for source in candidates:
+ if "agent_id" in source and isinstance(source.get("agent_id"), str):
+ merged["agent_id"] = source["agent_id"]
+ if "model" in source and isinstance(source.get("model"), str):
+ merged["model_used"] = source["model"]
+ if "model_used" in source and isinstance(source.get("model_used"), str):
+ merged["model_used"] = source["model_used"]
+
+ step_count = _as_int(source.get("step_count", source.get("step")))
+ if step_count is not None:
+ merged["step_count"] = step_count
+
+ tool_calls = _as_int(source.get("tool_calls"))
+ if tool_calls is not None:
+ merged["tool_calls"] = tool_calls
+
+ budget_remaining = _as_float(source.get("budget_remaining"))
+ if budget_remaining is not None:
+ merged["budget_remaining"] = budget_remaining
+
+ latency_used_ms = _as_float(source.get("latency_used_ms", source.get("latency_ms")))
+ if latency_used_ms is not None:
+ merged["latency_used_ms"] = latency_used_ms
+
+ energy_used = _as_float(source.get("energy_used", source.get("energy")))
+ if energy_used is not None:
+ merged["energy_used"] = energy_used
+
+ return merged
+
+
+def apply_langgraph_state(run_ctx: Any, state: Mapping[str, Any]) -> None:
+ """Apply extracted state fields onto an active HarnessRunContext."""
+ if run_ctx is None or not isinstance(state, Mapping):
+ return
+
+ step_count = _as_int(state.get("step_count"))
+ if step_count is not None and step_count > getattr(run_ctx, "step_count", 0):
+ run_ctx.step_count = step_count
+
+ tool_calls = _as_int(state.get("tool_calls"))
+ if tool_calls is not None and tool_calls > getattr(run_ctx, "tool_calls", 0):
+ run_ctx.tool_calls = tool_calls
+
+ latency_used_ms = _as_float(state.get("latency_used_ms"))
+ if latency_used_ms is not None and latency_used_ms > getattr(run_ctx, "latency_used_ms", 0.0):
+ run_ctx.latency_used_ms = latency_used_ms
+
+ energy_used = _as_float(state.get("energy_used"))
+ if energy_used is not None and energy_used > getattr(run_ctx, "energy_used", 0.0):
+ run_ctx.energy_used = energy_used
+
+ budget_remaining = _as_float(state.get("budget_remaining"))
+ if budget_remaining is not None:
+ run_ctx.budget_remaining = budget_remaining
+
+ model_used = state.get("model_used")
+ if isinstance(model_used, str) and model_used:
+ run_ctx.model_used = model_used
diff --git a/cascadeflow/integrations/langchain/tests/test_langchain_harness_callback.py b/cascadeflow/integrations/langchain/tests/test_langchain_harness_callback.py
new file mode 100644
index 00000000..9ba062e5
--- /dev/null
+++ b/cascadeflow/integrations/langchain/tests/test_langchain_harness_callback.py
@@ -0,0 +1,213 @@
+"""Tests for harness-aware LangChain callback integration."""
+
+from __future__ import annotations
+
+import pytest
+from langchain_core.messages import AIMessage
+from langchain_core.outputs import ChatGeneration, LLMResult
+
+from cascadeflow.harness import init, reset, run
+from cascadeflow.integrations.langchain.harness_callback import (
+ HarnessAwareCascadeFlowCallbackHandler,
+)
+from cascadeflow.integrations.langchain.harness_state import (
+ apply_langgraph_state,
+ extract_langgraph_state,
+)
+from cascadeflow.integrations.langchain.utils import extract_tool_calls
+from cascadeflow.schema.exceptions import BudgetExceededError, HarnessStopError
+
+
+@pytest.fixture(autouse=True)
+def _reset_harness_state() -> None:
+ reset()
+
+
+def _llm_result(model_name: str, prompt_tokens: int, completion_tokens: int) -> LLMResult:
+ generation = ChatGeneration(message=AIMessage(content="ok"), generation_info={})
+ return LLMResult(
+ generations=[[generation]],
+ llm_output={
+ "model_name": model_name,
+ "token_usage": {
+ "prompt_tokens": prompt_tokens,
+ "completion_tokens": completion_tokens,
+ "total_tokens": prompt_tokens + completion_tokens,
+ },
+ },
+ )
+
+
+def test_harness_callback_updates_active_run_metrics() -> None:
+ init(mode="observe", budget=1.0)
+ handler = HarnessAwareCascadeFlowCallbackHandler()
+
+ with run(budget=1.0) as ctx:
+ handler.on_llm_start(
+ serialized={},
+ prompts=["hello"],
+ invocation_params={"model": "gpt-4o-mini"},
+ )
+ handler.on_llm_end(_llm_result("gpt-4o-mini", 120, 80))
+
+ assert ctx.step_count == 1
+ assert ctx.cost > 0
+ assert ctx.energy_used > 0
+ assert ctx.budget_remaining is not None
+ assert ctx.budget_remaining < 1.0
+ assert ctx.last_action == "allow"
+ assert ctx.model_used == "gpt-4o-mini"
+
+
+def test_harness_callback_enforce_raises_when_budget_exhausted() -> None:
+ init(mode="enforce", budget=0.1)
+ handler = HarnessAwareCascadeFlowCallbackHandler(fail_open=False)
+
+ with run(budget=0.1) as ctx:
+ ctx.cost = 0.1
+ ctx.budget_remaining = 0.0
+
+ with pytest.raises(BudgetExceededError):
+ handler.on_llm_start(
+ serialized={},
+ prompts=["hello"],
+ invocation_params={"model": "gpt-4o-mini"},
+ )
+
+ trace = ctx.trace()
+ assert trace
+ assert trace[-1]["action"] == "stop"
+ assert trace[-1]["reason"] == "budget_exceeded"
+ assert trace[-1]["applied"] is True
+
+
+def test_harness_callback_observe_records_non_applied_decisions() -> None:
+ init(mode="observe", budget=1.0)
+ handler = HarnessAwareCascadeFlowCallbackHandler()
+
+ with run(budget=1.0) as ctx:
+ ctx.cost = 0.9
+ ctx.budget_remaining = 0.1
+
+ handler.on_llm_start(
+ serialized={},
+ prompts=["hello"],
+ invocation_params={"model": "gpt-4o", "tools": [{"name": "lookup"}]},
+ )
+
+ trace = ctx.trace()
+ assert trace
+ assert trace[-1]["action"] in {"switch_model", "deny_tool"}
+ assert trace[-1]["applied"] is False
+ assert trace[-1]["decision_mode"] == "observe"
+
+
+def test_harness_callback_enforce_denies_tool_when_limit_reached() -> None:
+ init(mode="enforce", max_tool_calls=0, budget=1.0)
+ handler = HarnessAwareCascadeFlowCallbackHandler(fail_open=False)
+
+ with run(max_tool_calls=0, budget=1.0) as ctx:
+ with pytest.raises(HarnessStopError, match="max tool calls"):
+ handler.on_tool_start(serialized={"name": "search"}, input_str="query")
+
+ trace = ctx.trace()
+ assert trace
+ assert trace[-1]["action"] == "deny_tool"
+ assert trace[-1]["applied"] is True
+ assert trace[-1]["decision_mode"] == "enforce"
+
+
+def test_on_llm_end_no_run_context_is_safe() -> None:
+ handler = HarnessAwareCascadeFlowCallbackHandler()
+ handler.on_llm_start(
+ serialized={},
+ prompts=["hello"],
+ invocation_params={"model": "gpt-4o-mini"},
+ )
+ handler.on_llm_end(_llm_result("gpt-4o-mini", 10, 5))
+
+
+def test_on_tool_start_no_run_context_is_safe() -> None:
+ handler = HarnessAwareCascadeFlowCallbackHandler()
+ handler.on_tool_start(serialized={"name": "search"}, input_str="query")
+
+
+def test_extract_state_ignores_plain_kwargs() -> None:
+ """Kwargs without a named state key should not leak into state."""
+ state = extract_langgraph_state({"model": "gpt-4o", "invocation_params": {"tools": []}})
+ assert state == {}
+
+
+def test_tool_deny_uses_run_ctx_tool_calls() -> None:
+ """Tool gating should use run_ctx.tool_calls, not a local counter."""
+ init(mode="enforce", max_tool_calls=2, budget=1.0)
+ handler = HarnessAwareCascadeFlowCallbackHandler(fail_open=False)
+
+ with run(max_tool_calls=2, budget=1.0) as ctx:
+ # Simulate tool calls already counted by on_llm_end or other integrations
+ ctx.tool_calls = 2
+
+ with pytest.raises(HarnessStopError, match="max tool calls"):
+ handler.on_tool_start(serialized={"name": "search"}, input_str="query")
+
+
+def test_tool_start_counts_executions_and_blocks_after_limit() -> None:
+ init(mode="enforce", max_tool_calls=1, budget=1.0)
+ handler = HarnessAwareCascadeFlowCallbackHandler(fail_open=False)
+
+ with run(max_tool_calls=1, budget=1.0) as ctx:
+ assert ctx.tool_calls == 0
+ assert handler.on_tool_start(serialized={"name": "search"}, input_str="first") is None
+ assert ctx.tool_calls == 1
+
+ with pytest.raises(HarnessStopError, match="max tool calls"):
+ handler.on_tool_start(serialized={"name": "search"}, input_str="second")
+
+ assert ctx.tool_calls == 1
+ trace = ctx.trace()
+ assert trace[-1]["action"] == "deny_tool"
+ assert trace[-1]["applied"] is True
+
+
+def test_extract_tool_calls_supports_llm_result_nested_generations() -> None:
+ generation = ChatGeneration(
+ message=AIMessage(
+ content="", tool_calls=[{"name": "search", "args": {"q": "x"}, "id": "t1"}]
+ ),
+ generation_info={},
+ )
+ llm_result = LLMResult(generations=[[generation]], llm_output={"model_name": "gpt-4o-mini"})
+ tool_calls = extract_tool_calls(llm_result)
+ assert len(tool_calls) == 1
+ assert tool_calls[0]["name"] == "search"
+
+
+def test_extract_and_apply_langgraph_state() -> None:
+ state = extract_langgraph_state(
+ {
+ "metadata": {
+ "langgraph_state": {
+ "step": 4,
+ "tool_calls": 3,
+ "budget_remaining": 0.42,
+ "latency_ms": 130.0,
+ "energy": 77.0,
+ "model": "gpt-4o-mini",
+ }
+ }
+ }
+ )
+
+ assert state["step_count"] == 4
+ assert state["tool_calls"] == 3
+ assert state["model_used"] == "gpt-4o-mini"
+
+ init(mode="observe", budget=1.0)
+ with run(budget=1.0) as ctx:
+ apply_langgraph_state(ctx, state)
+ assert ctx.step_count == 4
+ assert ctx.tool_calls == 3
+ assert ctx.budget_remaining == pytest.approx(0.42)
+ assert ctx.latency_used_ms == pytest.approx(130.0)
+ assert ctx.energy_used == pytest.approx(77.0)
+ assert ctx.model_used == "gpt-4o-mini"
diff --git a/cascadeflow/integrations/langchain/tests/test_langchain_integration_features.py b/cascadeflow/integrations/langchain/tests/test_langchain_integration_features.py
index fdbcff1d..0f051519 100644
--- a/cascadeflow/integrations/langchain/tests/test_langchain_integration_features.py
+++ b/cascadeflow/integrations/langchain/tests/test_langchain_integration_features.py
@@ -4,7 +4,11 @@
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
from langchain_core.outputs import ChatGeneration, ChatResult
+from cascadeflow.harness import init, reset, run
from cascadeflow.integrations.langchain import CascadeFlow
+from cascadeflow.integrations.langchain.harness_callback import (
+ HarnessAwareCascadeFlowCallbackHandler,
+)
class MockSequenceChatModel(BaseChatModel):
@@ -116,3 +120,38 @@ def test_domain_policy_direct_to_verifier_skips_drafter() -> None:
assert drafter.calls == 0
assert verifier.calls == 1
assert result.llm_output["cascade"]["routing_reason"] == "domain_policy_direct"
+
+
+def test_wrapper_only_auto_adds_harness_callback_inside_active_run_scope() -> None:
+ reset()
+ init(mode="observe")
+ drafter = MockSequenceChatModel("draft")
+ verifier = MockSequenceChatModel("verify")
+ cascade = CascadeFlow(drafter=drafter, verifier=verifier, enable_pre_router=False)
+
+ outside_callbacks = cascade._resolve_callbacks([])
+ assert not any(
+ isinstance(cb, HarnessAwareCascadeFlowCallbackHandler) for cb in outside_callbacks
+ )
+
+ with run():
+ inside_callbacks = cascade._resolve_callbacks([])
+ assert any(
+ isinstance(cb, HarnessAwareCascadeFlowCallbackHandler) for cb in inside_callbacks
+ )
+
+
+def test_wrapper_does_not_duplicate_harness_callback() -> None:
+ reset()
+ init(mode="observe")
+ drafter = MockSequenceChatModel("draft")
+ verifier = MockSequenceChatModel("verify")
+ cascade = CascadeFlow(drafter=drafter, verifier=verifier, enable_pre_router=False)
+ existing = HarnessAwareCascadeFlowCallbackHandler()
+
+ with run():
+ callbacks = cascade._resolve_callbacks([existing])
+ assert (
+ len([cb for cb in callbacks if isinstance(cb, HarnessAwareCascadeFlowCallbackHandler)])
+ == 1
+ )
diff --git a/cascadeflow/integrations/langchain/utils.py b/cascadeflow/integrations/langchain/utils.py
index fe47a353..04f3e4a5 100644
--- a/cascadeflow/integrations/langchain/utils.py
+++ b/cascadeflow/integrations/langchain/utils.py
@@ -195,6 +195,10 @@ def extract_tool_calls(response: Any) -> list[dict[str, Any]]:
msg = None
if hasattr(response, "generations") and response.generations:
generation = response.generations[0]
+ # LLMResult.generations is often list[list[Generation]], while ChatResult
+ # uses list[Generation]. Support both shapes.
+ if isinstance(generation, list) and generation:
+ generation = generation[0]
msg = getattr(generation, "message", None)
else:
msg = getattr(response, "message", None) or response
diff --git a/cascadeflow/integrations/langchain/wrapper.py b/cascadeflow/integrations/langchain/wrapper.py
index ed6d554b..f108d60f 100644
--- a/cascadeflow/integrations/langchain/wrapper.py
+++ b/cascadeflow/integrations/langchain/wrapper.py
@@ -169,6 +169,35 @@ def _split_runnable_config(
model_kwargs[key] = value
return model_kwargs, config
+ def _resolve_callbacks(self, raw_callbacks: Any) -> list[Any]:
+ if raw_callbacks is None:
+ callbacks: list[Any] = []
+ elif isinstance(raw_callbacks, list):
+ callbacks = list(raw_callbacks)
+ elif isinstance(raw_callbacks, tuple):
+ callbacks = list(raw_callbacks)
+ else:
+ callbacks = [raw_callbacks]
+
+ try:
+ from cascadeflow.harness import get_current_run, get_harness_config
+
+ harness_config = get_harness_config()
+ run_ctx = get_current_run()
+ if harness_config.mode == "off" or run_ctx is None or run_ctx.mode == "off":
+ return callbacks
+
+ from .harness_callback import HarnessAwareCascadeFlowCallbackHandler
+
+ if any(isinstance(cb, HarnessAwareCascadeFlowCallbackHandler) for cb in callbacks):
+ return callbacks
+
+ callbacks.append(HarnessAwareCascadeFlowCallbackHandler())
+ return callbacks
+ except Exception:
+ # Preserve existing behavior for users who do not enable harness flows.
+ return callbacks
+
def _generate(
self,
messages: list[BaseMessage],
@@ -202,7 +231,7 @@ def _generate(
merged_kwargs["stop"] = stop
# Extract callbacks before filtering (need to pass them explicitly to nested models)
- callbacks = merged_kwargs.get("callbacks", [])
+ callbacks = self._resolve_callbacks(merged_kwargs.get("callbacks", []))
existing_tags = merged_kwargs.get("tags", []) or []
base_tags = existing_tags + ["cascadeflow"] if existing_tags else ["cascadeflow"]
@@ -599,7 +628,7 @@ async def _agenerate(
merged_kwargs["stop"] = stop
# Extract callbacks before filtering (need to pass them explicitly to nested models)
- callbacks = merged_kwargs.get("callbacks", [])
+ callbacks = self._resolve_callbacks(merged_kwargs.get("callbacks", []))
existing_tags = merged_kwargs.get("tags", []) or []
base_tags = existing_tags + ["cascadeflow"] if existing_tags else ["cascadeflow"]
@@ -1001,7 +1030,7 @@ def _stream(
stream_kwargs, base_config = self._split_runnable_config(merged_kwargs)
base_tags = (base_config.get("tags") or []) + ["cascadeflow"]
existing_metadata = base_config.get("metadata", {}) or {}
- callbacks = base_config.get("callbacks", [])
+ callbacks = self._resolve_callbacks(base_config.get("callbacks", []))
resolved_domain = self._resolve_domain(messages, existing_metadata)
effective_quality_threshold = self._effective_quality_threshold(resolved_domain)
force_verifier_for_domain = self._domain_forces_verifier(resolved_domain)
@@ -1324,7 +1353,7 @@ async def _astream(
stream_kwargs, base_config = self._split_runnable_config(merged_kwargs)
base_tags = (base_config.get("tags") or []) + ["cascadeflow"]
existing_metadata = base_config.get("metadata", {}) or {}
- callbacks = base_config.get("callbacks", [])
+ callbacks = self._resolve_callbacks(base_config.get("callbacks", []))
safe_kwargs = {
k: v
for k, v in stream_kwargs.items()
diff --git a/docs-site/api-reference/python/agent-decorator.mdx b/docs-site/api-reference/python/agent-decorator.mdx
new file mode 100644
index 00000000..912a03fd
--- /dev/null
+++ b/docs-site/api-reference/python/agent-decorator.mdx
@@ -0,0 +1,79 @@
+---
+title: "@cascadeflow.agent()"
+description: Decorate agent functions with policy metadata including budget, compliance, and KPI weights.
+---
+
+# @cascadeflow.agent()
+
+Annotate agent functions with policy metadata. The decorator attaches budget, compliance, and KPI configuration to the function for the harness to use at runtime.
+
+## Signature
+
+```python
+def agent(
+ budget: Optional[float] = None,
+ compliance: Optional[str] = None,
+ kpi_weights: Optional[dict[str, float]] = None,
+ kpi_targets: Optional[dict[str, float]] = None,
+ max_tool_calls: Optional[int] = None,
+)
+```
+
+## Parameters
+
+| Parameter | Type | Default | Description |
+|---|---|---|---|
+| `budget` | `float \| None` | `None` | Max USD for this agent |
+| `compliance` | `str \| None` | `None` | Compliance mode |
+| `kpi_weights` | `dict \| None` | `None` | KPI dimension weights |
+| `kpi_targets` | `dict \| None` | `None` | KPI dimension targets |
+| `max_tool_calls` | `int \| None` | `None` | Max tool/function calls |
+
+## Usage
+
+### Basic
+
+```python
+@cascadeflow.agent(budget=0.20)
+async def my_agent(query: str):
+ return await llm.complete(query)
+```
+
+### With compliance
+
+```python
+@cascadeflow.agent(budget=0.50, compliance="gdpr")
+async def eu_agent(query: str):
+ return await llm.complete(query)
+```
+
+### With KPI weights
+
+```python
+@cascadeflow.agent(
+ budget=1.00,
+ kpi_weights={"quality": 0.8, "cost": 0.2},
+ kpi_targets={"quality": 0.9},
+)
+async def premium_agent(query: str):
+ return await llm.complete(query)
+```
+
+### Multiple agents with different policies
+
+```python
+@cascadeflow.agent(budget=0.10, kpi_weights={"cost": 0.9, "quality": 0.1})
+async def triage_agent(query: str):
+ return await llm.complete(query)
+
+@cascadeflow.agent(budget=2.00, kpi_weights={"quality": 0.9, "cost": 0.1})
+async def analysis_agent(query: str):
+ return await llm.complete(query)
+```
+
+## Notes
+
+- The decorator does not wrap or modify the function's execution. It attaches metadata that the harness reads at runtime.
+- Works with both sync and async functions.
+- Requires `init()` to have been called for the metadata to take effect.
+- Can be combined with `run()` — the run's constraints are checked in addition to the decorator's.
diff --git a/docs-site/api-reference/python/harness-config.mdx b/docs-site/api-reference/python/harness-config.mdx
new file mode 100644
index 00000000..42ae7a6d
--- /dev/null
+++ b/docs-site/api-reference/python/harness-config.mdx
@@ -0,0 +1,73 @@
+---
+title: HarnessConfig
+description: Full configuration dataclass for the cascadeflow harness with all fields, types, and defaults.
+---
+
+# HarnessConfig
+
+Configuration dataclass for the cascadeflow harness. Pass to `cascadeflow.init(config=...)` for full control.
+
+## Definition
+
+```python
+from dataclasses import dataclass
+from typing import Optional
+
+@dataclass
+class HarnessConfig:
+ mode: HarnessMode = "off"
+ verbose: bool = False
+ budget: Optional[float] = None
+ max_tool_calls: Optional[int] = None
+ max_latency_ms: Optional[float] = None
+ max_energy: Optional[float] = None
+ kpi_targets: Optional[dict[str, float]] = None
+ kpi_weights: Optional[dict[str, float]] = None
+ compliance: Optional[str] = None
+```
+
+## Fields
+
+| Field | Type | Default | Description |
+|---|---|---|---|
+| `mode` | `"off" \| "observe" \| "enforce"` | `"off"` | Harness mode |
+| `verbose` | `bool` | `False` | Print decisions to stderr |
+| `budget` | `float \| None` | `None` | Max USD for the run (None = unlimited) |
+| `max_tool_calls` | `int \| None` | `None` | Max tool/function calls (None = unlimited) |
+| `max_latency_ms` | `float \| None` | `None` | Max wall-clock ms per call (None = unlimited) |
+| `max_energy` | `float \| None` | `None` | Max energy units (None = unlimited) |
+| `kpi_targets` | `dict \| None` | `None` | Target values per KPI dimension |
+| `kpi_weights` | `dict \| None` | `None` | Relative weights per KPI dimension |
+| `compliance` | `str \| None` | `None` | Compliance mode: `"gdpr"`, `"hipaa"`, `"pci"`, `"strict"` |
+
+## HarnessMode
+
+```python
+HarnessMode = Literal["off", "observe", "enforce"]
+```
+
+## Usage
+
+```python
+from cascadeflow import HarnessConfig
+import cascadeflow
+
+config = HarnessConfig(
+ mode="enforce",
+ budget=1.00,
+ max_tool_calls=20,
+ max_energy=200.0,
+ compliance="gdpr",
+ kpi_weights={"quality": 0.6, "cost": 0.3, "latency": 0.1},
+ kpi_targets={"quality": 0.85},
+ verbose=True,
+)
+
+cascadeflow.init(config=config)
+```
+
+## Import
+
+```python
+from cascadeflow import HarnessConfig
+```
diff --git a/docs-site/api-reference/python/init.mdx b/docs-site/api-reference/python/init.mdx
new file mode 100644
index 00000000..b07a0e00
--- /dev/null
+++ b/docs-site/api-reference/python/init.mdx
@@ -0,0 +1,68 @@
+---
+title: cascadeflow.init()
+description: Activate the cascadeflow harness globally with a mode and optional configuration.
+---
+
+# cascadeflow.init()
+
+Activate the harness globally. All subsequent LLM calls (OpenAI, Anthropic) are automatically tracked.
+
+## Signature
+
+```python
+def init(
+ mode: HarnessMode = "off",
+ *,
+ config: Optional[HarnessConfig] = None,
+ verbose: bool = False,
+) -> HarnessInitReport
+```
+
+## Parameters
+
+| Parameter | Type | Default | Description |
+|---|---|---|---|
+| `mode` | `"off" \| "observe" \| "enforce"` | `"off"` | Harness mode |
+| `config` | `HarnessConfig \| None` | `None` | Full configuration (overrides mode) |
+| `verbose` | `bool` | `False` | Print decisions to stderr |
+
+## Returns
+
+`HarnessInitReport` — confirmation of harness activation with mode and configuration summary.
+
+## Usage
+
+### Minimal
+
+```python
+import cascadeflow
+cascadeflow.init(mode="observe")
+```
+
+### With config
+
+```python
+from cascadeflow import HarnessConfig
+
+config = HarnessConfig(
+ mode="enforce",
+ budget=1.00,
+ compliance="gdpr",
+ verbose=True,
+)
+cascadeflow.init(config=config)
+```
+
+### Environment-driven
+
+```python
+import os
+cascadeflow.init(mode=os.getenv("CASCADEFLOW_MODE", "observe"))
+```
+
+## Notes
+
+- Call `init()` once at application startup, before any LLM calls
+- Calling `init()` again replaces the previous configuration
+- Use `cascadeflow.reset()` to deactivate the harness
+- `init(mode="off")` is equivalent to not calling `init()` at all
diff --git a/docs-site/api-reference/python/run-context.mdx b/docs-site/api-reference/python/run-context.mdx
new file mode 100644
index 00000000..be9377a4
--- /dev/null
+++ b/docs-site/api-reference/python/run-context.mdx
@@ -0,0 +1,76 @@
+---
+title: HarnessRunContext
+description: Run context object yielded by cascadeflow.run() with summary(), trace(), and budget tracking methods.
+---
+
+# HarnessRunContext
+
+The context object yielded by `cascadeflow.run()`. Provides access to run metrics, decision traces, and budget state.
+
+## Methods
+
+### summary()
+
+Returns aggregate metrics for the run.
+
+```python
+summary = session.summary()
+```
+
+Returns a dict with:
+
+| Key | Type | Description |
+|---|---|---|
+| `cost_total` | `float` | Cumulative cost in USD |
+| `steps` | `int` | Number of LLM calls |
+| `tool_calls` | `int` | Number of tool/function calls |
+| `latency_total_ms` | `float` | Total wall-clock latency in ms |
+| `energy_used` | `float` | Total energy units consumed |
+| `budget_remaining` | `float \| None` | USD remaining (None if no budget set) |
+
+### trace()
+
+Returns the list of decision records for the run.
+
+```python
+records = session.trace()
+```
+
+Each record is a dict with:
+
+| Key | Type | Description |
+|---|---|---|
+| `action` | `str` | `"allow"`, `"switch_model"`, `"deny_tool"`, or `"stop"` |
+| `reason` | `str` | Human-readable explanation |
+| `model` | `str` | Model name |
+| `step` | `int` | Step number (1-indexed) |
+| `cost_total` | `float` | Cumulative cost at this step |
+| `budget_state` | `str` | `"ok"`, `"warning"`, or `"exceeded"` |
+| `applied` | `bool` | Whether the action was enforced |
+
+## Usage
+
+```python
+import cascadeflow
+
+cascadeflow.init(mode="enforce")
+
+with cascadeflow.run(budget=0.50) as session:
+ result = await agent.run("Analyze this dataset")
+
+ # Aggregate metrics
+ summary = session.summary()
+ print(f"Cost: ${summary['cost_total']:.4f}")
+ print(f"Steps: {summary['steps']}")
+ print(f"Budget remaining: ${summary['budget_remaining']:.4f}")
+
+ # Decision trace
+ for record in session.trace():
+ print(f"Step {record['step']}: {record['action']} — {record['reason']}")
+```
+
+## Import
+
+```python
+from cascadeflow import HarnessRunContext
+```
diff --git a/docs-site/api-reference/python/run.mdx b/docs-site/api-reference/python/run.mdx
new file mode 100644
index 00000000..72202a74
--- /dev/null
+++ b/docs-site/api-reference/python/run.mdx
@@ -0,0 +1,83 @@
+---
+title: cascadeflow.run()
+description: Create a scoped run context with budget caps, tool call limits, and metrics tracking.
+---
+
+# cascadeflow.run()
+
+Create a scoped run context manager that tracks metrics and optionally enforces constraints for a block of agent execution.
+
+## Signature
+
+```python
+def run(
+ budget: Optional[float] = None,
+ max_tool_calls: Optional[int] = None,
+ max_latency_ms: Optional[float] = None,
+ max_energy: Optional[float] = None,
+ compliance: Optional[str] = None,
+ kpi_weights: Optional[dict[str, float]] = None,
+ kpi_targets: Optional[dict[str, float]] = None,
+) -> ContextManager[HarnessRunContext]
+```
+
+## Parameters
+
+| Parameter | Type | Default | Description |
+|---|---|---|---|
+| `budget` | `float \| None` | `None` | Max USD for this run |
+| `max_tool_calls` | `int \| None` | `None` | Max tool/function calls |
+| `max_latency_ms` | `float \| None` | `None` | Max wall-clock ms per call |
+| `max_energy` | `float \| None` | `None` | Max energy units |
+| `compliance` | `str \| None` | `None` | `"gdpr"`, `"hipaa"`, `"pci"`, or `"strict"` |
+| `kpi_weights` | `dict \| None` | `None` | KPI dimension weights |
+| `kpi_targets` | `dict \| None` | `None` | KPI dimension targets |
+
+## Returns
+
+Context manager yielding `HarnessRunContext`. See [HarnessRunContext](/api-reference/python/run-context).
+
+## Usage
+
+### Basic budget
+
+```python
+with cascadeflow.run(budget=0.50) as session:
+ result = await agent.run("Analyze this data")
+ print(session.summary())
+```
+
+### Full configuration
+
+```python
+with cascadeflow.run(
+ budget=1.00,
+ max_tool_calls=10,
+ max_energy=100.0,
+ compliance="gdpr",
+ kpi_weights={"quality": 0.6, "cost": 0.3, "latency": 0.1},
+ kpi_targets={"quality": 0.9},
+) as session:
+ result = await agent.run("Process EU customer data")
+ print(session.summary())
+ for record in session.trace():
+ print(f"Step {record['step']}: {record['action']}")
+```
+
+### Nested runs
+
+Runs can be nested. Inner runs inherit the parent's remaining budget:
+
+```python
+with cascadeflow.run(budget=1.00) as outer:
+ with cascadeflow.run(budget=0.30) as inner:
+ await agent.run("Sub-task")
+ # outer.summary() includes inner costs
+```
+
+## Notes
+
+- `run()` requires `init()` to have been called first
+- Parameters override the global config for the duration of the block
+- Use `session.summary()` for aggregate metrics
+- Use `session.trace()` for per-step decision records
diff --git a/docs-site/api-reference/typescript/core.mdx b/docs-site/api-reference/typescript/core.mdx
new file mode 100644
index 00000000..ae8f8311
--- /dev/null
+++ b/docs-site/api-reference/typescript/core.mdx
@@ -0,0 +1,77 @@
+---
+title: "@cascadeflow/core"
+description: TypeScript core package with CascadeAgent for model routing, cost tracking, and quality validation.
+---
+
+# @cascadeflow/core
+
+The core TypeScript package for cascadeflow. Provides `CascadeAgent` for speculative model cascading with quality validation.
+
+## Install
+
+```bash
+npm install @cascadeflow/core
+```
+
+## CascadeAgent
+
+```typescript
+import { CascadeAgent, ModelConfig } from '@cascadeflow/core';
+
+const agent = new CascadeAgent({
+ models: [
+ { name: 'gpt-4o-mini', provider: 'openai', cost: 0.000375 },
+ { name: 'gpt-4o', provider: 'openai', cost: 0.00625 },
+ ],
+});
+
+const result = await agent.run('What is TypeScript?');
+console.log(`Model: ${result.modelUsed}`);
+console.log(`Cost: $${result.totalCost}`);
+console.log(`Saved: ${result.savingsPercentage}%`);
+```
+
+## ModelConfig
+
+```typescript
+interface ModelConfig {
+ name: string; // Model name (e.g. 'gpt-4o-mini')
+ provider: string; // Provider name (e.g. 'openai')
+ cost: number; // Cost per token (approximate)
+}
+```
+
+## CascadeAgentOptions
+
+```typescript
+interface CascadeAgentOptions {
+ models: ModelConfig[];
+ quality?: {
+ threshold?: number; // Confidence threshold (0-1)
+ requireMinimumTokens?: number; // Min response length
+ useSemanticValidation?: boolean; // Enable ML validation
+ semanticThreshold?: number; // Semantic similarity threshold
+ };
+}
+```
+
+## Result
+
+```typescript
+interface CascadeResult {
+ content: string;
+ modelUsed: string;
+ totalCost: number;
+ savingsPercentage: number;
+ cascadeDecision: string;
+}
+```
+
+## Features
+
+- Speculative execution with quality validation
+- Multi-provider support (OpenAI, Anthropic, Groq, Ollama, vLLM)
+- Streaming responses
+- Tool calling and structured output
+- Cost tracking and analytics
+- Works in Node.js, Browser, and Edge Functions
diff --git a/docs-site/api-reference/typescript/langchain.mdx b/docs-site/api-reference/typescript/langchain.mdx
new file mode 100644
index 00000000..9a9e3050
--- /dev/null
+++ b/docs-site/api-reference/typescript/langchain.mdx
@@ -0,0 +1,77 @@
+---
+title: "@cascadeflow/langchain"
+description: TypeScript LangChain integration with withCascade() for drop-in cascade routing and model discovery helpers.
+---
+
+# @cascadeflow/langchain
+
+LangChain integration for TypeScript. Provides `withCascade()` for drop-in cascade routing with any LangChain chat model.
+
+## Install
+
+```bash
+npm install @cascadeflow/langchain @langchain/core @langchain/openai
+```
+
+## withCascade
+
+Creates a cascade-enabled chat model from a drafter and verifier.
+
+```typescript
+import { ChatOpenAI } from '@langchain/openai';
+import { ChatAnthropic } from '@langchain/anthropic';
+import { withCascade } from '@cascadeflow/langchain';
+
+const cascade = withCascade({
+ drafter: new ChatOpenAI({ model: 'gpt-4o-mini' }),
+ verifier: new ChatAnthropic({ model: 'claude-sonnet-4' }),
+ qualityThreshold: 0.8,
+});
+
+// Use like any LangChain chat model
+const result = await cascade.invoke('Explain quantum computing');
+
+// With LCEL chains
+const chain = prompt.pipe(cascade).pipe(new StringOutputParser());
+```
+
+## Options
+
+```typescript
+interface CascadeOptions {
+ drafter: BaseChatModel; // Cheap, fast model
+ verifier: BaseChatModel; // Powerful fallback model
+ qualityThreshold?: number; // 0-1, default 0.4
+}
+```
+
+## Model Discovery
+
+```typescript
+import {
+ discoverCascadePairs,
+ findBestCascadePair,
+ analyzeModel,
+ validateCascadePair,
+} from '@cascadeflow/langchain';
+
+const models = [
+ new ChatOpenAI({ model: 'gpt-4o-mini' }),
+ new ChatOpenAI({ model: 'gpt-4o' }),
+ new ChatAnthropic({ model: 'claude-sonnet-4' }),
+];
+
+const best = findBestCascadePair(models);
+const cascade = withCascade({
+ drafter: best.drafter,
+ verifier: best.verifier,
+});
+```
+
+## Features
+
+- Full LCEL support (pipes, sequences, batch)
+- Streaming with pre-routing
+- Tool calling and structured output
+- LangSmith cost tracking metadata
+- Model discovery and pair validation
diff --git a/docs-site/api-reference/typescript/vercel-ai.mdx b/docs-site/api-reference/typescript/vercel-ai.mdx
new file mode 100644
index 00000000..ae9af949
--- /dev/null
+++ b/docs-site/api-reference/typescript/vercel-ai.mdx
@@ -0,0 +1,63 @@
+---
+title: "@cascadeflow/vercel-ai"
+description: Vercel AI SDK middleware integration for cascade routing with streaming, multi-turn chat, and tool execution.
+---
+
+# @cascadeflow/vercel-ai
+
+Middleware integration for the Vercel AI SDK. Adds cascade routing to AI SDK applications with streaming support.
+
+## Install
+
+```bash
+npm install @cascadeflow/vercel-ai
+```
+
+## createChatHandler
+
+Creates a request handler for AI SDK chat endpoints.
+
+```typescript
+import { createChatHandler } from '@cascadeflow/vercel-ai';
+import { CascadeAgent } from '@cascadeflow/core';
+
+const agent = new CascadeAgent({
+ models: [
+ { name: 'gpt-4o-mini', provider: 'openai', cost: 0.000375 },
+ { name: 'gpt-4o', provider: 'openai', cost: 0.00625 },
+ ],
+});
+
+const handler = createChatHandler(agent, {
+ protocol: 'data',
+ tools,
+ toolHandlers,
+ maxSteps: 5,
+});
+```
+
+## Options
+
+```typescript
+interface ChatHandlerOptions {
+ protocol: 'data' | 'ui'; // AI SDK stream protocol
+ tools?: ToolDefinition[]; // Tool definitions
+ toolHandlers?: Record