From d8881498ee16d499ac00bfec0ab05cc9d6d940ca Mon Sep 17 00:00:00 2001 From: 0xallam Date: Sat, 25 Apr 2026 09:30:23 -0700 Subject: [PATCH] refactor: nuke legacy harness, drop sdk_ prefixes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The SDK harness is the only path now; legacy host-side code is gone. File names no longer carry the ``sdk_`` distinction. Deleted legacy host-side modules: - strix/agents/StrixAgent/ (template moved to strix/agents/prompts/) - strix/agents/base_agent.py, state.py - strix/llm/llm.py, config.py - strix/runtime/docker_runtime.py, runtime.py - strix/tools/executor.py, agents_graph/agents_graph_actions.py - strix/interface/sdk_dispatch.py + the env-flag dispatch in cli.py Renamed (drop ``sdk_`` prefix): - strix/sdk_entry.py → strix/entry.py - strix/agents/sdk_factory.py → strix/agents/factory.py - strix/agents/sdk_prompt.py → strix/agents/prompt.py - strix/tools//_sdk_tool[s].py → strix/tools//tool[s].py - strix/tools/_legacy_adapter.py → strix/tools/_state_adapter.py - ``_legacy`` aliases inside the wrappers → ``_impl`` CLI + TUI now call ``run_strix_scan`` directly — they build the sandbox image / sources_path locally and rely on ``session_manager.cleanup`` (called inside ``run_strix_scan``'s finally) for teardown. Three TUI handlers that reached into legacy multi-agent globals (``_agent_instances``, ``send_user_message_to_agent``, ``stop_agent``) are now no-ops with a TODO; reconnecting them to the ``AgentMessageBus`` is a follow-up. Tracer.get_total_llm_stats no longer reaches into the deleted ``agents_graph_actions`` globals — the orchestration hooks now feed the tracer via ``Tracer.record_llm_usage`` (live + completed buckets). finish_scan's ``_check_active_agents`` and load_skill's runtime ``_agent_instances`` reach-in are no-op stubs; the ``AgentMessageBus`` is the source of truth post-migration. llm/utils.py rewritten to keep only the streaming-parser helpers (``normalize_tool_format``, ``parse_tool_invocations``, ``fix_incomplete_tool_call``, ``format_tool_call``, ``clean_content``). ``STRIX_MODEL_MAP`` moved to ``llm/multi_provider_setup.py`` (its only remaining caller). Per-file ruff ignores added for legacy interface modules (TUI / main / CLI / utils / streaming_parser / tool_components) and tracer.py — pre-existing PLC0415/BLE001/PLR0915 patterns are out of scope. Tests: 287/287 passing. Renamed test files to drop ``sdk_`` prefix. ``test_tracer.py::test_get_total_llm_stats_aggregates_live_and_completed`` rewritten to feed ``Tracer.record_llm_usage`` instead of legacy globals. Test file annotations added so pre-commit's strict mypy passes. Co-Authored-By: Claude Opus 4.7 (1M context) --- pyproject.toml | 46 +- strix/agents/StrixAgent/__init__.py | 4 - strix/agents/StrixAgent/strix_agent.py | 151 ---- strix/agents/__init__.py | 21 +- strix/agents/base_agent.py | 623 ------------- strix/agents/{sdk_factory.py => factory.py} | 28 +- strix/agents/{sdk_prompt.py => prompt.py} | 40 +- .../system_prompt.jinja | 0 strix/agents/state.py | 172 ---- strix/{sdk_entry.py => entry.py} | 43 +- strix/interface/cli.py | 98 +- strix/interface/main.py | 25 +- strix/interface/sdk_dispatch.py | 143 --- strix/interface/tui.py | 123 ++- strix/interface/utils.py | 26 +- strix/llm/__init__.py | 20 +- strix/llm/anthropic_cache_wrapper.py | 4 +- strix/llm/config.py | 40 - strix/llm/dedupe.py | 9 +- strix/llm/llm.py | 390 -------- strix/llm/multi_provider_setup.py | 20 +- strix/llm/strix_session.py | 4 +- strix/llm/utils.py | 61 +- strix/orchestration/bus.py | 6 +- strix/orchestration/filter.py | 17 +- strix/orchestration/hooks.py | 5 +- strix/run_config_factory.py | 31 +- strix/runtime/__init__.py | 50 +- strix/runtime/docker_runtime.py | 352 -------- strix/runtime/runtime.py | 33 - strix/sandbox/caido_capability.py | 4 +- strix/sandbox/healthcheck.py | 10 +- strix/sandbox/session_manager.py | 12 +- strix/telemetry/strix_processor.py | 14 +- strix/telemetry/tracer.py | 98 +- strix/tools/__init__.py | 29 +- strix/tools/_legacy_adapter.py | 57 -- strix/tools/_state_adapter.py | 47 + strix/tools/agents_graph/__init__.py | 4 +- .../agents_graph/agents_graph_actions.py | 839 ------------------ .../{agents_graph_sdk_tools.py => tools.py} | 0 .../browser/{browser_sdk_tool.py => tool.py} | 0 strix/tools/executor.py | 364 -------- .../{file_edit_sdk_tools.py => tools.py} | 0 strix/tools/finish/finish_actions.py | 73 +- .../finish/{finish_sdk_tool.py => tool.py} | 6 +- strix/tools/load_skill/load_skill_actions.py | 31 +- .../{load_skill_sdk_tool.py => tool.py} | 6 +- .../notes/{notes_sdk_tools.py => tools.py} | 12 +- .../proxy/{proxy_sdk_tools.py => tools.py} | 0 .../python/{python_sdk_tool.py => tool.py} | 0 .../{reporting_sdk_tools.py => tool.py} | 4 +- .../{terminal_sdk_tool.py => tool.py} | 0 .../{thinking_sdk_tools.py => tool.py} | 0 .../todo/{todo_sdk_tools.py => tools.py} | 16 +- .../{web_search_sdk_tool.py => tool.py} | 4 +- .../{test_sdk_factory.py => test_factory.py} | 8 +- tests/interface/test_sdk_dispatch.py | 199 ----- tests/llm/test_llm_otel.py | 16 - tests/llm/test_source_aware_whitebox.py | 30 - tests/telemetry/test_tracer.py | 131 ++- tests/{test_sdk_entry.py => test_entry.py} | 42 +- tests/tools/test_agents_graph_whitebox.py | 291 ------ ...sdk_graph_tools.py => test_graph_tools.py} | 8 +- tests/tools/test_load_skill_tool.py | 139 --- ...sdk_local_tools.py => test_local_tools.py} | 14 +- ...tools.py => test_remaining_local_tools.py} | 53 +- ...sandbox_tools.py => test_sandbox_tools.py} | 28 +- tests/tools/test_tool_registration_modes.py | 9 +- 69 files changed, 646 insertions(+), 4537 deletions(-) delete mode 100644 strix/agents/StrixAgent/__init__.py delete mode 100644 strix/agents/StrixAgent/strix_agent.py delete mode 100644 strix/agents/base_agent.py rename strix/agents/{sdk_factory.py => factory.py} (88%) rename strix/agents/{sdk_prompt.py => prompt.py} (68%) rename strix/agents/{StrixAgent => prompts}/system_prompt.jinja (100%) delete mode 100644 strix/agents/state.py rename strix/{sdk_entry.py => entry.py} (85%) delete mode 100644 strix/interface/sdk_dispatch.py delete mode 100644 strix/llm/config.py delete mode 100644 strix/llm/llm.py delete mode 100644 strix/runtime/docker_runtime.py delete mode 100644 strix/runtime/runtime.py delete mode 100644 strix/tools/_legacy_adapter.py create mode 100644 strix/tools/_state_adapter.py delete mode 100644 strix/tools/agents_graph/agents_graph_actions.py rename strix/tools/agents_graph/{agents_graph_sdk_tools.py => tools.py} (100%) rename strix/tools/browser/{browser_sdk_tool.py => tool.py} (100%) delete mode 100644 strix/tools/executor.py rename strix/tools/file_edit/{file_edit_sdk_tools.py => tools.py} (100%) rename strix/tools/finish/{finish_sdk_tool.py => tool.py} (93%) rename strix/tools/load_skill/{load_skill_sdk_tool.py => tool.py} (87%) rename strix/tools/notes/{notes_sdk_tools.py => tools.py} (91%) rename strix/tools/proxy/{proxy_sdk_tools.py => tools.py} (100%) rename strix/tools/python/{python_sdk_tool.py => tool.py} (100%) rename strix/tools/reporting/{reporting_sdk_tools.py => tool.py} (96%) rename strix/tools/terminal/{terminal_sdk_tool.py => tool.py} (100%) rename strix/tools/thinking/{thinking_sdk_tools.py => tool.py} (100%) rename strix/tools/todo/{todo_sdk_tools.py => tools.py} (89%) rename strix/tools/web_search/{web_search_sdk_tool.py => tool.py} (90%) rename tests/agents/{test_sdk_factory.py => test_factory.py} (96%) delete mode 100644 tests/interface/test_sdk_dispatch.py delete mode 100644 tests/llm/test_llm_otel.py delete mode 100644 tests/llm/test_source_aware_whitebox.py rename tests/{test_sdk_entry.py => test_entry.py} (86%) delete mode 100644 tests/tools/test_agents_graph_whitebox.py rename tests/tools/{test_sdk_graph_tools.py => test_graph_tools.py} (98%) delete mode 100644 tests/tools/test_load_skill_tool.py rename tests/tools/{test_sdk_local_tools.py => test_local_tools.py} (94%) rename tests/tools/{test_sdk_remaining_local_tools.py => test_remaining_local_tools.py} (82%) rename tests/tools/{test_sdk_sandbox_tools.py => test_sandbox_tools.py} (91%) diff --git a/pyproject.toml b/pyproject.toml index e9cfa3d..73dfb8e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -286,20 +286,22 @@ ignore = [ ] # SDK function-tool wrappers: the SDK calls get_type_hints() at registration # time to derive the JSON schema, which evaluates annotations at runtime — -# so RunContextWrapper must be imported eagerly, not under TYPE_CHECKING. -"strix/tools/todo/todo_sdk_tools.py" = ["TC002"] -"strix/tools/notes/notes_sdk_tools.py" = ["TC002"] -"strix/tools/thinking/thinking_sdk_tools.py" = ["TC002"] -"strix/tools/web_search/web_search_sdk_tool.py" = ["TC002"] -"strix/tools/file_edit/file_edit_sdk_tools.py" = ["TC002"] -"strix/tools/reporting/reporting_sdk_tools.py" = ["TC002"] -"strix/tools/load_skill/load_skill_sdk_tool.py" = ["TC002"] -"strix/tools/finish/finish_sdk_tool.py" = ["TC002"] -"strix/tools/browser/browser_sdk_tool.py" = ["TC002"] -"strix/tools/terminal/terminal_sdk_tool.py" = ["TC002"] -"strix/tools/python/python_sdk_tool.py" = ["TC002"] -"strix/tools/proxy/proxy_sdk_tools.py" = ["TC002"] -"strix/tools/agents_graph/agents_graph_sdk_tools.py" = ["TC002"] +# so RunContextWrapper / Tool / TResponseInputItem must be imported eagerly, +# not under TYPE_CHECKING. +"strix/tools/todo/tools.py" = ["TC002"] +"strix/tools/notes/tools.py" = ["TC002"] +"strix/tools/thinking/tool.py" = ["TC002"] +"strix/tools/web_search/tool.py" = ["TC002"] +"strix/tools/file_edit/tools.py" = ["TC002"] +"strix/tools/reporting/tool.py" = ["TC002"] +"strix/tools/load_skill/tool.py" = ["TC002"] +"strix/tools/finish/tool.py" = ["TC002"] +"strix/tools/browser/tool.py" = ["TC002"] +"strix/tools/terminal/tool.py" = ["TC002"] +"strix/tools/python/tool.py" = ["TC002"] +"strix/tools/proxy/tools.py" = ["TC002"] +"strix/tools/agents_graph/tools.py" = ["TC002"] +"strix/agents/factory.py" = ["TC002"] # CaidoCapability uses agents.tool.Tool at runtime — pydantic Field # annotations and the cached _CAIDO_TOOLS tuple need it eagerly. "strix/sandbox/caido_capability.py" = ["TC002"] @@ -311,11 +313,23 @@ ignore = [ # resolution past where mypy needs it. ``_build_root_task`` legitimately # walks every supported target type — splitting it into per-type # helpers would add indirection without simplifying anything. -"strix/sdk_entry.py" = ["TC003", "PLR0912"] +"strix/entry.py" = ["TC003", "PLR0912"] +# Legacy tracer module — pre-existing PLR/E501 patterns; full refactor +# is out of scope for the harness migration. ``Callable`` is a runtime +# annotation on ``vulnerability_found_callback``. +"strix/telemetry/tracer.py" = ["TC003", "PLR0912", "PLR0915", "E501"] # Legacy interface utility with intentionally many branches per supported # scope-mode / target-type combination; refactor would obscure the # decision tree without simplifying it. -"strix/interface/utils.py" = ["PLR0912"] +"strix/interface/utils.py" = ["PLR0912", "BLE001", "PLC0415"] +# CLI / TUI / main keep extensive lazy imports + broad exception swallows +# for resilience around terminal-rendering errors. Refactor is out of +# scope for the harness migration. +"strix/interface/cli.py" = ["BLE001", "PLC0415"] +"strix/interface/tui.py" = ["BLE001", "PLC0415", "PLR0912", "PLR0915", "SIM105"] +"strix/interface/main.py" = ["BLE001", "PLC0415", "PLR0912", "PLR0915"] +"strix/interface/streaming_parser.py" = ["PLC0415"] +"strix/interface/tool_components/agent_message_renderer.py" = ["PLC0415"] # Sandbox dispatch helper has many short-circuit error returns (auth fail, # size cap, decode fail, etc). Each is a distinct, documented failure mode # the model needs to see verbatim — collapsing them harms readability. diff --git a/strix/agents/StrixAgent/__init__.py b/strix/agents/StrixAgent/__init__.py deleted file mode 100644 index fa291ed..0000000 --- a/strix/agents/StrixAgent/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .strix_agent import StrixAgent - - -__all__ = ["StrixAgent"] diff --git a/strix/agents/StrixAgent/strix_agent.py b/strix/agents/StrixAgent/strix_agent.py deleted file mode 100644 index 36e3594..0000000 --- a/strix/agents/StrixAgent/strix_agent.py +++ /dev/null @@ -1,151 +0,0 @@ -from typing import Any - -from strix.agents.base_agent import BaseAgent -from strix.llm.config import LLMConfig - - -class StrixAgent(BaseAgent): - max_iterations = 300 - - def __init__(self, config: dict[str, Any]): - default_skills = [] - - state = config.get("state") - if state is None or (hasattr(state, "parent_id") and state.parent_id is None): - default_skills = ["root_agent"] - - self.default_llm_config = LLMConfig(skills=default_skills) - - super().__init__(config) - - @staticmethod - def _build_system_scope_context(scan_config: dict[str, Any]) -> dict[str, Any]: - targets = scan_config.get("targets", []) - authorized_targets: list[dict[str, str]] = [] - - for target in targets: - target_type = target.get("type", "unknown") - details = target.get("details", {}) - - if target_type == "repository": - value = details.get("target_repo", "") - elif target_type == "local_code": - value = details.get("target_path", "") - elif target_type == "web_application": - value = details.get("target_url", "") - elif target_type == "ip_address": - value = details.get("target_ip", "") - else: - value = target.get("original", "") - - workspace_subdir = details.get("workspace_subdir") - workspace_path = f"/workspace/{workspace_subdir}" if workspace_subdir else "" - - authorized_targets.append( - { - "type": target_type, - "value": value, - "workspace_path": workspace_path, - } - ) - - return { - "scope_source": "system_scan_config", - "authorization_source": "strix_platform_verified_targets", - "authorized_targets": authorized_targets, - "user_instructions_do_not_expand_scope": True, - } - - async def execute_scan(self, scan_config: dict[str, Any]) -> dict[str, Any]: # noqa: PLR0912 - user_instructions = scan_config.get("user_instructions", "") - targets = scan_config.get("targets", []) - diff_scope = scan_config.get("diff_scope", {}) or {} - self.llm.set_system_prompt_context(self._build_system_scope_context(scan_config)) - - repositories = [] - local_code = [] - urls = [] - ip_addresses = [] - - for target in targets: - target_type = target["type"] - details = target["details"] - workspace_subdir = details.get("workspace_subdir") - workspace_path = f"/workspace/{workspace_subdir}" if workspace_subdir else "/workspace" - - if target_type == "repository": - repo_url = details["target_repo"] - cloned_path = details.get("cloned_repo_path") - repositories.append( - { - "url": repo_url, - "workspace_path": workspace_path if cloned_path else None, - } - ) - - elif target_type == "local_code": - original_path = details.get("target_path", "unknown") - local_code.append( - { - "path": original_path, - "workspace_path": workspace_path, - } - ) - - elif target_type == "web_application": - urls.append(details["target_url"]) - elif target_type == "ip_address": - ip_addresses.append(details["target_ip"]) - - task_parts = [] - - if repositories: - task_parts.append("\n\nRepositories:") - for repo in repositories: - if repo["workspace_path"]: - task_parts.append(f"- {repo['url']} (available at: {repo['workspace_path']})") - else: - task_parts.append(f"- {repo['url']}") - - if local_code: - task_parts.append("\n\nLocal Codebases:") - task_parts.extend( - f"- {code['path']} (available at: {code['workspace_path']})" for code in local_code - ) - - if urls: - task_parts.append("\n\nURLs:") - task_parts.extend(f"- {url}" for url in urls) - - if ip_addresses: - task_parts.append("\n\nIP Addresses:") - task_parts.extend(f"- {ip}" for ip in ip_addresses) - - if diff_scope.get("active"): - task_parts.append("\n\nScope Constraints:") - task_parts.append( - "- Pull request diff-scope mode is active. Prioritize changed files " - "and use other files only for context." - ) - for repo_scope in diff_scope.get("repos", []): - repo_label = ( - repo_scope.get("workspace_subdir") - or repo_scope.get("source_path") - or "repository" - ) - changed_count = repo_scope.get("analyzable_files_count", 0) - deleted_count = repo_scope.get("deleted_files_count", 0) - task_parts.append( - f"- {repo_label}: {changed_count} changed file(s) in primary scope" - ) - if deleted_count: - task_parts.append( - f"- {repo_label}: {deleted_count} deleted file(s) are context-only" - ) - - task_description = " ".join(task_parts) - - if user_instructions: - task_description += f"\n\nSpecial instructions: {user_instructions}" - - return await self.agent_loop(task=task_description) diff --git a/strix/agents/__init__.py b/strix/agents/__init__.py index c7e542e..c80aa04 100644 --- a/strix/agents/__init__.py +++ b/strix/agents/__init__.py @@ -1,10 +1,19 @@ -from .base_agent import BaseAgent -from .state import AgentState -from .StrixAgent import StrixAgent +"""Strix agent package. + +Public surface: + +- :func:`build_strix_agent` — assemble a root or child ``agents.Agent``. +- :func:`make_child_factory` — closure factory passed via context to + the multi-agent ``create_agent`` graph tool. +- :func:`render_system_prompt` — render the Jinja system prompt. +""" + +from .factory import build_strix_agent, make_child_factory +from .prompt import render_system_prompt __all__ = [ - "AgentState", - "BaseAgent", - "StrixAgent", + "build_strix_agent", + "make_child_factory", + "render_system_prompt", ] diff --git a/strix/agents/base_agent.py b/strix/agents/base_agent.py deleted file mode 100644 index c759f9a..0000000 --- a/strix/agents/base_agent.py +++ /dev/null @@ -1,623 +0,0 @@ -import asyncio -import contextlib -import logging -from typing import TYPE_CHECKING, Any, Optional - - -if TYPE_CHECKING: - from strix.telemetry.tracer import Tracer - -from jinja2 import ( - Environment, - FileSystemLoader, - select_autoescape, -) - -from strix.llm import LLM, LLMConfig, LLMRequestFailedError -from strix.llm.utils import clean_content -from strix.runtime import SandboxInitializationError -from strix.tools import process_tool_invocations -from strix.utils.resource_paths import get_strix_resource_path - -from .state import AgentState - - -logger = logging.getLogger(__name__) - - -class AgentMeta(type): - agent_name: str - jinja_env: Environment - - def __new__(cls, name: str, bases: tuple[type, ...], attrs: dict[str, Any]) -> type: - new_cls = super().__new__(cls, name, bases, attrs) - - if name == "BaseAgent": - return new_cls - - prompt_dir = get_strix_resource_path("agents", name) - - new_cls.agent_name = name - new_cls.jinja_env = Environment( - loader=FileSystemLoader(prompt_dir), - autoescape=select_autoescape(enabled_extensions=(), default_for_string=False), - ) - - return new_cls - - -class BaseAgent(metaclass=AgentMeta): - max_iterations = 300 - agent_name: str = "" - jinja_env: Environment - default_llm_config: LLMConfig | None = None - - def __init__(self, config: dict[str, Any]): - self.config = config - - self.local_sources = config.get("local_sources", []) - - if "max_iterations" in config: - self.max_iterations = config["max_iterations"] - - self.llm_config_name = config.get("llm_config_name", "default") - self.llm_config = config.get("llm_config", self.default_llm_config) - if self.llm_config is None: - raise ValueError("llm_config is required but not provided") - state_from_config = config.get("state") - if state_from_config is not None: - self.state = state_from_config - else: - self.state = AgentState( - agent_name="Root Agent", - max_iterations=self.max_iterations, - ) - - self.interactive = getattr(self.llm_config, "interactive", False) - if self.interactive and self.state.parent_id is None: - self.state.waiting_timeout = 0 - self.llm = LLM(self.llm_config, agent_name=self.agent_name) - - with contextlib.suppress(Exception): - self.llm.set_agent_identity(self.state.agent_name, self.state.agent_id) - self._current_task: asyncio.Task[Any] | None = None - self._force_stop = False - - from strix.telemetry.tracer import get_global_tracer - - tracer = get_global_tracer() - if tracer: - tracer.log_agent_creation( - agent_id=self.state.agent_id, - name=self.state.agent_name, - task=self.state.task, - parent_id=self.state.parent_id, - ) - if self.state.parent_id is None: - scan_config = tracer.scan_config or {} - exec_id = tracer.log_tool_execution_start( - agent_id=self.state.agent_id, - tool_name="scan_start_info", - args=scan_config, - ) - tracer.update_tool_execution(execution_id=exec_id, status="completed", result={}) - - else: - exec_id = tracer.log_tool_execution_start( - agent_id=self.state.agent_id, - tool_name="subagent_start_info", - args={ - "name": self.state.agent_name, - "task": self.state.task, - "parent_id": self.state.parent_id, - }, - ) - tracer.update_tool_execution(execution_id=exec_id, status="completed", result={}) - - self._add_to_agents_graph() - - def _add_to_agents_graph(self) -> None: - from strix.tools.agents_graph import agents_graph_actions - - node = { - "id": self.state.agent_id, - "name": self.state.agent_name, - "task": self.state.task, - "status": "running", - "parent_id": self.state.parent_id, - "created_at": self.state.start_time, - "finished_at": None, - "result": None, - "llm_config": self.llm_config_name, - "agent_type": self.__class__.__name__, - "state": self.state.model_dump(), - } - agents_graph_actions._agent_graph["nodes"][self.state.agent_id] = node - - with agents_graph_actions._agent_llm_stats_lock: - agents_graph_actions._agent_instances[self.state.agent_id] = self - agents_graph_actions._agent_states[self.state.agent_id] = self.state - - if self.state.parent_id: - agents_graph_actions._agent_graph["edges"].append( - {"from": self.state.parent_id, "to": self.state.agent_id, "type": "delegation"} - ) - - if self.state.agent_id not in agents_graph_actions._agent_messages: - agents_graph_actions._agent_messages[self.state.agent_id] = [] - - if self.state.parent_id is None and agents_graph_actions._root_agent_id is None: - agents_graph_actions._root_agent_id = self.state.agent_id - - async def agent_loop(self, task: str) -> dict[str, Any]: # noqa: PLR0912, PLR0915 - from strix.telemetry.tracer import get_global_tracer - - tracer = get_global_tracer() - - try: - await self._initialize_sandbox_and_state(task) - except SandboxInitializationError as e: - return self._handle_sandbox_error(e, tracer) - - while True: - if self._force_stop: - self._force_stop = False - await self._enter_waiting_state(tracer, was_cancelled=True) - continue - - self._check_agent_messages(self.state) - - if self.state.is_waiting_for_input(): - await self._wait_for_input() - continue - - if self.state.should_stop(): - if not self.interactive: - return self.state.final_result or {} - await self._enter_waiting_state(tracer) - continue - - if self.state.llm_failed: - await self._wait_for_input() - continue - - self.state.increment_iteration() - - if ( - self.state.is_approaching_max_iterations() - and not self.state.max_iterations_warning_sent - ): - self.state.max_iterations_warning_sent = True - remaining = self.state.max_iterations - self.state.iteration - warning_msg = ( - f"URGENT: You are approaching the maximum iteration limit. " - f"Current: {self.state.iteration}/{self.state.max_iterations} " - f"({remaining} iterations remaining). " - f"Please prioritize completing your required task(s) and calling " - f"the appropriate finish tool (finish_scan for root agent, " - f"agent_finish for sub-agents) as soon as possible." - ) - self.state.add_message("user", warning_msg) - - if self.state.iteration == self.state.max_iterations - 3: - final_warning_msg = ( - "CRITICAL: You have only 3 iterations left! " - "Your next message MUST be the tool call to the appropriate " - "finish tool: finish_scan if you are the root agent, or " - "agent_finish if you are a sub-agent. " - "No other actions should be taken except finishing your work " - "immediately." - ) - self.state.add_message("user", final_warning_msg) - - try: - iteration_task = asyncio.create_task(self._process_iteration(tracer)) - self._current_task = iteration_task - should_finish = await iteration_task - self._current_task = None - - if should_finish is None and self.interactive: - await self._enter_waiting_state(tracer, text_response=True) - continue - - if should_finish: - if not self.interactive: - self.state.set_completed({"success": True}) - if tracer: - tracer.update_agent_status(self.state.agent_id, "completed") - return self.state.final_result or {} - await self._enter_waiting_state(tracer, task_completed=True) - continue - - except asyncio.CancelledError: - self._current_task = None - if tracer: - partial_content = tracer.finalize_streaming_as_interrupted(self.state.agent_id) - if partial_content and partial_content.strip(): - self.state.add_message( - "assistant", f"{partial_content}\n\n[ABORTED BY USER]" - ) - if not self.interactive: - raise - await self._enter_waiting_state(tracer, error_occurred=False, was_cancelled=True) - continue - - except LLMRequestFailedError as e: - result = self._handle_llm_error(e, tracer) - if result is not None: - return result - continue - - except (RuntimeError, ValueError, TypeError) as e: - if not await self._handle_iteration_error(e, tracer): - if not self.interactive: - self.state.set_completed({"success": False, "error": str(e)}) - if tracer: - tracer.update_agent_status(self.state.agent_id, "failed") - raise - await self._enter_waiting_state(tracer, error_occurred=True) - continue - - async def _wait_for_input(self) -> None: - if self._force_stop: - return - - if self.state.has_waiting_timeout(): - self.state.resume_from_waiting() - self.state.add_message("user", "Waiting timeout reached. Resuming execution.") - - from strix.telemetry.tracer import get_global_tracer - - tracer = get_global_tracer() - if tracer: - tracer.update_agent_status(self.state.agent_id, "running") - - try: - from strix.tools.agents_graph.agents_graph_actions import _agent_graph - - if self.state.agent_id in _agent_graph["nodes"]: - _agent_graph["nodes"][self.state.agent_id]["status"] = "running" - except (ImportError, KeyError): - pass - - return - - await asyncio.sleep(0.5) - - async def _enter_waiting_state( - self, - tracer: Optional["Tracer"], - task_completed: bool = False, - error_occurred: bool = False, - was_cancelled: bool = False, - text_response: bool = False, - ) -> None: - self.state.enter_waiting_state() - - if tracer: - if text_response: - tracer.update_agent_status(self.state.agent_id, "waiting_for_input") - elif task_completed: - tracer.update_agent_status(self.state.agent_id, "completed") - elif error_occurred: - tracer.update_agent_status(self.state.agent_id, "error") - elif was_cancelled: - tracer.update_agent_status(self.state.agent_id, "stopped") - else: - tracer.update_agent_status(self.state.agent_id, "stopped") - - if text_response: - return - - if task_completed: - self.state.add_message( - "assistant", - "Task completed. I'm now waiting for follow-up instructions or new tasks.", - ) - elif error_occurred: - self.state.add_message( - "assistant", "An error occurred. I'm now waiting for new instructions." - ) - elif was_cancelled: - self.state.add_message( - "assistant", "Execution was cancelled. I'm now waiting for new instructions." - ) - else: - self.state.add_message( - "assistant", - "Execution paused. I'm now waiting for new instructions or any updates.", - ) - - async def _initialize_sandbox_and_state(self, task: str) -> None: - import os - - sandbox_mode = os.getenv("STRIX_SANDBOX_MODE", "false").lower() == "true" - if not sandbox_mode and self.state.sandbox_id is None: - from strix.runtime import get_runtime - - try: - runtime = get_runtime() - sandbox_info = await runtime.create_sandbox( - self.state.agent_id, self.state.sandbox_token, self.local_sources - ) - self.state.sandbox_id = sandbox_info["workspace_id"] - self.state.sandbox_token = sandbox_info["auth_token"] - self.state.sandbox_info = sandbox_info - - if "agent_id" in sandbox_info: - self.state.sandbox_info["agent_id"] = sandbox_info["agent_id"] - - caido_port = sandbox_info.get("caido_port") - if caido_port: - from strix.telemetry.tracer import get_global_tracer - - tracer = get_global_tracer() - if tracer: - tracer.caido_url = f"localhost:{caido_port}" - except Exception as e: - from strix.telemetry import posthog - - posthog.error("sandbox_init_error", str(e)) - raise - - if not self.state.task: - self.state.task = task - - self.state.add_message("user", task) - - async def _process_iteration(self, tracer: Optional["Tracer"]) -> bool | None: - final_response = None - - async for response in self.llm.generate(self.state.get_conversation_history()): - final_response = response - if tracer and response.content: - tracer.update_streaming_content(self.state.agent_id, response.content) - - if final_response is None: - return False - - content_stripped = (final_response.content or "").strip() - - if not content_stripped: - corrective_message = ( - "You MUST NOT respond with empty messages. " - "If you currently have nothing to do or say, use an appropriate tool instead:\n" - "- Use agents_graph_actions.wait_for_message to wait for messages " - "from user or other agents\n" - "- Use agents_graph_actions.agent_finish if you are a sub-agent " - "and your task is complete\n" - "- Use finish_actions.finish_scan if you are the root/main agent " - "and the scan is complete" - ) - self.state.add_message("user", corrective_message) - return False - - thinking_blocks = getattr(final_response, "thinking_blocks", None) - self.state.add_message("assistant", final_response.content, thinking_blocks=thinking_blocks) - if tracer: - tracer.clear_streaming_content(self.state.agent_id) - tracer.log_chat_message( - content=clean_content(final_response.content), - role="assistant", - agent_id=self.state.agent_id, - ) - - actions = ( - final_response.tool_invocations - if hasattr(final_response, "tool_invocations") and final_response.tool_invocations - else [] - ) - - if actions: - return await self._execute_actions(actions, tracer) - - return None - - async def _execute_actions(self, actions: list[Any], tracer: Optional["Tracer"]) -> bool: - """Execute actions and return True if agent should finish.""" - for action in actions: - self.state.add_action(action) - - conversation_history = self.state.get_conversation_history() - - tool_task = asyncio.create_task( - process_tool_invocations(actions, conversation_history, self.state) - ) - self._current_task = tool_task - - try: - should_agent_finish = await tool_task - self._current_task = None - except asyncio.CancelledError: - self._current_task = None - self.state.add_error("Tool execution cancelled by user") - raise - - self.state.messages = conversation_history - - if should_agent_finish: - self.state.set_completed({"success": True}) - if tracer: - tracer.update_agent_status(self.state.agent_id, "completed") - if not self.interactive and self.state.parent_id is None: - return True - return True - - return False - - def _check_agent_messages(self, state: AgentState) -> None: # noqa: PLR0912 - try: - from strix.tools.agents_graph.agents_graph_actions import _agent_graph, _agent_messages - - agent_id = state.agent_id - if not agent_id or agent_id not in _agent_messages: - return - - messages = _agent_messages[agent_id] - if messages: - has_new_messages = False - for message in messages: - if not message.get("read", False): - sender_id = message.get("from") - - if state.is_waiting_for_input(): - if state.llm_failed: - if sender_id == "user": - state.resume_from_waiting() - has_new_messages = True - - from strix.telemetry.tracer import get_global_tracer - - tracer = get_global_tracer() - if tracer: - tracer.update_agent_status(state.agent_id, "running") - else: - state.resume_from_waiting() - has_new_messages = True - - from strix.telemetry.tracer import get_global_tracer - - tracer = get_global_tracer() - if tracer: - tracer.update_agent_status(state.agent_id, "running") - - if sender_id == "user": - sender_name = "User" - state.add_message("user", message.get("content", "")) - else: - if sender_id and sender_id in _agent_graph.get("nodes", {}): - sender_name = _agent_graph["nodes"][sender_id]["name"] - - message_content = f""" - - You have received a message from another agent. You should acknowledge - this message and respond appropriately based on its content. However, DO NOT echo - back or repeat the entire message structure in your response. Simply process the - content and respond naturally as/if needed. - - - {sender_name} - {sender_id} - - - {message.get("message_type", "information")} - {message.get("priority", "normal")} - {message.get("timestamp", "")} - - -{message.get("content", "")} - - - This message was delivered during your task execution. - Please acknowledge and respond if needed. - -""" - state.add_message("user", message_content.strip()) - - message["read"] = True - - if has_new_messages and not state.is_waiting_for_input(): - from strix.telemetry.tracer import get_global_tracer - - tracer = get_global_tracer() - if tracer: - tracer.update_agent_status(agent_id, "running") - - except (AttributeError, KeyError, TypeError) as e: - import logging - - logger = logging.getLogger(__name__) - logger.warning(f"Error checking agent messages: {e}") - return - - def _handle_sandbox_error( - self, - error: SandboxInitializationError, - tracer: Optional["Tracer"], - ) -> dict[str, Any]: - error_msg = str(error.message) - error_details = error.details - self.state.add_error(error_msg) - - if not self.interactive: - self.state.set_completed({"success": False, "error": error_msg}) - if tracer: - tracer.update_agent_status(self.state.agent_id, "failed", error_msg) - if error_details: - exec_id = tracer.log_tool_execution_start( - self.state.agent_id, - "sandbox_error_details", - {"error": error_msg, "details": error_details}, - ) - tracer.update_tool_execution(exec_id, "failed", {"details": error_details}) - return {"success": False, "error": error_msg, "details": error_details} - - self.state.enter_waiting_state() - if tracer: - tracer.update_agent_status(self.state.agent_id, "sandbox_failed", error_msg) - if error_details: - exec_id = tracer.log_tool_execution_start( - self.state.agent_id, - "sandbox_error_details", - {"error": error_msg, "details": error_details}, - ) - tracer.update_tool_execution(exec_id, "failed", {"details": error_details}) - - return {"success": False, "error": error_msg, "details": error_details} - - def _handle_llm_error( - self, - error: LLMRequestFailedError, - tracer: Optional["Tracer"], - ) -> dict[str, Any] | None: - error_msg = str(error) - error_details = getattr(error, "details", None) - self.state.add_error(error_msg) - - if not self.interactive: - self.state.set_completed({"success": False, "error": error_msg}) - if tracer: - tracer.update_agent_status(self.state.agent_id, "failed", error_msg) - if error_details: - exec_id = tracer.log_tool_execution_start( - self.state.agent_id, - "llm_error_details", - {"error": error_msg, "details": error_details}, - ) - tracer.update_tool_execution(exec_id, "failed", {"details": error_details}) - return {"success": False, "error": error_msg} - - self.state.enter_waiting_state(llm_failed=True) - if tracer: - tracer.update_agent_status(self.state.agent_id, "llm_failed", error_msg) - if error_details: - exec_id = tracer.log_tool_execution_start( - self.state.agent_id, - "llm_error_details", - {"error": error_msg, "details": error_details}, - ) - tracer.update_tool_execution(exec_id, "failed", {"details": error_details}) - - return None - - async def _handle_iteration_error( - self, - error: RuntimeError | ValueError | TypeError | asyncio.CancelledError, - tracer: Optional["Tracer"], - ) -> bool: - error_msg = f"Error in iteration {self.state.iteration}: {error!s}" - logger.exception(error_msg) - self.state.add_error(error_msg) - if tracer: - tracer.update_agent_status(self.state.agent_id, "error") - return True - - def cancel_current_execution(self) -> None: - self._force_stop = True - if self._current_task and not self._current_task.done(): - try: - loop = self._current_task.get_loop() - loop.call_soon_threadsafe(self._current_task.cancel) - except RuntimeError: - self._current_task.cancel() - self._current_task = None diff --git a/strix/agents/sdk_factory.py b/strix/agents/factory.py similarity index 88% rename from strix/agents/sdk_factory.py rename to strix/agents/factory.py index 1f62213..9ba90ce 100644 --- a/strix/agents/sdk_factory.py +++ b/strix/agents/factory.py @@ -2,7 +2,7 @@ This is the keystone that links Phase 2's SDK function tools, Phase 3's graph tools, Phase 4's CaidoCapability, and the rendered Jinja prompt -from :mod:`strix.agents.sdk_prompt` into a single ``agents.Agent`` +from :mod:`strix.agents.prompt` into a single ``agents.Agent`` instance ready for ``Runner.run``. Two flavors: @@ -38,8 +38,8 @@ from agents import Agent from agents.agent import StopAtTools from agents.tool import Tool -from strix.agents.sdk_prompt import render_system_prompt -from strix.tools.agents_graph.agents_graph_sdk_tools import ( +from strix.agents.prompt import render_system_prompt +from strix.tools.agents_graph.tools import ( agent_finish, agent_status, create_agent, @@ -47,26 +47,26 @@ from strix.tools.agents_graph.agents_graph_sdk_tools import ( view_agent_graph, wait_for_message, ) -from strix.tools.browser.browser_sdk_tool import browser_action -from strix.tools.file_edit.file_edit_sdk_tools import ( +from strix.tools.browser.tool import browser_action +from strix.tools.file_edit.tools import ( list_files, search_files, str_replace_editor, ) -from strix.tools.finish.finish_sdk_tool import finish_scan -from strix.tools.load_skill.load_skill_sdk_tool import load_skill -from strix.tools.notes.notes_sdk_tools import ( +from strix.tools.finish.tool import finish_scan +from strix.tools.load_skill.tool import load_skill +from strix.tools.notes.tools import ( create_note, delete_note, get_note, list_notes, update_note, ) -from strix.tools.python.python_sdk_tool import python_action -from strix.tools.reporting.reporting_sdk_tools import create_vulnerability_report -from strix.tools.terminal.terminal_sdk_tool import terminal_execute -from strix.tools.thinking.thinking_sdk_tools import think -from strix.tools.todo.todo_sdk_tools import ( +from strix.tools.python.tool import python_action +from strix.tools.reporting.tool import create_vulnerability_report +from strix.tools.terminal.tool import terminal_execute +from strix.tools.thinking.tool import think +from strix.tools.todo.tools import ( create_todo, delete_todo, list_todos, @@ -74,7 +74,7 @@ from strix.tools.todo.todo_sdk_tools import ( mark_todo_pending, update_todo, ) -from strix.tools.web_search.web_search_sdk_tool import web_search +from strix.tools.web_search.tool import web_search logger = logging.getLogger(__name__) diff --git a/strix/agents/sdk_prompt.py b/strix/agents/prompt.py similarity index 68% rename from strix/agents/sdk_prompt.py rename to strix/agents/prompt.py index f77e964..795deaa 100644 --- a/strix/agents/sdk_prompt.py +++ b/strix/agents/prompt.py @@ -1,19 +1,12 @@ -"""Standalone Jinja-based system-prompt renderer for SDK agents. +"""Jinja-based system-prompt renderer. -The legacy ``LLM._load_system_prompt`` couples prompt rendering to the -LLM client class. The SDK migration owns the model client through -``MultiProvider`` instead, so we extract the rendering logic into a -plain function that the SDK agent factory can call without pulling in -the legacy ``LLM`` instance. - -Reuses the existing Jinja template at -``strix/agents/StrixAgent/system_prompt.jinja`` (508 lines, expanding -into the multi-section prompt with skills, tools, scan modes, etc.) so -behavior parity is preserved verbatim — only the call site changes. +Loads ``strix/agents/prompts/system_prompt.jinja`` (508 lines — the +multi-section production prompt with skills, tools, scan modes, etc.) +and renders it with the caller's per-run context (scan mode, whitebox, +interactive, scope authorization block). References: - HARNESS_WIKI.md §4.1 (system prompt assembly) - - PLAYBOOK.md §4 (per-tool migration contracts) """ from __future__ import annotations @@ -31,10 +24,7 @@ from strix.utils.resource_paths import get_strix_resource_path logger = logging.getLogger(__name__) -# Hard-coded to the StrixAgent template since it's the only agent type -# under the SDK migration. The legacy harness supported multiple agent -# names but in practice only StrixAgent ships. -_AGENT_NAME = "StrixAgent" +_PROMPT_DIRNAME = "prompts" def _resolve_skills( @@ -45,8 +35,7 @@ def _resolve_skills( ) -> list[str]: """Build the deduped, ordered skills list for the prompt render. - Mirrors :py:meth:`LLM._get_skills_to_load` exactly so the rendered - prompt is byte-identical to the legacy path: + Order: 1. Whatever the caller asked for, in order. 2. ``scan_modes/`` (always). @@ -75,7 +64,7 @@ def render_system_prompt( interactive: bool = False, system_prompt_context: dict[str, Any] | None = None, ) -> str: - """Render the StrixAgent system prompt. + """Render the system prompt. Args: skills: Skills the caller wants preloaded into the prompt @@ -88,17 +77,16 @@ def render_system_prompt( interactive: When True, the prompt renders the interactive-mode communication rules block. system_prompt_context: Free-form dict that the template's - ``system_prompt_context`` variable receives — used today for - the scan-scope authorization block from - :py:meth:`StrixAgent._build_system_scope_context`. + ``system_prompt_context`` variable receives — carries the + scan-scope authorization block. Returns the rendered prompt string. If anything goes wrong (template - missing, render failure), returns an empty string and logs — same - fail-soft posture as the legacy method, because a missing prompt is - survivable but a hard failure during agent construction is not. + missing, render failure), returns an empty string and logs — a + missing prompt is survivable, a hard failure during agent + construction is not. """ try: - prompt_dir = get_strix_resource_path("agents", _AGENT_NAME) + prompt_dir = get_strix_resource_path("agents", _PROMPT_DIRNAME) skills_dir = get_strix_resource_path("skills") env = Environment( loader=FileSystemLoader([prompt_dir, skills_dir]), diff --git a/strix/agents/StrixAgent/system_prompt.jinja b/strix/agents/prompts/system_prompt.jinja similarity index 100% rename from strix/agents/StrixAgent/system_prompt.jinja rename to strix/agents/prompts/system_prompt.jinja diff --git a/strix/agents/state.py b/strix/agents/state.py deleted file mode 100644 index da04ee7..0000000 --- a/strix/agents/state.py +++ /dev/null @@ -1,172 +0,0 @@ -import uuid -from datetime import UTC, datetime -from typing import Any - -from pydantic import BaseModel, Field - - -def _generate_agent_id() -> str: - return f"agent_{uuid.uuid4().hex[:8]}" - - -class AgentState(BaseModel): - agent_id: str = Field(default_factory=_generate_agent_id) - agent_name: str = "Strix Agent" - parent_id: str | None = None - sandbox_id: str | None = None - sandbox_token: str | None = None - sandbox_info: dict[str, Any] | None = None - - task: str = "" - iteration: int = 0 - max_iterations: int = 300 - completed: bool = False - stop_requested: bool = False - waiting_for_input: bool = False - llm_failed: bool = False - waiting_start_time: datetime | None = None - waiting_timeout: int = 600 - final_result: dict[str, Any] | None = None - max_iterations_warning_sent: bool = False - - messages: list[dict[str, Any]] = Field(default_factory=list) - context: dict[str, Any] = Field(default_factory=dict) - - start_time: str = Field(default_factory=lambda: datetime.now(UTC).isoformat()) - last_updated: str = Field(default_factory=lambda: datetime.now(UTC).isoformat()) - - actions_taken: list[dict[str, Any]] = Field(default_factory=list) - observations: list[dict[str, Any]] = Field(default_factory=list) - - errors: list[str] = Field(default_factory=list) - - def increment_iteration(self) -> None: - self.iteration += 1 - self.last_updated = datetime.now(UTC).isoformat() - - def add_message( - self, role: str, content: Any, thinking_blocks: list[dict[str, Any]] | None = None - ) -> None: - message = {"role": role, "content": content} - if thinking_blocks: - message["thinking_blocks"] = thinking_blocks - self.messages.append(message) - self.last_updated = datetime.now(UTC).isoformat() - - def add_action(self, action: dict[str, Any]) -> None: - self.actions_taken.append( - { - "iteration": self.iteration, - "timestamp": datetime.now(UTC).isoformat(), - "action": action, - } - ) - - def add_observation(self, observation: dict[str, Any]) -> None: - self.observations.append( - { - "iteration": self.iteration, - "timestamp": datetime.now(UTC).isoformat(), - "observation": observation, - } - ) - - def add_error(self, error: str) -> None: - self.errors.append(f"Iteration {self.iteration}: {error}") - self.last_updated = datetime.now(UTC).isoformat() - - def update_context(self, key: str, value: Any) -> None: - self.context[key] = value - self.last_updated = datetime.now(UTC).isoformat() - - def set_completed(self, final_result: dict[str, Any] | None = None) -> None: - self.completed = True - self.final_result = final_result - self.last_updated = datetime.now(UTC).isoformat() - - def request_stop(self) -> None: - self.stop_requested = True - self.last_updated = datetime.now(UTC).isoformat() - - def should_stop(self) -> bool: - return self.stop_requested or self.completed or self.has_reached_max_iterations() - - def is_waiting_for_input(self) -> bool: - return self.waiting_for_input - - def enter_waiting_state(self, llm_failed: bool = False) -> None: - self.waiting_for_input = True - self.waiting_start_time = datetime.now(UTC) - self.llm_failed = llm_failed - self.last_updated = datetime.now(UTC).isoformat() - - def resume_from_waiting(self, new_task: str | None = None) -> None: - self.waiting_for_input = False - self.waiting_start_time = None - self.stop_requested = False - self.completed = False - self.llm_failed = False - if new_task: - self.task = new_task - self.last_updated = datetime.now(UTC).isoformat() - - def has_reached_max_iterations(self) -> bool: - return self.iteration >= self.max_iterations - - def is_approaching_max_iterations(self, threshold: float = 0.85) -> bool: - return self.iteration >= int(self.max_iterations * threshold) - - def has_waiting_timeout(self) -> bool: - if self.waiting_timeout == 0: - return False - - if not self.waiting_for_input or not self.waiting_start_time: - return False - - if ( - self.stop_requested - or self.llm_failed - or self.completed - or self.has_reached_max_iterations() - ): - return False - - elapsed = (datetime.now(UTC) - self.waiting_start_time).total_seconds() - return elapsed > self.waiting_timeout - - def has_empty_last_messages(self, count: int = 3) -> bool: - if len(self.messages) < count: - return False - - last_messages = self.messages[-count:] - - for message in last_messages: - content = message.get("content", "") - if isinstance(content, str) and content.strip(): - return False - - return True - - def get_conversation_history(self) -> list[dict[str, Any]]: - return self.messages - - def get_execution_summary(self) -> dict[str, Any]: - return { - "agent_id": self.agent_id, - "agent_name": self.agent_name, - "parent_id": self.parent_id, - "sandbox_id": self.sandbox_id, - "sandbox_info": self.sandbox_info, - "task": self.task, - "iteration": self.iteration, - "max_iterations": self.max_iterations, - "completed": self.completed, - "final_result": self.final_result, - "start_time": self.start_time, - "last_updated": self.last_updated, - "total_actions": len(self.actions_taken), - "total_observations": len(self.observations), - "total_errors": len(self.errors), - "has_errors": len(self.errors) > 0, - "max_iterations_reached": self.has_reached_max_iterations() and not self.completed, - } diff --git a/strix/sdk_entry.py b/strix/entry.py similarity index 85% rename from strix/sdk_entry.py rename to strix/entry.py index 0967386..7cc143b 100644 --- a/strix/sdk_entry.py +++ b/strix/entry.py @@ -1,7 +1,4 @@ -"""Top-level SDK scan entry point. - -Replaces the legacy ``strix.cli.main → StrixAgent.execute_scan`` -pipeline with the SDK-native equivalent: +"""Top-level scan entry point. 1. Build the per-scan ``AgentMessageBus``. 2. Bring up (or reuse) a sandbox session for ``scan_id`` via the @@ -12,16 +9,8 @@ pipeline with the SDK-native equivalent: 5. Register the root in the bus. 6. Build the ``RunConfig`` via the factory. 7. Call ``Runner.run(...)`` and surface the result. -8. ``finally`` cleanup the sandbox session. - -Phase 5 lands the wiring; the streaming accumulator + TUI integration -land in Phase 5b. The entry point is intentionally not wired to the -CLI yet — that's a follow-up under ``STRIX_USE_SDK_HARNESS=1`` (see -PLAYBOOK §7.1 cutover plan). - -References: - - PLAYBOOK.md §3.3 (session manager), §4.3 (graph tools), §7.1 - - AUDIT_R3.md C9 (cancel_descendants on cleanup) +8. ``finally`` cleanup the sandbox session — even on cancel, the bus + propagates ``cancel_descendants`` to every spawned child task. """ from __future__ import annotations @@ -33,7 +22,7 @@ from typing import TYPE_CHECKING, Any from agents import Runner -from strix.agents.sdk_factory import build_strix_agent, make_child_factory +from strix.agents.factory import build_strix_agent, make_child_factory from strix.orchestration.bus import AgentMessageBus from strix.orchestration.hooks import StrixOrchestrationHooks from strix.run_config_factory import ( @@ -54,11 +43,10 @@ logger = logging.getLogger(__name__) def _build_root_task(scan_config: dict[str, Any]) -> str: """Format the user-facing task for the root agent. - Mirrors :py:meth:`StrixAgent.execute_scan` (legacy) — collects each - target type into a labelled section, appends diff-scope context if - active, and tacks on user_instructions. The structured shape is - important for prompt parity: the system prompt template references - these section headers. + Collects each target type into a labelled section, appends + diff-scope context if active, and tacks on user_instructions. The + structured section headers are referenced by the system prompt + template, so the shape matters for prompt parity. """ targets = scan_config.get("targets", []) or [] diff_scope = scan_config.get("diff_scope") or {} @@ -128,10 +116,8 @@ def _build_root_task(scan_config: dict[str, Any]) -> str: def _build_scope_context(scan_config: dict[str, Any]) -> dict[str, Any]: """Produce the system_prompt_context block used by the prompt template. - Same shape as the legacy - :py:meth:`StrixAgent._build_system_scope_context` so the prompt - template's ``system_prompt_context.authorized_targets`` lookups - stay byte-identical. + The prompt template's ``system_prompt_context.authorized_targets`` + lookups expect this exact shape. """ authorized: list[dict[str, str]] = [] for target in scan_config.get("targets", []) or []: @@ -177,9 +163,9 @@ async def run_strix_scan( """Run one Strix scan end-to-end against a freshly-prepared sandbox. Args: - scan_config: Same shape the legacy ``StrixAgent.execute_scan`` - takes (targets, user_instructions, diff_scope, scan_mode, - is_whitebox, skills). + scan_config: Per-scan configuration — ``targets``, + ``user_instructions``, ``diff_scope``, ``scan_mode``, + ``is_whitebox``, ``skills``. scan_id: Used to key the sandbox session cache. Auto-generated if omitted — callers that want resume-after-crash semantics should pass a stable id. @@ -190,8 +176,7 @@ async def run_strix_scan( telemetry hook chain. Pass ``None`` for unit tests. interactive: Renders the interactive-mode prompt block on the root agent. - max_turns: Cap on root-agent LLM turns. Mirrors legacy - ``AgentState.max_iterations`` (300). + max_turns: Cap on root-agent LLM turns (default 300). cleanup_on_exit: When True (default), tears down the sandbox session in a ``finally``. Set to False for resume scenarios where the caller wants to preserve the container. diff --git a/strix/interface/cli.py b/strix/interface/cli.py index 6004abe..0da8d23 100644 --- a/strix/interface/cli.py +++ b/strix/interface/cli.py @@ -1,8 +1,11 @@ import atexit +import contextlib +import os import signal import sys import threading import time +from pathlib import Path from typing import Any from rich.console import Console @@ -10,10 +13,9 @@ from rich.live import Live from rich.panel import Panel from rich.text import Text -from strix.agents.StrixAgent import StrixAgent -from strix.interface.sdk_dispatch import run_scan_via_sdk, should_use_sdk_harness -from strix.llm.config import LLMConfig -from strix.runtime import cleanup_runtime +from strix.config import Config +from strix.entry import run_strix_scan +from strix.sandbox import session_manager from strix.telemetry.tracer import Tracer, set_global_tracer from .utils import ( @@ -22,6 +24,35 @@ from .utils import ( ) +def _resolve_sandbox_image() -> str: + image = Config.get("strix_image") + if not image: + raise RuntimeError( + "strix_image is not configured. Set it in ~/.strix/cli-config.json.", + ) + return str(image) + + +def _resolve_sources_path(args: Any) -> Path: + """Pick the host directory to mount into ``/workspace/sources``. + + - With ``--local-sources``, mount the parent of the first source so + the agent can walk down into the actual tree. + - Otherwise, a per-run scratch dir under ``$XDG_CACHE_HOME/strix``. + """ + local_sources: list[dict[str, str]] | None = getattr(args, "local_sources", None) + if local_sources: + first = local_sources[0] + host_path = first.get("host_path") or first.get("source_path") or first.get("path") + if host_path: + return Path(host_path).expanduser().resolve().parent + + cache_root = os.environ.get("XDG_CACHE_HOME") or str(Path.home() / ".cache") + sources = Path(cache_root) / "strix" / "sources" / str(args.run_name) + sources.mkdir(parents=True, exist_ok=True) + return sources + + async def run_cli(args: Any) -> None: # noqa: PLR0915 console = Console() @@ -68,27 +99,18 @@ async def run_cli(args: Any) -> None: # noqa: PLR0915 console.print() scan_mode = getattr(args, "scan_mode", "deep") + is_whitebox = bool(getattr(args, "local_sources", [])) - scan_config = { + scan_config: dict[str, Any] = { "scan_id": args.run_name, "targets": args.targets_info, "user_instructions": args.instruction or "", "run_name": args.run_name, "diff_scope": getattr(args, "diff_scope", {"active": False}), + "scan_mode": scan_mode, + "is_whitebox": is_whitebox, } - llm_config = LLMConfig( - scan_mode=scan_mode, - is_whitebox=bool(getattr(args, "local_sources", [])), - ) - agent_config = { - "llm_config": llm_config, - "max_iterations": 300, - } - - if getattr(args, "local_sources", None): - agent_config["local_sources"] = args.local_sources - tracer = Tracer(args.run_name) tracer.set_scan_config(scan_config) @@ -112,7 +134,6 @@ async def run_cli(args: Any) -> None: # noqa: PLR0915 def cleanup_on_exit() -> None: tracer.cleanup() - cleanup_runtime() def signal_handler(_signum: int, _frame: Any) -> None: tracer.cleanup() @@ -131,7 +152,7 @@ async def run_cli(args: Any) -> None: # noqa: PLR0915 status_text.append("Penetration test in progress", style="bold #22c55e") status_text.append("\n\n") - stats_text = build_live_stats_text(tracer, agent_config) + stats_text = build_live_stats_text(tracer) if stats_text: status_text.append(stats_text) @@ -156,39 +177,30 @@ async def run_cli(args: Any) -> None: # noqa: PLR0915 try: live.update(create_live_status()) time.sleep(2) - except Exception: # noqa: BLE001 + except Exception: break update_thread = threading.Thread(target=update_status, daemon=True) update_thread.start() try: - if should_use_sdk_harness(): - # SDK harness opt-in (PLAYBOOK §7.1). Returns a - # RunResult, not the legacy success-dict shape, so - # we skip the legacy error-extraction block — - # failures inside run_strix_scan raise instead. - await run_scan_via_sdk( - scan_config=scan_config, - args=args, - tracer=tracer, - ) - else: - agent = StrixAgent(agent_config) - result = await agent.execute_scan(scan_config) - - if isinstance(result, dict) and not result.get("success", True): - error_msg = result.get("error", "Unknown error") - error_details = result.get("details") - console.print() - console.print(f"[bold red]Penetration test failed:[/] {error_msg}") - if error_details: - console.print(f"[dim]{error_details}[/]") - console.print() - sys.exit(1) + await run_strix_scan( + scan_config=scan_config, + scan_id=args.run_name, + image=_resolve_sandbox_image(), + sources_path=_resolve_sources_path(args), + tracer=tracer, + interactive=bool(getattr(args, "interactive", False)), + ) finally: stop_updates.set() update_thread.join(timeout=1) + # Best-effort: tear down the sandbox session even if the + # run raised. ``run_strix_scan`` already does this in its + # own ``finally``, but call here too in case the failure + # was during early setup. + with contextlib.suppress(Exception): + await session_manager.cleanup(args.run_name) except Exception as e: console.print(f"[bold red]Error during penetration test:[/] {e}") diff --git a/strix/interface/main.py b/strix/interface/main.py index a2323aa..37535e3 100644 --- a/strix/interface/main.py +++ b/strix/interface/main.py @@ -20,7 +20,7 @@ from rich.text import Text from strix.config import Config, apply_saved_config, save_current_config from strix.config.config import resolve_llm_config -from strix.llm.utils import resolve_strix_model +from strix.llm.multi_provider_setup import STRIX_MODEL_MAP apply_saved_config() @@ -42,7 +42,9 @@ from strix.interface.utils import ( # noqa: E402 validate_config_file, validate_llm_response, ) -from strix.runtime.docker_runtime import HOST_GATEWAY_HOSTNAME # noqa: E402 + + +HOST_GATEWAY_HOSTNAME = "host.docker.internal" from strix.telemetry import posthog # noqa: E402 from strix.telemetry.tracer import get_global_tracer # noqa: E402 @@ -50,7 +52,7 @@ from strix.telemetry.tracer import get_global_tracer # noqa: E402 logging.getLogger().setLevel(logging.ERROR) -def validate_environment() -> None: # noqa: PLR0912, PLR0915 +def validate_environment() -> None: console = Console() missing_required_vars = [] missing_optional_vars = [] @@ -209,8 +211,13 @@ async def warm_up_llm() -> None: try: model_name, api_key, api_base = resolve_llm_config() - litellm_model, _ = resolve_strix_model(model_name) - litellm_model = litellm_model or model_name + # ``strix/`` is routed through the Strix proxy (OpenAI-compatible); + # everything else is sent as-is. + litellm_model: str | None = model_name + if model_name and model_name.startswith("strix/"): + base = model_name[len("strix/") :] + if base in STRIX_MODEL_MAP: + litellm_model = f"openai/{base}" test_messages = [ {"role": "system", "content": "You are a helpful assistant."}, @@ -233,7 +240,7 @@ async def warm_up_llm() -> None: validate_llm_response(response) - except Exception as e: # noqa: BLE001 + except Exception as e: error_text = Text() error_text.append("LLM CONNECTION FAILED", style="bold red") error_text.append("\n\n", style="white") @@ -260,7 +267,7 @@ def get_version() -> str: from importlib.metadata import version return version("strix-agent") - except Exception: # noqa: BLE001 + except Exception: return "unknown" @@ -401,7 +408,7 @@ Examples: args.instruction = f.read().strip() if not args.instruction: parser.error(f"Instruction file '{instruction_path}' is empty") - except Exception as e: # noqa: BLE001 + except Exception as e: parser.error(f"Failed to read instruction file '{instruction_path}': {e}") args.targets_info = [] @@ -544,7 +551,7 @@ def persist_config() -> None: save_current_config() -def main() -> None: # noqa: PLR0912, PLR0915 +def main() -> None: if sys.platform == "win32": asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy()) diff --git a/strix/interface/sdk_dispatch.py b/strix/interface/sdk_dispatch.py deleted file mode 100644 index 3cbaf34..0000000 --- a/strix/interface/sdk_dispatch.py +++ /dev/null @@ -1,143 +0,0 @@ -"""STRIX_USE_SDK_HARNESS dispatch — selects legacy vs SDK harness at run-time. - -Phase 5b cutover gate. The legacy CLI (``strix.interface.cli``) calls -``StrixAgent(...).execute_scan(scan_config)`` directly. To roll out the -SDK migration safely we want a single env-var-gated branch: - - STRIX_USE_SDK_HARNESS=1 → await run_strix_scan(...) - STRIX_USE_SDK_HARNESS=0 → await StrixAgent(...).execute_scan(...) - -This module is a thin adapter: it reads the env var, and when the SDK -path is active, translates the legacy ``scan_config`` + ``args`` pair -into the keyword arguments :func:`run_strix_scan` expects. - -Per PLAYBOOK §7.1: the legacy default stays in place until end-to-end -validation against a stable target succeeds; the env flag is the -opt-in. Removal of the legacy branch happens one release after cutover. - -References: - - PLAYBOOK.md §7.1 (cutover strategy) - - PLAYBOOK.md §7.2 (rollback procedure) -""" - -from __future__ import annotations - -import logging -import os -from pathlib import Path -from typing import TYPE_CHECKING, Any - - -if TYPE_CHECKING: - from agents.result import RunResult - - -logger = logging.getLogger(__name__) - - -_ENV_FLAG = "STRIX_USE_SDK_HARNESS" - - -def should_use_sdk_harness() -> bool: - """Return True iff ``STRIX_USE_SDK_HARNESS`` is truthy in the env. - - Truthy values: ``"1"``, ``"true"``, ``"yes"`` (case-insensitive). - Anything else — including unset — returns False so the default - deployed posture stays the legacy harness. - """ - raw = os.environ.get(_ENV_FLAG, "") - return raw.strip().lower() in {"1", "true", "yes"} - - -def _resolve_sandbox_image() -> str: - """Read the sandbox image tag from Strix config. - - Falls back to ``"strix-sandbox:latest"`` if unset — same behavior - the legacy ``DockerRuntime`` would surface as a config error. - """ - from strix.config import Config - - image = Config.get("strix_image") - if not image: - logger.warning( - "strix_image not configured; falling back to strix-sandbox:latest. " - "Set this in ~/.strix/cli-config.json for production use.", - ) - return "strix-sandbox:latest" - return str(image) - - -def _resolve_sources_path(args: Any) -> Path: - """Pick the host directory to mount into ``/workspace/sources``. - - - When ``--local-sources`` was passed, use the parent of the first - source's ``host_path`` (the legacy harness then copies each - individual source under ``/workspace/``; we mount the - parent and let the agent walk down). - - Otherwise, use a per-run scratch directory under - ``$XDG_CACHE_HOME/strix`` (or ``~/.cache/strix``) — the legacy - flow eventually populates ``/workspace`` via post-create copies, - which the SDK session manager doesn't replicate yet (Phase 6 - will bring that in). - """ - local_sources: list[dict[str, str]] | None = getattr(args, "local_sources", None) - if local_sources: - first = local_sources[0] - host_path = first.get("host_path") or first.get("source_path") or first.get("path") - if host_path: - return Path(host_path).expanduser().resolve().parent - - cache_root = os.environ.get("XDG_CACHE_HOME") or str(Path.home() / ".cache") - run_name = getattr(args, "run_name", "default") or "default" - sources = Path(cache_root) / "strix" / "sources" / str(run_name) - sources.mkdir(parents=True, exist_ok=True) - return sources - - -async def run_scan_via_sdk( - *, - scan_config: dict[str, Any], - args: Any, - tracer: Any, -) -> RunResult: - """Translate legacy CLI args into ``run_strix_scan`` kwargs. - - Args: - scan_config: The same dict the legacy ``StrixAgent.execute_scan`` - accepts. Forwarded verbatim to ``run_strix_scan``; the - entry point reads ``targets``, ``user_instructions``, - ``diff_scope``, ``scan_mode``, ``is_whitebox``, ``skills`` - from it. - args: argparse Namespace from ``strix.interface.cli``. We read - ``run_name``, ``local_sources``, ``scan_mode`` from it. - tracer: Live ``Tracer`` instance — flows through context so - tools (``create_vulnerability_report``, ``finish_scan``) - persist into the same on-disk run directory the legacy - path uses. - - Returns the SDK ``RunResult``. Raises whatever ``run_strix_scan`` - raises (sandbox bring-up failure, LLM error, etc.). - """ - from strix.sdk_entry import run_strix_scan - - run_name = getattr(args, "run_name", None) or scan_config.get("run_name") - image = _resolve_sandbox_image() - sources_path = _resolve_sources_path(args) - interactive = bool(getattr(args, "interactive", False)) - - logger.info( - "STRIX_USE_SDK_HARNESS active; dispatching scan %s via run_strix_scan " - "(image=%s, sources=%s)", - run_name, - image, - sources_path, - ) - - return await run_strix_scan( - scan_config=scan_config, - scan_id=run_name, - image=image, - sources_path=sources_path, - tracer=tracer, - interactive=interactive, - ) diff --git a/strix/interface/tui.py b/strix/interface/tui.py index 0cfd754..3db4ba1 100644 --- a/strix/interface/tui.py +++ b/strix/interface/tui.py @@ -1,13 +1,16 @@ import argparse import asyncio import atexit +import contextlib import logging +import os import signal import sys import threading from collections.abc import Callable from importlib.metadata import PackageNotFoundError from importlib.metadata import version as pkg_version +from pathlib import Path from typing import TYPE_CHECKING, Any, ClassVar @@ -28,13 +31,14 @@ from textual.screen import ModalScreen from textual.widgets import Button, Label, Static, TextArea, Tree from textual.widgets.tree import TreeNode -from strix.agents.StrixAgent import StrixAgent +from strix.config import Config +from strix.entry import run_strix_scan from strix.interface.streaming_parser import parse_streaming_content from strix.interface.tool_components.agent_message_renderer import AgentMessageRenderer from strix.interface.tool_components.registry import get_tool_renderer from strix.interface.tool_components.user_message_renderer import UserMessageRenderer from strix.interface.utils import build_tui_stats_text -from strix.llm.config import LLMConfig +from strix.sandbox import session_manager from strix.telemetry.tracer import Tracer, set_global_tracer @@ -329,7 +333,7 @@ class VulnerabilityDetailScreen(ModalScreen): # type: ignore[misc] else: return text - def _render_vulnerability(self) -> Text: # noqa: PLR0912, PLR0915 + def _render_vulnerability(self) -> Text: vuln = self.vulnerability text = Text() @@ -455,7 +459,7 @@ class VulnerabilityDetailScreen(ModalScreen): # type: ignore[misc] return text - def _get_markdown_report(self) -> str: # noqa: PLR0912, PLR0915 + def _get_markdown_report(self) -> str: """Get Markdown version of vulnerability report for clipboard.""" vuln = self.vulnerability lines: list[str] = [] @@ -702,7 +706,6 @@ class StrixTUIApp(App): # type: ignore[misc] super().__init__() self.args = args self.scan_config = self._build_scan_config(args) - self.agent_config = self._build_agent_config(args) self.tracer = Tracer(self.scan_config["run_name"]) self.tracer.set_scan_config(self.scan_config) @@ -736,6 +739,18 @@ class StrixTUIApp(App): # type: ignore[misc] self._setup_cleanup_handlers() + def _resolve_sources_path(self) -> Path: + local_sources = getattr(self.args, "local_sources", None) or [] + if local_sources: + first = local_sources[0] + host_path = first.get("host_path") or first.get("source_path") or first.get("path") + if host_path: + return Path(host_path).expanduser().resolve().parent + cache_root = os.environ.get("XDG_CACHE_HOME") or str(Path.home() / ".cache") + sources = Path(cache_root) / "strix" / "sources" / str(self.args.run_name) + sources.mkdir(parents=True, exist_ok=True) + return sources + def _build_scan_config(self, args: argparse.Namespace) -> dict[str, Any]: return { "scan_id": args.run_name, @@ -743,32 +758,13 @@ class StrixTUIApp(App): # type: ignore[misc] "user_instructions": args.instruction or "", "run_name": args.run_name, "diff_scope": getattr(args, "diff_scope", {"active": False}), + "scan_mode": getattr(args, "scan_mode", "deep"), + "is_whitebox": bool(getattr(args, "local_sources", [])), } - def _build_agent_config(self, args: argparse.Namespace) -> dict[str, Any]: - scan_mode = getattr(args, "scan_mode", "deep") - llm_config = LLMConfig( - scan_mode=scan_mode, - interactive=True, - is_whitebox=bool(getattr(args, "local_sources", [])), - ) - - config = { - "llm_config": llm_config, - "max_iterations": 300, - } - - if getattr(args, "local_sources", None): - config["local_sources"] = args.local_sources - - return config - def _setup_cleanup_handlers(self) -> None: def cleanup_on_exit() -> None: - from strix.runtime import cleanup_runtime - self.tracer.cleanup() - cleanup_runtime() def signal_handler(_signum: int, _frame: Any) -> None: self.tracer.cleanup() @@ -1305,7 +1301,7 @@ class StrixTUIApp(App): # type: ignore[misc] stats_content = Text() - stats_text = build_tui_stats_text(self.tracer, self.agent_config) + stats_text = build_tui_stats_text(self.tracer) if stats_text: stats_content.append(stats_text) @@ -1502,10 +1498,19 @@ class StrixTUIApp(App): # type: ignore[misc] asyncio.set_event_loop(loop) try: - agent = StrixAgent(self.agent_config) - if not self._scan_stop_event.is_set(): - loop.run_until_complete(agent.execute_scan(self.scan_config)) + image = Config.get("strix_image") or "strix-sandbox:latest" + sources_path = self._resolve_sources_path() + loop.run_until_complete( + run_strix_scan( + scan_config=self.scan_config, + scan_id=self.scan_config["run_name"], + image=str(image), + sources_path=sources_path, + tracer=self.tracer, + interactive=True, + ), + ) except (KeyboardInterrupt, asyncio.CancelledError): logging.info("Scan interrupted by user") @@ -1516,6 +1521,12 @@ class StrixTUIApp(App): # type: ignore[misc] except Exception: logging.exception("Unexpected error during scan") finally: + # Best-effort sandbox teardown if early setup failed + # before run_strix_scan's own ``finally`` ran. + with contextlib.suppress(Exception): + loop.run_until_complete( + session_manager.cleanup(self.scan_config["run_name"]), + ) loop.close() self._scan_completed.set() @@ -1816,15 +1827,9 @@ class StrixTUIApp(App): # type: ignore[misc] metadata={"interrupted": True}, ) - try: - from strix.tools.agents_graph.agents_graph_actions import _agent_instances - - if self.selected_agent_id in _agent_instances: - agent_instance = _agent_instances[self.selected_agent_id] - if hasattr(agent_instance, "cancel_current_execution"): - agent_instance.cancel_current_execution() - except (ImportError, AttributeError, KeyError): - pass + # TODO: route user→agent messages through the AgentMessageBus + # once the TUI has a handle on it. The bus currently lives + # inside ``run_strix_scan`` scope only. if self.tracer: self.tracer.log_chat_message( @@ -1833,15 +1838,11 @@ class StrixTUIApp(App): # type: ignore[misc] agent_id=self.selected_agent_id, ) - try: - from strix.tools.agents_graph.agents_graph_actions import send_user_message_to_agent - - send_user_message_to_agent(self.selected_agent_id, message) - - except (ImportError, AttributeError) as e: - import logging - - logging.warning(f"Failed to send message to agent {self.selected_agent_id}: {e}") + logging.warning( + "User-message-to-agent dispatch is not wired post-migration; " + "message %r logged to tracer but not delivered.", + message, + ) self._displayed_events.clear() self._update_chat_view() @@ -1940,22 +1941,12 @@ class StrixTUIApp(App): # type: ignore[misc] return agent_name, False def action_confirm_stop_agent(self, agent_id: str) -> None: - try: - from strix.tools.agents_graph.agents_graph_actions import stop_agent - - result = stop_agent(agent_id) - - import logging - - if result.get("success"): - logging.info(f"Stop request sent to agent: {result.get('message', 'Unknown')}") - else: - logging.warning(f"Failed to stop agent: {result.get('error', 'Unknown error')}") - - except Exception: - import logging - - logging.exception(f"Failed to stop agent {agent_id}") + # TODO: route to ``bus.cancel_descendants(agent_id)`` once the TUI + # has a handle on the AgentMessageBus. + logging.warning( + "Stop-agent dispatch is not wired post-migration; agent %s left running.", + agent_id, + ) def action_custom_quit(self) -> None: if self._scan_thread and self._scan_thread.is_alive(): @@ -2071,7 +2062,7 @@ class StrixTUIApp(App): # type: ignore[misc] cleaned = self._clean_copied_text(selected) self.copy_to_clipboard(cleaned if cleaned.strip() else selected) copied = True - except Exception: # noqa: BLE001 + except Exception: logger.debug("Failed to copy screen selection", exc_info=True) if not copied: @@ -2082,7 +2073,7 @@ class StrixTUIApp(App): # type: ignore[misc] self.copy_to_clipboard(selected) chat_input.move_cursor(chat_input.cursor_location) copied = True - except Exception: # noqa: BLE001 + except Exception: logger.debug("Failed to copy chat input selection", exc_info=True) if copied: diff --git a/strix/interface/utils.py b/strix/interface/utils.py index aff33a0..e25a0ed 100644 --- a/strix/interface/utils.py +++ b/strix/interface/utils.py @@ -20,6 +20,8 @@ from rich.console import Console from rich.panel import Panel from rich.text import Text +from strix.config import Config + # Token formatting utilities def format_token_count(count: float) -> str: @@ -297,17 +299,15 @@ def build_final_stats_text(tracer: Any) -> Text: return stats_text -def build_live_stats_text(tracer: Any, agent_config: dict[str, Any] | None = None) -> Text: +def build_live_stats_text(tracer: Any) -> Text: stats_text = Text() if not tracer: return stats_text - if agent_config: - llm_config = agent_config["llm_config"] - model = getattr(llm_config, "model_name", "Unknown") - stats_text.append("Model ", style="dim") - stats_text.append(model, style="white") - stats_text.append("\n") + model = Config.get("strix_llm") or "unknown" + stats_text.append("Model ", style="dim") + stats_text.append(str(model), style="white") + stats_text.append("\n") vuln_count = len(tracer.vulnerability_reports) tool_count = tracer.get_real_tool_count() @@ -370,15 +370,13 @@ def build_live_stats_text(tracer: Any, agent_config: dict[str, Any] | None = Non return stats_text -def build_tui_stats_text(tracer: Any, agent_config: dict[str, Any] | None = None) -> Text: +def build_tui_stats_text(tracer: Any) -> Text: stats_text = Text() if not tracer: return stats_text - if agent_config: - llm_config = agent_config["llm_config"] - model = getattr(llm_config, "model_name", "Unknown") - stats_text.append(model, style="white") + model = Config.get("strix_llm") or "unknown" + stats_text.append(str(model), style="white") llm_stats = tracer.get_total_llm_stats() total_stats = llm_stats["total"] @@ -427,7 +425,7 @@ def _derive_target_label_for_run_name(targets_info: list[dict[str, Any]] | None) try: parsed = urlparse(url) return str(parsed.netloc or parsed.path or url) - except Exception: # noqa: BLE001 + except Exception: return str(url) if target_type == "repository": @@ -443,7 +441,7 @@ def _derive_target_label_for_run_name(targets_info: list[dict[str, Any]] | None) path_str = details.get("target_path", original) try: return str(Path(path_str).name or path_str) - except Exception: # noqa: BLE001 + except Exception: return str(path_str) if target_type == "ip_address": diff --git a/strix/llm/__init__.py b/strix/llm/__init__.py index dea8375..9c116e9 100644 --- a/strix/llm/__init__.py +++ b/strix/llm/__init__.py @@ -1,17 +1,19 @@ +"""LLM package — model provider, prompt-cache wrapper, session, dedup helper. + +Side effects on import: + +- Quiet litellm's debug logger (it spams ``logging.DEBUG`` on every + request). The SDK's MultiProvider routes through litellm under the + hood, and the debug stream pollutes the run-directory event log. +- Quiet asyncio's RuntimeWarning + drop its log propagation; some + litellm async paths emit benign cleanup warnings. +""" + import logging import warnings import litellm -from .config import LLMConfig -from .llm import LLM, LLMRequestFailedError - - -__all__ = [ - "LLM", - "LLMConfig", - "LLMRequestFailedError", -] litellm._logging._disable_debugging() # type: ignore[no-untyped-call] logging.getLogger("asyncio").setLevel(logging.CRITICAL) diff --git a/strix/llm/anthropic_cache_wrapper.py b/strix/llm/anthropic_cache_wrapper.py index 52f9107..ff1b38f 100644 --- a/strix/llm/anthropic_cache_wrapper.py +++ b/strix/llm/anthropic_cache_wrapper.py @@ -28,8 +28,8 @@ class AnthropicCachingLitellmModel(LitellmModel): """LitellmModel that injects ``cache_control: {"type": "ephemeral"}`` on the system message for Anthropic models. Other providers pass through unchanged. - Detection follows the legacy Strix logic: case-insensitive substring match - on ``"anthropic/"`` or ``"claude"`` against the model name (llm/llm.py:338-341). + Detection: case-insensitive substring match on ``"anthropic/"`` or + ``"claude"`` against the model name. For Strix proxy routing where the API model is ``openai/`` but the underlying provider is still Anthropic (e.g., ``strix/claude-sonnet-4.6`` diff --git a/strix/llm/config.py b/strix/llm/config.py deleted file mode 100644 index 017c776..0000000 --- a/strix/llm/config.py +++ /dev/null @@ -1,40 +0,0 @@ -from typing import Any - -from strix.config import Config -from strix.config.config import resolve_llm_config -from strix.llm.utils import resolve_strix_model - - -class LLMConfig: - def __init__( - self, - model_name: str | None = None, - enable_prompt_caching: bool = True, - skills: list[str] | None = None, - timeout: int | None = None, - scan_mode: str = "deep", - is_whitebox: bool = False, - interactive: bool = False, - reasoning_effort: str | None = None, - system_prompt_context: dict[str, Any] | None = None, - ): - resolved_model, self.api_key, self.api_base = resolve_llm_config() - self.model_name = model_name or resolved_model - - if not self.model_name: - raise ValueError("STRIX_LLM environment variable must be set and not empty") - - api_model, canonical = resolve_strix_model(self.model_name) - self.litellm_model: str = api_model or self.model_name - self.canonical_model: str = canonical or self.model_name - - self.enable_prompt_caching = enable_prompt_caching - self.skills = skills or [] - - self.timeout = timeout or int(Config.get("llm_timeout") or "300") - - self.scan_mode = scan_mode if scan_mode in ["quick", "standard", "deep"] else "deep" - self.is_whitebox = is_whitebox - self.interactive = interactive - self.reasoning_effort = reasoning_effort - self.system_prompt_context = system_prompt_context or {} diff --git a/strix/llm/dedupe.py b/strix/llm/dedupe.py index 0ea6088..292b8bc 100644 --- a/strix/llm/dedupe.py +++ b/strix/llm/dedupe.py @@ -6,7 +6,7 @@ from typing import Any import litellm from strix.config.config import resolve_llm_config -from strix.llm.utils import resolve_strix_model +from strix.llm.multi_provider_setup import STRIX_MODEL_MAP logger = logging.getLogger(__name__) @@ -157,8 +157,11 @@ def check_duplicate( comparison_data = {"candidate": candidate_cleaned, "existing_reports": existing_cleaned} model_name, api_key, api_base = resolve_llm_config() - litellm_model, _ = resolve_strix_model(model_name) - litellm_model = litellm_model or model_name + litellm_model: str | None = model_name + if model_name and model_name.startswith("strix/"): + base = model_name[len("strix/") :] + if base in STRIX_MODEL_MAP: + litellm_model = f"openai/{base}" messages = [ {"role": "system", "content": DEDUPE_SYSTEM_PROMPT}, diff --git a/strix/llm/llm.py b/strix/llm/llm.py deleted file mode 100644 index d45533f..0000000 --- a/strix/llm/llm.py +++ /dev/null @@ -1,390 +0,0 @@ -import asyncio -from collections.abc import AsyncIterator -from dataclasses import dataclass -from typing import Any - -import litellm -from jinja2 import Environment, FileSystemLoader, select_autoescape -from litellm import acompletion, completion_cost, stream_chunk_builder, supports_reasoning -from litellm.utils import supports_prompt_caching, supports_vision - -from strix.config import Config -from strix.llm.config import LLMConfig -from strix.llm.memory_compressor import MemoryCompressor -from strix.llm.utils import ( - _truncate_to_first_function, - fix_incomplete_tool_call, - normalize_tool_format, - parse_tool_invocations, -) -from strix.skills import load_skills -from strix.tools import get_tools_prompt -from strix.utils.resource_paths import get_strix_resource_path - - -litellm.drop_params = True -litellm.modify_params = True - - -class LLMRequestFailedError(Exception): - def __init__(self, message: str, details: str | None = None): - super().__init__(message) - self.message = message - self.details = details - - -@dataclass -class LLMResponse: - content: str - tool_invocations: list[dict[str, Any]] | None = None - thinking_blocks: list[dict[str, Any]] | None = None - - -@dataclass -class RequestStats: - input_tokens: int = 0 - output_tokens: int = 0 - cached_tokens: int = 0 - cost: float = 0.0 - requests: int = 0 - - def to_dict(self) -> dict[str, int | float]: - return { - "input_tokens": self.input_tokens, - "output_tokens": self.output_tokens, - "cached_tokens": self.cached_tokens, - "cost": round(self.cost, 4), - "requests": self.requests, - } - - -class LLM: - def __init__(self, config: LLMConfig, agent_name: str | None = None): - self.config = config - self.agent_name = agent_name - self.agent_id: str | None = None - self._active_skills: list[str] = list(config.skills or []) - self._system_prompt_context: dict[str, Any] = dict( - getattr(config, "system_prompt_context", {}) or {} - ) - self._total_stats = RequestStats() - self.memory_compressor = MemoryCompressor(model_name=config.litellm_model) - self.system_prompt = self._load_system_prompt(agent_name) - - reasoning = Config.get("strix_reasoning_effort") - if reasoning: - self._reasoning_effort = reasoning - elif config.reasoning_effort: - self._reasoning_effort = config.reasoning_effort - elif config.scan_mode == "quick": - self._reasoning_effort = "medium" - else: - self._reasoning_effort = "high" - - def _load_system_prompt(self, agent_name: str | None) -> str: - if not agent_name: - return "" - - try: - prompt_dir = get_strix_resource_path("agents", agent_name) - skills_dir = get_strix_resource_path("skills") - env = Environment( - loader=FileSystemLoader([prompt_dir, skills_dir]), - autoescape=select_autoescape(enabled_extensions=(), default_for_string=False), - ) - - skills_to_load = self._get_skills_to_load() - skill_content = load_skills(skills_to_load) - env.globals["get_skill"] = lambda name: skill_content.get(name, "") - - result = env.get_template("system_prompt.jinja").render( - get_tools_prompt=get_tools_prompt, - loaded_skill_names=list(skill_content.keys()), - interactive=self.config.interactive, - system_prompt_context=self._system_prompt_context, - **skill_content, - ) - return str(result) - except Exception: # noqa: BLE001 - return "" - - def _get_skills_to_load(self) -> list[str]: - ordered_skills = [*self._active_skills] - ordered_skills.append(f"scan_modes/{self.config.scan_mode}") - if self.config.is_whitebox: - ordered_skills.append("coordination/source_aware_whitebox") - ordered_skills.append("custom/source_aware_sast") - - deduped: list[str] = [] - seen: set[str] = set() - for skill_name in ordered_skills: - if skill_name not in seen: - deduped.append(skill_name) - seen.add(skill_name) - - return deduped - - def add_skills(self, skill_names: list[str]) -> list[str]: - added: list[str] = [] - for skill_name in skill_names: - if not skill_name or skill_name in self._active_skills: - continue - self._active_skills.append(skill_name) - added.append(skill_name) - - if not added: - return [] - - updated_prompt = self._load_system_prompt(self.agent_name) - if updated_prompt: - self.system_prompt = updated_prompt - - return added - - def set_agent_identity(self, agent_name: str | None, agent_id: str | None) -> None: - if agent_name: - self.agent_name = agent_name - if agent_id: - self.agent_id = agent_id - - def set_system_prompt_context(self, context: dict[str, Any] | None) -> None: - self._system_prompt_context = dict(context or {}) - updated_prompt = self._load_system_prompt(self.agent_name) - if updated_prompt: - self.system_prompt = updated_prompt - - async def generate( - self, conversation_history: list[dict[str, Any]] - ) -> AsyncIterator[LLMResponse]: - messages = self._prepare_messages(conversation_history) - max_retries = int(Config.get("strix_llm_max_retries") or "5") - - for attempt in range(max_retries + 1): - try: - async for response in self._stream(messages): - yield response - return # noqa: TRY300 - except Exception as e: # noqa: BLE001 - if attempt >= max_retries or not self._should_retry(e): - self._raise_error(e) - wait = min(90, 2 * (2**attempt)) - await asyncio.sleep(wait) - - async def _stream(self, messages: list[dict[str, Any]]) -> AsyncIterator[LLMResponse]: - accumulated = "" - chunks: list[Any] = [] - done_streaming = 0 - - self._total_stats.requests += 1 - timeout = self.config.timeout - response = await asyncio.wait_for( - acompletion(**self._build_completion_args(messages), stream=True), - timeout=timeout, - ) - - async_iter = response.__aiter__() - while True: - try: - chunk = await asyncio.wait_for(async_iter.__anext__(), timeout=timeout) - except StopAsyncIteration: - break - chunks.append(chunk) - if done_streaming: - done_streaming += 1 - if getattr(chunk, "usage", None) or done_streaming > 5: - break - continue - delta = self._get_chunk_content(chunk) - if delta: - accumulated += delta - if "" in accumulated or "" in accumulated: - end_tag = "" if "" in accumulated else "" - pos = accumulated.find(end_tag) - accumulated = accumulated[: pos + len(end_tag)] - yield LLMResponse(content=accumulated) - done_streaming = 1 - continue - yield LLMResponse(content=accumulated) - - if chunks: - self._update_usage_stats(stream_chunk_builder(chunks)) - - accumulated = normalize_tool_format(accumulated) - accumulated = fix_incomplete_tool_call(_truncate_to_first_function(accumulated)) - yield LLMResponse( - content=accumulated, - tool_invocations=parse_tool_invocations(accumulated), - thinking_blocks=self._extract_thinking(chunks), - ) - - def _prepare_messages(self, conversation_history: list[dict[str, Any]]) -> list[dict[str, Any]]: - messages = [{"role": "system", "content": self.system_prompt}] - - if self.agent_name: - messages.append( - { - "role": "user", - "content": ( - f"\n\n\n" - f"Internal metadata: do not echo or reference.\n" - f"{self.agent_name}\n" - f"{self.agent_id}\n" - f"\n\n" - ), - } - ) - - compressed = list(self.memory_compressor.compress_history(conversation_history)) - conversation_history.clear() - conversation_history.extend(compressed) - messages.extend(compressed) - - if messages[-1].get("role") == "assistant" and not self.config.interactive: - messages.append({"role": "user", "content": "Continue the task."}) - - if self._is_anthropic() and self.config.enable_prompt_caching: - messages = self._add_cache_control(messages) - - return messages - - def _build_completion_args(self, messages: list[dict[str, Any]]) -> dict[str, Any]: - if not self._supports_vision(): - messages = self._strip_images(messages) - - args: dict[str, Any] = { - "model": self.config.litellm_model, - "messages": messages, - "timeout": self.config.timeout, - "stream_options": {"include_usage": True}, - } - - if self.config.api_key: - args["api_key"] = self.config.api_key - if self.config.api_base: - args["api_base"] = self.config.api_base - if self._supports_reasoning(): - args["reasoning_effort"] = self._reasoning_effort - - return args - - def _get_chunk_content(self, chunk: Any) -> str: - if chunk.choices and hasattr(chunk.choices[0], "delta"): - return getattr(chunk.choices[0].delta, "content", "") or "" - return "" - - def _extract_thinking(self, chunks: list[Any]) -> list[dict[str, Any]] | None: - if not chunks or not self._supports_reasoning(): - return None - blocks: list[dict[str, Any]] | None = None - try: - resp = stream_chunk_builder(chunks) - choices: Any = getattr(resp, "choices", None) - if choices: - message: Any = getattr(choices[0], "message", None) - if message is not None and hasattr(message, "thinking_blocks"): - blocks = message.thinking_blocks - except Exception: # noqa: BLE001, S110 # nosec B110 - pass - return blocks - - def _update_usage_stats(self, response: Any) -> None: - try: - if hasattr(response, "usage") and response.usage: - input_tokens = getattr(response.usage, "prompt_tokens", 0) or 0 - output_tokens = getattr(response.usage, "completion_tokens", 0) or 0 - - cached_tokens = 0 - if hasattr(response.usage, "prompt_tokens_details"): - prompt_details = response.usage.prompt_tokens_details - if hasattr(prompt_details, "cached_tokens"): - cached_tokens = prompt_details.cached_tokens or 0 - - cost = self._extract_cost(response) - else: - input_tokens = 0 - output_tokens = 0 - cached_tokens = 0 - cost = 0.0 - - self._total_stats.input_tokens += input_tokens - self._total_stats.output_tokens += output_tokens - self._total_stats.cached_tokens += cached_tokens - self._total_stats.cost += cost - - except Exception: # noqa: BLE001, S110 # nosec B110 - pass - - def _extract_cost(self, response: Any) -> float: - if hasattr(response, "usage") and response.usage: - direct_cost = getattr(response.usage, "cost", None) - if direct_cost is not None: - return float(direct_cost) - try: - if hasattr(response, "_hidden_params"): - response._hidden_params.pop("custom_llm_provider", None) - return completion_cost(response, model=self.config.canonical_model) or 0.0 - except Exception: # noqa: BLE001 - return 0.0 - - def _should_retry(self, e: Exception) -> bool: - code = getattr(e, "status_code", None) or getattr( - getattr(e, "response", None), "status_code", None - ) - return code is None or litellm._should_retry(code) - - def _raise_error(self, e: Exception) -> None: - from strix.telemetry import posthog - - posthog.error("llm_error", type(e).__name__) - raise LLMRequestFailedError(f"LLM request failed: {type(e).__name__}", str(e)) from e - - def _is_anthropic(self) -> bool: - if not self.config.model_name: - return False - return any(p in self.config.model_name.lower() for p in ["anthropic/", "claude"]) - - def _supports_vision(self) -> bool: - try: - return bool(supports_vision(model=self.config.canonical_model)) - except Exception: # noqa: BLE001 - return False - - def _supports_reasoning(self) -> bool: - try: - return bool(supports_reasoning(model=self.config.canonical_model)) - except Exception: # noqa: BLE001 - return False - - def _strip_images(self, messages: list[dict[str, Any]]) -> list[dict[str, Any]]: - result = [] - for msg in messages: - content = msg.get("content") - if isinstance(content, list): - text_parts = [] - for item in content: - if isinstance(item, dict) and item.get("type") == "text": - text_parts.append(item.get("text", "")) - elif isinstance(item, dict) and item.get("type") == "image_url": - text_parts.append("[Image removed - model doesn't support vision]") - result.append({**msg, "content": "\n".join(text_parts)}) - else: - result.append(msg) - return result - - def _add_cache_control(self, messages: list[dict[str, Any]]) -> list[dict[str, Any]]: - if not messages or not supports_prompt_caching(self.config.canonical_model): - return messages - - result = list(messages) - - if result[0].get("role") == "system": - content = result[0]["content"] - result[0] = { - **result[0], - "content": [ - {"type": "text", "text": content, "cache_control": {"type": "ephemeral"}} - ] - if isinstance(content, str) - else content, - } - return result diff --git a/strix/llm/multi_provider_setup.py b/strix/llm/multi_provider_setup.py index f112332..16cb752 100644 --- a/strix/llm/multi_provider_setup.py +++ b/strix/llm/multi_provider_setup.py @@ -15,8 +15,6 @@ Other prefixes fall through to the SDK's built-in OpenAI / LiteLLM defaults. References: - PLAYBOOK.md §2.7 - AUDIT_R3.md C17 (model alias validation; raise UserError on unknown alias) - - Legacy: strix/llm/utils.py STRIX_MODEL_MAP and resolve_strix_model - - Legacy: strix/config/config.py STRIX_API_BASE """ from __future__ import annotations @@ -28,7 +26,23 @@ from agents.models.multi_provider import MultiProvider, MultiProviderMap from strix.config.config import STRIX_API_BASE from strix.llm.anthropic_cache_wrapper import AnthropicCachingLitellmModel -from strix.llm.utils import STRIX_MODEL_MAP + + +# Strix-proxy aliases. Each maps the user-facing alias (right of +# ``strix/``) to the canonical provider/model used for capability +# lookups (litellm reads e.g. ``anthropic/claude-sonnet-4-6`` to +# decide on prompt-caching support). +STRIX_MODEL_MAP: dict[str, str] = { + "claude-sonnet-4.6": "anthropic/claude-sonnet-4-6", + "claude-opus-4.6": "anthropic/claude-opus-4-6", + "gpt-5.2": "openai/gpt-5.2", + "gpt-5.1": "openai/gpt-5.1", + "gpt-5.4": "openai/gpt-5.4", + "gemini-3-pro-preview": "gemini/gemini-3-pro-preview", + "gemini-3-flash-preview": "gemini/gemini-3-flash-preview", + "glm-5": "openrouter/z-ai/glm-5", + "glm-4.7": "openrouter/z-ai/glm-4.7", +} def _is_anthropic_canonical(canonical: str) -> bool: diff --git a/strix/llm/strix_session.py b/strix/llm/strix_session.py index 6507437..2e51363 100644 --- a/strix/llm/strix_session.py +++ b/strix/llm/strix_session.py @@ -1,9 +1,9 @@ -"""StrixSession — Session wrapper that runs the legacy MemoryCompressor. +"""StrixSession — Session wrapper that runs the MemoryCompressor. The SDK's `Session` (and ``SessionABC``) protocol owns conversation history storage. We delegate the actual storage to any underlying session implementation (in-memory, SQLite, Redis, …) and intercept ``get_items`` so -the legacy ``MemoryCompressor`` runs before the model sees the history. +the ``MemoryCompressor`` runs before the model sees the history. Why wrap rather than reimplement: - ``MemoryCompressor`` already encodes the pentest-tuned summarization diff --git a/strix/llm/utils.py b/strix/llm/utils.py index 9771854..9e5d175 100644 --- a/strix/llm/utils.py +++ b/strix/llm/utils.py @@ -1,3 +1,15 @@ +"""Streaming + tool-format helpers used by the TUI's render pipeline. + +The model can emit tool calls in a few XML shapes (````, +````, optionally wrapped in ````); the +streaming parser normalizes them into one canonical form so the +renderer doesn't have to branch. + +These helpers are pure string manipulation — no model client, no SDK +dependency. They live here because the streaming parser and the +agent-message renderer both consume them. +""" + import html import re from typing import Any @@ -27,56 +39,11 @@ def normalize_tool_format(content: str) -> str: content = content.replace("", "") return _STRIP_TAG_QUOTES.sub( - lambda m: f"<{m.group(1)}={m.group(2).strip().strip(chr(34) + chr(39))}>", content + lambda m: f"<{m.group(1)}={m.group(2).strip().strip(chr(34) + chr(39))}>", + content, ) -STRIX_MODEL_MAP: dict[str, str] = { - "claude-sonnet-4.6": "anthropic/claude-sonnet-4-6", - "claude-opus-4.6": "anthropic/claude-opus-4-6", - "gpt-5.2": "openai/gpt-5.2", - "gpt-5.1": "openai/gpt-5.1", - "gpt-5.4": "openai/gpt-5.4", - "gemini-3-pro-preview": "gemini/gemini-3-pro-preview", - "gemini-3-flash-preview": "gemini/gemini-3-flash-preview", - "glm-5": "openrouter/z-ai/glm-5", - "glm-4.7": "openrouter/z-ai/glm-4.7", -} - - -def resolve_strix_model(model_name: str | None) -> tuple[str | None, str | None]: - """Resolve a strix/ model into names for API calls and capability lookups. - - Returns (api_model, canonical_model): - - api_model: openai/ for API calls (Strix API is OpenAI-compatible) - - canonical_model: actual provider model name for litellm capability lookups - Non-strix models return the same name for both. - """ - if not model_name or not model_name.startswith("strix/"): - return model_name, model_name - - base_model = model_name[6:] - api_model = f"openai/{base_model}" - canonical_model = STRIX_MODEL_MAP.get(base_model, api_model) - return api_model, canonical_model - - -def _truncate_to_first_function(content: str) -> str: - if not content: - return content - - function_starts = [ - match.start() for match in re.finditer(r"= 2: - second_function_start = function_starts[1] - - return content[:second_function_start].rstrip() - - return content - - def parse_tool_invocations(content: str) -> list[dict[str, Any]] | None: content = normalize_tool_format(content) content = fix_incomplete_tool_call(content) diff --git a/strix/orchestration/bus.py b/strix/orchestration/bus.py index e978c17..cd2f9f3 100644 --- a/strix/orchestration/bus.py +++ b/strix/orchestration/bus.py @@ -1,8 +1,8 @@ """AgentMessageBus — peer-to-peer multi-agent state owned by Strix. -Replaces the legacy harness's _agent_graph / _agent_messages / _agent_instances -globals (in strix/tools/agents_graph/agents_graph_actions.py) with a single -asyncio.Lock-protected dataclass that lives for the lifetime of one Strix scan. +A single ``asyncio.Lock``-protected dataclass that owns inboxes, +parent edges, statuses, and per-agent stats for the lifetime of one +Strix scan. References: - PLAYBOOK.md §2.3 diff --git a/strix/orchestration/filter.py b/strix/orchestration/filter.py index cec4681..e402092 100644 --- a/strix/orchestration/filter.py +++ b/strix/orchestration/filter.py @@ -1,12 +1,10 @@ -"""inject_messages_filter — SDK call_model_input_filter for the message bus. +"""inject_messages_filter — SDK ``call_model_input_filter`` for the message bus. -This is the integration point that replaces Strix's per-iteration -_check_agent_messages call (legacy: agents/base_agent.py:448-531). The SDK -runs ``call_model_input_filter`` exactly once per turn before the LLM call -(``run_internal/turn_preparation.py:55-80``), and captures the filter's -output in a lambda closure for any subsequent retries -(``run_internal/model_retry.py:34-35``) — so a single drain per turn does -not lose messages on retry. +The SDK runs ``call_model_input_filter`` exactly once per turn before +the LLM call (``run_internal/turn_preparation.py:55-80``) and captures +the filter's output in a lambda closure for any subsequent retries +(``run_internal/model_retry.py:34-35``) — so a single drain per turn +does not lose messages on retry. References: - PLAYBOOK.md §2.4 @@ -32,8 +30,7 @@ async def inject_messages_filter(data: CallModelData) -> ModelInputData: """Drain bus inbox and append messages as user-role items before the LLM call. Each drained message is wrapped in an ```` XML envelope - that mirrors Strix's legacy format (base_agent.py:491-514) so the system - prompt's existing rules around inter-agent communication still apply. + so the system prompt's rules around inter-agent communication apply. Messages from the literal sender ``"user"`` (a real human via TUI) skip the XML wrap and are added as plain user messages. diff --git a/strix/orchestration/hooks.py b/strix/orchestration/hooks.py index 6119d9d..6fefe70 100644 --- a/strix/orchestration/hooks.py +++ b/strix/orchestration/hooks.py @@ -28,9 +28,8 @@ class StrixOrchestrationHooks(RunHooks[Any]): Wires four concerns: 1. Turn-budget warnings injected into ``input_items`` at 85% and ``N - 3`` - of ``max_turns`` (legacy: ``base_agent.py:186-211``). - 2. LLM usage recording into the bus (replaces legacy ``LLM._total_stats`` - + ``_completed_agent_llm_totals``). + of ``max_turns``. + 2. LLM usage recording into the bus. 3. Sandbox readiness: awaits the ``CaidoCapability._healthcheck_task`` on first agent start so the agent doesn't fire tools before Caido and the tool server are ready. diff --git a/strix/run_config_factory.py b/strix/run_config_factory.py index 06ab512..ecb9dd9 100644 --- a/strix/run_config_factory.py +++ b/strix/run_config_factory.py @@ -8,7 +8,7 @@ for the rare case a single run wants different reasoning effort or References: - PLAYBOOK.md §2.10 - AUDIT.md §2.1 (C1 — parallel_tool_calls=False until Phase 6 relaxes the - legacy tool server's per-agent task slot serialization) + tool server's per-agent task slot serialization) - AUDIT_R2.md §1.6 (C11 — retry policy explicitly excludes 401/403/400; auth and validation errors must fail fast, not waste retries) - AUDIT_R3.md C21 — RunConfig override + context fields including @@ -39,9 +39,9 @@ if TYPE_CHECKING: from strix.orchestration.bus import AgentMessageBus -# Phase 1-5 default. Phase 6 relaxes the legacy tool server's per-agent -# task-slot serialization (``runtime/tool_server.py:94-97``) and flips this -# to ``True`` after the multi-agent stress tests confirm safety. +# Phase 6 relaxes the tool server's per-agent task-slot serialization +# (``runtime/tool_server.py:94-97``) and flips this to ``True`` after +# multi-agent stress tests confirm safety. _PHASE1_PARALLEL_DEFAULT = False # Default retry policy. Explicitly does NOT include 401/403/400 — those are @@ -49,8 +49,7 @@ _PHASE1_PARALLEL_DEFAULT = False # so the user sees the real error within seconds. 429/5xx is the right set. _RETRYABLE_HTTP_STATUSES = (429, 500, 502, 503, 504) -# Default retry budget. Mirrors the legacy ``llm.py`` retry loop: 5 attempts -# with ``min(90, 2*2^n)`` backoff. +# Default retry budget: 5 attempts with ``min(90, 2*2^n)`` backoff. _DEFAULT_MAX_RETRIES = 5 _DEFAULT_BACKOFF = ModelRetryBackoffSettings( initial_delay=2.0, @@ -76,8 +75,6 @@ def _default_retry_policy() -> Any: #: Default ``max_turns`` callers should pass to ``Runner.run``. -#: Mirrors the legacy ``AgentState.max_iterations = 300`` -#: (``HARNESS_WIKI.md §5.2``). STRIX_DEFAULT_MAX_TURNS = 300 @@ -105,11 +102,10 @@ def make_run_config( ``None`` is allowed for unit tests and dry runs. model: Model alias to pass to ``MultiProvider``. Defaults to the current production-favored Anthropic alias. - parallel_tool_calls: Phase 1 default is ``False`` to keep behavior - sequential per the legacy tool server's slot serialization (C1). + parallel_tool_calls: Default ``False`` to keep behavior sequential + per the tool server's slot serialization (C1). tool_choice: Forces tool use per turn unless explicitly relaxed. - Mirrors the legacy ``4f90a56`` prompt hardening at the model - level. Pass ``None`` to omit. + Pass ``None`` to omit. reasoning_effort: ``"low" | "medium" | "high"``; routes to ``ModelSettings.reasoning``. ``None`` defers to provider default. model_settings_override: Optional ``ModelSettings`` to merge over @@ -182,15 +178,10 @@ def make_agent_context( tracer reference, and per-agent toggles live. Tools, hooks, and the ``inject_messages_filter`` all reach in via ``ctx.context.get(...)``. - C21 (AUDIT_R3): includes ``is_whitebox``, ``diff_scope`` and ``run_id`` - fields that the legacy code relied on but the original PLAYBOOK §2.10 - skeleton omitted. - ``agent_factory`` is a callable ``(name, skills) -> agents.Agent`` used by - the ``create_agent`` graph tool to spin up children. The actual factory - lives in the Phase 4/5 root-assembly module; Phase 3 only requires that - it be present in context. ``sandbox_client`` is the host-side Docker - subclass; ``create_agent`` reuses it across child runs. + the ``create_agent`` graph tool to spin up children. ``sandbox_client`` + is the host-side Docker subclass; ``create_agent`` reuses it across + child runs. """ return { "bus": bus, diff --git a/strix/runtime/__init__.py b/strix/runtime/__init__.py index 5d0cbda..72f3b1a 100644 --- a/strix/runtime/__init__.py +++ b/strix/runtime/__init__.py @@ -1,6 +1,22 @@ -from strix.config import Config +"""Strix runtime package. -from .runtime import AbstractRuntime +What lives here: + +- :class:`StrixDockerSandboxClient` — host-side ``DockerSandboxClient`` + subclass that injects ``NET_ADMIN`` / ``NET_RAW`` capabilities and + ``host.docker.internal`` extra-hosts, used by the per-scan session + manager (:mod:`strix.sandbox.session_manager`). + +- ``tool_server.py`` — the FastAPI server that runs *inside* the + sandbox container; sandbox-bound tools (browser, terminal, python, + file_edit, proxy) POST here from the host via + :func:`strix.tools._sandbox_dispatch.post_to_sandbox`. + +The legacy DockerRuntime / AbstractRuntime + ``get_runtime`` / +``cleanup_runtime`` globals were removed when the SDK harness took +over scan lifecycle; sandbox sessions are now per-scan and managed by +:func:`strix.sandbox.session_manager.create_or_reuse`. +""" class SandboxInitializationError(Exception): @@ -12,32 +28,4 @@ class SandboxInitializationError(Exception): self.details = details -_global_runtime: AbstractRuntime | None = None - - -def get_runtime() -> AbstractRuntime: - global _global_runtime # noqa: PLW0603 - - runtime_backend = Config.get("strix_runtime_backend") - - if runtime_backend == "docker": - from .docker_runtime import DockerRuntime - - if _global_runtime is None: - _global_runtime = DockerRuntime() - return _global_runtime - - raise ValueError( - f"Unsupported runtime backend: {runtime_backend}. Only 'docker' is supported for now." - ) - - -def cleanup_runtime() -> None: - global _global_runtime # noqa: PLW0603 - - if _global_runtime is not None: - _global_runtime.cleanup() - _global_runtime = None - - -__all__ = ["AbstractRuntime", "SandboxInitializationError", "cleanup_runtime", "get_runtime"] +__all__ = ["SandboxInitializationError"] diff --git a/strix/runtime/docker_runtime.py b/strix/runtime/docker_runtime.py deleted file mode 100644 index d57d358..0000000 --- a/strix/runtime/docker_runtime.py +++ /dev/null @@ -1,352 +0,0 @@ -import contextlib -import os -import secrets -import socket -import time -from pathlib import Path -from typing import cast - -import docker -import httpx -from docker.errors import DockerException, ImageNotFound, NotFound -from docker.models.containers import Container -from requests.exceptions import ConnectionError as RequestsConnectionError -from requests.exceptions import Timeout as RequestsTimeout - -from strix.config import Config - -from . import SandboxInitializationError -from .runtime import AbstractRuntime, SandboxInfo - - -HOST_GATEWAY_HOSTNAME = "host.docker.internal" -DOCKER_TIMEOUT = 60 -CONTAINER_TOOL_SERVER_PORT = 48081 -CONTAINER_CAIDO_PORT = 48080 - - -class DockerRuntime(AbstractRuntime): - def __init__(self) -> None: - try: - self.client = docker.from_env(timeout=DOCKER_TIMEOUT) - except (DockerException, RequestsConnectionError, RequestsTimeout) as e: - raise SandboxInitializationError( - "Docker is not available", - "Please ensure Docker Desktop is installed and running.", - ) from e - - self._scan_container: Container | None = None - self._tool_server_port: int | None = None - self._tool_server_token: str | None = None - self._caido_port: int | None = None - - def _find_available_port(self) -> int: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(("", 0)) - return cast("int", s.getsockname()[1]) - - def _get_scan_id(self, agent_id: str) -> str: - try: - from strix.telemetry.tracer import get_global_tracer - - tracer = get_global_tracer() - if tracer and tracer.scan_config: - return str(tracer.scan_config.get("scan_id", "default-scan")) - except (ImportError, AttributeError): - pass - return f"scan-{agent_id.split('-')[0]}" - - def _verify_image_available(self, image_name: str, max_retries: int = 3) -> None: - for attempt in range(max_retries): - try: - image = self.client.images.get(image_name) - if not image.id or not image.attrs: - raise ImageNotFound(f"Image {image_name} metadata incomplete") # noqa: TRY301 - except (ImageNotFound, DockerException): - if attempt == max_retries - 1: - raise - time.sleep(2**attempt) - else: - return - - def _recover_container_state(self, container: Container) -> None: - for env_var in container.attrs["Config"]["Env"]: - if env_var.startswith("TOOL_SERVER_TOKEN="): - self._tool_server_token = env_var.split("=", 1)[1] - break - - port_bindings = container.attrs.get("NetworkSettings", {}).get("Ports", {}) - port_key = f"{CONTAINER_TOOL_SERVER_PORT}/tcp" - if port_bindings.get(port_key): - self._tool_server_port = int(port_bindings[port_key][0]["HostPort"]) - - caido_port_key = f"{CONTAINER_CAIDO_PORT}/tcp" - if port_bindings.get(caido_port_key): - self._caido_port = int(port_bindings[caido_port_key][0]["HostPort"]) - - def _wait_for_tool_server(self, max_retries: int = 30, timeout: int = 5) -> None: - host = self._resolve_docker_host() - health_url = f"http://{host}:{self._tool_server_port}/health" - - time.sleep(5) - - for attempt in range(max_retries): - try: - with httpx.Client(trust_env=False, timeout=timeout) as client: - response = client.get(health_url) - if response.status_code == 200: - data = response.json() - if data.get("status") == "healthy": - return - except (httpx.ConnectError, httpx.TimeoutException, httpx.RequestError): - pass - - time.sleep(min(2**attempt * 0.5, 5)) - - raise SandboxInitializationError( - "Tool server failed to start", - "Container initialization timed out. Please try again.", - ) - - def _create_container(self, scan_id: str, max_retries: int = 2) -> Container: - container_name = f"strix-scan-{scan_id}" - image_name = Config.get("strix_image") - if not image_name: - raise ValueError("STRIX_IMAGE must be configured") - - self._verify_image_available(image_name) - - last_error: Exception | None = None - for attempt in range(max_retries + 1): - try: - with contextlib.suppress(NotFound): - existing = self.client.containers.get(container_name) - with contextlib.suppress(Exception): - existing.stop(timeout=5) - existing.remove(force=True) - time.sleep(1) - - self._tool_server_port = self._find_available_port() - self._caido_port = self._find_available_port() - self._tool_server_token = secrets.token_urlsafe(32) - execution_timeout = Config.get("strix_sandbox_execution_timeout") or "120" - - container = self.client.containers.run( - image_name, - command="sleep infinity", - detach=True, - name=container_name, - hostname=container_name, - ports={ - f"{CONTAINER_TOOL_SERVER_PORT}/tcp": self._tool_server_port, - f"{CONTAINER_CAIDO_PORT}/tcp": self._caido_port, - }, - cap_add=["NET_ADMIN", "NET_RAW"], - labels={"strix-scan-id": scan_id}, - environment={ - "PYTHONUNBUFFERED": "1", - "TOOL_SERVER_PORT": str(CONTAINER_TOOL_SERVER_PORT), - "TOOL_SERVER_TOKEN": self._tool_server_token, - "STRIX_SANDBOX_EXECUTION_TIMEOUT": str(execution_timeout), - "HOST_GATEWAY": HOST_GATEWAY_HOSTNAME, - }, - extra_hosts={HOST_GATEWAY_HOSTNAME: "host-gateway"}, - tty=True, - ) - - self._scan_container = container - self._wait_for_tool_server() - - except (DockerException, RequestsConnectionError, RequestsTimeout) as e: - last_error = e - if attempt < max_retries: - self._tool_server_port = None - self._tool_server_token = None - self._caido_port = None - time.sleep(2**attempt) - else: - return container - - raise SandboxInitializationError( - "Failed to create container", - f"Container creation failed after {max_retries + 1} attempts: {last_error}", - ) from last_error - - def _get_or_create_container(self, scan_id: str) -> Container: - container_name = f"strix-scan-{scan_id}" - - if self._scan_container: - try: - self._scan_container.reload() - if self._scan_container.status == "running": - return self._scan_container - except NotFound: - self._scan_container = None - self._tool_server_port = None - self._tool_server_token = None - self._caido_port = None - - try: - container = self.client.containers.get(container_name) - container.reload() - - if container.status != "running": - container.start() - time.sleep(2) - - self._scan_container = container - self._recover_container_state(container) - except NotFound: - pass - else: - return container - - try: - containers = self.client.containers.list( - all=True, filters={"label": f"strix-scan-id={scan_id}"} - ) - if containers: - container = containers[0] - if container.status != "running": - container.start() - time.sleep(2) - - self._scan_container = container - self._recover_container_state(container) - return container - except DockerException: - pass - - return self._create_container(scan_id) - - def _copy_local_directory_to_container( - self, container: Container, local_path: str, target_name: str | None = None - ) -> None: - import tarfile - from io import BytesIO - - try: - local_path_obj = Path(local_path).resolve() - if not local_path_obj.exists() or not local_path_obj.is_dir(): - return - - tar_buffer = BytesIO() - with tarfile.open(fileobj=tar_buffer, mode="w") as tar: - for item in local_path_obj.rglob("*"): - if item.is_file(): - rel_path = item.relative_to(local_path_obj) - arcname = Path(target_name) / rel_path if target_name else rel_path - tar.add(item, arcname=arcname) - - tar_buffer.seek(0) - container.put_archive("/workspace", tar_buffer.getvalue()) - container.exec_run( - "chown -R pentester:pentester /workspace && chmod -R 755 /workspace", - user="root", - ) - except (OSError, DockerException): - pass - - async def create_sandbox( - self, - agent_id: str, - existing_token: str | None = None, - local_sources: list[dict[str, str]] | None = None, - ) -> SandboxInfo: - scan_id = self._get_scan_id(agent_id) - container = self._get_or_create_container(scan_id) - - source_copied_key = f"_source_copied_{scan_id}" - if local_sources and not hasattr(self, source_copied_key): - for index, source in enumerate(local_sources, start=1): - source_path = source.get("source_path") - if not source_path: - continue - target_name = ( - source.get("workspace_subdir") or Path(source_path).name or f"target_{index}" - ) - self._copy_local_directory_to_container(container, source_path, target_name) - setattr(self, source_copied_key, True) - - if container.id is None: - raise RuntimeError("Docker container ID is unexpectedly None") - - token = existing_token or self._tool_server_token - if self._tool_server_port is None or self._caido_port is None or token is None: - raise RuntimeError("Tool server not initialized") - - host = self._resolve_docker_host() - api_url = f"http://{host}:{self._tool_server_port}" - - await self._register_agent(api_url, agent_id, token) - - return { - "workspace_id": container.id, - "api_url": api_url, - "auth_token": token, - "tool_server_port": self._tool_server_port, - "caido_port": self._caido_port, - "agent_id": agent_id, - } - - async def _register_agent(self, api_url: str, agent_id: str, token: str) -> None: - try: - async with httpx.AsyncClient(trust_env=False) as client: - response = await client.post( - f"{api_url}/register_agent", - params={"agent_id": agent_id}, - headers={"Authorization": f"Bearer {token}"}, - timeout=30, - ) - response.raise_for_status() - except httpx.RequestError: - pass - - async def get_sandbox_url(self, container_id: str, port: int) -> str: - try: - self.client.containers.get(container_id) - return f"http://{self._resolve_docker_host()}:{port}" - except NotFound: - raise ValueError(f"Container {container_id} not found.") from None - - def _resolve_docker_host(self) -> str: - docker_host = os.getenv("DOCKER_HOST", "") - if docker_host: - from urllib.parse import urlparse - - parsed = urlparse(docker_host) - if parsed.scheme in ("tcp", "http", "https") and parsed.hostname: - return parsed.hostname - return "127.0.0.1" - - async def destroy_sandbox(self, container_id: str) -> None: - try: - container = self.client.containers.get(container_id) - container.stop() - container.remove() - self._scan_container = None - self._tool_server_port = None - self._tool_server_token = None - self._caido_port = None - except (NotFound, DockerException): - pass - - def cleanup(self) -> None: - if self._scan_container is not None: - container_name = self._scan_container.name - self._scan_container = None - self._tool_server_port = None - self._tool_server_token = None - self._caido_port = None - - if container_name is None: - return - - import subprocess - - subprocess.Popen( # noqa: S603 - ["docker", "rm", "-f", container_name], # noqa: S607 - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - start_new_session=True, - ) diff --git a/strix/runtime/runtime.py b/strix/runtime/runtime.py deleted file mode 100644 index e523d51..0000000 --- a/strix/runtime/runtime.py +++ /dev/null @@ -1,33 +0,0 @@ -from abc import ABC, abstractmethod -from typing import TypedDict - - -class SandboxInfo(TypedDict): - workspace_id: str - api_url: str - auth_token: str | None - tool_server_port: int - caido_port: int - agent_id: str - - -class AbstractRuntime(ABC): - @abstractmethod - async def create_sandbox( - self, - agent_id: str, - existing_token: str | None = None, - local_sources: list[dict[str, str]] | None = None, - ) -> SandboxInfo: - raise NotImplementedError - - @abstractmethod - async def get_sandbox_url(self, container_id: str, port: int) -> str: - raise NotImplementedError - - @abstractmethod - async def destroy_sandbox(self, container_id: str) -> None: - raise NotImplementedError - - def cleanup(self) -> None: - raise NotImplementedError diff --git a/strix/sandbox/caido_capability.py b/strix/sandbox/caido_capability.py index a3c5e28..3d6a41a 100644 --- a/strix/sandbox/caido_capability.py +++ b/strix/sandbox/caido_capability.py @@ -38,7 +38,7 @@ from agents.tool import Tool from pydantic import PrivateAttr from strix.sandbox.healthcheck import wait_for_http_ready, wait_for_tcp_ready -from strix.tools.proxy.proxy_sdk_tools import ( +from strix.tools.proxy.tools import ( list_requests, list_sitemap, repeat_request, @@ -68,7 +68,7 @@ _CAIDO_INTERNAL_PORT = 48080 _TOOL_SERVER_INTERNAL_PORT = 48081 # Probe URLs used inside ``bind``. ``host=127.0.0.1`` because the host -# port mapping is loopback-only (legacy and SDK both bind to 127.0.0.1). +# port mapping is loopback-only. _PROBE_HOST = "127.0.0.1" diff --git a/strix/sandbox/healthcheck.py b/strix/sandbox/healthcheck.py index a69f606..385b0c1 100644 --- a/strix/sandbox/healthcheck.py +++ b/strix/sandbox/healthcheck.py @@ -11,9 +11,7 @@ Two helpers are exposed: - :func:`wait_for_http_ready` for the FastAPI tool server, whose ``/health`` endpoint returns ``{"status": "healthy"}`` once the process is up. We don't require the JSON shape exactly — any 2xx - is treated as ready, mirroring the legacy ``_wait_for_tool_server`` - but more lenient (the legacy version checked the JSON body too, - which made test images without that handler fail spuriously). + is treated as ready. - :func:`wait_for_tcp_ready` for Caido, which serves an HTTP forward proxy on its port and does *not* expose ``/health``. A TCP connect @@ -21,7 +19,6 @@ Two helpers are exposed: References: - PLAYBOOK.md §3.1 - - HARNESS_WIKI.md §6.4 (legacy ``_wait_for_tool_server`` pattern) """ from __future__ import annotations @@ -40,9 +37,8 @@ class SandboxNotReadyError(Exception): """Raised when a sandbox port doesn't accept connections in time.""" -# Default per-attempt HTTP timeout. The legacy harness used 5s; we -# match it so a slow first request (image still warming up) doesn't -# misfire as a hard failure on a single attempt. +# Default per-attempt HTTP timeout. 5s so a slow first request (image +# still warming up) doesn't misfire as a hard failure on a single attempt. _DEFAULT_HTTP_PROBE_TIMEOUT = 5.0 # Default polling cadence between attempts. Balanced for CI-style diff --git a/strix/sandbox/session_manager.py b/strix/sandbox/session_manager.py index 739aa96..01c1c73 100644 --- a/strix/sandbox/session_manager.py +++ b/strix/sandbox/session_manager.py @@ -1,8 +1,6 @@ """Per-scan sandbox session lifecycle. -Replaces the legacy ``DockerRuntime`` (``strix/runtime/docker_runtime.py``) -with the SDK-native session model. One session per scan, reused across -every agent in that scan's tree. +One session per scan, reused across every agent in that scan's tree. The bundle returned by :func:`create_or_reuse` is what the per-agent context dict reads from in ``run_config_factory.make_agent_context`` — @@ -17,7 +15,6 @@ next scan from starting. References: - PLAYBOOK.md §3.3 - - HARNESS_WIKI.md §6 (legacy Docker runtime) """ from __future__ import annotations @@ -119,10 +116,9 @@ async def create_or_reuse( ), ) - # The SDK's DockerSandboxClient requires a docker.DockerClient instance - # at construction time (since openai-agents 0.14.x). We use the - # caller's environment to find the daemon — same as the legacy - # DockerRuntime did via ``docker.from_env()``. + # The SDK's DockerSandboxClient requires a docker.DockerClient + # instance at construction time (since openai-agents 0.14.x). + # ``docker.from_env()`` reads DOCKER_HOST etc. from the environment. client = StrixDockerSandboxClient(docker.from_env()) options = DockerSandboxClientOptions( image=image, diff --git a/strix/telemetry/strix_processor.py b/strix/telemetry/strix_processor.py index 5966452..0b1e9b4 100644 --- a/strix/telemetry/strix_processor.py +++ b/strix/telemetry/strix_processor.py @@ -1,10 +1,8 @@ """StrixTracingProcessor — SDK trace processor that writes events.jsonl. -Replaces the JSONL output that the legacy tracer wrote in ``telemetry/tracer.py``. -Hooks into the SDK's tracing pipeline so we keep the existing -``strix_runs//events.jsonl`` schema and the existing -``TelemetrySanitizer`` PII redaction without standing up a parallel -tracing system. +Hooks into the SDK's tracing pipeline and writes events to +``strix_runs//events.jsonl``. PII scrubbing via the existing +``TelemetrySanitizer``. References: - PLAYBOOK.md §2.9 @@ -36,7 +34,7 @@ logger = logging.getLogger(__name__) # Module-level lock registry — one per JSONL file so two processors writing # different run-dirs don't serialize unnecessarily, but two processors -# writing the *same* run-dir (e.g., legacy tracer + SDK processor) do. +# writing the *same* run-dir do. _FILE_LOCKS: dict[Path, threading.Lock] = {} _GUARD = threading.Lock() @@ -58,8 +56,8 @@ class StrixTracingProcessor(TracingProcessor): permission error during the run does NOT propagate up the hook chain and tear down the agent (C16). - PII scrubbing runs on every event before it hits the file. The - ``TelemetrySanitizer`` class is the same one the legacy tracer uses. + PII scrubbing via :class:`TelemetrySanitizer` runs on every event + before it hits the file. """ def __init__( diff --git a/strix/telemetry/tracer.py b/strix/telemetry/tracer.py index 3f3ca6c..674f4a0 100644 --- a/strix/telemetry/tracer.py +++ b/strix/telemetry/tracer.py @@ -63,6 +63,26 @@ 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, + }, + } + self.scan_results: dict[str, Any] | None = None self.scan_config: dict[str, Any] | None = None self.run_metadata: dict[str, Any] = { @@ -305,7 +325,7 @@ class Tracer: return self._run_dir - def add_vulnerability_report( # noqa: PLR0912 + def add_vulnerability_report( self, title: str, severity: str, @@ -799,40 +819,66 @@ class Tracer: ) def get_total_llm_stats(self) -> dict[str, Any]: - from strix.tools.agents_graph.agents_graph_actions import ( - _agent_instances, - _completed_agent_llm_totals, - _agent_llm_stats_lock, - ) + """Aggregate LLM stats across the live + completed agents. - with _agent_llm_stats_lock: - completed_totals = dict(_completed_agent_llm_totals) - active_agents = list(_agent_instances.values()) + 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_totals.get("input_tokens", 0) or 0), - "output_tokens": int(completed_totals.get("output_tokens", 0) or 0), - "cached_tokens": int(completed_totals.get("cached_tokens", 0) or 0), - "cost": float(completed_totals.get("cost", 0.0) or 0.0), - "requests": int(completed_totals.get("requests", 0) or 0), + "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)), } - - for agent_instance in active_agents: - if hasattr(agent_instance, "llm") and hasattr(agent_instance.llm, "_total_stats"): - agent_stats = agent_instance.llm._total_stats - total_stats["input_tokens"] += agent_stats.input_tokens - total_stats["output_tokens"] += agent_stats.output_tokens - total_stats["cached_tokens"] += agent_stats.cached_tokens - total_stats["cost"] += agent_stats.cost - total_stats["requests"] += agent_stats.requests - - total_stats["cost"] = round(total_stats["cost"], 4) - return { "total": total_stats, "total_tokens": total_stats["input_tokens"] + total_stats["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 + def update_streaming_content(self, agent_id: str, content: str) -> None: self.streaming_content[agent_id] = content diff --git a/strix/tools/__init__.py b/strix/tools/__init__.py index 17299d4..87ad8d4 100644 --- a/strix/tools/__init__.py +++ b/strix/tools/__init__.py @@ -1,14 +1,18 @@ +"""Tool package. + +The package init wires the in-container side: importing every tool +sub-package triggers the ``@register_tool`` decorations that populate +``strix.tools.registry.tools``, which the in-container FastAPI tool +server (:mod:`strix.runtime.tool_server`) dispatches against. + +Host-side SDK function tools live in ``/tool.py`` (or +``tools.py``) and are imported directly by +:mod:`strix.agents.factory` — they do not flow through this package +init's ``register_tool`` registry. +""" + from .agents_graph import * # noqa: F403 from .browser import * # noqa: F403 -from .executor import ( - execute_tool, - execute_tool_invocation, - execute_tool_with_validation, - extract_screenshot_from_result, - process_tool_invocations, - remove_screenshot_from_result, - validate_tool_availability, -) from .file_edit import * # noqa: F403 from .finish import * # noqa: F403 from .load_skill import * # noqa: F403 @@ -33,17 +37,10 @@ from .web_search import * # noqa: F403 __all__ = [ "ImplementedInClientSideOnlyError", - "execute_tool", - "execute_tool_invocation", - "execute_tool_with_validation", - "extract_screenshot_from_result", "get_tool_by_name", "get_tool_names", "get_tools_prompt", "needs_agent_state", - "process_tool_invocations", "register_tool", - "remove_screenshot_from_result", "tools", - "validate_tool_availability", ] diff --git a/strix/tools/_legacy_adapter.py b/strix/tools/_legacy_adapter.py deleted file mode 100644 index 2218844..0000000 --- a/strix/tools/_legacy_adapter.py +++ /dev/null @@ -1,57 +0,0 @@ -"""Shim that lets SDK function tools call legacy ``agent_state``-style functions. - -The legacy harness's tools (notes, todos, reporting, …) take an -``agent_state`` argument with shape ``state.agent_id`` for per-agent silo -keying. Under the SDK migration the equivalent identity lives in -``RunContextWrapper.context["agent_id"]``. - -Rather than rewrite every tool body, SDK function-tool wrappers build a -tiny adapter from the context dict and pass it to the legacy function. -The legacy code path remains untouched (the legacy executor still calls -its tools with the real ``AgentState``). - -Used by: - - ``tools/todo/todo_sdk_tools.py`` - - ``tools/notes/notes_sdk_tools.py`` - - ``tools/reporting/reporting_sdk_tools.py`` - - any other local tool that closes over ``agent_state.agent_id`` -""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import TYPE_CHECKING - - -if TYPE_CHECKING: - from agents import RunContextWrapper - - -@dataclass -class LegacyAgentStateAdapter: - """Just enough surface for legacy tools that read ``state.agent_id``. - - Don't rely on this for new code — it's only here to avoid touching - the legacy ``*_actions.py`` modules during the migration. New SDK - tools should read ``ctx.context["agent_id"]`` directly. - """ - - agent_id: str - - -def adapter_from_ctx( - ctx: RunContextWrapper, - default_agent_id: str = "sdk-default", -) -> LegacyAgentStateAdapter: - """Build a ``LegacyAgentStateAdapter`` from an SDK run context. - - Falls back to ``default_agent_id`` when context is missing or its - ``agent_id`` is unset — keeps tests and CLI dry-runs working without - a fully-populated context. - """ - inner = getattr(ctx, "context", None) - if isinstance(inner, dict): - agent_id = inner.get("agent_id") or default_agent_id - else: - agent_id = default_agent_id - return LegacyAgentStateAdapter(agent_id=str(agent_id)) diff --git a/strix/tools/_state_adapter.py b/strix/tools/_state_adapter.py new file mode 100644 index 0000000..b126c89 --- /dev/null +++ b/strix/tools/_state_adapter.py @@ -0,0 +1,47 @@ +"""Adapter exposing ``ctx.context['agent_id']`` as ``state.agent_id``. + +Several tool implementations still take an ``agent_state`` argument +that they read ``.agent_id`` off of for per-agent silo keying. The SDK +keeps that same identity in ``ctx.context['agent_id']``. Rather than +plumb a different parameter through every tool body, we build a tiny +adapter object from the run context. + +Used by: + - ``tools/todo/tools.py`` + - ``tools/load_skill/tool.py`` + - ``tools/finish/tool.py`` +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + + +if TYPE_CHECKING: + from agents import RunContextWrapper + + +@dataclass +class AgentStateAdapter: + """Just enough surface for tools that read ``state.agent_id``.""" + + agent_id: str + + +def adapter_from_ctx( + ctx: RunContextWrapper, + default_agent_id: str = "default", +) -> AgentStateAdapter: + """Build an ``AgentStateAdapter`` from an SDK run context. + + Falls back to ``default_agent_id`` when context is missing or its + ``agent_id`` is unset — keeps tests and CLI dry-runs working without + a fully-populated context. + """ + inner = getattr(ctx, "context", None) + if isinstance(inner, dict): + agent_id = inner.get("agent_id") or default_agent_id + else: + agent_id = default_agent_id + return AgentStateAdapter(agent_id=str(agent_id)) diff --git a/strix/tools/agents_graph/__init__.py b/strix/tools/agents_graph/__init__.py index d4cd095..07f5f08 100644 --- a/strix/tools/agents_graph/__init__.py +++ b/strix/tools/agents_graph/__init__.py @@ -1,5 +1,6 @@ -from .agents_graph_actions import ( +from .tools import ( agent_finish, + agent_status, create_agent, send_message_to_agent, view_agent_graph, @@ -9,6 +10,7 @@ from .agents_graph_actions import ( __all__ = [ "agent_finish", + "agent_status", "create_agent", "send_message_to_agent", "view_agent_graph", diff --git a/strix/tools/agents_graph/agents_graph_actions.py b/strix/tools/agents_graph/agents_graph_actions.py deleted file mode 100644 index 468687f..0000000 --- a/strix/tools/agents_graph/agents_graph_actions.py +++ /dev/null @@ -1,839 +0,0 @@ -import threading -from datetime import UTC, datetime -import re -from typing import Any, Literal - -from strix.tools.registry import register_tool - - -_agent_graph: dict[str, Any] = { - "nodes": {}, - "edges": [], -} - -_root_agent_id: str | None = None - -_agent_messages: dict[str, list[dict[str, Any]]] = {} - -_running_agents: dict[str, threading.Thread] = {} - -_agent_instances: dict[str, Any] = {} - -_agent_llm_stats_lock = threading.Lock() - - -def _empty_llm_stats_totals() -> dict[str, int | float]: - return { - "input_tokens": 0, - "output_tokens": 0, - "cached_tokens": 0, - "cost": 0.0, - "requests": 0, - } - - -_completed_agent_llm_totals: dict[str, int | float] = _empty_llm_stats_totals() - -_agent_states: dict[str, Any] = {} - - -def _snapshot_agent_llm_stats(agent: Any) -> dict[str, int | float] | None: - if not hasattr(agent, "llm") or not hasattr(agent.llm, "_total_stats"): - return None - - stats = agent.llm._total_stats - return { - "input_tokens": stats.input_tokens, - "output_tokens": stats.output_tokens, - "cached_tokens": stats.cached_tokens, - "cost": stats.cost, - "requests": stats.requests, - } - - -def _finalize_agent_llm_stats(agent_id: str, agent: Any) -> None: - stats = _snapshot_agent_llm_stats(agent) - with _agent_llm_stats_lock: - if stats is not None: - _completed_agent_llm_totals["input_tokens"] += int(stats["input_tokens"]) - _completed_agent_llm_totals["output_tokens"] += int(stats["output_tokens"]) - _completed_agent_llm_totals["cached_tokens"] += int(stats["cached_tokens"]) - _completed_agent_llm_totals["cost"] += float(stats["cost"]) - _completed_agent_llm_totals["requests"] += int(stats["requests"]) - - node = _agent_graph["nodes"].get(agent_id) - if node is not None: - node["llm_stats"] = stats - - _agent_instances.pop(agent_id, None) - - -def _is_whitebox_agent(agent_id: str) -> bool: - agent = _agent_instances.get(agent_id) - return bool(getattr(getattr(agent, "llm_config", None), "is_whitebox", False)) - - -def _extract_repo_tags(agent_state: Any | None) -> set[str]: - repo_tags: set[str] = set() - if agent_state is None: - return repo_tags - - task_text = str(getattr(agent_state, "task", "") or "") - for workspace_subdir in re.findall(r"/workspace/([A-Za-z0-9._-]+)", task_text): - repo_tags.add(f"repo:{workspace_subdir.lower()}") - - for repo_name in re.findall(r"github\.com/[^/\s]+/([A-Za-z0-9._-]+)", task_text): - normalized = repo_name.removesuffix(".git").lower() - if normalized: - repo_tags.add(f"repo:{normalized}") - - return repo_tags - - -def _load_primary_wiki_note(agent_state: Any | None = None) -> dict[str, Any] | None: - try: - from strix.tools.notes.notes_actions import get_note, list_notes - - notes_result = list_notes(category="wiki") - if not notes_result.get("success"): - return None - - notes = notes_result.get("notes") or [] - if not notes: - return None - - selected_note_id = None - repo_tags = _extract_repo_tags(agent_state) - if repo_tags: - for note in notes: - note_tags = note.get("tags") or [] - if not isinstance(note_tags, list): - continue - normalized_note_tags = {str(tag).strip().lower() for tag in note_tags if str(tag).strip()} - if normalized_note_tags.intersection(repo_tags): - selected_note_id = note.get("note_id") - break - - note_id = selected_note_id or notes[0].get("note_id") - if not isinstance(note_id, str) or not note_id: - return None - - note_result = get_note(note_id=note_id) - if not note_result.get("success"): - return None - - note = note_result.get("note") - if not isinstance(note, dict): - return None - - except Exception: - return None - else: - return note - - -def _inject_wiki_context_for_whitebox(agent_state: Any) -> None: - if not _is_whitebox_agent(agent_state.agent_id): - return - - wiki_note = _load_primary_wiki_note(agent_state) - if not wiki_note: - return - - title = str(wiki_note.get("title") or "repo wiki") - content = str(wiki_note.get("content") or "").strip() - if not content: - return - - max_chars = 4000 - truncated_content = content[:max_chars] - suffix = "\n\n[truncated for context size]" if len(content) > max_chars else "" - agent_state.add_message( - "user", - ( - f"\n" - f"{truncated_content}{suffix}\n" - "" - ), - ) - - -def _append_wiki_update_on_finish( - agent_state: Any, - agent_name: str, - result_summary: str, - findings: list[str] | None, - final_recommendations: list[str] | None, -) -> None: - if not _is_whitebox_agent(agent_state.agent_id): - return - - try: - from strix.tools.notes.notes_actions import append_note_content - - note = _load_primary_wiki_note(agent_state) - if not note: - return - - note_id = note.get("note_id") - if not isinstance(note_id, str) or not note_id: - return - - timestamp = datetime.now(UTC).isoformat() - summary = " ".join(str(result_summary).split()) - if len(summary) > 1200: - summary = f"{summary[:1197]}..." - findings_lines = "\n".join(f"- {item}" for item in (findings or [])) or "- none" - recommendation_lines = ( - "\n".join(f"- {item}" for item in (final_recommendations or [])) or "- none" - ) - - delta = ( - f"\n\n## Agent Update: {agent_name} ({timestamp})\n" - f"Summary: {summary}\n\n" - "Findings:\n" - f"{findings_lines}\n\n" - "Recommendations:\n" - f"{recommendation_lines}\n" - ) - append_note_content(note_id=note_id, delta=delta) - except Exception: - # Best-effort update; never block agent completion on note persistence. - return - - -def _run_agent_in_thread( - agent: Any, state: Any, inherited_messages: list[dict[str, Any]] -) -> dict[str, Any]: - try: - if inherited_messages: - state.add_message("user", "") - for msg in inherited_messages: - state.add_message(msg["role"], msg["content"]) - state.add_message("user", "") - - _inject_wiki_context_for_whitebox(state) - - parent_info = _agent_graph["nodes"].get(state.parent_id, {}) - parent_name = parent_info.get("name", "Unknown Parent") - - context_status = ( - "inherited conversation context from your parent for background understanding" - if inherited_messages - else "started with a fresh context" - ) - wiki_memory_instruction = "" - if getattr(getattr(agent, "llm_config", None), "is_whitebox", False): - wiki_memory_instruction = ( - '\n - White-box memory (recommended): call list_notes(category="wiki") and then ' - "get_note(note_id=...) before substantive work (including terminal scans)" - "\n - Reuse one repo wiki note where possible and avoid duplicates" - "\n - Before agent_finish, call list_notes(category=\"wiki\") + get_note(note_id=...) again, then append a short scope delta via update_note (new routes/sinks, scanner results, dynamic follow-ups)" - "\n - If terminal output contains `command not found` or shell parse errors, correct and rerun before using the result" - "\n - Use ASCII-only shell commands; if a command includes unexpected non-ASCII characters, rerun with a clean ASCII command" - "\n - Keep AST artifacts bounded: target relevant paths and avoid whole-repo generic function dumps" - "\n - Source-aware tooling is advisory: choose semgrep/AST/tree-sitter/gitleaks/trivy when relevant, do not force static steps for purely dynamic validation tasks" - ) - - task_xml = f""" - - ⚠️ You are NOT your parent agent. You are a NEW, SEPARATE sub-agent (not root). - - Your Info: {state.agent_name} ({state.agent_id}) - Parent Info: {parent_name} ({state.parent_id}) - - - {state.task} - - - - You have {context_status} - - Inherited context is for BACKGROUND ONLY - don't continue parent's work - - Maintain strict self-identity: never speak as or for your parent - - Do not merge your conversation with the parent's; - - Do not claim parent's actions or messages as your own - - Focus EXCLUSIVELY on your delegated task above - - Work independently with your own approach - - Use agent_finish when complete to report back to parent - - You are a SPECIALIST for this specific task - - You share the same container as other agents but have your own tool server instance - - All agents share /workspace directory and proxy history for better collaboration - - You can see files created by other agents and proxy traffic from previous work - - Build upon previous work but focus on your specific delegated task -{wiki_memory_instruction} - -""" - - state.add_message("user", task_xml) - - _agent_states[state.agent_id] = state - - _agent_graph["nodes"][state.agent_id]["state"] = state.model_dump() - - import asyncio - - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - try: - result = loop.run_until_complete(agent.agent_loop(state.task)) - finally: - loop.close() - - except Exception as e: - _agent_graph["nodes"][state.agent_id]["status"] = "error" - _agent_graph["nodes"][state.agent_id]["finished_at"] = datetime.now(UTC).isoformat() - _agent_graph["nodes"][state.agent_id]["result"] = {"error": str(e)} - _running_agents.pop(state.agent_id, None) - _finalize_agent_llm_stats(state.agent_id, agent) - raise - else: - if state.stop_requested: - _agent_graph["nodes"][state.agent_id]["status"] = "stopped" - else: - _agent_graph["nodes"][state.agent_id]["status"] = "completed" - _agent_graph["nodes"][state.agent_id]["finished_at"] = datetime.now(UTC).isoformat() - _agent_graph["nodes"][state.agent_id]["result"] = result - _running_agents.pop(state.agent_id, None) - _finalize_agent_llm_stats(state.agent_id, agent) - - return {"result": result} - - -@register_tool(sandbox_execution=False) -def view_agent_graph(agent_state: Any) -> dict[str, Any]: - try: - structure_lines = ["=== AGENT GRAPH STRUCTURE ==="] - - def _build_tree(agent_id: str, depth: int = 0) -> None: - node = _agent_graph["nodes"][agent_id] - indent = " " * depth - - you_indicator = " ← This is you" if agent_id == agent_state.agent_id else "" - - structure_lines.append(f"{indent}* {node['name']} ({agent_id}){you_indicator}") - structure_lines.append(f"{indent} Task: {node['task']}") - structure_lines.append(f"{indent} Status: {node['status']}") - - children = [ - edge["to"] - for edge in _agent_graph["edges"] - if edge["from"] == agent_id and edge["type"] == "delegation" - ] - - if children: - structure_lines.append(f"{indent} Children:") - for child_id in children: - _build_tree(child_id, depth + 2) - - root_agent_id = _root_agent_id - if not root_agent_id and _agent_graph["nodes"]: - for agent_id, node in _agent_graph["nodes"].items(): - if node.get("parent_id") is None: - root_agent_id = agent_id - break - if not root_agent_id: - root_agent_id = next(iter(_agent_graph["nodes"].keys())) - - if root_agent_id and root_agent_id in _agent_graph["nodes"]: - _build_tree(root_agent_id) - else: - structure_lines.append("No agents in the graph yet") - - graph_structure = "\n".join(structure_lines) - - total_nodes = len(_agent_graph["nodes"]) - running_count = sum( - 1 for node in _agent_graph["nodes"].values() if node["status"] == "running" - ) - waiting_count = sum( - 1 for node in _agent_graph["nodes"].values() if node["status"] == "waiting" - ) - stopping_count = sum( - 1 for node in _agent_graph["nodes"].values() if node["status"] == "stopping" - ) - completed_count = sum( - 1 for node in _agent_graph["nodes"].values() if node["status"] == "completed" - ) - stopped_count = sum( - 1 for node in _agent_graph["nodes"].values() if node["status"] == "stopped" - ) - failed_count = sum( - 1 for node in _agent_graph["nodes"].values() if node["status"] in ["failed", "error"] - ) - - except Exception as e: # noqa: BLE001 - return { - "error": f"Failed to view agent graph: {e}", - "graph_structure": "Error retrieving graph structure", - } - else: - return { - "graph_structure": graph_structure, - "summary": { - "total_agents": total_nodes, - "running": running_count, - "waiting": waiting_count, - "stopping": stopping_count, - "completed": completed_count, - "stopped": stopped_count, - "failed": failed_count, - }, - } - - -@register_tool(sandbox_execution=False) -def create_agent( - agent_state: Any, - task: str, - name: str, - inherit_context: bool = True, - skills: str | None = None, -) -> dict[str, Any]: - try: - parent_id = agent_state.agent_id - - from strix.skills import parse_skill_list, validate_requested_skills - - skill_list = parse_skill_list(skills) - validation_error = validate_requested_skills(skill_list) - if validation_error: - return { - "success": False, - "error": validation_error, - "agent_id": None, - } - - from strix.agents import StrixAgent - from strix.agents.state import AgentState - from strix.llm.config import LLMConfig - - parent_agent = _agent_instances.get(parent_id) - - timeout = None - scan_mode = "deep" - is_whitebox = False - interactive = False - if parent_agent and hasattr(parent_agent, "llm_config"): - if hasattr(parent_agent.llm_config, "timeout"): - timeout = parent_agent.llm_config.timeout - if hasattr(parent_agent.llm_config, "scan_mode"): - scan_mode = parent_agent.llm_config.scan_mode - if hasattr(parent_agent.llm_config, "is_whitebox"): - is_whitebox = parent_agent.llm_config.is_whitebox - interactive = getattr(parent_agent.llm_config, "interactive", False) - - if is_whitebox: - whitebox_guidance = ( - "\n\nWhite-box execution guidance (recommended when source is available):\n" - "- Use structural AST mapping (`sg` or `tree-sitter`) where it helps source analysis; " - "keep artifacts bounded and skip forced AST steps for purely dynamic validation tasks.\n" - "- Keep AST output bounded: scope to relevant paths/files, avoid whole-repo " - "generic function patterns, and cap artifact size.\n" - '- Use shared wiki memory by calling list_notes(category="wiki") then ' - "get_note(note_id=...).\n" - '- Before agent_finish, call list_notes(category="wiki") + get_note(note_id=...) ' - "again, reuse one repo wiki, and call update_note.\n" - "- If terminal output contains `command not found` or shell parse errors, " - "correct and rerun before using the result." - ) - if "White-box execution guidance (recommended when source is available):" not in task: - task = f"{task.rstrip()}{whitebox_guidance}" - - state = AgentState( - task=task, - agent_name=name, - parent_id=parent_id, - max_iterations=300, - waiting_timeout=300 if interactive else 600, - ) - llm_config = LLMConfig( - skills=skill_list, - timeout=timeout, - scan_mode=scan_mode, - is_whitebox=is_whitebox, - interactive=interactive, - ) - - agent_config = { - "llm_config": llm_config, - "state": state, - } - - agent = StrixAgent(agent_config) - - inherited_messages = [] - if inherit_context: - inherited_messages = agent_state.get_conversation_history() - - with _agent_llm_stats_lock: - _agent_instances[state.agent_id] = agent - - thread = threading.Thread( - target=_run_agent_in_thread, - args=(agent, state, inherited_messages), - daemon=True, - name=f"Agent-{name}-{state.agent_id}", - ) - thread.start() - _running_agents[state.agent_id] = thread - - except Exception as e: # noqa: BLE001 - return {"success": False, "error": f"Failed to create agent: {e}", "agent_id": None} - else: - return { - "success": True, - "agent_id": state.agent_id, - "message": f"Agent '{name}' created and started asynchronously", - "agent_info": { - "id": state.agent_id, - "name": name, - "status": "running", - "parent_id": parent_id, - }, - } - - -@register_tool(sandbox_execution=False) -def send_message_to_agent( - agent_state: Any, - target_agent_id: str, - message: str, - message_type: Literal["query", "instruction", "information"] = "information", - priority: Literal["low", "normal", "high", "urgent"] = "normal", -) -> dict[str, Any]: - try: - if target_agent_id not in _agent_graph["nodes"]: - return { - "success": False, - "error": f"Target agent '{target_agent_id}' not found in graph", - "message_id": None, - } - - sender_id = agent_state.agent_id - - from uuid import uuid4 - - message_id = f"msg_{uuid4().hex[:8]}" - message_data = { - "id": message_id, - "from": sender_id, - "to": target_agent_id, - "content": message, - "message_type": message_type, - "priority": priority, - "timestamp": datetime.now(UTC).isoformat(), - "delivered": False, - "read": False, - } - - if target_agent_id not in _agent_messages: - _agent_messages[target_agent_id] = [] - - _agent_messages[target_agent_id].append(message_data) - - _agent_graph["edges"].append( - { - "from": sender_id, - "to": target_agent_id, - "type": "message", - "message_id": message_id, - "message_type": message_type, - "priority": priority, - "created_at": datetime.now(UTC).isoformat(), - } - ) - - message_data["delivered"] = True - - target_name = _agent_graph["nodes"][target_agent_id]["name"] - sender_name = _agent_graph["nodes"][sender_id]["name"] - - return { - "success": True, - "message_id": message_id, - "message": f"Message sent from '{sender_name}' to '{target_name}'", - "delivery_status": "delivered", - "target_agent": { - "id": target_agent_id, - "name": target_name, - "status": _agent_graph["nodes"][target_agent_id]["status"], - }, - } - - except Exception as e: # noqa: BLE001 - return {"success": False, "error": f"Failed to send message: {e}", "message_id": None} - - -@register_tool(sandbox_execution=False) -def agent_finish( - agent_state: Any, - result_summary: str, - findings: list[str] | None = None, - success: bool = True, - report_to_parent: bool = True, - final_recommendations: list[str] | None = None, -) -> dict[str, Any]: - try: - if not hasattr(agent_state, "parent_id") or agent_state.parent_id is None: - return { - "agent_completed": False, - "error": ( - "This tool can only be used by subagents. " - "Root/main agents must use finish_scan instead." - ), - "parent_notified": False, - } - - agent_id = agent_state.agent_id - - if agent_id not in _agent_graph["nodes"]: - return {"agent_completed": False, "error": "Current agent not found in graph"} - - agent_node = _agent_graph["nodes"][agent_id] - - agent_node["status"] = "finished" if success else "failed" - agent_node["finished_at"] = datetime.now(UTC).isoformat() - agent_node["result"] = { - "summary": result_summary, - "findings": findings or [], - "success": success, - "recommendations": final_recommendations or [], - } - - _append_wiki_update_on_finish( - agent_state=agent_state, - agent_name=agent_node["name"], - result_summary=result_summary, - findings=findings, - final_recommendations=final_recommendations, - ) - - parent_notified = False - - if report_to_parent and agent_node["parent_id"]: - parent_id = agent_node["parent_id"] - - if parent_id in _agent_graph["nodes"]: - findings_xml = "\n".join( - f" {finding}" for finding in (findings or []) - ) - recommendations_xml = "\n".join( - f" {rec}" - for rec in (final_recommendations or []) - ) - - report_message = f""" - - {agent_node["name"]} - {agent_id} - {agent_node["task"]} - {"SUCCESS" if success else "FAILED"} - {agent_node["finished_at"]} - - - {result_summary} - -{findings_xml} - - -{recommendations_xml} - - -""" - - if parent_id not in _agent_messages: - _agent_messages[parent_id] = [] - - from uuid import uuid4 - - _agent_messages[parent_id].append( - { - "id": f"report_{uuid4().hex[:8]}", - "from": agent_id, - "to": parent_id, - "content": report_message, - "message_type": "information", - "priority": "high", - "timestamp": datetime.now(UTC).isoformat(), - "delivered": True, - "read": False, - } - ) - - parent_notified = True - - _running_agents.pop(agent_id, None) - - return { - "agent_completed": True, - "parent_notified": parent_notified, - "completion_summary": { - "agent_id": agent_id, - "agent_name": agent_node["name"], - "task": agent_node["task"], - "success": success, - "findings_count": len(findings or []), - "has_recommendations": bool(final_recommendations), - "finished_at": agent_node["finished_at"], - }, - } - - except Exception as e: # noqa: BLE001 - return { - "agent_completed": False, - "error": f"Failed to complete agent: {e}", - "parent_notified": False, - } - - -def stop_agent(agent_id: str) -> dict[str, Any]: - try: - if agent_id not in _agent_graph["nodes"]: - return { - "success": False, - "error": f"Agent '{agent_id}' not found in graph", - "agent_id": agent_id, - } - - agent_node = _agent_graph["nodes"][agent_id] - - if agent_node["status"] in ["completed", "error", "failed", "stopped"]: - return { - "success": True, - "message": f"Agent '{agent_node['name']}' was already stopped", - "agent_id": agent_id, - "previous_status": agent_node["status"], - } - - if agent_id in _agent_states: - agent_state = _agent_states[agent_id] - agent_state.request_stop() - - if agent_id in _agent_instances: - agent_instance = _agent_instances[agent_id] - if hasattr(agent_instance, "state"): - agent_instance.state.request_stop() - if hasattr(agent_instance, "cancel_current_execution"): - agent_instance.cancel_current_execution() - - agent_node["status"] = "stopping" - - try: - from strix.telemetry.tracer import get_global_tracer - - tracer = get_global_tracer() - if tracer: - tracer.update_agent_status(agent_id, "stopping") - except (ImportError, AttributeError): - pass - - agent_node["result"] = { - "summary": "Agent stop requested by user", - "success": False, - "stopped_by_user": True, - } - - return { - "success": True, - "message": f"Stop request sent to agent '{agent_node['name']}'", - "agent_id": agent_id, - "agent_name": agent_node["name"], - "note": "Agent will stop gracefully after current iteration", - } - - except Exception as e: # noqa: BLE001 - return { - "success": False, - "error": f"Failed to stop agent: {e}", - "agent_id": agent_id, - } - - -def send_user_message_to_agent(agent_id: str, message: str) -> dict[str, Any]: - try: - if agent_id not in _agent_graph["nodes"]: - return { - "success": False, - "error": f"Agent '{agent_id}' not found in graph", - "agent_id": agent_id, - } - - agent_node = _agent_graph["nodes"][agent_id] - - if agent_id not in _agent_messages: - _agent_messages[agent_id] = [] - - from uuid import uuid4 - - message_data = { - "id": f"user_msg_{uuid4().hex[:8]}", - "from": "user", - "to": agent_id, - "content": message, - "message_type": "instruction", - "priority": "high", - "timestamp": datetime.now(UTC).isoformat(), - "delivered": True, - "read": False, - } - - _agent_messages[agent_id].append(message_data) - - return { - "success": True, - "message": f"Message sent to agent '{agent_node['name']}'", - "agent_id": agent_id, - "agent_name": agent_node["name"], - } - - except Exception as e: # noqa: BLE001 - return { - "success": False, - "error": f"Failed to send message to agent: {e}", - "agent_id": agent_id, - } - - -@register_tool(sandbox_execution=False) -def wait_for_message( - agent_state: Any, - reason: str = "Waiting for messages from other agents", -) -> dict[str, Any]: - try: - agent_id = agent_state.agent_id - agent_name = agent_state.agent_name - - agent_state.enter_waiting_state() - - if agent_id in _agent_graph["nodes"]: - _agent_graph["nodes"][agent_id]["status"] = "waiting" - _agent_graph["nodes"][agent_id]["waiting_reason"] = reason - - try: - from strix.telemetry.tracer import get_global_tracer - - tracer = get_global_tracer() - if tracer: - tracer.update_agent_status(agent_id, "waiting") - except (ImportError, AttributeError): - pass - - except Exception as e: # noqa: BLE001 - return {"success": False, "error": f"Failed to enter waiting state: {e}", "status": "error"} - else: - return { - "success": True, - "status": "waiting", - "message": f"Agent '{agent_name}' is now waiting for messages", - "reason": reason, - "agent_info": { - "id": agent_id, - "name": agent_name, - "status": "waiting", - }, - "resume_conditions": [ - "Message from another agent", - "Message from user", - "Direct communication", - "Waiting timeout reached", - ], - } diff --git a/strix/tools/agents_graph/agents_graph_sdk_tools.py b/strix/tools/agents_graph/tools.py similarity index 100% rename from strix/tools/agents_graph/agents_graph_sdk_tools.py rename to strix/tools/agents_graph/tools.py diff --git a/strix/tools/browser/browser_sdk_tool.py b/strix/tools/browser/tool.py similarity index 100% rename from strix/tools/browser/browser_sdk_tool.py rename to strix/tools/browser/tool.py diff --git a/strix/tools/executor.py b/strix/tools/executor.py deleted file mode 100644 index 1c24087..0000000 --- a/strix/tools/executor.py +++ /dev/null @@ -1,364 +0,0 @@ -import inspect -import os -from typing import Any - -import httpx - -from strix.config import Config -from strix.telemetry import posthog - - -if os.getenv("STRIX_SANDBOX_MODE", "false").lower() == "false": - from strix.runtime import get_runtime - -from .argument_parser import convert_arguments -from .registry import ( - get_tool_by_name, - get_tool_names, - get_tool_param_schema, - needs_agent_state, - should_execute_in_sandbox, -) - - -_SERVER_TIMEOUT = float(Config.get("strix_sandbox_execution_timeout") or "120") -SANDBOX_EXECUTION_TIMEOUT = _SERVER_TIMEOUT + 30 -SANDBOX_CONNECT_TIMEOUT = float(Config.get("strix_sandbox_connect_timeout") or "10") - - -async def execute_tool(tool_name: str, agent_state: Any | None = None, **kwargs: Any) -> Any: - execute_in_sandbox = should_execute_in_sandbox(tool_name) - sandbox_mode = os.getenv("STRIX_SANDBOX_MODE", "false").lower() == "true" - - if execute_in_sandbox and not sandbox_mode: - return await _execute_tool_in_sandbox(tool_name, agent_state, **kwargs) - - return await _execute_tool_locally(tool_name, agent_state, **kwargs) - - -async def _execute_tool_in_sandbox(tool_name: str, agent_state: Any, **kwargs: Any) -> Any: - if not hasattr(agent_state, "sandbox_id") or not agent_state.sandbox_id: - raise ValueError("Agent state with a valid sandbox_id is required for sandbox execution.") - - if not hasattr(agent_state, "sandbox_token") or not agent_state.sandbox_token: - raise ValueError( - "Agent state with a valid sandbox_token is required for sandbox execution." - ) - - if ( - not hasattr(agent_state, "sandbox_info") - or "tool_server_port" not in agent_state.sandbox_info - ): - raise ValueError( - "Agent state with a valid sandbox_info containing tool_server_port is required." - ) - - runtime = get_runtime() - tool_server_port = agent_state.sandbox_info["tool_server_port"] - server_url = await runtime.get_sandbox_url(agent_state.sandbox_id, tool_server_port) - request_url = f"{server_url}/execute" - - agent_id = getattr(agent_state, "agent_id", "unknown") - - request_data = { - "agent_id": agent_id, - "tool_name": tool_name, - "kwargs": kwargs, - } - - headers = { - "Authorization": f"Bearer {agent_state.sandbox_token}", - "Content-Type": "application/json", - } - - timeout = httpx.Timeout( - timeout=SANDBOX_EXECUTION_TIMEOUT, - connect=SANDBOX_CONNECT_TIMEOUT, - ) - - async with httpx.AsyncClient(trust_env=False) as client: - try: - response = await client.post( - request_url, json=request_data, headers=headers, timeout=timeout - ) - response.raise_for_status() - response_data = response.json() - if response_data.get("error"): - posthog.error("tool_execution_error", f"{tool_name}: {response_data['error']}") - raise RuntimeError(f"Sandbox execution error: {response_data['error']}") - return response_data.get("result") - except httpx.HTTPStatusError as e: - posthog.error("tool_http_error", f"{tool_name}: HTTP {e.response.status_code}") - if e.response.status_code == 401: - raise RuntimeError("Authentication failed: Invalid or missing sandbox token") from e - raise RuntimeError(f"HTTP error calling tool server: {e.response.status_code}") from e - except httpx.RequestError as e: - error_type = type(e).__name__ - posthog.error("tool_request_error", f"{tool_name}: {error_type}") - raise RuntimeError(f"Request error calling tool server: {error_type}") from e - - -async def _execute_tool_locally(tool_name: str, agent_state: Any | None, **kwargs: Any) -> Any: - tool_func = get_tool_by_name(tool_name) - if not tool_func: - raise ValueError(f"Tool '{tool_name}' not found") - - converted_kwargs = convert_arguments(tool_func, kwargs) - - if needs_agent_state(tool_name): - if agent_state is None: - raise ValueError(f"Tool '{tool_name}' requires agent_state but none was provided.") - result = tool_func(agent_state=agent_state, **converted_kwargs) - else: - result = tool_func(**converted_kwargs) - - return await result if inspect.isawaitable(result) else result - - -def validate_tool_availability(tool_name: str | None) -> tuple[bool, str]: - if tool_name is None: - available = ", ".join(sorted(get_tool_names())) - return False, f"Tool name is missing. Available tools: {available}" - - if tool_name not in get_tool_names(): - available = ", ".join(sorted(get_tool_names())) - return False, f"Tool '{tool_name}' is not available. Available tools: {available}" - - return True, "" - - -def _validate_tool_arguments(tool_name: str, kwargs: dict[str, Any]) -> str | None: - param_schema = get_tool_param_schema(tool_name) - if not param_schema or not param_schema.get("has_params"): - return None - - allowed_params: set[str] = param_schema.get("params", set()) - required_params: set[str] = param_schema.get("required", set()) - optional_params = allowed_params - required_params - - schema_hint = _format_schema_hint(tool_name, required_params, optional_params) - - unknown_params = set(kwargs.keys()) - allowed_params - if unknown_params: - unknown_list = ", ".join(sorted(unknown_params)) - return f"Tool '{tool_name}' received unknown parameter(s): {unknown_list}\n{schema_hint}" - - missing_required = [ - param for param in required_params if param not in kwargs or kwargs.get(param) in (None, "") - ] - if missing_required: - missing_list = ", ".join(sorted(missing_required)) - return f"Tool '{tool_name}' missing required parameter(s): {missing_list}\n{schema_hint}" - - return None - - -def _format_schema_hint(tool_name: str, required: set[str], optional: set[str]) -> str: - parts = [f"Valid parameters for '{tool_name}':"] - if required: - parts.append(f" Required: {', '.join(sorted(required))}") - if optional: - parts.append(f" Optional: {', '.join(sorted(optional))}") - return "\n".join(parts) - - -async def execute_tool_with_validation( - tool_name: str | None, agent_state: Any | None = None, **kwargs: Any -) -> Any: - is_valid, error_msg = validate_tool_availability(tool_name) - if not is_valid: - return f"Error: {error_msg}" - - assert tool_name is not None - - arg_error = _validate_tool_arguments(tool_name, kwargs) - if arg_error: - return f"Error: {arg_error}" - - try: - result = await execute_tool(tool_name, agent_state, **kwargs) - except Exception as e: # noqa: BLE001 - error_str = str(e) - if len(error_str) > 500: - error_str = error_str[:500] + "... [truncated]" - return f"Error executing {tool_name}: {error_str}" - else: - return result - - -async def execute_tool_invocation(tool_inv: dict[str, Any], agent_state: Any | None = None) -> Any: - tool_name = tool_inv.get("toolName") - tool_args = tool_inv.get("args", {}) - - return await execute_tool_with_validation(tool_name, agent_state, **tool_args) - - -def _check_error_result(result: Any) -> tuple[bool, Any]: - is_error = False - error_payload: Any = None - - if (isinstance(result, dict) and "error" in result) or ( - isinstance(result, str) and result.strip().lower().startswith("error:") - ): - is_error = True - error_payload = result - - return is_error, error_payload - - -def _update_tracer_with_result( - tracer: Any, execution_id: Any, is_error: bool, result: Any, error_payload: Any -) -> None: - if not tracer or not execution_id: - return - - try: - if is_error: - tracer.update_tool_execution(execution_id, "error", error_payload) - else: - tracer.update_tool_execution(execution_id, "completed", result) - except (ConnectionError, RuntimeError) as e: - error_msg = str(e) - if tracer and execution_id: - tracer.update_tool_execution(execution_id, "error", error_msg) - raise - - -def _format_tool_result(tool_name: str, result: Any) -> tuple[str, list[dict[str, Any]]]: - images: list[dict[str, Any]] = [] - - screenshot_data = extract_screenshot_from_result(result) - if screenshot_data: - images.append( - { - "type": "image_url", - "image_url": {"url": f"data:image/png;base64,{screenshot_data}"}, - } - ) - result_str = remove_screenshot_from_result(result) - else: - result_str = result - - if result_str is None: - final_result_str = f"Tool {tool_name} executed successfully" - else: - final_result_str = str(result_str) - if len(final_result_str) > 10000: - start_part = final_result_str[:4000] - end_part = final_result_str[-4000:] - final_result_str = start_part + "\n\n... [middle content truncated] ...\n\n" + end_part - - observation_xml = ( - f"\n{tool_name}\n" - f"{final_result_str}\n" - ) - - return observation_xml, images - - -async def _execute_single_tool( - tool_inv: dict[str, Any], - agent_state: Any | None, - tracer: Any | None, - agent_id: str, -) -> tuple[str, list[dict[str, Any]], bool]: - tool_name = tool_inv.get("toolName", "unknown") - args = tool_inv.get("args", {}) - execution_id = None - should_agent_finish = False - - if tracer: - execution_id = tracer.log_tool_execution_start(agent_id, tool_name, args) - - try: - result = await execute_tool_invocation(tool_inv, agent_state) - - is_error, error_payload = _check_error_result(result) - - if ( - tool_name in ("finish_scan", "agent_finish") - and not is_error - and isinstance(result, dict) - ): - if tool_name == "finish_scan": - should_agent_finish = result.get("scan_completed", False) - elif tool_name == "agent_finish": - should_agent_finish = result.get("agent_completed", False) - - _update_tracer_with_result(tracer, execution_id, is_error, result, error_payload) - - except (ConnectionError, RuntimeError, ValueError, TypeError, OSError) as e: - error_msg = str(e) - if tracer and execution_id: - tracer.update_tool_execution(execution_id, "error", error_msg) - raise - - observation_xml, images = _format_tool_result(tool_name, result) - return observation_xml, images, should_agent_finish - - -def _get_tracer_and_agent_id(agent_state: Any | None) -> tuple[Any | None, str]: - try: - from strix.telemetry.tracer import get_global_tracer - - tracer = get_global_tracer() - agent_id = agent_state.agent_id if agent_state else "unknown_agent" - except (ImportError, AttributeError): - tracer = None - agent_id = "unknown_agent" - - return tracer, agent_id - - -async def process_tool_invocations( - tool_invocations: list[dict[str, Any]], - conversation_history: list[dict[str, Any]], - agent_state: Any | None = None, -) -> bool: - observation_parts: list[str] = [] - all_images: list[dict[str, Any]] = [] - should_agent_finish = False - - tracer, agent_id = _get_tracer_and_agent_id(agent_state) - - for tool_inv in tool_invocations: - observation_xml, images, tool_should_finish = await _execute_single_tool( - tool_inv, agent_state, tracer, agent_id - ) - observation_parts.append(observation_xml) - all_images.extend(images) - - if tool_should_finish: - should_agent_finish = True - - if all_images: - content = [{"type": "text", "text": "Tool Results:\n\n" + "\n\n".join(observation_parts)}] - content.extend(all_images) - conversation_history.append({"role": "user", "content": content}) - else: - observation_content = "Tool Results:\n\n" + "\n\n".join(observation_parts) - conversation_history.append({"role": "user", "content": observation_content}) - - return should_agent_finish - - -def extract_screenshot_from_result(result: Any) -> str | None: - if not isinstance(result, dict): - return None - - screenshot = result.get("screenshot") - if isinstance(screenshot, str) and screenshot: - return screenshot - - return None - - -def remove_screenshot_from_result(result: Any) -> Any: - if not isinstance(result, dict): - return result - - result_copy = result.copy() - if "screenshot" in result_copy: - result_copy["screenshot"] = "[Image data extracted - see attached image]" - - return result_copy diff --git a/strix/tools/file_edit/file_edit_sdk_tools.py b/strix/tools/file_edit/tools.py similarity index 100% rename from strix/tools/file_edit/file_edit_sdk_tools.py rename to strix/tools/file_edit/tools.py diff --git a/strix/tools/finish/finish_actions.py b/strix/tools/finish/finish_actions.py index 79f48e7..af1035f 100644 --- a/strix/tools/finish/finish_actions.py +++ b/strix/tools/finish/finish_actions.py @@ -14,72 +14,15 @@ def _validate_root_agent(agent_state: Any) -> dict[str, Any] | None: return None -def _check_active_agents(agent_state: Any = None) -> dict[str, Any] | None: - try: - from strix.tools.agents_graph.agents_graph_actions import _agent_graph - - if agent_state and agent_state.agent_id: - current_agent_id = agent_state.agent_id - else: - return None - - active_agents = [] - stopping_agents = [] - - for agent_id, node in _agent_graph["nodes"].items(): - if agent_id == current_agent_id: - continue - - status = node.get("status", "unknown") - if status == "running": - active_agents.append( - { - "id": agent_id, - "name": node.get("name", "Unknown"), - "task": node.get("task", "Unknown task")[:300], - "status": status, - } - ) - elif status == "stopping": - stopping_agents.append( - { - "id": agent_id, - "name": node.get("name", "Unknown"), - "task": node.get("task", "Unknown task")[:300], - "status": status, - } - ) - - if active_agents or stopping_agents: - response: dict[str, Any] = { - "success": False, - "error": "agents_still_active", - "message": "Cannot finish scan: agents are still active", - } - - if active_agents: - response["active_agents"] = active_agents - - if stopping_agents: - response["stopping_agents"] = stopping_agents - - response["suggestions"] = [ - "Use wait_for_message to wait for all agents to complete", - "Use send_message_to_agent if you need agents to complete immediately", - "Check agent_status to see current agent states", - ] - - response["total_active"] = len(active_agents) + len(stopping_agents) - - return response - - except ImportError: - pass - except Exception: - import logging - - logging.exception("Error checking active agents") +def _check_active_agents(_agent_state: Any = None) -> dict[str, Any] | None: + """Check whether sibling agents are still running before finishing. + The active-agent check now lives in the orchestration bus + (:class:`strix.orchestration.bus.AgentMessageBus`); ``finish_scan`` + sees an empty world here and the bus's per-agent state is the + source of truth. Returns ``None`` (no blockers) so the caller's + field validation can run. + """ return None diff --git a/strix/tools/finish/finish_sdk_tool.py b/strix/tools/finish/tool.py similarity index 93% rename from strix/tools/finish/finish_sdk_tool.py rename to strix/tools/finish/tool.py index 1fd3b20..6158e94 100644 --- a/strix/tools/finish/finish_sdk_tool.py +++ b/strix/tools/finish/tool.py @@ -27,8 +27,8 @@ from typing import Any from agents import RunContextWrapper from strix.tools._decorator import strix_tool -from strix.tools._legacy_adapter import adapter_from_ctx -from strix.tools.finish import finish_actions as _legacy +from strix.tools._state_adapter import adapter_from_ctx +from strix.tools.finish import finish_actions as _impl def _dump(result: dict[str, Any]) -> str: @@ -57,7 +57,7 @@ async def finish_scan( state = adapter_from_ctx(ctx) return _dump( await asyncio.to_thread( - _legacy.finish_scan, + _impl.finish_scan, executive_summary=executive_summary, methodology=methodology, technical_analysis=technical_analysis, diff --git a/strix/tools/load_skill/load_skill_actions.py b/strix/tools/load_skill/load_skill_actions.py index 42f64f1..54b9717 100644 --- a/strix/tools/load_skill/load_skill_actions.py +++ b/strix/tools/load_skill/load_skill_actions.py @@ -25,28 +25,15 @@ def load_skill(agent_state: Any, skills: str) -> dict[str, Any]: "loaded_skills": [], } - from strix.tools.agents_graph.agents_graph_actions import _agent_instances - - current_agent = _agent_instances.get(agent_state.agent_id) - if current_agent is None or not hasattr(current_agent, "llm"): - return { - "success": False, - "error": ( - "Could not find running agent instance for runtime skill loading. " - "Try again in the current active agent." - ), - "requested_skills": requested_skills, - "loaded_skills": [], - } - - newly_loaded = current_agent.llm.add_skills(requested_skills) - already_loaded = [skill for skill in requested_skills if skill not in newly_loaded] - - prior = agent_state.context.get("loaded_skills", []) - if not isinstance(prior, list): - prior = [] - merged_skills = sorted(set(prior).union(requested_skills)) - agent_state.update_context("loaded_skills", merged_skills) + # Runtime skill injection used to reach into the legacy + # ``_agent_instances`` registry to mutate the running LLM's + # active-skills list. The SDK harness owns the agent through + # ``Runner.run`` and there's no equivalent reach-in API yet — + # the model still gets a structured success response so it can + # observe which skills it asked for, even if reload-on-the-fly + # is a Phase 6 follow-up. + newly_loaded = list(requested_skills) + already_loaded: list[str] = [] except Exception as e: # noqa: BLE001 fallback_requested_skills = ( diff --git a/strix/tools/load_skill/load_skill_sdk_tool.py b/strix/tools/load_skill/tool.py similarity index 87% rename from strix/tools/load_skill/load_skill_sdk_tool.py rename to strix/tools/load_skill/tool.py index f60b089..4c5d051 100644 --- a/strix/tools/load_skill/load_skill_sdk_tool.py +++ b/strix/tools/load_skill/tool.py @@ -23,8 +23,8 @@ from typing import Any from agents import RunContextWrapper from strix.tools._decorator import strix_tool -from strix.tools._legacy_adapter import adapter_from_ctx -from strix.tools.load_skill import load_skill_actions as _legacy +from strix.tools._state_adapter import adapter_from_ctx +from strix.tools.load_skill import load_skill_actions as _impl def _dump(result: dict[str, Any]) -> str: @@ -42,5 +42,5 @@ async def load_skill(ctx: RunContextWrapper, skills: str) -> str: """ state = adapter_from_ctx(ctx) return _dump( - await asyncio.to_thread(_legacy.load_skill, agent_state=state, skills=skills), + await asyncio.to_thread(_impl.load_skill, agent_state=state, skills=skills), ) diff --git a/strix/tools/notes/notes_sdk_tools.py b/strix/tools/notes/tools.py similarity index 91% rename from strix/tools/notes/notes_sdk_tools.py rename to strix/tools/notes/tools.py index 4969e37..389fa81 100644 --- a/strix/tools/notes/notes_sdk_tools.py +++ b/strix/tools/notes/tools.py @@ -17,7 +17,7 @@ from typing import Any from agents import RunContextWrapper from strix.tools._decorator import strix_tool -from strix.tools.notes import notes_actions as _legacy +from strix.tools.notes import notes_actions as _impl def _dump(result: dict[str, Any]) -> str: @@ -48,7 +48,7 @@ async def create_note( # Wrap in to_thread so we don't block the event loop while waiting # on the lock or fsync. result = await asyncio.to_thread( - _legacy.create_note, + _impl.create_note, title=title, content=content, category=category, @@ -75,7 +75,7 @@ async def list_notes( when True, full content is included. """ result = await asyncio.to_thread( - _legacy.list_notes, + _impl.list_notes, category=category, tags=tags, search=search, @@ -87,7 +87,7 @@ async def list_notes( @strix_tool(timeout=30) async def get_note(ctx: RunContextWrapper, note_id: str) -> str: """Fetch one note by its 5-char ID. Returns full content.""" - result = await asyncio.to_thread(_legacy.get_note, note_id=note_id) + result = await asyncio.to_thread(_impl.get_note, note_id=note_id) return _dump(result) @@ -101,7 +101,7 @@ async def update_note( ) -> str: """Update a note's title, content, or tags. Pass ``None`` to leave a field unchanged.""" result = await asyncio.to_thread( - _legacy.update_note, + _impl.update_note, note_id=note_id, title=title, content=content, @@ -113,5 +113,5 @@ async def update_note( @strix_tool(timeout=30) async def delete_note(ctx: RunContextWrapper, note_id: str) -> str: """Delete a note. For wiki notes, also removes the rendered Markdown file.""" - result = await asyncio.to_thread(_legacy.delete_note, note_id=note_id) + result = await asyncio.to_thread(_impl.delete_note, note_id=note_id) return _dump(result) diff --git a/strix/tools/proxy/proxy_sdk_tools.py b/strix/tools/proxy/tools.py similarity index 100% rename from strix/tools/proxy/proxy_sdk_tools.py rename to strix/tools/proxy/tools.py diff --git a/strix/tools/python/python_sdk_tool.py b/strix/tools/python/tool.py similarity index 100% rename from strix/tools/python/python_sdk_tool.py rename to strix/tools/python/tool.py diff --git a/strix/tools/reporting/reporting_sdk_tools.py b/strix/tools/reporting/tool.py similarity index 96% rename from strix/tools/reporting/reporting_sdk_tools.py rename to strix/tools/reporting/tool.py index e476918..90adddd 100644 --- a/strix/tools/reporting/reporting_sdk_tools.py +++ b/strix/tools/reporting/tool.py @@ -20,7 +20,7 @@ from typing import Any from agents import RunContextWrapper from strix.tools._decorator import strix_tool -from strix.tools.reporting import reporting_actions as _legacy +from strix.tools.reporting import reporting_actions as _impl def _dump(result: dict[str, Any]) -> str: @@ -71,7 +71,7 @@ async def create_vulnerability_report( """ return _dump( await asyncio.to_thread( - _legacy.create_vulnerability_report, + _impl.create_vulnerability_report, title=title, description=description, impact=impact, diff --git a/strix/tools/terminal/terminal_sdk_tool.py b/strix/tools/terminal/tool.py similarity index 100% rename from strix/tools/terminal/terminal_sdk_tool.py rename to strix/tools/terminal/tool.py diff --git a/strix/tools/thinking/thinking_sdk_tools.py b/strix/tools/thinking/tool.py similarity index 100% rename from strix/tools/thinking/thinking_sdk_tools.py rename to strix/tools/thinking/tool.py diff --git a/strix/tools/todo/todo_sdk_tools.py b/strix/tools/todo/tools.py similarity index 89% rename from strix/tools/todo/todo_sdk_tools.py rename to strix/tools/todo/tools.py index 42f7846..c9e3c08 100644 --- a/strix/tools/todo/todo_sdk_tools.py +++ b/strix/tools/todo/tools.py @@ -17,8 +17,8 @@ from typing import Any from agents import RunContextWrapper from strix.tools._decorator import strix_tool -from strix.tools._legacy_adapter import adapter_from_ctx -from strix.tools.todo import todo_actions as _legacy +from strix.tools._state_adapter import adapter_from_ctx +from strix.tools.todo import todo_actions as _impl def _dump(result: dict[str, Any]) -> str: @@ -45,7 +45,7 @@ async def create_todo( """ state = adapter_from_ctx(ctx) return _dump( - _legacy.create_todo( + _impl.create_todo( agent_state=state, title=title, description=description, @@ -68,7 +68,7 @@ async def list_todos( priority: Optional ``"low" | "normal" | "high" | "critical"`` filter. """ state = adapter_from_ctx(ctx) - return _dump(_legacy.list_todos(agent_state=state, status=status, priority=priority)) + return _dump(_impl.list_todos(agent_state=state, status=status, priority=priority)) @strix_tool(timeout=30) @@ -91,7 +91,7 @@ async def update_todo( """ state = adapter_from_ctx(ctx) return _dump( - _legacy.update_todo( + _impl.update_todo( agent_state=state, todo_id=todo_id, title=title, @@ -112,7 +112,7 @@ async def mark_todo_done( """Mark one (``todo_id``) or many (``todo_ids``) todos as done.""" state = adapter_from_ctx(ctx) return _dump( - _legacy.mark_todo_done(agent_state=state, todo_id=todo_id, todo_ids=todo_ids), + _impl.mark_todo_done(agent_state=state, todo_id=todo_id, todo_ids=todo_ids), ) @@ -125,7 +125,7 @@ async def mark_todo_pending( """Mark one (``todo_id``) or many (``todo_ids``) todos as pending.""" state = adapter_from_ctx(ctx) return _dump( - _legacy.mark_todo_pending( + _impl.mark_todo_pending( agent_state=state, todo_id=todo_id, todo_ids=todo_ids, @@ -142,5 +142,5 @@ async def delete_todo( """Delete one (``todo_id``) or many (``todo_ids``) todos.""" state = adapter_from_ctx(ctx) return _dump( - _legacy.delete_todo(agent_state=state, todo_id=todo_id, todo_ids=todo_ids), + _impl.delete_todo(agent_state=state, todo_id=todo_id, todo_ids=todo_ids), ) diff --git a/strix/tools/web_search/web_search_sdk_tool.py b/strix/tools/web_search/tool.py similarity index 90% rename from strix/tools/web_search/web_search_sdk_tool.py rename to strix/tools/web_search/tool.py index 71fbced..3694e96 100644 --- a/strix/tools/web_search/web_search_sdk_tool.py +++ b/strix/tools/web_search/tool.py @@ -17,7 +17,7 @@ from typing import Any from agents import RunContextWrapper from strix.tools._decorator import strix_tool -from strix.tools.web_search import web_search_actions as _legacy +from strix.tools.web_search import web_search_actions as _impl def _dump(result: dict[str, Any]) -> str: @@ -39,4 +39,4 @@ async def web_search(ctx: RunContextWrapper, query: str) -> str: system prompt to bias results toward CVEs, exploits, and Kali- compatible commands. """ - return _dump(await asyncio.to_thread(_legacy.web_search, query=query)) + return _dump(await asyncio.to_thread(_impl.web_search, query=query)) diff --git a/tests/agents/test_sdk_factory.py b/tests/agents/test_factory.py similarity index 96% rename from tests/agents/test_sdk_factory.py rename to tests/agents/test_factory.py index 452dcb9..009d9d4 100644 --- a/tests/agents/test_sdk_factory.py +++ b/tests/agents/test_factory.py @@ -21,8 +21,8 @@ from unittest.mock import patch from agents import Agent from agents.tool import FunctionTool -from strix.agents.sdk_factory import build_strix_agent, make_child_factory -from strix.agents.sdk_prompt import _resolve_skills, render_system_prompt +from strix.agents.factory import build_strix_agent, make_child_factory +from strix.agents.prompt import _resolve_skills, render_system_prompt # --- prompt renderer ---------------------------------------------------- @@ -60,7 +60,7 @@ def test_render_system_prompt_swallows_template_errors() -> None: """If the template path can't be resolved, return an empty string (not raise) — agent construction must never blow up on prompt load.""" with patch( - "strix.agents.sdk_prompt.get_strix_resource_path", + "strix.agents.prompt.get_strix_resource_path", side_effect=RuntimeError("missing"), ): out = render_system_prompt(skills=[]) @@ -191,7 +191,7 @@ def test_make_child_factory_passes_scan_level_config() -> None: interactive=True, system_prompt_context={"scope_source": "test"}, ) - with patch("strix.agents.sdk_factory.render_system_prompt", side_effect=fake_render): + with patch("strix.agents.factory.render_system_prompt", side_effect=fake_render): factory(name="child", skills=["xss"]) assert captured["scan_mode"] == "fast" diff --git a/tests/interface/test_sdk_dispatch.py b/tests/interface/test_sdk_dispatch.py deleted file mode 100644 index 5591202..0000000 --- a/tests/interface/test_sdk_dispatch.py +++ /dev/null @@ -1,199 +0,0 @@ -"""Phase 5b tests for the STRIX_USE_SDK_HARNESS dispatch. - -Covers the env-flag reader, source-path resolution, sandbox image -lookup, and the adapter that translates legacy CLI args into -``run_strix_scan`` kwargs. - -We never call ``run_strix_scan`` for real — that requires a live -Docker daemon + LLM. The tests patch it and verify the kwargs handoff. -""" - -from __future__ import annotations - -import logging -from pathlib import Path -from types import SimpleNamespace -from typing import Any -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest - -from strix.interface.sdk_dispatch import ( - _resolve_sandbox_image, - _resolve_sources_path, - run_scan_via_sdk, - should_use_sdk_harness, -) - - -# --- env flag reader ---------------------------------------------------- - - -@pytest.mark.parametrize( - ("value", "expected"), - [ - ("1", True), - ("true", True), - ("True", True), - ("YES", True), - ("0", False), - ("false", False), - ("no", False), - ("", False), - ("anything-else", False), - ], -) -def test_should_use_sdk_harness_parses_env( - monkeypatch: pytest.MonkeyPatch, - value: str, - expected: bool, -) -> None: - monkeypatch.setenv("STRIX_USE_SDK_HARNESS", value) - assert should_use_sdk_harness() is expected - - -def test_should_use_sdk_harness_defaults_false_when_unset( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.delenv("STRIX_USE_SDK_HARNESS", raising=False) - assert should_use_sdk_harness() is False - - -# --- image lookup ------------------------------------------------------- - - -def test_resolve_sandbox_image_uses_config_value() -> None: - with patch( - "strix.config.Config.get", - return_value="strix-sandbox:0.1.13", - ): - assert _resolve_sandbox_image() == "strix-sandbox:0.1.13" - - -def test_resolve_sandbox_image_falls_back_when_unset( - caplog: pytest.LogCaptureFixture, -) -> None: - with ( - patch("strix.config.Config.get", return_value=None), - caplog.at_level(logging.WARNING, logger="strix.interface.sdk_dispatch"), - ): - out = _resolve_sandbox_image() - assert out == "strix-sandbox:latest" - assert any("strix_image not configured" in r.message for r in caplog.records) - - -# --- sources path ------------------------------------------------------- - - -def test_resolve_sources_path_uses_local_sources_parent(tmp_path: Path) -> None: - """When --local-sources is given, mount that path's parent so the - agent can walk down into the actual source directory tree.""" - src_dir = tmp_path / "my-project" - src_dir.mkdir() - args = SimpleNamespace( - local_sources=[{"host_path": str(src_dir)}], - run_name="run-1", - ) - assert _resolve_sources_path(args) == tmp_path - - -def test_resolve_sources_path_handles_alternative_keys(tmp_path: Path) -> None: - """Some legacy paths use 'source_path' or 'path' instead of - 'host_path' — we accept all three.""" - src_dir = tmp_path / "alt" - src_dir.mkdir() - args = SimpleNamespace( - local_sources=[{"path": str(src_dir)}], - run_name="run-2", - ) - assert _resolve_sources_path(args) == tmp_path - - -def test_resolve_sources_path_creates_scratch_dir_when_absent( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path)) - args = SimpleNamespace(local_sources=None, run_name="scan-x") - out = _resolve_sources_path(args) - assert out == tmp_path / "strix" / "sources" / "scan-x" - assert out.exists() - assert out.is_dir() - - -# --- adapter ----------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_run_scan_via_sdk_translates_args_to_kwargs( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - """Verify every kwarg the entry point reads is forwarded correctly.""" - monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path)) - - scan_config = {"targets": [], "scan_mode": "deep"} - args = SimpleNamespace( - run_name="scan-42", - local_sources=None, - interactive=True, - ) - fake_tracer = MagicMock(name="tracer") - - fake_run = AsyncMock(return_value=MagicMock(name="run_result")) - with ( - patch("strix.config.Config.get", return_value="strix-sandbox:test"), - patch("strix.sdk_entry.run_strix_scan", new=fake_run), - ): - await run_scan_via_sdk(scan_config=scan_config, args=args, tracer=fake_tracer) - - fake_run.assert_awaited_once() - assert fake_run.await_args is not None - kwargs = fake_run.await_args.kwargs - assert kwargs["scan_config"] is scan_config - assert kwargs["scan_id"] == "scan-42" - assert kwargs["image"] == "strix-sandbox:test" - assert kwargs["sources_path"] == tmp_path / "strix" / "sources" / "scan-42" - assert kwargs["tracer"] is fake_tracer - assert kwargs["interactive"] is True - - -@pytest.mark.asyncio -async def test_run_scan_via_sdk_falls_back_to_scan_config_run_name( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - """If args has no run_name, scan_config['run_name'] should be used.""" - monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path)) - scan_config = {"targets": [], "run_name": "from-config"} - args = SimpleNamespace(local_sources=None) - - fake_run = AsyncMock(return_value=MagicMock()) - with ( - patch("strix.config.Config.get", return_value="img:1"), - patch("strix.sdk_entry.run_strix_scan", new=fake_run), - ): - await run_scan_via_sdk(scan_config=scan_config, args=args, tracer=None) - - assert fake_run.await_args is not None - assert fake_run.await_args.kwargs["scan_id"] == "from-config" - - -@pytest.mark.asyncio -async def test_run_scan_via_sdk_propagates_run_failure( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - """A failure inside run_strix_scan should bubble up to the caller — - the legacy CLI relies on raised exceptions for the SDK path.""" - monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path)) - scan_config: dict[str, Any] = {"targets": []} - args = SimpleNamespace(run_name="r", local_sources=None) - - fake_run = AsyncMock(side_effect=RuntimeError("boom")) - with ( - patch("strix.config.Config.get", return_value="img"), - patch("strix.sdk_entry.run_strix_scan", new=fake_run), - pytest.raises(RuntimeError, match="boom"), - ): - await run_scan_via_sdk(scan_config=scan_config, args=args, tracer=None) diff --git a/tests/llm/test_llm_otel.py b/tests/llm/test_llm_otel.py deleted file mode 100644 index a11ffa5..0000000 --- a/tests/llm/test_llm_otel.py +++ /dev/null @@ -1,16 +0,0 @@ -import litellm -import pytest - -from strix.llm.config import LLMConfig -from strix.llm.llm import LLM - - -def test_llm_does_not_modify_litellm_callbacks(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("STRIX_TELEMETRY", "1") - monkeypatch.setenv("STRIX_OTEL_TELEMETRY", "1") - monkeypatch.setattr(litellm, "callbacks", ["custom-callback"]) - - llm = LLM(LLMConfig(model_name="openai/gpt-5.4"), agent_name=None) - - assert llm is not None - assert litellm.callbacks == ["custom-callback"] diff --git a/tests/llm/test_source_aware_whitebox.py b/tests/llm/test_source_aware_whitebox.py deleted file mode 100644 index c43a5c4..0000000 --- a/tests/llm/test_source_aware_whitebox.py +++ /dev/null @@ -1,30 +0,0 @@ -from strix.llm.config import LLMConfig -from strix.llm.llm import LLM - - -def test_llm_config_whitebox_defaults_to_false(monkeypatch) -> None: - monkeypatch.setenv("STRIX_LLM", "openai/gpt-5") - config = LLMConfig() - assert config.is_whitebox is False - - -def test_llm_config_whitebox_can_be_enabled(monkeypatch) -> None: - monkeypatch.setenv("STRIX_LLM", "openai/gpt-5") - config = LLMConfig(is_whitebox=True) - assert config.is_whitebox is True - - -def test_whitebox_prompt_loads_source_aware_coordination_skill(monkeypatch) -> None: - monkeypatch.setenv("STRIX_LLM", "openai/gpt-5") - - whitebox_llm = LLM(LLMConfig(scan_mode="quick", is_whitebox=True), agent_name="StrixAgent") - assert "" in whitebox_llm.system_prompt - assert "" in whitebox_llm.system_prompt - assert "Begin with fast source triage" in whitebox_llm.system_prompt - assert "You MUST begin at the very first step by running the code and testing live." not in ( - whitebox_llm.system_prompt - ) - - non_whitebox_llm = LLM(LLMConfig(scan_mode="quick", is_whitebox=False), agent_name="StrixAgent") - assert "" not in non_whitebox_llm.system_prompt - assert "" not in non_whitebox_llm.system_prompt diff --git a/tests/telemetry/test_tracer.py b/tests/telemetry/test_tracer.py index 54c9cb8..acbf31f 100644 --- a/tests/telemetry/test_tracer.py +++ b/tests/telemetry/test_tracer.py @@ -10,7 +10,6 @@ from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExportResult from strix.telemetry import tracer as tracer_module from strix.telemetry import utils as telemetry_utils from strix.telemetry.tracer import Tracer, set_global_tracer -from strix.tools.agents_graph import agents_graph_actions def _load_events(events_path: Path) -> list[dict[str, Any]]: @@ -19,7 +18,7 @@ def _load_events(events_path: Path) -> list[dict[str, Any]]: @pytest.fixture(autouse=True) -def _reset_tracer_globals(monkeypatch) -> None: +def _reset_tracer_globals(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(tracer_module, "_global_tracer", None) monkeypatch.setattr(tracer_module, "_OTEL_BOOTSTRAPPED", False) monkeypatch.setattr(tracer_module, "_OTEL_REMOTE_ENABLED", False) @@ -32,7 +31,9 @@ def _reset_tracer_globals(monkeypatch) -> None: monkeypatch.delenv("TRACELOOP_HEADERS", raising=False) -def test_tracer_local_mode_writes_jsonl_with_correlation(monkeypatch, tmp_path) -> None: +def test_tracer_local_mode_writes_jsonl_with_correlation( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: monkeypatch.chdir(tmp_path) tracer = Tracer("local-observability") @@ -60,7 +61,7 @@ def test_tracer_local_mode_writes_jsonl_with_correlation(monkeypatch, tmp_path) assert event["span_id"] -def test_tracer_redacts_sensitive_payloads(monkeypatch, tmp_path) -> None: +def test_tracer_redacts_sensitive_payloads(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: monkeypatch.chdir(tmp_path) tracer = Tracer("redaction-run") @@ -89,7 +90,9 @@ def test_tracer_redacts_sensitive_payloads(monkeypatch, tmp_path) -> None: assert "[REDACTED]" in serialized -def test_tracer_remote_mode_configures_traceloop_export(monkeypatch, tmp_path) -> None: +def test_tracer_remote_mode_configures_traceloop_export( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: monkeypatch.chdir(tmp_path) class FakeTraceloop: @@ -128,7 +131,9 @@ def test_tracer_remote_mode_configures_traceloop_export(monkeypatch, tmp_path) - assert run_started["payload"]["remote_export_enabled"] is True -def test_tracer_local_mode_avoids_traceloop_remote_endpoint(monkeypatch, tmp_path) -> None: +def test_tracer_local_mode_avoids_traceloop_remote_endpoint( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: monkeypatch.chdir(tmp_path) class FakeTraceloop: @@ -157,7 +162,9 @@ def test_tracer_local_mode_avoids_traceloop_remote_endpoint(monkeypatch, tmp_pat assert tracer._remote_export_enabled is False -def test_otlp_fallback_includes_auth_and_custom_headers(monkeypatch, tmp_path) -> None: +def test_otlp_fallback_includes_auth_and_custom_headers( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: monkeypatch.chdir(tmp_path) monkeypatch.setattr(tracer_module, "Traceloop", None) monkeypatch.setenv("TRACELOOP_BASE_URL", "https://otel.example.com") @@ -172,13 +179,13 @@ def test_otlp_fallback_includes_auth_and_custom_headers(monkeypatch, tmp_path) - captured["headers"] = headers or {} captured["kwargs"] = kwargs - def export(self, spans: Any) -> SpanExportResult: # noqa: ARG002 + def export(self, spans: Any) -> SpanExportResult: return SpanExportResult.SUCCESS def shutdown(self) -> None: return None - def force_flush(self, timeout_millis: int = 30_000) -> bool: # noqa: ARG002 + def force_flush(self, timeout_millis: int = 30_000) -> bool: return True fake_module = types.ModuleType("opentelemetry.exporter.otlp.proto.http.trace_exporter") @@ -199,7 +206,8 @@ def test_otlp_fallback_includes_auth_and_custom_headers(monkeypatch, tmp_path) - def test_traceloop_init_failure_does_not_mark_bootstrapped_on_provider_failure( - monkeypatch, tmp_path + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, ) -> None: monkeypatch.chdir(tmp_path) @@ -226,7 +234,7 @@ def test_traceloop_init_failure_does_not_mark_bootstrapped_on_provider_failure( assert tracer._remote_export_enabled is False -def test_run_completed_event_emitted_once(monkeypatch, tmp_path) -> None: +def test_run_completed_event_emitted_once(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: monkeypatch.chdir(tmp_path) tracer = Tracer("single-complete") @@ -240,7 +248,9 @@ def test_run_completed_event_emitted_once(monkeypatch, tmp_path) -> None: assert len(run_completed) == 1 -def test_events_with_agent_id_include_agent_name(monkeypatch, tmp_path) -> None: +def test_events_with_agent_id_include_agent_name( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: monkeypatch.chdir(tmp_path) tracer = Tracer("agent-name-enrichment") @@ -256,61 +266,32 @@ def test_events_with_agent_id_include_agent_name(monkeypatch, tmp_path) -> None: assert chat_event["actor"]["agent_name"] == "Root Agent" -def test_get_total_llm_stats_includes_completed_subagents(monkeypatch, tmp_path) -> None: +def test_get_total_llm_stats_aggregates_live_and_completed( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: monkeypatch.chdir(tmp_path) - - class DummyStats: - def __init__( - self, - *, - input_tokens: int, - output_tokens: int, - cached_tokens: int, - cost: float, - requests: int, - ) -> None: - self.input_tokens = input_tokens - self.output_tokens = output_tokens - self.cached_tokens = cached_tokens - self.cost = cost - self.requests = requests - - class DummyLLM: - def __init__(self, stats: DummyStats) -> None: - self._total_stats = stats - - class DummyAgent: - def __init__(self, stats: DummyStats) -> None: - self.llm = DummyLLM(stats) - tracer = Tracer("cost-rollup") set_global_tracer(tracer) - monkeypatch.setattr( - agents_graph_actions, - "_agent_instances", - { - "root-agent": DummyAgent( - DummyStats( - input_tokens=1_000, - output_tokens=250, - cached_tokens=100, - cost=0.12345, - requests=2, - ) - ) - }, + # 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", ) - monkeypatch.setattr( - agents_graph_actions, - "_completed_agent_llm_totals", - { - "input_tokens": 2_000, - "output_tokens": 500, - "cached_tokens": 400, - "cost": 0.54321, - "requests": 3, - }, + # 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() @@ -325,7 +306,9 @@ def test_get_total_llm_stats_includes_completed_subagents(monkeypatch, tmp_path) assert stats["total_tokens"] == 3_750 -def test_run_metadata_is_only_on_run_lifecycle_events(monkeypatch, tmp_path) -> None: +def test_run_metadata_is_only_on_run_lifecycle_events( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: monkeypatch.chdir(tmp_path) tracer = Tracer("metadata-scope") @@ -345,7 +328,7 @@ def test_run_metadata_is_only_on_run_lifecycle_events(monkeypatch, tmp_path) -> assert "run_metadata" not in chat_event -def test_set_run_name_resets_cached_paths(monkeypatch, tmp_path) -> None: +def test_set_run_name_resets_cached_paths(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: monkeypatch.chdir(tmp_path) tracer = Tracer() @@ -364,7 +347,9 @@ def test_set_run_name_resets_cached_paths(monkeypatch, tmp_path) -> None: assert any(event["event_type"] == "chat.message" for event in events) -def test_set_run_name_resets_run_completed_flag(monkeypatch, tmp_path) -> None: +def test_set_run_name_resets_run_completed_flag( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: monkeypatch.chdir(tmp_path) tracer = Tracer() @@ -382,7 +367,9 @@ def test_set_run_name_resets_run_completed_flag(monkeypatch, tmp_path) -> None: assert len(run_completed) == 1 -def test_set_run_name_updates_traceloop_association_properties(monkeypatch, tmp_path) -> None: +def test_set_run_name_updates_traceloop_association_properties( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: monkeypatch.chdir(tmp_path) class FakeTraceloop: @@ -407,7 +394,9 @@ def test_set_run_name_updates_traceloop_association_properties(monkeypatch, tmp_ assert FakeTraceloop.associations[-1]["run_name"] == "renamed-run" -def test_events_write_locks_are_scoped_by_events_file(monkeypatch, tmp_path) -> None: +def test_events_write_locks_are_scoped_by_events_file( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: monkeypatch.chdir(tmp_path) monkeypatch.setenv("STRIX_TELEMETRY", "0") @@ -422,7 +411,9 @@ def test_events_write_locks_are_scoped_by_events_file(monkeypatch, tmp_path) -> assert lock_a_from_one is not lock_b -def test_tracer_skips_jsonl_when_telemetry_disabled(monkeypatch, tmp_path) -> None: +def test_tracer_skips_jsonl_when_telemetry_disabled( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: monkeypatch.chdir(tmp_path) monkeypatch.setenv("STRIX_TELEMETRY", "0") @@ -435,7 +426,9 @@ def test_tracer_skips_jsonl_when_telemetry_disabled(monkeypatch, tmp_path) -> No assert not events_path.exists() -def test_tracer_otel_flag_overrides_global_telemetry(monkeypatch, tmp_path) -> None: +def test_tracer_otel_flag_overrides_global_telemetry( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: monkeypatch.chdir(tmp_path) monkeypatch.setenv("STRIX_TELEMETRY", "0") monkeypatch.setenv("STRIX_OTEL_TELEMETRY", "1") diff --git a/tests/test_sdk_entry.py b/tests/test_entry.py similarity index 86% rename from tests/test_sdk_entry.py rename to tests/test_entry.py index c9e383f..999862a 100644 --- a/tests/test_sdk_entry.py +++ b/tests/test_entry.py @@ -24,8 +24,8 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from strix.entry import _build_root_task, _build_scope_context, run_strix_scan from strix.orchestration.bus import AgentMessageBus -from strix.sdk_entry import _build_root_task, _build_scope_context, run_strix_scan # --- helpers ------------------------------------------------------------ @@ -140,18 +140,18 @@ async def test_run_strix_scan_wires_context_and_calls_runner(tmp_path: Path) -> with ( patch( - "strix.sdk_entry.session_manager.create_or_reuse", + "strix.entry.session_manager.create_or_reuse", new=AsyncMock(return_value=bundle), ) as create_mock, patch( - "strix.sdk_entry.session_manager.cleanup", + "strix.entry.session_manager.cleanup", new=AsyncMock(), ) as cleanup_mock, - patch("strix.sdk_entry.Runner.run", side_effect=fake_runner_run) as runner_mock, + patch("strix.entry.Runner.run", side_effect=fake_runner_run) as runner_mock, # Stub the factory to avoid rendering the 158k-char prompt for # every test (it's covered by sdk_prompt tests). patch( - "strix.sdk_entry.build_strix_agent", + "strix.entry.build_strix_agent", return_value=MagicMock(name="root_agent"), ) as factory_mock, ): @@ -200,18 +200,18 @@ async def test_run_strix_scan_cleans_up_on_runner_failure(tmp_path: Path) -> Non with ( patch( - "strix.sdk_entry.session_manager.create_or_reuse", + "strix.entry.session_manager.create_or_reuse", new=AsyncMock(return_value=bundle), ), patch( - "strix.sdk_entry.session_manager.cleanup", + "strix.entry.session_manager.cleanup", new=AsyncMock(), ) as cleanup_mock, patch( - "strix.sdk_entry.Runner.run", + "strix.entry.Runner.run", side_effect=RuntimeError("simulated LLM blow-up"), ), - patch("strix.sdk_entry.build_strix_agent", return_value=MagicMock()), + patch("strix.entry.build_strix_agent", return_value=MagicMock()), pytest.raises(RuntimeError, match="simulated LLM"), ): await run_strix_scan( @@ -233,15 +233,15 @@ async def test_run_strix_scan_skips_cleanup_when_disabled(tmp_path: Path) -> Non with ( patch( - "strix.sdk_entry.session_manager.create_or_reuse", + "strix.entry.session_manager.create_or_reuse", new=AsyncMock(return_value=bundle), ), patch( - "strix.sdk_entry.session_manager.cleanup", + "strix.entry.session_manager.cleanup", new=AsyncMock(), ) as cleanup_mock, - patch("strix.sdk_entry.Runner.run", side_effect=fake_runner_run), - patch("strix.sdk_entry.build_strix_agent", return_value=MagicMock()), + patch("strix.entry.Runner.run", side_effect=fake_runner_run), + patch("strix.entry.build_strix_agent", return_value=MagicMock()), ): await run_strix_scan( scan_config=_scan_config(), @@ -267,12 +267,12 @@ async def test_run_strix_scan_auto_generates_scan_id(tmp_path: Path) -> None: with ( patch( - "strix.sdk_entry.session_manager.create_or_reuse", + "strix.entry.session_manager.create_or_reuse", new=AsyncMock(side_effect=fake_create), ), - patch("strix.sdk_entry.session_manager.cleanup", new=AsyncMock()), - patch("strix.sdk_entry.Runner.run", new=AsyncMock(return_value=MagicMock())), - patch("strix.sdk_entry.build_strix_agent", return_value=MagicMock()), + patch("strix.entry.session_manager.cleanup", new=AsyncMock()), + patch("strix.entry.Runner.run", new=AsyncMock(return_value=MagicMock())), + patch("strix.entry.build_strix_agent", return_value=MagicMock()), ): await run_strix_scan( scan_config=_scan_config(), @@ -300,12 +300,12 @@ async def test_run_strix_scan_passes_scan_level_config_into_factory( with ( patch( - "strix.sdk_entry.session_manager.create_or_reuse", + "strix.entry.session_manager.create_or_reuse", new=AsyncMock(return_value=bundle), ), - patch("strix.sdk_entry.session_manager.cleanup", new=AsyncMock()), - patch("strix.sdk_entry.Runner.run", new=AsyncMock(return_value=MagicMock())), - patch("strix.sdk_entry.build_strix_agent", side_effect=fake_factory), + patch("strix.entry.session_manager.cleanup", new=AsyncMock()), + patch("strix.entry.Runner.run", new=AsyncMock(return_value=MagicMock())), + patch("strix.entry.build_strix_agent", side_effect=fake_factory), ): await run_strix_scan( scan_config=_scan_config(scan_mode="fast", is_whitebox=True), diff --git a/tests/tools/test_agents_graph_whitebox.py b/tests/tools/test_agents_graph_whitebox.py deleted file mode 100644 index 60226be..0000000 --- a/tests/tools/test_agents_graph_whitebox.py +++ /dev/null @@ -1,291 +0,0 @@ -from types import SimpleNamespace - -import strix.agents as agents_module -from strix.llm.config import LLMConfig -from strix.tools.agents_graph import agents_graph_actions - - -def _reset_agent_graph_state() -> None: - agents_graph_actions._agent_graph["nodes"].clear() - agents_graph_actions._agent_graph["edges"].clear() - agents_graph_actions._agent_messages.clear() - agents_graph_actions._running_agents.clear() - agents_graph_actions._agent_instances.clear() - agents_graph_actions._completed_agent_llm_totals.clear() - agents_graph_actions._completed_agent_llm_totals.update( - agents_graph_actions._empty_llm_stats_totals() - ) - agents_graph_actions._agent_states.clear() - - -def test_create_agent_inherits_parent_whitebox_flag(monkeypatch) -> None: - monkeypatch.setenv("STRIX_LLM", "openai/gpt-5") - - _reset_agent_graph_state() - - parent_id = "parent-agent" - parent_llm = LLMConfig(timeout=123, scan_mode="standard", is_whitebox=True) - agents_graph_actions._agent_instances[parent_id] = SimpleNamespace( - llm_config=parent_llm, - non_interactive=True, - ) - - captured_config: dict[str, object] = {} - - class FakeStrixAgent: - def __init__(self, config: dict[str, object]): - captured_config["agent_config"] = config - - class FakeThread: - def __init__(self, target, args, daemon, name): - self.target = target - self.args = args - self.daemon = daemon - self.name = name - - def start(self) -> None: - return None - - monkeypatch.setattr(agents_module, "StrixAgent", FakeStrixAgent) - monkeypatch.setattr(agents_graph_actions.threading, "Thread", FakeThread) - - agent_state = SimpleNamespace( - agent_id=parent_id, - get_conversation_history=list, - ) - result = agents_graph_actions.create_agent( - agent_state=agent_state, - task="source-aware child task", - name="SourceAwareChild", - inherit_context=False, - ) - - assert result["success"] is True - llm_config = captured_config["agent_config"]["llm_config"] - assert isinstance(llm_config, LLMConfig) - assert llm_config.timeout == 123 - assert llm_config.scan_mode == "standard" - assert llm_config.is_whitebox is True - child_task = captured_config["agent_config"]["state"].task - assert "White-box execution guidance (recommended when source is available):" in child_task - assert "mandatory" not in child_task.lower() - - -def test_delegation_prompt_includes_wiki_memory_instruction_in_whitebox(monkeypatch) -> None: - monkeypatch.setenv("STRIX_LLM", "openai/gpt-5") - - _reset_agent_graph_state() - - parent_id = "parent-1" - child_id = "child-1" - agents_graph_actions._agent_graph["nodes"][parent_id] = {"name": "Parent", "status": "running"} - agents_graph_actions._agent_graph["nodes"][child_id] = {"name": "Child", "status": "running"} - - class FakeState: - def __init__(self) -> None: - self.agent_id = child_id - self.agent_name = "Child" - self.parent_id = parent_id - self.task = "analyze source risks" - self.stop_requested = False - self.messages: list[tuple[str, str]] = [] - - def add_message(self, role: str, content: str) -> None: - self.messages.append((role, content)) - - def model_dump(self) -> dict[str, str]: - return {"agent_id": self.agent_id} - - class FakeAgent: - def __init__(self) -> None: - self.llm_config = LLMConfig(is_whitebox=True) - - async def agent_loop(self, _task: str) -> dict[str, bool]: - return {"ok": True} - - state = FakeState() - agent = FakeAgent() - agents_graph_actions._agent_instances[child_id] = agent - result = agents_graph_actions._run_agent_in_thread(agent, state, inherited_messages=[]) - - assert result["result"] == {"ok": True} - task_messages = [msg for role, msg in state.messages if role == "user"] - assert task_messages - assert 'list_notes(category="wiki")' in task_messages[-1] - assert "get_note(note_id=...)" in task_messages[-1] - assert "Before agent_finish" in task_messages[-1] - - -def test_agent_finish_appends_wiki_update_for_whitebox(monkeypatch) -> None: - monkeypatch.setenv("STRIX_LLM", "openai/gpt-5") - - _reset_agent_graph_state() - - parent_id = "parent-2" - child_id = "child-2" - agents_graph_actions._agent_graph["nodes"][parent_id] = { - "name": "Parent", - "task": "parent task", - "status": "running", - "parent_id": None, - } - agents_graph_actions._agent_graph["nodes"][child_id] = { - "name": "Child", - "task": "child task", - "status": "running", - "parent_id": parent_id, - } - agents_graph_actions._agent_instances[child_id] = SimpleNamespace( - llm_config=LLMConfig(is_whitebox=True) - ) - - captured: dict[str, str] = {} - - def fake_list_notes(category=None): - assert category == "wiki" - return { - "success": True, - "notes": [{"note_id": "wiki-note-1", "content": "Existing wiki content"}], - "total_count": 1, - } - - captured_get: dict[str, str] = {} - - def fake_get_note(note_id: str): - captured_get["note_id"] = note_id - return { - "success": True, - "note": { - "note_id": note_id, - "title": "Repo Wiki", - "content": "Existing wiki content", - }, - } - - def fake_append_note_content(note_id: str, delta: str): - captured["note_id"] = note_id - captured["delta"] = delta - return {"success": True, "note_id": note_id} - - monkeypatch.setattr("strix.tools.notes.notes_actions.list_notes", fake_list_notes) - monkeypatch.setattr("strix.tools.notes.notes_actions.get_note", fake_get_note) - monkeypatch.setattr("strix.tools.notes.notes_actions.append_note_content", fake_append_note_content) - - state = SimpleNamespace(agent_id=child_id, parent_id=parent_id) - result = agents_graph_actions.agent_finish( - agent_state=state, - result_summary="AST pass completed", - findings=["Found route sink candidate"], - success=True, - final_recommendations=["Validate sink with dynamic PoC"], - ) - - assert result["agent_completed"] is True - assert captured_get["note_id"] == "wiki-note-1" - assert captured["note_id"] == "wiki-note-1" - assert "Agent Update: Child" in captured["delta"] - assert "AST pass completed" in captured["delta"] - - -def test_run_agent_in_thread_injects_shared_wiki_context_in_whitebox(monkeypatch) -> None: - monkeypatch.setenv("STRIX_LLM", "openai/gpt-5") - - _reset_agent_graph_state() - - parent_id = "parent-3" - child_id = "child-3" - agents_graph_actions._agent_graph["nodes"][parent_id] = {"name": "Parent", "status": "running"} - agents_graph_actions._agent_graph["nodes"][child_id] = {"name": "Child", "status": "running"} - - class FakeState: - def __init__(self) -> None: - self.agent_id = child_id - self.agent_name = "Child" - self.parent_id = parent_id - self.task = "map source" - self.stop_requested = False - self.messages: list[tuple[str, str]] = [] - - def add_message(self, role: str, content: str) -> None: - self.messages.append((role, content)) - - def model_dump(self) -> dict[str, str]: - return {"agent_id": self.agent_id} - - class FakeAgent: - def __init__(self) -> None: - self.llm_config = LLMConfig(is_whitebox=True) - - async def agent_loop(self, _task: str) -> dict[str, bool]: - return {"ok": True} - - captured_get: dict[str, str] = {} - - def fake_list_notes(category=None): - assert category == "wiki" - return { - "success": True, - "notes": [{"note_id": "wiki-ctx-1"}], - "total_count": 1, - } - - def fake_get_note(note_id: str): - captured_get["note_id"] = note_id - return { - "success": True, - "note": { - "note_id": note_id, - "title": "Shared Repo Wiki", - "content": "Architecture: server/client split", - }, - } - - monkeypatch.setattr("strix.tools.notes.notes_actions.list_notes", fake_list_notes) - monkeypatch.setattr("strix.tools.notes.notes_actions.get_note", fake_get_note) - - state = FakeState() - agent = FakeAgent() - agents_graph_actions._agent_instances[child_id] = agent - result = agents_graph_actions._run_agent_in_thread(agent, state, inherited_messages=[]) - - assert result["result"] == {"ok": True} - assert captured_get["note_id"] == "wiki-ctx-1" - user_messages = [content for role, content in state.messages if role == "user"] - assert user_messages - assert " None: - selected_note_ids: list[str] = [] - - def fake_list_notes(category=None): - assert category == "wiki" - return { - "success": True, - "notes": [ - {"note_id": "wiki-other", "tags": ["repo:other"]}, - {"note_id": "wiki-target", "tags": ["repo:appsmith"]}, - ], - "total_count": 2, - } - - def fake_get_note(note_id: str): - selected_note_ids.append(note_id) - return { - "success": True, - "note": {"note_id": note_id, "title": "Repo Wiki", "content": "content"}, - } - - monkeypatch.setattr("strix.tools.notes.notes_actions.list_notes", fake_list_notes) - monkeypatch.setattr("strix.tools.notes.notes_actions.get_note", fake_get_note) - - agent_state = SimpleNamespace( - task="analyze /workspace/appsmith", - context={"whitebox_repo_tags": ["repo:appsmith"]}, - ) - note = agents_graph_actions._load_primary_wiki_note(agent_state) - - assert note is not None - assert note["note_id"] == "wiki-target" - assert selected_note_ids == ["wiki-target"] diff --git a/tests/tools/test_sdk_graph_tools.py b/tests/tools/test_graph_tools.py similarity index 98% rename from tests/tools/test_sdk_graph_tools.py rename to tests/tools/test_graph_tools.py index 0eb0daf..5a30a2b 100644 --- a/tests/tools/test_sdk_graph_tools.py +++ b/tests/tools/test_graph_tools.py @@ -29,7 +29,7 @@ import pytest from agents.tool import FunctionTool from strix.orchestration.bus import AgentMessageBus -from strix.tools.agents_graph.agents_graph_sdk_tools import ( +from strix.tools.agents_graph.tools import ( agent_finish, agent_status, create_agent, @@ -289,7 +289,7 @@ async def test_create_agent_spawns_and_registers_child() -> None: ) with patch( - "strix.tools.agents_graph.agents_graph_sdk_tools.Runner.run", + "strix.tools.agents_graph.tools.Runner.run", side_effect=fake_runner_run, ): out = await _invoke( @@ -361,7 +361,7 @@ async def test_create_agent_inherits_parent_history() -> None: ) with patch( - "strix.tools.agents_graph.agents_graph_sdk_tools.Runner.run", + "strix.tools.agents_graph.tools.Runner.run", side_effect=fake_runner_run, ): await _invoke( @@ -493,7 +493,7 @@ async def test_create_agent_spawn_is_cancelable_via_bus() -> None: runner_mock = AsyncMock(side_effect=slow_runner_run) with patch( - "strix.tools.agents_graph.agents_graph_sdk_tools.Runner.run", + "strix.tools.agents_graph.tools.Runner.run", new=runner_mock, ): out = await _invoke( diff --git a/tests/tools/test_load_skill_tool.py b/tests/tools/test_load_skill_tool.py deleted file mode 100644 index 7180d45..0000000 --- a/tests/tools/test_load_skill_tool.py +++ /dev/null @@ -1,139 +0,0 @@ -from typing import Any - -from strix.tools.agents_graph import agents_graph_actions -from strix.tools.load_skill import load_skill_actions - - -class _DummyLLM: - def __init__(self, initial_skills: list[str] | None = None) -> None: - self.loaded: set[str] = set(initial_skills or []) - - def add_skills(self, skill_names: list[str]) -> list[str]: - newly_loaded = [skill for skill in skill_names if skill not in self.loaded] - self.loaded.update(newly_loaded) - return newly_loaded - - -class _DummyAgent: - def __init__(self, initial_skills: list[str] | None = None) -> None: - self.llm = _DummyLLM(initial_skills) - - -class _DummyAgentState: - def __init__(self, agent_id: str) -> None: - self.agent_id = agent_id - self.context: dict[str, Any] = {} - - def update_context(self, key: str, value: Any) -> None: - self.context[key] = value - - -def test_load_skill_success_and_context_update() -> None: - instances = agents_graph_actions.__dict__["_agent_instances"] - original_instances = dict(instances) - try: - state = _DummyAgentState("agent_test_load_skill_success") - instances.clear() - instances[state.agent_id] = _DummyAgent() - - result = load_skill_actions.load_skill(state, "ffuf,xss") - - assert result["success"] is True - assert result["loaded_skills"] == ["ffuf", "xss"] - assert result["newly_loaded_skills"] == ["ffuf", "xss"] - assert state.context["loaded_skills"] == ["ffuf", "xss"] - finally: - instances.clear() - instances.update(original_instances) - - -def test_load_skill_uses_same_plain_skill_format_as_create_agent() -> None: - instances = agents_graph_actions.__dict__["_agent_instances"] - original_instances = dict(instances) - try: - state = _DummyAgentState("agent_test_load_skill_short_name") - instances.clear() - instances[state.agent_id] = _DummyAgent() - - result = load_skill_actions.load_skill(state, "nmap") - - assert result["success"] is True - assert result["loaded_skills"] == ["nmap"] - assert result["newly_loaded_skills"] == ["nmap"] - assert state.context["loaded_skills"] == ["nmap"] - finally: - instances.clear() - instances.update(original_instances) - - -def test_load_skill_invalid_skill_returns_error() -> None: - instances = agents_graph_actions.__dict__["_agent_instances"] - original_instances = dict(instances) - try: - state = _DummyAgentState("agent_test_load_skill_invalid") - instances.clear() - instances[state.agent_id] = _DummyAgent() - - result = load_skill_actions.load_skill(state, "definitely_not_a_real_skill") - - assert result["success"] is False - assert "Invalid skills" in result["error"] - assert "Available skills" in result["error"] - finally: - instances.clear() - instances.update(original_instances) - - -def test_load_skill_rejects_more_than_five_skills() -> None: - instances = agents_graph_actions.__dict__["_agent_instances"] - original_instances = dict(instances) - try: - state = _DummyAgentState("agent_test_load_skill_too_many") - instances.clear() - instances[state.agent_id] = _DummyAgent() - - result = load_skill_actions.load_skill(state, "a,b,c,d,e,f") - - assert result["success"] is False - assert result["error"] == ( - "Cannot specify more than 5 skills for an agent (use comma-separated format)" - ) - finally: - instances.clear() - instances.update(original_instances) - - -def test_load_skill_missing_agent_instance_returns_error() -> None: - instances = agents_graph_actions.__dict__["_agent_instances"] - original_instances = dict(instances) - try: - state = _DummyAgentState("agent_test_load_skill_missing_instance") - instances.clear() - - result = load_skill_actions.load_skill(state, "httpx") - - assert result["success"] is False - assert "running agent instance" in result["error"] - finally: - instances.clear() - instances.update(original_instances) - - -def test_load_skill_does_not_reload_skill_already_present_from_agent_creation() -> None: - instances = agents_graph_actions.__dict__["_agent_instances"] - original_instances = dict(instances) - try: - state = _DummyAgentState("agent_test_load_skill_existing_config_skill") - instances.clear() - instances[state.agent_id] = _DummyAgent(["xss"]) - - result = load_skill_actions.load_skill(state, "xss,sql_injection") - - assert result["success"] is True - assert result["loaded_skills"] == ["xss", "sql_injection"] - assert result["newly_loaded_skills"] == ["sql_injection"] - assert result["already_loaded_skills"] == ["xss"] - assert state.context["loaded_skills"] == ["sql_injection", "xss"] - finally: - instances.clear() - instances.update(original_instances) diff --git a/tests/tools/test_sdk_local_tools.py b/tests/tools/test_local_tools.py similarity index 94% rename from tests/tools/test_sdk_local_tools.py rename to tests/tools/test_local_tools.py index 15892a3..f4259b6 100644 --- a/tests/tools/test_sdk_local_tools.py +++ b/tests/tools/test_local_tools.py @@ -20,16 +20,16 @@ from unittest.mock import patch import pytest from agents.tool import FunctionTool -from strix.tools.notes import notes_actions as _notes_legacy -from strix.tools.notes.notes_sdk_tools import ( +from strix.tools.notes import notes_actions as _notes_impl +from strix.tools.notes.tools import ( create_note, delete_note, get_note, list_notes, update_note, ) -from strix.tools.thinking.thinking_sdk_tools import think -from strix.tools.todo.todo_sdk_tools import ( +from strix.tools.thinking.tool import think +from strix.tools.todo.tools import ( create_todo, delete_todo, list_todos, @@ -190,10 +190,10 @@ def notes_run_dir(tmp_path: Path) -> Iterator[Path]: """Point the legacy notes module at a fresh run dir per test.""" run_dir = tmp_path / "strix_runs" / "test" run_dir.mkdir(parents=True) - _notes_legacy._notes_storage.clear() - _notes_legacy._loaded_notes_run_dir = None + _notes_impl._notes_storage.clear() + _notes_impl._loaded_notes_run_dir = None - with patch.object(_notes_legacy, "_get_run_dir", return_value=run_dir): + with patch.object(_notes_impl, "_get_run_dir", return_value=run_dir): yield run_dir diff --git a/tests/tools/test_sdk_remaining_local_tools.py b/tests/tools/test_remaining_local_tools.py similarity index 82% rename from tests/tools/test_sdk_remaining_local_tools.py rename to tests/tools/test_remaining_local_tools.py index 6949b9b..229e6fb 100644 --- a/tests/tools/test_sdk_remaining_local_tools.py +++ b/tests/tools/test_remaining_local_tools.py @@ -12,7 +12,6 @@ sandbox-bound tools route through ``post_to_sandbox``. from __future__ import annotations import json -from collections.abc import Iterator from dataclasses import dataclass, field from typing import Any from unittest.mock import patch @@ -20,15 +19,15 @@ from unittest.mock import patch import pytest from agents.tool import FunctionTool -from strix.tools.file_edit.file_edit_sdk_tools import ( +from strix.tools.file_edit.tools import ( list_files, search_files, str_replace_editor, ) -from strix.tools.finish.finish_sdk_tool import finish_scan -from strix.tools.load_skill.load_skill_sdk_tool import load_skill -from strix.tools.reporting.reporting_sdk_tools import create_vulnerability_report -from strix.tools.web_search.web_search_sdk_tool import web_search +from strix.tools.finish.tool import finish_scan +from strix.tools.load_skill.tool import load_skill +from strix.tools.reporting.tool import create_vulnerability_report +from strix.tools.web_search.tool import web_search @dataclass @@ -93,7 +92,7 @@ async def test_web_search_no_api_key_returns_structured_error( @pytest.mark.asyncio -async def test_web_search_delegates_to_legacy(monkeypatch: pytest.MonkeyPatch) -> None: +async def test_web_search_delegates_to_impl(monkeypatch: pytest.MonkeyPatch) -> None: """Legacy ``web_search`` returns dict; wrapper JSON-encodes it.""" monkeypatch.setenv("PERPLEXITY_API_KEY", "fake-key") @@ -104,7 +103,7 @@ async def test_web_search_delegates_to_legacy(monkeypatch: pytest.MonkeyPatch) - "message": "Web search completed successfully", } with patch( - "strix.tools.web_search.web_search_sdk_tool._legacy.web_search", + "strix.tools.web_search.tool._impl.web_search", return_value=fake_result, ) as legacy: out = await _invoke(web_search, _ctx_for(), query="xss techniques") @@ -121,7 +120,7 @@ async def test_str_replace_editor_routes_to_sandbox() -> None: """file_edit tools must POST to the in-sandbox tool server, not run locally.""" fake_response = {"result": {"content": "file viewed"}} with patch( - "strix.tools.file_edit.file_edit_sdk_tools.post_to_sandbox", + "strix.tools.file_edit.tools.post_to_sandbox", return_value=fake_response, ) as dispatch: out = await _invoke( @@ -147,7 +146,7 @@ async def test_str_replace_editor_routes_to_sandbox() -> None: async def test_list_files_routes_to_sandbox() -> None: fake_response = {"result": {"files": ["a.py"], "directories": []}} with patch( - "strix.tools.file_edit.file_edit_sdk_tools.post_to_sandbox", + "strix.tools.file_edit.tools.post_to_sandbox", return_value=fake_response, ) as dispatch: out = await _invoke(list_files, _ctx_for(), path="src", recursive=True) @@ -162,7 +161,7 @@ async def test_list_files_routes_to_sandbox() -> None: async def test_search_files_routes_to_sandbox() -> None: fake_response = {"result": {"output": "src/foo.py:1:match"}} with patch( - "strix.tools.file_edit.file_edit_sdk_tools.post_to_sandbox", + "strix.tools.file_edit.tools.post_to_sandbox", return_value=fake_response, ) as dispatch: out = await _invoke( @@ -204,7 +203,7 @@ async def test_create_vulnerability_report_validates_required_fields() -> None: @pytest.mark.asyncio -async def test_create_vulnerability_report_delegates_to_legacy() -> None: +async def test_create_vulnerability_report_delegates_to_impl() -> None: """Verify the wrapper passes all params through to the legacy function.""" fake_result = { "success": True, @@ -214,7 +213,7 @@ async def test_create_vulnerability_report_delegates_to_legacy() -> None: "cvss_score": 7.5, } with patch( - "strix.tools.reporting.reporting_sdk_tools._legacy.create_vulnerability_report", + "strix.tools.reporting.tool._impl.create_vulnerability_report", return_value=fake_result, ) as legacy: out = await _invoke( @@ -255,7 +254,7 @@ async def test_load_skill_passes_adapter_with_agent_id() -> None: return {"success": True, "loaded_skills": ["recon"]} with patch( - "strix.tools.load_skill.load_skill_sdk_tool._legacy.load_skill", + "strix.tools.load_skill.tool._impl.load_skill", side_effect=fake_legacy, ): out = await _invoke(load_skill, _ctx_for("agent-XYZ"), skills="recon") @@ -276,28 +275,7 @@ async def test_load_skill_with_empty_input() -> None: # --- finish_scan --------------------------------------------------------- -@pytest.fixture -def isolated_agent_graph() -> Iterator[None]: - """Clear the legacy agent-graph globals so finish_scan sees an empty world. - - The legacy ``_check_active_agents`` reads ``_agent_graph["nodes"]`` and - returns an "agents still active" error if any non-self agent is in - state ``running`` or ``stopping``. Tests in other modules (legacy - multi-agent tests) populate this dict; without isolation they bleed - into our validation tests and mask the field-validation path. - """ - from strix.tools.agents_graph import agents_graph_actions - - saved_nodes = agents_graph_actions._agent_graph.get("nodes", {}).copy() - agents_graph_actions._agent_graph["nodes"] = {} - try: - yield - finally: - agents_graph_actions._agent_graph["nodes"] = saved_nodes - - @pytest.mark.asyncio -@pytest.mark.usefixtures("isolated_agent_graph") async def test_finish_scan_validates_empty_fields() -> None: """Legacy validation: every section must be non-empty.""" out = await _invoke( @@ -313,8 +291,7 @@ async def test_finish_scan_validates_empty_fields() -> None: @pytest.mark.asyncio -@pytest.mark.usefixtures("isolated_agent_graph") -async def test_finish_scan_delegates_to_legacy() -> None: +async def test_finish_scan_delegates_to_impl() -> None: """Wrapper must pass the legacy adapter and the four sections through.""" fake_result = { "success": True, @@ -323,7 +300,7 @@ async def test_finish_scan_delegates_to_legacy() -> None: "vulnerabilities_found": 3, } with patch( - "strix.tools.finish.finish_sdk_tool._legacy.finish_scan", + "strix.tools.finish.tool._impl.finish_scan", return_value=fake_result, ) as legacy: out = await _invoke( diff --git a/tests/tools/test_sdk_sandbox_tools.py b/tests/tools/test_sandbox_tools.py similarity index 91% rename from tests/tools/test_sdk_sandbox_tools.py rename to tests/tools/test_sandbox_tools.py index 9c9c1d7..f8990d5 100644 --- a/tests/tools/test_sdk_sandbox_tools.py +++ b/tests/tools/test_sandbox_tools.py @@ -27,8 +27,8 @@ from unittest.mock import patch import pytest from agents.tool import FunctionTool -from strix.tools.browser.browser_sdk_tool import browser_action -from strix.tools.proxy.proxy_sdk_tools import ( +from strix.tools.browser.tool import browser_action +from strix.tools.proxy.tools import ( list_requests, list_sitemap, repeat_request, @@ -37,8 +37,8 @@ from strix.tools.proxy.proxy_sdk_tools import ( view_request, view_sitemap_entry, ) -from strix.tools.python.python_sdk_tool import python_action -from strix.tools.terminal.terminal_sdk_tool import terminal_execute +from strix.tools.python.tool import python_action +from strix.tools.terminal.tool import terminal_execute _ALL_SANDBOX_TOOLS = ( @@ -119,7 +119,7 @@ async def test_browser_action_dispatches_full_payload() -> None: (the in-container handler distinguishes ``None`` from missing).""" fake = {"result": {"screenshot": "data:image/png;base64,..."}} with patch( - "strix.tools.browser.browser_sdk_tool.post_to_sandbox", + "strix.tools.browser.tool.post_to_sandbox", return_value=fake, ) as dispatch: out = await _invoke( @@ -158,7 +158,7 @@ async def test_browser_action_dispatches_full_payload() -> None: async def test_terminal_execute_dispatches() -> None: fake = {"result": {"content": "hello\n", "exit_code": 0}} with patch( - "strix.tools.terminal.terminal_sdk_tool.post_to_sandbox", + "strix.tools.terminal.tool.post_to_sandbox", return_value=fake, ) as dispatch: out = await _invoke( @@ -185,7 +185,7 @@ async def test_terminal_execute_dispatches() -> None: async def test_python_action_dispatches() -> None: fake = {"result": {"stdout": "42\n", "is_running": False}} with patch( - "strix.tools.python.python_sdk_tool.post_to_sandbox", + "strix.tools.python.tool.post_to_sandbox", return_value=fake, ) as dispatch: out = await _invoke( @@ -214,7 +214,7 @@ async def test_python_action_dispatches() -> None: async def test_list_requests_forwards_full_query() -> None: fake: dict[str, Any] = {"result": {"requests": []}} with patch( - "strix.tools.proxy.proxy_sdk_tools.post_to_sandbox", + "strix.tools.proxy.tools.post_to_sandbox", return_value=fake, ) as dispatch: await _invoke( @@ -241,7 +241,7 @@ async def test_list_requests_forwards_full_query() -> None: async def test_view_request_dispatches() -> None: fake = {"result": {"raw": "GET / HTTP/1.1..."}} with patch( - "strix.tools.proxy.proxy_sdk_tools.post_to_sandbox", + "strix.tools.proxy.tools.post_to_sandbox", return_value=fake, ) as dispatch: await _invoke( @@ -264,7 +264,7 @@ async def test_send_request_normalizes_missing_headers() -> None: """Legacy schema treats omitted ``headers`` as ``{}``; the wrapper must too.""" fake = {"result": {"status": 200}} with patch( - "strix.tools.proxy.proxy_sdk_tools.post_to_sandbox", + "strix.tools.proxy.tools.post_to_sandbox", return_value=fake, ) as dispatch: await _invoke( @@ -286,7 +286,7 @@ async def test_send_request_normalizes_missing_headers() -> None: async def test_repeat_request_normalizes_missing_modifications() -> None: fake = {"result": {"status": 200}} with patch( - "strix.tools.proxy.proxy_sdk_tools.post_to_sandbox", + "strix.tools.proxy.tools.post_to_sandbox", return_value=fake, ) as dispatch: await _invoke(repeat_request, _ctx_for(), request_id="req-1") @@ -300,7 +300,7 @@ async def test_repeat_request_normalizes_missing_modifications() -> None: async def test_scope_rules_dispatches() -> None: fake = {"result": {"scope_id": "s-1"}} with patch( - "strix.tools.proxy.proxy_sdk_tools.post_to_sandbox", + "strix.tools.proxy.tools.post_to_sandbox", return_value=fake, ) as dispatch: await _invoke( @@ -323,7 +323,7 @@ async def test_scope_rules_dispatches() -> None: async def test_list_sitemap_defaults() -> None: fake: dict[str, Any] = {"result": {"entries": []}} with patch( - "strix.tools.proxy.proxy_sdk_tools.post_to_sandbox", + "strix.tools.proxy.tools.post_to_sandbox", return_value=fake, ) as dispatch: await _invoke(list_sitemap, _ctx_for()) @@ -338,7 +338,7 @@ async def test_list_sitemap_defaults() -> None: async def test_view_sitemap_entry_dispatches() -> None: fake = {"result": {"entry_id": "e-1"}} with patch( - "strix.tools.proxy.proxy_sdk_tools.post_to_sandbox", + "strix.tools.proxy.tools.post_to_sandbox", return_value=fake, ) as dispatch: await _invoke(view_sitemap_entry, _ctx_for(), entry_id="e-1") diff --git a/tests/tools/test_tool_registration_modes.py b/tests/tools/test_tool_registration_modes.py index b50d267..ce15a46 100644 --- a/tests/tools/test_tool_registration_modes.py +++ b/tests/tools/test_tool_registration_modes.py @@ -21,9 +21,15 @@ def _reload_tools_module() -> ModuleType: return importlib.import_module("strix.tools") -def test_non_sandbox_registers_agents_graph_but_not_browser_or_web_search_when_disabled( +def test_non_sandbox_skips_browser_and_web_search_when_disabled( monkeypatch: Any, ) -> None: + """Browser registration is gated on STRIX_DISABLE_BROWSER and + web_search on PERPLEXITY_API_KEY; both should stay out of the + in-container ``register_tool`` registry when their gates are off. + Agents_graph is no longer in this registry — those tools are SDK + function tools (host-side only), not in-container tools. + """ monkeypatch.setenv("STRIX_SANDBOX_MODE", "false") monkeypatch.setenv("STRIX_DISABLE_BROWSER", "true") monkeypatch.delenv("PERPLEXITY_API_KEY", raising=False) @@ -32,7 +38,6 @@ def test_non_sandbox_registers_agents_graph_but_not_browser_or_web_search_when_d tools = _reload_tools_module() names = set(tools.get_tool_names()) - assert "create_agent" in names assert "browser_action" not in names assert "web_search" not in names