mirror of
https://github.com/usestrix/strix.git
synced 2026-11-09 13:31:41 -10:00
Compare commits
7
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
8faf47de96 | ||
|
|
fb60e104f6 | ||
|
|
8267d7b311 | ||
|
|
8685c53b55 | ||
|
|
b4113e662f | ||
|
|
93dae3db07 | ||
|
|
1e1339e494 |
No files matched your search
+106
-46
@@ -58,7 +58,7 @@ if TYPE_CHECKING:
|
||||
from openai import AsyncOpenAI
|
||||
from openai.types.responses.response_prompt_param import ResponsePromptParam
|
||||
|
||||
from strix.config.settings import ApiType, LlmSettings, ReasoningEffort, Settings
|
||||
from strix.config.settings import LlmSettings, ReasoningEffort, Settings
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -632,6 +632,61 @@ DEFAULT_MODEL_RETRY = ModelRetrySettings(
|
||||
),
|
||||
)
|
||||
|
||||
RECOMMENDED_MODEL_NAMES = (
|
||||
"zai/glm-5.3",
|
||||
"zai/glm-5.3-flash",
|
||||
"openai/gpt-5.6-sol",
|
||||
"openai/gpt-5.6-terra",
|
||||
"openai/gpt-5.6-luna",
|
||||
"openai/gpt-5.6",
|
||||
"openai/gpt-5.5-pro",
|
||||
"openai/gpt-5.5",
|
||||
"openai/gpt-5.4",
|
||||
"openai/gpt-5.3-codex",
|
||||
"anthropic/claude-fable-5-1",
|
||||
"anthropic/claude-fable-5",
|
||||
"anthropic/claude-opus-5",
|
||||
"anthropic/claude-opus-4-8",
|
||||
"anthropic/claude-sonnet-5",
|
||||
"anthropic/claude-sonnet-4-6",
|
||||
"vertex_ai/gemini-3.1-pro-preview",
|
||||
"gemini/gemini-3.1-pro-preview",
|
||||
"vertex_ai/gemini-3.7-flash",
|
||||
"gemini/gemini-3.7-flash",
|
||||
"gemini/gemini-3.6-flash",
|
||||
"deepseek/deepseek-v4-pro",
|
||||
"deepseek/deepseek-v4-flash",
|
||||
"dashscope/qwen3.8-max",
|
||||
"dashscope/qwen3.7-max-2026-06-08",
|
||||
"moonshot/kimi-k3",
|
||||
"moonshot/kimi-k2.7-code",
|
||||
)
|
||||
|
||||
_RECOMMENDED_MODEL_NAME_SET = frozenset(name.lower() for name in RECOMMENDED_MODEL_NAMES)
|
||||
|
||||
# Matched against the bare model name only: the route (``openai/``, ``openrouter/``,
|
||||
# a local gateway, ...) says nothing about the model's quality.
|
||||
FRONTIER_MODEL_PREFIXES = (
|
||||
"gpt-5",
|
||||
"claude-fable-5",
|
||||
"claude-opus-5",
|
||||
"claude-opus-4",
|
||||
"claude-sonnet-5",
|
||||
"claude-sonnet-4",
|
||||
"gemini-3",
|
||||
"deepseek-v4",
|
||||
"deepseek-r1",
|
||||
"deepseek-reasoner",
|
||||
"qwen3.8",
|
||||
"qwen3.7",
|
||||
"qwen3-max",
|
||||
"kimi-k3",
|
||||
"kimi-k2.7",
|
||||
"kimi-k2.6",
|
||||
"glm-5.3",
|
||||
"glm-5.2",
|
||||
)
|
||||
|
||||
|
||||
def configure_sdk_model_defaults(settings: Settings) -> None:
|
||||
"""Apply Strix config to SDK-native defaults."""
|
||||
@@ -649,38 +704,12 @@ def configure_sdk_model_defaults(settings: Settings) -> None:
|
||||
if llm.api_base:
|
||||
os.environ["OPENAI_BASE_URL"] = llm.api_base
|
||||
_configure_litellm_default("api_base", llm.api_base)
|
||||
configure_sdk_api_route(llm.model or "", settings)
|
||||
_configure_extra_headers(llm)
|
||||
api_type = llm.api_type
|
||||
if api_type is None:
|
||||
api_type = "chat_completions" if llm.api_base else "responses"
|
||||
|
||||
|
||||
def configure_sdk_api_route(model_name: str, settings: Settings) -> None:
|
||||
"""Point SDK-native OpenAI requests for ``model_name`` at Responses or chat completions."""
|
||||
api_type = resolve_api_type(model_name, settings)
|
||||
logger.info("OpenAI API route for %s: %s", model_name, api_type)
|
||||
set_default_openai_api(api_type)
|
||||
|
||||
|
||||
_RESPONSES_ENDPOINT = "/v1/responses"
|
||||
|
||||
|
||||
def resolve_api_type(model_name: str, settings: Settings) -> ApiType:
|
||||
"""The SDK-native OpenAI route for ``model_name``: Responses or chat completions.
|
||||
|
||||
An explicit ``STRIX_API_TYPE`` wins. Otherwise the model decides: Responses
|
||||
when LiteLLM's catalog lists ``/v1/responses`` for it, chat completions for
|
||||
everything else.
|
||||
"""
|
||||
if settings.llm.api_type is not None:
|
||||
return settings.llm.api_type
|
||||
if _RESPONSES_ENDPOINT in _catalog_supported_endpoints(model_name):
|
||||
return "responses"
|
||||
return "chat_completions"
|
||||
|
||||
|
||||
def _catalog_supported_endpoints(model_name: str) -> list[str]:
|
||||
entry = _catalog_entry(model_name)
|
||||
endpoints = entry.get("supported_endpoints") if entry else None
|
||||
return [str(e) for e in endpoints] if isinstance(endpoints, list) else []
|
||||
_configure_extra_headers(llm)
|
||||
|
||||
|
||||
def _mirror_api_key_to_provider_env(model_name: str | None, api_key: str) -> None:
|
||||
@@ -885,7 +914,11 @@ def uses_chat_completions_tool_schema(model_name: str, settings: Settings) -> bo
|
||||
model = model_name.strip().lower()
|
||||
if "/" in model and not model.startswith("openai/"):
|
||||
return True
|
||||
return resolve_api_type(model_name, settings) == "chat_completions"
|
||||
if settings.llm.api_type is not None:
|
||||
return settings.llm.api_type == "chat_completions"
|
||||
if settings.llm.api_base:
|
||||
return True
|
||||
return not model_supports_reasoning(model_name)
|
||||
|
||||
|
||||
def supports_strict_tool_schemas(model_name: str) -> bool:
|
||||
@@ -901,27 +934,54 @@ def supports_strict_tool_schemas(model_name: str) -> bool:
|
||||
|
||||
|
||||
def model_supports_reasoning(model_name: str) -> bool:
|
||||
entry = _catalog_entry(model_name)
|
||||
return bool(entry and entry.get("supports_reasoning"))
|
||||
import litellm
|
||||
|
||||
|
||||
def _bare_openai_name(model_name: str) -> str:
|
||||
name = model_name.strip().lower()
|
||||
for prefix in ("litellm/", "any-llm/", "openai/"):
|
||||
if name.startswith(prefix):
|
||||
return name[len(prefix) :]
|
||||
return name
|
||||
|
||||
|
||||
def _catalog_entry(model_name: str) -> dict[str, Any] | None:
|
||||
"""LiteLLM's cost-map entry for the model, looked up as it would route it."""
|
||||
import litellm
|
||||
|
||||
name = _bare_openai_name(model_name)
|
||||
name = name[len(prefix) :]
|
||||
break
|
||||
entry = litellm.model_cost.get(name)
|
||||
if entry is None and "/" in name:
|
||||
entry = litellm.model_cost.get(name.rsplit("/", 1)[1])
|
||||
return entry if isinstance(entry, dict) else None
|
||||
return bool(entry and entry.get("supports_reasoning"))
|
||||
|
||||
|
||||
def is_recommended_or_frontier_model(model_name: str) -> bool:
|
||||
"""Return whether a model is recommended or in a frontier model family."""
|
||||
name = _normalized_model_name(model_name)
|
||||
if not name:
|
||||
return False
|
||||
if name in _RECOMMENDED_MODEL_NAME_SET:
|
||||
return True
|
||||
bare_model_name = name.rsplit("/", 1)[-1]
|
||||
return _matches_model_prefix(bare_model_name, FRONTIER_MODEL_PREFIXES)
|
||||
|
||||
|
||||
def _normalized_model_name(model_name: str) -> str:
|
||||
name = model_name.strip().lower()
|
||||
for prefix in ("litellm/", "any-llm/"):
|
||||
if name.startswith(prefix):
|
||||
name = name[len(prefix) :]
|
||||
break
|
||||
return name
|
||||
|
||||
|
||||
def _matches_model_prefix(model_name: str, model_prefixes: tuple[str, ...]) -> bool:
|
||||
return any(
|
||||
candidate.startswith(prefix)
|
||||
for candidate in _model_name_candidates(model_name)
|
||||
for prefix in model_prefixes
|
||||
)
|
||||
|
||||
|
||||
def _model_name_candidates(model_name: str) -> tuple[str, ...]:
|
||||
if "." not in model_name:
|
||||
return (model_name,)
|
||||
suffixes = tuple(
|
||||
model_name.split(".", index)[-1] for index in range(1, model_name.count(".") + 1)
|
||||
)
|
||||
return (model_name, *suffixes)
|
||||
|
||||
|
||||
def is_known_openai_bare_model(model_name: str) -> bool:
|
||||
|
||||
@@ -4,7 +4,7 @@ from __future__ import annotations
|
||||
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import AliasChoices, Field, field_validator
|
||||
from pydantic import AliasChoices, Field
|
||||
from pydantic_settings import BaseSettings, SettingsConfigDict
|
||||
|
||||
|
||||
@@ -13,12 +13,6 @@ ApiType = Literal["responses", "chat_completions"]
|
||||
|
||||
DEFAULT_MAX_TURNS = 500
|
||||
|
||||
|
||||
def _lowercase(value: object) -> object:
|
||||
"""Enum-like env values are matched case-insensitively (``None`` is ``none``)."""
|
||||
return value.strip().lower() if isinstance(value, str) else value
|
||||
|
||||
|
||||
_BASE_CONFIG = SettingsConfigDict(
|
||||
case_sensitive=False,
|
||||
populate_by_name=True,
|
||||
@@ -84,11 +78,6 @@ class LlmSettings(BaseSettings):
|
||||
alias="LLM_MAX_TOOL_CALLS_PER_TURN",
|
||||
)
|
||||
|
||||
@field_validator("api_type", "reasoning_effort", mode="before")
|
||||
@classmethod
|
||||
def _normalize_case(cls, value: object) -> object:
|
||||
return _lowercase(value)
|
||||
|
||||
|
||||
class DedupeSettings(BaseSettings):
|
||||
model_config = _BASE_CONFIG
|
||||
@@ -106,11 +95,6 @@ class DedupeSettings(BaseSettings):
|
||||
repr=False,
|
||||
)
|
||||
|
||||
@field_validator("reasoning_effort", mode="before")
|
||||
@classmethod
|
||||
def _normalize_case(cls, value: object) -> object:
|
||||
return _lowercase(value)
|
||||
|
||||
|
||||
class ContextSettings(BaseSettings):
|
||||
"""Context-window management: per-tool-output caps and history compaction."""
|
||||
|
||||
+13
-112
@@ -24,14 +24,8 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
Status = Literal["running", "waiting", "completed", "stopped", "crashed", "failed", "budget_paused"]
|
||||
|
||||
BudgetPolicy = Literal["stop", "pause"]
|
||||
|
||||
TERMINAL_STATUSES: frozenset[str] = frozenset({"completed", "stopped", "crashed", "failed"})
|
||||
|
||||
# Agents that still have work to do or to wake up for. A parked agent is mid-task:
|
||||
# a scan is not finished while one exists, and it can be stopped like any other.
|
||||
ACTIVE_STATUSES: frozenset[str] = frozenset({"running", "waiting", "budget_paused"})
|
||||
|
||||
# Why an agent parked. The user can message any agent, so this - not the agent's
|
||||
# position in the tree - decides whether waiting is bounded: only an agent waiting
|
||||
# on other agents is re-checked on a timer.
|
||||
@@ -74,10 +68,7 @@ class AgentCoordinator:
|
||||
self._budget_stopped = False
|
||||
self._reserve_stopped = False
|
||||
self._budget_paused = False
|
||||
self._resume_epoch = 0
|
||||
self._budget_policy: BudgetPolicy = "stop"
|
||||
self._extend_budget: Callable[[], None] | None = None
|
||||
self._set_budget_limit: Callable[[float | None], None] | None = None
|
||||
|
||||
def set_snapshot_path(self, path: Path) -> None:
|
||||
self._snapshot_path = path
|
||||
@@ -104,100 +95,15 @@ class AgentCoordinator:
|
||||
def budget_paused(self) -> bool:
|
||||
return self._budget_paused
|
||||
|
||||
@property
|
||||
def resume_epoch(self) -> int:
|
||||
"""Bumped by every ``resume_budget``; an agent parks against the value it read."""
|
||||
return self._resume_epoch
|
||||
|
||||
@property
|
||||
def budget_policy(self) -> BudgetPolicy:
|
||||
return self._budget_policy
|
||||
|
||||
def set_budget_policy(self, policy: BudgetPolicy) -> None:
|
||||
self._budget_policy = policy
|
||||
|
||||
def set_budget_extender(self, extend: Callable[[], None]) -> None:
|
||||
self._extend_budget = extend
|
||||
|
||||
def set_budget_limit_setter(self, setter: Callable[[float | None], None]) -> None:
|
||||
self._set_budget_limit = setter
|
||||
|
||||
async def pause_for_budget(self, agent_id: str) -> None:
|
||||
async with self._lock:
|
||||
self._budget_paused = True
|
||||
await self.set_status(agent_id, "budget_paused")
|
||||
|
||||
async def park_for_budget(self, agent_id: str) -> bool:
|
||||
"""Park ``agent_id`` before an LLM call (pause policy).
|
||||
|
||||
Only a ``running`` agent parks; returns False when it was stopped in the
|
||||
meantime, so the stop is not overwritten.
|
||||
"""
|
||||
async with self._lock:
|
||||
if self.statuses.get(agent_id) != "running":
|
||||
return False
|
||||
self._set_status_locked(agent_id, "budget_paused")
|
||||
logger.info("agent.status %s=budget_paused", agent_id)
|
||||
await self._maybe_snapshot()
|
||||
return True
|
||||
|
||||
async def pause_budget(self) -> None:
|
||||
"""Operator pause: every agent parks before its next LLM call.
|
||||
|
||||
Agents mid-call or mid-tool finish that step first, so their spend still
|
||||
lands; nothing is cancelled.
|
||||
"""
|
||||
async with self._lock:
|
||||
self._budget_paused = True
|
||||
logger.info("scan paused by the operator")
|
||||
await self._maybe_snapshot()
|
||||
|
||||
async def resume_budget(self, *, max_budget_usd: float | None = None) -> list[str]:
|
||||
"""Lift the pause and wake every parked agent; returns the woken agent ids.
|
||||
|
||||
With ``max_budget_usd`` the scan's limit is replaced first (``None`` keeps
|
||||
the current one). Agents continue with the LLM call they parked on; no
|
||||
message is added to any session. An agent that parks again on its next
|
||||
call (the new limit is already spent) is not an error.
|
||||
"""
|
||||
if max_budget_usd is not None and self._set_budget_limit is not None:
|
||||
self._set_budget_limit(max_budget_usd)
|
||||
async with self._lock:
|
||||
self._budget_paused = False
|
||||
self._resume_epoch += 1
|
||||
woken = [aid for aid, status in self.statuses.items() if status == "budget_paused"]
|
||||
for aid in woken:
|
||||
self.runtimes.setdefault(aid, AgentRuntime()).wake.set()
|
||||
logger.info("scan resumed; woke %d parked agent(s)", len(woken))
|
||||
await self._maybe_snapshot()
|
||||
return woken
|
||||
|
||||
async def wait_for_budget_resume(self, agent_id: str, *, parked_epoch: int) -> None:
|
||||
"""Block until a resume newer than ``parked_epoch``, a scan-wide stop, or
|
||||
the agent itself being stopped while parked.
|
||||
|
||||
``parked_epoch`` is the ``resume_epoch`` the agent read when it decided to
|
||||
park, so a resume that lands between that decision and this wait is not
|
||||
missed. The switch back to ``running`` happens under the lock, so a stop
|
||||
can never be overwritten by it; on return the agent is ``running`` unless
|
||||
it was stopped.
|
||||
"""
|
||||
while True:
|
||||
async with self._lock:
|
||||
runtime = self.runtimes.setdefault(agent_id, AgentRuntime())
|
||||
if self._budget_stopped or self.statuses.get(agent_id) != "budget_paused":
|
||||
return
|
||||
if self._resume_epoch != parked_epoch:
|
||||
self._set_status_locked(agent_id, "running")
|
||||
break
|
||||
wake = runtime.wake
|
||||
wake.clear()
|
||||
await wake.wait()
|
||||
logger.info("agent.status %s=running", agent_id)
|
||||
await self._maybe_snapshot()
|
||||
|
||||
async def resume_from_budget_pause(self, *, exclude: str | None = None) -> None:
|
||||
"""Legacy interactive resume: extend by the original budget and nudge agents."""
|
||||
async with self._lock:
|
||||
if not self._budget_paused:
|
||||
return
|
||||
@@ -352,25 +258,20 @@ class AgentCoordinator:
|
||||
async with self._lock:
|
||||
if agent_id not in self.statuses:
|
||||
return
|
||||
self._set_status_locked(agent_id, status, error=error)
|
||||
self.statuses[agent_id] = status # type: ignore[assignment]
|
||||
if error is not None:
|
||||
self.errors[agent_id] = error
|
||||
elif status == "running":
|
||||
self.errors.pop(agent_id, None)
|
||||
if status == "running":
|
||||
# Running again means a fresh stint that owes its parent its own notice.
|
||||
self._parent_notified.discard(agent_id)
|
||||
runtime = self.runtimes.setdefault(agent_id, AgentRuntime())
|
||||
runtime.user_wake_required = status in {"failed", "crashed"}
|
||||
runtime.wake.set()
|
||||
logger.info("agent.status %s=%s", agent_id, status)
|
||||
await self._maybe_snapshot()
|
||||
|
||||
def _set_status_locked(
|
||||
self, agent_id: str, status: Status | str, *, error: str | None = None
|
||||
) -> None:
|
||||
self.statuses[agent_id] = status # type: ignore[assignment]
|
||||
if error is not None:
|
||||
self.errors[agent_id] = error
|
||||
elif status == "running":
|
||||
self.errors.pop(agent_id, None)
|
||||
if status == "running":
|
||||
# Running again means a fresh stint that owes its parent its own notice.
|
||||
self._parent_notified.discard(agent_id)
|
||||
runtime = self.runtimes.setdefault(agent_id, AgentRuntime())
|
||||
runtime.user_wake_required = status in {"failed", "crashed"}
|
||||
runtime.wake.set()
|
||||
|
||||
async def claim_parent_notice(self, agent_id: str) -> bool:
|
||||
"""Reserve the one notice a child owes its parent when it stops running.
|
||||
|
||||
@@ -407,7 +308,7 @@ class AgentCoordinator:
|
||||
unknown, or it is terminal and its loop does not park for wake-ups.
|
||||
"""
|
||||
from_user = message.get("from") == "user"
|
||||
if from_user and self._budget_paused and self._budget_policy != "pause":
|
||||
if from_user and self._budget_paused:
|
||||
await self.resume_from_budget_pause(exclude=target_agent_id)
|
||||
async with self._lock:
|
||||
if target_agent_id not in self.statuses:
|
||||
@@ -561,7 +462,7 @@ class AgentCoordinator:
|
||||
"parent_id": self.parent_of.get(aid),
|
||||
}
|
||||
for aid, status in self.statuses.items()
|
||||
if aid != agent_id and status in ACTIVE_STATUSES
|
||||
if aid != agent_id and status in {"running", "waiting"}
|
||||
]
|
||||
|
||||
async def graph_snapshot(
|
||||
|
||||
+28
-49
@@ -522,51 +522,33 @@ async def _run_until_lifecycle(
|
||||
await coordinator.set_status(agent_id, "stopped")
|
||||
raise SubagentBudgetReservedError("scan reached the sub-agent budget reserve")
|
||||
|
||||
try:
|
||||
if interactive:
|
||||
result = await _run_cycle_parked(
|
||||
agent,
|
||||
coordinator,
|
||||
agent_id,
|
||||
input_data=input_data,
|
||||
run_config=run_config,
|
||||
context=context,
|
||||
max_turns=max_turns,
|
||||
session=session,
|
||||
event_sink=event_sink,
|
||||
hooks=hooks,
|
||||
)
|
||||
else:
|
||||
result = await _run_cycle(
|
||||
agent,
|
||||
coordinator,
|
||||
agent_id,
|
||||
input_data=input_data,
|
||||
run_config=run_config,
|
||||
context=context,
|
||||
max_turns=max_turns,
|
||||
session=session,
|
||||
interactive=False,
|
||||
event_sink=event_sink,
|
||||
hooks=hooks,
|
||||
)
|
||||
except BudgetPausedError as exc:
|
||||
if coordinator.budget_policy != "pause":
|
||||
raise
|
||||
# The agent parked right before an LLM call; everything up to that
|
||||
# point is already in its session. Once resumed, the same call goes
|
||||
# out with nothing added to the conversation.
|
||||
await coordinator.wait_for_budget_resume(agent_id, parked_epoch=exc.resume_epoch)
|
||||
if (
|
||||
not coordinator.budget_stopped
|
||||
and await _agent_status(coordinator, agent_id) != "running"
|
||||
):
|
||||
# Stopped while parked (operator stop or a parent's stop_agent).
|
||||
await coordinator.reset_recovery(agent_id)
|
||||
return result
|
||||
if session is not None:
|
||||
input_data = []
|
||||
continue
|
||||
if interactive:
|
||||
result = await _run_cycle_parked(
|
||||
agent,
|
||||
coordinator,
|
||||
agent_id,
|
||||
input_data=input_data,
|
||||
run_config=run_config,
|
||||
context=context,
|
||||
max_turns=max_turns,
|
||||
session=session,
|
||||
event_sink=event_sink,
|
||||
hooks=hooks,
|
||||
)
|
||||
else:
|
||||
result = await _run_cycle(
|
||||
agent,
|
||||
coordinator,
|
||||
agent_id,
|
||||
input_data=input_data,
|
||||
run_config=run_config,
|
||||
context=context,
|
||||
max_turns=max_turns,
|
||||
session=session,
|
||||
interactive=False,
|
||||
event_sink=event_sink,
|
||||
hooks=hooks,
|
||||
)
|
||||
|
||||
status = await _agent_status(coordinator, agent_id)
|
||||
if status != "running":
|
||||
@@ -777,10 +759,7 @@ async def _run_cycle( # noqa: PLR0912, PLR0915
|
||||
await coordinator.detach_stream(agent_id, stream)
|
||||
except BudgetPausedError as exc:
|
||||
logger.info("agent %s paused at the scan budget limit: %s", agent_id, exc)
|
||||
if coordinator.budget_policy == "pause":
|
||||
await coordinator.park_for_budget(agent_id)
|
||||
else:
|
||||
await coordinator.pause_for_budget(agent_id)
|
||||
await coordinator.pause_for_budget(agent_id)
|
||||
raise
|
||||
except SubagentBudgetReservedError as exc:
|
||||
logger.info("sub-agent %s stopped at the budget reserve: %s", agent_id, exc)
|
||||
|
||||
+7
-89
@@ -8,7 +8,6 @@ from typing import TYPE_CHECKING, Any
|
||||
|
||||
from agents.lifecycle import RunHooks
|
||||
|
||||
from strix.core.agents import BudgetPolicy, coordinator_from_context
|
||||
from strix.report.state import get_global_report_state
|
||||
|
||||
|
||||
@@ -23,26 +22,6 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
LLM_TURN_KEY = "llm_turn"
|
||||
|
||||
# ``BudgetPolicy`` decides what happens when the accumulated LLM cost reaches
|
||||
# ``max_budget_usd``.
|
||||
#
|
||||
# ``stop``: the agents are warned as the limit approaches, sub-agents are cut at
|
||||
# a reserve so the root can write its report, and the scan ends at the limit.
|
||||
#
|
||||
# ``pause``: the agents are never told a limit exists. Every agent parks right
|
||||
# before its next LLM call once the limit is reached (or an operator pauses the
|
||||
# scan), keeping its session, context and sandbox alive, and continues with that
|
||||
# same call when the operator raises the limit or resumes.
|
||||
__all__ = [
|
||||
"LLM_TURN_KEY",
|
||||
"BudgetExceededError",
|
||||
"BudgetPausedError",
|
||||
"BudgetPolicy",
|
||||
"ReportUsageHooks",
|
||||
"SubagentBudgetReservedError",
|
||||
"recomputed_budget_flags",
|
||||
]
|
||||
|
||||
_STAGE_LABELS: tuple[str, ...] = ("NOTICE", "URGENT", "CRITICAL")
|
||||
_TURN_WARN_BANDS: tuple[float, ...] = (0.70, 0.85, 0.95)
|
||||
_ROOT_BUDGET_WARN_BANDS: tuple[float, ...] = (0.70, 0.85, 0.95)
|
||||
@@ -59,20 +38,7 @@ class SubagentBudgetReservedError(RuntimeError):
|
||||
|
||||
|
||||
class BudgetPausedError(RuntimeError):
|
||||
"""Raised to park one agent until the scan budget is raised or the pause lifted.
|
||||
|
||||
``resume_epoch`` is the coordinator's ``resume_epoch`` at the moment the agent
|
||||
decided to park; the agent waits for a resume newer than that.
|
||||
"""
|
||||
|
||||
def __init__(self, message: str, *, resume_epoch: int = 0) -> None:
|
||||
super().__init__(message)
|
||||
self.resume_epoch = resume_epoch
|
||||
|
||||
|
||||
def _validate_budget(max_budget_usd: float | 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")
|
||||
"""Raised to park one agent when an interactive scan reaches its budget."""
|
||||
|
||||
|
||||
def recomputed_budget_flags(
|
||||
@@ -80,12 +46,11 @@ def recomputed_budget_flags(
|
||||
max_budget_usd: float | None,
|
||||
*,
|
||||
interactive: bool,
|
||||
budget_policy: BudgetPolicy = "stop",
|
||||
) -> tuple[bool, bool]:
|
||||
"""Return the (budget_stopped, reserve_stopped) flags a resumed scan should carry."""
|
||||
if max_budget_usd is None:
|
||||
return False, False
|
||||
if interactive or budget_policy == "pause":
|
||||
if interactive:
|
||||
return False, False
|
||||
budget_stopped = cost >= max_budget_usd
|
||||
reserve_stopped = cost >= max_budget_usd * _SUBAGENT_BUDGET_RESERVE
|
||||
@@ -156,32 +121,18 @@ class ReportUsageHooks(RunHooks[dict[str, Any]]):
|
||||
max_budget_usd: float | None = None,
|
||||
max_turns: int | None = None,
|
||||
interactive: bool = False,
|
||||
budget_policy: BudgetPolicy = "stop",
|
||||
) -> None:
|
||||
_validate_budget(max_budget_usd)
|
||||
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")
|
||||
if max_turns is not None and max_turns <= 0:
|
||||
raise ValueError("max_turns must be a positive integer")
|
||||
if budget_policy not in ("stop", "pause"):
|
||||
raise ValueError(f"unknown budget_policy: {budget_policy!r}")
|
||||
self._model = model
|
||||
self._max_budget_usd = max_budget_usd
|
||||
self._budget_increment = max_budget_usd
|
||||
self._max_turns = max_turns
|
||||
self._interactive = interactive
|
||||
self._budget_policy: BudgetPolicy = budget_policy
|
||||
|
||||
@property
|
||||
def max_budget_usd(self) -> float | None:
|
||||
return self._max_budget_usd
|
||||
|
||||
@property
|
||||
def budget_policy(self) -> BudgetPolicy:
|
||||
return self._budget_policy
|
||||
|
||||
def set_max_budget_usd(self, max_budget_usd: float | None) -> None:
|
||||
"""Replace the scan's cost limit; ``None`` removes it."""
|
||||
_validate_budget(max_budget_usd)
|
||||
self._max_budget_usd = max_budget_usd
|
||||
|
||||
def extend_budget(self) -> None:
|
||||
if self._max_budget_usd is None or self._budget_increment is None:
|
||||
@@ -195,8 +146,6 @@ class ReportUsageHooks(RunHooks[dict[str, Any]]):
|
||||
system_prompt: str | None, # noqa: ARG002
|
||||
input_items: list[TResponseInputItem],
|
||||
) -> None:
|
||||
if self._budget_policy == "pause":
|
||||
self._pause_if_limited(context)
|
||||
context.context[LLM_TURN_KEY] = int(context.context.get(LLM_TURN_KEY, 0)) + 1
|
||||
try:
|
||||
self._maybe_warn_turns(context, input_items)
|
||||
@@ -204,32 +153,6 @@ class ReportUsageHooks(RunHooks[dict[str, Any]]):
|
||||
except Exception:
|
||||
logger.exception("budget/turn warning injection failed")
|
||||
|
||||
def _pause_if_limited(self, context: RunContextWrapper[dict[str, Any]]) -> None:
|
||||
"""Park the agent before a paid call when the scan is at its limit or paused.
|
||||
|
||||
Only calls that have already returned are counted, so calls in flight on
|
||||
other agents still land and are paid for: ``spent`` may end up above the
|
||||
limit, which is expected and never an error under this policy.
|
||||
"""
|
||||
coordinator = coordinator_from_context(context.context)
|
||||
epoch = coordinator.resume_epoch if coordinator is not None else 0
|
||||
if coordinator is not None and coordinator.budget_paused:
|
||||
raise BudgetPausedError(
|
||||
"scan paused; waiting for the operator to resume", resume_epoch=epoch
|
||||
)
|
||||
if self._max_budget_usd is None:
|
||||
return
|
||||
report_state = get_global_report_state()
|
||||
if report_state is None:
|
||||
return
|
||||
cost = report_state.get_total_llm_cost()
|
||||
if cost >= self._max_budget_usd:
|
||||
raise BudgetPausedError(
|
||||
f"Scan budget of ${self._max_budget_usd:.2f} reached (spent ${cost:.4f}); "
|
||||
"pausing until the operator raises the limit",
|
||||
resume_epoch=epoch,
|
||||
)
|
||||
|
||||
def _maybe_warn_turns(
|
||||
self,
|
||||
context: RunContextWrapper[dict[str, Any]],
|
||||
@@ -259,7 +182,7 @@ class ReportUsageHooks(RunHooks[dict[str, Any]]):
|
||||
context: RunContextWrapper[dict[str, Any]],
|
||||
input_items: list[TResponseInputItem],
|
||||
) -> None:
|
||||
if self._max_budget_usd is None or self._budget_policy == "pause":
|
||||
if self._max_budget_usd is None:
|
||||
return
|
||||
report_state = get_global_report_state()
|
||||
if report_state is None:
|
||||
@@ -327,11 +250,6 @@ class ReportUsageHooks(RunHooks[dict[str, Any]]):
|
||||
except Exception:
|
||||
logger.exception("failed to record SDK usage for agent %s", agent_id)
|
||||
|
||||
if self._budget_policy == "pause":
|
||||
# The finished call is paid for and its tool calls still run for free;
|
||||
# the agent parks before its next call, in ``on_llm_start``.
|
||||
return
|
||||
|
||||
if self._max_budget_usd is not None:
|
||||
cost = report_state.get_total_llm_cost()
|
||||
if cost >= self._max_budget_usd:
|
||||
|
||||
@@ -267,7 +267,9 @@ def make_model_settings(
|
||||
and reasoning_effort != "none"
|
||||
and model_supports_reasoning(model_name)
|
||||
):
|
||||
model_settings = model_settings.resolve(_reasoning_settings(reasoning_effort))
|
||||
model_settings = model_settings.resolve(
|
||||
_reasoning_settings(reasoning_effort),
|
||||
)
|
||||
if force_required_tool_choice and _accepts_required_tool_choice(model_name):
|
||||
model_settings = model_settings.resolve(ModelSettings(tool_choice="required"))
|
||||
|
||||
|
||||
+4
-27
@@ -18,16 +18,15 @@ from openai import RateLimitError
|
||||
|
||||
from strix.agents.factory import build_strix_agent, make_child_factory
|
||||
from strix.agents.prompt import render_scope_prompt, render_system_prompt
|
||||
from strix.config import codex, load_settings
|
||||
from strix.config import load_settings
|
||||
from strix.config.models import (
|
||||
StrixProvider,
|
||||
configure_sdk_api_route,
|
||||
configure_sdk_model_defaults,
|
||||
supports_strict_tool_schemas,
|
||||
uses_chat_completions_tool_schema,
|
||||
)
|
||||
from strix.config.settings import DEFAULT_MAX_TURNS
|
||||
from strix.core.agents import AgentCoordinator, BudgetPolicy
|
||||
from strix.core.agents import AgentCoordinator
|
||||
from strix.core.execution import (
|
||||
respawn_subagents,
|
||||
run_agent_loop,
|
||||
@@ -190,7 +189,6 @@ async def run_strix_scan(
|
||||
interactive: bool = False,
|
||||
max_turns: int = DEFAULT_MAX_TURNS,
|
||||
max_budget_usd: float | None = None,
|
||||
budget_policy: BudgetPolicy = "stop",
|
||||
model: str | None = None,
|
||||
cleanup_on_exit: bool = True,
|
||||
event_sink: StreamEventSink | None = None,
|
||||
@@ -210,12 +208,6 @@ async def run_strix_scan(
|
||||
``extra_system_prompt_context`` is merged into the root agent's scan
|
||||
context before prompt rendering. Child agents keep the standard scan prompt
|
||||
and context.
|
||||
``budget_policy`` decides what happens when the LLM spend reaches
|
||||
``max_budget_usd``: ``"stop"`` warns the agents as the limit approaches and
|
||||
ends the scan at it; ``"pause"`` tells the agents nothing and parks every
|
||||
agent before its next LLM call until the caller resumes the scan through
|
||||
``coordinator.resume_budget()`` (optionally with a higher limit) or cancels
|
||||
it. ``coordinator.pause_budget()`` parks a running scan the same way.
|
||||
``mcp_connection_requests`` supplies the run's MCP connections from any
|
||||
source: when given, the engine connects those requests; when ``None`` (the
|
||||
command-line default) it reads ``~/.strix/mcp-servers.json`` itself. Either
|
||||
@@ -258,22 +250,15 @@ async def run_strix_scan(
|
||||
raise RuntimeError(
|
||||
"No LLM model configured. Set STRIX_LLM env or pass model= to run_strix_scan().",
|
||||
)
|
||||
if resolved_model != (settings.llm.model or "").strip() and not codex.subscription_model(
|
||||
resolved_model
|
||||
):
|
||||
configure_sdk_api_route(resolved_model, settings)
|
||||
logger.info("LLM model resolved: %s", resolved_model)
|
||||
chat_completions_tools = uses_chat_completions_tool_schema(resolved_model, settings)
|
||||
strict_tool_schemas = supports_strict_tool_schemas(resolved_model)
|
||||
if not strict_tool_schemas:
|
||||
logger.info("Sending non-strict tool schemas: %s caps strict tools", resolved_model)
|
||||
|
||||
if budget_policy not in ("stop", "pause"):
|
||||
raise ValueError(f"unknown budget_policy: {budget_policy!r}")
|
||||
if coordinator is None:
|
||||
coordinator = AgentCoordinator()
|
||||
coordinator.set_snapshot_path(agents_path)
|
||||
coordinator.set_budget_policy(budget_policy)
|
||||
|
||||
from strix.tools.coverage.tools import hydrate_coverage_from_disk
|
||||
from strix.tools.notes.tools import hydrate_notes_from_disk
|
||||
@@ -304,17 +289,11 @@ async def run_strix_scan(
|
||||
report_state.get_total_llm_cost(),
|
||||
max_budget_usd,
|
||||
interactive=interactive,
|
||||
budget_policy=budget_policy,
|
||||
)
|
||||
# Under the pause policy the hooks re-park at the first call if the
|
||||
# spend is still at the limit, so a restored pause flag would only
|
||||
# hold agents back after the limit was raised.
|
||||
await coordinator.reset_budget_stops(
|
||||
budget_stopped=budget_stopped,
|
||||
reserve_stopped=reserve_stopped,
|
||||
budget_paused=(
|
||||
interactive and budget_policy != "pause" and coordinator.budget_paused
|
||||
),
|
||||
budget_paused=interactive and coordinator.budget_paused,
|
||||
)
|
||||
for aid, parent in coordinator.parent_of.items():
|
||||
if parent is None:
|
||||
@@ -393,10 +372,8 @@ async def run_strix_scan(
|
||||
max_budget_usd=max_budget_usd,
|
||||
max_turns=max_turns,
|
||||
interactive=interactive,
|
||||
budget_policy=budget_policy,
|
||||
)
|
||||
coordinator.set_budget_limit_setter(hooks.set_max_budget_usd)
|
||||
if interactive and budget_policy != "pause":
|
||||
if interactive:
|
||||
coordinator.set_budget_extender(hooks.extend_budget)
|
||||
|
||||
scope_context = build_scope_context(scan_config)
|
||||
|
||||
+30
-2
@@ -135,12 +135,14 @@ def _subscription_error_hint(exc: BaseException) -> str | None:
|
||||
return None
|
||||
|
||||
|
||||
async def warm_up_llm() -> None:
|
||||
async def warm_up_llm(show_model_warning: bool = True) -> None:
|
||||
from agents.models.interface import ModelTracing
|
||||
|
||||
from strix.config.models import (
|
||||
RECOMMENDED_MODEL_NAMES,
|
||||
configure_sdk_model_defaults,
|
||||
is_known_openai_bare_model,
|
||||
is_recommended_or_frontier_model,
|
||||
)
|
||||
from strix.core.inputs import make_model_settings
|
||||
|
||||
@@ -184,6 +186,32 @@ async def warm_up_llm() -> None:
|
||||
)
|
||||
sys.exit(1)
|
||||
|
||||
if show_model_warning and raw_model and not is_recommended_or_frontier_model(raw_model):
|
||||
warn_text = Text()
|
||||
warn_text.append("MODEL QUALITY WARNING", style="bold yellow")
|
||||
warn_text.append("\n\n", style="white")
|
||||
warn_text.append(f"'{raw_model}'", style="bold cyan")
|
||||
warn_text.append(
|
||||
" is not a recommended frontier model for Strix.\nSecurity scans work best with:\n",
|
||||
style="white",
|
||||
)
|
||||
for recommended_model in RECOMMENDED_MODEL_NAMES:
|
||||
warn_text.append(f"• {recommended_model}\n", style="bold cyan")
|
||||
warn_text.append(
|
||||
"\nYou can continue, but weaker models may miss vulnerabilities "
|
||||
"or produce lower-quality findings.",
|
||||
style="white",
|
||||
)
|
||||
console.print(
|
||||
Panel(
|
||||
warn_text,
|
||||
title="[bold white]STRIX",
|
||||
title_align="left",
|
||||
border_style="yellow",
|
||||
padding=(1, 2),
|
||||
),
|
||||
)
|
||||
|
||||
await preflight_model_connection(raw_model, settings=settings)
|
||||
logger.info("LLM warm-up succeeded for model %s", (llm.model or "").strip())
|
||||
|
||||
@@ -375,7 +403,7 @@ def _bootstrap_scan(args: argparse.Namespace) -> None:
|
||||
"""
|
||||
set_scan_phase("preflight")
|
||||
try:
|
||||
asyncio.run(warm_up_llm())
|
||||
asyncio.run(warm_up_llm(show_model_warning=True))
|
||||
except ModelConnectionError as exc:
|
||||
report_error("model_connection_failed", exc)
|
||||
_print_model_connection_error(exc, exc.model_name)
|
||||
|
||||
@@ -11,6 +11,7 @@ from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from strix.config import load_settings
|
||||
from strix.config.models import is_recommended_or_frontier_model
|
||||
from strix.config.settings import DEFAULT_MAX_TURNS
|
||||
from strix.interface.tui.backend.live_view import TuiLiveView
|
||||
from strix.interface.tui.backend.projection import (
|
||||
@@ -188,6 +189,11 @@ class TuiController:
|
||||
subscription = False
|
||||
with contextlib.suppress(Exception):
|
||||
subscription = is_subscription_run(self.report_state)
|
||||
model_warning = ""
|
||||
if model and not is_recommended_or_frontier_model(model):
|
||||
model_warning = (
|
||||
f"{model} is not a recommended frontier model. Pentest quality could be degraded."
|
||||
)
|
||||
state = {
|
||||
"setup_mode": self.setup_mode,
|
||||
"scan_started": self.scan_started,
|
||||
@@ -205,6 +211,7 @@ class TuiController:
|
||||
"scope_mode": self.scope_mode,
|
||||
"diff_base": terminal_projection(self.diff_base, max_string=256),
|
||||
"model": terminal_projection(model, max_string=256),
|
||||
"model_warning": terminal_projection(model_warning, max_string=512),
|
||||
"caido_url": terminal_projection(
|
||||
getattr(self.report_state, "caido_url", None), max_string=1024
|
||||
),
|
||||
|
||||
@@ -150,6 +150,7 @@ def bounded_state_projection(state: dict[str, Any]) -> dict[str, Any]:
|
||||
key: state["usage"][key] for key in ("total_tokens", "cost") if key in state["usage"]
|
||||
}
|
||||
state["error"] = terminal_projection(state["error"], max_string=512)
|
||||
state["model_warning"] = terminal_projection(state["model_warning"], max_string=256)
|
||||
state["caido_url"] = terminal_projection(state["caido_url"], max_string=256)
|
||||
state["viewer_url"] = terminal_projection(state["viewer_url"], max_string=256)
|
||||
if encoded_size(state) <= STATE_TARGET_BYTES:
|
||||
@@ -172,6 +173,7 @@ def bounded_state_projection(state: dict[str, Any]) -> dict[str, Any]:
|
||||
"scope_mode": state["scope_mode"],
|
||||
"diff_base": state["diff_base"],
|
||||
"model": state["model"],
|
||||
"model_warning": "",
|
||||
"caido_url": None,
|
||||
"messages": [],
|
||||
"usage": state["usage"],
|
||||
|
||||
@@ -370,6 +370,17 @@ func TestStartedSnapshotTransitionsToLiveView(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplashModelWarningRendersTheBackendSentenceOnce(t *testing.T) {
|
||||
warning := "openai/glm-5.3 is not a recommended frontier model. Pentest quality could be degraded."
|
||||
got := ansi.Strip(splashModelWarning("openai/glm-5.3", warning))
|
||||
if got != "⚠ "+warning {
|
||||
t.Fatalf("splash warning = %q, want %q", got, "⚠ "+warning)
|
||||
}
|
||||
if got := ansi.Strip(splashModelWarning("other/model", warning)); got != "⚠ "+warning {
|
||||
t.Fatalf("splash warning with unrelated model = %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetupStartScreenFitsNarrowTerminal(t *testing.T) {
|
||||
model := New(nil)
|
||||
model.width, model.height = 40, 18
|
||||
|
||||
@@ -423,12 +423,27 @@ func (m Model) splashView() string {
|
||||
content := wordmark() + "\n\n" +
|
||||
welcome + "\n" + version + "\n" + tagline + "\n\n" +
|
||||
start.String() + "\n\n" + url
|
||||
if warn := m.snapshot.ModelWarning; warn != "" {
|
||||
content += "\n\n" + splashModelWarning(m.snapshot.Model, warn)
|
||||
}
|
||||
panel := lipgloss.NewStyle().Border(lipgloss.RoundedBorder()).BorderForeground(green).Padding(1, 6).Align(lipgloss.Center).Render(content)
|
||||
// #splash_screen background is solid black.
|
||||
return lipgloss.Place(m.width, m.height, lipgloss.Center, lipgloss.Center, panel,
|
||||
lipgloss.WithWhitespaceBackground(black))
|
||||
}
|
||||
|
||||
// splashModelWarning renders the backend's full warning sentence, with the
|
||||
// model name highlighted when the sentence leads with it.
|
||||
func splashModelWarning(model, warning string) string {
|
||||
yellow := lipgloss.Color("#eab308")
|
||||
out := lipgloss.NewStyle().Bold(true).Foreground(yellow).Render("⚠ ")
|
||||
if model != "" && strings.HasPrefix(warning, model) {
|
||||
out += lipgloss.NewStyle().Bold(true).Foreground(render.Cyan).Render(model)
|
||||
warning = strings.TrimPrefix(warning, model)
|
||||
}
|
||||
return out + lipgloss.NewStyle().Foreground(yellow).Render(warning)
|
||||
}
|
||||
|
||||
// chatPaneKey identifies everything the bordered trace depends on.
|
||||
type chatPaneKey struct {
|
||||
offset int
|
||||
|
||||
@@ -71,6 +71,7 @@ type Snapshot struct {
|
||||
ScopeMode string `json:"scope_mode"`
|
||||
DiffBase string `json:"diff_base"`
|
||||
Model string `json:"model"`
|
||||
ModelWarning string `json:"model_warning"`
|
||||
CaidoURL string `json:"caido_url"`
|
||||
Messages []Message `json:"messages"`
|
||||
Agents []Agent `json:"-"`
|
||||
|
||||
@@ -12,13 +12,16 @@ from typing import Any, Literal, get_args
|
||||
|
||||
from agents import RunContextWrapper, function_tool
|
||||
|
||||
from strix.core.agents import ACTIVE_STATUSES, Status, coordinator_from_context
|
||||
from strix.core.agents import Status, coordinator_from_context
|
||||
from strix.core.execution import notify_parent_on_terminal
|
||||
from strix.core.hooks import LLM_TURN_KEY
|
||||
from strix.report.state import get_global_report_state
|
||||
from strix.skills import validate_requested_skills
|
||||
|
||||
|
||||
_ACTIVE_STATUSES: frozenset[str] = frozenset({"running", "waiting"})
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@@ -807,13 +810,13 @@ async def stop_agent(
|
||||
)
|
||||
|
||||
current_status = statuses[target_agent_id]
|
||||
if current_status not in ACTIVE_STATUSES:
|
||||
if current_status not in _ACTIVE_STATUSES:
|
||||
return json.dumps(
|
||||
{
|
||||
"success": False,
|
||||
"error": (
|
||||
f"Agent {target_agent_id} is already '{current_status}'; "
|
||||
"stop_agent only acts on running/waiting/paused agents — use "
|
||||
"stop_agent only acts on running/waiting agents — use "
|
||||
"view_agent_graph to find still-active descendants and "
|
||||
"stop them individually, or send_message_to_agent if you "
|
||||
"want to wake this one with new instructions"
|
||||
|
||||
@@ -1075,11 +1075,11 @@ async def create_vulnerability_report(
|
||||
"Techniques" that read like an engineering runbook rather than a
|
||||
client deliverable.
|
||||
- **Use markdown in every text field**: ``**bold**`` for emphasis,
|
||||
``inline code`` for identifiers/values/parameters, and
|
||||
language-tagged fenced code blocks for any code/payload/HTTP
|
||||
``inline code`` for identifiers/values/parameters, and fenced
|
||||
code blocks (```` ```language ````) for any code/payload/HTTP
|
||||
excerpt. Never leave code bare/unformatted. When referencing a
|
||||
file, annotate the opening fence with
|
||||
``title=app.py startLineNumber=42 endLineNumber=50`` after the language.
|
||||
file, annotate the fence, e.g.
|
||||
```` ```python title=app.py startLineNumber=42 endLineNumber=50 ````.
|
||||
- Field discipline: ``poc_description`` is steps only — NO code (all
|
||||
code goes in ``poc_script_code``); ``remediation_steps`` is prose
|
||||
only — NO code/diffs (code fixes go in ``code_locations``).
|
||||
@@ -1191,57 +1191,6 @@ async def create_vulnerability_report(
|
||||
Broken / Risky Crypto, CWE-311 Missing Encryption, CWE-916 Weak
|
||||
Password Hashing.
|
||||
|
||||
Example (abbreviated — mirror this structure)::
|
||||
|
||||
title: "Reflected XSS in /search q parameter"
|
||||
description:
|
||||
The **`q`** parameter of `/search` reflects user input into
|
||||
the HTML response without encoding, allowing script
|
||||
injection.
|
||||
technical_analysis:
|
||||
The handler interpolates `q` directly into the page body:
|
||||
|
||||
```python title=views.py startLineNumber=42 endLineNumber=44
|
||||
html = f"<h2>Results for {q}</h2>"
|
||||
return HttpResponse(html)
|
||||
```
|
||||
|
||||
No output encoding is applied, so `<script>` executes.
|
||||
poc_description:
|
||||
1. Navigate to `/search?q=<payload>`.
|
||||
2. Observe the payload executes in the victim's browser.
|
||||
poc_script_code:
|
||||
```
|
||||
GET /search?q=<script>alert(document.domain)</script>
|
||||
```
|
||||
evidence:
|
||||
Response echoes the payload verbatim:
|
||||
|
||||
```html
|
||||
<h2>Results for <script>alert(document.domain)</script></h2>
|
||||
```
|
||||
assumptions:
|
||||
Assumes a victim can be induced to open a crafted link.
|
||||
remediation_steps:
|
||||
Context-encode all user input rendered into HTML; prefer the
|
||||
template engine's auto-escaping over string interpolation.
|
||||
counterevidence:
|
||||
No output encoding, CSP, or WAF observed on this response;
|
||||
payload executed in a current browser. The parameter is
|
||||
reflected on an unauthenticated route, so no privileged
|
||||
position is required.
|
||||
confidence: "high"
|
||||
severity_change_conditions:
|
||||
A restrictive CSP that blocks inline script execution would
|
||||
reduce impact and lower the severity.
|
||||
fix_effort: "low"
|
||||
|
||||
Nice to have: for code findings, if the checkout has git history, a quick
|
||||
``git blame`` (quote the paths) on the vulnerable line is worth weaving into
|
||||
``technical_analysis`` — who last touched it, when, and in which commit, as
|
||||
part of the prose, not a separate section. Skip it if the line is
|
||||
uncommitted or the command fails.
|
||||
|
||||
Args:
|
||||
title: Specific finding title (e.g.
|
||||
``"SQL Injection in /api/users login parameter"``). Don't
|
||||
@@ -1410,6 +1359,57 @@ async def create_vulnerability_report(
|
||||
fix (summary + rationale). Prose/markdown only — the code
|
||||
change itself belongs in ``code_locations``. Omit for
|
||||
black-box findings.
|
||||
|
||||
Example (abbreviated — mirror this structure)::
|
||||
|
||||
title: "Reflected XSS in /search q parameter"
|
||||
description:
|
||||
The **`q`** parameter of `/search` reflects user input into
|
||||
the HTML response without encoding, allowing script
|
||||
injection.
|
||||
technical_analysis:
|
||||
The handler interpolates `q` directly into the page body:
|
||||
|
||||
```python title=views.py startLineNumber=42 endLineNumber=44
|
||||
html = f"<h2>Results for {q}</h2>"
|
||||
return HttpResponse(html)
|
||||
```
|
||||
|
||||
No output encoding is applied, so `<script>` executes.
|
||||
poc_description:
|
||||
1. Navigate to `/search?q=<payload>`.
|
||||
2. Observe the payload executes in the victim's browser.
|
||||
poc_script_code:
|
||||
```
|
||||
GET /search?q=<script>alert(document.domain)</script>
|
||||
```
|
||||
evidence:
|
||||
Response echoes the payload verbatim:
|
||||
|
||||
```html
|
||||
<h2>Results for <script>alert(document.domain)</script></h2>
|
||||
```
|
||||
assumptions:
|
||||
Assumes a victim can be induced to open a crafted link.
|
||||
remediation_steps:
|
||||
Context-encode all user input rendered into HTML; prefer the
|
||||
template engine's auto-escaping over string interpolation.
|
||||
counterevidence:
|
||||
No output encoding, CSP, or WAF observed on this response;
|
||||
payload executed in a current browser. The parameter is
|
||||
reflected on an unauthenticated route, so no privileged
|
||||
position is required.
|
||||
confidence: "high"
|
||||
severity_change_conditions:
|
||||
A restrictive CSP that blocks inline script execution would
|
||||
reduce impact and lower the severity.
|
||||
fix_effort: "low"
|
||||
|
||||
Nice to have: for code findings, if the checkout has git history, a quick
|
||||
``git blame`` (quote the paths) on the vulnerable line is worth weaving into
|
||||
``technical_analysis`` — who last touched it, when, and in which commit, as
|
||||
part of the prose, not a separate section. Skip it if the line is
|
||||
uncommitted or the command fails.
|
||||
"""
|
||||
(
|
||||
http_exchange_ids,
|
||||
|
||||
@@ -1,540 +0,0 @@
|
||||
"""``budget_policy="pause"``: agents park before a paid call and never hear about budgets."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from strix.core import execution
|
||||
from strix.core.agents import AgentCoordinator
|
||||
from strix.core.execution import _start_child_runner, run_agent_loop
|
||||
from strix.core.hooks import (
|
||||
BudgetExceededError,
|
||||
BudgetPausedError,
|
||||
ReportUsageHooks,
|
||||
recomputed_budget_flags,
|
||||
)
|
||||
from strix.core.sessions import open_agent_session
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import AsyncIterator, Callable
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
COST_PER_CALL = 1.0
|
||||
_CALL_LATENCY_S = 0.005
|
||||
|
||||
|
||||
class _FakeLedger:
|
||||
def __init__(self) -> None:
|
||||
self.cost = 0.0
|
||||
self.calls: list[str] = []
|
||||
self.remaining: dict[str, int] = {}
|
||||
self.warned_inputs: list[list[Any]] = []
|
||||
self.gate: asyncio.Event | None = None
|
||||
self.in_flight = 0
|
||||
|
||||
def record_sdk_usage(self, **_kwargs: Any) -> None:
|
||||
return
|
||||
|
||||
def get_total_llm_cost(self) -> float:
|
||||
return self.cost
|
||||
|
||||
|
||||
class _FakeStream:
|
||||
"""One ``Runner.run_streamed`` call: several LLM turns, each guarded by the hooks.
|
||||
|
||||
Mirrors the SDK's ordering: ``on_llm_start`` runs before the paid request,
|
||||
``on_llm_end`` after it. A ``BudgetPausedError`` from ``on_llm_start`` ends
|
||||
the stream without spending, exactly like the SDK surfacing a hook error.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
ledger: _FakeLedger,
|
||||
hooks: ReportUsageHooks,
|
||||
context: dict[str, Any],
|
||||
agent: Any,
|
||||
coordinator: AgentCoordinator,
|
||||
) -> None:
|
||||
self._ledger = ledger
|
||||
self._hooks = hooks
|
||||
self._context = context
|
||||
self._agent = agent
|
||||
self._coordinator = coordinator
|
||||
self.run_loop_exception: BaseException | None = None
|
||||
self.final_output = None
|
||||
|
||||
async def stream_events(self) -> AsyncIterator[Any]:
|
||||
agent_id = str(self._context.get("agent_id"))
|
||||
ctx_wrapper = MagicMock()
|
||||
ctx_wrapper.context = self._context
|
||||
while self._ledger.remaining.get(agent_id, 0) > 0:
|
||||
input_items: list[Any] = []
|
||||
try:
|
||||
await self._hooks.on_llm_start(ctx_wrapper, self._agent, None, input_items)
|
||||
except BudgetPausedError as exc:
|
||||
self.run_loop_exception = exc
|
||||
return
|
||||
self._ledger.warned_inputs.append(input_items)
|
||||
if self._ledger.gate is not None:
|
||||
self._ledger.in_flight += 1
|
||||
await self._ledger.gate.wait()
|
||||
self._ledger.in_flight -= 1
|
||||
self._ledger.cost += COST_PER_CALL
|
||||
self._ledger.calls.append(agent_id)
|
||||
self._ledger.remaining[agent_id] -= 1
|
||||
await self._hooks.on_llm_end(ctx_wrapper, self._agent, MagicMock())
|
||||
await asyncio.sleep(_CALL_LATENCY_S)
|
||||
if self._coordinator.statuses.get(agent_id) == "running":
|
||||
await self._coordinator.set_status(agent_id, "completed")
|
||||
items: tuple[Any, ...] = ()
|
||||
for item in items:
|
||||
yield item
|
||||
|
||||
def cancel(self, mode: str = "immediate") -> None: # noqa: ARG002
|
||||
return
|
||||
|
||||
|
||||
def _fake_runner(ledger: _FakeLedger, coordinator: AgentCoordinator) -> Any:
|
||||
class _FakeRunner:
|
||||
@staticmethod
|
||||
def run_streamed(
|
||||
agent: Any,
|
||||
input: Any, # noqa: A002, ARG004
|
||||
*,
|
||||
run_config: Any, # noqa: ARG004
|
||||
context: dict[str, Any],
|
||||
max_turns: int, # noqa: ARG004
|
||||
session: Any, # noqa: ARG004
|
||||
hooks: ReportUsageHooks,
|
||||
) -> _FakeStream:
|
||||
return _FakeStream(
|
||||
ledger=ledger,
|
||||
hooks=hooks,
|
||||
context=context,
|
||||
agent=agent,
|
||||
coordinator=coordinator,
|
||||
)
|
||||
|
||||
return _FakeRunner
|
||||
|
||||
|
||||
async def _noop_compact(*_args: Any, **_kwargs: Any) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
async def _wait_until(predicate: Callable[[], bool], *, timeout: float = 5.0) -> None:
|
||||
async def _poll() -> None:
|
||||
while not predicate():
|
||||
await asyncio.sleep(0.001)
|
||||
|
||||
await asyncio.wait_for(_poll(), timeout=timeout)
|
||||
|
||||
|
||||
def _all_parked(coordinator: AgentCoordinator, *agent_ids: str) -> bool:
|
||||
return all(coordinator.statuses.get(aid) == "budget_paused" for aid in agent_ids)
|
||||
|
||||
|
||||
class _Scan:
|
||||
"""Root + children driven through the real non-interactive loops."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
*,
|
||||
max_budget_usd: float | None,
|
||||
) -> None:
|
||||
self.ledger = _FakeLedger()
|
||||
self.hooks = ReportUsageHooks(
|
||||
model="test-model", max_budget_usd=max_budget_usd, budget_policy="pause"
|
||||
)
|
||||
self.coordinator = AgentCoordinator()
|
||||
self.coordinator.set_budget_policy("pause")
|
||||
self.coordinator.set_budget_limit_setter(self.hooks.set_max_budget_usd)
|
||||
monkeypatch.setattr(execution, "Runner", _fake_runner(self.ledger, self.coordinator))
|
||||
monkeypatch.setattr(execution, "_compact_session", _noop_compact)
|
||||
self.db_path = tmp_path / "agents.sqlite"
|
||||
self.sessions: list[Any] = []
|
||||
self.run_config = MagicMock()
|
||||
self.root_ctx: dict[str, Any] = {
|
||||
"agent_id": "root",
|
||||
"parent_id": None,
|
||||
"coordinator": self.coordinator,
|
||||
}
|
||||
self.root_task: asyncio.Task[Any] | None = None
|
||||
|
||||
async def start_root(self, *, calls: int) -> None:
|
||||
self.ledger.remaining["root"] = calls
|
||||
await self.coordinator.register("root", "strix", parent_id=None)
|
||||
session = open_agent_session("root", self.db_path)
|
||||
self.sessions.append(session)
|
||||
self.root_task = asyncio.create_task(
|
||||
run_agent_loop(
|
||||
agent=MagicMock(),
|
||||
initial_input=[],
|
||||
run_config=self.run_config,
|
||||
context=self.root_ctx,
|
||||
max_turns=500,
|
||||
coordinator=self.coordinator,
|
||||
agent_id="root",
|
||||
interactive=False,
|
||||
session=session,
|
||||
hooks=self.hooks,
|
||||
)
|
||||
)
|
||||
|
||||
async def start_child(self, child_id: str, *, calls: int) -> None:
|
||||
self.ledger.remaining[child_id] = calls
|
||||
await self.coordinator.register(child_id, "recon", parent_id="root")
|
||||
await _start_child_runner(
|
||||
parent_ctx=self.root_ctx,
|
||||
coordinator=self.coordinator,
|
||||
agents_db_path=self.db_path,
|
||||
sessions_to_close=self.sessions,
|
||||
run_config=self.run_config,
|
||||
max_turns=500,
|
||||
interactive=False,
|
||||
child_agent=MagicMock(),
|
||||
child_id=child_id,
|
||||
name=f"recon-{child_id}",
|
||||
parent_id="root",
|
||||
task="probe things",
|
||||
initial_input=[],
|
||||
hooks=self.hooks,
|
||||
)
|
||||
|
||||
def tasks(self) -> list[asyncio.Task[Any]]:
|
||||
tasks = [self.root_task] if self.root_task is not None else []
|
||||
tasks.extend(rt.task for rt in self.coordinator.runtimes.values() if rt.task is not None)
|
||||
return tasks
|
||||
|
||||
async def teardown(self) -> None:
|
||||
for task in self.tasks():
|
||||
task.cancel()
|
||||
await asyncio.gather(*self.tasks(), return_exceptions=True)
|
||||
for session in self.sessions:
|
||||
session.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pause_policy_parks_every_agent_at_the_limit_and_resumes_in_place(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
scan = _Scan(tmp_path, monkeypatch, max_budget_usd=5.0)
|
||||
agents = ("root", "child-a", "child-b")
|
||||
|
||||
with patch("strix.core.hooks.get_global_report_state", return_value=scan.ledger):
|
||||
await scan.start_root(calls=100)
|
||||
await scan.start_child("child-a", calls=100)
|
||||
await scan.start_child("child-b", calls=100)
|
||||
|
||||
await _wait_until(lambda: _all_parked(scan.coordinator, *agents))
|
||||
assert scan.ledger.cost == pytest.approx(5.0)
|
||||
assert len(scan.ledger.calls) == 5
|
||||
assert scan.coordinator.budget_paused is False
|
||||
assert scan.coordinator.budget_stopped is False
|
||||
assert scan.coordinator.reserve_stopped is False
|
||||
assert all(not task.done() for task in scan.tasks())
|
||||
|
||||
woken = await scan.coordinator.resume_budget(max_budget_usd=8.0)
|
||||
assert sorted(woken) == sorted(agents)
|
||||
assert scan.hooks.max_budget_usd == 8.0
|
||||
await _wait_until(lambda: scan.ledger.cost >= 8.0)
|
||||
await _wait_until(lambda: _all_parked(scan.coordinator, *agents))
|
||||
assert scan.ledger.cost == pytest.approx(8.0)
|
||||
|
||||
await scan.coordinator.resume_budget(max_budget_usd=9.0)
|
||||
await _wait_until(lambda: scan.ledger.cost >= 9.0)
|
||||
await _wait_until(lambda: _all_parked(scan.coordinator, *agents))
|
||||
assert scan.ledger.cost == pytest.approx(9.0)
|
||||
assert len(scan.ledger.calls) == 9
|
||||
assert all(not task.done() for task in scan.tasks())
|
||||
|
||||
# The model never saw a budget message: no warning band, no resume note.
|
||||
assert scan.ledger.warned_inputs
|
||||
assert all(items == [] for items in scan.ledger.warned_inputs)
|
||||
for session in scan.sessions:
|
||||
assert await session.get_items() == []
|
||||
|
||||
# Stop while parked: an individual stop wakes that loop and it exits.
|
||||
child_a_task = scan.coordinator.runtimes["child-a"].task
|
||||
assert child_a_task is not None
|
||||
await scan.coordinator.request_stop("child-a")
|
||||
await asyncio.wait_for(child_a_task, timeout=5.0)
|
||||
assert scan.coordinator.statuses["child-a"] == "stopped"
|
||||
assert scan.coordinator.statuses["root"] == "budget_paused"
|
||||
assert scan.coordinator.statuses["child-b"] == "budget_paused"
|
||||
|
||||
# A scan-wide cancel while parked tears the rest down cleanly.
|
||||
await scan.teardown()
|
||||
assert all(task.done() for task in scan.tasks())
|
||||
assert scan.ledger.cost == pytest.approx(9.0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pause_policy_operator_pause_and_resume_without_a_limit(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
scan = _Scan(tmp_path, monkeypatch, max_budget_usd=None)
|
||||
agents = ("root", "child-a", "child-b")
|
||||
|
||||
with patch("strix.core.hooks.get_global_report_state", return_value=scan.ledger):
|
||||
await scan.start_root(calls=6)
|
||||
await scan.start_child("child-a", calls=6)
|
||||
await scan.start_child("child-b", calls=6)
|
||||
|
||||
await _wait_until(lambda: scan.ledger.cost >= 3.0)
|
||||
await scan.coordinator.pause_budget()
|
||||
spent_at_pause = scan.ledger.cost
|
||||
await _wait_until(lambda: _all_parked(scan.coordinator, *agents))
|
||||
await _wait_until(lambda: scan.coordinator.budget_paused, timeout=0.1)
|
||||
assert scan.ledger.cost == pytest.approx(spent_at_pause)
|
||||
assert all(not task.done() for task in scan.tasks())
|
||||
|
||||
await asyncio.sleep(0.05)
|
||||
assert scan.ledger.cost == pytest.approx(spent_at_pause)
|
||||
|
||||
woken = await scan.coordinator.resume_budget()
|
||||
assert sorted(woken) == sorted(agents)
|
||||
await _wait_until(lambda: not scan.coordinator.budget_paused, timeout=0.1)
|
||||
await asyncio.wait_for(asyncio.gather(*scan.tasks(), return_exceptions=True), timeout=5.0)
|
||||
assert scan.ledger.cost == pytest.approx(18.0)
|
||||
assert {aid: str(s) for aid, s in scan.coordinator.statuses.items()} == {
|
||||
"root": "completed",
|
||||
"child-a": "completed",
|
||||
"child-b": "completed",
|
||||
}
|
||||
assert all(items == [] for items in scan.ledger.warned_inputs)
|
||||
|
||||
for session in scan.sessions:
|
||||
session.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pause_policy_keeps_in_flight_overshoot(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
scan = _Scan(tmp_path, monkeypatch, max_budget_usd=1.0)
|
||||
scan.ledger.gate = asyncio.Event()
|
||||
await scan.coordinator.register("root", "strix", parent_id=None)
|
||||
|
||||
with patch("strix.core.hooks.get_global_report_state", return_value=scan.ledger):
|
||||
await scan.start_child("child-a", calls=100)
|
||||
await scan.start_child("child-b", calls=100)
|
||||
|
||||
# Both calls were dispatched under the limit; neither is cancelled.
|
||||
await _wait_until(lambda: scan.ledger.in_flight == 2)
|
||||
assert scan.ledger.cost == 0.0
|
||||
scan.ledger.gate.set()
|
||||
|
||||
await _wait_until(lambda: _all_parked(scan.coordinator, "child-a", "child-b"))
|
||||
assert scan.ledger.cost == pytest.approx(2.0)
|
||||
assert scan.hooks.max_budget_usd is not None
|
||||
assert scan.ledger.cost > scan.hooks.max_budget_usd
|
||||
assert scan.coordinator.budget_stopped is False
|
||||
assert all(not task.done() for task in scan.tasks())
|
||||
|
||||
# Resuming at a limit that is already spent parks again without a call.
|
||||
await scan.coordinator.resume_budget(max_budget_usd=1.5)
|
||||
await asyncio.sleep(0.05)
|
||||
await _wait_until(lambda: _all_parked(scan.coordinator, "child-a", "child-b"))
|
||||
assert scan.ledger.cost == pytest.approx(2.0)
|
||||
|
||||
await scan.teardown()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resume_between_park_decision_and_wait_is_not_missed() -> None:
|
||||
coordinator = AgentCoordinator()
|
||||
coordinator.set_budget_policy("pause")
|
||||
await coordinator.register("a", "strix", parent_id=None)
|
||||
|
||||
parked_epoch = coordinator.resume_epoch
|
||||
await coordinator.park_for_budget("a")
|
||||
assert _all_parked(coordinator, "a")
|
||||
await coordinator.resume_budget()
|
||||
|
||||
await asyncio.wait_for(
|
||||
coordinator.wait_for_budget_resume("a", parked_epoch=parked_epoch), timeout=1.0
|
||||
)
|
||||
assert coordinator.statuses["a"] == "running"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parked_wait_returns_on_stop_signals() -> None:
|
||||
coordinator = AgentCoordinator()
|
||||
coordinator.set_budget_policy("pause")
|
||||
await coordinator.register("a", "strix", parent_id=None)
|
||||
await coordinator.register("b", "recon", parent_id="a")
|
||||
await coordinator.park_for_budget("a")
|
||||
await coordinator.park_for_budget("b")
|
||||
|
||||
wait_a = asyncio.create_task(
|
||||
coordinator.wait_for_budget_resume("a", parked_epoch=coordinator.resume_epoch)
|
||||
)
|
||||
wait_b = asyncio.create_task(
|
||||
coordinator.wait_for_budget_resume("b", parked_epoch=coordinator.resume_epoch)
|
||||
)
|
||||
await asyncio.sleep(0.02)
|
||||
assert not wait_a.done()
|
||||
assert not wait_b.done()
|
||||
|
||||
await coordinator.request_stop("b")
|
||||
await asyncio.wait_for(wait_b, timeout=1.0)
|
||||
assert coordinator.statuses["b"] == "stopped"
|
||||
assert not wait_a.done()
|
||||
|
||||
await coordinator.trigger_budget_stop()
|
||||
await asyncio.wait_for(wait_a, timeout=1.0)
|
||||
assert coordinator.statuses["a"] == "budget_paused"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_park_never_overwrites_a_stop_that_landed_first() -> None:
|
||||
coordinator = AgentCoordinator()
|
||||
coordinator.set_budget_policy("pause")
|
||||
await coordinator.register("a", "strix", parent_id=None)
|
||||
|
||||
parked_epoch = coordinator.resume_epoch
|
||||
await coordinator.request_stop("a")
|
||||
assert await coordinator.park_for_budget("a") is False
|
||||
assert coordinator.statuses["a"] == "stopped"
|
||||
|
||||
await asyncio.wait_for(
|
||||
coordinator.wait_for_budget_resume("a", parked_epoch=parked_epoch), timeout=1.0
|
||||
)
|
||||
assert coordinator.statuses["a"] == "stopped"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parked_children_keep_the_scan_open() -> None:
|
||||
coordinator = AgentCoordinator()
|
||||
coordinator.set_budget_policy("pause")
|
||||
await coordinator.register("root", "strix", parent_id=None)
|
||||
await coordinator.register("child", "recon", parent_id="root")
|
||||
assert await coordinator.park_for_budget("child") is True
|
||||
|
||||
active = await coordinator.active_agents_except("root")
|
||||
assert [a["agent_id"] for a in active] == ["child"]
|
||||
assert active[0]["status"] == "budget_paused"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resume_budget_replaces_the_limit_and_validates_it() -> None:
|
||||
hooks = ReportUsageHooks(model="m", max_budget_usd=10.0, budget_policy="pause")
|
||||
coordinator = AgentCoordinator()
|
||||
coordinator.set_budget_policy("pause")
|
||||
coordinator.set_budget_limit_setter(hooks.set_max_budget_usd)
|
||||
|
||||
await coordinator.resume_budget(max_budget_usd=25.0)
|
||||
assert hooks.max_budget_usd == 25.0
|
||||
await coordinator.resume_budget()
|
||||
assert hooks.max_budget_usd == 25.0
|
||||
with pytest.raises(ValueError, match="greater than 0"):
|
||||
await coordinator.resume_budget(max_budget_usd=0.0)
|
||||
with pytest.raises(ValueError, match="finite"):
|
||||
await coordinator.resume_budget(max_budget_usd=float("inf"))
|
||||
assert hooks.max_budget_usd == 25.0
|
||||
|
||||
|
||||
def _ctx(coordinator: AgentCoordinator | None, *, parent_id: str | None = None) -> MagicMock:
|
||||
wrapper = MagicMock()
|
||||
wrapper.context = {"agent_id": "x", "parent_id": parent_id}
|
||||
if coordinator is not None:
|
||||
wrapper.context["coordinator"] = coordinator
|
||||
return wrapper
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pause_hooks_never_warn_and_park_only_at_the_limit() -> None:
|
||||
ledger = _FakeLedger()
|
||||
hooks = ReportUsageHooks(model="m", max_budget_usd=10.0, budget_policy="pause")
|
||||
coordinator = AgentCoordinator()
|
||||
coordinator.set_budget_policy("pause")
|
||||
|
||||
with patch("strix.core.hooks.get_global_report_state", return_value=ledger):
|
||||
for cost in (7.0, 8.5, 9.5, 9.99):
|
||||
ledger.cost = cost
|
||||
for parent_id in (None, "root"):
|
||||
items: list[Any] = []
|
||||
await hooks.on_llm_start(
|
||||
_ctx(coordinator, parent_id=parent_id), MagicMock(), None, items
|
||||
)
|
||||
assert items == []
|
||||
await hooks.on_llm_end(
|
||||
_ctx(coordinator, parent_id=parent_id), MagicMock(), MagicMock()
|
||||
)
|
||||
|
||||
ledger.cost = 10.0
|
||||
await hooks.on_llm_end(_ctx(coordinator, parent_id="root"), MagicMock(), MagicMock())
|
||||
with pytest.raises(BudgetPausedError) as at_limit:
|
||||
await hooks.on_llm_start(_ctx(coordinator), MagicMock(), None, [])
|
||||
assert at_limit.value.resume_epoch == coordinator.resume_epoch
|
||||
|
||||
ledger.cost = 13.7
|
||||
with pytest.raises(BudgetPausedError):
|
||||
await hooks.on_llm_start(_ctx(coordinator, parent_id="root"), MagicMock(), None, [])
|
||||
|
||||
hooks.set_max_budget_usd(20.0)
|
||||
items = []
|
||||
await hooks.on_llm_start(_ctx(coordinator), MagicMock(), None, items)
|
||||
assert items == []
|
||||
|
||||
await coordinator.pause_budget()
|
||||
with pytest.raises(BudgetPausedError, match="paused"):
|
||||
await hooks.on_llm_start(_ctx(coordinator), MagicMock(), None, [])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pause_hooks_do_not_count_a_parked_turn() -> None:
|
||||
ledger = _FakeLedger()
|
||||
ledger.cost = 10.0
|
||||
hooks = ReportUsageHooks(model="m", max_budget_usd=10.0, budget_policy="pause")
|
||||
coordinator = AgentCoordinator()
|
||||
ctx = _ctx(coordinator)
|
||||
with patch("strix.core.hooks.get_global_report_state", return_value=ledger):
|
||||
with pytest.raises(BudgetPausedError):
|
||||
await hooks.on_llm_start(ctx, MagicMock(), None, [])
|
||||
assert "llm_turn" not in ctx.context
|
||||
hooks.set_max_budget_usd(11.0)
|
||||
await hooks.on_llm_start(ctx, MagicMock(), None, [])
|
||||
assert ctx.context["llm_turn"] == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stop_policy_is_unchanged() -> None:
|
||||
ledger = _FakeLedger()
|
||||
hooks = ReportUsageHooks(model="m", max_budget_usd=10.0)
|
||||
assert hooks.budget_policy == "stop"
|
||||
|
||||
with patch("strix.core.hooks.get_global_report_state", return_value=ledger):
|
||||
ledger.cost = 7.0
|
||||
items: list[Any] = []
|
||||
await hooks.on_llm_start(_ctx(None), MagicMock(), None, items)
|
||||
assert len(items) == 1
|
||||
assert "Scan cost budget" in str(items[0])
|
||||
|
||||
ledger.cost = 10.0
|
||||
with pytest.raises(BudgetExceededError):
|
||||
await hooks.on_llm_end(_ctx(None), MagicMock(), MagicMock())
|
||||
|
||||
assert recomputed_budget_flags(10.0, 10.0, interactive=False, budget_policy="stop") == (
|
||||
True,
|
||||
True,
|
||||
)
|
||||
assert recomputed_budget_flags(10.0, 10.0, interactive=False, budget_policy="pause") == (
|
||||
False,
|
||||
False,
|
||||
)
|
||||
|
||||
|
||||
def test_budget_policy_is_validated() -> None:
|
||||
with pytest.raises(ValueError, match="budget_policy"):
|
||||
ReportUsageHooks(model="m", budget_policy="later") # type: ignore[arg-type]
|
||||
@@ -430,10 +430,3 @@ def test_user_headers_override_openrouter_attribution() -> None:
|
||||
assert headers["X-Title"] == "Custom"
|
||||
assert headers["X-Tenant"] == "acme"
|
||||
assert headers["HTTP-Referer"] == "https://strix.ai"
|
||||
|
||||
|
||||
def test_reasoning_effort_sent_as_configured_and_none_omitted() -> None:
|
||||
assert make_model_settings("none", model_name="gpt-5.6-sol").reasoning is None
|
||||
settings = make_model_settings("high", model_name="gpt-daybreak-blue-latest")
|
||||
assert settings.reasoning is not None
|
||||
assert settings.reasoning.effort == "high"
|
||||
+84
-67
@@ -1,4 +1,4 @@
|
||||
"""Tests for LLM model configuration helpers."""
|
||||
"""Tests for LLM model recommendation helpers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -10,14 +10,14 @@ from agents.models import _openai_shared
|
||||
from agents.models.openai_chatcompletions import OpenAIChatCompletionsModel
|
||||
from agents.models.openai_responses import OpenAIResponsesModel
|
||||
|
||||
from strix.config import models
|
||||
from strix.config.models import (
|
||||
RECOMMENDED_MODEL_NAMES,
|
||||
StrixProvider,
|
||||
_NonStreamingModel,
|
||||
_TurnGuardModel,
|
||||
configure_sdk_model_defaults,
|
||||
is_recommended_or_frontier_model,
|
||||
request_timeout_extra_args,
|
||||
resolve_api_type,
|
||||
routes_through_litellm,
|
||||
supports_strict_tool_schemas,
|
||||
uses_chat_completions_tool_schema,
|
||||
@@ -26,6 +26,11 @@ from strix.config.settings import Settings
|
||||
from strix.llm.request_log import RequestLoggingModel
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name", RECOMMENDED_MODEL_NAMES)
|
||||
def test_recommended_models_are_accepted(model_name: str) -> None:
|
||||
assert is_recommended_or_frontier_model(model_name)
|
||||
|
||||
|
||||
def test_request_timeout_extra_args_positive() -> None:
|
||||
assert request_timeout_extra_args(300) == {"timeout": 300}
|
||||
assert request_timeout_extra_args(10) == {"timeout": 10}
|
||||
@@ -43,6 +48,80 @@ def test_request_timeout_extra_args_disabled(value: float | None) -> None:
|
||||
assert request_timeout_extra_args(value) is None
|
||||
|
||||
|
||||
def test_recommended_models_are_matched_case_insensitively() -> None:
|
||||
assert is_recommended_or_frontier_model("Vertex_AI/Gemini-3-Pro-Preview")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[
|
||||
"gpt-5.5",
|
||||
"chatgpt/gpt-5.4",
|
||||
"litellm/openai/gpt-5.4-pro",
|
||||
"azure_ai/gpt-5.5-pro",
|
||||
"bedrock_mantle/openai.gpt-5.5",
|
||||
"anthropic/claude-opus-5",
|
||||
"anthropic/claude-opus-4-8",
|
||||
"anthropic.claude-opus-4-8",
|
||||
"anthropic/claude-opus-4-7",
|
||||
"anthropic/claude-fable-5",
|
||||
"anthropic/claude-sonnet-5",
|
||||
"vertex_ai/claude-sonnet-5@default",
|
||||
"vertex_ai/claude-sonnet-4-6@default",
|
||||
"any-llm/anthropic/claude-sonnet-4-6",
|
||||
"vertex_ai/gemini-3.1-pro-preview",
|
||||
"openrouter/google/gemini-3.1-pro-preview",
|
||||
"deepseek/deepseek-v4-pro",
|
||||
"deepseek/deepseek-r1-0528",
|
||||
"deepseek/deepseek-reasoner",
|
||||
"dashscope/qwen3-max-2026-01-23",
|
||||
"qwen3.7-max",
|
||||
"dashscope/qwen3.8-max",
|
||||
"moonshot/kimi-k2.6",
|
||||
"kimi-k2.7-code",
|
||||
"moonshot/kimi-k3",
|
||||
"anthropic/claude-fable-5-1",
|
||||
"vertex_ai/claude-fable-5-1@default",
|
||||
"gemini/gemini-3.7-flash",
|
||||
"glm-5.3",
|
||||
"zai/glm-5.3-flash",
|
||||
"openrouter/z-ai/glm-5.3",
|
||||
"novita/zai-org/glm-5.2",
|
||||
"openai/glm-5.3",
|
||||
"openai/zai-org/glm-5.3",
|
||||
"hosted_vllm/glm-5.3",
|
||||
"openai/claude-opus-4-8",
|
||||
"openai/deepseek-v4-pro",
|
||||
"custom-ollama/gpt-5-mini-local",
|
||||
"custom-provider/claude-opus-4-local",
|
||||
"custom-provider/glm-5.3-local",
|
||||
],
|
||||
)
|
||||
def test_frontier_model_families_are_accepted(model_name: str) -> None:
|
||||
assert is_recommended_or_frontier_model(model_name)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[
|
||||
"",
|
||||
"openai/gpt-4.1",
|
||||
"anthropic/claude-3-5-sonnet-latest",
|
||||
"ollama/llama3.1",
|
||||
"deepseek/deepseek-chat",
|
||||
"xai/grok-4.5",
|
||||
"openrouter/x-ai/grok-4",
|
||||
"mistral/mistral-medium-3-5",
|
||||
"mistral/magistral-medium-latest",
|
||||
"zai/glm-4.7",
|
||||
"openai/glm-4.7",
|
||||
"openrouter/z-ai/glm-5",
|
||||
],
|
||||
)
|
||||
def test_non_frontier_models_are_rejected(model_name: str) -> None:
|
||||
assert not is_recommended_or_frontier_model(model_name)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[
|
||||
@@ -113,7 +192,7 @@ def test_api_type_override_settings(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
@pytest.mark.parametrize(
|
||||
("api_type", "expected"),
|
||||
[
|
||||
(None, OpenAIResponsesModel),
|
||||
(None, OpenAIChatCompletionsModel),
|
||||
("chat_completions", OpenAIChatCompletionsModel),
|
||||
("responses", OpenAIResponsesModel),
|
||||
],
|
||||
@@ -121,7 +200,7 @@ def test_api_type_override_settings(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
def test_api_type_overrides_the_api_base_route(
|
||||
monkeypatch: pytest.MonkeyPatch, api_type: str | None, expected: type
|
||||
) -> None:
|
||||
"""gpt-5 is catalogued on /v1/responses, so a base URL alone changes nothing."""
|
||||
"""``LLM_API_BASE`` defaults to chat completions. ``STRIX_API_TYPE`` must win."""
|
||||
monkeypatch.setattr(_openai_shared, "_use_responses_by_default", True)
|
||||
monkeypatch.setattr(_openai_shared, "_default_openai_client", None)
|
||||
monkeypatch.setattr(_openai_shared, "_default_openai_key", None)
|
||||
@@ -140,65 +219,3 @@ def test_api_type_overrides_the_api_base_route(
|
||||
while isinstance(model, _NonStreamingModel | _TurnGuardModel | RequestLoggingModel):
|
||||
model = model._inner
|
||||
assert isinstance(model, expected)
|
||||
|
||||
|
||||
def _settings(monkeypatch: pytest.MonkeyPatch, model: str, api_base: str | None) -> Settings:
|
||||
monkeypatch.setenv("STRIX_LLM", model)
|
||||
monkeypatch.delenv("STRIX_API_TYPE", raising=False)
|
||||
for name in ("LLM_API_BASE", "OPENAI_API_BASE", "OPENAI_BASE_URL"):
|
||||
monkeypatch.delenv(name, raising=False)
|
||||
if api_base is not None:
|
||||
monkeypatch.setenv("LLM_API_BASE", api_base)
|
||||
return Settings()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"api_base", [None, "https://api.openai.com/v1", "https://gateway.example/v1"]
|
||||
)
|
||||
def test_resolve_api_type_follows_the_catalog_not_the_base_url(
|
||||
monkeypatch: pytest.MonkeyPatch, api_base: str | None
|
||||
) -> None:
|
||||
"""Responses when LiteLLM lists /v1/responses for the model, chat completions otherwise."""
|
||||
settings = _settings(monkeypatch, "gpt-5", api_base)
|
||||
for model in ("gpt-5", "gpt-5.6-sol", "openai/gpt-5.4", "gpt-daybreak-blue-latest"):
|
||||
assert resolve_api_type(model, settings) == "responses", model
|
||||
assert uses_chat_completions_tool_schema(model, settings) is False, model
|
||||
for model in ("gpt-4o", "my-private-model"):
|
||||
assert resolve_api_type(model, settings) == "chat_completions", model
|
||||
assert uses_chat_completions_tool_schema(model, settings) is True, model
|
||||
|
||||
|
||||
def test_resolve_api_type_explicit_override_wins(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
_settings(monkeypatch, "gpt-daybreak-blue-latest", "https://gateway.example/v1")
|
||||
monkeypatch.setenv("STRIX_API_TYPE", "chat_completions")
|
||||
assert resolve_api_type("gpt-daybreak-blue-latest", Settings()) == "chat_completions"
|
||||
monkeypatch.setenv("STRIX_API_TYPE", "Responses")
|
||||
assert resolve_api_type("gpt-5", Settings()) == "responses"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("raw", "expected"),
|
||||
[("None", "none"), ("HIGH", "high"), (" xhigh ", "xhigh"), ("Max", "max")],
|
||||
)
|
||||
def test_reasoning_effort_is_case_insensitive(
|
||||
monkeypatch: pytest.MonkeyPatch, raw: str, expected: str
|
||||
) -> None:
|
||||
monkeypatch.setenv("STRIX_REASONING_EFFORT", raw)
|
||||
monkeypatch.setenv("STRIX_DEDUPE_REASONING_EFFORT", raw)
|
||||
settings = Settings()
|
||||
assert settings.llm.reasoning_effort == expected
|
||||
assert settings.dedupe.reasoning_effort == expected
|
||||
|
||||
|
||||
def test_configure_sdk_api_route_follows_the_given_model(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""A ``model=`` override picks its own route, not ``STRIX_LLM``'s."""
|
||||
routes: list[str] = []
|
||||
monkeypatch.setattr(models, "set_default_openai_api", routes.append)
|
||||
settings = _settings(monkeypatch, "gpt-5", "https://gateway.example/v1")
|
||||
|
||||
models.configure_sdk_api_route("gpt-5", settings)
|
||||
models.configure_sdk_api_route("my-private-model", settings)
|
||||
|
||||
assert routes == ["responses", "chat_completions"]
|
||||
@@ -1,45 +0,0 @@
|
||||
"""Every parameter of every scan-agent tool carries a description in its JSON schema.
|
||||
|
||||
The model only sees a parameter's docstring text through the schema. An ``Args:``
|
||||
section that the docstring parser fails to recognise (for example because an
|
||||
earlier line opens a markdown code fence it never closes) silently drops every
|
||||
parameter description, so optional-but-conditionally-required fields such as
|
||||
``confidence_rationale`` become anonymous nullable strings the model never fills.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from agents.tool import FunctionTool
|
||||
|
||||
from strix.agents import factory
|
||||
from strix.tools.agents_graph.tools import agent_finish
|
||||
from strix.tools.finish.tool import finish_scan
|
||||
from strix.tools.reporting.tool import create_vulnerability_report
|
||||
from strix.tools.respond.tool import respond_to_user
|
||||
|
||||
|
||||
_SCAN_AGENT_TOOLS = [
|
||||
tool
|
||||
for tool in (*factory._BASE_TOOLS, finish_scan, agent_finish, respond_to_user)
|
||||
if isinstance(tool, FunctionTool)
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("tool", _SCAN_AGENT_TOOLS, ids=lambda tool: tool.name)
|
||||
def test_every_parameter_has_a_description(tool: FunctionTool) -> None:
|
||||
properties = tool.params_json_schema.get("properties", {})
|
||||
missing = sorted(
|
||||
name
|
||||
for name, schema in properties.items()
|
||||
if not (isinstance(schema.get("description"), str) and schema["description"].strip())
|
||||
)
|
||||
assert not missing, f"{tool.name}: parameters without a description: {missing}"
|
||||
|
||||
|
||||
def test_create_vulnerability_report_explains_confidence_rationale() -> None:
|
||||
schema = create_vulnerability_report.params_json_schema["properties"]
|
||||
assert "confidence" in create_vulnerability_report.params_json_schema["required"]
|
||||
rationale = schema["confidence_rationale"]["description"]
|
||||
assert "Required when" in rationale
|
||||
assert "high" in rationale
|
||||
@@ -129,6 +129,16 @@ async def test_large_target_list_reports_truncated_snapshot_count() -> None:
|
||||
assert len(snapshot["targets"]) == 16
|
||||
|
||||
|
||||
def test_state_populates_model_warning_for_non_frontier_model() -> None:
|
||||
os.environ["STRIX_LLM"] = "openai/gpt-3.5-turbo"
|
||||
loader._cached = None
|
||||
|
||||
warning = TuiController(args()).snapshot()["model_warning"]
|
||||
|
||||
assert "openai/gpt-3.5-turbo" in warning
|
||||
assert "not a recommended frontier model" in warning
|
||||
|
||||
|
||||
def test_setup_restores_prepared_cli_targets() -> None:
|
||||
setup_args = args()
|
||||
setup_args.targets_info = [
|
||||
|
||||
Reference in new issue
Block a user