From 41444871858156e8e708a1862321e950ed788c76 Mon Sep 17 00:00:00 2001 From: Namrata Ghadi Date: Sat, 31 Oct 2026 14:32:01 -0700 Subject: [PATCH 1/5] gateway routing based on AC policies --- .../budget/__init__.py | 2 + .../budget/config.py | 25 ++ .../budget/evaluator.py | 109 +++++- .../gateway_controls.json | 88 +++++ models/src/agent_control_models/__init__.py | 19 ++ models/src/agent_control_models/actions.py | 11 +- models/src/agent_control_models/controls.py | 14 + models/src/agent_control_models/gateway.py | 101 ++++++ models/tests/test_actions.py | 24 +- server/README.md | 100 ++++++ server/pyproject.toml | 2 +- .../agent_control_server/endpoints/gateway.py | 309 ++++++++++++++++++ server/src/agent_control_server/main.py | 6 + server/tests/test_config.py | 46 +++ server/tests/test_gateway.py | 277 ++++++++++++++++ 15 files changed, 1111 insertions(+), 22 deletions(-) create mode 100644 examples/inference_gateway_poc/gateway_controls.json create mode 100644 models/src/agent_control_models/gateway.py create mode 100644 server/src/agent_control_server/endpoints/gateway.py create mode 100644 server/tests/test_gateway.py diff --git a/evaluators/contrib/budget/src/agent_control_evaluator_budget/budget/__init__.py b/evaluators/contrib/budget/src/agent_control_evaluator_budget/budget/__init__.py index c747c4f83..f8437ab82 100644 --- a/evaluators/contrib/budget/src/agent_control_evaluator_budget/budget/__init__.py +++ b/evaluators/contrib/budget/src/agent_control_evaluator_budget/budget/__init__.py @@ -3,6 +3,7 @@ from agent_control_evaluator_budget.budget.config import ( BudgetEvaluatorConfig, BudgetLimitRule, + GatewayRoutingConfig, ModelPricing, ) from agent_control_evaluator_budget.budget.evaluator import BudgetEvaluator @@ -17,6 +18,7 @@ "BudgetEvaluator", "BudgetEvaluatorConfig", "BudgetLimitRule", + "GatewayRoutingConfig", "BudgetSnapshot", "BudgetStore", "InMemoryBudgetStore", diff --git a/evaluators/contrib/budget/src/agent_control_evaluator_budget/budget/config.py b/evaluators/contrib/budget/src/agent_control_evaluator_budget/budget/config.py index 795044be4..a4d161953 100644 --- a/evaluators/contrib/budget/src/agent_control_evaluator_budget/budget/config.py +++ b/evaluators/contrib/budget/src/agent_control_evaluator_budget/budget/config.py @@ -5,6 +5,7 @@ from typing import Literal from agent_control_evaluators._base import EvaluatorConfig +from agent_control_models import GatewayModelTier from pydantic import Field, field_validator, model_validator # --------------------------------------------------------------------------- @@ -68,6 +69,25 @@ def validate_window_seconds(cls, v: int | None) -> int | None: return v +class GatewayRoutingConfig(EvaluatorConfig): + """Optional model-tier routing behavior stored on a budget control.""" + + tiers: dict[str, GatewayModelTier] = Field(min_length=1) + task_tiers: dict[str, str] = Field(min_length=1) + threshold: float = Field(default=0.8, gt=0.0, le=1.0) + threshold_tier: str | None = None + + @model_validator(mode="after") + def validate_tier_references(self) -> GatewayRoutingConfig: + configured = set(self.tiers) + unknown = set(self.task_tiers.values()) - configured + if unknown: + raise ValueError(f"task_tiers references unknown tiers: {sorted(unknown)}") + if self.threshold_tier is not None and self.threshold_tier not in configured: + raise ValueError(f"threshold_tier references unknown tier: {self.threshold_tier}") + return self + + class BudgetEvaluatorConfig(EvaluatorConfig): """Configuration for the budget evaluator. @@ -109,6 +129,7 @@ class BudgetEvaluatorConfig(EvaluatorConfig): token_path: str | None = None model_path: str | None = None metadata_paths: dict[str, str] = Field(default_factory=dict) + gateway_routing: GatewayRoutingConfig | None = None @model_validator(mode="after") def require_pricing_for_cost_rules(self) -> BudgetEvaluatorConfig: @@ -117,4 +138,8 @@ def require_pricing_for_cost_rules(self) -> BudgetEvaluatorConfig: raise ValueError('pricing is required when any rule uses limit_unit="usd_cents"') if has_cost_rule and not (self.model_path or "").strip(): raise ValueError('model_path is required when any rule uses limit_unit="usd_cents"') + if self.gateway_routing is not None and not any( + rule.limit_unit == "tokens" for rule in self.limits + ): + raise ValueError("gateway_routing requires at least one token budget limit") return self diff --git a/evaluators/contrib/budget/src/agent_control_evaluator_budget/budget/evaluator.py b/evaluators/contrib/budget/src/agent_control_evaluator_budget/budget/evaluator.py index c43809031..b28db158c 100644 --- a/evaluators/contrib/budget/src/agent_control_evaluator_budget/budget/evaluator.py +++ b/evaluators/contrib/budget/src/agent_control_evaluator_budget/budget/evaluator.py @@ -20,9 +20,9 @@ from agent_control_evaluators._base import Evaluator, EvaluatorMetadata from agent_control_evaluators._registry import register_evaluator -from agent_control_models import EvaluatorResult +from agent_control_models import EvaluatorResult, GatewayModelCandidate -from .config import BudgetEvaluatorConfig, ModelPricing +from .config import BudgetEvaluatorConfig, GatewayRoutingConfig, ModelPricing from .memory_store import InMemoryBudgetStore, _scope_matches from .store import BudgetStore @@ -174,6 +174,92 @@ def _extract_metadata(data: Any, metadata_paths: dict[str, str]) -> dict[str, st return result +def _gateway_route_metadata( + data: Any, + routing: GatewayRoutingConfig | None, + snapshots: list[dict[str, Any]], + *, + over_limit: bool, +) -> dict[str, Any] | None: + """Select a configured candidate using task tier and accumulated token use.""" + if routing is None: + return None + request_data = data if isinstance(data, dict) else {} + task_type = request_data.get("task_type", "default") + if not isinstance(task_type, str): + task_type = "default" + budget_snapshot = next( + (snapshot for snapshot in snapshots if snapshot.get("limit_unit") == "tokens"), {} + ) + used_tokens = int(budget_snapshot.get("spent_tokens", 0)) + limit_tokens = int(budget_snapshot.get("limit", 0)) + utilization = float(budget_snapshot.get("utilization", 0.0)) + threshold_reached = utilization >= routing.threshold + task_tier = routing.task_tiers.get(task_type, routing.task_tiers.get("default")) + selected_tier = ( + routing.threshold_tier + if threshold_reached and routing.threshold_tier is not None + else task_tier + ) + reason = "" + candidate: GatewayModelCandidate | None = None + if over_limit: + reason = "The configured token budget is exhausted." + elif selected_tier is None: + reason = f"No model tier is configured for task type '{task_type}'." + else: + raw_candidates = request_data.get("candidate_models", []) + candidates: set[tuple[str, str, str]] = set() + legacy_model = request_data.get("model") + if isinstance(raw_candidates, list): + for raw_candidate in raw_candidates: + if isinstance(raw_candidate, dict): + try: + parsed = GatewayModelCandidate.model_validate(raw_candidate) + except ValueError: + continue + candidates.add((parsed.provider, parsed.model, parsed.protocol)) + if not candidates and isinstance(legacy_model, str): + candidates = { + (configured.provider, configured.model, configured.protocol) + for configured in routing.tiers[selected_tier].models + if configured.model == legacy_model + } + candidate = next( + ( + configured + for configured in routing.tiers[selected_tier].models + if (configured.provider, configured.model, configured.protocol) in candidates + and (not isinstance(legacy_model, str) or configured.model == legacy_model) + ), + None, + ) + reason = ( + f"Budget threshold reached; routed to the '{selected_tier}' tier." + if threshold_reached + else f"Task type '{task_type}' routed to the '{selected_tier}' tier." + ) if candidate is not None else ( + f"No compatible model candidate is available in tier '{selected_tier}'." + ) + return { + "allowed": candidate is not None and not over_limit, + "provider": candidate.provider if candidate is not None and not over_limit else None, + "model": candidate.model if candidate is not None and not over_limit else None, + "protocol": candidate.protocol if candidate is not None and not over_limit else None, + "task_type": task_type, + "tier": selected_tier, + "budget": { + "used_tokens": used_tokens, + "limit_tokens": limit_tokens, + "utilization": min(utilization, 1.0), + "threshold": routing.threshold, + "threshold_reached": threshold_reached, + "over_limit": over_limit, + }, + "reason": reason, + } + + # --------------------------------------------------------------------------- # Evaluator # --------------------------------------------------------------------------- @@ -281,6 +367,9 @@ async def evaluate(self, data: Any) -> EvaluatorResult: if breached: first = breached[0] + route_metadata = _gateway_route_metadata( + data, self.config.gateway_routing, all_snaps, over_limit=True + ) return EvaluatorResult( matched=True, confidence=1.0, @@ -291,19 +380,31 @@ async def evaluate(self, data: Any) -> EvaluatorResult: "input_tokens": input_tokens, "output_tokens": output_tokens, "cost": round(cost, 6), + **({"gateway_route": route_metadata} if route_metadata is not None else {}), }, ) max_util = max((s["utilization"] for s in all_snaps), default=0.0) + route_metadata = _gateway_route_metadata( + data, self.config.gateway_routing, all_snaps, over_limit=False + ) return EvaluatorResult( - matched=False, + matched=( + self.config.gateway_routing is not None + or bool(route_metadata and route_metadata.get("allowed")) + ), confidence=1.0, - message=f"Within budget (utilization={max_util:.0%})", + message=( + str(route_metadata.get("reason")) + if route_metadata + else f"Within budget (utilization={max_util:.0%})" + ), metadata={ "all_snapshots": all_snaps, "input_tokens": input_tokens, "output_tokens": output_tokens, "cost": round(cost, 6), "max_utilization": round(max_util, 4), + **({"gateway_route": route_metadata} if route_metadata is not None else {}), }, ) diff --git a/examples/inference_gateway_poc/gateway_controls.json b/examples/inference_gateway_poc/gateway_controls.json new file mode 100644 index 000000000..93348428d --- /dev/null +++ b/examples/inference_gateway_poc/gateway_controls.json @@ -0,0 +1,88 @@ +{ + "customer-success-agent": { + "name": "customer-success-agent-gateway-routing", + "binding": {"target_type": "agent", "target_id": "customer-success-agent"}, + "data": { + "description": "Route customer-success tasks by tier and accumulated token budget", + "enabled": true, + "execution": "server", + "scope": {"step_types": ["llm"], "stages": ["pre"]}, + "condition": { + "selector": {"path": "input"}, + "evaluator": { + "name": "budget", + "config": { + "limits": [{"limit": 100000, "limit_unit": "tokens"}], + "budget_id": "gateway-demo:customer-success-agent", + "gateway_routing": { + "tiers": { + "economy": { + "models": [ + {"provider": "openai", "model": "gpt-4.1-mini", "protocol": "openai"} + ] + }, + "premium": { + "models": [ + {"provider": "openai", "model": "gpt-4.1", "protocol": "openai"} + ] + } + }, + "task_tiers": { + "default": "economy", + "classification": "economy", + "analysis": "premium" + }, + "threshold": 0.8, + "threshold_tier": "economy" + } + } + } + }, + "action": {"decision": "route"}, + "tags": ["gateway"] + } + }, + "claude-code": { + "name": "claude-code-gateway-routing", + "binding": {"target_type": "agent", "target_id": "claude-code"}, + "data": { + "description": "Route Claude Code only to Anthropic-compatible models", + "enabled": true, + "execution": "server", + "scope": {"step_types": ["llm"], "stages": ["pre"]}, + "condition": { + "selector": {"path": "input"}, + "evaluator": { + "name": "budget", + "config": { + "limits": [{"limit": 250000, "limit_unit": "tokens"}], + "budget_id": "gateway-demo:claude-code", + "gateway_routing": { + "tiers": { + "economy": { + "models": [ + {"provider": "anthropic", "model": "claude-haiku-4-5", "protocol": "anthropic"} + ] + }, + "premium": { + "models": [ + {"provider": "anthropic", "model": "claude-sonnet-4-5", "protocol": "anthropic"} + ] + } + }, + "task_tiers": { + "default": "premium", + "code-review": "premium", + "simple-edit": "economy" + }, + "threshold": 0.9, + "threshold_tier": "economy" + } + } + } + }, + "action": {"decision": "route"}, + "tags": ["gateway"] + } + } +} diff --git a/models/src/agent_control_models/__init__.py b/models/src/agent_control_models/__init__.py index 6b5562a73..da3470e05 100644 --- a/models/src/agent_control_models/__init__.py +++ b/models/src/agent_control_models/__init__.py @@ -65,6 +65,16 @@ EvaluationResponse, EvaluationResult, ) +from .gateway import ( + GatewayBudgetState, + GatewayModelCandidate, + GatewayModelTier, + GatewayPreflightRequest, + GatewayPreflightResponse, + GatewayRouteDecisionResponse, + GatewayUsageSettlementRequest, + GatewayUsageSettlementResponse, +) from .health import HealthResponse from .observability import ( BatchEventsRequest, @@ -141,6 +151,15 @@ "EvaluationRequest", "EvaluationResponse", "EvaluationResult", + # Inference gateway + "GatewayModelCandidate", + "GatewayModelTier", + "GatewayPreflightRequest", + "GatewayPreflightResponse", + "GatewayRouteDecisionResponse", + "GatewayUsageSettlementRequest", + "GatewayUsageSettlementResponse", + "GatewayBudgetState", # Controls "ControlDefinition", "ConditionNode", diff --git a/models/src/agent_control_models/actions.py b/models/src/agent_control_models/actions.py index 972f50087..c148ca5cc 100644 --- a/models/src/agent_control_models/actions.py +++ b/models/src/agent_control_models/actions.py @@ -5,14 +5,15 @@ from collections.abc import Sequence from typing import Literal, cast -type ActionDecision = Literal["deny", "steer", "observe"] +type ActionDecision = Literal["deny", "steer", "observe", "route"] -_CANONICAL_ACTIONS = frozenset({"deny", "steer", "observe"}) +_CANONICAL_ACTIONS = frozenset({"deny", "steer", "observe", "route"}) _OBSERVE_ACTION_ALIASES = frozenset({"allow", "observe", "warn", "log"}) _ACTION_QUERY_EXPANSION: dict[ActionDecision, tuple[str, ...]] = { "deny": ("deny",), "steer": ("steer",), "observe": ("observe", "allow", "warn", "log"), + "route": ("route",), } @@ -25,7 +26,7 @@ def validate_action(action: str) -> ActionDecision: if action in _CANONICAL_ACTIONS: return cast(ActionDecision, action) raise ValueError( - f"Invalid action {action!r}. Must be one of: deny, steer, observe." + f"Invalid action {action!r}. Must be one of: deny, steer, observe, route." ) @@ -50,10 +51,10 @@ def normalize_action(action: str) -> ActionDecision: """ if action in _OBSERVE_ACTION_ALIASES: return "observe" - if action in ("deny", "steer"): + if action in ("deny", "steer", "route"): return cast(ActionDecision, action) raise ValueError( - f"Invalid action {action!r}. Expected one of: deny, steer, observe." + f"Invalid action {action!r}. Expected one of: deny, steer, observe, route." ) diff --git a/models/src/agent_control_models/controls.py b/models/src/agent_control_models/controls.py index 46e8fd8ec..66541b9cd 100644 --- a/models/src/agent_control_models/controls.py +++ b/models/src/agent_control_models/controls.py @@ -542,6 +542,20 @@ def _validate_common_control_constraints( ): raise ValueError("Composite steer controls require action.steering_context") + leaf_parts = condition.leaf_parts() if condition.is_leaf() else None + has_gateway_budget_evaluator = bool( + leaf_parts is not None + and leaf_parts[1].name == "budget" + and "gateway_routing" in leaf_parts[1].config + ) + if action.decision == "route" and not has_gateway_budget_evaluator: + raise ValueError( + "The route action requires a single budget evaluator with " + "gateway_routing configured" + ) + if has_gateway_budget_evaluator and action.decision != "route": + raise ValueError("A budget evaluator with gateway_routing requires the route action") + class ConditionNode(BaseModel): """Recursive boolean condition tree for control evaluation.""" diff --git a/models/src/agent_control_models/gateway.py b/models/src/agent_control_models/gateway.py new file mode 100644 index 000000000..16a8b8920 --- /dev/null +++ b/models/src/agent_control_models/gateway.py @@ -0,0 +1,101 @@ +"""Typed contracts for Agent Control's inference gateway POC.""" + +from typing import Literal, Self + +from pydantic import Field, field_validator, model_validator + +from .actions import ActionDecision +from .agent import AGENT_NAME_MIN_LENGTH, AGENT_NAME_PATTERN, JSONObject, normalize_agent_name +from .base import BaseModel + +GatewayProtocol = Literal["openai", "anthropic"] + + +class GatewayModelCandidate(BaseModel): + """A model and wire protocol the gateway can route to.""" + + provider: str = Field(min_length=1) + model: str = Field(min_length=1) + protocol: GatewayProtocol + + +class GatewayModelTier(BaseModel): + """Ordered model candidates configured for one capability or budget tier.""" + + models: list[GatewayModelCandidate] = Field(min_length=1) + + +class GatewayPreflightRequest(BaseModel): + """Candidate routes and task information supplied before inference.""" + + # Legacy field retained for old POC clients. Policy lookup still comes + # from the authenticated route target, and the server checks this value. + agent_name: str | None = Field( + default=None, + min_length=AGENT_NAME_MIN_LENGTH, + pattern=AGENT_NAME_PATTERN, + ) + task_type: str = Field(default="default", min_length=1) + task_context: JSONObject | None = None + candidate_models: list[GatewayModelCandidate] = Field(default_factory=list) + # Compatibility input for clients using the original single-model POC. + model: str | None = Field(default=None, min_length=1) + prompt: str | None = None + + @field_validator("agent_name", mode="before") + @classmethod + def normalize_legacy_agent_name(cls, value: str | None) -> str | None: + return None if value is None else normalize_agent_name(value) + + @model_validator(mode="after") + def require_candidate(self) -> Self: + if not self.candidate_models and self.model is None: + raise ValueError("Provide candidate_models or the legacy model field.") + return self + + +class GatewayUsageSettlementRequest(BaseModel): + """Actual token usage reported by the gateway after inference.""" + + model: str = Field(min_length=1) + input_tokens: int = Field(ge=0) + output_tokens: int = Field(ge=0) + + +class GatewayBudgetState(BaseModel): + """Accumulated token use and configured budget threshold state.""" + + used_tokens: int = Field(ge=0) + limit_tokens: int = Field(gt=0) + utilization: float = Field(ge=0.0, le=1.0) + threshold: float = Field(gt=0.0, le=1.0) + threshold_reached: bool + over_limit: bool + + +class GatewayPreflightResponse(BaseModel): + """Control and budget decision returned before inference.""" + + allowed: bool + policy: str + budget: GatewayBudgetState | None + reason: str + + +class GatewayRouteDecisionResponse(GatewayPreflightResponse): + """Route result directly consumable by an inference gateway adapter.""" + + provider: str | None + model: str | None + protocol: GatewayProtocol | None + task_type: str + decision: Literal["route", "deny"] + action: ActionDecision | None = None + + +class GatewayUsageSettlementResponse(BaseModel): + """Budget state after actual usage has been accumulated.""" + + policy: str + budget: GatewayBudgetState + reason: str diff --git a/models/tests/test_actions.py b/models/tests/test_actions.py index 2b0c57759..13cdc421b 100644 --- a/models/tests/test_actions.py +++ b/models/tests/test_actions.py @@ -7,14 +7,13 @@ ControlAction, ControlExecutionEvent, ControlMatch, - EventQueryRequest, EvaluatorResult, + EventQueryRequest, expand_action_filter, ) from agent_control_models.actions import normalize_action, validate_action from pydantic import ValidationError - # --------------------------------------------------------------------------- # validate_action (strict, for API boundaries) # --------------------------------------------------------------------------- @@ -23,7 +22,7 @@ class TestValidateAction: """Tests for the strict validate_action used on public API boundaries.""" - @pytest.mark.parametrize("action", ["deny", "steer", "observe"]) + @pytest.mark.parametrize("action", ["deny", "steer", "observe", "route"]) def test_accepts_canonical_actions(self, action: str) -> None: # Given: a canonical action name # When: validating the action @@ -54,7 +53,7 @@ def test_rejects_unknown_action(self) -> None: class TestNormalizeAction: """Tests for the lenient normalize_action used on read paths.""" - @pytest.mark.parametrize("action", ["deny", "steer", "observe"]) + @pytest.mark.parametrize("action", ["deny", "steer", "observe", "route"]) def test_passes_canonical_actions(self, action: str) -> None: # Given: a canonical action name # When: normalizing @@ -87,7 +86,7 @@ def test_rejects_unknown_action(self) -> None: class TestControlActionValidation: """ControlAction.decision uses strict validation (rejects legacy values).""" - @pytest.mark.parametrize("action", ["deny", "steer", "observe"]) + @pytest.mark.parametrize("action", ["deny", "steer", "observe", "route"]) def test_accepts_canonical_actions(self, action: str) -> None: # Given: a control action payload with a canonical decision # When: validating via Pydantic @@ -234,17 +233,18 @@ def test_observe_expands_to_include_legacy(self) -> None: # Then: it includes all legacy advisory action values stored in old rows assert expanded == ["observe", "allow", "warn", "log"] - def test_deny_and_steer_do_not_expand(self) -> None: - # Given: deny and steer filters (no legacy aliases) + def test_deny_steer_and_route_do_not_expand(self) -> None: + # Given: canonical actions without legacy aliases # When: expanding - # Then: they map only to themselves + # Then: each maps only to itself assert expand_action_filter(["deny"]) == ["deny"] assert expand_action_filter(["steer"]) == ["steer"] + assert expand_action_filter(["route"]) == ["route"] def test_full_expansion(self) -> None: - # Given: all three canonical actions + # Given: all canonical actions # When: expanding - expanded = expand_action_filter(["deny", "steer", "observe"]) + expanded = expand_action_filter(["deny", "steer", "observe", "route"]) - # Then: deny and steer are unchanged, observe expands to include legacy - assert expanded == ["deny", "steer", "observe", "allow", "warn", "log"] + # Then: only observe expands to include legacy values + assert expanded == ["deny", "steer", "observe", "allow", "warn", "log", "route"] diff --git a/server/README.md b/server/README.md index 7c56e4462..f66438d7a 100644 --- a/server/README.md +++ b/server/README.md @@ -7,6 +7,7 @@ FastAPI server that powers Agent Control. It manages agents and controls, evalua - Agent registration and control association - Control CRUD and evaluator configuration - Runtime evaluation (`/api/v1/evaluation`) with pre/post stages +- Authenticated inference gateway preflight, route decision, and usage settlement - Observability endpoints for events and stats - API key authentication for production deployments @@ -56,4 +57,103 @@ psycopg's bundled binary package. Server configuration is driven by environment variables (database, auth, observability, evaluators). For the full list and examples, see the docs. +### Inference gateway POC API + +Create a regular Agent Control with the budget evaluator, set gateway_routing +in its evaluator config, choose the route action, and bind the control to the +agent through the existing control-binding flow. Routing configuration is +stored with the control and uses the same budget evaluator store as usage +settlement. Example payloads for the distinct customer-success-agent and +claude-code controls are in +[examples/inference_gateway_poc/gateway_controls.json](../examples/inference_gateway_poc/gateway_controls.json). +The route action is accepted only with a single budget evaluator configured +with gateway_routing. It is non-blocking in ordinary ControlEngine safety +results; the gateway maps its matched evaluator metadata to a route or deny. +The Claude Code route is restricted at runtime to the Anthropic provider and +protocol. The request path must match the authenticated runtime identity; the +caller cannot choose another agent's bound controls in the body. + +`POST /api/v1/gateway/agents/{agent_name}/route-decision` combines control +preflight and tier selection. A gateway that needs a separate eligibility check +can call `/preflight` with the same request first. The route path carries the +agent identity; the optional body `agent_name` is retained only for legacy +clients and must match it. A route-decision request looks like: + +```json +{ + "task_type": "analysis", + "task_context": {"classification": "multi-step reasoning"}, + "candidate_models": [ + {"provider": "openai", "model": "gpt-4.1-mini", "protocol": "openai"}, + {"provider": "openai", "model": "gpt-4.1", "protocol": "openai"} + ] +} +``` + +For that request, the route response has this shape: + +```json +{ + "allowed": true, + "policy": "customer-success-agent-gateway-routing", + "budget": { + "used_tokens": 0, + "limit_tokens": 100000, + "utilization": 0.0, + "threshold": 0.8, + "threshold_reached": false, + "over_limit": false + }, + "reason": "Task type 'analysis' routed to the 'premium' tier.", + "provider": "openai", + "model": "gpt-4.1", + "protocol": "openai", + "task_type": "analysis", + "decision": "route", + "action": "route" +} +``` + +`budget.used_tokens` and `budget.limit_tokens` are token counts. `utilization` +and `threshold` are ratios from 0 to 1; this API does not report USD cost. A +denied decision sets `provider`, `model`, and `protocol` to null. The legacy +single-model `model` request field and optional body `agent_name` remain +accepted. The route response returns decision=route only when a bound +gateway-routing control selects a candidate. Denials return decision=deny and +null provider/model/protocol values. + +After the provider responds, post the actual usage to +`POST /api/v1/gateway/agents/{agent_name}/usage-settlement`: + +```json +{ + "model": "gpt-4.1", + "input_tokens": 40000, + "output_tokens": 20000 +} +``` + +The settlement response uses the same budget fields and identifies the bound +control: + +```json +{ + "policy": "customer-success-agent-gateway-routing", + "budget": { + "used_tokens": 60000, + "limit_tokens": 100000, + "utilization": 0.6, + "threshold": 0.8, + "threshold_reached": false, + "over_limit": false + }, + "reason": "Usage settled against the authenticated agent budget." +} +``` + +Budget snapshots and settlement use the budget evaluator's process-local +store. The demo records actual token usage after inference; it does not reserve +tokens across concurrent requests, so requests already in flight can exceed the +final configured budget. + Full guide: https://docs.agentcontrol.dev/components/server diff --git a/server/pyproject.toml b/server/pyproject.toml index 82b5a2af8..fdbe225b4 100644 --- a/server/pyproject.toml +++ b/server/pyproject.toml @@ -24,7 +24,7 @@ dependencies = [ "jsonschema-rs>=0.22.0", "PyJWT>=2.8.0", "google-re2>=1.1", # For engine (bundled) - "agent-control-evaluators>=8.6.0", # NOT vendored - avoid conflict with galileo + "agent-control-evaluators[budget]>=8.6.0", # NOT vendored - avoid conflict with galileo ] authors = [ {name = "Agent Control Team"} diff --git a/server/src/agent_control_server/endpoints/gateway.py b/server/src/agent_control_server/endpoints/gateway.py new file mode 100644 index 000000000..2d42a2164 --- /dev/null +++ b/server/src/agent_control_server/endpoints/gateway.py @@ -0,0 +1,309 @@ +"""Inference gateway policy endpoints for the serial POC.""" + +from __future__ import annotations + +from typing import Literal + +from agent_control_engine.core import ControlEngine +from agent_control_evaluator_budget.budget import ( # type: ignore[import-untyped] + BudgetEvaluator, + BudgetEvaluatorConfig, +) +from agent_control_models import ( + ActionDecision, + EvaluationRequest, + EvaluationResponse, + GatewayBudgetState, + GatewayModelCandidate, + GatewayPreflightRequest, + GatewayPreflightResponse, + GatewayRouteDecisionResponse, + GatewayUsageSettlementRequest, + GatewayUsageSettlementResponse, + JSONObject, + Step, +) +from agent_control_models.errors import ErrorCode +from fastapi import APIRouter, Depends, Request +from sqlalchemy import select + +from ..auth_framework import Operation, Principal, require_operation +from ..db import AsyncSessionLocal +from ..errors import ForbiddenError, NotFoundError +from ..models import Agent +from ..services.controls import ControlService +from .evaluation import ControlAdapter + +router = APIRouter(prefix="/gateway/agents", tags=["inference-gateway"]) + + +def _agent_context(request: Request) -> dict[str, str]: + """Bind runtime authorization to the agent named by the route.""" + return {"target_type": "agent", "target_id": request.path_params["agent_name"]} + + +_runtime_agent = require_operation(Operation.RUNTIME_USE, context_builder=_agent_context) + + +async def _get_gateway_controls(agent_name: str, principal: Principal) -> list[ControlAdapter]: + """Load controls effective for this authenticated runtime identity.""" + if principal.target_type != "agent" or principal.target_id != agent_name: + raise ForbiddenError( + error_code=ErrorCode.AUTH_INSUFFICIENT_PRIVILEGES, + detail="Gateway requests must use a runtime identity bound to this agent.", + ) + async with AsyncSessionLocal() as db: + agent_result = await db.execute( + select(Agent).where( + Agent.name == agent_name, + Agent.namespace_key == principal.namespace_key, + ) + ) + if agent_result.scalar_one_or_none() is None: + raise NotFoundError( + error_code=ErrorCode.AGENT_NOT_FOUND, + detail=f"Agent '{agent_name}' not found", + resource="Agent", + resource_id=agent_name, + ) + runtime_controls = await ControlService(db).list_runtime_controls_for_agent( + agent_name, + namespace_key=principal.namespace_key, + target_type="agent", + target_id=agent_name, + allow_invalid_step_name_regex=True, + ) + return [ + ControlAdapter(control.id, control.name, control.control) + for control in runtime_controls + ] + + +def _gateway_budget_control( + controls: list[ControlAdapter], +) -> tuple[ControlAdapter, BudgetEvaluatorConfig] | None: + """Find an existing bound budget control configured for gateway routing.""" + for control in controls: + for _, evaluator_spec in control.control.iter_condition_leaf_parts(): + if evaluator_spec.name != "budget" or "gateway_routing" not in evaluator_spec.config: + continue + config = BudgetEvaluatorConfig.model_validate(evaluator_spec.config) + return control, config + return None + + +def _budget_state_from_metadata(metadata: object) -> GatewayBudgetState | None: + if not isinstance(metadata, dict): + return None + route = metadata.get("gateway_route", metadata) + if not isinstance(route, dict): + return None + budget = route.get("budget") + if not isinstance(budget, dict): + return None + return GatewayBudgetState.model_validate(budget) + + +async def _evaluate_controls( + agent_name: str, + request: GatewayPreflightRequest, + principal: Principal, +) -> EvaluationResponse: + """Evaluate current agent controls, including enabled agent target bindings.""" + controls = await _get_gateway_controls(agent_name, principal) + input_data: JSONObject = { + "task_type": request.task_type, + "task_context": request.task_context, + "prompt": request.prompt, + "candidate_models": [ + { + "provider": candidate.provider, + "model": candidate.model, + "protocol": candidate.protocol, + } + for candidate in request.candidate_models + ], + } + if request.model is not None: + input_data["model"] = request.model + evaluation = EvaluationRequest( + agent_name=agent_name, + stage="pre", + step=Step( + type="llm", + name="inference-gateway-preflight", + input=input_data, + output=None, + context=None, + tools=None, + ground_truth=None, + children=None, + ), + ) + return await ControlEngine(controls).process(evaluation) + + +def _find_gateway_route( + evaluation: EvaluationResponse, +) -> tuple[str, dict[str, object], ActionDecision] | None: + """Extract gateway routing metadata from the normal control match flow.""" + for match in evaluation.matches or []: + if match.action != "route": + continue + route = (match.result.metadata or {}).get("gateway_route") + if isinstance(route, dict): + return match.control_name, route, match.action + return None + + +async def _decide( + agent_name: str, + request: GatewayPreflightRequest, + principal: Principal, +) -> tuple[GatewayPreflightResponse, GatewayModelCandidate | None, ActionDecision | None]: + if request.agent_name is not None and request.agent_name != agent_name: + raise ForbiddenError( + error_code=ErrorCode.AUTH_INSUFFICIENT_PRIVILEGES, + detail="Request agent_name must match the authenticated agent target.", + ) + evaluation = await _evaluate_controls(agent_name, request, principal) + route_result = _find_gateway_route(evaluation) + policy_name, route, action = route_result if route_result is not None else (None, {}, None) + budget = _budget_state_from_metadata(route) + candidate: GatewayModelCandidate | None = None + provider = route.get("provider") + model = route.get("model") + protocol = route.get("protocol") + route_protocol: Literal["openai", "anthropic"] | None + if protocol == "openai": + route_protocol = "openai" + elif protocol == "anthropic": + route_protocol = "anthropic" + else: + route_protocol = None + if isinstance(provider, str) and isinstance(model, str) and route_protocol is not None: + candidate = GatewayModelCandidate( + provider=provider, model=model, protocol=route_protocol + ) + allowed = bool(route.get("allowed", False)) and evaluation.is_safe and candidate is not None + if agent_name == "claude-code" and candidate is not None: + allowed = ( + allowed + and candidate.provider == "anthropic" + and candidate.protocol == "anthropic" + ) + if not allowed: + candidate = None + denied_match = next( + (match for match in evaluation.matches or [] if match.action == "deny"), None + ) + if denied_match is not None: + policy_name = denied_match.control_name + action = denied_match.action + reason_value = route.get("reason") + reason = ( + reason_value + if isinstance(reason_value, str) + else "No gateway routing evaluator is configured on a bound control." + ) + if denied_match is not None or not evaluation.is_safe: + reason = ( + f"Control '{denied_match.control_name}' denied the inference request." + if denied_match is not None + else evaluation.reason or "An Agent Control control denied this inference request." + ) + elif not allowed and agent_name == "claude-code": + reason = "Claude Code routes must use an Anthropic-compatible provider and protocol." + if not allowed: + return ( + GatewayPreflightResponse( + allowed=False, + policy=policy_name or "agent-controls", + budget=budget, + reason=reason, + ), + None, + action, + ) + return ( + GatewayPreflightResponse( + allowed=True, + policy=policy_name or "agent-controls", + budget=budget, + reason=reason, + ), + candidate, + action, + ) + + +@router.post("/{agent_name}/preflight", response_model=GatewayPreflightResponse) +async def preflight( + request: GatewayPreflightRequest, + agent_name: str, + principal: Principal = Depends(_runtime_agent), +) -> GatewayPreflightResponse: + """Evaluate controls and policy eligibility before inference.""" + response, _, _ = await _decide(agent_name, request, principal) + return response + + +@router.post("/{agent_name}/route-decision", response_model=GatewayRouteDecisionResponse) +async def route_decision( + request: GatewayPreflightRequest, + agent_name: str, + principal: Principal = Depends(_runtime_agent), +) -> GatewayRouteDecisionResponse: + """Select an allowed model for the authenticated agent and classified task.""" + response, candidate, action = await _decide(agent_name, request, principal) + return GatewayRouteDecisionResponse( + **response.model_dump(), + provider=candidate.provider if candidate is not None else None, + model=candidate.model if candidate is not None else None, + protocol=candidate.protocol if candidate is not None else None, + task_type=request.task_type, + decision="route" if candidate is not None and response.allowed else "deny", + action=action, + ) + + +@router.post("/{agent_name}/usage-settlement", response_model=GatewayUsageSettlementResponse) +async def usage_settlement( + request: GatewayUsageSettlementRequest, + agent_name: str, + principal: Principal = Depends(_runtime_agent), +) -> GatewayUsageSettlementResponse: + """Accumulate completed inference usage against the authenticated agent budget.""" + controls = await _get_gateway_controls(agent_name, principal) + budget_control = _gateway_budget_control(controls) + if budget_control is None: + raise NotFoundError( + error_code=ErrorCode.AGENT_NOT_FOUND, + detail="No bound budget evaluator with gateway routing is configured for this agent.", + resource="GatewayRoutingControl", + resource_id=agent_name, + ) + control, config = budget_control + result = await BudgetEvaluator(config).evaluate( + { + "input_tokens": request.input_tokens, + "output_tokens": request.output_tokens, + "model": request.model, + } + ) + metadata = result.metadata or {} + budget = _budget_state_from_metadata(metadata) + if budget is None: + raise NotFoundError( + error_code=ErrorCode.AGENT_NOT_FOUND, + detail="The bound budget evaluator did not return gateway budget state.", + resource="GatewayRoutingControl", + resource_id=control.name, + ) + over_limit = budget.over_limit + reason = ( + "Usage settled; the authenticated agent budget is exhausted." + if over_limit + else "Usage settled against the authenticated agent budget." + ) + return GatewayUsageSettlementResponse(policy=control.name, budget=budget, reason=reason) diff --git a/server/src/agent_control_server/main.py b/server/src/agent_control_server/main.py index bd0b9efdc..ae7e7fed5 100644 --- a/server/src/agent_control_server/main.py +++ b/server/src/agent_control_server/main.py @@ -32,6 +32,7 @@ from .endpoints.controls import template_router as control_template_router from .endpoints.evaluation import router as evaluation_router from .endpoints.evaluators import router as evaluator_router +from .endpoints.gateway import router as gateway_router from .endpoints.observability import router as observability_router from .endpoints.policies import router as policy_router from .endpoints.system import router as system_router @@ -356,6 +357,11 @@ async def attach_version_header(request, call_next): # type: ignore[no-untyped- prefix=api_v1_prefix, dependencies=[Depends(get_api_key_from_header)], ) +app.include_router( + gateway_router, + prefix=api_v1_prefix, + dependencies=[Depends(get_api_key_from_header)], +) app.include_router( evaluator_router, diff --git a/server/tests/test_config.py b/server/tests/test_config.py index 7cbb31860..70495c3de 100644 --- a/server/tests/test_config.py +++ b/server/tests/test_config.py @@ -1,11 +1,57 @@ """Tests for server configuration helpers.""" +import pytest +from agent_control_evaluator_budget.budget import BudgetEvaluatorConfig from agent_control_server.config import ( AgentControlServerDatabaseConfig, LoggingSettings, ObservabilitySettings, Settings, ) +from pydantic import ValidationError + + +def _gateway_budget_config() -> dict[str, object]: + return { + "limits": [{"limit": 1000, "limit_unit": "tokens"}], + "budget_id": "gateway-config-test", + "gateway_routing": { + "tiers": { + "economy": { + "models": [ + {"provider": "openai", "model": "small", "protocol": "openai"} + ] + } + }, + "task_tiers": {"default": "economy"}, + "threshold": 0.8, + "threshold_tier": "economy", + }, + } + + +def test_budget_evaluator_parses_gateway_routing_configuration() -> None: + config = BudgetEvaluatorConfig.model_validate(_gateway_budget_config()) + assert config.gateway_routing is not None + assert config.gateway_routing.tiers["economy"].models[0].model == "small" + + +def test_budget_evaluator_rejects_unknown_gateway_tier() -> None: + raw = _gateway_budget_config() + routing = raw["gateway_routing"] + assert isinstance(routing, dict) + routing["task_tiers"] = {"default": "missing"} + with pytest.raises(ValidationError, match="unknown tiers"): + BudgetEvaluatorConfig.model_validate(raw) + + +def test_gateway_routing_requires_token_budget_limit() -> None: + raw = _gateway_budget_config() + raw["limits"] = [{"limit": 50, "limit_unit": "usd_cents"}] + raw["pricing"] = {"small": {"input_per_1k": 0.1, "output_per_1k": 0.2}} + raw["model_path"] = "model" + with pytest.raises(ValidationError, match="requires at least one token budget limit"): + BudgetEvaluatorConfig.model_validate(raw) def test_db_config_prefers_explicit_url() -> None: diff --git a/server/tests/test_gateway.py b/server/tests/test_gateway.py new file mode 100644 index 000000000..ff10bf7aa --- /dev/null +++ b/server/tests/test_gateway.py @@ -0,0 +1,277 @@ +"""Gateway routing tests through existing bound controls and evaluators.""" + +from __future__ import annotations + +from typing import Any + +import pytest +from agent_control_evaluator_budget.budget.evaluator import clear_budget_stores +from agent_control_models import ControlMatch, EvaluationResponse, EvaluatorResult +from agent_control_server.auth_framework import Operation, Principal, set_authorizer +from agent_control_server.endpoints.gateway import _find_gateway_route +from fastapi.testclient import TestClient + + +def _routing_config(agent_name: str) -> dict[str, Any]: + if agent_name == "claude-code": + provider = protocol = "anthropic" + economy_model, premium_model = "claude-haiku-4-5", "claude-sonnet-4-5" + else: + provider = protocol = "openai" + economy_model, premium_model = "gpt-4.1-mini", "gpt-4.1" + return { + "limits": [{"limit": 100, "limit_unit": "tokens"}], + "budget_id": f"gateway-demo:{agent_name}", + "gateway_routing": { + "tiers": { + "economy": { + "models": [ + {"provider": provider, "model": economy_model, "protocol": protocol} + ] + }, + "premium": { + "models": [ + {"provider": provider, "model": premium_model, "protocol": protocol} + ] + }, + }, + "task_tiers": { + "default": "economy", + "classification": "economy", + "analysis": "premium", + "code-review": "premium", + }, + "threshold": 0.8, + "threshold_tier": "economy", + }, + } + + +def _create_control(client: TestClient, name: str, data: dict[str, Any]) -> int: + response = client.put("/api/v1/controls", json={"name": name, "data": data}) + assert response.status_code == 200, response.text + return int(response.json()["control_id"]) + + +def _bind_control(client: TestClient, agent_name: str, control_id: int) -> None: + response = client.put( + "/api/v1/control-bindings", + json={"target_type": "agent", "target_id": agent_name, "control_id": control_id}, + ) + assert response.status_code == 200, response.text + + +def _create_route_control(client: TestClient, agent_name: str) -> int: + data = { + "description": "Task-sensitive gateway model routing and token budget", + "enabled": True, + "execution": "server", + "scope": {"step_types": ["llm"], "stages": ["pre"]}, + "condition": { + "selector": {"path": "input"}, + "evaluator": {"name": "budget", "config": _routing_config(agent_name)}, + }, + "action": {"decision": "route"}, + "tags": ["gateway"], + } + return _create_control(client, f"{agent_name}-gateway-routing", data) + + +class _AgentRuntimeAuthorizer: + """Test authorizer that returns a target-bound runtime identity.""" + + def __init__(self, authenticated_agent: str) -> None: + self.authenticated_agent = authenticated_agent + + async def authorize(self, request, operation, context=None): + assert operation is Operation.RUNTIME_USE + assert context == { + "target_type": "agent", + "target_id": request.path_params["agent_name"], + } + return Principal( + namespace_key="default", + target_type="agent", + target_id=self.authenticated_agent, + scopes=(Operation.RUNTIME_USE.value,), + ) + + +def _register_agent(client: TestClient, agent_name: str) -> None: + response = client.post( + "/api/v1/agents/initAgent", + json={"agent": {"agent_name": agent_name}, "steps": []}, + ) + assert response.status_code == 200, response.text + _bind_control(client, agent_name, _create_route_control(client, agent_name)) + + +def _candidates(agent_name: str) -> list[dict[str, str]]: + if agent_name == "claude-code": + return [ + {"provider": "anthropic", "model": "claude-haiku-4-5", "protocol": "anthropic"}, + {"provider": "anthropic", "model": "claude-sonnet-4-5", "protocol": "anthropic"}, + {"provider": "openai", "model": "gpt-4.1", "protocol": "openai"}, + ] + return [ + {"provider": "openai", "model": "gpt-4.1-mini", "protocol": "openai"}, + {"provider": "openai", "model": "gpt-4.1", "protocol": "openai"}, + {"provider": "anthropic", "model": "claude-sonnet-4-5", "protocol": "anthropic"}, + ] + + +def _route(client: TestClient, agent_name: str, task_type: str) -> dict[str, Any]: + response = client.post( + f"/api/v1/gateway/agents/{agent_name}/route-decision", + json={"task_type": task_type, "candidate_models": _candidates(agent_name)}, + ) + assert response.status_code == 200, response.text + return response.json() + + +@pytest.fixture(autouse=True) +def reset_gateway_budgets() -> None: + clear_budget_stores() + yield + clear_budget_stores() + + +def test_bound_budget_control_selects_task_tier_and_threshold_route(client: TestClient) -> None: + agent_name = "customer-success-agent" + _register_agent(client, agent_name) + set_authorizer(_AgentRuntimeAuthorizer(agent_name), operation=Operation.RUNTIME_USE) + + premium = _route(client, agent_name, "analysis") + settlement = client.post( + f"/api/v1/gateway/agents/{agent_name}/usage-settlement", + json={"model": "gpt-4.1", "input_tokens": 50, "output_tokens": 30}, + ) + threshold = _route(client, agent_name, "analysis") + + assert premium["decision"] == "route" + assert premium["action"] == "route" + assert premium["model"] == "gpt-4.1" + assert premium["policy"] == f"{agent_name}-gateway-routing" + assert settlement.status_code == 200, settlement.text + assert settlement.json()["budget"]["used_tokens"] == 80 + assert threshold["model"] == "gpt-4.1-mini" + assert threshold["budget"]["threshold_reached"] is True + + +def test_budget_exhaustion_returns_deny_without_a_usable_route(client: TestClient) -> None: + agent_name = "customer-success-agent" + _register_agent(client, agent_name) + set_authorizer(_AgentRuntimeAuthorizer(agent_name), operation=Operation.RUNTIME_USE) + client.post( + f"/api/v1/gateway/agents/{agent_name}/usage-settlement", + json={"model": "gpt-4.1", "input_tokens": 60, "output_tokens": 40}, + ) + + result = _route(client, agent_name, "classification") + + assert result["decision"] == "deny" + assert result["allowed"] is False + assert result["provider"] is None + assert result["model"] is None + assert result["protocol"] is None + assert result["budget"]["over_limit"] is True + + +def test_controls_keep_separate_agent_budgets_and_claude_stays_anthropic( + client: TestClient, +) -> None: + for name in ("customer-success-agent", "claude-code"): + _register_agent(client, name) + + set_authorizer( + _AgentRuntimeAuthorizer("customer-success-agent"), operation=Operation.RUNTIME_USE + ) + customer = _route(client, "customer-success-agent", "analysis") + customer_settlement = client.post( + "/api/v1/gateway/agents/customer-success-agent/usage-settlement", + json={"model": "gpt-4.1", "input_tokens": 31, "output_tokens": 9}, + ) + set_authorizer(_AgentRuntimeAuthorizer("claude-code"), operation=Operation.RUNTIME_USE) + claude = _route(client, "claude-code", "code-review") + claude_after_customer_usage = _route(client, "claude-code", "code-review") + claude_openai_only = client.post( + "/api/v1/gateway/agents/claude-code/route-decision", + json={ + "task_type": "code-review", + "candidate_models": [ + {"provider": "openai", "model": "gpt-4.1", "protocol": "openai"} + ], + }, + ) + + assert customer["model"] == "gpt-4.1" + assert claude["provider"] == "anthropic" + assert claude["protocol"] == "anthropic" + assert claude["model"] == "claude-sonnet-4-5" + assert claude["budget"]["used_tokens"] == 0 + assert customer_settlement.status_code == 200 + assert customer_settlement.json()["budget"]["used_tokens"] == 40 + assert claude_after_customer_usage["budget"]["used_tokens"] == 0 + assert claude_openai_only.status_code == 200 + assert claude_openai_only.json()["decision"] == "deny" + assert claude_openai_only.json()["model"] is None + + +def test_bound_deny_control_blocks_gateway_route(client: TestClient) -> None: + agent_name = "customer-success-agent" + _register_agent(client, agent_name) + deny_data = { + "enabled": True, + "execution": "server", + "scope": {"step_types": ["llm"], "stages": ["pre"]}, + "condition": { + "selector": {"path": "input.task_type"}, + "evaluator": {"name": "regex", "config": {"pattern": "blocked"}}, + }, + "action": {"decision": "deny"}, + } + _bind_control(client, agent_name, _create_control(client, "gateway-block", deny_data)) + set_authorizer(_AgentRuntimeAuthorizer(agent_name), operation=Operation.RUNTIME_USE) + + result = _route(client, agent_name, "blocked") + + assert result["decision"] == "deny" + assert result["action"] == "deny" + assert result["provider"] is None + assert result["model"] is None + assert "denied" in result["reason"].lower() + + +def test_gateway_rejects_identity_mismatch(client: TestClient) -> None: + set_authorizer(_AgentRuntimeAuthorizer("claude-code"), operation=Operation.RUNTIME_USE) + response = client.post( + "/api/v1/gateway/agents/customer-success-agent/route-decision", + json={ + "task_type": "classification", + "candidate_models": _candidates("customer-success-agent"), + }, + ) + assert response.status_code == 403 + + +def test_non_matching_routing_evaluator_cannot_authorize_a_route() -> None: + non_match = ControlMatch( + control_id=1, + control_name="stale-route-result", + action="route", + result=EvaluatorResult( + matched=False, + confidence=1.0, + metadata={ + "gateway_route": { + "allowed": True, + "provider": "openai", + "model": "gpt-4.1", + "protocol": "openai", + } + }, + ), + ) + evaluation = EvaluationResponse(is_safe=True, confidence=1.0, non_matches=[non_match]) + + assert _find_gateway_route(evaluation) is None From b24bf3b3db074d0bc74eed0ac04ee0233911b92d Mon Sep 17 00:00:00 2001 From: Namrata Ghadi Date: Sat, 31 Oct 2026 15:03:36 -0700 Subject: [PATCH 2/5] fix: update gateway SDK generation and Galileo types --- .../records/factory.py | 3 +- .../overlays/method-names.overlay.yaml | 15 ++ .../funcs/inference-gateway-preflight.ts | 185 +++++++++++++++++ .../funcs/inference-gateway-route-decision.ts | 188 ++++++++++++++++++ .../inference-gateway-usage-settlement.ts | 188 ++++++++++++++++++ sdks/typescript/src/generated/hooks/hooks.ts | 3 + .../src/generated/hooks/registration.ts | 14 ++ .../src/generated/models/action-decision.ts | 1 + .../generated/models/gateway-budget-state.ts | 55 +++++ .../models/gateway-model-candidate.ts | 53 +++++ .../models/gateway-preflight-request.ts | 64 ++++++ .../models/gateway-preflight-response.ts | 44 ++++ .../models/gateway-route-decision-response.ts | 90 +++++++++ .../gateway-usage-settlement-request.ts | 50 +++++ .../gateway-usage-settlement-response.ts | 45 +++++ sdks/typescript/src/generated/models/index.ts | 7 + .../src/generated/models/operations/index.ts | 3 + ...ateway-agents-agent-name-preflight-post.ts | 46 +++++ ...y-agents-agent-name-route-decision-post.ts | 46 +++++ ...agents-agent-name-usage-settlement-post.ts | 49 +++++ .../src/generated/sdk/inference-gateway.ts | 67 +++++++ sdks/typescript/src/generated/sdk/sdk.ts | 6 + 22 files changed, 1221 insertions(+), 1 deletion(-) create mode 100644 sdks/typescript/src/generated/funcs/inference-gateway-preflight.ts create mode 100644 sdks/typescript/src/generated/funcs/inference-gateway-route-decision.ts create mode 100644 sdks/typescript/src/generated/funcs/inference-gateway-usage-settlement.ts create mode 100644 sdks/typescript/src/generated/hooks/registration.ts create mode 100644 sdks/typescript/src/generated/models/gateway-budget-state.ts create mode 100644 sdks/typescript/src/generated/models/gateway-model-candidate.ts create mode 100644 sdks/typescript/src/generated/models/gateway-preflight-request.ts create mode 100644 sdks/typescript/src/generated/models/gateway-preflight-response.ts create mode 100644 sdks/typescript/src/generated/models/gateway-route-decision-response.ts create mode 100644 sdks/typescript/src/generated/models/gateway-usage-settlement-request.ts create mode 100644 sdks/typescript/src/generated/models/gateway-usage-settlement-response.ts create mode 100644 sdks/typescript/src/generated/models/operations/preflight-api-v1-gateway-agents-agent-name-preflight-post.ts create mode 100644 sdks/typescript/src/generated/models/operations/route-decision-api-v1-gateway-agents-agent-name-route-decision-post.ts create mode 100644 sdks/typescript/src/generated/models/operations/usage-settlement-api-v1-gateway-agents-agent-name-usage-settlement-post.ts create mode 100644 sdks/typescript/src/generated/sdk/inference-gateway.ts diff --git a/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/records/factory.py b/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/records/factory.py index 75a3e5b68..486b1a0c3 100644 --- a/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/records/factory.py +++ b/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/records/factory.py @@ -12,6 +12,7 @@ ControlSpan, LlmSpan, RetrieverSpan, + Span, ToolSpan, WorkflowSpan, ) @@ -21,7 +22,7 @@ type GalileoRecord = LlmSpan | ToolSpan | RetrieverSpan | Trace | Session type GalileoSpan = LlmSpan | ToolSpan | RetrieverSpan -type GalileoCoreSpan = AgentSpan | WorkflowSpan | LlmSpan | RetrieverSpan | ToolSpan | ControlSpan +type GalileoCoreSpan = Span _MISSING = object() _SPAN_TYPES = {"llm", "tool", "retriever"} _RECORD_TYPES = _SPAN_TYPES | {"trace", "session"} diff --git a/sdks/typescript/overlays/method-names.overlay.yaml b/sdks/typescript/overlays/method-names.overlay.yaml index 967847c65..516b82bfc 100644 --- a/sdks/typescript/overlays/method-names.overlay.yaml +++ b/sdks/typescript/overlays/method-names.overlay.yaml @@ -220,6 +220,21 @@ actions: x-speakeasy-group: evaluators x-speakeasy-name-override: list + - target: $["paths"]["/api/v1/gateway/agents/{agent_name}/preflight"]["post"] + update: + x-speakeasy-group: inferenceGateway + x-speakeasy-name-override: preflight + + - target: $["paths"]["/api/v1/gateway/agents/{agent_name}/route-decision"]["post"] + update: + x-speakeasy-group: inferenceGateway + x-speakeasy-name-override: routeDecision + + - target: $["paths"]["/api/v1/gateway/agents/{agent_name}/usage-settlement"]["post"] + update: + x-speakeasy-group: inferenceGateway + x-speakeasy-name-override: usageSettlement + - target: $["paths"]["/api/v1/observability/events"]["post"] update: x-speakeasy-group: observability diff --git a/sdks/typescript/src/generated/funcs/inference-gateway-preflight.ts b/sdks/typescript/src/generated/funcs/inference-gateway-preflight.ts new file mode 100644 index 000000000..86d454dab --- /dev/null +++ b/sdks/typescript/src/generated/funcs/inference-gateway-preflight.ts @@ -0,0 +1,185 @@ +/* + * Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT. + */ + +import * as z from "zod/v4-mini"; +import { AgentControlSDKCore } from "../core.js"; +import { encodeJSON, encodeSimple } from "../lib/encodings.js"; +import * as M from "../lib/matchers.js"; +import { compactMap } from "../lib/primitives.js"; +import { safeParse } from "../lib/schemas.js"; +import { RequestOptions } from "../lib/sdks.js"; +import { extractSecurity, resolveGlobalSecurity } from "../lib/security.js"; +import { pathToFunc } from "../lib/url.js"; +import { AgentControlSDKError } from "../models/errors/agent-control-sdk-error.js"; +import { + ConnectionError, + InvalidRequestError, + RequestAbortedError, + RequestTimeoutError, + UnexpectedClientError, +} from "../models/errors/http-client-errors.js"; +import * as errors from "../models/errors/index.js"; +import { ResponseValidationError } from "../models/errors/response-validation-error.js"; +import { SDKValidationError } from "../models/errors/sdk-validation-error.js"; +import * as models from "../models/index.js"; +import * as operations from "../models/operations/index.js"; +import { APICall, APIPromise } from "../types/async.js"; +import { Result } from "../types/fp.js"; + +/** + * Preflight + * + * @remarks + * Evaluate controls and policy eligibility before inference. + */ +export function inferenceGatewayPreflight( + client: AgentControlSDKCore, + request: operations.PreflightApiV1GatewayAgentsAgentNamePreflightPostRequest, + options?: RequestOptions, +): APIPromise< + Result< + models.GatewayPreflightResponse, + | errors.HTTPValidationError + | AgentControlSDKError + | ResponseValidationError + | ConnectionError + | RequestAbortedError + | RequestTimeoutError + | InvalidRequestError + | UnexpectedClientError + | SDKValidationError + > +> { + return new APIPromise($do( + client, + request, + options, + )); +} + +async function $do( + client: AgentControlSDKCore, + request: operations.PreflightApiV1GatewayAgentsAgentNamePreflightPostRequest, + options?: RequestOptions, +): Promise< + [ + Result< + models.GatewayPreflightResponse, + | errors.HTTPValidationError + | AgentControlSDKError + | ResponseValidationError + | ConnectionError + | RequestAbortedError + | RequestTimeoutError + | InvalidRequestError + | UnexpectedClientError + | SDKValidationError + >, + APICall, + ] +> { + const parsed = safeParse( + request, + (value) => + z.parse( + operations + .PreflightApiV1GatewayAgentsAgentNamePreflightPostRequest$outboundSchema, + value, + ), + "Input validation failed", + ); + if (!parsed.ok) { + return [parsed, { status: "invalid" }]; + } + const payload = parsed.value; + const body = encodeJSON("body", payload.body, { explode: true }); + + const pathParams = { + agent_name: encodeSimple("agent_name", payload.agent_name, { + explode: false, + charEncoding: "percent", + }), + }; + + const path = pathToFunc("/api/v1/gateway/agents/{agent_name}/preflight")( + pathParams, + ); + + const headers = new Headers(compactMap({ + "Content-Type": "application/json", + Accept: "application/json", + })); + + const secConfig = await extractSecurity(client._options.apiKeyHeader); + const securityInput = secConfig == null ? {} : { apiKeyHeader: secConfig }; + const requestSecurity = resolveGlobalSecurity(securityInput); + + const context = { + options: client._options, + baseURL: options?.serverURL ?? client._baseURL ?? "", + operationID: "preflight_api_v1_gateway_agents__agent_name__preflight_post", + oAuth2Scopes: null, + + resolvedSecurity: requestSecurity, + + securitySource: client._options.apiKeyHeader, + retryConfig: options?.retries + || client._options.retryConfig + || { strategy: "none" }, + retryCodes: options?.retryCodes || ["429", "500", "502", "503", "504"], + }; + + const requestRes = client._createRequest(context, { + security: requestSecurity, + method: "POST", + baseURL: options?.serverURL, + path: path, + headers: headers, + body: body, + userAgent: client._options.userAgent, + timeoutMs: options?.timeoutMs || client._options.timeoutMs || -1, + }, options); + if (!requestRes.ok) { + return [requestRes, { status: "invalid" }]; + } + const req = requestRes.value; + + const doResult = await client._do(req, { + context, + errorCodes: ["422", "4XX", "5XX"], + retryConfig: context.retryConfig, + retryCodes: context.retryCodes, + }); + if (!doResult.ok) { + return [doResult, { status: "request-error", request: req }]; + } + const response = doResult.value; + + const responseFields = { + HttpMeta: { Response: response, Request: req }, + }; + + const [result] = await M.match< + models.GatewayPreflightResponse, + | errors.HTTPValidationError + | AgentControlSDKError + | ResponseValidationError + | ConnectionError + | RequestAbortedError + | RequestTimeoutError + | InvalidRequestError + | UnexpectedClientError + | SDKValidationError + >( + M.json(200, models.GatewayPreflightResponse$inboundSchema), + M.jsonErr(422, errors.HTTPValidationError$inboundSchema), + M.fail("4XX"), + M.fail("5XX"), + )(response, req, { extraFields: responseFields }); + if (!result.ok) { + return [result, { status: "complete", request: req, response }]; + } + + return [result, { status: "complete", request: req, response }]; +} diff --git a/sdks/typescript/src/generated/funcs/inference-gateway-route-decision.ts b/sdks/typescript/src/generated/funcs/inference-gateway-route-decision.ts new file mode 100644 index 000000000..6060e7e9b --- /dev/null +++ b/sdks/typescript/src/generated/funcs/inference-gateway-route-decision.ts @@ -0,0 +1,188 @@ +/* + * Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT. + */ + +import * as z from "zod/v4-mini"; +import { AgentControlSDKCore } from "../core.js"; +import { encodeJSON, encodeSimple } from "../lib/encodings.js"; +import * as M from "../lib/matchers.js"; +import { compactMap } from "../lib/primitives.js"; +import { safeParse } from "../lib/schemas.js"; +import { RequestOptions } from "../lib/sdks.js"; +import { extractSecurity, resolveGlobalSecurity } from "../lib/security.js"; +import { pathToFunc } from "../lib/url.js"; +import { AgentControlSDKError } from "../models/errors/agent-control-sdk-error.js"; +import { + ConnectionError, + InvalidRequestError, + RequestAbortedError, + RequestTimeoutError, + UnexpectedClientError, +} from "../models/errors/http-client-errors.js"; +import * as errors from "../models/errors/index.js"; +import { ResponseValidationError } from "../models/errors/response-validation-error.js"; +import { SDKValidationError } from "../models/errors/sdk-validation-error.js"; +import * as models from "../models/index.js"; +import * as operations from "../models/operations/index.js"; +import { APICall, APIPromise } from "../types/async.js"; +import { Result } from "../types/fp.js"; + +/** + * Route Decision + * + * @remarks + * Select an allowed model for the authenticated agent and classified task. + */ +export function inferenceGatewayRouteDecision( + client: AgentControlSDKCore, + request: + operations.RouteDecisionApiV1GatewayAgentsAgentNameRouteDecisionPostRequest, + options?: RequestOptions, +): APIPromise< + Result< + models.GatewayRouteDecisionResponse, + | errors.HTTPValidationError + | AgentControlSDKError + | ResponseValidationError + | ConnectionError + | RequestAbortedError + | RequestTimeoutError + | InvalidRequestError + | UnexpectedClientError + | SDKValidationError + > +> { + return new APIPromise($do( + client, + request, + options, + )); +} + +async function $do( + client: AgentControlSDKCore, + request: + operations.RouteDecisionApiV1GatewayAgentsAgentNameRouteDecisionPostRequest, + options?: RequestOptions, +): Promise< + [ + Result< + models.GatewayRouteDecisionResponse, + | errors.HTTPValidationError + | AgentControlSDKError + | ResponseValidationError + | ConnectionError + | RequestAbortedError + | RequestTimeoutError + | InvalidRequestError + | UnexpectedClientError + | SDKValidationError + >, + APICall, + ] +> { + const parsed = safeParse( + request, + (value) => + z.parse( + operations + .RouteDecisionApiV1GatewayAgentsAgentNameRouteDecisionPostRequest$outboundSchema, + value, + ), + "Input validation failed", + ); + if (!parsed.ok) { + return [parsed, { status: "invalid" }]; + } + const payload = parsed.value; + const body = encodeJSON("body", payload.body, { explode: true }); + + const pathParams = { + agent_name: encodeSimple("agent_name", payload.agent_name, { + explode: false, + charEncoding: "percent", + }), + }; + + const path = pathToFunc("/api/v1/gateway/agents/{agent_name}/route-decision")( + pathParams, + ); + + const headers = new Headers(compactMap({ + "Content-Type": "application/json", + Accept: "application/json", + })); + + const secConfig = await extractSecurity(client._options.apiKeyHeader); + const securityInput = secConfig == null ? {} : { apiKeyHeader: secConfig }; + const requestSecurity = resolveGlobalSecurity(securityInput); + + const context = { + options: client._options, + baseURL: options?.serverURL ?? client._baseURL ?? "", + operationID: + "route_decision_api_v1_gateway_agents__agent_name__route_decision_post", + oAuth2Scopes: null, + + resolvedSecurity: requestSecurity, + + securitySource: client._options.apiKeyHeader, + retryConfig: options?.retries + || client._options.retryConfig + || { strategy: "none" }, + retryCodes: options?.retryCodes || ["429", "500", "502", "503", "504"], + }; + + const requestRes = client._createRequest(context, { + security: requestSecurity, + method: "POST", + baseURL: options?.serverURL, + path: path, + headers: headers, + body: body, + userAgent: client._options.userAgent, + timeoutMs: options?.timeoutMs || client._options.timeoutMs || -1, + }, options); + if (!requestRes.ok) { + return [requestRes, { status: "invalid" }]; + } + const req = requestRes.value; + + const doResult = await client._do(req, { + context, + errorCodes: ["422", "4XX", "5XX"], + retryConfig: context.retryConfig, + retryCodes: context.retryCodes, + }); + if (!doResult.ok) { + return [doResult, { status: "request-error", request: req }]; + } + const response = doResult.value; + + const responseFields = { + HttpMeta: { Response: response, Request: req }, + }; + + const [result] = await M.match< + models.GatewayRouteDecisionResponse, + | errors.HTTPValidationError + | AgentControlSDKError + | ResponseValidationError + | ConnectionError + | RequestAbortedError + | RequestTimeoutError + | InvalidRequestError + | UnexpectedClientError + | SDKValidationError + >( + M.json(200, models.GatewayRouteDecisionResponse$inboundSchema), + M.jsonErr(422, errors.HTTPValidationError$inboundSchema), + M.fail("4XX"), + M.fail("5XX"), + )(response, req, { extraFields: responseFields }); + if (!result.ok) { + return [result, { status: "complete", request: req, response }]; + } + + return [result, { status: "complete", request: req, response }]; +} diff --git a/sdks/typescript/src/generated/funcs/inference-gateway-usage-settlement.ts b/sdks/typescript/src/generated/funcs/inference-gateway-usage-settlement.ts new file mode 100644 index 000000000..d9b653c89 --- /dev/null +++ b/sdks/typescript/src/generated/funcs/inference-gateway-usage-settlement.ts @@ -0,0 +1,188 @@ +/* + * Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT. + */ + +import * as z from "zod/v4-mini"; +import { AgentControlSDKCore } from "../core.js"; +import { encodeJSON, encodeSimple } from "../lib/encodings.js"; +import * as M from "../lib/matchers.js"; +import { compactMap } from "../lib/primitives.js"; +import { safeParse } from "../lib/schemas.js"; +import { RequestOptions } from "../lib/sdks.js"; +import { extractSecurity, resolveGlobalSecurity } from "../lib/security.js"; +import { pathToFunc } from "../lib/url.js"; +import { AgentControlSDKError } from "../models/errors/agent-control-sdk-error.js"; +import { + ConnectionError, + InvalidRequestError, + RequestAbortedError, + RequestTimeoutError, + UnexpectedClientError, +} from "../models/errors/http-client-errors.js"; +import * as errors from "../models/errors/index.js"; +import { ResponseValidationError } from "../models/errors/response-validation-error.js"; +import { SDKValidationError } from "../models/errors/sdk-validation-error.js"; +import * as models from "../models/index.js"; +import * as operations from "../models/operations/index.js"; +import { APICall, APIPromise } from "../types/async.js"; +import { Result } from "../types/fp.js"; + +/** + * Usage Settlement + * + * @remarks + * Accumulate completed inference usage against the authenticated agent budget. + */ +export function inferenceGatewayUsageSettlement( + client: AgentControlSDKCore, + request: + operations.UsageSettlementApiV1GatewayAgentsAgentNameUsageSettlementPostRequest, + options?: RequestOptions, +): APIPromise< + Result< + models.GatewayUsageSettlementResponse, + | errors.HTTPValidationError + | AgentControlSDKError + | ResponseValidationError + | ConnectionError + | RequestAbortedError + | RequestTimeoutError + | InvalidRequestError + | UnexpectedClientError + | SDKValidationError + > +> { + return new APIPromise($do( + client, + request, + options, + )); +} + +async function $do( + client: AgentControlSDKCore, + request: + operations.UsageSettlementApiV1GatewayAgentsAgentNameUsageSettlementPostRequest, + options?: RequestOptions, +): Promise< + [ + Result< + models.GatewayUsageSettlementResponse, + | errors.HTTPValidationError + | AgentControlSDKError + | ResponseValidationError + | ConnectionError + | RequestAbortedError + | RequestTimeoutError + | InvalidRequestError + | UnexpectedClientError + | SDKValidationError + >, + APICall, + ] +> { + const parsed = safeParse( + request, + (value) => + z.parse( + operations + .UsageSettlementApiV1GatewayAgentsAgentNameUsageSettlementPostRequest$outboundSchema, + value, + ), + "Input validation failed", + ); + if (!parsed.ok) { + return [parsed, { status: "invalid" }]; + } + const payload = parsed.value; + const body = encodeJSON("body", payload.body, { explode: true }); + + const pathParams = { + agent_name: encodeSimple("agent_name", payload.agent_name, { + explode: false, + charEncoding: "percent", + }), + }; + + const path = pathToFunc( + "/api/v1/gateway/agents/{agent_name}/usage-settlement", + )(pathParams); + + const headers = new Headers(compactMap({ + "Content-Type": "application/json", + Accept: "application/json", + })); + + const secConfig = await extractSecurity(client._options.apiKeyHeader); + const securityInput = secConfig == null ? {} : { apiKeyHeader: secConfig }; + const requestSecurity = resolveGlobalSecurity(securityInput); + + const context = { + options: client._options, + baseURL: options?.serverURL ?? client._baseURL ?? "", + operationID: + "usage_settlement_api_v1_gateway_agents__agent_name__usage_settlement_post", + oAuth2Scopes: null, + + resolvedSecurity: requestSecurity, + + securitySource: client._options.apiKeyHeader, + retryConfig: options?.retries + || client._options.retryConfig + || { strategy: "none" }, + retryCodes: options?.retryCodes || ["429", "500", "502", "503", "504"], + }; + + const requestRes = client._createRequest(context, { + security: requestSecurity, + method: "POST", + baseURL: options?.serverURL, + path: path, + headers: headers, + body: body, + userAgent: client._options.userAgent, + timeoutMs: options?.timeoutMs || client._options.timeoutMs || -1, + }, options); + if (!requestRes.ok) { + return [requestRes, { status: "invalid" }]; + } + const req = requestRes.value; + + const doResult = await client._do(req, { + context, + errorCodes: ["422", "4XX", "5XX"], + retryConfig: context.retryConfig, + retryCodes: context.retryCodes, + }); + if (!doResult.ok) { + return [doResult, { status: "request-error", request: req }]; + } + const response = doResult.value; + + const responseFields = { + HttpMeta: { Response: response, Request: req }, + }; + + const [result] = await M.match< + models.GatewayUsageSettlementResponse, + | errors.HTTPValidationError + | AgentControlSDKError + | ResponseValidationError + | ConnectionError + | RequestAbortedError + | RequestTimeoutError + | InvalidRequestError + | UnexpectedClientError + | SDKValidationError + >( + M.json(200, models.GatewayUsageSettlementResponse$inboundSchema), + M.jsonErr(422, errors.HTTPValidationError$inboundSchema), + M.fail("4XX"), + M.fail("5XX"), + )(response, req, { extraFields: responseFields }); + if (!result.ok) { + return [result, { status: "complete", request: req, response }]; + } + + return [result, { status: "complete", request: req, response }]; +} diff --git a/sdks/typescript/src/generated/hooks/hooks.ts b/sdks/typescript/src/generated/hooks/hooks.ts index 94ad0a323..7ed9e0445 100644 --- a/sdks/typescript/src/generated/hooks/hooks.ts +++ b/sdks/typescript/src/generated/hooks/hooks.ts @@ -18,6 +18,8 @@ import { SDKInitHook, } from "./types.js"; +import { initHooks } from "./registration.js"; + export class SDKHooks implements Hooks { sdkInitHooks: SDKInitHook[] = []; beforeCreateRequestHooks: BeforeCreateRequestHook[] = []; @@ -45,6 +47,7 @@ export class SDKHooks implements Hooks { this.registerAfterErrorHook(hook); } } + initHooks(this); } registerSDKInitHook(hook: SDKInitHook) { diff --git a/sdks/typescript/src/generated/hooks/registration.ts b/sdks/typescript/src/generated/hooks/registration.ts new file mode 100644 index 000000000..3a8fdeec1 --- /dev/null +++ b/sdks/typescript/src/generated/hooks/registration.ts @@ -0,0 +1,14 @@ +import { Hooks } from "./types.js"; + +/* + * This file is only ever generated once on the first generation and then is free to be modified. + * Any hooks you wish to add should be registered in the initHooks function. Feel free to define them + * in this file or in separate files in the hooks folder. + */ + +// @ts-ignore remove this line when you add your first hook and hooks is used +export function initHooks(hooks: Hooks) { + // Add hooks by calling hooks.register{ClientInit/BeforeCreateRequest/BeforeRequest/AfterSuccess/AfterError}Hook + // with an instance of a hook that implements that specific Hook interface + // Hooks are registered per SDK instance, and are valid for the lifetime of the SDK instance +} diff --git a/sdks/typescript/src/generated/models/action-decision.ts b/sdks/typescript/src/generated/models/action-decision.ts index 6be26dba1..15e9fc3e1 100644 --- a/sdks/typescript/src/generated/models/action-decision.ts +++ b/sdks/typescript/src/generated/models/action-decision.ts @@ -10,6 +10,7 @@ export const ActionDecision = { Deny: "deny", Steer: "steer", Observe: "observe", + Route: "route", } as const; export type ActionDecision = OpenEnum; diff --git a/sdks/typescript/src/generated/models/gateway-budget-state.ts b/sdks/typescript/src/generated/models/gateway-budget-state.ts new file mode 100644 index 000000000..6466c2cae --- /dev/null +++ b/sdks/typescript/src/generated/models/gateway-budget-state.ts @@ -0,0 +1,55 @@ +/* + * Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT. + */ + +import * as z from "zod/v4-mini"; +import { remap as remap$ } from "../lib/primitives.js"; +import { safeParse } from "../lib/schemas.js"; +import { Result as SafeParseResult } from "../types/fp.js"; +import * as types from "../types/primitives.js"; +import { SDKValidationError } from "./errors/sdk-validation-error.js"; + +/** + * Accumulated token use and configured budget threshold state. + */ +export type GatewayBudgetState = { + limitTokens: number; + overLimit: boolean; + threshold: number; + thresholdReached: boolean; + usedTokens: number; + utilization: number; +}; + +/** @internal */ +export const GatewayBudgetState$inboundSchema: z.ZodMiniType< + GatewayBudgetState, + unknown +> = z.pipe( + z.object({ + limit_tokens: types.number(), + over_limit: types.boolean(), + threshold: types.number(), + threshold_reached: types.boolean(), + used_tokens: types.number(), + utilization: types.number(), + }), + z.transform((v) => { + return remap$(v, { + "limit_tokens": "limitTokens", + "over_limit": "overLimit", + "threshold_reached": "thresholdReached", + "used_tokens": "usedTokens", + }); + }), +); + +export function gatewayBudgetStateFromJSON( + jsonString: string, +): SafeParseResult { + return safeParse( + jsonString, + (x) => GatewayBudgetState$inboundSchema.parse(JSON.parse(x)), + `Failed to parse 'GatewayBudgetState' from JSON`, + ); +} diff --git a/sdks/typescript/src/generated/models/gateway-model-candidate.ts b/sdks/typescript/src/generated/models/gateway-model-candidate.ts new file mode 100644 index 000000000..0c04625a5 --- /dev/null +++ b/sdks/typescript/src/generated/models/gateway-model-candidate.ts @@ -0,0 +1,53 @@ +/* + * Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT. + */ + +import * as z from "zod/v4-mini"; +import { ClosedEnum } from "../types/enums.js"; + +export const GatewayModelCandidateProtocol = { + Openai: "openai", + Anthropic: "anthropic", +} as const; +export type GatewayModelCandidateProtocol = ClosedEnum< + typeof GatewayModelCandidateProtocol +>; + +/** + * A model and wire protocol the gateway can route to. + */ +export type GatewayModelCandidate = { + model: string; + protocol: GatewayModelCandidateProtocol; + provider: string; +}; + +/** @internal */ +export const GatewayModelCandidateProtocol$outboundSchema: z.ZodMiniEnum< + typeof GatewayModelCandidateProtocol +> = z.enum(GatewayModelCandidateProtocol); + +/** @internal */ +export type GatewayModelCandidate$Outbound = { + model: string; + protocol: string; + provider: string; +}; + +/** @internal */ +export const GatewayModelCandidate$outboundSchema: z.ZodMiniType< + GatewayModelCandidate$Outbound, + GatewayModelCandidate +> = z.object({ + model: z.string(), + protocol: GatewayModelCandidateProtocol$outboundSchema, + provider: z.string(), +}); + +export function gatewayModelCandidateToJSON( + gatewayModelCandidate: GatewayModelCandidate, +): string { + return JSON.stringify( + GatewayModelCandidate$outboundSchema.parse(gatewayModelCandidate), + ); +} diff --git a/sdks/typescript/src/generated/models/gateway-preflight-request.ts b/sdks/typescript/src/generated/models/gateway-preflight-request.ts new file mode 100644 index 000000000..c8b21ea53 --- /dev/null +++ b/sdks/typescript/src/generated/models/gateway-preflight-request.ts @@ -0,0 +1,64 @@ +/* + * Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT. + */ + +import * as z from "zod/v4-mini"; +import { remap as remap$ } from "../lib/primitives.js"; +import { + GatewayModelCandidate, + GatewayModelCandidate$Outbound, + GatewayModelCandidate$outboundSchema, +} from "./gateway-model-candidate.js"; + +/** + * Candidate routes and task information supplied before inference. + */ +export type GatewayPreflightRequest = { + agentName?: string | null | undefined; + candidateModels?: Array | undefined; + model?: string | null | undefined; + prompt?: string | null | undefined; + taskContext?: { [k: string]: any } | null | undefined; + taskType?: string | undefined; +}; + +/** @internal */ +export type GatewayPreflightRequest$Outbound = { + agent_name?: string | null | undefined; + candidate_models?: Array | undefined; + model?: string | null | undefined; + prompt?: string | null | undefined; + task_context?: { [k: string]: any } | null | undefined; + task_type: string; +}; + +/** @internal */ +export const GatewayPreflightRequest$outboundSchema: z.ZodMiniType< + GatewayPreflightRequest$Outbound, + GatewayPreflightRequest +> = z.pipe( + z.object({ + agentName: z.optional(z.nullable(z.string())), + candidateModels: z.optional(z.array(GatewayModelCandidate$outboundSchema)), + model: z.optional(z.nullable(z.string())), + prompt: z.optional(z.nullable(z.string())), + taskContext: z.optional(z.nullable(z.record(z.string(), z.any()))), + taskType: z._default(z.string(), "default"), + }), + z.transform((v) => { + return remap$(v, { + agentName: "agent_name", + candidateModels: "candidate_models", + taskContext: "task_context", + taskType: "task_type", + }); + }), +); + +export function gatewayPreflightRequestToJSON( + gatewayPreflightRequest: GatewayPreflightRequest, +): string { + return JSON.stringify( + GatewayPreflightRequest$outboundSchema.parse(gatewayPreflightRequest), + ); +} diff --git a/sdks/typescript/src/generated/models/gateway-preflight-response.ts b/sdks/typescript/src/generated/models/gateway-preflight-response.ts new file mode 100644 index 000000000..1ad1d77f1 --- /dev/null +++ b/sdks/typescript/src/generated/models/gateway-preflight-response.ts @@ -0,0 +1,44 @@ +/* + * Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT. + */ + +import * as z from "zod/v4-mini"; +import { safeParse } from "../lib/schemas.js"; +import { Result as SafeParseResult } from "../types/fp.js"; +import * as types from "../types/primitives.js"; +import { SDKValidationError } from "./errors/sdk-validation-error.js"; +import { + GatewayBudgetState, + GatewayBudgetState$inboundSchema, +} from "./gateway-budget-state.js"; + +/** + * Control and budget decision returned before inference. + */ +export type GatewayPreflightResponse = { + allowed: boolean; + budget: GatewayBudgetState | null; + policy: string; + reason: string; +}; + +/** @internal */ +export const GatewayPreflightResponse$inboundSchema: z.ZodMiniType< + GatewayPreflightResponse, + unknown +> = z.object({ + allowed: types.boolean(), + budget: types.nullable(GatewayBudgetState$inboundSchema), + policy: types.string(), + reason: types.string(), +}); + +export function gatewayPreflightResponseFromJSON( + jsonString: string, +): SafeParseResult { + return safeParse( + jsonString, + (x) => GatewayPreflightResponse$inboundSchema.parse(JSON.parse(x)), + `Failed to parse 'GatewayPreflightResponse' from JSON`, + ); +} diff --git a/sdks/typescript/src/generated/models/gateway-route-decision-response.ts b/sdks/typescript/src/generated/models/gateway-route-decision-response.ts new file mode 100644 index 000000000..481f6836b --- /dev/null +++ b/sdks/typescript/src/generated/models/gateway-route-decision-response.ts @@ -0,0 +1,90 @@ +/* + * Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT. + */ + +import * as z from "zod/v4-mini"; +import { remap as remap$ } from "../lib/primitives.js"; +import { safeParse } from "../lib/schemas.js"; +import * as openEnums from "../types/enums.js"; +import { OpenEnum } from "../types/enums.js"; +import { Result as SafeParseResult } from "../types/fp.js"; +import * as types from "../types/primitives.js"; +import { + ActionDecision, + ActionDecision$inboundSchema, +} from "./action-decision.js"; +import { SDKValidationError } from "./errors/sdk-validation-error.js"; +import { + GatewayBudgetState, + GatewayBudgetState$inboundSchema, +} from "./gateway-budget-state.js"; + +export const Decision = { + Route: "route", + Deny: "deny", +} as const; +export type Decision = OpenEnum; + +export const Protocol = { + Openai: "openai", + Anthropic: "anthropic", +} as const; +export type Protocol = OpenEnum; + +/** + * Route result directly consumable by an inference gateway adapter. + */ +export type GatewayRouteDecisionResponse = { + action?: ActionDecision | null | undefined; + allowed: boolean; + budget: GatewayBudgetState | null; + decision: Decision; + model: string | null; + policy: string; + protocol: Protocol | null; + provider: string | null; + reason: string; + taskType: string; +}; + +/** @internal */ +export const Decision$inboundSchema: z.ZodMiniType = + openEnums.inboundSchema(Decision); + +/** @internal */ +export const Protocol$inboundSchema: z.ZodMiniType = + openEnums.inboundSchema(Protocol); + +/** @internal */ +export const GatewayRouteDecisionResponse$inboundSchema: z.ZodMiniType< + GatewayRouteDecisionResponse, + unknown +> = z.pipe( + z.object({ + action: z.optional(z.nullable(ActionDecision$inboundSchema)), + allowed: types.boolean(), + budget: types.nullable(GatewayBudgetState$inboundSchema), + decision: Decision$inboundSchema, + model: types.nullable(types.string()), + policy: types.string(), + protocol: types.nullable(Protocol$inboundSchema), + provider: types.nullable(types.string()), + reason: types.string(), + task_type: types.string(), + }), + z.transform((v) => { + return remap$(v, { + "task_type": "taskType", + }); + }), +); + +export function gatewayRouteDecisionResponseFromJSON( + jsonString: string, +): SafeParseResult { + return safeParse( + jsonString, + (x) => GatewayRouteDecisionResponse$inboundSchema.parse(JSON.parse(x)), + `Failed to parse 'GatewayRouteDecisionResponse' from JSON`, + ); +} diff --git a/sdks/typescript/src/generated/models/gateway-usage-settlement-request.ts b/sdks/typescript/src/generated/models/gateway-usage-settlement-request.ts new file mode 100644 index 000000000..958f3dc24 --- /dev/null +++ b/sdks/typescript/src/generated/models/gateway-usage-settlement-request.ts @@ -0,0 +1,50 @@ +/* + * Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT. + */ + +import * as z from "zod/v4-mini"; +import { remap as remap$ } from "../lib/primitives.js"; + +/** + * Actual token usage reported by the gateway after inference. + */ +export type GatewayUsageSettlementRequest = { + inputTokens: number; + model: string; + outputTokens: number; +}; + +/** @internal */ +export type GatewayUsageSettlementRequest$Outbound = { + input_tokens: number; + model: string; + output_tokens: number; +}; + +/** @internal */ +export const GatewayUsageSettlementRequest$outboundSchema: z.ZodMiniType< + GatewayUsageSettlementRequest$Outbound, + GatewayUsageSettlementRequest +> = z.pipe( + z.object({ + inputTokens: z.int(), + model: z.string(), + outputTokens: z.int(), + }), + z.transform((v) => { + return remap$(v, { + inputTokens: "input_tokens", + outputTokens: "output_tokens", + }); + }), +); + +export function gatewayUsageSettlementRequestToJSON( + gatewayUsageSettlementRequest: GatewayUsageSettlementRequest, +): string { + return JSON.stringify( + GatewayUsageSettlementRequest$outboundSchema.parse( + gatewayUsageSettlementRequest, + ), + ); +} diff --git a/sdks/typescript/src/generated/models/gateway-usage-settlement-response.ts b/sdks/typescript/src/generated/models/gateway-usage-settlement-response.ts new file mode 100644 index 000000000..7a27fb587 --- /dev/null +++ b/sdks/typescript/src/generated/models/gateway-usage-settlement-response.ts @@ -0,0 +1,45 @@ +/* + * Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT. + */ + +import * as z from "zod/v4-mini"; +import { safeParse } from "../lib/schemas.js"; +import { Result as SafeParseResult } from "../types/fp.js"; +import * as types from "../types/primitives.js"; +import { SDKValidationError } from "./errors/sdk-validation-error.js"; +import { + GatewayBudgetState, + GatewayBudgetState$inboundSchema, +} from "./gateway-budget-state.js"; + +/** + * Budget state after actual usage has been accumulated. + */ +export type GatewayUsageSettlementResponse = { + /** + * Accumulated token use and configured budget threshold state. + */ + budget: GatewayBudgetState; + policy: string; + reason: string; +}; + +/** @internal */ +export const GatewayUsageSettlementResponse$inboundSchema: z.ZodMiniType< + GatewayUsageSettlementResponse, + unknown +> = z.object({ + budget: GatewayBudgetState$inboundSchema, + policy: types.string(), + reason: types.string(), +}); + +export function gatewayUsageSettlementResponseFromJSON( + jsonString: string, +): SafeParseResult { + return safeParse( + jsonString, + (x) => GatewayUsageSettlementResponse$inboundSchema.parse(JSON.parse(x)), + `Failed to parse 'GatewayUsageSettlementResponse' from JSON`, + ); +} diff --git a/sdks/typescript/src/generated/models/index.ts b/sdks/typescript/src/generated/models/index.ts index 8043233b0..f6b9521fa 100644 --- a/sdks/typescript/src/generated/models/index.ts +++ b/sdks/typescript/src/generated/models/index.ts @@ -53,6 +53,13 @@ export * from "./evaluator-schema.js"; export * from "./evaluator-spec.js"; export * from "./event-query-request.js"; export * from "./event-query-response.js"; +export * from "./gateway-budget-state.js"; +export * from "./gateway-model-candidate.js"; +export * from "./gateway-preflight-request.js"; +export * from "./gateway-preflight-response.js"; +export * from "./gateway-route-decision-response.js"; +export * from "./gateway-usage-settlement-request.js"; +export * from "./gateway-usage-settlement-response.js"; export * from "./get-agent-policies-response.js"; export * from "./get-agent-response.js"; export * from "./get-control-binding-response.js"; diff --git a/sdks/typescript/src/generated/models/operations/index.ts b/sdks/typescript/src/generated/models/operations/index.ts index b031df327..bc0145348 100644 --- a/sdks/typescript/src/generated/models/operations/index.ts +++ b/sdks/typescript/src/generated/models/operations/index.ts @@ -29,9 +29,12 @@ export * from "./list-policy-controls-api-v1-policies-policy-id-controls-get.js" export * from "./patch-agent-api-v1-agents-agent-name-patch.js"; export * from "./patch-control-api-v1-controls-control-id-patch.js"; export * from "./patch-control-binding-api-v1-control-bindings-binding-id-patch.js"; +export * from "./preflight-api-v1-gateway-agents-agent-name-preflight-post.js"; export * from "./remove-agent-control-api-v1-agents-agent-name-controls-control-id-delete.js"; export * from "./remove-agent-policy-api-v1-agents-agent-name-policies-policy-id-delete.js"; export * from "./remove-all-agent-policies-api-v1-agents-agent-name-policies-delete.js"; export * from "./remove-control-from-policy-api-v1-policies-policy-id-controls-control-id-delete.js"; +export * from "./route-decision-api-v1-gateway-agents-agent-name-route-decision-post.js"; export * from "./set-agent-policy-api-v1-agents-agent-name-policy-policy-id-post.js"; export * from "./set-control-data-api-v1-controls-control-id-data-put.js"; +export * from "./usage-settlement-api-v1-gateway-agents-agent-name-usage-settlement-post.js"; diff --git a/sdks/typescript/src/generated/models/operations/preflight-api-v1-gateway-agents-agent-name-preflight-post.ts b/sdks/typescript/src/generated/models/operations/preflight-api-v1-gateway-agents-agent-name-preflight-post.ts new file mode 100644 index 000000000..07d6924d7 --- /dev/null +++ b/sdks/typescript/src/generated/models/operations/preflight-api-v1-gateway-agents-agent-name-preflight-post.ts @@ -0,0 +1,46 @@ +/* + * Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT. + */ + +import * as z from "zod/v4-mini"; +import { remap as remap$ } from "../../lib/primitives.js"; +import * as models from "../index.js"; + +export type PreflightApiV1GatewayAgentsAgentNamePreflightPostRequest = { + agentName: string; + body: models.GatewayPreflightRequest; +}; + +/** @internal */ +export type PreflightApiV1GatewayAgentsAgentNamePreflightPostRequest$Outbound = + { + agent_name: string; + body: models.GatewayPreflightRequest$Outbound; + }; + +/** @internal */ +export const PreflightApiV1GatewayAgentsAgentNamePreflightPostRequest$outboundSchema: + z.ZodMiniType< + PreflightApiV1GatewayAgentsAgentNamePreflightPostRequest$Outbound, + PreflightApiV1GatewayAgentsAgentNamePreflightPostRequest + > = z.pipe( + z.object({ + agentName: z.string(), + body: models.GatewayPreflightRequest$outboundSchema, + }), + z.transform((v) => { + return remap$(v, { + agentName: "agent_name", + }); + }), + ); + +export function preflightApiV1GatewayAgentsAgentNamePreflightPostRequestToJSON( + preflightApiV1GatewayAgentsAgentNamePreflightPostRequest: + PreflightApiV1GatewayAgentsAgentNamePreflightPostRequest, +): string { + return JSON.stringify( + PreflightApiV1GatewayAgentsAgentNamePreflightPostRequest$outboundSchema + .parse(preflightApiV1GatewayAgentsAgentNamePreflightPostRequest), + ); +} diff --git a/sdks/typescript/src/generated/models/operations/route-decision-api-v1-gateway-agents-agent-name-route-decision-post.ts b/sdks/typescript/src/generated/models/operations/route-decision-api-v1-gateway-agents-agent-name-route-decision-post.ts new file mode 100644 index 000000000..26a5a4d7a --- /dev/null +++ b/sdks/typescript/src/generated/models/operations/route-decision-api-v1-gateway-agents-agent-name-route-decision-post.ts @@ -0,0 +1,46 @@ +/* + * Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT. + */ + +import * as z from "zod/v4-mini"; +import { remap as remap$ } from "../../lib/primitives.js"; +import * as models from "../index.js"; + +export type RouteDecisionApiV1GatewayAgentsAgentNameRouteDecisionPostRequest = { + agentName: string; + body: models.GatewayPreflightRequest; +}; + +/** @internal */ +export type RouteDecisionApiV1GatewayAgentsAgentNameRouteDecisionPostRequest$Outbound = + { + agent_name: string; + body: models.GatewayPreflightRequest$Outbound; + }; + +/** @internal */ +export const RouteDecisionApiV1GatewayAgentsAgentNameRouteDecisionPostRequest$outboundSchema: + z.ZodMiniType< + RouteDecisionApiV1GatewayAgentsAgentNameRouteDecisionPostRequest$Outbound, + RouteDecisionApiV1GatewayAgentsAgentNameRouteDecisionPostRequest + > = z.pipe( + z.object({ + agentName: z.string(), + body: models.GatewayPreflightRequest$outboundSchema, + }), + z.transform((v) => { + return remap$(v, { + agentName: "agent_name", + }); + }), + ); + +export function routeDecisionApiV1GatewayAgentsAgentNameRouteDecisionPostRequestToJSON( + routeDecisionApiV1GatewayAgentsAgentNameRouteDecisionPostRequest: + RouteDecisionApiV1GatewayAgentsAgentNameRouteDecisionPostRequest, +): string { + return JSON.stringify( + RouteDecisionApiV1GatewayAgentsAgentNameRouteDecisionPostRequest$outboundSchema + .parse(routeDecisionApiV1GatewayAgentsAgentNameRouteDecisionPostRequest), + ); +} diff --git a/sdks/typescript/src/generated/models/operations/usage-settlement-api-v1-gateway-agents-agent-name-usage-settlement-post.ts b/sdks/typescript/src/generated/models/operations/usage-settlement-api-v1-gateway-agents-agent-name-usage-settlement-post.ts new file mode 100644 index 000000000..58a0fbb1b --- /dev/null +++ b/sdks/typescript/src/generated/models/operations/usage-settlement-api-v1-gateway-agents-agent-name-usage-settlement-post.ts @@ -0,0 +1,49 @@ +/* + * Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT. + */ + +import * as z from "zod/v4-mini"; +import { remap as remap$ } from "../../lib/primitives.js"; +import * as models from "../index.js"; + +export type UsageSettlementApiV1GatewayAgentsAgentNameUsageSettlementPostRequest = + { + agentName: string; + body: models.GatewayUsageSettlementRequest; + }; + +/** @internal */ +export type UsageSettlementApiV1GatewayAgentsAgentNameUsageSettlementPostRequest$Outbound = + { + agent_name: string; + body: models.GatewayUsageSettlementRequest$Outbound; + }; + +/** @internal */ +export const UsageSettlementApiV1GatewayAgentsAgentNameUsageSettlementPostRequest$outboundSchema: + z.ZodMiniType< + UsageSettlementApiV1GatewayAgentsAgentNameUsageSettlementPostRequest$Outbound, + UsageSettlementApiV1GatewayAgentsAgentNameUsageSettlementPostRequest + > = z.pipe( + z.object({ + agentName: z.string(), + body: models.GatewayUsageSettlementRequest$outboundSchema, + }), + z.transform((v) => { + return remap$(v, { + agentName: "agent_name", + }); + }), + ); + +export function usageSettlementApiV1GatewayAgentsAgentNameUsageSettlementPostRequestToJSON( + usageSettlementApiV1GatewayAgentsAgentNameUsageSettlementPostRequest: + UsageSettlementApiV1GatewayAgentsAgentNameUsageSettlementPostRequest, +): string { + return JSON.stringify( + UsageSettlementApiV1GatewayAgentsAgentNameUsageSettlementPostRequest$outboundSchema + .parse( + usageSettlementApiV1GatewayAgentsAgentNameUsageSettlementPostRequest, + ), + ); +} diff --git a/sdks/typescript/src/generated/sdk/inference-gateway.ts b/sdks/typescript/src/generated/sdk/inference-gateway.ts new file mode 100644 index 000000000..c31b46aef --- /dev/null +++ b/sdks/typescript/src/generated/sdk/inference-gateway.ts @@ -0,0 +1,67 @@ +/* + * Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT. + */ + +import { inferenceGatewayPreflight } from "../funcs/inference-gateway-preflight.js"; +import { inferenceGatewayRouteDecision } from "../funcs/inference-gateway-route-decision.js"; +import { inferenceGatewayUsageSettlement } from "../funcs/inference-gateway-usage-settlement.js"; +import { ClientSDK, RequestOptions } from "../lib/sdks.js"; +import * as models from "../models/index.js"; +import * as operations from "../models/operations/index.js"; +import { unwrapAsync } from "../types/fp.js"; + +export class InferenceGateway extends ClientSDK { + /** + * Preflight + * + * @remarks + * Evaluate controls and policy eligibility before inference. + */ + async preflight( + request: + operations.PreflightApiV1GatewayAgentsAgentNamePreflightPostRequest, + options?: RequestOptions, + ): Promise { + return unwrapAsync(inferenceGatewayPreflight( + this, + request, + options, + )); + } + + /** + * Route Decision + * + * @remarks + * Select an allowed model for the authenticated agent and classified task. + */ + async routeDecision( + request: + operations.RouteDecisionApiV1GatewayAgentsAgentNameRouteDecisionPostRequest, + options?: RequestOptions, + ): Promise { + return unwrapAsync(inferenceGatewayRouteDecision( + this, + request, + options, + )); + } + + /** + * Usage Settlement + * + * @remarks + * Accumulate completed inference usage against the authenticated agent budget. + */ + async usageSettlement( + request: + operations.UsageSettlementApiV1GatewayAgentsAgentNameUsageSettlementPostRequest, + options?: RequestOptions, + ): Promise { + return unwrapAsync(inferenceGatewayUsageSettlement( + this, + request, + options, + )); + } +} diff --git a/sdks/typescript/src/generated/sdk/sdk.ts b/sdks/typescript/src/generated/sdk/sdk.ts index 7a66b2af1..1a2d64271 100644 --- a/sdks/typescript/src/generated/sdk/sdk.ts +++ b/sdks/typescript/src/generated/sdk/sdk.ts @@ -9,6 +9,7 @@ import { ControlBindings } from "./control-bindings.js"; import { Controls } from "./controls.js"; import { Evaluation } from "./evaluation.js"; import { Evaluators } from "./evaluators.js"; +import { InferenceGateway } from "./inference-gateway.js"; import { Observability } from "./observability.js"; import { Policies } from "./policies.js"; import { System } from "./system.js"; @@ -49,6 +50,11 @@ export class AgentControlSDK extends ClientSDK { return (this._evaluators ??= new Evaluators(this._options)); } + private _inferenceGateway?: InferenceGateway; + get inferenceGateway(): InferenceGateway { + return (this._inferenceGateway ??= new InferenceGateway(this._options)); + } + private _observability?: Observability; get observability(): Observability { return (this._observability ??= new Observability(this._options)); From 0abd6fc419300f0d3d1370ef9cb1d1a8f8aba639 Mon Sep 17 00:00:00 2001 From: Namrata Ghadi Date: Sat, 31 Oct 2026 15:11:18 -0700 Subject: [PATCH 3/5] fix: stabilize generated gateway SDK output --- .../records/factory.py | 3 --- sdks/typescript/scripts/generate-sdk.sh | 14 +++++++++----- sdks/typescript/src/generated/hooks/hooks.ts | 3 --- .../typescript/src/generated/hooks/registration.ts | 14 -------------- 4 files changed, 9 insertions(+), 25 deletions(-) delete mode 100644 sdks/typescript/src/generated/hooks/registration.ts diff --git a/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/records/factory.py b/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/records/factory.py index 486b1a0c3..e340cbeef 100644 --- a/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/records/factory.py +++ b/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/records/factory.py @@ -8,13 +8,10 @@ from agent_control_models import Step from galileo_core.schemas.logging.session import Session from galileo_core.schemas.logging.span import ( - AgentSpan, - ControlSpan, LlmSpan, RetrieverSpan, Span, ToolSpan, - WorkflowSpan, ) from galileo_core.schemas.logging.trace import Trace diff --git a/sdks/typescript/scripts/generate-sdk.sh b/sdks/typescript/scripts/generate-sdk.sh index d0b059ea1..bf50ee7b5 100755 --- a/sdks/typescript/scripts/generate-sdk.sh +++ b/sdks/typescript/scripts/generate-sdk.sh @@ -67,12 +67,16 @@ mkdir -p "${GENERATED_DIR}" rsync -a --delete "${TMP_OUTPUT_DIR}/src/" "${GENERATED_DIR}/" rm -rf "${TMP_OUTPUT_DIR}" -# Speakeasy seeds hooks/registration.ts with @ts-expect-error, which fails -# strict typecheck when no hook is registered. Normalize to @ts-ignore so -# committed and regenerated output stay deterministic. +# Speakeasy's empty hook-registration scaffold varies by generator account. +# It is not part of the client API, so drop it when it contains only the +# generated placeholder to keep checked-in output deterministic. REGISTRATION_FILE="${GENERATED_DIR}/hooks/registration.ts" -if [[ -f "${REGISTRATION_FILE}" ]]; then - perl -0pi -e 's/@ts-expect-error/@ts-ignore/g' "${REGISTRATION_FILE}" +HOOKS_FILE="${GENERATED_DIR}/hooks/hooks.ts" +if [[ -f "${REGISTRATION_FILE}" ]] && \ + grep -Fq 'Add hooks by calling hooks.register' "${REGISTRATION_FILE}" && \ + ! grep -Eq '^[[:space:]]+hooks\.register[A-Za-z]+Hook\(' "${REGISTRATION_FILE}"; then + rm "${REGISTRATION_FILE}" + perl -0pi -e 's/\nimport \{ initHooks \} from "\.\/registration\.js";\n//; s/\n[[:space:]]*initHooks\(this\);\n/\n/' "${HOOKS_FILE}" fi # Default initAgent conflict mode to overwrite for SDK ergonomics while diff --git a/sdks/typescript/src/generated/hooks/hooks.ts b/sdks/typescript/src/generated/hooks/hooks.ts index 7ed9e0445..94ad0a323 100644 --- a/sdks/typescript/src/generated/hooks/hooks.ts +++ b/sdks/typescript/src/generated/hooks/hooks.ts @@ -18,8 +18,6 @@ import { SDKInitHook, } from "./types.js"; -import { initHooks } from "./registration.js"; - export class SDKHooks implements Hooks { sdkInitHooks: SDKInitHook[] = []; beforeCreateRequestHooks: BeforeCreateRequestHook[] = []; @@ -47,7 +45,6 @@ export class SDKHooks implements Hooks { this.registerAfterErrorHook(hook); } } - initHooks(this); } registerSDKInitHook(hook: SDKInitHook) { diff --git a/sdks/typescript/src/generated/hooks/registration.ts b/sdks/typescript/src/generated/hooks/registration.ts deleted file mode 100644 index 3a8fdeec1..000000000 --- a/sdks/typescript/src/generated/hooks/registration.ts +++ /dev/null @@ -1,14 +0,0 @@ -import { Hooks } from "./types.js"; - -/* - * This file is only ever generated once on the first generation and then is free to be modified. - * Any hooks you wish to add should be registered in the initHooks function. Feel free to define them - * in this file or in separate files in the hooks folder. - */ - -// @ts-ignore remove this line when you add your first hook and hooks is used -export function initHooks(hooks: Hooks) { - // Add hooks by calling hooks.register{ClientInit/BeforeCreateRequest/BeforeRequest/AfterSuccess/AfterError}Hook - // with an instance of a hook that implements that specific Hook interface - // Hooks are registered per SDK instance, and are valid for the lifetime of the SDK instance -} From ab8889485aefc322ba871e62c88c80d365f34236 Mon Sep 17 00:00:00 2001 From: Namrata Ghadi Date: Sat, 31 Oct 2026 15:23:34 -0700 Subject: [PATCH 4/5] test: cover gateway policy edge cases --- .../budget/tests/budget/test_budget.py | 168 +++++++++++++++++- models/tests/test_controls.py | 40 +++++ models/tests/test_gateway.py | 34 ++++ server/tests/test_gateway.py | 92 +++++++++- 4 files changed, 329 insertions(+), 5 deletions(-) create mode 100644 models/tests/test_gateway.py diff --git a/evaluators/contrib/budget/tests/budget/test_budget.py b/evaluators/contrib/budget/tests/budget/test_budget.py index 68e5cf249..8b5210790 100644 --- a/evaluators/contrib/budget/tests/budget/test_budget.py +++ b/evaluators/contrib/budget/tests/budget/test_budget.py @@ -5,25 +5,25 @@ from __future__ import annotations -from importlib.metadata import PackageNotFoundError, version import threading +from importlib.metadata import PackageNotFoundError, version from typing import Any -import pytest -from pydantic import ValidationError - import agent_control_evaluator_budget.budget.evaluator as budget_evaluator_module +import pytest from agent_control_evaluator_budget.budget.config import ( WINDOW_DAILY, WINDOW_MONTHLY, WINDOW_WEEKLY, BudgetEvaluatorConfig, BudgetLimitRule, + GatewayRoutingConfig, ModelPricing, ) from agent_control_evaluator_budget.budget.evaluator import ( BudgetEvaluator, _extract_tokens, + _gateway_route_metadata, clear_budget_stores, get_or_create_store, ) @@ -33,6 +33,8 @@ _compute_utilization, _derive_period_key, ) +from agent_control_models import GatewayModelCandidate, GatewayModelTier +from pydantic import ValidationError @pytest.fixture(autouse=True) @@ -2103,3 +2105,161 @@ async def test_inf_clock_does_not_crash(self) -> None: # Then: no crash assert len(snaps) == 1 assert snaps[0].spent == 100 + + +def _gateway_routing_config(*, threshold_tier: str | None = "economy") -> GatewayRoutingConfig: + return GatewayRoutingConfig( + tiers={ + "economy": GatewayModelTier( + models=[GatewayModelCandidate(provider="openai", model="small", protocol="openai")] + ), + "premium": GatewayModelTier( + models=[GatewayModelCandidate(provider="openai", model="large", protocol="openai")] + ), + }, + task_tiers={"default": "economy", "analysis": "premium"}, + threshold=0.8, + threshold_tier=threshold_tier, + ) + + +def _route_candidates(*models: str) -> list[dict[str, str]]: + return [ + {"provider": "openai", "model": model, "protocol": "openai"} for model in models + ] + + +def test_gateway_routing_config_rejects_unknown_tier_references() -> None: + economy_tier = {"economy": _gateway_routing_config().tiers["economy"]} + + with pytest.raises(ValidationError, match="task_tiers references unknown tiers"): + GatewayRoutingConfig(tiers=economy_tier, task_tiers={"default": "premium"}) + with pytest.raises(ValidationError, match="threshold_tier references unknown tier"): + GatewayRoutingConfig( + tiers=economy_tier, + task_tiers={"default": "economy"}, + threshold_tier="premium", + ) + + +def test_budget_config_requires_token_limit_for_gateway_routing() -> None: + with pytest.raises(ValidationError, match="requires at least one token budget limit"): + BudgetEvaluatorConfig( + limits=[BudgetLimitRule(limit=100)], + pricing={"small": ModelPricing()}, + model_path="model", + gateway_routing=_gateway_routing_config(), + ) + + +def test_gateway_routing_none_returns_no_route_metadata() -> None: + assert _gateway_route_metadata({}, None, [], over_limit=False) is None + + +def test_gateway_route_uses_default_task_and_zero_budget_snapshot() -> None: + route = _gateway_route_metadata( + {"candidate_models": _route_candidates("small")}, + _gateway_routing_config(), + [], + over_limit=False, + ) + + assert route is not None + assert route["allowed"] is True + assert route["model"] == "small" + assert route["task_type"] == "default" + assert route["budget"] == { + "used_tokens": 0, + "limit_tokens": 0, + "utilization": 0.0, + "threshold": 0.8, + "threshold_reached": False, + "over_limit": False, + } + + +def test_gateway_threshold_without_override_keeps_task_tier() -> None: + route = _gateway_route_metadata( + {"task_type": "analysis", "candidate_models": _route_candidates("large")}, + _gateway_routing_config(threshold_tier=None), + [{"limit_unit": "tokens", "spent_tokens": 8, "limit": 10, "utilization": 0.8}], + over_limit=False, + ) + + assert route is not None + assert route["model"] == "large" + assert route["budget"]["threshold_reached"] is True + assert "Budget threshold reached" in route["reason"] + + +def test_gateway_route_supports_legacy_model_and_rejects_unknown_tier() -> None: + legacy_route = _gateway_route_metadata( + { + "task_type": "analysis", + "candidate_models": "not-a-list", + "model": "large", + }, + _gateway_routing_config(), + [], + over_limit=False, + ) + unknown_route = _gateway_route_metadata( + {"task_type": "unknown", "candidate_models": _route_candidates("small")}, + GatewayRoutingConfig( + tiers=_gateway_routing_config().tiers, + task_tiers={"analysis": "premium"}, + ), + [], + over_limit=False, + ) + + assert legacy_route is not None and legacy_route["model"] == "large" + assert unknown_route is not None + assert unknown_route["allowed"] is False + assert unknown_route["tier"] is None + assert "No model tier" in unknown_route["reason"] + + +def test_gateway_route_skips_invalid_candidates_and_denies_incompatible_tier() -> None: + route = _gateway_route_metadata( + { + "task_type": "analysis", + "candidate_models": [ + "not-an-object", + {"provider": "openai", "model": "small", "protocol": "unknown"}, + ], + }, + _gateway_routing_config(), + [], + over_limit=False, + ) + + assert route is not None + assert route["allowed"] is False + assert route["model"] is None + assert "No compatible model candidate" in route["reason"] + + +def test_gateway_over_limit_denies_candidate_and_non_mapping_input_defaults_task() -> None: + over_limit = _gateway_route_metadata( + {"candidate_models": _route_candidates("small")}, + _gateway_routing_config(), + [], + over_limit=True, + ) + non_mapping = _gateway_route_metadata( + None, _gateway_routing_config(), [], over_limit=False + ) + non_string_task = _gateway_route_metadata( + {"task_type": 42}, _gateway_routing_config(), [], over_limit=False + ) + + assert over_limit is not None + assert over_limit["allowed"] is False + assert over_limit["provider"] is None + assert "exhausted" in over_limit["reason"] + assert non_mapping is not None + assert non_mapping["task_type"] == "default" + assert non_mapping["allowed"] is False + assert non_string_task is not None + assert non_string_task["task_type"] == "default" diff --git a/models/tests/test_controls.py b/models/tests/test_controls.py index dc36e67f8..76e0fcce0 100644 --- a/models/tests/test_controls.py +++ b/models/tests/test_controls.py @@ -107,6 +107,46 @@ def test_runtime_legacy_leaf_payload_is_canonicalized() -> None: assert dumped["condition"]["evaluator"]["name"] == "regex" +def _gateway_control_payload(*, decision: str, with_routing: bool) -> dict[str, object]: + evaluator_config: dict[str, object] = { + "limits": [{"limit": 1000, "limit_unit": "tokens"}], + } + if with_routing: + evaluator_config["gateway_routing"] = { + "tiers": { + "economy": { + "models": [ + {"provider": "openai", "model": "gpt-small", "protocol": "openai"} + ] + } + }, + "task_tiers": {"default": "economy"}, + } + return { + "execution": "server", + "scope": {"step_types": ["llm"], "stages": ["pre"]}, + "condition": { + "selector": {"path": "input"}, + "evaluator": {"name": "budget", "config": evaluator_config}, + }, + "action": {"decision": decision}, + } + + +def test_route_action_requires_gateway_routing_budget_evaluator() -> None: + with pytest.raises(ValidationError, match="route action requires a single budget evaluator"): + ControlDefinition.model_validate( + _gateway_control_payload(decision="route", with_routing=False) + ) + + +def test_gateway_routing_budget_requires_route_action() -> None: + with pytest.raises(ValidationError, match="requires the route action"): + ControlDefinition.model_validate( + _gateway_control_payload(decision="deny", with_routing=True) + ) + + def test_mixed_legacy_and_condition_fields_are_rejected() -> None: # Given: a payload that mixes canonical condition and legacy flat fields payload = { diff --git a/models/tests/test_gateway.py b/models/tests/test_gateway.py new file mode 100644 index 000000000..73d2c8d8f --- /dev/null +++ b/models/tests/test_gateway.py @@ -0,0 +1,34 @@ +"""Tests for the typed inference-gateway contracts.""" + +from __future__ import annotations + +import pytest +from agent_control_models import GatewayModelCandidate, GatewayModelTier, GatewayPreflightRequest +from pydantic import ValidationError + + +def test_gateway_preflight_request_normalizes_legacy_agent_and_accepts_candidate() -> None: + request = GatewayPreflightRequest.model_validate( + { + "agent_name": " Customer-Success-Agent ", + "task_type": "analysis", + "candidate_models": [ + {"provider": "openai", "model": "gpt-small", "protocol": "openai"} + ], + } + ) + + assert request.agent_name == "customer-success-agent" + assert request.candidate_models[0].model == "gpt-small" + + +def test_gateway_preflight_request_requires_a_candidate_or_legacy_model() -> None: + with pytest.raises(ValidationError, match="Provide candidate_models or the legacy model"): + GatewayPreflightRequest() + + +def test_gateway_model_candidate_rejects_unknown_protocol_and_empty_tier() -> None: + with pytest.raises(ValidationError): + GatewayModelCandidate(provider="openai", model="gpt-small", protocol="custom") # type: ignore[arg-type] + with pytest.raises(ValidationError): + GatewayModelTier(models=[]) diff --git a/server/tests/test_gateway.py b/server/tests/test_gateway.py index ff10bf7aa..1f31a091d 100644 --- a/server/tests/test_gateway.py +++ b/server/tests/test_gateway.py @@ -8,7 +8,10 @@ from agent_control_evaluator_budget.budget.evaluator import clear_budget_stores from agent_control_models import ControlMatch, EvaluationResponse, EvaluatorResult from agent_control_server.auth_framework import Operation, Principal, set_authorizer -from agent_control_server.endpoints.gateway import _find_gateway_route +from agent_control_server.endpoints.gateway import ( + _budget_state_from_metadata, + _find_gateway_route, +) from fastapi.testclient import TestClient @@ -242,6 +245,93 @@ def test_bound_deny_control_blocks_gateway_route(client: TestClient) -> None: assert "denied" in result["reason"].lower() +def test_preflight_returns_policy_decision_and_accepts_legacy_model(client: TestClient) -> None: + agent_name = "customer-success-agent" + _register_agent(client, agent_name) + set_authorizer(_AgentRuntimeAuthorizer(agent_name), operation=Operation.RUNTIME_USE) + + response = client.post( + f"/api/v1/gateway/agents/{agent_name}/preflight", + json={"task_type": "analysis", "model": "gpt-4.1", "prompt": "summarize"}, + ) + + assert response.status_code == 200, response.text + assert response.json()["allowed"] is True + assert response.json()["policy"] == f"{agent_name}-gateway-routing" + + +def test_preflight_without_bound_routing_control_denies(client: TestClient) -> None: + agent_name = "customer-success-agent" + response = client.post( + "/api/v1/agents/initAgent", + json={"agent": {"agent_name": agent_name}, "steps": []}, + ) + assert response.status_code == 200, response.text + set_authorizer(_AgentRuntimeAuthorizer(agent_name), operation=Operation.RUNTIME_USE) + + preflight = client.post( + f"/api/v1/gateway/agents/{agent_name}/preflight", + json={"task_type": "classification", "candidate_models": _candidates(agent_name)}, + ) + + assert preflight.status_code == 200, preflight.text + assert preflight.json()["allowed"] is False + assert preflight.json()["policy"] == "agent-controls" + assert preflight.json()["budget"] is None + + +def test_claude_code_denies_misconfigured_openai_route(client: TestClient) -> None: + agent_name = "claude-code" + response = client.post( + "/api/v1/agents/initAgent", + json={"agent": {"agent_name": agent_name}, "steps": []}, + ) + assert response.status_code == 200, response.text + # Bind the first-party policy deliberately to this agent to exercise the + # server-side Anthropic protocol guard independently of candidate filtering. + _bind_control(client, agent_name, _create_route_control(client, "customer-success-agent")) + set_authorizer(_AgentRuntimeAuthorizer(agent_name), operation=Operation.RUNTIME_USE) + + route = client.post( + f"/api/v1/gateway/agents/{agent_name}/route-decision", + json={ + "task_type": "analysis", + "candidate_models": [ + {"provider": "openai", "model": "gpt-4.1", "protocol": "openai"} + ], + }, + ) + + assert route.status_code == 200, route.text + assert route.json()["decision"] == "deny" + assert route.json()["provider"] is None + assert "Anthropic-compatible" in route.json()["reason"] + + +def test_usage_settlement_without_bound_routing_policy_returns_not_found( + client: TestClient, +) -> None: + agent_name = "customer-success-agent" + response = client.post( + "/api/v1/agents/initAgent", + json={"agent": {"agent_name": agent_name}, "steps": []}, + ) + assert response.status_code == 200, response.text + set_authorizer(_AgentRuntimeAuthorizer(agent_name), operation=Operation.RUNTIME_USE) + + settlement = client.post( + f"/api/v1/gateway/agents/{agent_name}/usage-settlement", + json={"model": "gpt-4.1", "input_tokens": 10, "output_tokens": 5}, + ) + + assert settlement.status_code == 404 + + +@pytest.mark.parametrize("metadata", [None, {"gateway_route": None}, {"budget": None}]) +def test_budget_state_parser_ignores_unusable_metadata(metadata: object) -> None: + assert _budget_state_from_metadata(metadata) is None + + def test_gateway_rejects_identity_mismatch(client: TestClient) -> None: set_authorizer(_AgentRuntimeAuthorizer("claude-code"), operation=Operation.RUNTIME_USE) response = client.post( From fc12987f9b178bd70542b88252c53b79dd171b33 Mon Sep 17 00:00:00 2001 From: Namrata Ghadi Date: Sun, 1 Nov 2026 16:24:26 -0800 Subject: [PATCH 5/5] fix auth header --- models/src/agent_control_models/gateway.py | 5 + models/tests/test_gateway.py | 39 ++++++- server/README.md | 24 +++- .../auth_framework/providers/http_upstream.py | 20 +++- .../agent_control_server/endpoints/gateway.py | 46 ++++++-- server/tests/test_auth_framework.py | 40 +++++++ server/tests/test_gateway.py | 103 +++++++++++++++++- 7 files changed, 261 insertions(+), 16 deletions(-) diff --git a/models/src/agent_control_models/gateway.py b/models/src/agent_control_models/gateway.py index 16a8b8920..1895adf5b 100644 --- a/models/src/agent_control_models/gateway.py +++ b/models/src/agent_control_models/gateway.py @@ -1,6 +1,7 @@ """Typed contracts for Agent Control's inference gateway POC.""" from typing import Literal, Self +from uuid import UUID from pydantic import Field, field_validator, model_validator @@ -35,6 +36,9 @@ class GatewayPreflightRequest(BaseModel): min_length=AGENT_NAME_MIN_LENGTH, pattern=AGENT_NAME_PATTERN, ) + # When present, runtime authorization is bound to this Galileo stream UUID. + # Agent name remains the policy and attribution identity in the route. + log_stream_id: UUID | None = None task_type: str = Field(default="default", min_length=1) task_context: JSONObject | None = None candidate_models: list[GatewayModelCandidate] = Field(default_factory=list) @@ -57,6 +61,7 @@ def require_candidate(self) -> Self: class GatewayUsageSettlementRequest(BaseModel): """Actual token usage reported by the gateway after inference.""" + log_stream_id: UUID | None = None model: str = Field(min_length=1) input_tokens: int = Field(ge=0) output_tokens: int = Field(ge=0) diff --git a/models/tests/test_gateway.py b/models/tests/test_gateway.py index 73d2c8d8f..4521deb96 100644 --- a/models/tests/test_gateway.py +++ b/models/tests/test_gateway.py @@ -3,7 +3,12 @@ from __future__ import annotations import pytest -from agent_control_models import GatewayModelCandidate, GatewayModelTier, GatewayPreflightRequest +from agent_control_models import ( + GatewayModelCandidate, + GatewayModelTier, + GatewayPreflightRequest, + GatewayUsageSettlementRequest, +) from pydantic import ValidationError @@ -27,6 +32,38 @@ def test_gateway_preflight_request_requires_a_candidate_or_legacy_model() -> Non GatewayPreflightRequest() +def test_gateway_requests_accept_typed_log_stream_uuid() -> None: + stream_id = "00000000-0000-4000-8000-000000000001" + preflight = GatewayPreflightRequest.model_validate( + { + "log_stream_id": stream_id, + "candidate_models": [ + {"provider": "openai", "model": "gpt-small", "protocol": "openai"} + ], + } + ) + settlement = GatewayUsageSettlementRequest.model_validate( + { + "log_stream_id": stream_id, + "model": "gpt-small", + "input_tokens": 1, + "output_tokens": 1, + } + ) + + assert str(preflight.log_stream_id) == stream_id + assert str(settlement.log_stream_id) == stream_id + with pytest.raises(ValidationError): + GatewayPreflightRequest.model_validate( + { + "log_stream_id": "not-a-uuid", + "candidate_models": [ + {"provider": "openai", "model": "gpt-small", "protocol": "openai"} + ], + } + ) + + def test_gateway_model_candidate_rejects_unknown_protocol_and_empty_tier() -> None: with pytest.raises(ValidationError): GatewayModelCandidate(provider="openai", model="gpt-small", protocol="custom") # type: ignore[arg-type] diff --git a/server/README.md b/server/README.md index f66438d7a..c6d933122 100644 --- a/server/README.md +++ b/server/README.md @@ -70,8 +70,18 @@ The route action is accepted only with a single budget evaluator configured with gateway_routing. It is non-blocking in ordinary ControlEngine safety results; the gateway maps its matched evaluator metadata to a route or deny. The Claude Code route is restricted at runtime to the Anthropic provider and -protocol. The request path must match the authenticated runtime identity; the -caller cannot choose another agent's bound controls in the body. +protocol. The route path selects the agent policy and attribution identity; the +runtime token binds either that agent or a log stream. A body-supplied agent +name cannot select another agent's controls. + +Gateway runtime tokens may be bound to a Galileo log stream rather than the +gateway agent name. In that mode the adapter includes the stream's UUID as +`log_stream_id` on preflight, route-decision, and usage-settlement requests; +Agent Control checks that UUID against the runtime token and includes controls +bound to `(log_stream, log_stream_id)` when evaluating the controls effective +for the path agent. The `{agent_name}` route remains the policy-selection and +attribution identity. Legacy agent-bound runtime tokens remain supported when +`log_stream_id` is omitted. `POST /api/v1/gateway/agents/{agent_name}/route-decision` combines control preflight and tier selection. A gateway that needs a separate eligibility check @@ -81,6 +91,7 @@ clients and must match it. A route-decision request looks like: ```json { + "log_stream_id": "00000000-0000-4000-8000-000000000001", "task_type": "analysis", "task_context": {"classification": "multi-step reasoning"}, "candidate_models": [ @@ -127,12 +138,21 @@ After the provider responds, post the actual usage to ```json { + "log_stream_id": "00000000-0000-4000-8000-000000000001", "model": "gpt-4.1", "input_tokens": 40000, "output_tokens": 20000 } ``` +Exchange the caller's request-time Galileo API key with +`POST /api/v1/auth/runtime-token-exchange` using +`{"target_type":"log_stream","target_id":""}`. When +Agent Control uses the HTTP-upstream authorization provider, it forwards the +incoming `X-API-Key` as both `X-API-Key` and `Galileo-API-Key` so Galileo AuthZ +can authorize the stream target. The key is forwarded per request and is not +stored by Agent Control. + The settlement response uses the same budget fields and identifies the bound control: diff --git a/server/src/agent_control_server/auth_framework/providers/http_upstream.py b/server/src/agent_control_server/auth_framework/providers/http_upstream.py index 79dcf4838..703e64ee4 100644 --- a/server/src/agent_control_server/auth_framework/providers/http_upstream.py +++ b/server/src/agent_control_server/auth_framework/providers/http_upstream.py @@ -355,7 +355,10 @@ async def authorize( operation: Operation, context: dict[str, Any] | None = None, ) -> Principal: - headers = self._forward_headers(request) + headers = self._forward_headers( + request, + map_galileo_api_key=operation is Operation.RUNTIME_TOKEN_EXCHANGE, + ) payload: dict[str, Any] = {"operation": operation.value} if context: payload["context"] = context @@ -428,17 +431,30 @@ async def _post_upstream( ) return response - def _forward_headers(self, request: Request) -> dict[str, str]: + def _forward_headers( + self, + request: Request, + *, + map_galileo_api_key: bool = False, + ) -> dict[str, str]: headers: dict[str, str] = {} seen: set[str] = set() + caller_api_key = request.headers.get("X-API-Key") if map_galileo_api_key else None for name in (*_DEFAULT_FORWARDED_HEADERS, *self._config.extra_forward_headers): lower = name.lower() if lower in seen: continue seen.add(lower) + # If an extra-forward configuration mentions Galileo's native + # header, the explicit X-API-Key mapping below remains the source + # of truth rather than allowing a second caller value to override it. + if lower == "galileo-api-key" and caller_api_key is not None: + continue value = request.headers.get(name) if value is not None: headers[name] = value + if caller_api_key is not None: + headers["Galileo-API-Key"] = caller_api_key if self._config.service_token is not None: headers[self._config.service_token_header] = self._config.service_token return headers diff --git a/server/src/agent_control_server/endpoints/gateway.py b/server/src/agent_control_server/endpoints/gateway.py index 2d42a2164..0814760e4 100644 --- a/server/src/agent_control_server/endpoints/gateway.py +++ b/server/src/agent_control_server/endpoints/gateway.py @@ -3,6 +3,7 @@ from __future__ import annotations from typing import Literal +from uuid import UUID from agent_control_engine.core import ControlEngine from agent_control_evaluator_budget.budget import ( # type: ignore[import-untyped] @@ -37,20 +38,49 @@ router = APIRouter(prefix="/gateway/agents", tags=["inference-gateway"]) -def _agent_context(request: Request) -> dict[str, str]: - """Bind runtime authorization to the agent named by the route.""" - return {"target_type": "agent", "target_id": request.path_params["agent_name"]} +async def _agent_context(request: Request) -> dict[str, str]: + """Bind runtime authorization to the route agent or supplied stream UUID.""" + target_id = request.path_params["agent_name"] + target_type = "agent" + try: + body = await request.json() + except (TypeError, ValueError): + body = None + if isinstance(body, dict) and body.get("log_stream_id") is not None: + raw_stream_id = body["log_stream_id"] + target_type = "log_stream" + try: + target_id = str(UUID(raw_stream_id)) + except (ValueError, TypeError, AttributeError): + target_id = str(raw_stream_id) + return {"target_type": target_type, "target_id": target_id} _runtime_agent = require_operation(Operation.RUNTIME_USE, context_builder=_agent_context) async def _get_gateway_controls(agent_name: str, principal: Principal) -> list[ControlAdapter]: - """Load controls effective for this authenticated runtime identity.""" - if principal.target_type != "agent" or principal.target_id != agent_name: + """Load agent controls plus bindings for the authenticated runtime target.""" + if principal.target_type == "agent" and principal.target_id == agent_name: + target_type = "agent" + target_id = agent_name + elif principal.target_type == "log_stream" and principal.target_id is not None: + try: + target_id = str(UUID(principal.target_id)) + except ValueError: + target_id = "" + if not target_id: + raise ForbiddenError( + error_code=ErrorCode.AUTH_INSUFFICIENT_PRIVILEGES, + detail="Gateway runtime identity must be bound to a valid log stream UUID.", + ) + target_type = "log_stream" + else: raise ForbiddenError( error_code=ErrorCode.AUTH_INSUFFICIENT_PRIVILEGES, - detail="Gateway requests must use a runtime identity bound to this agent.", + detail=( + "Gateway requests must use a runtime identity bound to this agent or log stream." + ), ) async with AsyncSessionLocal() as db: agent_result = await db.execute( @@ -69,8 +99,8 @@ async def _get_gateway_controls(agent_name: str, principal: Principal) -> list[C runtime_controls = await ControlService(db).list_runtime_controls_for_agent( agent_name, namespace_key=principal.namespace_key, - target_type="agent", - target_id=agent_name, + target_type=target_type, + target_id=target_id, allow_invalid_step_name_regex=True, ) return [ diff --git a/server/tests/test_auth_framework.py b/server/tests/test_auth_framework.py index 9e717b141..a0d80b0e8 100644 --- a/server/tests/test_auth_framework.py +++ b/server/tests/test_auth_framework.py @@ -281,6 +281,46 @@ def factory(request: httpx.Request) -> httpx.Response: assert principal == Principal(namespace_key="org-7", is_admin=True, caller_id="user-42") assert captured["url"] == "https://upstream.example/check" assert captured["headers"]["x-api-key"] == "caller-key" + assert "galileo-api-key" not in captured["headers"] + + +@pytest.mark.asyncio +async def test_http_upstream_maps_caller_api_key_to_galileo_authz_header( + caplog: pytest.LogCaptureFixture, +): + """The gateway's request-time key reaches Galileo AuthZ under its header name.""" + captured: dict[str, Any] = {} + stream_id = "00000000-0000-4000-8000-000000000001" + + def factory(request: httpx.Request) -> httpx.Response: + captured["headers"] = dict(request.headers) + captured["payload"] = json.loads(request.content) + return httpx.Response( + 200, + json={ + "namespace_key": "org-7", + "caller_id": "user-42", + "target_type": "log_stream", + "target_id": stream_id, + "scopes": ["runtime.use"], + }, + ) + + provider = _build_upstream(factory) + principal = await provider.authorize( + _build_request(headers={"X-API-Key": "request-time-key"}), + Operation.RUNTIME_TOKEN_EXCHANGE, + context={"target_type": "log_stream", "target_id": stream_id}, + ) + + assert captured["headers"]["x-api-key"] == "request-time-key" + assert captured["headers"]["galileo-api-key"] == "request-time-key" + assert captured["payload"] == { + "operation": "runtime.token_exchange", + "context": {"target_type": "log_stream", "target_id": stream_id}, + } + assert principal.scopes == ("runtime.use",) + assert "request-time-key" not in caplog.text @pytest.mark.asyncio diff --git a/server/tests/test_gateway.py b/server/tests/test_gateway.py index 1f31a091d..c97d08415 100644 --- a/server/tests/test_gateway.py +++ b/server/tests/test_gateway.py @@ -3,11 +3,18 @@ from __future__ import annotations from typing import Any +from uuid import UUID, uuid4 import pytest from agent_control_evaluator_budget.budget.evaluator import clear_budget_stores from agent_control_models import ControlMatch, EvaluationResponse, EvaluatorResult -from agent_control_server.auth_framework import Operation, Principal, set_authorizer +from agent_control_server.auth_framework import ( + Operation, + Principal, + set_authorizer, +) +from agent_control_server.auth_framework.providers import LocalJwtVerifyProvider +from agent_control_server.auth_framework.runtime_token import mint_runtime_token from agent_control_server.endpoints.gateway import ( _budget_state_from_metadata, _find_gateway_route, @@ -123,10 +130,24 @@ def _candidates(agent_name: str) -> list[dict[str, str]]: ] -def _route(client: TestClient, agent_name: str, task_type: str) -> dict[str, Any]: +def _route( + client: TestClient, + agent_name: str, + task_type: str, + *, + log_stream_id: UUID | None = None, + headers: dict[str, str] | None = None, +) -> dict[str, Any]: + payload: dict[str, Any] = { + "task_type": task_type, + "candidate_models": _candidates(agent_name), + } + if log_stream_id is not None: + payload["log_stream_id"] = str(log_stream_id) response = client.post( f"/api/v1/gateway/agents/{agent_name}/route-decision", - json={"task_type": task_type, "candidate_models": _candidates(agent_name)}, + json=payload, + headers=headers, ) assert response.status_code == 200, response.text return response.json() @@ -220,6 +241,82 @@ def test_controls_keep_separate_agent_budgets_and_claude_stays_anthropic( assert claude_openai_only.json()["model"] is None +def test_log_stream_runtime_token_loads_controls_bound_to_stream(client: TestClient) -> None: + # Given: an agent with a route control bound only to its Galileo log stream + agent_name = "customer-success-agent" + stream_id = uuid4() + registered = client.post( + "/api/v1/agents/initAgent", + json={"agent": {"agent_name": agent_name}, "steps": []}, + ) + assert registered.status_code == 200, registered.text + control_id = _create_route_control(client, agent_name) + binding = client.put( + "/api/v1/control-bindings", + json={ + "target_type": "log_stream", + "target_id": str(stream_id), + "control_id": control_id, + }, + ) + assert binding.status_code == 200, binding.text + + secret = "gateway-log-stream-runtime-secret-32-bytes" + token, _ = mint_runtime_token( + namespace_key="default", + actor_id="gateway-caller", + target_type="log_stream", + target_id=str(stream_id), + scopes=(Operation.RUNTIME_USE.value,), + secret=secret, + ttl_seconds=300, + ) + set_authorizer( + LocalJwtVerifyProvider(secret=secret), + operation=Operation.RUNTIME_USE, + ) + headers = {"Authorization": f"Bearer {token}"} + + # When: the gateway identifies the agent by route and the stream by request body + route = _route( + client, + agent_name, + "analysis", + log_stream_id=stream_id, + headers=headers, + ) + settlement = client.post( + f"/api/v1/gateway/agents/{agent_name}/usage-settlement", + headers=headers, + json={ + "log_stream_id": str(stream_id), + "model": "gpt-4.1", + "input_tokens": 10, + "output_tokens": 5, + }, + ) + + # Then: the stream-bound grant authorizes both calls and stream controls apply + assert route["decision"] == "route" + assert route["model"] == "gpt-4.1" + assert route["policy"] == f"{agent_name}-gateway-routing" + assert settlement.status_code == 200, settlement.text + assert settlement.json()["budget"]["used_tokens"] == 15 + + # And: a different stream cannot reuse the stream-bound token + mismatch = client.post( + f"/api/v1/gateway/agents/{agent_name}/route-decision", + headers=headers, + json={ + "task_type": "analysis", + "log_stream_id": str(uuid4()), + "candidate_models": _candidates(agent_name), + }, + ) + assert mismatch.status_code == 403 + assert mismatch.json()["detail"] == "Runtime token target_id does not match the request." + + def test_bound_deny_control_blocks_gateway_route(client: TestClient) -> None: agent_name = "customer-success-agent" _register_agent(client, agent_name)