Compare commits
31 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 4c9f56fcd5 | |||
| 8ca83e4c16 | |||
| 94c361cbb6 | |||
| 7141ccff62 | |||
| 962d4459d9 | |||
| 11e5d1c2b3 | |||
| cc23eeb65d | |||
| 7217abfe23 | |||
| f7e3af49bd | |||
| 6202131028 | |||
| 45409cef0d | |||
| 250fe2cf3e | |||
| 6c99829325 | |||
| 1c9ab993bb | |||
| 04eb03febe | |||
| ac0fef2ed7 | |||
| dcf3155a9a | |||
| 36b374bd1b | |||
| 1a329e8972 | |||
| 143b9e7040 | |||
| 3665a7899f | |||
| 232711be8c | |||
| 712c64f630 | |||
| dee2a03d07 | |||
| 1473fc7336 | |||
| dd1f816f7c | |||
| 9ab70c6d61 | |||
| 13046cc74a | |||
| 1aad460f6e | |||
| d0321510d2 | |||
| 3bd9d56814 |
@@ -11,14 +11,16 @@ repos:
|
||||
|
||||
# MyPy for static type checking
|
||||
- repo: https://github.com/pre-commit/mirrors-mypy
|
||||
rev: v1.16.0
|
||||
rev: v1.19.1
|
||||
hooks:
|
||||
- id: mypy
|
||||
additional_dependencies: [
|
||||
types-requests,
|
||||
types-python-dateutil,
|
||||
types-Pygments,
|
||||
pydantic,
|
||||
fastapi,
|
||||
pytest,
|
||||
"openai-agents[litellm]==0.14.6",
|
||||
]
|
||||
args: [--install-types, --non-interactive]
|
||||
@@ -46,7 +48,7 @@ repos:
|
||||
|
||||
# Additional Python code quality checks
|
||||
- repo: https://github.com/asottile/pyupgrade
|
||||
rev: v3.20.0
|
||||
rev: v3.21.2
|
||||
hooks:
|
||||
- id: pyupgrade
|
||||
args: [--py312-plus]
|
||||
|
||||
@@ -79,6 +79,10 @@ When remote vars are set, Strix dual-writes telemetry to both local JSONL and th
|
||||
Runtime backend for the sandbox environment.
|
||||
</ParamField>
|
||||
|
||||
<ParamField path="STRIX_MAX_LOCAL_COPY_MB" default="1024" type="integer">
|
||||
Maximum size (in MB) of a local directory target that Strix will copy into the sandbox file-by-file. Larger targets exit early with a suggestion to use `--mount` instead. Set to `0` to disable the check.
|
||||
</ParamField>
|
||||
|
||||
## Sandbox Configuration
|
||||
|
||||
<ParamField path="STRIX_SANDBOX_EXECUTION_TIMEOUT" default="120" type="integer">
|
||||
|
||||
@@ -15,6 +15,20 @@ strix --target <target> [options]
|
||||
Target to test. Accepts URLs, repositories, local directories, domains, or IP addresses. Can be specified multiple times.
|
||||
</ParamField>
|
||||
|
||||
<ParamField path="--mount" type="string">
|
||||
Bind-mount a local directory into the sandbox (read-only) instead of copying it in file-by-file. Use this for large repositories that are too big to stream into the container. Can be specified multiple times.
|
||||
|
||||
Strix copies local `--target` directories into the sandbox one file at a time, which stalls on very large trees. When a local target exceeds the copy limit (see `STRIX_MAX_LOCAL_COPY_MB`, default 1024 MB) Strix exits early and asks you to re-run with `--mount`.
|
||||
|
||||
<Note>
|
||||
The mount is read-only to protect your source from accidental modification. This is not a hard security boundary: a root process inside the container can remount it writable, so treat `--mount` as "scan my own code", not as isolation from untrusted code.
|
||||
</Note>
|
||||
|
||||
<Note>
|
||||
The size pre-flight only covers local directory targets. Remote repositories (cloned at scan time) are not size-checked.
|
||||
</Note>
|
||||
</ParamField>
|
||||
|
||||
<ParamField path="--instruction" type="string">
|
||||
Custom instructions for the scan. Use for credentials, focus areas, or specific testing approaches.
|
||||
</ParamField>
|
||||
@@ -43,6 +57,24 @@ strix --target <target> [options]
|
||||
Path to a custom config file (JSON) to use instead of `~/.strix/cli-config.json`.
|
||||
</ParamField>
|
||||
|
||||
<ParamField path="--max-budget-usd" type="number">
|
||||
Maximum LLM spend in USD for the whole scan, counted cumulatively across the
|
||||
root agent and every child agent. The budget is checked after each model
|
||||
response; once the running cost reaches the threshold, the scan stops cleanly
|
||||
with a `stopped` status (not a failure) and the sandbox is torn down.
|
||||
|
||||
Must be greater than `0`. Omit the flag for no limit.
|
||||
|
||||
**Limitations**
|
||||
|
||||
- The check fires *after* a response is returned, so the final spend can
|
||||
slightly overshoot the limit by any calls already in flight when the
|
||||
threshold is crossed (most relevant with several child agents running
|
||||
concurrently).
|
||||
- Cost is a best-effort estimate derived from token usage and model pricing;
|
||||
providers that do not expose priced usage may under-count.
|
||||
</ParamField>
|
||||
|
||||
## Examples
|
||||
|
||||
```bash
|
||||
@@ -63,6 +95,9 @@ strix -n --target ./ --scan-mode quick --scope-mode diff --diff-base origin/main
|
||||
|
||||
# Multi-target white-box testing
|
||||
strix -t https://github.com/org/app -t https://staging.example.com
|
||||
|
||||
# Large local repository — bind-mount instead of copying it in
|
||||
strix --mount ./huge-monorepo
|
||||
```
|
||||
|
||||
## Exit Codes
|
||||
|
||||
+12
-1
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "strix-agent"
|
||||
version = "1.0.0"
|
||||
version = "1.0.4"
|
||||
description = "Open-source AI Hackers for your apps"
|
||||
readme = "README.md"
|
||||
license = "Apache-2.0"
|
||||
@@ -55,8 +55,13 @@ dev = [
|
||||
"bandit>=1.8.3",
|
||||
"pre-commit>=4.2.0",
|
||||
"pyinstaller>=6.17.0; python_version >= '3.12' and python_version < '3.15'",
|
||||
"pytest>=8.3",
|
||||
"pytest-asyncio>=0.24",
|
||||
]
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
asyncio_mode = "auto"
|
||||
|
||||
[build-system]
|
||||
requires = ["hatchling"]
|
||||
build-backend = "hatchling.build"
|
||||
@@ -104,6 +109,10 @@ module = [
|
||||
ignore_missing_imports = true
|
||||
disable_error_code = ["import-untyped"]
|
||||
|
||||
[[tool.mypy.overrides]]
|
||||
module = ["tests.*"]
|
||||
disallow_untyped_decorators = false
|
||||
|
||||
# ============================================================================
|
||||
# Ruff Configuration (Fast Python Linter & Formatter)
|
||||
# ============================================================================
|
||||
@@ -219,6 +228,8 @@ ignore = [
|
||||
# ReportState carries scan artifact/report fields and
|
||||
# a runtime ``Callable`` annotation on ``vulnerability_found_callback``.
|
||||
"strix/report/state.py" = ["TC003", "PLR0912", "PLR0915", "E501", "PERF401"]
|
||||
"strix/report/usage.py" = ["PLC0415"]
|
||||
"strix/config/models.py" = ["PLC0415"]
|
||||
# Interface utility branches per scope-mode / target-type combination;
|
||||
# splitting would obscure the decision tree without simplifying it.
|
||||
"strix/interface/utils.py" = ["PLR0912", "BLE001", "PLC0415"]
|
||||
|
||||
+2
-3
@@ -32,6 +32,8 @@ datas += collect_data_files('tiktoken_ext')
|
||||
|
||||
datas += collect_data_files('litellm')
|
||||
|
||||
datas += collect_data_files('agents', includes=['**/*.md', '**/*.jinja', '**/*.json'])
|
||||
|
||||
hiddenimports = [
|
||||
# Core dependencies
|
||||
'litellm',
|
||||
@@ -188,9 +190,6 @@ excludes = [
|
||||
'pyte',
|
||||
'openhands_aci',
|
||||
'openhands-aci',
|
||||
'gql',
|
||||
'fastapi',
|
||||
'uvicorn',
|
||||
'numpydoc',
|
||||
|
||||
# Google Cloud / Vertex AI
|
||||
|
||||
@@ -395,7 +395,6 @@ def build_strix_agent(
|
||||
instructions=instructions,
|
||||
tools=tools,
|
||||
tool_use_behavior=_finish_tool_use_behavior,
|
||||
reset_tool_choice=interactive,
|
||||
model=None,
|
||||
capabilities=[
|
||||
Filesystem(
|
||||
|
||||
@@ -43,6 +43,7 @@ AUTONOMOUS BEHAVIOR:
|
||||
- NEVER send an empty or blank message. If you have no content to output or need to wait (for user input, subagent results, or any other reason), you MUST call the wait_for_message tool (or another appropriate tool) instead of emitting an empty response.
|
||||
- If there is nothing to execute and no user query to answer any more: do NOT send filler/repetitive text — either call wait_for_message or finish your work (subagents: agent_finish; root: finish_scan)
|
||||
- While the agent loop is running, almost every output MUST be a tool call. Do NOT send plain text messages; act via tools. If idle, use wait_for_message; when done, use agent_finish (subagents) or finish_scan (root)
|
||||
- A text-only turn — even one — IMMEDIATELY ends the scan/run with no report written. The lifecycle tools (``finish_scan`` for root, ``agent_finish`` for subagents) are the ONLY valid way to terminate. If you find yourself wanting to say "Done!" or "Scan complete" without a tool call, call the lifecycle tool instead — the report and termination signal both flow through it.
|
||||
{% endif %}
|
||||
</communication_rules>
|
||||
|
||||
|
||||
+204
-32
@@ -5,7 +5,8 @@ from __future__ import annotations
|
||||
import os
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from agents import set_default_openai_api, set_default_openai_key
|
||||
from agents import set_default_openai_api, set_default_openai_key, set_tracing_disabled
|
||||
from agents.models.multi_provider import MultiProvider
|
||||
from agents.retry import (
|
||||
ModelRetryBackoffSettings,
|
||||
ModelRetrySettings,
|
||||
@@ -14,10 +15,33 @@ from agents.retry import (
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from agents.models.interface import ModelProvider
|
||||
|
||||
from strix.config.settings import Settings
|
||||
|
||||
|
||||
_SDK_PREFIXES = {"any-llm", "litellm", "openai"}
|
||||
class StrixProvider(MultiProvider):
|
||||
"""Route any non-OpenAI prefix through LiteLLM with the prefix preserved,
|
||||
so users type ``deepseek/deepseek-chat`` rather than
|
||||
``litellm/deepseek/deepseek-chat``.
|
||||
"""
|
||||
|
||||
def _resolve_prefixed_model(
|
||||
self,
|
||||
*,
|
||||
original_model_name: str,
|
||||
prefix: str,
|
||||
stripped_model_name: str | None,
|
||||
) -> tuple[ModelProvider, str | None]:
|
||||
if prefix in {"openai", "litellm", "any-llm"}:
|
||||
return super()._resolve_prefixed_model(
|
||||
original_model_name=original_model_name,
|
||||
prefix=prefix,
|
||||
stripped_model_name=stripped_model_name,
|
||||
)
|
||||
if prefix == "ollama" and stripped_model_name:
|
||||
return self._get_fallback_provider("litellm"), f"ollama_chat/{stripped_model_name}"
|
||||
return self._get_fallback_provider("litellm"), original_model_name
|
||||
|
||||
|
||||
DEFAULT_MODEL_RETRY = ModelRetrySettings(
|
||||
@@ -35,19 +59,52 @@ DEFAULT_MODEL_RETRY = ModelRetrySettings(
|
||||
),
|
||||
)
|
||||
|
||||
RECOMMENDED_MODEL_NAMES = (
|
||||
"openai/gpt-5.5",
|
||||
"openai/gpt-5.5-pro",
|
||||
"openai/gpt-5.4",
|
||||
"openai/gpt-5.4-pro",
|
||||
"openai/gpt-5.3-codex",
|
||||
"anthropic/claude-opus-4-8",
|
||||
"anthropic/claude-sonnet-4-6",
|
||||
"vertex_ai/gemini-3.1-pro-preview",
|
||||
"gemini/gemini-3.1-pro-preview",
|
||||
"xai/grok-4.3",
|
||||
"deepseek/deepseek-v4-pro",
|
||||
"deepseek/deepseek-reasoner",
|
||||
"dashscope/qwen3-max-2026-01-23",
|
||||
"moonshot/kimi-k2.7-code",
|
||||
"moonshot/kimi-k2.6",
|
||||
"mistral/mistral-medium-3-5",
|
||||
"mistral/magistral-medium-latest",
|
||||
)
|
||||
|
||||
_RECOMMENDED_MODEL_NAME_SET = frozenset(name.lower() for name in RECOMMENDED_MODEL_NAMES)
|
||||
|
||||
FRONTIER_MODEL_FAMILIES = (
|
||||
(("azure", "azure_ai", "bedrock_mantle", "openai"), ("gpt-5",)),
|
||||
(
|
||||
("anthropic", "azure_ai", "bedrock", "claude", "databricks", "snowflake", "vertex_ai"),
|
||||
("claude-opus-4", "claude-sonnet-4"),
|
||||
),
|
||||
(("google", "gemini", "vertex_ai"), ("gemini-3",)),
|
||||
(("xai", "x-ai"), ("grok-4",)),
|
||||
(("deepseek",), ("deepseek-v4", "deepseek-r1", "deepseek-reasoner")),
|
||||
(("alibaba", "dashscope", "qwen"), ("qwen3.7", "qwen3.5", "qwen3-max")),
|
||||
(("moonshot", "moonshotai", "kimi"), ("kimi-k2.7", "kimi-k2.6", "kimi-k2.5")),
|
||||
(("mistral", "mistralai"), ("mistral-medium-3-5", "magistral-medium")),
|
||||
)
|
||||
|
||||
|
||||
def configure_sdk_model_defaults(settings: Settings) -> None:
|
||||
"""Apply Strix config to SDK-native defaults.
|
||||
|
||||
OpenAI-compatible base URLs are handled by the SDK OpenAI provider.
|
||||
Non-OpenAI providers should use the SDK's native ``litellm/`` or
|
||||
``any-llm/`` routing, produced by :func:`normalize_model_name`.
|
||||
"""
|
||||
"""Apply Strix config to SDK-native defaults."""
|
||||
llm = settings.llm
|
||||
set_tracing_disabled(True)
|
||||
_configure_litellm_compatibility()
|
||||
if llm.api_key:
|
||||
set_default_openai_key(llm.api_key, use_for_tracing=False)
|
||||
_configure_litellm_default("api_key", llm.api_key)
|
||||
_mirror_api_key_to_provider_env(llm.model, llm.api_key)
|
||||
if llm.api_base:
|
||||
os.environ["OPENAI_BASE_URL"] = llm.api_base
|
||||
_configure_litellm_default("api_base", llm.api_base)
|
||||
@@ -56,12 +113,50 @@ def configure_sdk_model_defaults(settings: Settings) -> None:
|
||||
set_default_openai_api("responses")
|
||||
|
||||
|
||||
def _mirror_api_key_to_provider_env(model_name: str | None, api_key: str) -> None:
|
||||
if not model_name:
|
||||
return
|
||||
import litellm
|
||||
|
||||
name = model_name.strip()
|
||||
for prefix in ("litellm/", "any-llm/"):
|
||||
if name.lower().startswith(prefix):
|
||||
name = name[len(prefix) :]
|
||||
break
|
||||
try:
|
||||
report = litellm.validate_environment(model=name.lower())
|
||||
except Exception: # noqa: BLE001
|
||||
return
|
||||
for env_key in report.get("missing_keys") or []:
|
||||
if env_key.endswith("_API_KEY"):
|
||||
os.environ.setdefault(env_key, api_key)
|
||||
|
||||
|
||||
def _configure_litellm_compatibility() -> None:
|
||||
"""Enable LiteLLM's permissive param-handling mode."""
|
||||
"""Enable LiteLLM's permissive param handling and disable its callbacks."""
|
||||
import litellm
|
||||
|
||||
litellm.drop_params = True
|
||||
litellm.modify_params = True
|
||||
litellm.turn_off_message_logging = True
|
||||
litellm.disable_streaming_logging = True
|
||||
litellm.suppress_debug_info = True
|
||||
|
||||
_register_litellm_cost_callback()
|
||||
|
||||
|
||||
def _register_litellm_cost_callback() -> None:
|
||||
import litellm
|
||||
|
||||
from strix.report.state import litellm_cost_callback
|
||||
|
||||
for bucket_name in ("success_callback", "_async_success_callback"):
|
||||
bucket = getattr(litellm, bucket_name, None)
|
||||
if not isinstance(bucket, list):
|
||||
continue
|
||||
if litellm_cost_callback in bucket:
|
||||
continue
|
||||
bucket.append(litellm_cost_callback)
|
||||
|
||||
|
||||
def _configure_litellm_default(name: str, value: str) -> None:
|
||||
@@ -71,30 +166,107 @@ def _configure_litellm_default(name: str, value: str) -> None:
|
||||
setattr(litellm, name, value)
|
||||
|
||||
|
||||
def normalize_model_name(model_name: str) -> str:
|
||||
"""Normalize friendly Strix model names to SDK-native model ids."""
|
||||
model = model_name.strip()
|
||||
if not model:
|
||||
return model
|
||||
|
||||
if "/" in model:
|
||||
prefix = model.split("/", 1)[0].lower()
|
||||
if prefix in _SDK_PREFIXES:
|
||||
return model
|
||||
return f"litellm/{model}"
|
||||
|
||||
lower = model.lower()
|
||||
if lower.startswith("claude"):
|
||||
return f"litellm/anthropic/{model}"
|
||||
if lower.startswith("gemini"):
|
||||
return f"litellm/gemini/{model}"
|
||||
|
||||
return model
|
||||
|
||||
|
||||
def uses_chat_completions_tool_schema(model_name: str, settings: Settings) -> bool:
|
||||
"""Return whether the resolved SDK route can only receive JSON function tools."""
|
||||
model = model_name.strip().lower()
|
||||
if model.startswith(("litellm/", "any-llm/")):
|
||||
if "/" in model and not model.startswith("openai/"):
|
||||
return True
|
||||
return bool(settings.llm.api_base)
|
||||
if settings.llm.api_base:
|
||||
return True
|
||||
return not model_supports_reasoning(model_name)
|
||||
|
||||
|
||||
def model_supports_reasoning(model_name: str) -> bool:
|
||||
import litellm
|
||||
|
||||
name = model_name.strip().lower()
|
||||
for prefix in ("litellm/", "any-llm/", "openai/"):
|
||||
if name.startswith(prefix):
|
||||
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 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
|
||||
provider_name, bare_model_name = _split_model_provider(name)
|
||||
return any(
|
||||
_matches_frontier_family(provider_name, bare_model_name, provider_markers, prefixes)
|
||||
for provider_markers, prefixes in FRONTIER_MODEL_FAMILIES
|
||||
)
|
||||
|
||||
|
||||
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 _split_model_provider(model_name: str) -> tuple[str | None, str]:
|
||||
if "/" not in model_name:
|
||||
return None, model_name
|
||||
provider_name, bare_model_name = model_name.rsplit("/", 1)
|
||||
return provider_name, bare_model_name
|
||||
|
||||
|
||||
def _matches_frontier_family(
|
||||
provider_name: str | None,
|
||||
model_name: str,
|
||||
provider_markers: tuple[str, ...],
|
||||
model_prefixes: tuple[str, ...],
|
||||
) -> bool:
|
||||
if not _matches_model_prefix(model_name, model_prefixes):
|
||||
return False
|
||||
if provider_name is None:
|
||||
return True
|
||||
return _contains_provider_marker(
|
||||
provider_name, provider_markers, split_compound_names=True
|
||||
) or _contains_provider_marker(model_name, provider_markers)
|
||||
|
||||
|
||||
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 _contains_provider_marker(
|
||||
value: str, provider_markers: tuple[str, ...], *, split_compound_names: bool = False
|
||||
) -> bool:
|
||||
parts = set(value.replace(".", "/").split("/"))
|
||||
if split_compound_names:
|
||||
for separator in ("_", "-"):
|
||||
parts.update(piece for part in tuple(parts) for piece in part.split(separator))
|
||||
return any(marker in parts for marker in provider_markers)
|
||||
|
||||
|
||||
def is_known_openai_bare_model(model_name: str) -> bool:
|
||||
import litellm
|
||||
|
||||
name = model_name.strip().lower()
|
||||
if not name or "/" in name:
|
||||
return False
|
||||
entry = litellm.model_cost.get(name)
|
||||
return bool(entry and entry.get("litellm_provider") == "openai")
|
||||
|
||||
@@ -47,6 +47,11 @@ class RuntimeSettings(BaseSettings):
|
||||
alias="STRIX_IMAGE",
|
||||
)
|
||||
backend: str = Field(default="docker", alias="STRIX_RUNTIME_BACKEND")
|
||||
# Hard cap on a local target's size before we refuse to stream it into the
|
||||
# sandbox file-by-file (the SDK copies every file individually, which stalls
|
||||
# on large repos). Above this, the user must bind-mount via ``--mount``.
|
||||
# Set to 0 (or less) to disable the pre-flight check entirely.
|
||||
max_local_copy_mb: int = Field(default=1024, alias="STRIX_MAX_LOCAL_COPY_MB")
|
||||
|
||||
|
||||
class TelemetrySettings(BaseSettings):
|
||||
|
||||
+17
-1
@@ -42,10 +42,26 @@ class AgentCoordinator:
|
||||
self.runtimes: dict[str, AgentRuntime] = {}
|
||||
self._lock = asyncio.Lock()
|
||||
self._snapshot_path: Path | None = None
|
||||
self.is_shutting_down = False
|
||||
self._budget_stopped = False
|
||||
|
||||
def set_snapshot_path(self, path: Path) -> None:
|
||||
self._snapshot_path = path
|
||||
|
||||
def mark_shutting_down(self) -> None:
|
||||
self.is_shutting_down = True
|
||||
|
||||
@property
|
||||
def budget_stopped(self) -> bool:
|
||||
return self._budget_stopped
|
||||
|
||||
async def trigger_budget_stop(self) -> None:
|
||||
"""Signal a scan-wide budget stop and wake every parked agent so it exits."""
|
||||
async with self._lock:
|
||||
self._budget_stopped = True
|
||||
for runtime in self.runtimes.values():
|
||||
runtime.wake.set()
|
||||
|
||||
async def register(
|
||||
self,
|
||||
agent_id: str,
|
||||
@@ -139,7 +155,7 @@ class AgentCoordinator:
|
||||
async def wait_for_message(self, agent_id: str) -> None:
|
||||
while True:
|
||||
async with self._lock:
|
||||
if self.pending_counts.get(agent_id, 0) > 0:
|
||||
if self._budget_stopped or self.pending_counts.get(agent_id, 0) > 0:
|
||||
return
|
||||
wake = self.runtimes.setdefault(agent_id, AgentRuntime()).wake
|
||||
wake.clear()
|
||||
|
||||
+75
-29
@@ -11,10 +11,13 @@ from typing import TYPE_CHECKING, Any, cast
|
||||
|
||||
from agents import RunConfig, Runner
|
||||
from agents.exceptions import AgentsException, MaxTurnsExceeded, UserError
|
||||
from agents.sandbox.errors import ExecTransportError
|
||||
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.sessions import open_agent_session, strip_latest_image_from_session
|
||||
from strix.core.sessions import open_agent_session, strip_all_images_from_session
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -95,6 +98,10 @@ async def run_agent_loop(
|
||||
except asyncio.CancelledError:
|
||||
return result
|
||||
|
||||
if coordinator.budget_stopped:
|
||||
await coordinator.set_status(agent_id, "stopped")
|
||||
raise BudgetExceededError("scan budget reached")
|
||||
|
||||
await coordinator.consume_pending(agent_id)
|
||||
result = await _run_cycle(
|
||||
agent,
|
||||
@@ -276,6 +283,10 @@ async def _run_noninteractive_until_lifecycle(
|
||||
invalid_final_output_limit = max(1, max_turns)
|
||||
|
||||
while True:
|
||||
if coordinator.budget_stopped:
|
||||
await coordinator.set_status(agent_id, "stopped")
|
||||
raise BudgetExceededError("scan budget reached")
|
||||
|
||||
result = await _run_cycle(
|
||||
agent,
|
||||
coordinator,
|
||||
@@ -320,7 +331,7 @@ async def _run_noninteractive_until_lifecycle(
|
||||
)
|
||||
|
||||
|
||||
async def _run_cycle( # noqa: PLR0912
|
||||
async def _run_cycle( # noqa: PLR0912, PLR0915
|
||||
agent: Any,
|
||||
coordinator: AgentCoordinator,
|
||||
agent_id: str,
|
||||
@@ -349,16 +360,43 @@ async def _run_cycle( # noqa: PLR0912
|
||||
)
|
||||
await coordinator.attach_stream(agent_id, stream)
|
||||
try:
|
||||
async for event in stream.stream_events():
|
||||
if event_sink is not None:
|
||||
try:
|
||||
event_sink(agent_id, event)
|
||||
except Exception:
|
||||
logger.exception("stream event sink failed for %s", agent_id)
|
||||
if stream.run_loop_exception is not None:
|
||||
raise stream.run_loop_exception
|
||||
try:
|
||||
async for event in stream.stream_events():
|
||||
if event_sink is not None:
|
||||
try:
|
||||
event_sink(agent_id, event)
|
||||
except Exception:
|
||||
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
|
||||
# mistaken for the LiteLLM "after shutdown" race below.
|
||||
raise
|
||||
except RuntimeError as stream_exc:
|
||||
if "after shutdown" not in str(stream_exc):
|
||||
raise
|
||||
logger.warning(
|
||||
"Ignoring LiteLLM end-of-stream shutdown race for %s",
|
||||
agent_id,
|
||||
)
|
||||
except (ExecTransportError, docker_errors.NotFound):
|
||||
if not coordinator.is_shutting_down:
|
||||
raise
|
||||
logger.warning(
|
||||
"Ignoring sandbox container error during teardown for %s",
|
||||
agent_id,
|
||||
exc_info=True,
|
||||
)
|
||||
finally:
|
||||
await coordinator.detach_stream(agent_id, stream)
|
||||
except BudgetExceededError as exc:
|
||||
logger.info(
|
||||
"agent %s reached the scan budget limit; stopping the scan: %s", agent_id, exc
|
||||
)
|
||||
await coordinator.set_status(agent_id, "stopped")
|
||||
await coordinator.trigger_budget_stop()
|
||||
raise
|
||||
except Exception as exc:
|
||||
if (
|
||||
image_strips < 3
|
||||
@@ -366,14 +404,14 @@ async def _run_cycle( # noqa: PLR0912
|
||||
and getattr(exc, "status_code", None) in _INPUT_REJECTION_CODES
|
||||
):
|
||||
try:
|
||||
stripped = await strip_latest_image_from_session(session)
|
||||
stripped = await strip_all_images_from_session(session)
|
||||
except Exception:
|
||||
logger.exception("image-strip recovery failed for %s", agent_id)
|
||||
stripped = False
|
||||
if stripped:
|
||||
image_strips += 1
|
||||
logger.info(
|
||||
"Stripped latest image from %s session after rejection; retrying (%d)",
|
||||
"Stripped images from %s session after rejection; retrying (%d)",
|
||||
agent_id,
|
||||
image_strips,
|
||||
)
|
||||
@@ -509,21 +547,29 @@ async def _start_child_runner(
|
||||
child_ctx["parent_id"] = parent_id
|
||||
child_ctx["task"] = task
|
||||
|
||||
task_handle = asyncio.create_task(
|
||||
run_agent_loop(
|
||||
agent=child_agent,
|
||||
initial_input=initial_input,
|
||||
run_config=run_config,
|
||||
context=child_ctx,
|
||||
max_turns=max_turns,
|
||||
coordinator=coordinator,
|
||||
agent_id=child_id,
|
||||
interactive=interactive,
|
||||
session=session,
|
||||
start_parked=start_parked,
|
||||
event_sink=event_sink,
|
||||
hooks=hooks,
|
||||
),
|
||||
name=f"agent-{name}-{child_id}",
|
||||
)
|
||||
async def _child_loop() -> None:
|
||||
# A budget stop is a clean scan-wide shutdown, not a child failure: the
|
||||
# child's status and parent notification are already settled in
|
||||
# ``_run_cycle``. Swallow it here so the detached task does not surface a
|
||||
# spurious "Task exception was never retrieved" warning. The root agent
|
||||
# hits the same limit on its next call and tears the scan down.
|
||||
try:
|
||||
await run_agent_loop(
|
||||
agent=child_agent,
|
||||
initial_input=initial_input,
|
||||
run_config=run_config,
|
||||
context=child_ctx,
|
||||
max_turns=max_turns,
|
||||
coordinator=coordinator,
|
||||
agent_id=child_id,
|
||||
interactive=interactive,
|
||||
session=session,
|
||||
start_parked=start_parked,
|
||||
event_sink=event_sink,
|
||||
hooks=hooks,
|
||||
)
|
||||
except BudgetExceededError:
|
||||
logger.info("child %s stopped after reaching the scan budget limit", child_id)
|
||||
|
||||
task_handle = asyncio.create_task(_child_loop(), name=f"agent-{name}-{child_id}")
|
||||
await coordinator.attach_runtime(child_id, task=task_handle)
|
||||
|
||||
+19
-1
@@ -19,11 +19,22 @@ if TYPE_CHECKING:
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BudgetExceededError(RuntimeError):
|
||||
"""Raised when the accumulated LLM cost reaches the configured budget."""
|
||||
|
||||
|
||||
class ReportUsageHooks(RunHooks[dict[str, Any]]):
|
||||
"""Persist SDK-native usage after every model response."""
|
||||
|
||||
def __init__(self, *, model: str) -> None:
|
||||
def __init__(self, *, model: str, max_budget_usd: float | None = None) -> None:
|
||||
import math
|
||||
|
||||
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
|
||||
|
||||
async def on_llm_end(
|
||||
self,
|
||||
@@ -52,3 +63,10 @@ class ReportUsageHooks(RunHooks[dict[str, Any]]):
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("failed to record SDK usage for agent %s", agent_id)
|
||||
|
||||
if self._max_budget_usd is not None:
|
||||
cost = report_state.get_total_llm_cost()
|
||||
if cost >= self._max_budget_usd:
|
||||
raise BudgetExceededError(
|
||||
f"Token budget of ${self._max_budget_usd:.2f} exceeded (spent ${cost:.4f})"
|
||||
)
|
||||
|
||||
+10
-4
@@ -8,7 +8,7 @@ from typing import TYPE_CHECKING, Any
|
||||
from agents.model_settings import ModelSettings
|
||||
from openai.types.shared import Reasoning
|
||||
|
||||
from strix.config.models import DEFAULT_MODEL_RETRY
|
||||
from strix.config.models import DEFAULT_MODEL_RETRY, model_supports_reasoning
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -44,7 +44,8 @@ def build_root_task(scan_config: dict[str, Any]) -> str:
|
||||
)
|
||||
elif ttype == "local_code":
|
||||
path = details.get("target_path", "unknown")
|
||||
sections["Local Codebases"].append(f"- {path} (available at: {workspace_path})")
|
||||
suffix = ", read-only mount" if details.get("mount") else ""
|
||||
sections["Local Codebases"].append(f"- {path} (available at: {workspace_path}{suffix})")
|
||||
elif ttype == "web_application":
|
||||
sections["URLs"].append(f"- {details.get('target_url', '')}")
|
||||
elif ttype == "ip_address":
|
||||
@@ -108,14 +109,19 @@ def build_scope_context(scan_config: dict[str, Any]) -> dict[str, Any]:
|
||||
|
||||
def make_model_settings(
|
||||
reasoning_effort: ReasoningEffort | None,
|
||||
*,
|
||||
model_name: str,
|
||||
) -> ModelSettings:
|
||||
model_settings = ModelSettings(
|
||||
parallel_tool_calls=False,
|
||||
tool_choice="required",
|
||||
retry=DEFAULT_MODEL_RETRY,
|
||||
include_usage=True,
|
||||
)
|
||||
if reasoning_effort is not None:
|
||||
if (
|
||||
reasoning_effort is not None
|
||||
and reasoning_effort != "none"
|
||||
and model_supports_reasoning(model_name)
|
||||
):
|
||||
model_settings = model_settings.resolve(
|
||||
ModelSettings(reasoning=Reasoning(effort=reasoning_effort)),
|
||||
)
|
||||
|
||||
+40
-7
@@ -15,8 +15,8 @@ from agents.sandbox import SandboxRunConfig
|
||||
from strix.agents.factory import build_strix_agent, make_child_factory
|
||||
from strix.config import load_settings
|
||||
from strix.config.models import (
|
||||
StrixProvider,
|
||||
configure_sdk_model_defaults,
|
||||
normalize_model_name,
|
||||
uses_chat_completions_tool_schema,
|
||||
)
|
||||
from strix.core.agents import AgentCoordinator
|
||||
@@ -27,7 +27,7 @@ from strix.core.execution import (
|
||||
from strix.core.execution import (
|
||||
spawn_child_agent as start_child_agent,
|
||||
)
|
||||
from strix.core.hooks import ReportUsageHooks
|
||||
from strix.core.hooks import BudgetExceededError, ReportUsageHooks
|
||||
from strix.core.inputs import (
|
||||
DEFAULT_MAX_TURNS,
|
||||
build_root_task,
|
||||
@@ -55,10 +55,11 @@ async def run_strix_scan(
|
||||
scan_config: dict[str, Any],
|
||||
scan_id: str | None = None,
|
||||
image: str,
|
||||
local_sources: list[dict[str, str]] | None = None,
|
||||
local_sources: list[dict[str, Any]] | None = None,
|
||||
coordinator: AgentCoordinator | None = None,
|
||||
interactive: bool = False,
|
||||
max_turns: int = DEFAULT_MAX_TURNS,
|
||||
max_budget_usd: float | None = None,
|
||||
model: str | None = None,
|
||||
cleanup_on_exit: bool = True,
|
||||
event_sink: StreamEventSink | None = None,
|
||||
@@ -90,7 +91,7 @@ async def run_strix_scan(
|
||||
|
||||
settings = load_settings()
|
||||
configure_sdk_model_defaults(settings)
|
||||
resolved_model = normalize_model_name(model or settings.llm.model or "")
|
||||
resolved_model = (model or settings.llm.model or "").strip()
|
||||
if not resolved_model:
|
||||
raise RuntimeError(
|
||||
"No LLM model configured. Set STRIX_LLM env or pass model= to run_strix_scan().",
|
||||
@@ -153,14 +154,18 @@ async def run_strix_scan(
|
||||
is_whitebox = any(t.get("type") == "local_code" for t in targets)
|
||||
skills = list(scan_config.get("skills") or [])
|
||||
root_task = build_root_task(scan_config)
|
||||
model_settings = make_model_settings(settings.llm.reasoning_effort)
|
||||
model_settings = make_model_settings(
|
||||
settings.llm.reasoning_effort,
|
||||
model_name=resolved_model,
|
||||
)
|
||||
run_config = RunConfig(
|
||||
model=resolved_model,
|
||||
model_provider=StrixProvider(),
|
||||
model_settings=model_settings,
|
||||
sandbox=SandboxRunConfig(client=bundle["client"], session=bundle["session"]),
|
||||
trace_include_sensitive_data=False,
|
||||
)
|
||||
hooks = ReportUsageHooks(model=resolved_model)
|
||||
hooks = ReportUsageHooks(model=resolved_model, max_budget_usd=max_budget_usd)
|
||||
|
||||
scope_context = build_scope_context(scan_config)
|
||||
|
||||
@@ -261,7 +266,7 @@ async def run_strix_scan(
|
||||
async with coordinator._lock:
|
||||
root_status = coordinator.statuses.get(root_id)
|
||||
|
||||
return await run_agent_loop(
|
||||
result = await run_agent_loop(
|
||||
agent=root_agent,
|
||||
initial_input=initial_input,
|
||||
run_config=run_config,
|
||||
@@ -275,6 +280,34 @@ async def run_strix_scan(
|
||||
event_sink=event_sink,
|
||||
hooks=hooks,
|
||||
)
|
||||
if not interactive and result is not None:
|
||||
final = getattr(result, "final_output", None)
|
||||
scan_completed = False
|
||||
if isinstance(final, str):
|
||||
try:
|
||||
parsed = json.loads(final)
|
||||
scan_completed = bool(isinstance(parsed, dict) and parsed.get("scan_completed"))
|
||||
except (ValueError, TypeError):
|
||||
scan_completed = False
|
||||
elif isinstance(final, dict):
|
||||
scan_completed = bool(final.get("scan_completed"))
|
||||
if not scan_completed:
|
||||
logger.error(
|
||||
"Scan %s ended without calling finish_scan. The agent "
|
||||
"emitted a text-only turn instead of a lifecycle tool call, "
|
||||
"so no executive report was written. Final output (first "
|
||||
"300 chars): %r",
|
||||
scan_id,
|
||||
str(final)[:300],
|
||||
)
|
||||
return result # noqa: TRY300
|
||||
except BudgetExceededError as exc:
|
||||
logger.info("Scan %s stopped: %s", scan_id, exc)
|
||||
if root_id is not None:
|
||||
await coordinator.cancel_descendants(root_id)
|
||||
with contextlib.suppress(Exception):
|
||||
await coordinator.set_status(root_id, "stopped")
|
||||
return None
|
||||
except BaseException:
|
||||
logger.exception("Strix scan %s failed", scan_id)
|
||||
if root_id is not None:
|
||||
|
||||
+34
-19
@@ -2,7 +2,8 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, cast
|
||||
import contextlib
|
||||
from typing import TYPE_CHECKING, Any, cast
|
||||
|
||||
from agents.memory import SQLiteSession
|
||||
|
||||
@@ -22,29 +23,43 @@ def open_agent_session(agent_id: str, path: Path) -> SQLiteSession:
|
||||
_IMAGE_REJECTED_TEXT = "[image rejected by the model]"
|
||||
|
||||
|
||||
async def strip_latest_image_from_session(session: Session) -> bool:
|
||||
async def strip_all_images_from_session(session: Session) -> bool:
|
||||
items = await session.get_items()
|
||||
if not items:
|
||||
return False
|
||||
latest = items[-1]
|
||||
if not isinstance(latest, dict) or latest.get("type") != "function_call_output":
|
||||
return False
|
||||
output = latest.get("output")
|
||||
if not isinstance(output, list):
|
||||
return False
|
||||
if not any(isinstance(b, dict) and b.get("type") == "input_image" for b in output):
|
||||
return False
|
||||
await session.pop_item()
|
||||
await session.add_items(
|
||||
cast(
|
||||
"list[TResponseInputItem]",
|
||||
[
|
||||
|
||||
rebuilt: list[Any] = []
|
||||
changed = False
|
||||
for item in items:
|
||||
item_dict = cast("dict[str, Any]", item) if isinstance(item, dict) else None
|
||||
if (
|
||||
item_dict is not None
|
||||
and item_dict.get("type") == "function_call_output"
|
||||
and isinstance(item_dict.get("output"), list)
|
||||
and any(
|
||||
isinstance(b, dict) and b.get("type") == "input_image" for b in item_dict["output"]
|
||||
)
|
||||
):
|
||||
rebuilt.append(
|
||||
{
|
||||
"type": "function_call_output",
|
||||
"call_id": latest.get("call_id"),
|
||||
"call_id": item_dict.get("call_id"),
|
||||
"output": [{"type": "input_text", "text": _IMAGE_REJECTED_TEXT}],
|
||||
},
|
||||
],
|
||||
),
|
||||
)
|
||||
)
|
||||
changed = True
|
||||
else:
|
||||
rebuilt.append(item)
|
||||
|
||||
if not changed:
|
||||
return False
|
||||
|
||||
rebuilt_items = cast("list[TResponseInputItem]", rebuilt)
|
||||
await session.clear_session()
|
||||
try:
|
||||
await session.add_items(rebuilt_items)
|
||||
except Exception:
|
||||
with contextlib.suppress(Exception):
|
||||
await session.add_items(rebuilt_items)
|
||||
raise
|
||||
return True
|
||||
|
||||
@@ -183,6 +183,7 @@ async def run_cli(args: Any) -> None: # noqa: PLR0915
|
||||
image=_resolve_sandbox_image(),
|
||||
local_sources=getattr(args, "local_sources", None) or [],
|
||||
interactive=bool(getattr(args, "interactive", False)),
|
||||
max_budget_usd=getattr(args, "max_budget_usd", None),
|
||||
)
|
||||
finally:
|
||||
stop_updates.set()
|
||||
|
||||
+130
-11
@@ -12,7 +12,6 @@ from pathlib import Path
|
||||
|
||||
from agents.model_settings import ModelSettings
|
||||
from agents.models.interface import ModelTracing
|
||||
from agents.models.multi_provider import MultiProvider
|
||||
from docker.errors import DockerException
|
||||
from rich.console import Console
|
||||
from rich.panel import Panel
|
||||
@@ -23,16 +22,25 @@ from strix.config import (
|
||||
load_settings,
|
||||
persist_current,
|
||||
)
|
||||
from strix.config.models import configure_sdk_model_defaults, normalize_model_name
|
||||
from strix.config.models import (
|
||||
RECOMMENDED_MODEL_NAMES,
|
||||
StrixProvider,
|
||||
configure_sdk_model_defaults,
|
||||
is_known_openai_bare_model,
|
||||
is_recommended_or_frontier_model,
|
||||
)
|
||||
from strix.core.paths import run_dir_for, runtime_state_dir
|
||||
from strix.interface.cli import run_cli
|
||||
from strix.interface.tui import run_tui
|
||||
from strix.interface.utils import (
|
||||
assign_workspace_subdirs,
|
||||
build_final_stats_text,
|
||||
build_mount_targets_info,
|
||||
check_docker_connection,
|
||||
clone_repository,
|
||||
collect_local_sources,
|
||||
dedupe_local_targets,
|
||||
find_oversized_local_targets,
|
||||
generate_run_name,
|
||||
image_exists,
|
||||
infer_target_type,
|
||||
@@ -98,7 +106,8 @@ def validate_environment() -> None:
|
||||
error_text.append("• ", style="white")
|
||||
error_text.append("STRIX_LLM", style="bold cyan")
|
||||
error_text.append(
|
||||
" - Model name to use (e.g., 'gpt-5.4' or 'claude-sonnet-4-6')\n",
|
||||
" - Model name to use (e.g., 'openai/gpt-5.4' or "
|
||||
"'anthropic/claude-opus-4-7')\n",
|
||||
style="white",
|
||||
)
|
||||
|
||||
@@ -137,7 +146,7 @@ def validate_environment() -> None:
|
||||
)
|
||||
|
||||
error_text.append("\nExample setup:\n", style="white")
|
||||
error_text.append("export STRIX_LLM='gpt-5.4'\n", style="dim white")
|
||||
error_text.append("export STRIX_LLM='openai/gpt-5.4'\n", style="dim white")
|
||||
|
||||
if missing_optional_vars:
|
||||
for var in missing_optional_vars:
|
||||
@@ -215,7 +224,65 @@ async def warm_up_llm() -> None:
|
||||
configure_sdk_model_defaults(settings)
|
||||
llm = settings.llm
|
||||
|
||||
model = MultiProvider().get_model(normalize_model_name(llm.model or ""))
|
||||
raw_model = (llm.model or "").strip()
|
||||
if (
|
||||
raw_model
|
||||
and "/" not in raw_model
|
||||
and not is_known_openai_bare_model(raw_model)
|
||||
and not llm.api_base
|
||||
):
|
||||
warn_text = Text()
|
||||
warn_text.append("UNKNOWN MODEL NAME", 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 known OpenAI model. Bare names route to OpenAI by default.\n"
|
||||
"If you meant a non-OpenAI provider, use the '",
|
||||
style="white",
|
||||
)
|
||||
warn_text.append("<provider>/<model>", style="bold cyan")
|
||||
warn_text.append(
|
||||
"' form, e.g. 'anthropic/claude-opus-4-7', 'deepseek/deepseek-v4-pro'.",
|
||||
style="white",
|
||||
)
|
||||
console.print(
|
||||
Panel(
|
||||
warn_text,
|
||||
title="[bold white]STRIX",
|
||||
title_align="left",
|
||||
border_style="yellow",
|
||||
padding=(1, 2),
|
||||
),
|
||||
)
|
||||
sys.exit(1)
|
||||
|
||||
if 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),
|
||||
),
|
||||
)
|
||||
|
||||
model = StrixProvider().get_model(raw_model)
|
||||
await asyncio.wait_for(
|
||||
model.get_response(
|
||||
system_instructions="You are a helpful assistant.",
|
||||
@@ -231,7 +298,7 @@ async def warm_up_llm() -> None:
|
||||
),
|
||||
timeout=llm.timeout,
|
||||
)
|
||||
logger.info("LLM warm-up succeeded for model %s", normalize_model_name(llm.model or ""))
|
||||
logger.info("LLM warm-up succeeded for model %s", (llm.model or "").strip())
|
||||
|
||||
except Exception as e:
|
||||
logger.exception("LLM warm-up failed")
|
||||
@@ -265,6 +332,18 @@ def get_version() -> str:
|
||||
return "unknown"
|
||||
|
||||
|
||||
def _positive_budget(value: str) -> float:
|
||||
try:
|
||||
budget = float(value)
|
||||
except ValueError as exc:
|
||||
raise argparse.ArgumentTypeError(f"invalid float value: {value!r}") from exc
|
||||
import math
|
||||
|
||||
if not math.isfinite(budget) or budget <= 0:
|
||||
raise argparse.ArgumentTypeError("must be a finite number greater than 0")
|
||||
return budget
|
||||
|
||||
|
||||
def parse_arguments() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Strix Multi-Agent Cybersecurity Penetration Testing Tool",
|
||||
@@ -281,6 +360,9 @@ Examples:
|
||||
# Local code analysis
|
||||
strix --target ./my-project
|
||||
|
||||
# Large local repository (bind-mounted read-only instead of copied)
|
||||
strix --mount ./huge-monorepo
|
||||
|
||||
# Domain penetration test
|
||||
strix --target example.com
|
||||
|
||||
@@ -316,6 +398,15 @@ Examples:
|
||||
"Can be specified multiple times for multi-target scans. "
|
||||
"Required for fresh runs; loaded from disk when ``--resume`` is set.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--mount",
|
||||
type=str,
|
||||
action="append",
|
||||
metavar="PATH",
|
||||
help="Bind-mount a local directory into the sandbox (read-only) instead of "
|
||||
"copying it file-by-file. Use this for large repositories that are too big to "
|
||||
"stream into the container. Can be specified multiple times.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--instruction",
|
||||
type=str,
|
||||
@@ -388,6 +479,13 @@ Examples:
|
||||
help="Path to a custom config file (JSON) to use instead of ~/.strix/cli-config.json",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--max-budget-usd",
|
||||
type=_positive_budget,
|
||||
default=None,
|
||||
help="Maximum LLM cost in USD (> 0). The scan stops cleanly when this limit is reached.",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--resume",
|
||||
type=str,
|
||||
@@ -419,9 +517,9 @@ Examples:
|
||||
args.user_explicit_instruction = args.instruction if args.resume else None
|
||||
|
||||
if args.resume:
|
||||
if args.target:
|
||||
if args.target or args.mount:
|
||||
parser.error(
|
||||
"Cannot combine --resume with --target. --resume picks up where "
|
||||
"Cannot combine --resume with --target/--mount. --resume picks up where "
|
||||
"the prior run left off, including the original target list."
|
||||
)
|
||||
_load_resume_state(args, parser)
|
||||
@@ -434,13 +532,13 @@ Examples:
|
||||
f"or remove --resume to start over with the same targets."
|
||||
)
|
||||
else:
|
||||
if not args.target:
|
||||
if not args.target and not args.mount:
|
||||
parser.error(
|
||||
"the following arguments are required: -t/--target "
|
||||
"the following arguments are required: -t/--target or --mount "
|
||||
"(or use --resume <run_name> to continue a prior scan)"
|
||||
)
|
||||
args.targets_info = []
|
||||
for target in args.target:
|
||||
for target in args.target or []:
|
||||
try:
|
||||
target_type, target_dict = infer_target_type(target)
|
||||
|
||||
@@ -455,9 +553,30 @@ Examples:
|
||||
except ValueError:
|
||||
parser.error(f"Invalid target '{target}'")
|
||||
|
||||
try:
|
||||
args.targets_info.extend(build_mount_targets_info(args.mount or []))
|
||||
except ValueError as e:
|
||||
parser.error(str(e))
|
||||
|
||||
args.targets_info = dedupe_local_targets(args.targets_info)
|
||||
|
||||
assign_workspace_subdirs(args.targets_info)
|
||||
rewrite_localhost_targets(args.targets_info, HOST_GATEWAY_HOSTNAME)
|
||||
|
||||
max_local_copy_mb = load_settings().runtime.max_local_copy_mb
|
||||
max_copy_bytes = max_local_copy_mb * 1024 * 1024
|
||||
oversized = find_oversized_local_targets(args.targets_info, max_copy_bytes)
|
||||
if oversized:
|
||||
details = "; ".join(
|
||||
f"{path} ({size / (1024 * 1024):.0f} MB)" for path, size in oversized
|
||||
)
|
||||
parser.error(
|
||||
f"Local target too large to stream into the sandbox: {details}. "
|
||||
f"The limit is {max_local_copy_mb} MB "
|
||||
"(set STRIX_MAX_LOCAL_COPY_MB to change it). Re-run with "
|
||||
"--mount <path> to bind-mount the directory instead of copying it."
|
||||
)
|
||||
|
||||
return args
|
||||
|
||||
|
||||
|
||||
+33
-14
@@ -31,6 +31,7 @@ from textual.widgets import Button, Label, Static, TextArea, Tree
|
||||
from textual.widgets.tree import TreeNode
|
||||
|
||||
from strix.config import load_settings
|
||||
from strix.core.hooks import BudgetExceededError
|
||||
from strix.core.runner import run_strix_scan
|
||||
from strix.interface.tui.live_view import TuiLiveView
|
||||
from strix.interface.tui.messages import send_user_message_to_agent
|
||||
@@ -82,7 +83,7 @@ class ChatTextArea(TextArea): # type: ignore[misc]
|
||||
|
||||
super()._on_key(event)
|
||||
|
||||
@on(TextArea.Changed) # type: ignore[misc]
|
||||
@on(TextArea.Changed) # type: ignore[untyped-decorator]
|
||||
def _update_height(self, _event: TextArea.Changed | None = None) -> None:
|
||||
if not self.parent:
|
||||
return
|
||||
@@ -663,9 +664,9 @@ class QuitScreen(ModalScreen): # type: ignore[misc]
|
||||
self.app.pop_screen()
|
||||
event.prevent_default()
|
||||
|
||||
def on_button_pressed(self, event: Button.Pressed) -> None:
|
||||
async def on_button_pressed(self, event: Button.Pressed) -> None:
|
||||
if event.button.id == "quit":
|
||||
self.app.action_custom_quit()
|
||||
await self.app.action_custom_quit()
|
||||
else:
|
||||
self.app.pop_screen()
|
||||
|
||||
@@ -751,6 +752,7 @@ class StrixTUIApp(App): # type: ignore[misc]
|
||||
self.report_state.cleanup()
|
||||
|
||||
def signal_handler(_signum: int, _frame: Any) -> None:
|
||||
self._fire_sandbox_cleanup()
|
||||
self.report_state.cleanup(status="interrupted")
|
||||
sys.exit(0)
|
||||
|
||||
@@ -1142,9 +1144,9 @@ class StrixTUIApp(App): # type: ignore[misc]
|
||||
return (text, Text(), False)
|
||||
|
||||
if status == "waiting":
|
||||
keymap = Text()
|
||||
keymap.append("Send message to resume", style="dim")
|
||||
return (Text(" "), keymap, False)
|
||||
text = Text()
|
||||
text.append("Send message to resume", style="dim")
|
||||
return (text, Text(), False)
|
||||
|
||||
if status == "running":
|
||||
if self._agent_has_real_activity(agent_id):
|
||||
@@ -1368,12 +1370,18 @@ class StrixTUIApp(App): # type: ignore[misc]
|
||||
local_sources=getattr(self.args, "local_sources", None) or [],
|
||||
coordinator=self.coordinator,
|
||||
interactive=True,
|
||||
max_budget_usd=getattr(self.args, "max_budget_usd", None),
|
||||
event_sink=self._capture_sdk_event,
|
||||
),
|
||||
)
|
||||
|
||||
except (KeyboardInterrupt, asyncio.CancelledError):
|
||||
logger.info("Scan interrupted by user")
|
||||
except BudgetExceededError:
|
||||
# Defensive: the runner stops the scan cleanly on budget and
|
||||
# returns, so this normally never propagates. Treat it as a
|
||||
# graceful stop, not a scan error, if it ever does.
|
||||
logger.info("Scan stopped: --max-budget-usd limit reached")
|
||||
except (ConnectionError, TimeoutError) as e:
|
||||
logging.exception("Network error during scan")
|
||||
self._scan_error = e
|
||||
@@ -1541,7 +1549,7 @@ class StrixTUIApp(App): # type: ignore[misc]
|
||||
|
||||
return AgentMessageRenderer.render_simple(content)
|
||||
|
||||
@on(Tree.NodeHighlighted) # type: ignore[misc]
|
||||
@on(Tree.NodeHighlighted) # type: ignore[untyped-decorator]
|
||||
def handle_tree_highlight(self, event: Tree.NodeHighlighted) -> None:
|
||||
if len(self.screen_stack) > 1 or self.show_splash:
|
||||
return
|
||||
@@ -1561,7 +1569,7 @@ class StrixTUIApp(App): # type: ignore[misc]
|
||||
if agent_id:
|
||||
self.selected_agent_id = agent_id
|
||||
|
||||
@on(Tree.NodeSelected) # type: ignore[misc]
|
||||
@on(Tree.NodeSelected) # type: ignore[untyped-decorator]
|
||||
def handle_tree_node_selected(self, event: Tree.NodeSelected) -> None:
|
||||
if len(self.screen_stack) > 1 or self.show_splash:
|
||||
return
|
||||
@@ -1632,9 +1640,9 @@ class StrixTUIApp(App): # type: ignore[misc]
|
||||
|
||||
self.push_screen(HelpScreen())
|
||||
|
||||
def action_request_quit(self) -> None:
|
||||
async def action_request_quit(self) -> None:
|
||||
if self.show_splash or not self.is_mounted:
|
||||
self.action_custom_quit()
|
||||
await self.action_custom_quit()
|
||||
return
|
||||
|
||||
if len(self.screen_stack) > 1:
|
||||
@@ -1643,7 +1651,7 @@ class StrixTUIApp(App): # type: ignore[misc]
|
||||
try:
|
||||
self.query_one("#main_container")
|
||||
except (ValueError, Exception):
|
||||
self.action_custom_quit()
|
||||
await self.action_custom_quit()
|
||||
return
|
||||
|
||||
self.push_screen(QuitScreen())
|
||||
@@ -1703,16 +1711,27 @@ class StrixTUIApp(App): # type: ignore[misc]
|
||||
self._scan_loop,
|
||||
)
|
||||
|
||||
def action_custom_quit(self) -> None:
|
||||
async def action_custom_quit(self) -> None:
|
||||
self._fire_sandbox_cleanup()
|
||||
|
||||
if self._scan_thread and self._scan_thread.is_alive():
|
||||
self._scan_stop_event.set()
|
||||
|
||||
self._scan_thread.join(timeout=1.0)
|
||||
|
||||
self.report_state.cleanup()
|
||||
|
||||
self.exit()
|
||||
|
||||
def _fire_sandbox_cleanup(self) -> None:
|
||||
self.coordinator.mark_shutting_down()
|
||||
loop = self._scan_loop
|
||||
if loop is None or loop.is_closed():
|
||||
return
|
||||
run_name = self.scan_config.get("run_name")
|
||||
if not run_name:
|
||||
return
|
||||
with contextlib.suppress(Exception):
|
||||
asyncio.run_coroutine_threadsafe(session_manager.cleanup(run_name), loop)
|
||||
|
||||
def _is_widget_safe(self, widget: Any) -> bool:
|
||||
try:
|
||||
_ = widget.screen
|
||||
|
||||
@@ -26,6 +26,11 @@ _EXIT_RE = re.compile(r"Process exited with code (-?\d+)")
|
||||
_SESSION_RE = re.compile(r"Process running with session ID (\d+)")
|
||||
_OUTPUT_HEADER = "\nOutput:\n"
|
||||
|
||||
_CONTROL_BYTES_TO_DROP = dict.fromkeys(
|
||||
[b for b in range(0x20) if b not in (0x09, 0x0A)] + [0x7F],
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
@cache
|
||||
def _get_style_colors() -> dict[Any, str]:
|
||||
@@ -60,14 +65,13 @@ def _parse_sdk_shell_result(result: Any) -> dict[str, Any]:
|
||||
|
||||
|
||||
def _truncate_line(line: str) -> str:
|
||||
clean_line = re.sub(r"\x1b\[[0-9;]*m", "", line)
|
||||
if len(clean_line) > MAX_LINE_LENGTH:
|
||||
if len(line) > MAX_LINE_LENGTH:
|
||||
return line[: MAX_LINE_LENGTH - 3] + "..."
|
||||
return line
|
||||
|
||||
|
||||
def _clean_output(output: str) -> str:
|
||||
cleaned = output
|
||||
cleaned = Text.from_ansi(output).plain.translate(_CONTROL_BYTES_TO_DROP)
|
||||
for pattern in STRIP_PATTERNS:
|
||||
cleaned = re.sub(pattern, "", cleaned, flags=re.MULTILINE)
|
||||
|
||||
|
||||
+121
-2
@@ -1,5 +1,6 @@
|
||||
import ipaddress
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import secrets
|
||||
@@ -23,6 +24,9 @@ from rich.text import Text
|
||||
from strix.config import load_settings
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def get_severity_color(severity: str) -> str:
|
||||
severity_colors = {
|
||||
"critical": "#dc2626",
|
||||
@@ -1185,8 +1189,8 @@ def is_whitebox_scan(targets_info: list[dict[str, Any]]) -> bool:
|
||||
return any(t.get("type") == "local_code" for t in targets_info or [])
|
||||
|
||||
|
||||
def collect_local_sources(targets_info: list[dict[str, Any]]) -> list[dict[str, str]]:
|
||||
local_sources: list[dict[str, str]] = []
|
||||
def collect_local_sources(targets_info: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
local_sources: list[dict[str, Any]] = []
|
||||
|
||||
for target_info in targets_info:
|
||||
details = target_info["details"]
|
||||
@@ -1197,6 +1201,7 @@ def collect_local_sources(targets_info: list[dict[str, Any]]) -> list[dict[str,
|
||||
{
|
||||
"source_path": details["target_path"],
|
||||
"workspace_subdir": workspace_subdir,
|
||||
"mount": bool(details.get("mount", False)),
|
||||
}
|
||||
)
|
||||
|
||||
@@ -1205,12 +1210,126 @@ def collect_local_sources(targets_info: list[dict[str, Any]]) -> list[dict[str,
|
||||
{
|
||||
"source_path": details["cloned_repo_path"],
|
||||
"workspace_subdir": workspace_subdir,
|
||||
"mount": False,
|
||||
}
|
||||
)
|
||||
|
||||
return local_sources
|
||||
|
||||
|
||||
def directory_size_bytes(path: Path) -> int:
|
||||
"""Total size in bytes of regular files under ``path`` (symlinks not followed).
|
||||
|
||||
Best-effort: files that disappear or can't be stat'd mid-walk are skipped.
|
||||
Used as a cheap (stat-only) pre-flight to estimate the cost of streaming a
|
||||
local target into the sandbox before we actually try to copy it.
|
||||
|
||||
Directories that can't be listed (e.g. permission denied) are logged and
|
||||
skipped rather than silently dropped — so an under-count is at least
|
||||
visible — but the returned total then excludes their contents.
|
||||
"""
|
||||
|
||||
def _on_walk_error(error: OSError) -> None:
|
||||
logger.warning("Could not read %s while measuring size: %s", error.filename, error)
|
||||
|
||||
total = 0
|
||||
for root, _dirs, files in os.walk(path, followlinks=False, onerror=_on_walk_error):
|
||||
for name in files:
|
||||
file_path = os.path.join(root, name) # noqa: PTH118
|
||||
try:
|
||||
if os.path.islink(file_path): # noqa: PTH114
|
||||
continue
|
||||
total += os.path.getsize(file_path) # noqa: PTH202
|
||||
except OSError:
|
||||
continue
|
||||
return total
|
||||
|
||||
|
||||
def find_oversized_local_targets(
|
||||
targets_info: list[dict[str, Any]], max_bytes: int
|
||||
) -> list[tuple[str, int]]:
|
||||
"""Return ``(path, size_bytes)`` for non-mounted local targets over ``max_bytes``.
|
||||
|
||||
Mounted targets are bind-mounted rather than copied, so their size is
|
||||
irrelevant and they are excluded. A ``max_bytes`` of zero or less disables
|
||||
the check entirely (returns no targets).
|
||||
"""
|
||||
if max_bytes <= 0:
|
||||
return []
|
||||
oversized: list[tuple[str, int]] = []
|
||||
for target in targets_info:
|
||||
if target.get("type") != "local_code":
|
||||
continue
|
||||
details = target.get("details") or {}
|
||||
if details.get("mount"):
|
||||
continue
|
||||
target_path = details.get("target_path")
|
||||
if not target_path:
|
||||
continue
|
||||
size = directory_size_bytes(Path(target_path))
|
||||
if size > max_bytes:
|
||||
oversized.append((target_path, size))
|
||||
return oversized
|
||||
|
||||
|
||||
def build_mount_targets_info(mount_paths: list[str]) -> list[dict[str, Any]]:
|
||||
"""Build ``targets_info`` entries for ``--mount`` directories.
|
||||
|
||||
Each path must be an existing local directory; it is bind-mounted into the
|
||||
sandbox (read-only) instead of being copied file-by-file. Raises
|
||||
``ValueError`` for an empty path, or one that does not exist or is not a
|
||||
directory.
|
||||
"""
|
||||
targets_info: list[dict[str, Any]] = []
|
||||
for raw in mount_paths:
|
||||
if not raw or not raw.strip():
|
||||
raise ValueError("--mount path must not be empty.")
|
||||
path = Path(raw).expanduser()
|
||||
try:
|
||||
resolved = path.resolve()
|
||||
is_dir = resolved.is_dir()
|
||||
except (OSError, RuntimeError) as e:
|
||||
raise ValueError(f"Invalid mount path '{raw}': {e!s}") from e
|
||||
if not is_dir:
|
||||
raise ValueError(
|
||||
f"Mount path '{raw}' is not an existing directory. "
|
||||
"--mount requires a path to a local directory."
|
||||
)
|
||||
targets_info.append(
|
||||
{
|
||||
"type": "local_code",
|
||||
"details": {"target_path": str(resolved), "mount": True},
|
||||
"original": str(resolved),
|
||||
}
|
||||
)
|
||||
return targets_info
|
||||
|
||||
|
||||
def dedupe_local_targets(targets_info: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
"""Collapse local_code targets that resolve to the same path.
|
||||
|
||||
When a directory is supplied both as a copied ``--target`` and via
|
||||
``--mount`` (or as duplicate values of either), keep one entry and prefer
|
||||
the bind-mounted one — so the same tree is never both streamed in and
|
||||
mounted. Order is preserved; non-local targets pass through untouched.
|
||||
"""
|
||||
result: list[dict[str, Any]] = []
|
||||
index_by_path: dict[str, int] = {}
|
||||
for target in targets_info:
|
||||
details = target.get("details") or {}
|
||||
path = details.get("target_path")
|
||||
if target.get("type") != "local_code" or not path:
|
||||
result.append(target)
|
||||
continue
|
||||
existing = index_by_path.get(path)
|
||||
if existing is None:
|
||||
index_by_path[path] = len(result)
|
||||
result.append(target)
|
||||
elif details.get("mount") and not (result[existing].get("details") or {}).get("mount"):
|
||||
result[existing] = target # bind mount supersedes the copied entry
|
||||
return result
|
||||
|
||||
|
||||
def _is_localhost_host(host: str) -> bool:
|
||||
host_lower = host.lower().strip("[]")
|
||||
|
||||
|
||||
@@ -8,14 +8,13 @@ from typing import TYPE_CHECKING, Any
|
||||
|
||||
from agents.model_settings import ModelSettings
|
||||
from agents.models.interface import ModelTracing
|
||||
from agents.models.multi_provider import MultiProvider
|
||||
from openai.types.responses import ResponseOutputMessage
|
||||
|
||||
from strix.config import load_settings
|
||||
from strix.config.models import (
|
||||
DEFAULT_MODEL_RETRY,
|
||||
StrixProvider,
|
||||
configure_sdk_model_defaults,
|
||||
normalize_model_name,
|
||||
)
|
||||
from strix.report.state import get_global_report_state
|
||||
|
||||
@@ -188,8 +187,8 @@ async def check_duplicate(
|
||||
)
|
||||
|
||||
configure_sdk_model_defaults(settings)
|
||||
resolved_model = normalize_model_name(model_name)
|
||||
model = MultiProvider().get_model(resolved_model)
|
||||
resolved_model = model_name.strip()
|
||||
model = StrixProvider().get_model(resolved_model)
|
||||
response = await model.get_response(
|
||||
system_instructions=DEDUPE_SYSTEM_PROMPT,
|
||||
input=user_msg,
|
||||
|
||||
@@ -230,9 +230,16 @@ class ReportState:
|
||||
):
|
||||
self.save_run_data()
|
||||
|
||||
def record_observed_llm_cost(self, cost: float) -> None:
|
||||
self._llm_usage.record_observed_cost(cost)
|
||||
|
||||
def get_total_llm_usage(self) -> dict[str, Any]:
|
||||
return dict(self.run_record.get("llm_usage") or self._build_llm_usage_record())
|
||||
|
||||
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 update_scan_final_fields(
|
||||
self,
|
||||
executive_summary: str,
|
||||
@@ -343,3 +350,45 @@ class ReportState:
|
||||
def _hydrate_llm_usage(self, raw_usage: Any) -> None:
|
||||
self._llm_usage.hydrate(raw_usage)
|
||||
self._sync_llm_usage_record()
|
||||
|
||||
|
||||
def litellm_cost_callback(
|
||||
kwargs: Any,
|
||||
completion_response: Any,
|
||||
_start_time: Any = None,
|
||||
_end_time: Any = None,
|
||||
) -> None:
|
||||
"""LiteLLM ``success_callback`` adapter; forwards observed cost to the active scan."""
|
||||
cost: float | None = None
|
||||
raw = kwargs.get("response_cost") if isinstance(kwargs, dict) else None
|
||||
if isinstance(raw, int | float) and raw > 0:
|
||||
cost = float(raw)
|
||||
|
||||
if cost is None:
|
||||
hidden = getattr(completion_response, "_hidden_params", None) or {}
|
||||
candidate = hidden.get("response_cost") if isinstance(hidden, dict) else None
|
||||
if isinstance(candidate, int | float) and candidate > 0:
|
||||
cost = float(candidate)
|
||||
else:
|
||||
headers = hidden.get("additional_headers") or {} if isinstance(hidden, dict) else {}
|
||||
raw = (
|
||||
headers.get("llm_provider-x-litellm-response-cost")
|
||||
if isinstance(headers, dict)
|
||||
else None
|
||||
)
|
||||
try:
|
||||
value = float(raw) if raw is not None else None
|
||||
except (TypeError, ValueError):
|
||||
value = None
|
||||
if value is not None and value > 0:
|
||||
cost = value
|
||||
|
||||
if cost is None or cost <= 0:
|
||||
return
|
||||
report_state = get_global_report_state()
|
||||
if report_state is None:
|
||||
return
|
||||
try:
|
||||
report_state.record_observed_llm_cost(cost)
|
||||
except Exception:
|
||||
logger.exception("Failed to record observed LiteLLM cost")
|
||||
|
||||
+56
-28
@@ -19,7 +19,6 @@ class LLMUsageLedger:
|
||||
self._agent_usage: dict[str, Usage] = {}
|
||||
self._agent_metadata: dict[str, dict[str, str]] = {}
|
||||
self._total_cost = 0.0
|
||||
self._agent_cost: dict[str, float] = {}
|
||||
|
||||
def record(
|
||||
self,
|
||||
@@ -33,8 +32,6 @@ class LLMUsageLedger:
|
||||
return False
|
||||
|
||||
normalized_agent_id = str(agent_id or "unknown")
|
||||
estimated_cost = _estimate_litellm_cost(usage, model)
|
||||
|
||||
self._total_usage.add(usage)
|
||||
self._agent_usage.setdefault(normalized_agent_id, Usage()).add(usage)
|
||||
|
||||
@@ -44,31 +41,42 @@ class LLMUsageLedger:
|
||||
if model:
|
||||
metadata["model"] = model
|
||||
|
||||
if estimated_cost is not None:
|
||||
self._total_cost += estimated_cost
|
||||
self._agent_cost[normalized_agent_id] = (
|
||||
self._agent_cost.get(normalized_agent_id, 0.0) + estimated_cost
|
||||
)
|
||||
if not _is_litellm_routed(model):
|
||||
estimated = _estimate_litellm_cost(usage, model)
|
||||
if estimated:
|
||||
self._total_cost += estimated
|
||||
|
||||
return True
|
||||
|
||||
def record_observed_cost(self, cost: float) -> None:
|
||||
if isinstance(cost, int | float) and cost > 0:
|
||||
self._total_cost += float(cost)
|
||||
|
||||
@property
|
||||
def total_cost(self) -> float:
|
||||
return _round_cost(self._total_cost)
|
||||
|
||||
def to_record(self) -> dict[str, Any]:
|
||||
record = serialize_usage(self._total_usage)
|
||||
record["cost"] = _round_cost(self._total_cost)
|
||||
record["cost_source"] = "litellm_estimate"
|
||||
record["agents"] = []
|
||||
|
||||
agent_tokens = {aid: _resolve_total_tokens(u) for aid, u in self._agent_usage.items()}
|
||||
total_tokens = sum(agent_tokens.values())
|
||||
for agent_id in sorted(self._agent_usage):
|
||||
usage = self._agent_usage[agent_id]
|
||||
metadata = self._agent_metadata.get(agent_id, {})
|
||||
agent_cost = (
|
||||
self._total_cost * (agent_tokens[agent_id] / total_tokens) if total_tokens else 0.0
|
||||
)
|
||||
|
||||
agent_record = serialize_usage(usage)
|
||||
agent_record.update(
|
||||
{
|
||||
"agent_id": agent_id,
|
||||
"agent_name": metadata.get("agent_name") or agent_id,
|
||||
"model": metadata.get("model"),
|
||||
"cost": _round_cost(self._agent_cost.get(agent_id, 0.0)),
|
||||
"cost_source": "litellm_estimate",
|
||||
"cost": _round_cost(agent_cost),
|
||||
}
|
||||
)
|
||||
record["agents"].append(agent_record)
|
||||
@@ -80,7 +88,6 @@ class LLMUsageLedger:
|
||||
self._agent_usage.clear()
|
||||
self._agent_metadata.clear()
|
||||
self._total_cost = 0.0
|
||||
self._agent_cost.clear()
|
||||
|
||||
if not isinstance(raw_usage, dict):
|
||||
return
|
||||
@@ -92,11 +99,8 @@ class LLMUsageLedger:
|
||||
self._total_usage = Usage()
|
||||
|
||||
self._total_cost = _float_or_zero(raw_usage.get("cost"))
|
||||
agents = raw_usage.get("agents") or []
|
||||
if not isinstance(agents, list):
|
||||
return
|
||||
|
||||
for raw_agent in agents:
|
||||
for raw_agent in raw_usage.get("agents") or []:
|
||||
if not isinstance(raw_agent, dict):
|
||||
continue
|
||||
agent_id = str(raw_agent.get("agent_id") or "").strip()
|
||||
@@ -116,7 +120,24 @@ class LLMUsageLedger:
|
||||
if isinstance(model, str) and model:
|
||||
metadata["model"] = model
|
||||
self._agent_metadata[agent_id] = metadata
|
||||
self._agent_cost[agent_id] = _float_or_zero(raw_agent.get("cost"))
|
||||
|
||||
|
||||
def _resolve_total_tokens(usage: Usage) -> int:
|
||||
total = max(0, int(usage.total_tokens or 0))
|
||||
if total > 0:
|
||||
return total
|
||||
prompt = _int_or_zero(getattr(usage, "input_tokens", 0))
|
||||
completion = _int_or_zero(getattr(usage, "output_tokens", 0))
|
||||
return prompt + completion
|
||||
|
||||
|
||||
def _is_litellm_routed(model: str | None) -> bool:
|
||||
if not model:
|
||||
return False
|
||||
name = model.strip().lower()
|
||||
if "/" not in name:
|
||||
return False
|
||||
return not name.startswith("openai/")
|
||||
|
||||
|
||||
def _usage_has_activity(usage: Usage) -> bool:
|
||||
@@ -171,18 +192,25 @@ def _estimate_litellm_entry_cost(entry: Any, model: str) -> float | None:
|
||||
if completion_details:
|
||||
usage_payload["completion_tokens_details"] = completion_details
|
||||
|
||||
try:
|
||||
from litellm import completion_cost
|
||||
from litellm import completion_cost
|
||||
|
||||
cost = completion_cost(
|
||||
completion_response={
|
||||
"model": model.split("/", 1)[-1],
|
||||
"usage": usage_payload,
|
||||
},
|
||||
model=model,
|
||||
)
|
||||
except Exception: # noqa: BLE001 - LiteLLM raises plain Exception for unknown model prices.
|
||||
logger.debug("LiteLLM cost estimate unavailable for model %s", model, exc_info=True)
|
||||
candidates = [model]
|
||||
if "/" in model:
|
||||
candidates.append(model.split("/", 1)[-1])
|
||||
|
||||
cost: Any = None
|
||||
for candidate in candidates:
|
||||
try:
|
||||
cost = completion_cost(
|
||||
completion_response={"model": candidate, "usage": usage_payload},
|
||||
model=model,
|
||||
)
|
||||
break
|
||||
except Exception: # nosec B112 # noqa: BLE001, S112
|
||||
continue
|
||||
|
||||
if cost is None:
|
||||
logger.debug("LiteLLM cost estimate unavailable for model %s", model)
|
||||
return None
|
||||
|
||||
return cost if isinstance(cost, int | float) and cost >= 0 else None
|
||||
|
||||
@@ -27,7 +27,7 @@ def read_run_record(run_dir: Path) -> dict[str, Any]:
|
||||
except (OSError, json.JSONDecodeError) as exc:
|
||||
raise RuntimeError(f"run.json at {path} is unreadable: {exc}") from exc
|
||||
if not isinstance(data, dict):
|
||||
raise RuntimeError(f"run.json at {path} is not an object")
|
||||
raise TypeError(f"run.json at {path} is not an object")
|
||||
return data
|
||||
|
||||
|
||||
|
||||
@@ -22,6 +22,7 @@ async def _docker_backend(
|
||||
image: str,
|
||||
manifest: Manifest,
|
||||
exposed_ports: tuple[int, ...],
|
||||
bind_mounts: list[dict[str, Any]] | None = None,
|
||||
) -> tuple[Any, Any]:
|
||||
"""Bring up a session backed by the local Docker daemon.
|
||||
|
||||
@@ -31,11 +32,15 @@ async def _docker_backend(
|
||||
backend don't need the docker-py library installed.
|
||||
|
||||
``session.start()`` is what materializes the manifest entries
|
||||
(LocalDir copies, mount setup, etc.) into the running container —
|
||||
the SDK's ``client.create()`` only builds the inner session object
|
||||
without applying the manifest. ``async with session:`` would call it
|
||||
too, but Strix manages session lifetime explicitly via
|
||||
(LocalDir copies and manifest-declared volume/FUSE mounts) into the
|
||||
running container — the SDK's ``client.create()`` only builds the inner
|
||||
session object without applying the manifest. ``async with session:``
|
||||
would call it too, but Strix manages session lifetime explicitly via
|
||||
``client.delete()`` so we trigger ``start()`` ourselves.
|
||||
|
||||
``bind_mounts`` are host directories (e.g. large repos passed via
|
||||
``--mount``) bind-mounted read-only; unlike manifest entries they are
|
||||
applied by Docker at container-create time, not by ``start()``.
|
||||
"""
|
||||
import docker
|
||||
from agents.sandbox.sandboxes.docker import DockerSandboxClientOptions
|
||||
@@ -43,6 +48,7 @@ async def _docker_backend(
|
||||
from strix.runtime.docker_client import StrixDockerSandboxClient
|
||||
|
||||
client = StrixDockerSandboxClient(docker.from_env())
|
||||
client.strix_bind_mounts = bind_mounts or []
|
||||
options = DockerSandboxClientOptions(image=image, exposed_ports=exposed_ports)
|
||||
session = await client.create(options=options, manifest=manifest)
|
||||
await session.start()
|
||||
|
||||
@@ -22,6 +22,7 @@ re-merging the parent body. Track upstream for an injection hook.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import logging
|
||||
import uuid
|
||||
from typing import Any
|
||||
@@ -34,7 +35,10 @@ from agents.sandbox.sandboxes.docker import (
|
||||
_manifest_requires_fuse,
|
||||
_manifest_requires_sys_admin,
|
||||
)
|
||||
from agents.sandbox.session.sandbox_session import SandboxSession
|
||||
from docker import errors as docker_errors # type: ignore[import-untyped, unused-ignore]
|
||||
from docker.models.containers import Container # type: ignore[import-untyped, unused-ignore]
|
||||
from docker.types import Mount as DockerSDKMount # type: ignore[import-untyped, unused-ignore]
|
||||
from docker.utils import parse_repository_tag # type: ignore[import-untyped, unused-ignore]
|
||||
|
||||
|
||||
@@ -42,6 +46,10 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class StrixDockerSandboxClient(DockerSandboxClient):
|
||||
# Host directories to bind-mount into the container, set by the docker
|
||||
# backend before ``create()``. Each item is ``{source, target, read_only}``.
|
||||
strix_bind_mounts: list[dict[str, Any]] | None = None
|
||||
|
||||
async def _create_container(
|
||||
self,
|
||||
image: str,
|
||||
@@ -108,6 +116,21 @@ class StrixDockerSandboxClient(DockerSandboxClient):
|
||||
extra_hosts = create_kwargs.setdefault("extra_hosts", {})
|
||||
extra_hosts["host.docker.internal"] = "host-gateway"
|
||||
|
||||
# Strix injection: host bind mounts (e.g. large repos passed via --mount)
|
||||
# that bypass the SDK's file-by-file LocalDir copy.
|
||||
bind_mounts = getattr(self, "strix_bind_mounts", ())
|
||||
if bind_mounts:
|
||||
mounts = create_kwargs.setdefault("mounts", [])
|
||||
for spec in bind_mounts:
|
||||
mounts.append(
|
||||
DockerSDKMount(
|
||||
target=spec["target"],
|
||||
source=spec["source"],
|
||||
type="bind",
|
||||
read_only=spec.get("read_only", True),
|
||||
)
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
"Creating sandbox container: image=%s caps=%s exposed_ports=%s",
|
||||
image,
|
||||
@@ -121,3 +144,10 @@ class StrixDockerSandboxClient(DockerSandboxClient):
|
||||
image,
|
||||
)
|
||||
return container
|
||||
|
||||
async def delete(self, session: SandboxSession) -> SandboxSession:
|
||||
container_id = getattr(getattr(session._inner, "state", None), "container_id", None)
|
||||
if container_id:
|
||||
with contextlib.suppress(docker_errors.NotFound, docker_errors.APIError):
|
||||
self.docker_client.containers.get(container_id).kill()
|
||||
return await super().delete(session)
|
||||
|
||||
@@ -23,30 +23,59 @@ _CONTAINER_CAIDO_PORT = 48080
|
||||
|
||||
_SESSION_CACHE: dict[str, dict[str, Any]] = {}
|
||||
|
||||
# Manifest root inside the container; entry keys hang off this path.
|
||||
_WORKSPACE_ROOT = "/workspace"
|
||||
|
||||
|
||||
def build_session_entries(
|
||||
local_sources: list[dict[str, Any]],
|
||||
) -> tuple[dict[str | Path, BaseEntry], list[dict[str, Any]]]:
|
||||
"""Split local sources into copied manifest entries and host bind mounts.
|
||||
|
||||
Sources flagged ``mount`` are bind-mounted read-only at
|
||||
``/workspace/<workspace_subdir>`` (not added to the manifest, so the SDK
|
||||
does not stream them in file-by-file). Every other source becomes a
|
||||
``LocalDir`` entry copied into the container as before.
|
||||
"""
|
||||
entries: dict[str | Path, BaseEntry] = {}
|
||||
bind_mounts: list[dict[str, Any]] = []
|
||||
for src in local_sources:
|
||||
ws_subdir = src.get("workspace_subdir") or ""
|
||||
host_path = src.get("source_path") or ""
|
||||
if not ws_subdir or not host_path:
|
||||
continue
|
||||
resolved = Path(host_path).expanduser().resolve()
|
||||
if src.get("mount"):
|
||||
bind_mounts.append(
|
||||
{
|
||||
"source": str(resolved),
|
||||
"target": f"{_WORKSPACE_ROOT}/{ws_subdir}",
|
||||
"read_only": True,
|
||||
}
|
||||
)
|
||||
else:
|
||||
entries[ws_subdir] = LocalDir(src=resolved)
|
||||
return entries, bind_mounts
|
||||
|
||||
|
||||
async def create_or_reuse(
|
||||
scan_id: str,
|
||||
*,
|
||||
image: str,
|
||||
local_sources: list[dict[str, str]],
|
||||
local_sources: list[dict[str, Any]],
|
||||
) -> dict[str, Any]:
|
||||
"""Return the existing session bundle for ``scan_id`` or create a new one.
|
||||
|
||||
Each ``local_sources`` entry mounts its host ``source_path`` at
|
||||
``/workspace/<workspace_subdir>`` inside the container.
|
||||
Each ``local_sources`` entry exposes its host ``source_path`` at
|
||||
``/workspace/<workspace_subdir>`` inside the container — copied in, or
|
||||
bind-mounted read-only when the entry is flagged ``mount``.
|
||||
"""
|
||||
cached = _SESSION_CACHE.get(scan_id)
|
||||
if cached is not None:
|
||||
logger.info("Reusing existing sandbox session for scan %s", scan_id)
|
||||
return cached
|
||||
|
||||
entries: dict[str | Path, BaseEntry] = {}
|
||||
for src in local_sources:
|
||||
ws_subdir = src.get("workspace_subdir") or ""
|
||||
host_path = src.get("source_path") or ""
|
||||
if not ws_subdir or not host_path:
|
||||
continue
|
||||
entries[ws_subdir] = LocalDir(src=Path(host_path).expanduser().resolve())
|
||||
entries, bind_mounts = build_session_entries(local_sources)
|
||||
|
||||
# Caido runs as an in-container sidecar; HTTP(S) traffic from any
|
||||
# process started via ``session.exec`` (the SDK's Shell tool, etc.)
|
||||
@@ -81,6 +110,7 @@ async def create_or_reuse(
|
||||
image=image,
|
||||
manifest=manifest,
|
||||
exposed_ports=(_CONTAINER_CAIDO_PORT,),
|
||||
bind_mounts=bind_mounts,
|
||||
)
|
||||
|
||||
caido_endpoint = await session.resolve_exposed_port(_CONTAINER_CAIDO_PORT)
|
||||
|
||||
@@ -13,5 +13,5 @@ returns it as an image content block for vision-capable models.
|
||||
`containers/docker-entrypoint.sh`).
|
||||
- **Skill:** screenshot workflow lives in `strix/skills/tooling/agent_browser.md`.
|
||||
- **Recovery:** vision-not-supported model rejections are auto-recovered
|
||||
via `strix.core.sessions.strip_latest_image_from_session`, invoked from
|
||||
via `strix.core.sessions.strip_all_images_from_session`, invoked from
|
||||
`strix/core/execution.py`.
|
||||
|
||||
@@ -0,0 +1,44 @@
|
||||
"""Tests for the scan-wide budget-stop signal on the agent coordinator."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
|
||||
from strix.core.agents import AgentCoordinator
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_budget_stop_sets_flag() -> None:
|
||||
coordinator = AgentCoordinator()
|
||||
await coordinator.register("root", "strix", parent_id=None)
|
||||
|
||||
assert coordinator.budget_stopped is False
|
||||
await coordinator.trigger_budget_stop()
|
||||
assert coordinator.budget_stopped is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_budget_stop_unblocks_parked_agent() -> None:
|
||||
# A parent parked in wait_for_message (awaiting a child) must be released so
|
||||
# it can exit, no matter where in the tree the budget limit was hit.
|
||||
coordinator = AgentCoordinator()
|
||||
await coordinator.register("parent", "strix", parent_id=None)
|
||||
|
||||
waiter = asyncio.create_task(coordinator.wait_for_message("parent"))
|
||||
await asyncio.sleep(0) # let the waiter park
|
||||
assert not waiter.done()
|
||||
|
||||
await coordinator.trigger_budget_stop()
|
||||
await asyncio.wait_for(waiter, timeout=1.0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_wait_for_message_returns_immediately_after_budget_stop() -> None:
|
||||
coordinator = AgentCoordinator()
|
||||
await coordinator.register("agent", "recon", parent_id="parent")
|
||||
await coordinator.trigger_budget_stop()
|
||||
|
||||
# No pending messages, but the stop flag short-circuits the wait.
|
||||
await asyncio.wait_for(coordinator.wait_for_message("agent"), timeout=1.0)
|
||||
@@ -0,0 +1,108 @@
|
||||
"""Tests for budget enforcement in ReportUsageHooks."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from strix.core.hooks import BudgetExceededError, ReportUsageHooks
|
||||
|
||||
|
||||
def _make_hooks(max_budget: float | None) -> ReportUsageHooks:
|
||||
return ReportUsageHooks(model="test-model", max_budget_usd=max_budget)
|
||||
|
||||
|
||||
def _make_report_state(cost: float) -> MagicMock:
|
||||
state = MagicMock()
|
||||
state.get_total_llm_cost.return_value = cost
|
||||
state.record_sdk_usage = MagicMock()
|
||||
return state
|
||||
|
||||
|
||||
def _make_context(agent_id: str = "test-agent") -> MagicMock:
|
||||
ctx: MagicMock = MagicMock()
|
||||
ctx.context = {"agent_id": agent_id}
|
||||
return ctx
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_budget_never_raises() -> None:
|
||||
hooks = _make_hooks(None)
|
||||
state = _make_report_state(9999.0)
|
||||
with patch("strix.core.hooks.get_global_report_state", return_value=state):
|
||||
await hooks.on_llm_end(_make_context(), MagicMock(), MagicMock())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_under_budget_does_not_raise() -> None:
|
||||
hooks = _make_hooks(10.0)
|
||||
state = _make_report_state(9.99)
|
||||
with patch("strix.core.hooks.get_global_report_state", return_value=state):
|
||||
await hooks.on_llm_end(_make_context(), MagicMock(), MagicMock())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_at_budget_raises() -> None:
|
||||
hooks = _make_hooks(10.0)
|
||||
state = _make_report_state(10.0)
|
||||
with (
|
||||
patch("strix.core.hooks.get_global_report_state", return_value=state),
|
||||
pytest.raises(BudgetExceededError),
|
||||
):
|
||||
await hooks.on_llm_end(_make_context(), MagicMock(), MagicMock())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_over_budget_raises() -> None:
|
||||
hooks = _make_hooks(10.0)
|
||||
state = _make_report_state(10.01)
|
||||
with (
|
||||
patch("strix.core.hooks.get_global_report_state", return_value=state),
|
||||
pytest.raises(BudgetExceededError),
|
||||
):
|
||||
await hooks.on_llm_end(_make_context(), MagicMock(), MagicMock())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_budget_check_uses_live_cost_accessor() -> None:
|
||||
# The check must read the live ledger, not the persisted run-record snapshot,
|
||||
# so it stays accurate even when a save fails after a usage record.
|
||||
hooks = _make_hooks(5.0)
|
||||
state = _make_report_state(6.0)
|
||||
with (
|
||||
patch("strix.core.hooks.get_global_report_state", return_value=state),
|
||||
pytest.raises(BudgetExceededError),
|
||||
):
|
||||
await hooks.on_llm_end(_make_context(), MagicMock(), MagicMock())
|
||||
state.get_total_llm_cost.assert_called_once()
|
||||
state.get_total_llm_usage.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_error_message_includes_amounts() -> None:
|
||||
hooks = _make_hooks(5.0)
|
||||
state = _make_report_state(7.1234)
|
||||
with patch("strix.core.hooks.get_global_report_state", return_value=state):
|
||||
with pytest.raises(BudgetExceededError, match=r"\$5\.00") as exc_info:
|
||||
await hooks.on_llm_end(_make_context(), MagicMock(), MagicMock())
|
||||
assert "7.1234" in str(exc_info.value)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_raise_when_report_state_none() -> None:
|
||||
hooks = _make_hooks(1.0)
|
||||
with patch("strix.core.hooks.get_global_report_state", return_value=None):
|
||||
# Should return early without raising, even with budget set
|
||||
await hooks.on_llm_end(_make_context(), MagicMock(), MagicMock())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad_budget", [0.0, -0.01, -5.0])
|
||||
def test_non_positive_budget_rejected(bad_budget: float) -> None:
|
||||
with pytest.raises(ValueError, match="greater than 0"):
|
||||
ReportUsageHooks(model="test-model", max_budget_usd=bad_budget)
|
||||
|
||||
|
||||
def test_budget_exceeded_error_is_runtime_error() -> None:
|
||||
err = BudgetExceededError("test")
|
||||
assert isinstance(err, RuntimeError)
|
||||
@@ -0,0 +1,188 @@
|
||||
"""Tests for local-source sizing and ``--mount`` target helpers in interface.utils."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pathlib import Path
|
||||
|
||||
from strix.interface.utils import (
|
||||
build_mount_targets_info,
|
||||
collect_local_sources,
|
||||
dedupe_local_targets,
|
||||
directory_size_bytes,
|
||||
find_oversized_local_targets,
|
||||
)
|
||||
|
||||
|
||||
def _write_file(path: Path, size: int) -> None:
|
||||
path.write_bytes(b"x" * size)
|
||||
|
||||
|
||||
def _local_target(target_path: str, *, mount: bool = False) -> dict[str, Any]:
|
||||
details: dict[str, Any] = {"target_path": target_path, "workspace_subdir": "repo"}
|
||||
if mount:
|
||||
details["mount"] = True
|
||||
return {"type": "local_code", "details": details, "original": target_path}
|
||||
|
||||
|
||||
def test_directory_size_empty_dir_is_zero(tmp_path: Path) -> None:
|
||||
assert directory_size_bytes(tmp_path) == 0
|
||||
|
||||
|
||||
def test_directory_size_sums_flat_and_nested_files(tmp_path: Path) -> None:
|
||||
_write_file(tmp_path / "a.txt", 100)
|
||||
nested = tmp_path / "sub" / "deep"
|
||||
nested.mkdir(parents=True)
|
||||
_write_file(nested / "b.txt", 250)
|
||||
assert directory_size_bytes(tmp_path) == 350
|
||||
|
||||
|
||||
def test_directory_size_skips_symlinks(tmp_path: Path) -> None:
|
||||
_write_file(tmp_path / "real.txt", 100)
|
||||
(tmp_path / "link.txt").symlink_to(tmp_path / "real.txt")
|
||||
# The symlink target is counted once via the real file, not doubled.
|
||||
assert directory_size_bytes(tmp_path) == 100
|
||||
|
||||
|
||||
@pytest.mark.skipif(sys.platform == "win32", reason="relies on POSIX permissions")
|
||||
def test_directory_size_logs_and_skips_unreadable_subdir(
|
||||
tmp_path: Path, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
if hasattr(os, "geteuid") and os.geteuid() == 0:
|
||||
pytest.skip("root bypasses directory permissions")
|
||||
_write_file(tmp_path / "top.txt", 100)
|
||||
locked = tmp_path / "locked"
|
||||
locked.mkdir()
|
||||
_write_file(locked / "secret.bin", 9999)
|
||||
locked.chmod(0o000)
|
||||
try:
|
||||
with caplog.at_level(logging.WARNING):
|
||||
size = directory_size_bytes(tmp_path)
|
||||
finally:
|
||||
locked.chmod(0o755)
|
||||
# The unreadable subtree is excluded (not silently treated as readable) and
|
||||
# the omission is logged rather than vanishing without a trace.
|
||||
assert size == 100
|
||||
assert any("Could not read" in record.message for record in caplog.records)
|
||||
|
||||
|
||||
def test_find_oversized_returns_nothing_under_limit(tmp_path: Path) -> None:
|
||||
_write_file(tmp_path / "a.txt", 100)
|
||||
targets = [_local_target(str(tmp_path))]
|
||||
assert find_oversized_local_targets(targets, max_bytes=1000) == []
|
||||
|
||||
|
||||
def test_find_oversized_returns_target_over_limit(tmp_path: Path) -> None:
|
||||
_write_file(tmp_path / "big.bin", 500)
|
||||
targets = [_local_target(str(tmp_path))]
|
||||
result = find_oversized_local_targets(targets, max_bytes=100)
|
||||
assert result == [(str(tmp_path), 500)]
|
||||
|
||||
|
||||
def test_find_oversized_ignores_mounted_targets(tmp_path: Path) -> None:
|
||||
_write_file(tmp_path / "big.bin", 500)
|
||||
targets = [_local_target(str(tmp_path), mount=True)]
|
||||
assert find_oversized_local_targets(targets, max_bytes=100) == []
|
||||
|
||||
|
||||
def test_find_oversized_ignores_non_local_targets() -> None:
|
||||
targets = [{"type": "web_application", "details": {"target_url": "https://x"}}]
|
||||
assert find_oversized_local_targets(targets, max_bytes=1) == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("disabled", [0, -1])
|
||||
def test_find_oversized_disabled_for_non_positive_limit(tmp_path: Path, disabled: int) -> None:
|
||||
_write_file(tmp_path / "big.bin", 500)
|
||||
targets = [_local_target(str(tmp_path))]
|
||||
assert find_oversized_local_targets(targets, max_bytes=disabled) == []
|
||||
|
||||
|
||||
def test_collect_local_sources_propagates_mount_flag() -> None:
|
||||
copied = _local_target("/copied")
|
||||
copied["details"]["workspace_subdir"] = "copied"
|
||||
mounted = _local_target("/mounted", mount=True)
|
||||
mounted["details"]["workspace_subdir"] = "mounted"
|
||||
|
||||
sources = collect_local_sources([copied, mounted])
|
||||
|
||||
by_path = {s["source_path"]: s for s in sources}
|
||||
assert by_path["/copied"]["mount"] is False
|
||||
assert by_path["/mounted"]["mount"] is True
|
||||
|
||||
|
||||
def test_collect_local_sources_repository_is_never_mounted() -> None:
|
||||
repo = {
|
||||
"type": "repository",
|
||||
"details": {"cloned_repo_path": "/clone", "workspace_subdir": "clone"},
|
||||
}
|
||||
sources = collect_local_sources([repo])
|
||||
assert sources == [{"source_path": "/clone", "workspace_subdir": "clone", "mount": False}]
|
||||
|
||||
|
||||
def test_build_mount_targets_info_for_valid_dir(tmp_path: Path) -> None:
|
||||
result = build_mount_targets_info([str(tmp_path)])
|
||||
assert len(result) == 1
|
||||
entry = result[0]
|
||||
assert entry["type"] == "local_code"
|
||||
assert entry["details"]["mount"] is True
|
||||
assert entry["details"]["target_path"] == str(tmp_path.resolve())
|
||||
|
||||
|
||||
def test_build_mount_targets_info_rejects_missing_path(tmp_path: Path) -> None:
|
||||
missing = tmp_path / "does-not-exist"
|
||||
with pytest.raises(ValueError, match="not an existing directory"):
|
||||
build_mount_targets_info([str(missing)])
|
||||
|
||||
|
||||
def test_build_mount_targets_info_rejects_file(tmp_path: Path) -> None:
|
||||
file_path = tmp_path / "a-file.txt"
|
||||
_write_file(file_path, 10)
|
||||
with pytest.raises(ValueError, match="not an existing directory"):
|
||||
build_mount_targets_info([str(file_path)])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("empty", ["", " "])
|
||||
def test_build_mount_targets_info_rejects_empty_path(empty: str) -> None:
|
||||
# An empty path would otherwise resolve to the current working directory
|
||||
# and silently bind-mount it into the sandbox.
|
||||
with pytest.raises(ValueError, match="must not be empty"):
|
||||
build_mount_targets_info([empty])
|
||||
|
||||
|
||||
def test_dedupe_keeps_distinct_targets_in_order() -> None:
|
||||
targets = [
|
||||
_local_target("/a"),
|
||||
{"type": "web_application", "details": {"target_url": "https://x"}},
|
||||
_local_target("/b", mount=True),
|
||||
]
|
||||
assert dedupe_local_targets(targets) == targets
|
||||
|
||||
|
||||
def test_dedupe_mount_supersedes_copied_same_path() -> None:
|
||||
copied = _local_target("/repo")
|
||||
mounted = _local_target("/repo", mount=True)
|
||||
|
||||
# Copied first, then mounted: the single surviving entry is the mount.
|
||||
result = dedupe_local_targets([copied, mounted])
|
||||
assert len(result) == 1
|
||||
assert result[0]["details"]["mount"] is True
|
||||
|
||||
# Order-independent: mounted first, copied second also yields the mount.
|
||||
result_rev = dedupe_local_targets([mounted, copied])
|
||||
assert len(result_rev) == 1
|
||||
assert result_rev[0]["details"]["mount"] is True
|
||||
|
||||
|
||||
def test_dedupe_collapses_duplicate_mounts() -> None:
|
||||
result = dedupe_local_targets(
|
||||
[_local_target("/repo", mount=True), _local_target("/repo", mount=True)]
|
||||
)
|
||||
assert len(result) == 1
|
||||
@@ -0,0 +1,62 @@
|
||||
"""Tests for LLM model recommendation helpers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from strix.config.models import RECOMMENDED_MODEL_NAMES, is_recommended_or_frontier_model
|
||||
|
||||
|
||||
@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_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",
|
||||
"litellm/openai/gpt-5.4-pro",
|
||||
"azure_ai/gpt-5.5-pro",
|
||||
"bedrock_mantle/openai.gpt-5.5",
|
||||
"anthropic/claude-opus-4-8",
|
||||
"anthropic.claude-opus-4-8",
|
||||
"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",
|
||||
"xai/grok-4.3",
|
||||
"openrouter/x-ai/grok-4",
|
||||
"deepseek/deepseek-v4-pro",
|
||||
"deepseek/deepseek-r1-0528",
|
||||
"deepseek/deepseek-reasoner",
|
||||
"dashscope/qwen3-max-2026-01-23",
|
||||
"qwen3.7-max",
|
||||
"moonshot/kimi-k2.6",
|
||||
"kimi-k2.7-code",
|
||||
"mistral/mistral-medium-3-5",
|
||||
"mistral/magistral-medium-latest",
|
||||
],
|
||||
)
|
||||
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",
|
||||
"custom-ollama/gpt-5-mini-local",
|
||||
"custom-provider/claude-opus-4-local",
|
||||
],
|
||||
)
|
||||
def test_non_frontier_models_are_rejected(model_name: str) -> None:
|
||||
assert not is_recommended_or_frontier_model(model_name)
|
||||
@@ -0,0 +1,67 @@
|
||||
"""Tests for build_session_entries: splitting copied vs bind-mounted sources."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from agents.sandbox.entries import LocalDir
|
||||
|
||||
from strix.runtime.session_manager import build_session_entries
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def _source(subdir: str, path: str, *, mount: bool = False) -> dict[str, Any]:
|
||||
return {"source_path": path, "workspace_subdir": subdir, "mount": mount}
|
||||
|
||||
|
||||
def test_copied_source_becomes_localdir_entry(tmp_path: Path) -> None:
|
||||
entries, bind_mounts = build_session_entries([_source("repo", str(tmp_path))])
|
||||
|
||||
assert bind_mounts == []
|
||||
assert isinstance(entries["repo"], LocalDir)
|
||||
assert entries["repo"].src == tmp_path.resolve()
|
||||
|
||||
|
||||
def test_mounted_source_becomes_bind_mount(tmp_path: Path) -> None:
|
||||
entries, bind_mounts = build_session_entries([_source("repo", str(tmp_path), mount=True)])
|
||||
|
||||
assert entries == {}
|
||||
assert bind_mounts == [
|
||||
{
|
||||
"source": str(tmp_path.resolve()),
|
||||
"target": "/workspace/repo",
|
||||
"read_only": True,
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_mixed_sources_split_correctly(tmp_path: Path) -> None:
|
||||
copied = tmp_path / "copied"
|
||||
mounted = tmp_path / "mounted"
|
||||
copied.mkdir()
|
||||
mounted.mkdir()
|
||||
|
||||
entries, bind_mounts = build_session_entries(
|
||||
[
|
||||
_source("copied", str(copied)),
|
||||
_source("mounted", str(mounted), mount=True),
|
||||
]
|
||||
)
|
||||
|
||||
assert list(entries) == ["copied"]
|
||||
assert isinstance(entries["copied"], LocalDir)
|
||||
assert [m["target"] for m in bind_mounts] == ["/workspace/mounted"]
|
||||
|
||||
|
||||
def test_incomplete_sources_are_skipped() -> None:
|
||||
entries, bind_mounts = build_session_entries(
|
||||
[
|
||||
{"source_path": "", "workspace_subdir": "x"},
|
||||
{"source_path": "/p", "workspace_subdir": ""},
|
||||
]
|
||||
)
|
||||
assert entries == {}
|
||||
assert bind_mounts == []
|
||||
@@ -775,6 +775,15 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/a0/d9/a1e041c5e7caa9a05c925f4bdbdfb7f006d1f74996af53467bc394c97be7/importlib_metadata-8.5.0-py3-none-any.whl", hash = "sha256:45e54197d28b7a7f1559e60b95e7c567032b602131fbd588f1497f47880aa68b", size = 26514, upload-time = "2024-09-11T14:56:07.019Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "iniconfig"
|
||||
version = "2.3.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "jinja2"
|
||||
version = "3.1.6"
|
||||
@@ -956,7 +965,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "litellm"
|
||||
version = "1.83.7"
|
||||
version = "1.88.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
@@ -972,9 +981,9 @@ dependencies = [
|
||||
{ name = "tiktoken" },
|
||||
{ name = "tokenizers" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/77/2b/b58bf6bbcbc3d0e55d0a84fdf9128e5b1436517f46fce89b1cd8948ebb81/litellm-1.83.7.tar.gz", hash = "sha256:e2f2cb99df2e2b2eab63f1354faa45c88dd7c8d40c18eb648afb1b349c689633", size = 17791694, upload-time = "2026-04-13T17:35:01.606Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/9b/8c/6cfce5d15554e076b8438a955436e9e6e2a2bd39c76539656c1c3861c369/litellm-1.88.0.tar.gz", hash = "sha256:4ff794493e40bd86c6f13e91dcb3e1aad697403fd46a96902196d93356ba48f4", size = 13886026, upload-time = "2026-06-06T23:25:17.93Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/75/80/caeb4cdcad96451ba83ad3ba2a9da08b1e1a915fa845c489f56ea044488b/litellm-1.83.7-py3-none-any.whl", hash = "sha256:5784a1d9a9a4a8acd6ca1e347003a5e2e1b3c749b4d41e7da4904577adade111", size = 16069807, upload-time = "2026-04-13T17:34:58.36Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/da/71/deb6637475253b83eb4577c058863eec4b4ddaef2d08dcd93981b1aa4209/litellm-1.88.0-py3-none-any.whl", hash = "sha256:abd3037e0bf5703f833f5565c87bdfd93578d3de46cbcb36dfa108c3ef58021c", size = 15276203, upload-time = "2026-06-06T23:25:07.257Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1347,6 +1356,15 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/63/d7/97f7e3a6abb67d8080dd406fd4df842c2be0efaf712d1c899c32a075027c/platformdirs-4.9.4-py3-none-any.whl", hash = "sha256:68a9a4619a666ea6439f2ff250c12a853cd1cbd5158d258bd824a7df6be2f868", size = 21216, upload-time = "2026-03-05T18:34:12.172Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pluggy"
|
||||
version = "1.6.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pre-commit"
|
||||
version = "4.5.1"
|
||||
@@ -1633,6 +1651,35 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/0c/82/a2c93e32800940d9573fb28c346772a14778b84ba7524e691b324620ab89/pyright-1.1.408-py3-none-any.whl", hash = "sha256:090b32865f4fdb1e0e6cd82bf5618480d48eecd2eb2e70f960982a3d9a4c17c1", size = 6399144, upload-time = "2026-01-08T08:07:37.082Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pytest"
|
||||
version = "9.1.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "colorama", marker = "sys_platform == 'win32'" },
|
||||
{ name = "iniconfig" },
|
||||
{ name = "packaging" },
|
||||
{ name = "pluggy" },
|
||||
{ name = "pygments" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/84/0e/b5858858d74958632c49b72cb25a3976ff9f632397626715be71c89d3971/pytest-9.1.0.tar.gz", hash = "sha256:41dd9148c08072446394cefd3d79701701335a9f4cae69ba92e39f6c7f5c061c", size = 1634181, upload-time = "2026-06-13T18:52:45.983Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/8b/5a/ba30a81239b909821b3153e303e7def45178bf353da4f72380e6c5e8793b/pytest-9.1.0-py3-none-any.whl", hash = "sha256:8ebb0e7888bdf2bdfc602ec51f8f62d50200af37356c74e503c79a94f5c81f32", size = 386453, upload-time = "2026-06-13T18:52:44.045Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pytest-asyncio"
|
||||
version = "1.4.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "pytest" },
|
||||
{ name = "typing-extensions", marker = "python_full_version < '3.13'" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/43/7c/d36d04db312ecf4298932ef77e6e4a9e8ad017906e24e34f0b0c361a2473/pytest_asyncio-1.4.0.tar.gz", hash = "sha256:c6c0d2259945122819f171a32ecea2c349ead889ee28176caaf492143424be42", size = 58514, upload-time = "2026-05-26T09:56:04.083Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/03/e2/08a497ef684b88559c9cc5f4ad53a37e7b99e727094a86d6ea32536d5d3c/pytest_asyncio-1.4.0-py3-none-any.whl", hash = "sha256:933ca923a23075a87fb7070c0ec272a6848489824d887c85c812670932835aa1", size = 16930, upload-time = "2026-05-26T09:56:02.576Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "python-discovery"
|
||||
version = "1.2.0"
|
||||
@@ -2035,7 +2082,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "strix-agent"
|
||||
version = "1.0.0"
|
||||
version = "1.0.4"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "caido-sdk-client" },
|
||||
@@ -2056,6 +2103,8 @@ dev = [
|
||||
{ name = "pre-commit" },
|
||||
{ name = "pyinstaller", marker = "python_full_version < '3.15'" },
|
||||
{ name = "pyright" },
|
||||
{ name = "pytest" },
|
||||
{ name = "pytest-asyncio" },
|
||||
{ name = "ruff" },
|
||||
]
|
||||
|
||||
@@ -2079,6 +2128,8 @@ dev = [
|
||||
{ name = "pre-commit", specifier = ">=4.2.0" },
|
||||
{ name = "pyinstaller", marker = "python_full_version >= '3.12' and python_full_version < '3.15'", specifier = ">=6.17.0" },
|
||||
{ name = "pyright", specifier = ">=1.1.401" },
|
||||
{ name = "pytest", specifier = ">=8.3" },
|
||||
{ name = "pytest-asyncio", specifier = ">=0.24" },
|
||||
{ name = "ruff", specifier = ">=0.11.13" },
|
||||
]
|
||||
|
||||
|
||||
Reference in New Issue
Block a user