diff --git a/strix/entry.py b/strix/entry.py index bc4e4c9..8ac3548 100644 --- a/strix/entry.py +++ b/strix/entry.py @@ -18,12 +18,13 @@ from __future__ import annotations import logging import uuid from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Literal from agents import Runner from agents.tracing import add_trace_processor from strix.agents.factory import build_strix_agent, make_child_factory +from strix.config.config import Config from strix.orchestration.bus import AgentMessageBus from strix.orchestration.hooks import StrixOrchestrationHooks from strix.run_config_factory import ( @@ -254,7 +255,6 @@ async def run_strix_scan( caido_host_port=bundle["caido_host_port"], caido_capability=bundle.get("capability"), agent_id=root_id, - agent_name="strix", parent_id=None, tracer=tracer, model=model, @@ -265,10 +265,15 @@ async def run_strix_scan( agent_factory=agent_factory, ) + reasoning = Config.get("strix_reasoning_effort") + reasoning_effort: Literal["low", "medium", "high"] | None = ( + reasoning if reasoning in ("low", "medium", "high") else None # type: ignore[assignment] + ) run_config = make_run_config( sandbox_session=bundle["session"], sandbox_client=bundle["client"], model=model, + reasoning_effort=reasoning_effort, ) task_text = _build_root_task(scan_config) diff --git a/strix/orchestration/hooks.py b/strix/orchestration/hooks.py index 9293247..c479a9e 100644 --- a/strix/orchestration/hooks.py +++ b/strix/orchestration/hooks.py @@ -92,13 +92,9 @@ class StrixOrchestrationHooks(RunHooks[Any]): if details is not None: cached = int(getattr(details, "cached_tokens", 0) or 0) tracer.record_llm_usage( - agent_id=str(agent_id or "unknown"), input_tokens=int(getattr(usage, "input_tokens", 0) or 0), output_tokens=int(getattr(usage, "output_tokens", 0) or 0), cached_tokens=cached, - cost=0.0, - requests=1, - bucket="live", ) ctx["turn_count"] = int(ctx.get("turn_count", 0)) + 1 except Exception: diff --git a/strix/run_config_factory.py b/strix/run_config_factory.py index 42436a3..9c14f8b 100644 --- a/strix/run_config_factory.py +++ b/strix/run_config_factory.py @@ -69,8 +69,6 @@ def make_run_config( *, sandbox_session: BaseSandboxSession | None, model: str = "anthropic/claude-sonnet-4-6", - parallel_tool_calls: bool = _PARALLEL_TOOL_CALLS_DEFAULT, - tool_choice: Literal["auto", "required", "none"] | None = "required", reasoning_effort: Literal["low", "medium", "high"] | None = None, model_settings_override: ModelSettings | None = None, sandbox_client: Any | None = None, @@ -88,9 +86,6 @@ def make_run_config( for unit tests and dry runs. model: Model alias passed to ``MultiProvider``. Defaults to the production Anthropic alias. - parallel_tool_calls: Default ``False`` — the tool server - serializes one task per agent. - tool_choice: Forces tool use per turn unless explicitly relaxed. reasoning_effort: ``"low" | "medium" | "high"``; routes to ``ModelSettings.reasoning``. model_settings_override: Optional per-run ``ModelSettings`` @@ -100,8 +95,8 @@ def make_run_config( supplied without a client. """ base_settings = ModelSettings( - parallel_tool_calls=parallel_tool_calls, - tool_choice=tool_choice, + parallel_tool_calls=_PARALLEL_TOOL_CALLS_DEFAULT, + tool_choice="required", retry=ModelRetrySettings( max_retries=_DEFAULT_MAX_RETRIES, backoff=_DEFAULT_BACKOFF, @@ -113,8 +108,6 @@ def make_run_config( ModelSettings(reasoning=Reasoning(effort=reasoning_effort)), ) if model_settings_override is not None: - # ``ModelSettings.resolve`` merges another ModelSettings into self - # with override-wins semantics — exactly what we want. base_settings = base_settings.resolve(model_settings_override) sandbox_config = ( @@ -142,7 +135,6 @@ def make_agent_context( tool_server_host_port: int | None, caido_host_port: int | None, agent_id: str, - agent_name: str, parent_id: str | None, tracer: Any | None, model: str = "anthropic/claude-sonnet-4-6", @@ -175,7 +167,6 @@ def make_agent_context( "caido_host_port": caido_host_port, "caido_capability": caido_capability, "agent_id": agent_id, - "agent_name": agent_name, "parent_id": parent_id, "tracer": tracer, "model": model, diff --git a/strix/telemetry/tracer.py b/strix/telemetry/tracer.py index f0770bf..88edbde 100644 --- a/strix/telemetry/tracer.py +++ b/strix/telemetry/tracer.py @@ -61,24 +61,13 @@ class Tracer: self.vulnerability_reports: list[dict[str, Any]] = [] self.final_scan_result: str | None = None - # LLM usage roll-up. Two buckets: ``live`` (active agents) and - # ``completed`` (finalized agents — moved here on on_agent_end). - # The orchestration hook chain feeds both via ``record_llm_usage``. - self._llm_stats: dict[str, dict[str, Any]] = { - "live": { - "input_tokens": 0, - "output_tokens": 0, - "cached_tokens": 0, - "cost": 0.0, - "requests": 0, - }, - "completed": { - "input_tokens": 0, - "output_tokens": 0, - "cached_tokens": 0, - "cost": 0.0, - "requests": 0, - }, + # LLM usage roll-up across all agents in this run. + self._llm_stats: dict[str, Any] = { + "input_tokens": 0, + "output_tokens": 0, + "cached_tokens": 0, + "cost": 0.0, + "requests": 0, } self.scan_results: dict[str, Any] | None = None @@ -679,65 +668,35 @@ class Tracer: ) def get_total_llm_stats(self) -> dict[str, Any]: - """Aggregate LLM stats across the live + completed agents. - - Reads ``self._llm_stats`` which the orchestration hooks update - per turn via :meth:`record_llm_usage`. The legacy reach-into- - ``agents_graph_actions`` globals is gone. - """ - completed = self._llm_stats.get("completed", {}) or {} - live = self._llm_stats.get("live", {}) or {} - - total_stats = { - "input_tokens": int(completed.get("input_tokens", 0)) - + int(live.get("input_tokens", 0)), - "output_tokens": int(completed.get("output_tokens", 0)) - + int(live.get("output_tokens", 0)), - "cached_tokens": int(completed.get("cached_tokens", 0)) - + int(live.get("cached_tokens", 0)), - "cost": round( - float(completed.get("cost", 0.0)) + float(live.get("cost", 0.0)), - 4, - ), - "requests": int(completed.get("requests", 0)) + int(live.get("requests", 0)), + """Snapshot the run's aggregated LLM usage.""" + stats = self._llm_stats + total = { + "input_tokens": int(stats["input_tokens"]), + "output_tokens": int(stats["output_tokens"]), + "cached_tokens": int(stats["cached_tokens"]), + "cost": round(float(stats["cost"]), 4), + "requests": int(stats["requests"]), } return { - "total": total_stats, - "total_tokens": total_stats["input_tokens"] + total_stats["output_tokens"], + "total": total, + "total_tokens": total["input_tokens"] + total["output_tokens"], } def record_llm_usage( self, *, - agent_id: str, # noqa: ARG002 input_tokens: int = 0, output_tokens: int = 0, cached_tokens: int = 0, cost: float = 0.0, requests: int = 1, - bucket: str = "live", ) -> None: - """Accumulate LLM usage. Called by the orchestration hooks. - - ``bucket`` is ``"live"`` for in-flight agents and ``"completed"`` - for finalized ones — the SDK's on_agent_end hook moves a child's - running totals from live to completed when it terminates. - """ - target = self._llm_stats.setdefault( - bucket, - { - "input_tokens": 0, - "output_tokens": 0, - "cached_tokens": 0, - "cost": 0.0, - "requests": 0, - }, - ) - target["input_tokens"] += input_tokens - target["output_tokens"] += output_tokens - target["cached_tokens"] += cached_tokens - target["cost"] += cost - target["requests"] += requests + """Accumulate LLM usage from the orchestration hooks.""" + self._llm_stats["input_tokens"] += input_tokens + self._llm_stats["output_tokens"] += output_tokens + self._llm_stats["cached_tokens"] += cached_tokens + self._llm_stats["cost"] += cost + self._llm_stats["requests"] += requests def cleanup(self) -> None: self.save_run_data(mark_complete=True) diff --git a/strix/tools/agents_graph/tools.py b/strix/tools/agents_graph/tools.py index a6d224a..5a76bef 100644 --- a/strix/tools/agents_graph/tools.py +++ b/strix/tools/agents_graph/tools.py @@ -42,6 +42,10 @@ def _dump(result: dict[str, Any]) -> str: return json.dumps(result, ensure_ascii=False, default=str) +def _ctx(ctx: RunContextWrapper) -> dict[str, Any]: + return ctx.context if isinstance(ctx.context, dict) else {} + + @strix_tool(timeout=30) async def view_agent_graph(ctx: RunContextWrapper) -> str: """Print the multi-agent tree — every agent, its parent, its status. @@ -53,7 +57,7 @@ async def view_agent_graph(ctx: RunContextWrapper) -> str: bullet list with status in brackets; the agent that called this tool is marked ``← you``. """ - inner = ctx.context if isinstance(ctx.context, dict) else {} + inner = _ctx(ctx) bus = inner.get("bus") me = inner.get("agent_id") if bus is None: @@ -108,7 +112,7 @@ async def agent_status(ctx: RunContextWrapper, agent_id: str) -> str: agent_id: The 8-char id from ``view_agent_graph`` / ``create_agent``. """ - inner = ctx.context if isinstance(ctx.context, dict) else {} + inner = _ctx(ctx) bus = inner.get("bus") if bus is None: return _dump({"success": False, "error": "Bus not initialized in context."}) @@ -165,7 +169,7 @@ async def send_message_to_agent( expected). Default ``information``. priority: ``low`` / ``normal`` / ``high`` / ``urgent``. """ - inner = ctx.context if isinstance(ctx.context, dict) else {} + inner = _ctx(ctx) bus = inner.get("bus") me = inner.get("agent_id") if bus is None or me is None: @@ -247,7 +251,7 @@ async def wait_for_message( returns and you decide whether to keep working or wait again. """ - inner = ctx.context if isinstance(ctx.context, dict) else {} + inner = _ctx(ctx) bus = inner.get("bus") me = inner.get("agent_id") if bus is None or me is None: @@ -338,7 +342,7 @@ async def create_agent( when starting a clean-slate task. skills: Comma-separated skill names. Max 5; prefer 1-3. """ - inner = ctx.context if isinstance(ctx.context, dict) else {} + inner = _ctx(ctx) bus = inner.get("bus") parent_id = inner.get("agent_id") factory: Callable[..., SDKAgent] | None = inner.get("agent_factory") @@ -408,7 +412,6 @@ async def create_agent( caido_host_port=inner.get("caido_host_port"), caido_capability=inner.get("caido_capability"), agent_id=child_id, - agent_name=name, parent_id=parent_id, tracer=inner.get("tracer"), model=inner.get("model", "anthropic/claude-sonnet-4-6"), @@ -495,7 +498,7 @@ async def agent_finish( parent (e.g., "prioritize testing X", "spawn an agent to cover Y"). """ - inner = ctx.context if isinstance(ctx.context, dict) else {} + inner = _ctx(ctx) bus = inner.get("bus") me = inner.get("agent_id") if bus is None or me is None: diff --git a/strix/tools/notes/__init__.py b/strix/tools/notes/__init__.py index 4f2c09a..287e6e4 100644 --- a/strix/tools/notes/__init__.py +++ b/strix/tools/notes/__init__.py @@ -1,15 +1,7 @@ -from .tools import ( - append_note_content, - create_note, - delete_note, - get_note, - list_notes, - update_note, -) +from .tools import create_note, delete_note, get_note, list_notes, update_note __all__ = [ - "append_note_content", "create_note", "delete_note", "get_note", diff --git a/strix/tools/notes/tools.py b/strix/tools/notes/tools.py index 92cd975..b8ead6c 100644 --- a/strix/tools/notes/tools.py +++ b/strix/tools/notes/tools.py @@ -246,7 +246,7 @@ def _create_note_impl( # noqa: PLR0911 category: str = "general", tags: list[str] | None = None, ) -> dict[str, Any]: - """Create one note. Public — used by ``append_note_content`` and tests.""" + """Create one note. Public — used by tests.""" with _notes_lock: try: _ensure_notes_loaded() @@ -398,22 +398,6 @@ def _delete_note_impl(note_id: str) -> dict[str, Any]: } -def append_note_content(note_id: str, delta: str) -> dict[str, Any]: - """Append text to an existing note's content. Used by the agents-graph - wiki-update hook on agent_finish.""" - with _notes_lock: - try: - _ensure_notes_loaded() - if note_id not in _notes_storage: - return {"success": False, "error": f"Note with ID '{note_id}' not found"} - note = _notes_storage[note_id] - existing = str(note.get("content") or "") - updated = f"{existing.rstrip()}{delta}" - return _update_note_impl(note_id=note_id, content=updated) - except (ValueError, TypeError) as e: - return {"success": False, "error": f"Failed to append note content: {e}"} - - # --- public tools --------------------------------------------------------- @@ -457,7 +441,6 @@ async def create_note( category: One of the categories above. Default ``"general"``. tags: Optional free-form tags. """ - del ctx return _dump( await asyncio.to_thread(_create_note_impl, title, content, category, tags), ) @@ -489,7 +472,6 @@ async def list_notes( include_content: When False (default) entries have a preview; when True the full ``content`` is included. """ - del ctx return _dump( await asyncio.to_thread( _list_notes_impl, @@ -508,7 +490,6 @@ async def get_note(ctx: RunContextWrapper, note_id: str) -> str: Args: note_id: Note id from ``create_note`` or a ``list_notes`` entry. """ - del ctx return _dump(await asyncio.to_thread(_get_note_impl, note_id)) @@ -532,7 +513,6 @@ async def update_note( content: New content, or ``None`` to keep. tags: New tags list, or ``None`` to keep. """ - del ctx return _dump( await asyncio.to_thread( _update_note_impl, @@ -551,5 +531,4 @@ async def delete_note(ctx: RunContextWrapper, note_id: str) -> str: Args: note_id: Note id to delete. """ - del ctx return _dump(await asyncio.to_thread(_delete_note_impl, note_id)) diff --git a/strix/tools/reporting/tool.py b/strix/tools/reporting/tool.py index 2c34fa8..a9d6566 100644 --- a/strix/tools/reporting/tool.py +++ b/strix/tools/reporting/tool.py @@ -399,7 +399,6 @@ async def create_vulnerability_report( ``fix_before`` (verbatim source), ``fix_after`` (suggested replacement). """ - del ctx result = await asyncio.to_thread( _do_create, title=title, diff --git a/strix/tools/todo/tools.py b/strix/tools/todo/tools.py index 04986a2..990e3cb 100644 --- a/strix/tools/todo/tools.py +++ b/strix/tools/todo/tools.py @@ -21,6 +21,17 @@ from strix.tools._decorator import strix_tool VALID_PRIORITIES = ["low", "normal", "high", "critical"] VALID_STATUSES = ["pending", "in_progress", "done"] +_PRIORITY_RANK = {"critical": 0, "high": 1, "normal": 2, "low": 3} +_STATUS_RANK = {"done": 0, "in_progress": 1, "pending": 2} + + +def _todo_sort_key(todo: dict[str, Any]) -> tuple[int, int, str]: + return ( + _STATUS_RANK.get(todo.get("status", "pending"), 99), + _PRIORITY_RANK.get(todo.get("priority", "normal"), 99), + todo.get("created_at", ""), + ) + # Per-agent silo: ``_todos_storage[agent_id][todo_id] = todo_dict``. # Keyed by ``ctx.context['agent_id']`` so two agents in the same scan @@ -49,22 +60,10 @@ def _normalize_priority(priority: str | None, default: str = "normal") -> str: def _sorted_todos(agent_id: str) -> list[dict[str, Any]]: - agent_todos = _get_agent_todos(agent_id) - todos_list: list[dict[str, Any]] = [] - for todo_id, todo in agent_todos.items(): - entry = todo.copy() - entry["todo_id"] = todo_id - todos_list.append(entry) - - priority_order = {"critical": 0, "high": 1, "normal": 2, "low": 3} - status_order = {"done": 0, "in_progress": 1, "pending": 2} - todos_list.sort( - key=lambda x: ( - status_order.get(x.get("status", "pending"), 99), - priority_order.get(x.get("priority", "normal"), 99), - x.get("created_at", ""), - ), - ) + todos_list = [ + {**todo, "todo_id": todo_id} for todo_id, todo in _get_agent_todos(agent_id).items() + ] + todos_list.sort(key=_todo_sort_key) return todos_list @@ -323,15 +322,7 @@ async def list_todos( entry["todo_id"] = todo_id todos_list.append(entry) - priority_order = {"critical": 0, "high": 1, "normal": 2, "low": 3} - status_order = {"done": 0, "in_progress": 1, "pending": 2} - todos_list.sort( - key=lambda x: ( - status_order.get(x.get("status", "pending"), 99), - priority_order.get(x.get("priority", "normal"), 99), - x.get("created_at", ""), - ), - ) + todos_list.sort(key=_todo_sort_key) summary: dict[str, int] = {"pending": 0, "in_progress": 0, "done": 0} for todo in todos_list: diff --git a/strix/tools/web_search/tool.py b/strix/tools/web_search/tool.py index 9722b5a..be0628c 100644 --- a/strix/tools/web_search/tool.py +++ b/strix/tools/web_search/tool.py @@ -121,6 +121,5 @@ async def web_search(ctx: RunContextWrapper, query: str) -> str: target tech, and the specific question. Treat it like a ticket title for a senior security engineer. """ - del ctx result = await asyncio.to_thread(_do_search, query) return json.dumps(result, ensure_ascii=False, default=str) diff --git a/tests/telemetry/test_tracer.py b/tests/telemetry/test_tracer.py index 107a542..96dff63 100644 --- a/tests/telemetry/test_tracer.py +++ b/tests/telemetry/test_tracer.py @@ -258,32 +258,26 @@ def test_events_with_agent_id_include_agent_name( assert chat_event["actor"]["agent_name"] == "Root Agent" -def test_get_total_llm_stats_aggregates_live_and_completed( +def test_get_total_llm_stats_aggregates_recordings( monkeypatch: pytest.MonkeyPatch, tmp_path: Path ) -> None: monkeypatch.chdir(tmp_path) tracer = Tracer("cost-rollup") set_global_tracer(tracer) - # Live agent (still running). tracer.record_llm_usage( - agent_id="root-agent", input_tokens=1_000, output_tokens=250, cached_tokens=100, cost=0.12345, requests=2, - bucket="live", ) - # Completed agents (finalized — moved by on_agent_end hook). tracer.record_llm_usage( - agent_id="child-1", input_tokens=2_000, output_tokens=500, cached_tokens=400, cost=0.54321, requests=3, - bucket="completed", ) stats = tracer.get_total_llm_stats() diff --git a/tests/test_run_config_factory.py b/tests/test_run_config_factory.py index d3fc903..05693e9 100644 --- a/tests/test_run_config_factory.py +++ b/tests/test_run_config_factory.py @@ -107,7 +107,6 @@ def test_max_turns_default_is_300() -> None: tool_server_host_port=None, caido_host_port=None, agent_id="root", - agent_name="root", parent_id=None, tracer=None, ) @@ -124,7 +123,6 @@ def test_make_agent_context_full_shape() -> None: tool_server_host_port=48081, caido_host_port=48080, agent_id="agent-1", - agent_name="root", parent_id=None, tracer="not-a-real-tracer", is_whitebox=True, @@ -154,7 +152,6 @@ def test_make_agent_context_is_whitebox_defaults_false() -> None: tool_server_host_port=None, caido_host_port=None, agent_id="r", - agent_name="root", parent_id=None, tracer=None, ) diff --git a/tests/tools/test_notes_wiki.py b/tests/tools/test_notes_wiki.py index 581be72..2db4ab7 100644 --- a/tests/tools/test_notes_wiki.py +++ b/tests/tools/test_notes_wiki.py @@ -1,5 +1,7 @@ from pathlib import Path +import pytest + from strix.telemetry.tracer import Tracer, get_global_tracer, set_global_tracer from strix.tools.notes import tools as notes_actions @@ -9,7 +11,9 @@ def _reset_notes_state() -> None: notes_actions._loaded_notes_run_dir = None -def test_wiki_notes_are_persisted_and_removed(tmp_path: Path, monkeypatch) -> None: +def test_wiki_notes_are_persisted_and_removed( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: monkeypatch.chdir(tmp_path) _reset_notes_state() @@ -48,10 +52,12 @@ def test_wiki_notes_are_persisted_and_removed(tmp_path: Path, monkeypatch) -> No assert wiki_path.exists() is False finally: _reset_notes_state() - set_global_tracer(previous_tracer) # type: ignore[arg-type] + set_global_tracer(previous_tracer) -def test_notes_jsonl_replay_survives_memory_reset(tmp_path: Path, monkeypatch) -> None: +def test_notes_jsonl_replay_survives_memory_reset( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: monkeypatch.chdir(tmp_path) _reset_notes_state() @@ -108,10 +114,10 @@ def test_notes_jsonl_replay_survives_memory_reset(tmp_path: Path, monkeypatch) - assert listed_after_delete["total_count"] == 0 finally: _reset_notes_state() - set_global_tracer(previous_tracer) # type: ignore[arg-type] + set_global_tracer(previous_tracer) -def test_get_note_returns_full_note(tmp_path: Path, monkeypatch) -> None: +def test_get_note_returns_full_note(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.chdir(tmp_path) _reset_notes_state() @@ -136,44 +142,11 @@ def test_get_note_returns_full_note(tmp_path: Path, monkeypatch) -> None: assert result["note"]["content"] == "entrypoints and sinks" finally: _reset_notes_state() - set_global_tracer(previous_tracer) # type: ignore[arg-type] - - -def test_append_note_content_appends_delta(tmp_path: Path, monkeypatch) -> None: - monkeypatch.chdir(tmp_path) - _reset_notes_state() - - previous_tracer = get_global_tracer() - tracer = Tracer("append-note-run") - set_global_tracer(tracer) - - try: - created = notes_actions._create_note_impl( - title="Repo wiki", - content="base", - category="wiki", - tags=["repo:demo"], - ) - assert created["success"] is True - note_id = created["note_id"] - assert isinstance(note_id, str) - - appended = notes_actions.append_note_content( - note_id=note_id, - delta="\n\n## Agent Update: worker\nSummary: done", - ) - assert appended["success"] is True - - loaded = notes_actions._get_note_impl(note_id=note_id) - assert loaded["success"] is True - assert loaded["note"]["content"] == "base\n\n## Agent Update: worker\nSummary: done" - finally: - _reset_notes_state() - set_global_tracer(previous_tracer) # type: ignore[arg-type] + set_global_tracer(previous_tracer) def test_list_and_get_note_handle_wiki_repersist_oserror_gracefully( - tmp_path: Path, monkeypatch + tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.chdir(tmp_path) _reset_notes_state() @@ -195,7 +168,7 @@ def test_list_and_get_note_handle_wiki_repersist_oserror_gracefully( _reset_notes_state() - def _raise_oserror(*_args, **_kwargs) -> None: + def _raise_oserror(*_args: object, **_kwargs: object) -> None: raise OSError("disk full") monkeypatch.setattr(notes_actions, "_persist_wiki_note", _raise_oserror) @@ -211,4 +184,4 @@ def test_list_and_get_note_handle_wiki_repersist_oserror_gracefully( assert fetched["note"]["content"] == "initial wiki content" finally: _reset_notes_state() - set_global_tracer(previous_tracer) # type: ignore[arg-type] + set_global_tracer(previous_tracer)