diff --git a/docs/advanced/configuration.mdx b/docs/advanced/configuration.mdx index ef007b790..8b08af4b8 100644 --- a/docs/advanced/configuration.mdx +++ b/docs/advanced/configuration.mdx @@ -74,6 +74,45 @@ skill. An explicit model passed to `create_agent` takes precedence over routes. } ``` +### Per-model budgets and fallback chains + +Set `llm.model_budgets_usd` and `llm.model_fallbacks` to move every agent using +an exhausted model to its next configured model. Budgets are cumulative per +model across the whole scan, so the same rule covers the orchestrator, default +children, explicit child routes, and skill routes. Fallback targets can have +their own budget and fallback, creating any number of layers. + +```json +{ + "llm": { + "model": "openai/gpt-5.4", + "subagent_model": "z-ai/glm-4.7", + "skill_model_routes": [ + {"skill": "xss", "model": "openai/gpt-5.4"} + ], + "model_budgets_usd": { + "openai/gpt-5.4": 20.0, + "z-ai/glm-4.7": 10.0 + }, + "model_fallbacks": { + "openai/gpt-5.4": "z-ai/glm-4.7", + "z-ai/glm-4.7": "deepseek/deepseek-chat" + } + } +} +``` + +A fallback is checked after a response returns, so a call can slightly +overshoot its model budget. Strix discards that boundary response before any of +its tool calls execute, updates the agent's persisted active model, and replays +the unchanged session with the fallback. This avoids duplicate side effects and +keeps scan resume on the selected fallback layer. A model budget without a +fallback only tracks usage; it does not stop the scan. The scan-wide +`--max-budget-usd` remains a hard stop and takes precedence. + +The TUI's bottom-right panel displays cumulative tokens and estimated cost for +each model actually used, including fallback models. + ## Optional Features diff --git a/strix/agents/factory.py b/strix/agents/factory.py index a13135671..f8995de4b 100644 --- a/strix/agents/factory.py +++ b/strix/agents/factory.py @@ -17,8 +17,12 @@ from pydantic import ValidationError from strix.agents.prompt import render_system_prompt -from strix.config.models import _same_provider, uses_chat_completions_tool_schema from strix.core.inputs import make_model_settings +from strix.core.model_routing import ( + chain_uses_chat_completions_tools, + resolve_budget_model, + resolve_route_api_key, +) from strix.tools.agents_graph.tools import ( agent_finish, create_agent, @@ -63,7 +67,7 @@ from agents import RunContextWrapper from agents.tool import FunctionToolResult - from strix.config.settings import LlmSettings, ReasoningEffort, Settings, SkillModelRoute + from strix.config.settings import ReasoningEffort, Settings, SkillModelRoute logger = logging.getLogger(__name__) @@ -482,22 +486,6 @@ def build_strix_agent( ) -def _resolve_child_api_key( - llm: LlmSettings, resolved_model: str, default_model: str -) -> str | None: - """Pick the credential a child should send as a per-call ``api_key``. - - Prefer an explicit ``SUBAGENT_LLM_API_KEY``. Otherwise reuse the - orchestrator key only when the child shares the orchestrator's provider; - cross-provider children fall back to their provider's ambient env/auth. - """ - if llm.subagent_api_key and llm.subagent_api_key.strip(): - return llm.subagent_api_key.strip() - if llm.api_key and _same_provider(resolved_model, default_model): - return llm.api_key - return None - - def make_child_factory( *, settings: Settings, @@ -526,12 +514,13 @@ def _factory( model: str | None = None, ) -> SandboxAgent[Any]: route = _matching_route(skills) if model is None else None - resolved_model = ( + configured_model = ( model or (route.model if route is not None else None) or llm.subagent_model or default_model ).strip() + resolved_model = resolve_budget_model(configured_model, llm) reasoning: ReasoningEffort | None = ( route.reasoning_effort if route is not None and route.reasoning_effort is not None @@ -539,7 +528,7 @@ def _factory( ) if reasoning is None: reasoning = llm.reasoning_effort - child_api_key = _resolve_child_api_key(llm, resolved_model, default_model) + child_api_key = resolve_route_api_key(llm, resolved_model, default_model) child_model_settings = make_model_settings( reasoning, model_name=resolved_model, @@ -553,8 +542,8 @@ def _factory( scan_mode=scan_mode, is_whitebox=is_whitebox, interactive=interactive, - chat_completions_tools=uses_chat_completions_tool_schema( - resolved_model, + chat_completions_tools=chain_uses_chat_completions_tools( + configured_model, settings, ), model=resolved_model, diff --git a/strix/config/__init__.py b/strix/config/__init__.py index 27c84008b..97d80ca75 100644 --- a/strix/config/__init__.py +++ b/strix/config/__init__.py @@ -19,6 +19,7 @@ from strix.config.settings import ( IntegrationSettings, LlmSettings, + ReasoningEffort, RuntimeSettings, Settings, SkillModelRoute, @@ -29,6 +30,7 @@ __all__ = [ "IntegrationSettings", "LlmSettings", + "ReasoningEffort", "RuntimeSettings", "Settings", "SkillModelRoute", diff --git a/strix/config/models.py b/strix/config/models.py index 98d4520a2..bdc8bc117 100644 --- a/strix/config/models.py +++ b/strix/config/models.py @@ -125,7 +125,7 @@ def configure_sdk_model_defaults(settings: Settings) -> None: # orchestrator credential into differently-routed child providers. # Provider-specific environment variables keep concurrent routes apart. _mirror_api_key_to_provider_env(llm.model, llm.api_key) - for model_name in _subagent_model_names(settings): + for model_name in _configured_model_names(settings): key = llm.subagent_api_key or ( llm.api_key if _same_provider(model_name, llm.model) else None ) @@ -152,14 +152,24 @@ def _same_provider(first: str | None, second: str | None) -> bool: return _provider_prefix(first) == _provider_prefix(second) -def _subagent_model_names(settings: Settings) -> set[str]: +def _configured_model_names(settings: Settings) -> set[str]: llm = settings.llm names = {route.model for route in llm.skill_model_routes if route.model} + names.update(llm.model_budgets_usd) + names.update(llm.model_fallbacks) + names.update(llm.model_fallbacks.values()) + if llm.model: + names.add(llm.model) if llm.subagent_model: names.add(llm.subagent_model) return names +# Kept as a private compatibility alias for callers/tests from earlier releases. +def _subagent_model_names(settings: Settings) -> set[str]: + return _configured_model_names(settings) + + def _mirror_api_key_to_provider_env(model_name: str | None, api_key: str) -> None: if not model_name: return diff --git a/strix/config/settings.py b/strix/config/settings.py index d9a1fda99..84ede1446 100644 --- a/strix/config/settings.py +++ b/strix/config/settings.py @@ -2,9 +2,10 @@ from __future__ import annotations +import math from typing import Literal -from pydantic import AliasChoices, BaseModel, Field, model_validator +from pydantic import AliasChoices, BaseModel, Field, field_validator, model_validator from pydantic_settings import BaseSettings, SettingsConfigDict @@ -73,6 +74,14 @@ class LlmSettings(BaseSettings): ) subagent_model: str | None = Field(default=None, alias="STRIX_SUBAGENT_MODEL") skill_model_routes: list[SkillModelRoute] = Field(default_factory=list) + model_budgets_usd: dict[str, float] = Field( + default_factory=dict, + alias="STRIX_MODEL_BUDGETS_USD", + ) + model_fallbacks: dict[str, str] = Field( + default_factory=dict, + alias="STRIX_MODEL_FALLBACKS", + ) api_key: str | None = Field( default=None, validation_alias=AliasChoices("LLM_API_KEY", "OPENAI_API_KEY"), @@ -99,6 +108,95 @@ class LlmSettings(BaseSettings): ) timeout: int = Field(default=300, alias="LLM_TIMEOUT") + @field_validator("model_budgets_usd", mode="before") + @classmethod + def validate_model_budgets(cls, value: object) -> object: + if value is None: + return {} + if not isinstance(value, dict): + raise TypeError("model_budgets_usd must be a model-to-USD mapping") + normalized: dict[str, float] = {} + for raw_model, raw_budget in value.items(): + model = str(raw_model).strip() + if not model: + raise ValueError("model_budgets_usd contains an empty model name") + try: + budget = float(raw_budget) + except (TypeError, ValueError) as exc: + raise ValueError(f"budget for {model!r} must be a number") from exc + if not math.isfinite(budget) or budget <= 0: + raise ValueError(f"budget for {model!r} must be finite and greater than 0") + normalized[model] = budget + return normalized + + @field_validator("model_fallbacks", mode="before") + @classmethod + def validate_model_fallbacks(cls, value: object) -> object: + if value is None: + return {} + if not isinstance(value, dict): + raise TypeError("model_fallbacks must be a model-to-model mapping") + normalized: dict[str, str] = {} + for raw_model, raw_fallback in value.items(): + model = str(raw_model).strip() + fallback = str(raw_fallback).strip() + if not model or not fallback: + raise ValueError("model_fallbacks cannot contain empty model names") + if model.lower() == fallback.lower(): + raise ValueError(f"model {model!r} cannot fall back to itself") + normalized[model] = fallback + return normalized + + @model_validator(mode="after") + def validate_fallback_graph(self) -> LlmSettings: + budget_keys = {name.lower() for name in self.model_budgets_usd} + fallbacks = {name.lower(): target.lower() for name, target in self.model_fallbacks.items()} + display_names = {name.lower(): name for name in self.model_fallbacks} + for source in fallbacks: + if source not in budget_keys: + raise ValueError( + f"model_fallbacks source {display_names[source]!r} requires a " + "model_budgets_usd entry" + ) + seen: set[str] = set() + current = source + while current in fallbacks: + if current in seen: + raise ValueError("model_fallbacks contains a cycle") + seen.add(current) + current = fallbacks[current] + return self + + def budget_for(self, model: str) -> float | None: + normalized = model.strip().lower() + return next( + ( + budget + for name, budget in self.model_budgets_usd.items() + if name.lower() == normalized + ), + None, + ) + + def fallback_for(self, model: str) -> str | None: + normalized = model.strip().lower() + return next( + ( + fallback + for name, fallback in self.model_fallbacks.items() + if name.lower() == normalized + ), + None, + ) + + def model_chain(self, model: str) -> list[str]: + chain: list[str] = [] + current: str | None = model.strip() + while current: + chain.append(current) + current = self.fallback_for(current) + return chain + class RuntimeSettings(BaseSettings): model_config = _BASE_CONFIG diff --git a/strix/core/agents.py b/strix/core/agents.py index c789f9b43..aa5d2af6f 100644 --- a/strix/core/agents.py +++ b/strix/core/agents.py @@ -10,6 +10,8 @@ from pathlib import Path from typing import TYPE_CHECKING, Any, Literal, cast +from strix.core.inputs import make_model_settings +from strix.core.model_routing import resolve_route_api_key from strix.core.sessions import session_write_lock @@ -25,6 +27,7 @@ @dataclass(slots=True) class AgentRuntime: + agent: Any | None = None session: Session | None = None task: asyncio.Task[Any] | None = None stream: Any | None = None @@ -92,12 +95,15 @@ async def attach_runtime( self, agent_id: str, *, + agent: Any | None = None, session: Session | None = None, task: asyncio.Task[Any] | None = None, interrupt_on_message: bool | None = None, ) -> None: async with self._lock: runtime = self.runtimes.setdefault(agent_id, AgentRuntime()) + if agent is not None: + runtime.agent = agent if session is not None: runtime.session = session if task is not None: @@ -111,6 +117,52 @@ async def mark_running(self, agent_id: str) -> None: self.statuses[agent_id] = "running" await self._maybe_snapshot() + async def transition_model( + self, + previous_model: str, + next_model: str, + *, + llm_settings: Any | None = None, + ) -> None: + """Persist and apply a scan-wide model fallback to every matching agent.""" + previous = previous_model.strip().lower() + changed: list[str] = [] + async with self._lock: + for agent_id, metadata in self.metadata.items(): + current = metadata.get("model") + if not isinstance(current, str) or current.strip().lower() != previous: + continue + metadata["model"] = next_model + runtime_agent = self.runtimes.setdefault(agent_id, AgentRuntime()).agent + if runtime_agent is not None: + runtime_agent.model = next_model + if llm_settings is not None: + reasoning = ( + llm_settings.reasoning_effort + if self.parent_of.get(agent_id) is None + else ( + llm_settings.subagent_reasoning_effort + or llm_settings.reasoning_effort + ) + ) + runtime_agent.model_settings = make_model_settings( + reasoning, + model_name=next_model, + force_required_tool_choice=llm_settings.force_required_tool_choice, + api_key=resolve_route_api_key( + llm_settings, next_model, llm_settings.model or next_model + ), + ) + changed.append(agent_id) + if changed: + logger.info( + "agent.model fallback %s -> %s (%d agent(s))", + previous_model, + next_model, + len(changed), + ) + await self._maybe_snapshot() + async def park_waiting(self, agent_id: str) -> None: await self.set_status(agent_id, "waiting") diff --git a/strix/core/execution.py b/strix/core/execution.py index ab72243b0..06aef348a 100644 --- a/strix/core/execution.py +++ b/strix/core/execution.py @@ -15,8 +15,8 @@ from docker import errors as docker_errors # type: ignore[import-untyped, unused-ignore] from openai import APIError -from strix.core.hooks import BudgetExceededError -from strix.core.inputs import child_initial_input +from strix.core.hooks import BudgetExceededError, ModelFallbackError +from strix.core.inputs import child_initial_input, make_model_settings from strix.core.sessions import ( enforce_image_budget, open_agent_session, @@ -32,6 +32,7 @@ from agents.memory import Session, SQLiteSession from agents.result import RunResultBase + from strix.config.settings import Settings from strix.core.agents import AgentCoordinator, Status @@ -56,9 +57,11 @@ async def run_agent_loop( start_parked: bool = False, event_sink: StreamEventSink | None = None, hooks: RunHooks[dict[str, Any]] | None = None, + settings: Settings | None = None, ) -> RunResultBase | None: await coordinator.attach_runtime( agent_id, + agent=agent, session=session, interrupt_on_message=interactive, ) @@ -78,6 +81,7 @@ async def run_agent_loop( interactive=interactive, event_sink=event_sink, hooks=hooks, + settings=settings, ) else: result = await _run_noninteractive_until_lifecycle( @@ -91,6 +95,7 @@ async def run_agent_loop( session=session, event_sink=event_sink, hooks=hooks, + settings=settings, ) if not interactive: @@ -119,6 +124,7 @@ async def run_agent_loop( interactive=interactive, event_sink=event_sink, hooks=hooks, + settings=settings, ) @@ -139,6 +145,7 @@ async def spawn_child_agent( model: str | None = None, event_sink: StreamEventSink | None = None, hooks: RunHooks[dict[str, Any]] | None = None, + settings: Settings | None = None, ) -> dict[str, Any]: parent_id = parent_ctx.get("agent_id") if not isinstance(parent_id, str): @@ -182,6 +189,7 @@ async def spawn_child_agent( ), event_sink=event_sink, hooks=hooks, + settings=settings, ) return { @@ -206,6 +214,7 @@ async def respawn_subagents( root_id: str, event_sink: StreamEventSink | None = None, hooks: RunHooks[dict[str, Any]] | None = None, + settings: Settings | None = None, ) -> None: async with coordinator._lock: agents_snapshot = [ @@ -265,6 +274,7 @@ async def respawn_subagents( start_parked=start_parked, event_sink=event_sink, hooks=hooks, + settings=settings, ) logger.info( "respawned %s (%s) parent=%s task_len=%d", @@ -291,6 +301,7 @@ async def _run_noninteractive_until_lifecycle( session: Session | None, event_sink: StreamEventSink | None, hooks: RunHooks[dict[str, Any]] | None, + settings: Settings | None, ) -> RunResultBase | None: """Non-chat mode keeps running until finish_scan / agent_finish settles status.""" result: RunResultBase | None = None @@ -315,6 +326,7 @@ async def _run_noninteractive_until_lifecycle( interactive=False, event_sink=event_sink, hooks=hooks, + settings=settings, ) status = await _agent_status(coordinator, agent_id) @@ -360,6 +372,7 @@ async def _run_cycle( # noqa: PLR0912, PLR0915 interactive: bool, event_sink: StreamEventSink | None, hooks: RunHooks[dict[str, Any]] | None, + settings: Settings | None, ) -> RunResultBase | None: image_strips = 0 while True: @@ -392,8 +405,8 @@ async def _run_cycle( # noqa: PLR0912, PLR0915 logger.exception("stream event sink failed for %s", agent_id) if stream.run_loop_exception is not None: raise stream.run_loop_exception - except BudgetExceededError: - # A RuntimeError subclass: re-raise explicitly so it is never + except (BudgetExceededError, ModelFallbackError): + # RuntimeError subclasses: re-raise explicitly so neither is # mistaken for the LiteLLM "after shutdown" race below. raise except RuntimeError as stream_exc: @@ -413,6 +426,18 @@ async def _run_cycle( # noqa: PLR0912, PLR0915 ) finally: await coordinator.detach_stream(agent_id, stream) + except ModelFallbackError as exc: + if settings is None: + raise RuntimeError("model fallback requires scan settings") from exc + _refresh_agent_for_model( + agent, + exc.next_model, + settings, + is_root=context.get("parent_id") is None, + ) + logger.info("agent %s %s", agent_id, exc) + input_data = [] if session is not None else input_data + continue except BudgetExceededError as exc: logger.info( "agent %s reached the scan budget limit; stopping the scan: %s", agent_id, exc @@ -459,6 +484,28 @@ async def _run_cycle( # noqa: PLR0912, PLR0915 return stream +def _refresh_agent_for_model( + agent: Any, + model: str, + settings: Settings, + *, + is_root: bool, +) -> None: + """Apply a fallback route before replaying the unmodified SDK session.""" + llm = settings.llm + reasoning = ( + llm.reasoning_effort + if is_root + else (llm.subagent_reasoning_effort or llm.reasoning_effort) + ) + agent.model = model + agent.model_settings = make_model_settings( + reasoning, + model_name=model, + force_required_tool_choice=llm.force_required_tool_choice, + ) + + async def _settle_run_result( coordinator: AgentCoordinator, agent_id: str, @@ -560,10 +607,11 @@ async def _start_child_runner( start_parked: bool = False, event_sink: StreamEventSink | None = None, hooks: RunHooks[dict[str, Any]] | None = None, + settings: Settings | None = None, ) -> None: session = open_agent_session(child_id, agents_db_path) sessions_to_close.append(session) - await coordinator.attach_runtime(child_id, session=session) + await coordinator.attach_runtime(child_id, agent=child_agent, session=session) child_ctx: dict[str, Any] = dict(parent_ctx) child_ctx["agent_id"] = child_id @@ -590,6 +638,7 @@ async def _child_loop() -> None: start_parked=start_parked, event_sink=event_sink, hooks=hooks, + settings=settings, ) except BudgetExceededError: logger.info("child %s stopped after reaching the scan budget limit", child_id) diff --git a/strix/core/hooks.py b/strix/core/hooks.py index 8edf1b8d6..32b6b91e3 100644 --- a/strix/core/hooks.py +++ b/strix/core/hooks.py @@ -8,6 +8,8 @@ from agents.lifecycle import RunHooks +from strix.core.inputs import make_model_settings +from strix.core.model_routing import resolve_budget_model, resolve_route_api_key from strix.report.state import get_global_report_state @@ -16,24 +18,77 @@ from agents.agent import Agent from agents.items import ModelResponse + from strix.config.settings import LlmSettings + logger = logging.getLogger(__name__) class BudgetExceededError(RuntimeError): - """Raised when the accumulated LLM cost reaches the configured budget.""" + """Raised when the accumulated scan-wide LLM cost reaches its hard limit.""" + + +class ModelFallbackError(RuntimeError): + """Interrupt the current cycle after atomically moving an agent to its next model.""" + + def __init__( + self, + *, + previous_model: str, + next_model: str, + spent: float, + budget: float, + ) -> None: + self.previous_model = previous_model + self.next_model = next_model + self.spent = spent + self.budget = budget + super().__init__( + f"model {previous_model} reached its ${budget:.2f} budget " + f"(spent ${spent:.4f}); continuing with {next_model}" + ) class ReportUsageHooks(RunHooks[dict[str, Any]]): """Persist SDK-native usage after every model response.""" - def __init__(self, *, model: str, max_budget_usd: float | None = None) -> None: + def __init__( + self, + *, + model: str, + max_budget_usd: float | None = None, + llm_settings: LlmSettings | None = None, + ) -> None: if max_budget_usd is not None and ( not math.isfinite(max_budget_usd) or max_budget_usd <= 0 ): raise ValueError("max_budget_usd must be a finite number greater than 0") self._model = model self._max_budget_usd = max_budget_usd + self._llm_settings = llm_settings + self._inflight_models: dict[str, str] = {} + + @staticmethod + def _agent_id(context: RunContextWrapper[dict[str, Any]], agent: Agent[dict[str, Any]]) -> str: + ctx = context.context if isinstance(context.context, dict) else {} + agent_id = ctx.get("agent_id") + if isinstance(agent_id, str) and agent_id: + return agent_id + agent_name = getattr(agent, "name", None) + return agent_name if isinstance(agent_name, str) and agent_name else "unknown" + + async def on_llm_start( + self, + context: RunContextWrapper[dict[str, Any]], + agent: Agent[dict[str, Any]], + system_prompt: str | None, + input_items: list[Any], + ) -> None: + del system_prompt, input_items + model = getattr(agent, "model", None) + if not isinstance(model, str) or not model: + model = self._model + self._inflight_models[self._agent_id(context, agent)] = model async def on_llm_end( self, @@ -49,12 +104,16 @@ async def on_llm_end( agent_name = getattr(agent, "name", None) if not isinstance(agent_name, str): agent_name = None - agent_id = ctx.get("agent_id") - if not isinstance(agent_id, str) or not agent_id: - agent_id = agent_name or "unknown" + agent_id = self._agent_id(context, agent) + # A different agent can exhaust this model while this call is still in + # flight and mutate all matching public agents. Attribute the completed + # response to the route captured at on_llm_start, not that newer route. agent_model = getattr(agent, "model", None) - model = agent_model if isinstance(agent_model, str) and agent_model else self._model + fallback_model = ( + agent_model if isinstance(agent_model, str) and agent_model else self._model + ) + model = self._inflight_models.pop(agent_id, fallback_model) try: report_state.record_sdk_usage( agent_id=agent_id, @@ -71,3 +130,56 @@ async def on_llm_end( raise BudgetExceededError( f"Token budget of ${self._max_budget_usd:.2f} exceeded (spent ${cost:.4f})" ) + + llm = self._llm_settings + if llm is None: + return + model_budget = llm.budget_for(model) + fallback = llm.fallback_for(model) + if model_budget is None or fallback is None: + return + model_cost = report_state.get_model_llm_cost(model) + if model_cost < model_budget: + return + + current_agent_model = getattr(agent, "model", None) + if ( + isinstance(current_agent_model, str) + and current_agent_model.strip().lower() != model.strip().lower() + ): + # Another concurrent response already transitioned this shared + # model route while this request was in flight. Discard this stale + # response and retry on the route already selected for the agent. + raise ModelFallbackError( + previous_model=model, + next_model=current_agent_model, + spent=model_cost, + budget=model_budget, + ) + + # Abort at an LLM-response boundary. The SDK has not executed any tool + # calls from this response yet, so retrying the same persisted session + # with the fallback cannot duplicate side effects or expose a partial + # assistant turn to the new model. + next_model = resolve_budget_model(fallback, llm) + agent.model = next_model + reasoning = ( + llm.reasoning_effort + if ctx.get("parent_id") is None + else (llm.subagent_reasoning_effort or llm.reasoning_effort) + ) + agent.model_settings = make_model_settings( + reasoning, + model_name=next_model, + force_required_tool_choice=llm.force_required_tool_choice, + api_key=resolve_route_api_key(llm, next_model, self._model), + ) + coordinator = ctx.get("coordinator") + if coordinator is not None and hasattr(coordinator, "transition_model"): + await coordinator.transition_model(model, next_model, llm_settings=llm) + raise ModelFallbackError( + previous_model=model, + next_model=next_model, + spent=model_cost, + budget=model_budget, + ) diff --git a/strix/core/model_routing.py b/strix/core/model_routing.py new file mode 100644 index 000000000..96a4c7701 --- /dev/null +++ b/strix/core/model_routing.py @@ -0,0 +1,63 @@ +"""Runtime helpers for cumulative per-model budget fallback routing.""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +from strix.config.models import _same_provider, uses_chat_completions_tool_schema +from strix.report.state import get_global_report_state + + +if TYPE_CHECKING: + from strix.config.settings import LlmSettings, Settings + + +logger = logging.getLogger(__name__) + + +def resolve_route_api_key(llm: LlmSettings, resolved_model: str, default_model: str) -> str | None: + """Pick the credential a routed agent should send as a per-call ``api_key``. + + Prefer an explicit ``SUBAGENT_LLM_API_KEY``. Otherwise reuse the + orchestrator key only when the routed model shares the orchestrator's + provider; cross-provider routes fall back to ambient env/auth. Env vars are + global per provider, so per-call keys are the only way to keep same-provider + routes (root vs. child/fallback) on distinct credentials. + """ + if llm.subagent_api_key and llm.subagent_api_key.strip(): + return llm.subagent_api_key.strip() + if llm.api_key and _same_provider(resolved_model, default_model): + return llm.api_key + return None + + +def resolve_budget_model(model: str, llm: LlmSettings) -> str: + """Skip already-exhausted fallback layers when creating or resuming an agent.""" + if not hasattr(llm, "budget_for") or not hasattr(llm, "fallback_for"): + return model.strip() + state = get_global_report_state() + current = model.strip() + while state is not None: + budget = llm.budget_for(current) + fallback = llm.fallback_for(current) + if budget is None or fallback is None: + break + if state.get_model_llm_cost(current) < budget: + break + logger.info( + "model route skips exhausted layer %s ($%.4f / $%.2f) -> %s", + current, + state.get_model_llm_cost(current), + budget, + fallback, + ) + current = fallback + return current + + +def chain_uses_chat_completions_tools(model: str, settings: Settings) -> bool: + """Use a tool representation accepted by every possible fallback layer.""" + model_chain = getattr(settings.llm, "model_chain", None) + chain = model_chain(model) if callable(model_chain) else [model] + return any(uses_chat_completions_tool_schema(candidate, settings) for candidate in chain) diff --git a/strix/core/runner.py b/strix/core/runner.py index 0e32ac04c..c3b89c8f1 100644 --- a/strix/core/runner.py +++ b/strix/core/runner.py @@ -16,11 +16,7 @@ from strix.agents.factory import build_strix_agent, make_child_factory from strix.agents.prompt import render_system_prompt from strix.config import load_settings -from strix.config.models import ( - StrixProvider, - configure_sdk_model_defaults, - uses_chat_completions_tool_schema, -) +from strix.config.models import StrixProvider, configure_sdk_model_defaults from strix.core.agents import AgentCoordinator from strix.core.execution import ( respawn_subagents, @@ -36,6 +32,7 @@ build_scope_context, make_model_settings, ) +from strix.core.model_routing import chain_uses_chat_completions_tools, resolve_budget_model from strix.core.paths import run_dir_for, runtime_state_dir from strix.core.sessions import open_agent_session from strix.runtime import session_manager @@ -153,8 +150,10 @@ async def run_strix_scan( raise RuntimeError( "No LLM model configured. Set STRIX_LLM env or pass model= to run_strix_scan().", ) + configured_root_model = resolved_model + resolved_model = resolve_budget_model(configured_root_model, settings.llm) logger.info("LLM model resolved: %s", resolved_model) - chat_completions_tools = uses_chat_completions_tool_schema(resolved_model, settings) + chat_completions_tools = chain_uses_chat_completions_tools(configured_root_model, settings) if coordinator is None: coordinator = AgentCoordinator() @@ -187,10 +186,19 @@ async def run_strix_scan( raise RuntimeError( f"Cannot resume scan {scan_id}: agents.json has no root agent (parent=None)", ) + restored_root_model = coordinator.metadata.get(root_id, {}).get("model") + if isinstance(restored_root_model, str) and restored_root_model.strip(): + configured_root_model = restored_root_model.strip() + resolved_model = resolve_budget_model(configured_root_model, settings.llm) + chat_completions_tools = chain_uses_chat_completions_tools( + configured_root_model, + settings, + ) logger.info( - "Resume: restored coordinator with %d agent(s); root=%s", + "Resume: restored coordinator with %d agent(s); root=%s model=%s", len(coordinator.statuses), root_id, + resolved_model, ) else: root_id = uuid.uuid4().hex[:8] @@ -226,7 +234,11 @@ async def run_strix_scan( sandbox=SandboxRunConfig(client=bundle["client"], session=bundle["session"]), trace_include_sensitive_data=False, ) - hooks = ReportUsageHooks(model=resolved_model, max_budget_usd=max_budget_usd) + hooks = ReportUsageHooks( + model=resolved_model, + max_budget_usd=max_budget_usd, + llm_settings=settings.llm, + ) scope_context = build_scope_context(scan_config) root_context = _merge_root_prompt_context(scope_context, extra_system_prompt_context) @@ -260,6 +272,7 @@ async def run_strix_scan( parent_id=None, task=root_task, skills=skills, + model=resolved_model, ) child_agent_builder = make_child_factory( @@ -282,6 +295,7 @@ async def spawn_child_agent(**kwargs: Any) -> dict[str, Any]: interactive=interactive, event_sink=event_sink, hooks=hooks, + settings=settings, **kwargs, ) @@ -298,7 +312,7 @@ async def spawn_child_agent(**kwargs: Any) -> dict[str, Any]: root_session = open_agent_session(root_id, agents_db) sessions_to_close.append(root_session) - await coordinator.attach_runtime(root_id, session=root_session) + await coordinator.attach_runtime(root_id, agent=root_agent, session=root_session) if is_resume: await respawn_subagents( @@ -313,6 +327,7 @@ async def spawn_child_agent(**kwargs: Any) -> dict[str, Any]: root_id=root_id, event_sink=event_sink, hooks=hooks, + settings=settings, ) initial_input: Any = [] if is_resume else root_task @@ -354,6 +369,7 @@ async def spawn_child_agent(**kwargs: Any) -> dict[str, Any]: start_parked=bool(interactive and is_resume and root_status != "running"), event_sink=event_sink, hooks=hooks, + settings=settings, ) if not interactive and result is not None: final = getattr(result, "final_output", None) diff --git a/strix/interface/tui/app.py b/strix/interface/tui/app.py index 73a7ec0de..a994c64e3 100644 --- a/strix/interface/tui/app.py +++ b/strix/interface/tui/app.py @@ -899,6 +899,8 @@ def watch_show_splash(self, show_splash: bool) -> None: agents_tree.guide_style = "dashed" stats_display = Static("", id="stats_display") + # This is the bottom-right status panel. Usage is grouped by the + # actual model route so fallback transitions remain visible. stats_scroll = VerticalScroll(stats_display, id="stats_scroll") vulnerabilities_panel = VulnerabilitiesPanel(id="vulnerabilities_panel") diff --git a/strix/interface/utils.py b/strix/interface/utils.py index 27c10a468..286f58b67 100644 --- a/strix/interface/utils.py +++ b/strix/interface/utils.py @@ -377,20 +377,36 @@ def build_tui_stats_text(report_state: Any) -> Text: if not report_state: return stats_text - model = load_settings().llm.model or "unknown" - stats_text.append(str(model), style="white") - usage = _llm_usage(report_state) - if usage and _int_stat(usage, "total_tokens") > 0: - stats_text.append("\n") - stats_text.append( - f"{format_token_count(_int_stat(usage, 'total_tokens'))} tokens", - style="white", - ) - cost = _float_stat(usage, "cost") - if cost > 0: - stats_text.append(" · ", style="white") - stats_text.append(f"${cost:.2f}", style="white") + models = usage.get("models") if isinstance(usage, dict) else None + if isinstance(models, list) and models: + stats_text.append("Tokens by model", style="dim") + for model_usage in models: + if not isinstance(model_usage, dict): + continue + stats_text.append("\n") + stats_text.append(str(model_usage.get("model") or "unknown"), style="white") + stats_text.append(" ", style="dim") + stats_text.append( + f"{format_token_count(_int_stat(model_usage, 'total_tokens'))} tokens", + style="white", + ) + cost = _float_stat(model_usage, "cost") + if cost > 0: + stats_text.append(f" · ${cost:.2f}", style="white") + else: + model = load_settings().llm.model or "unknown" + stats_text.append(str(model), style="white") + if usage and _int_stat(usage, "total_tokens") > 0: + stats_text.append("\n") + stats_text.append( + f"{format_token_count(_int_stat(usage, 'total_tokens'))} tokens", + style="white", + ) + cost = _float_stat(usage, "cost") + if cost > 0: + stats_text.append(" · ", style="white") + stats_text.append(f"${cost:.2f}", style="white") caido_url = getattr(report_state, "caido_url", None) if caido_url: diff --git a/strix/report/state.py b/strix/report/state.py index 217b266d5..b4a91df7f 100644 --- a/strix/report/state.py +++ b/strix/report/state.py @@ -316,8 +316,8 @@ def record_sdk_usage( ): self.save_run_data() - def record_observed_llm_cost(self, cost: float) -> None: - self._llm_usage.record_observed_cost(cost) + def record_observed_llm_cost(self, cost: float, *, model: str | None = None) -> None: + self._llm_usage.record_observed_cost(cost, model=model) def get_total_llm_usage(self) -> dict[str, Any]: return dict(self.run_record.get("llm_usage") or self._build_llm_usage_record()) @@ -326,6 +326,10 @@ def get_total_llm_cost(self) -> float: """Live accumulated LLM cost, independent of the persisted run-record snapshot.""" return self._llm_usage.total_cost + def get_model_llm_cost(self, model: str) -> float: + """Live cumulative spend attributed to one exact model route.""" + return self._llm_usage.model_cost(model) + def update_scan_final_fields( self, executive_summary: str, @@ -544,8 +548,18 @@ def litellm_cost_callback( report_state = get_global_report_state() if report_state is None: return + model: str | None = None + if isinstance(kwargs, dict): + raw_model = kwargs.get("model") + if isinstance(raw_model, str) and raw_model.strip(): + model = raw_model.strip() + if model is None: + raw_model = getattr(completion_response, "model", None) + if isinstance(raw_model, str) and raw_model.strip(): + model = raw_model.strip() + try: - report_state.record_observed_llm_cost(cost) + report_state.record_observed_llm_cost(cost, model=model) except Exception: logger.exception("Failed to record observed LiteLLM cost") diff --git a/strix/report/usage.py b/strix/report/usage.py index 7b223a6e1..87309da6e 100644 --- a/strix/report/usage.py +++ b/strix/report/usage.py @@ -19,6 +19,10 @@ def __init__(self) -> None: self._agent_usage: dict[str, Usage] = {} self._agent_metadata: dict[str, dict[str, str]] = {} self._agent_cost: dict[str, float] = {} + self._model_usage: dict[str, Usage] = {} + self._model_estimated_cost: dict[str, float] = {} + self._model_observed_cost: dict[str, float] = {} + self._model_base_cost: dict[str, float] = {} self._total_cost = 0.0 def record( @@ -35,6 +39,8 @@ def record( normalized_agent_id = str(agent_id or "unknown") self._total_usage.add(usage) self._agent_usage.setdefault(normalized_agent_id, Usage()).add(usage) + normalized_model = str(model or "unknown").strip() or "unknown" + self._model_usage.setdefault(normalized_model, Usage()).add(usage) metadata = self._agent_metadata.setdefault(normalized_agent_id, {}) if agent_name: @@ -42,9 +48,12 @@ def record( if model: metadata["model"] = model - if not _is_litellm_routed(model): - estimated = _estimate_litellm_cost(usage, model) - if estimated: + estimated = _estimate_litellm_cost(usage, model) + if estimated: + self._model_estimated_cost[normalized_model] = ( + self._model_estimated_cost.get(normalized_model, 0.0) + estimated + ) + if not _is_litellm_routed(model): self._total_cost += estimated self._agent_cost[normalized_agent_id] = ( self._agent_cost.get(normalized_agent_id, 0.0) + estimated @@ -52,18 +61,44 @@ def record( return True - def record_observed_cost(self, cost: float) -> None: + def record_observed_cost(self, cost: float, *, model: str | None = None) -> None: if isinstance(cost, int | float) and cost > 0: self._total_cost += float(cost) + normalized_model = str(model or "").strip() + if normalized_model: + self._model_observed_cost[normalized_model] = ( + self._model_observed_cost.get(normalized_model, 0.0) + float(cost) + ) @property def total_cost(self) -> float: return _round_cost(self._total_cost) + def model_cost(self, model: str) -> float: + """Return best available cumulative spend for one exact model route.""" + normalized = model.strip().lower() + estimated = sum( + cost for name, cost in self._model_estimated_cost.items() if name.lower() == normalized + ) + observed = sum( + cost for name, cost in self._model_observed_cost.items() if name.lower() == normalized + ) + # SDK token pricing and provider callbacks describe the same calls. Use + # the larger cumulative figure rather than double-counting them. + base = sum( + cost for name, cost in self._model_base_cost.items() if name.lower() == normalized + ) + return _round_cost(base + max(estimated, observed)) + def to_record(self) -> dict[str, Any]: record = serialize_usage(self._total_usage) record["cost"] = _round_cost(self._total_cost) record["agents"] = [] + record["models"] = [] + for model in sorted(self._model_usage, key=str.lower): + model_record = serialize_usage(self._model_usage[model]) + model_record.update({"model": model, "cost": self.model_cost(model)}) + record["models"].append(model_record) agent_tokens = {aid: _resolve_total_tokens(u) for aid, u in self._agent_usage.items()} # Native provider estimates are tied to an agent. Any remainder comes @@ -99,11 +134,15 @@ def to_record(self) -> dict[str, Any]: return record - def hydrate(self, raw_usage: Any) -> None: + def hydrate(self, raw_usage: Any) -> None: # noqa: PLR0912, PLR0915 self._total_usage = Usage() self._agent_usage.clear() self._agent_metadata.clear() self._agent_cost.clear() + self._model_usage.clear() + self._model_estimated_cost.clear() + self._model_observed_cost.clear() + self._model_base_cost.clear() self._total_cost = 0.0 if not isinstance(raw_usage, dict): @@ -117,6 +156,23 @@ def hydrate(self, raw_usage: Any) -> None: self._total_cost = _float_or_zero(raw_usage.get("cost")) + raw_models_value = raw_usage.get("models") or [] + raw_models = raw_models_value if isinstance(raw_models_value, list) else [] + persisted_model_cost: dict[str, float] = {} + if raw_models: + for raw_model in raw_models: + if not isinstance(raw_model, dict): + continue + model = str(raw_model.get("model") or "").strip() + if not model: + continue + try: + self._model_usage[model] = deserialize_usage(raw_model) + except Exception: + logger.exception("Failed to hydrate llm_usage for model %s", model) + self._model_usage[model] = Usage() + persisted_model_cost[model] = _float_or_zero(raw_model.get("cost")) + for raw_agent in raw_usage.get("agents") or []: if not isinstance(raw_agent, dict): continue @@ -131,11 +187,11 @@ def hydrate(self, raw_usage: Any) -> None: metadata: dict[str, str] = {} agent_name = raw_agent.get("agent_name") - model = raw_agent.get("model") + agent_model = raw_agent.get("model") if isinstance(agent_name, str) and agent_name: metadata["agent_name"] = agent_name - if isinstance(model, str) and model: - metadata["model"] = model + if isinstance(agent_model, str) and agent_model: + metadata["model"] = agent_model self._agent_metadata[agent_id] = metadata # Persisted per-agent costs may include proportional allocation of # provider callback costs. Keep only native estimates here so a @@ -143,6 +199,24 @@ def hydrate(self, raw_usage: Any) -> None: if not _is_litellm_routed(metadata.get("model")): self._agent_cost[agent_id] = _float_or_zero(raw_agent.get("cost")) + # Backward compatibility for run records written before per-model + # usage was persisted. This cannot split agents that switched + # models, but old records never supported model switching. + if not raw_models and metadata.get("model"): + model_name = metadata["model"] + self._model_usage.setdefault(model_name, Usage()).add( + self._agent_usage[agent_id] + ) + persisted_model_cost[model_name] = ( + persisted_model_cost.get(model_name, 0.0) + + _float_or_zero(raw_agent.get("cost")) + ) + + # Hydrated values are an immutable historical baseline. New SDK + # estimates and provider callbacks are accumulated separately so resume + # neither loses nor double-counts the persisted model spend. + self._model_base_cost.update(persisted_model_cost) + def _resolve_total_tokens(usage: Usage) -> int: total = max(0, int(usage.total_tokens or 0)) diff --git a/tests/test_budget_fallbacks.py b/tests/test_budget_fallbacks.py new file mode 100644 index 000000000..4a82e9df1 --- /dev/null +++ b/tests/test_budget_fallbacks.py @@ -0,0 +1,254 @@ +"""Per-model budget fallback and TUI usage aggregation tests.""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +from agents.model_settings import ModelSettings +from agents.usage import Usage +from pydantic import ValidationError + +from strix.config.settings import LlmSettings +from strix.core.agents import AgentCoordinator +from strix.core.hooks import ModelFallbackError, ReportUsageHooks +from strix.core.model_routing import chain_uses_chat_completions_tools, resolve_budget_model +from strix.interface.utils import build_tui_stats_text +from strix.report.usage import LLMUsageLedger + + +def test_fallback_chain_and_case_insensitive_lookup() -> None: + llm = LlmSettings( + model_budgets_usd={"openai/GPT": 5, "z-ai/GLM": 3}, + model_fallbacks={ + "openai/GPT": "z-ai/GLM", + "z-ai/GLM": "deepseek/final", + }, + ) + + assert llm.budget_for("OPENAI/gpt") == 5 + assert llm.fallback_for("Z-AI/glm") == "deepseek/final" + assert llm.model_chain("openai/GPT") == [ + "openai/GPT", + "z-ai/GLM", + "deepseek/final", + ] + + +def test_fallback_source_requires_budget_and_cycles_are_rejected() -> None: + with pytest.raises(ValidationError, match="requires a model_budgets_usd entry"): + LlmSettings(model_fallbacks={"a": "b"}) + with pytest.raises(ValidationError, match="cycle"): + LlmSettings( + model_budgets_usd={"a": 1, "b": 1}, + model_fallbacks={"a": "b", "b": "a"}, + ) + + +def test_new_agent_skips_an_already_exhausted_model() -> None: + llm = LlmSettings( + model_budgets_usd={"gpt": 1, "glm": 2}, + model_fallbacks={"gpt": "glm", "glm": "final"}, + ) + state = MagicMock() + state.get_model_llm_cost.side_effect = lambda model: {"gpt": 1.5, "glm": 1.0}[model] + + with patch("strix.core.model_routing.get_global_report_state", return_value=state): + assert resolve_budget_model("gpt", llm) == "glm" + + +def test_tool_schema_is_compatible_with_entire_chain(monkeypatch: pytest.MonkeyPatch) -> None: + settings = SimpleNamespace( + llm=LlmSettings( + model_budgets_usd={"openai/gpt": 1}, + model_fallbacks={"openai/gpt": "z-ai/glm"}, + ) + ) + monkeypatch.setattr( + "strix.core.model_routing.uses_chat_completions_tool_schema", + lambda model, _settings: model == "z-ai/glm", + ) + + assert chain_uses_chat_completions_tools("openai/gpt", settings) + + +def test_usage_ledger_groups_tokens_by_actual_model() -> None: + ledger = LLMUsageLedger() + ledger.record( + agent_id="root", + model="openai/gpt", + usage=Usage(input_tokens=100, output_tokens=20, total_tokens=120), + ) + ledger.record( + agent_id="child", + model="z-ai/glm", + usage=Usage(input_tokens=50, output_tokens=10, total_tokens=60), + ) + + models = {item["model"]: item for item in ledger.to_record()["models"]} + assert models["openai/gpt"]["total_tokens"] == 120 + assert models["z-ai/glm"]["total_tokens"] == 60 + + +def test_usage_hydration_preserves_model_cost_without_double_counting() -> None: + ledger = LLMUsageLedger() + ledger.hydrate( + { + "total_tokens": 10, + "cost": 1.5, + "models": [{"model": "gpt", "total_tokens": 10, "cost": 1.5}], + "agents": [], + } + ) + + assert ledger.model_cost("gpt") == 1.5 + assert ledger.to_record()["models"][0]["cost"] == 1.5 + + +def test_tui_stats_lists_tokens_for_every_model() -> None: + state = MagicMock() + state.caido_url = None + state.get_total_llm_usage.return_value = { + "models": [ + {"model": "openai/gpt", "total_tokens": 1500, "cost": 1.25}, + {"model": "z-ai/glm", "total_tokens": 700, "cost": 0.2}, + ] + } + + rendered = build_tui_stats_text(state).plain + assert "openai/gpt 1.5K tokens · $1.25" in rendered + assert "z-ai/glm 700 tokens · $0.20" in rendered + + +@pytest.mark.asyncio +async def test_hook_switches_all_matching_agents_at_response_boundary() -> None: + llm = LlmSettings( + reasoning_effort="none", + model_budgets_usd={"openai/gpt": 1}, + model_fallbacks={"openai/gpt": "z-ai/glm"}, + ) + hooks = ReportUsageHooks(model="openai/gpt", llm_settings=llm) + coordinator = AgentCoordinator() + root = SimpleNamespace(name="strix", model="openai/gpt", model_settings=ModelSettings()) + child = SimpleNamespace(name="xss", model="openai/gpt", model_settings=ModelSettings()) + await coordinator.register("root", "strix", None, model="openai/gpt") + await coordinator.register("child", "xss", "root", model="openai/gpt") + await coordinator.attach_runtime("root", agent=root) + await coordinator.attach_runtime("child", agent=child) + + state = MagicMock() + state.get_total_llm_cost.return_value = 1.1 + state.get_model_llm_cost.return_value = 1.1 + context = SimpleNamespace( + context={"agent_id": "root", "parent_id": None, "coordinator": coordinator} + ) + response = SimpleNamespace(usage=Usage(input_tokens=10, output_tokens=5, total_tokens=15)) + + with ( + patch("strix.core.hooks.get_global_report_state", return_value=state), + patch("strix.core.hooks.make_model_settings", return_value=ModelSettings()), + pytest.raises(ModelFallbackError) as exc_info, + ): + await hooks.on_llm_end(context, root, response) + + assert exc_info.value.next_model == "z-ai/glm" + assert root.model == "z-ai/glm" + assert child.model == "z-ai/glm" + assert coordinator.metadata["root"]["model"] == "z-ai/glm" + assert coordinator.metadata["child"]["model"] == "z-ai/glm" + + +@pytest.mark.asyncio +async def test_fallback_carries_subagent_credential_per_call() -> None: + # A cross-provider fallback must re-resolve credentials so it doesn't run on + # the exhausted model's key (env vars are global per provider). + llm = LlmSettings( + reasoning_effort="none", + api_key="root-key", + subagent_api_key="fallback-key", + model_budgets_usd={"openai/gpt": 1}, + model_fallbacks={"openai/gpt": "anthropic/claude"}, + ) + hooks = ReportUsageHooks(model="openai/gpt", llm_settings=llm) + coordinator = AgentCoordinator() + agent = SimpleNamespace(name="strix", model="openai/gpt", model_settings=ModelSettings()) + await coordinator.register("root", "strix", None, model="openai/gpt") + await coordinator.attach_runtime("root", agent=agent) + + state = MagicMock() + state.get_total_llm_cost.return_value = 1.1 + state.get_model_llm_cost.return_value = 1.1 + context = SimpleNamespace( + context={"agent_id": "root", "parent_id": None, "coordinator": coordinator} + ) + response = SimpleNamespace(usage=Usage(input_tokens=10, output_tokens=5, total_tokens=15)) + + with ( + patch("strix.core.hooks.get_global_report_state", return_value=state), + pytest.raises(ModelFallbackError), + ): + await hooks.on_llm_end(context, agent, response) + + assert agent.model == "anthropic/claude" + assert agent.model_settings.extra_args["api_key"] == "fallback-key" + + +@pytest.mark.asyncio +async def test_stale_inflight_response_retries_existing_fallback_without_advancing() -> None: + llm = LlmSettings( + reasoning_effort="none", + model_budgets_usd={"gpt": 1, "glm": 2}, + model_fallbacks={"gpt": "glm", "glm": "final"}, + ) + hooks = ReportUsageHooks(model="gpt", llm_settings=llm) + coordinator = AgentCoordinator() + agent = SimpleNamespace(name="child", model="gpt", model_settings=ModelSettings()) + await coordinator.register("child", "child", "root", model="gpt") + await coordinator.attach_runtime("child", agent=agent) + context = SimpleNamespace( + context={"agent_id": "child", "parent_id": "root", "coordinator": coordinator} + ) + await hooks.on_llm_start(context, agent, None, []) + agent.model = "glm" # A concurrent gpt response already switched the shared route. + state = MagicMock() + state.get_model_llm_cost.return_value = 1.2 + + with ( + patch("strix.core.hooks.get_global_report_state", return_value=state), + pytest.raises(ModelFallbackError) as exc_info, + ): + await hooks.on_llm_end(context, agent, SimpleNamespace(usage=Usage(total_tokens=1))) + + assert exc_info.value.next_model == "glm" + assert agent.model == "glm" + + +@pytest.mark.asyncio +async def test_hook_can_continue_to_second_fallback_layer() -> None: + llm = LlmSettings( + reasoning_effort="none", + model_budgets_usd={"gpt": 1, "glm": 2}, + model_fallbacks={"gpt": "glm", "glm": "final"}, + ) + hooks = ReportUsageHooks(model="gpt", llm_settings=llm) + coordinator = AgentCoordinator() + agent = SimpleNamespace(name="strix", model="glm", model_settings=ModelSettings()) + await coordinator.register("root", "strix", None, model="glm") + await coordinator.attach_runtime("root", agent=agent) + state = MagicMock() + state.get_total_llm_cost.return_value = 3.0 + state.get_model_llm_cost.return_value = 2.1 + context = SimpleNamespace( + context={"agent_id": "root", "parent_id": None, "coordinator": coordinator} + ) + + with ( + patch("strix.core.hooks.get_global_report_state", return_value=state), + patch("strix.core.hooks.make_model_settings", return_value=ModelSettings()), + pytest.raises(ModelFallbackError) as exc_info, + ): + await hooks.on_llm_end(context, agent, SimpleNamespace(usage=Usage(total_tokens=1))) + + assert exc_info.value.next_model == "final" + assert agent.model == "final" diff --git a/tests/test_config_loader.py b/tests/test_config_loader.py index be0bfb8f7..ca0f9c14e 100644 --- a/tests/test_config_loader.py +++ b/tests/test_config_loader.py @@ -26,6 +26,8 @@ "LITELLM_BASE_URL", "OLLAMA_API_BASE", "STRIX_REASONING_EFFORT", + "STRIX_MODEL_BUDGETS_USD", + "STRIX_MODEL_FALLBACKS", "STRIX_FORCE_REQUIRED_TOOL_CHOICE", "LLM_TIMEOUT", "PERPLEXITY_API_KEY", @@ -168,6 +170,28 @@ def test_apply_override_and_load_settings_round_trip(tmp_path: Path) -> None: assert loader.load_settings() is settings +def test_structured_model_budgets_and_fallbacks_load(tmp_path: Path) -> None: + path = tmp_path / "fallbacks.json" + path.write_text( + json.dumps( + { + "llm": { + "model": "gpt", + "model_budgets_usd": {"gpt": 5, "glm": 2}, + "model_fallbacks": {"gpt": "glm", "glm": "final"}, + } + } + ), + encoding="utf-8", + ) + + loader.apply_config_override(path) + settings = loader.load_settings() + + assert settings.llm.model_chain("gpt") == ["gpt", "glm", "final"] + assert settings.llm.budget_for("glm") == 2 + + def test_apply_config_override_invalidates_cache(tmp_path: Path) -> None: first = tmp_path / "first.json" first.write_text(json.dumps({"env": {"STRIX_LLM": "first-model"}}), encoding="utf-8") diff --git a/tests/test_cost_tracking.py b/tests/test_cost_tracking.py index 543d3fd5c..7734635b4 100644 --- a/tests/test_cost_tracking.py +++ b/tests/test_cost_tracking.py @@ -32,7 +32,7 @@ def test_cost_callback_reads_openrouter_stream_usage_cost() -> None: with patch("strix.report.state.get_global_report_state", return_value=report_state): litellm_cost_callback({"response_cost": None}, response) - report_state.record_observed_llm_cost.assert_called_once_with(1.2345) + report_state.record_observed_llm_cost.assert_called_once_with(1.2345, model=None) def test_cost_callback_reads_usage_cost_from_mapping_response() -> None: @@ -42,7 +42,7 @@ def test_cost_callback_reads_usage_cost_from_mapping_response() -> None: with patch("strix.report.state.get_global_report_state", return_value=report_state): litellm_cost_callback({}, response) - report_state.record_observed_llm_cost.assert_called_once_with(0.125) + report_state.record_observed_llm_cost.assert_called_once_with(0.125, model=None) def test_cost_callback_reads_byok_upstream_inference_cost() -> None: @@ -59,7 +59,7 @@ def test_cost_callback_reads_byok_upstream_inference_cost() -> None: with patch("strix.report.state.get_global_report_state", return_value=report_state): litellm_cost_callback({"response_cost": None}, response) - report_state.record_observed_llm_cost.assert_called_once_with(6.75e-06) + report_state.record_observed_llm_cost.assert_called_once_with(6.75e-06, model=None) def test_cost_callback_sums_usage_cost_and_upstream_inference_cost() -> None: @@ -75,7 +75,9 @@ def test_cost_callback_sums_usage_cost_and_upstream_inference_cost() -> None: with patch("strix.report.state.get_global_report_state", return_value=report_state): litellm_cost_callback({}, response) - report_state.record_observed_llm_cost.assert_called_once_with(pytest.approx(0.21)) + report_state.record_observed_llm_cost.assert_called_once_with( + pytest.approx(0.21), model=None + ) def test_cost_callback_ignores_upstream_cost_for_non_byok_responses() -> None: @@ -91,7 +93,7 @@ def test_cost_callback_ignores_upstream_cost_for_non_byok_responses() -> None: with patch("strix.report.state.get_global_report_state", return_value=report_state): litellm_cost_callback({}, response) - report_state.record_observed_llm_cost.assert_called_once_with(0.05) + report_state.record_observed_llm_cost.assert_called_once_with(0.05, model=None) def test_cost_callback_estimates_cost_with_provider_prefixed_model() -> None: @@ -114,7 +116,9 @@ def fake_completion_cost(**kwargs: object) -> float: ): litellm_cost_callback(kwargs, response) - report_state.record_observed_llm_cost.assert_called_once_with(0.5) + report_state.record_observed_llm_cost.assert_called_once_with( + 0.5, model="anthropic/claude-sonnet-4.5" + ) def test_cost_callback_estimates_cost_with_bare_model_fallback() -> None: @@ -137,7 +141,9 @@ def fake_completion_cost(**kwargs: object) -> float: ): litellm_cost_callback(kwargs, response) - report_state.record_observed_llm_cost.assert_called_once_with(0.025) + report_state.record_observed_llm_cost.assert_called_once_with( + 0.025, model="openai/gpt-4o-mini" + ) def test_cost_callback_records_nothing_when_no_cost_available() -> None: diff --git a/tests/test_model_routing.py b/tests/test_model_routing.py index f26fd873b..2c56f0c41 100644 --- a/tests/test_model_routing.py +++ b/tests/test_model_routing.py @@ -9,9 +9,10 @@ from agents.model_settings import ModelSettings from agents.usage import Usage -from strix.agents.factory import _resolve_child_api_key, build_strix_agent, make_child_factory +from strix.agents.factory import build_strix_agent, make_child_factory from strix.config.settings import LlmSettings, Settings, SkillModelRoute from strix.core.hooks import ReportUsageHooks +from strix.core.model_routing import resolve_route_api_key def test_child_api_key_prefers_explicit_subagent_key() -> None: @@ -21,14 +22,14 @@ def test_child_api_key_prefers_explicit_subagent_key() -> None: subagent_api_key="child-key", ) # Explicit subagent key wins even when the child shares the root provider. - assert _resolve_child_api_key(llm, "openai/child", "openai/root") == "child-key" + assert resolve_route_api_key(llm, "openai/child", "openai/root") == "child-key" def test_child_api_key_reuses_root_key_only_for_same_provider() -> None: llm = LlmSettings(model="openai/root", api_key="root-key") - assert _resolve_child_api_key(llm, "openai/mini", "openai/root") == "root-key" + assert resolve_route_api_key(llm, "openai/mini", "openai/root") == "root-key" # A cross-provider child must not receive the root provider's credential. - assert _resolve_child_api_key(llm, "anthropic/opus", "openai/root") is None + assert resolve_route_api_key(llm, "anthropic/opus", "openai/root") is None def test_child_api_key_flows_into_model_settings_extra_args() -> None: @@ -83,7 +84,7 @@ def test_child_factory_precedence_and_per_model_schema(monkeypatch: pytest.Monke ) seen_schema_models: list[str] = [] monkeypatch.setattr( - "strix.agents.factory.uses_chat_completions_tool_schema", + "strix.core.model_routing.uses_chat_completions_tool_schema", lambda model, _settings: seen_schema_models.append(model) or model != "openai/escalated", ) monkeypatch.setattr( diff --git a/tests/test_runner_rate_limit.py b/tests/test_runner_rate_limit.py index 62242339c..938da1831 100644 --- a/tests/test_runner_rate_limit.py +++ b/tests/test_runner_rate_limit.py @@ -44,7 +44,7 @@ async def test_persistent_rate_limit_stops_gracefully( monkeypatch.setattr(runner, "load_settings", lambda: settings) monkeypatch.setattr(runner, "configure_sdk_model_defaults", lambda _settings: None) monkeypatch.setattr( - runner, "uses_chat_completions_tool_schema", lambda _model, _settings: False + runner, "chain_uses_chat_completions_tools", lambda _model, _settings: False ) monkeypatch.setattr(todo_tools, "hydrate_todos_from_disk", lambda _state_dir: None) diff --git a/tests/test_runner_root_prompt.py b/tests/test_runner_root_prompt.py index 670177d04..dbede0e8c 100644 --- a/tests/test_runner_root_prompt.py +++ b/tests/test_runner_root_prompt.py @@ -53,7 +53,7 @@ def _patch_engine_scaffold( monkeypatch.setattr(runner, "configure_sdk_model_defaults", lambda _settings: None) monkeypatch.setattr( runner, - "uses_chat_completions_tool_schema", + "chain_uses_chat_completions_tools", lambda _model, _settings: False, )