mirror of
https://github.com/usestrix/strix.git
synced 2026-08-16 09:26:39 +02:00
Add configurable token / cost usage limits (#576)
* fix: resolve pre-commit check failures - Change RuntimeError to TypeError for type validation in report/writer.py - Update pyupgrade to v3.21.2 for Python 3.14 compatibility * feat(cli): add --max-budget-usd flag Raises BudgetExceededError in ReportUsageHooks after each LLM call when accumulated cost reaches the limit, with clean "stopped" status and child-agent cancellation in non-interactive mode. * test: add budget enforcement unit tests 7 tests covering no-budget, under-budget, at-limit, over-limit, error message content, None report state, and exception hierarchy. Also adds pytest/pytest-asyncio to dev deps and a mypy override for tests. * fix(budget): validate positive budget and check the live cost ledger Two hardening fixes for --max-budget-usd enforcement: - Reject non-positive budgets. ReportUsageHooks now raises ValueError for max_budget_usd <= 0, and the CLI validates the flag via a custom argparse type so '--max-budget-usd 0' fails fast with a friendly message instead of silently killing the scan on the first model response. - Read the live cost. The budget check now reads ReportState.get_total_llm_cost() (the live ledger) instead of the persisted run-record snapshot, so it stays accurate even when a usage save fails after a model call. * fix(budget): stop the entire scan deterministically when the limit is hit Previously a BudgetExceededError was handled per-agent: it was swallowed in interactive mode (the loop kept waiting), a child's error escaped its detached task as an unretrieved-exception warning, the parent was never released from wait_for_message, and the stop was logged at ERROR with a traceback as if the agent had failed. Replace that with a single scan-wide signal on the coordinator: - AgentCoordinator.trigger_budget_stop() sets a flag and wakes every parked agent; wait_for_message returns as soon as the flag is set. - The run loops check coordinator.budget_stopped and raise to exit cleanly, marking themselves 'stopped'. The root's exception reaches run_strix_scan's handler, which cancels descendants and tears the scan down once; child exceptions are swallowed in their detached task. - The budget stop is logged at INFO, not as a failure. This is deterministic regardless of tree depth or which agent first sees the limit, fixing the interactive/TUI hang where a deep agent's stop never reached a parked root. Also re-raises BudgetExceededError explicitly in the stream handler so it can't be mistaken for the LiteLLM 'after shutdown' race. * fix(budget): treat a budget stop as a clean stop in the TUI Add an explicit BudgetExceededError handler in the TUI scan thread so that, if the error ever reaches it, the budget stop is logged as a graceful stop rather than surfaced as a red scan error by the broad 'except Exception'. The runner normally absorbs the error and returns cleanly, so this is defensive depth for a money-spending feature. * docs(cli): document --max-budget-usd behavior and limitations Clarify that the budget is cumulative across all agents, checked after each model response, that the scan stops cleanly (not as a failure), that the value must be > 0, and that spend can slightly overshoot due to in-flight calls and best-effort cost estimation. * Apply suggestions from code review Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --------- Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
This commit is contained in:
co-authored by
greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
parent
11e5d1c2b3
commit
962d4459d9
@@ -19,6 +19,7 @@ repos:
|
|||||||
types-python-dateutil,
|
types-python-dateutil,
|
||||||
pydantic,
|
pydantic,
|
||||||
fastapi,
|
fastapi,
|
||||||
|
pytest,
|
||||||
"openai-agents[litellm]==0.14.6",
|
"openai-agents[litellm]==0.14.6",
|
||||||
]
|
]
|
||||||
args: [--install-types, --non-interactive]
|
args: [--install-types, --non-interactive]
|
||||||
@@ -46,7 +47,7 @@ repos:
|
|||||||
|
|
||||||
# Additional Python code quality checks
|
# Additional Python code quality checks
|
||||||
- repo: https://github.com/asottile/pyupgrade
|
- repo: https://github.com/asottile/pyupgrade
|
||||||
rev: v3.20.0
|
rev: v3.21.2
|
||||||
hooks:
|
hooks:
|
||||||
- id: pyupgrade
|
- id: pyupgrade
|
||||||
args: [--py312-plus]
|
args: [--py312-plus]
|
||||||
|
|||||||
@@ -43,6 +43,24 @@ strix --target <target> [options]
|
|||||||
Path to a custom config file (JSON) to use instead of `~/.strix/cli-config.json`.
|
Path to a custom config file (JSON) to use instead of `~/.strix/cli-config.json`.
|
||||||
</ParamField>
|
</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
|
## Examples
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
|
|||||||
@@ -55,8 +55,13 @@ dev = [
|
|||||||
"bandit>=1.8.3",
|
"bandit>=1.8.3",
|
||||||
"pre-commit>=4.2.0",
|
"pre-commit>=4.2.0",
|
||||||
"pyinstaller>=6.17.0; python_version >= '3.12' and python_version < '3.15'",
|
"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]
|
[build-system]
|
||||||
requires = ["hatchling"]
|
requires = ["hatchling"]
|
||||||
build-backend = "hatchling.build"
|
build-backend = "hatchling.build"
|
||||||
@@ -104,6 +109,10 @@ module = [
|
|||||||
ignore_missing_imports = true
|
ignore_missing_imports = true
|
||||||
disable_error_code = ["import-untyped"]
|
disable_error_code = ["import-untyped"]
|
||||||
|
|
||||||
|
[[tool.mypy.overrides]]
|
||||||
|
module = ["tests.*"]
|
||||||
|
disallow_untyped_decorators = false
|
||||||
|
|
||||||
# ============================================================================
|
# ============================================================================
|
||||||
# Ruff Configuration (Fast Python Linter & Formatter)
|
# Ruff Configuration (Fast Python Linter & Formatter)
|
||||||
# ============================================================================
|
# ============================================================================
|
||||||
|
|||||||
+13
-1
@@ -43,6 +43,7 @@ class AgentCoordinator:
|
|||||||
self._lock = asyncio.Lock()
|
self._lock = asyncio.Lock()
|
||||||
self._snapshot_path: Path | None = None
|
self._snapshot_path: Path | None = None
|
||||||
self.is_shutting_down = False
|
self.is_shutting_down = False
|
||||||
|
self._budget_stopped = False
|
||||||
|
|
||||||
def set_snapshot_path(self, path: Path) -> None:
|
def set_snapshot_path(self, path: Path) -> None:
|
||||||
self._snapshot_path = path
|
self._snapshot_path = path
|
||||||
@@ -50,6 +51,17 @@ class AgentCoordinator:
|
|||||||
def mark_shutting_down(self) -> None:
|
def mark_shutting_down(self) -> None:
|
||||||
self.is_shutting_down = True
|
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(
|
async def register(
|
||||||
self,
|
self,
|
||||||
agent_id: str,
|
agent_id: str,
|
||||||
@@ -143,7 +155,7 @@ class AgentCoordinator:
|
|||||||
async def wait_for_message(self, agent_id: str) -> None:
|
async def wait_for_message(self, agent_id: str) -> None:
|
||||||
while True:
|
while True:
|
||||||
async with self._lock:
|
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
|
return
|
||||||
wake = self.runtimes.setdefault(agent_id, AgentRuntime()).wake
|
wake = self.runtimes.setdefault(agent_id, AgentRuntime()).wake
|
||||||
wake.clear()
|
wake.clear()
|
||||||
|
|||||||
+45
-17
@@ -15,6 +15,7 @@ from agents.sandbox.errors import ExecTransportError
|
|||||||
from docker import errors as docker_errors # type: ignore[import-untyped, unused-ignore]
|
from docker import errors as docker_errors # type: ignore[import-untyped, unused-ignore]
|
||||||
from openai import APIError
|
from openai import APIError
|
||||||
|
|
||||||
|
from strix.core.hooks import BudgetExceededError
|
||||||
from strix.core.inputs import child_initial_input
|
from strix.core.inputs import child_initial_input
|
||||||
from strix.core.sessions import open_agent_session, strip_all_images_from_session
|
from strix.core.sessions import open_agent_session, strip_all_images_from_session
|
||||||
|
|
||||||
@@ -97,6 +98,10 @@ async def run_agent_loop(
|
|||||||
except asyncio.CancelledError:
|
except asyncio.CancelledError:
|
||||||
return result
|
return result
|
||||||
|
|
||||||
|
if coordinator.budget_stopped:
|
||||||
|
await coordinator.set_status(agent_id, "stopped")
|
||||||
|
raise BudgetExceededError("scan budget reached")
|
||||||
|
|
||||||
await coordinator.consume_pending(agent_id)
|
await coordinator.consume_pending(agent_id)
|
||||||
result = await _run_cycle(
|
result = await _run_cycle(
|
||||||
agent,
|
agent,
|
||||||
@@ -278,6 +283,10 @@ async def _run_noninteractive_until_lifecycle(
|
|||||||
invalid_final_output_limit = max(1, max_turns)
|
invalid_final_output_limit = max(1, max_turns)
|
||||||
|
|
||||||
while True:
|
while True:
|
||||||
|
if coordinator.budget_stopped:
|
||||||
|
await coordinator.set_status(agent_id, "stopped")
|
||||||
|
raise BudgetExceededError("scan budget reached")
|
||||||
|
|
||||||
result = await _run_cycle(
|
result = await _run_cycle(
|
||||||
agent,
|
agent,
|
||||||
coordinator,
|
coordinator,
|
||||||
@@ -360,6 +369,10 @@ async def _run_cycle( # noqa: PLR0912, PLR0915
|
|||||||
logger.exception("stream event sink failed for %s", agent_id)
|
logger.exception("stream event sink failed for %s", agent_id)
|
||||||
if stream.run_loop_exception is not None:
|
if stream.run_loop_exception is not None:
|
||||||
raise stream.run_loop_exception
|
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:
|
except RuntimeError as stream_exc:
|
||||||
if "after shutdown" not in str(stream_exc):
|
if "after shutdown" not in str(stream_exc):
|
||||||
raise
|
raise
|
||||||
@@ -377,6 +390,13 @@ async def _run_cycle( # noqa: PLR0912, PLR0915
|
|||||||
)
|
)
|
||||||
finally:
|
finally:
|
||||||
await coordinator.detach_stream(agent_id, stream)
|
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:
|
except Exception as exc:
|
||||||
if (
|
if (
|
||||||
image_strips < 3
|
image_strips < 3
|
||||||
@@ -527,21 +547,29 @@ async def _start_child_runner(
|
|||||||
child_ctx["parent_id"] = parent_id
|
child_ctx["parent_id"] = parent_id
|
||||||
child_ctx["task"] = task
|
child_ctx["task"] = task
|
||||||
|
|
||||||
task_handle = asyncio.create_task(
|
async def _child_loop() -> None:
|
||||||
run_agent_loop(
|
# A budget stop is a clean scan-wide shutdown, not a child failure: the
|
||||||
agent=child_agent,
|
# child's status and parent notification are already settled in
|
||||||
initial_input=initial_input,
|
# ``_run_cycle``. Swallow it here so the detached task does not surface a
|
||||||
run_config=run_config,
|
# spurious "Task exception was never retrieved" warning. The root agent
|
||||||
context=child_ctx,
|
# hits the same limit on its next call and tears the scan down.
|
||||||
max_turns=max_turns,
|
try:
|
||||||
coordinator=coordinator,
|
await run_agent_loop(
|
||||||
agent_id=child_id,
|
agent=child_agent,
|
||||||
interactive=interactive,
|
initial_input=initial_input,
|
||||||
session=session,
|
run_config=run_config,
|
||||||
start_parked=start_parked,
|
context=child_ctx,
|
||||||
event_sink=event_sink,
|
max_turns=max_turns,
|
||||||
hooks=hooks,
|
coordinator=coordinator,
|
||||||
),
|
agent_id=child_id,
|
||||||
name=f"agent-{name}-{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)
|
await coordinator.attach_runtime(child_id, task=task_handle)
|
||||||
|
|||||||
+16
-1
@@ -19,11 +19,19 @@ if TYPE_CHECKING:
|
|||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class BudgetExceededError(RuntimeError):
|
||||||
|
"""Raised when the accumulated LLM cost reaches the configured budget."""
|
||||||
|
|
||||||
|
|
||||||
class ReportUsageHooks(RunHooks[dict[str, Any]]):
|
class ReportUsageHooks(RunHooks[dict[str, Any]]):
|
||||||
"""Persist SDK-native usage after every model response."""
|
"""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._model = model
|
||||||
|
self._max_budget_usd = max_budget_usd
|
||||||
|
|
||||||
async def on_llm_end(
|
async def on_llm_end(
|
||||||
self,
|
self,
|
||||||
@@ -52,3 +60,10 @@ class ReportUsageHooks(RunHooks[dict[str, Any]]):
|
|||||||
)
|
)
|
||||||
except Exception:
|
except Exception:
|
||||||
logger.exception("failed to record SDK usage for agent %s", agent_id)
|
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
-2
@@ -27,7 +27,7 @@ from strix.core.execution import (
|
|||||||
from strix.core.execution import (
|
from strix.core.execution import (
|
||||||
spawn_child_agent as start_child_agent,
|
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 (
|
from strix.core.inputs import (
|
||||||
DEFAULT_MAX_TURNS,
|
DEFAULT_MAX_TURNS,
|
||||||
build_root_task,
|
build_root_task,
|
||||||
@@ -59,6 +59,7 @@ async def run_strix_scan(
|
|||||||
coordinator: AgentCoordinator | None = None,
|
coordinator: AgentCoordinator | None = None,
|
||||||
interactive: bool = False,
|
interactive: bool = False,
|
||||||
max_turns: int = DEFAULT_MAX_TURNS,
|
max_turns: int = DEFAULT_MAX_TURNS,
|
||||||
|
max_budget_usd: float | None = None,
|
||||||
model: str | None = None,
|
model: str | None = None,
|
||||||
cleanup_on_exit: bool = True,
|
cleanup_on_exit: bool = True,
|
||||||
event_sink: StreamEventSink | None = None,
|
event_sink: StreamEventSink | None = None,
|
||||||
@@ -164,7 +165,7 @@ async def run_strix_scan(
|
|||||||
sandbox=SandboxRunConfig(client=bundle["client"], session=bundle["session"]),
|
sandbox=SandboxRunConfig(client=bundle["client"], session=bundle["session"]),
|
||||||
trace_include_sensitive_data=False,
|
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)
|
scope_context = build_scope_context(scan_config)
|
||||||
|
|
||||||
@@ -300,6 +301,13 @@ async def run_strix_scan(
|
|||||||
str(final)[:300],
|
str(final)[:300],
|
||||||
)
|
)
|
||||||
return result # noqa: TRY300
|
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:
|
except BaseException:
|
||||||
logger.exception("Strix scan %s failed", scan_id)
|
logger.exception("Strix scan %s failed", scan_id)
|
||||||
if root_id is not None:
|
if root_id is not None:
|
||||||
|
|||||||
@@ -183,6 +183,7 @@ async def run_cli(args: Any) -> None: # noqa: PLR0915
|
|||||||
image=_resolve_sandbox_image(),
|
image=_resolve_sandbox_image(),
|
||||||
local_sources=getattr(args, "local_sources", None) or [],
|
local_sources=getattr(args, "local_sources", None) or [],
|
||||||
interactive=bool(getattr(args, "interactive", False)),
|
interactive=bool(getattr(args, "interactive", False)),
|
||||||
|
max_budget_usd=getattr(args, "max_budget_usd", None),
|
||||||
)
|
)
|
||||||
finally:
|
finally:
|
||||||
stop_updates.set()
|
stop_updates.set()
|
||||||
|
|||||||
@@ -301,6 +301,17 @@ def get_version() -> str:
|
|||||||
return "unknown"
|
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:
|
def parse_arguments() -> argparse.Namespace:
|
||||||
parser = argparse.ArgumentParser(
|
parser = argparse.ArgumentParser(
|
||||||
description="Strix Multi-Agent Cybersecurity Penetration Testing Tool",
|
description="Strix Multi-Agent Cybersecurity Penetration Testing Tool",
|
||||||
@@ -424,6 +435,13 @@ Examples:
|
|||||||
help="Path to a custom config file (JSON) to use instead of ~/.strix/cli-config.json",
|
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(
|
parser.add_argument(
|
||||||
"--resume",
|
"--resume",
|
||||||
type=str,
|
type=str,
|
||||||
|
|||||||
@@ -31,6 +31,7 @@ from textual.widgets import Button, Label, Static, TextArea, Tree
|
|||||||
from textual.widgets.tree import TreeNode
|
from textual.widgets.tree import TreeNode
|
||||||
|
|
||||||
from strix.config import load_settings
|
from strix.config import load_settings
|
||||||
|
from strix.core.hooks import BudgetExceededError
|
||||||
from strix.core.runner import run_strix_scan
|
from strix.core.runner import run_strix_scan
|
||||||
from strix.interface.tui.live_view import TuiLiveView
|
from strix.interface.tui.live_view import TuiLiveView
|
||||||
from strix.interface.tui.messages import send_user_message_to_agent
|
from strix.interface.tui.messages import send_user_message_to_agent
|
||||||
@@ -1369,12 +1370,18 @@ class StrixTUIApp(App): # type: ignore[misc]
|
|||||||
local_sources=getattr(self.args, "local_sources", None) or [],
|
local_sources=getattr(self.args, "local_sources", None) or [],
|
||||||
coordinator=self.coordinator,
|
coordinator=self.coordinator,
|
||||||
interactive=True,
|
interactive=True,
|
||||||
|
max_budget_usd=getattr(self.args, "max_budget_usd", None),
|
||||||
event_sink=self._capture_sdk_event,
|
event_sink=self._capture_sdk_event,
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
except (KeyboardInterrupt, asyncio.CancelledError):
|
except (KeyboardInterrupt, asyncio.CancelledError):
|
||||||
logger.info("Scan interrupted by user")
|
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:
|
except (ConnectionError, TimeoutError) as e:
|
||||||
logging.exception("Network error during scan")
|
logging.exception("Network error during scan")
|
||||||
self._scan_error = e
|
self._scan_error = e
|
||||||
|
|||||||
@@ -236,6 +236,10 @@ class ReportState:
|
|||||||
def get_total_llm_usage(self) -> dict[str, Any]:
|
def get_total_llm_usage(self) -> dict[str, Any]:
|
||||||
return dict(self.run_record.get("llm_usage") or self._build_llm_usage_record())
|
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(
|
def update_scan_final_fields(
|
||||||
self,
|
self,
|
||||||
executive_summary: str,
|
executive_summary: str,
|
||||||
|
|||||||
@@ -52,6 +52,10 @@ class LLMUsageLedger:
|
|||||||
if isinstance(cost, int | float) and cost > 0:
|
if isinstance(cost, int | float) and cost > 0:
|
||||||
self._total_cost += float(cost)
|
self._total_cost += float(cost)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def total_cost(self) -> float:
|
||||||
|
return _round_cost(self._total_cost)
|
||||||
|
|
||||||
def to_record(self) -> dict[str, Any]:
|
def to_record(self) -> dict[str, Any]:
|
||||||
record = serialize_usage(self._total_usage)
|
record = serialize_usage(self._total_usage)
|
||||||
record["cost"] = _round_cost(self._total_cost)
|
record["cost"] = _round_cost(self._total_cost)
|
||||||
|
|||||||
@@ -27,7 +27,7 @@ def read_run_record(run_dir: Path) -> dict[str, Any]:
|
|||||||
except (OSError, json.JSONDecodeError) as exc:
|
except (OSError, json.JSONDecodeError) as exc:
|
||||||
raise RuntimeError(f"run.json at {path} is unreadable: {exc}") from exc
|
raise RuntimeError(f"run.json at {path} is unreadable: {exc}") from exc
|
||||||
if not isinstance(data, dict):
|
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
|
return data
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -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)
|
||||||
@@ -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" },
|
{ 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]]
|
[[package]]
|
||||||
name = "jinja2"
|
name = "jinja2"
|
||||||
version = "3.1.6"
|
version = "3.1.6"
|
||||||
@@ -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" },
|
{ 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]]
|
[[package]]
|
||||||
name = "pre-commit"
|
name = "pre-commit"
|
||||||
version = "4.5.1"
|
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" },
|
{ 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]]
|
[[package]]
|
||||||
name = "python-discovery"
|
name = "python-discovery"
|
||||||
version = "1.2.0"
|
version = "1.2.0"
|
||||||
@@ -2056,6 +2103,8 @@ dev = [
|
|||||||
{ name = "pre-commit" },
|
{ name = "pre-commit" },
|
||||||
{ name = "pyinstaller", marker = "python_full_version < '3.15'" },
|
{ name = "pyinstaller", marker = "python_full_version < '3.15'" },
|
||||||
{ name = "pyright" },
|
{ name = "pyright" },
|
||||||
|
{ name = "pytest" },
|
||||||
|
{ name = "pytest-asyncio" },
|
||||||
{ name = "ruff" },
|
{ name = "ruff" },
|
||||||
]
|
]
|
||||||
|
|
||||||
@@ -2079,6 +2128,8 @@ dev = [
|
|||||||
{ name = "pre-commit", specifier = ">=4.2.0" },
|
{ 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 = "pyinstaller", marker = "python_full_version >= '3.12' and python_full_version < '3.15'", specifier = ">=6.17.0" },
|
||||||
{ name = "pyright", specifier = ">=1.1.401" },
|
{ name = "pyright", specifier = ">=1.1.401" },
|
||||||
|
{ name = "pytest", specifier = ">=8.3" },
|
||||||
|
{ name = "pytest-asyncio", specifier = ">=0.24" },
|
||||||
{ name = "ruff", specifier = ">=0.11.13" },
|
{ name = "ruff", specifier = ">=0.11.13" },
|
||||||
]
|
]
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user