strix-agent 0.1.8__py3-none-any.whl → 0.1.10__py3-none-any.whl

strix/runtime/tool_server.py

@@ -1,7 +1,15 @@
+from __future__ import annotations
+
+import argparse
+import asyncio
 import logging
 import os
+import signal
+import sys
+from multiprocessing import Process, Queue
 from typing import Any
 
+import uvicorn
 from fastapi import Depends, FastAPI, HTTPException, status
 from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
 from pydantic import BaseModel, ValidationError
@@ -11,20 +19,25 @@ SANDBOX_MODE = os.getenv("STRIX_SANDBOX_MODE", "false").lower() == "true"
 if not SANDBOX_MODE:
     raise RuntimeError("Tool server should only run in sandbox mode (STRIX_SANDBOX_MODE=true)")
 
-EXPECTED_TOKEN = os.getenv("STRIX_SANDBOX_TOKEN")
-if not EXPECTED_TOKEN:
-    raise RuntimeError("STRIX_SANDBOX_TOKEN environment variable is required in sandbox mode")
+parser = argparse.ArgumentParser(description="Start Strix tool server")
+parser.add_argument("--token", required=True, help="Authentication token")
+parser.add_argument("--host", default="0.0.0.0", help="Host to bind to")  # nosec
+parser.add_argument("--port", type=int, required=True, help="Port to bind to")
+
+args = parser.parse_args()
+EXPECTED_TOKEN = args.token
 
 app = FastAPI()
-logger = logging.getLogger(__name__)
 security = HTTPBearer()
 
 security_dependency = Depends(security)
 
+agent_processes: dict[str, dict[str, Any]] = {}
+agent_queues: dict[str, dict[str, Queue[Any]]] = {}
+
 
 def verify_token(credentials: HTTPAuthorizationCredentials) -> str:
     if not credentials or credentials.scheme != "Bearer":
-        logger.warning("Authentication failed: Invalid or missing Bearer token scheme")
         raise HTTPException(
             status_code=status.HTTP_401_UNAUTHORIZED,
             detail="Invalid authentication scheme. Bearer token required.",
@@ -32,18 +45,17 @@ def verify_token(credentials: HTTPAuthorizationCredentials) -> str:
         )
 
     if credentials.credentials != EXPECTED_TOKEN:
-        logger.warning("Authentication failed: Invalid token provided from remote host")
         raise HTTPException(
             status_code=status.HTTP_401_UNAUTHORIZED,
             detail="Invalid authentication token",
             headers={"WWW-Authenticate": "Bearer"},
         )
 
-    logger.debug("Authentication successful for tool execution request")
     return credentials.credentials
 
 
 class ToolExecutionRequest(BaseModel):
+    agent_id: str
     tool_name: str
     kwargs: dict[str, Any]
 
@@ -53,45 +65,141 @@ class ToolExecutionResponse(BaseModel):
     error: str | None = None
 
 
+def agent_worker(_agent_id: str, request_queue: Queue[Any], response_queue: Queue[Any]) -> None:
+    null_handler = logging.NullHandler()
+
+    root_logger = logging.getLogger()
+    root_logger.handlers = [null_handler]
+    root_logger.setLevel(logging.CRITICAL)
+
+    from strix.tools.argument_parser import ArgumentConversionError, convert_arguments
+    from strix.tools.registry import get_tool_by_name
+
+    while True:
+        try:
+            request = request_queue.get()
+
+            if request is None:
+                break
+
+            tool_name = request["tool_name"]
+            kwargs = request["kwargs"]
+
+            try:
+                tool_func = get_tool_by_name(tool_name)
+                if not tool_func:
+                    response_queue.put({"error": f"Tool '{tool_name}' not found"})
+                    continue
+
+                converted_kwargs = convert_arguments(tool_func, kwargs)
+                result = tool_func(**converted_kwargs)
+
+                response_queue.put({"result": result})
+
+            except (ArgumentConversionError, ValidationError) as e:
+                response_queue.put({"error": f"Invalid arguments: {e}"})
+            except (RuntimeError, ValueError, ImportError) as e:
+                response_queue.put({"error": f"Tool execution error: {e}"})
+
+        except (RuntimeError, ValueError, ImportError) as e:
+            response_queue.put({"error": f"Worker error: {e}"})
+
+
+def ensure_agent_process(agent_id: str) -> tuple[Queue[Any], Queue[Any]]:
+    if agent_id not in agent_processes:
+        request_queue: Queue[Any] = Queue()
+        response_queue: Queue[Any] = Queue()
+
+        process = Process(
+            target=agent_worker, args=(agent_id, request_queue, response_queue), daemon=True
+        )
+        process.start()
+
+        agent_processes[agent_id] = {"process": process, "pid": process.pid}
+        agent_queues[agent_id] = {"request": request_queue, "response": response_queue}
+
+    return agent_queues[agent_id]["request"], agent_queues[agent_id]["response"]
+
+
 @app.post("/execute", response_model=ToolExecutionResponse)
 async def execute_tool(
     request: ToolExecutionRequest, credentials: HTTPAuthorizationCredentials = security_dependency
 ) -> ToolExecutionResponse:
     verify_token(credentials)
 
-    from strix.tools.argument_parser import ArgumentConversionError, convert_arguments
-    from strix.tools.registry import get_tool_by_name
+    request_queue, response_queue = ensure_agent_process(request.agent_id)
+
+    request_queue.put({"tool_name": request.tool_name, "kwargs": request.kwargs})
 
     try:
-        tool_func = get_tool_by_name(request.tool_name)
-        if not tool_func:
-            return ToolExecutionResponse(error=f"Tool '{request.tool_name}' not found")
+        loop = asyncio.get_event_loop()
+        response = await loop.run_in_executor(None, response_queue.get)
 
-        converted_kwargs = convert_arguments(tool_func, request.kwargs)
+        if "error" in response:
+            return ToolExecutionResponse(error=response["error"])
+        return ToolExecutionResponse(result=response.get("result"))
 
-        result = tool_func(**converted_kwargs)
+    except (RuntimeError, ValueError, OSError) as e:
+        return ToolExecutionResponse(error=f"Worker error: {e}")
 
-        return ToolExecutionResponse(result=result)
 
-    except (ArgumentConversionError, ValidationError) as e:
-        logger.warning("Invalid tool arguments: %s", e)
-        return ToolExecutionResponse(error=f"Invalid arguments: {e}")
-    except TypeError as e:
-        logger.warning("Tool execution type error: %s", e)
-        return ToolExecutionResponse(error=f"Tool execution error: {e}")
-    except ValueError as e:
-        logger.warning("Tool execution value error: %s", e)
-        return ToolExecutionResponse(error=f"Tool execution error: {e}")
-    except Exception:
-        logger.exception("Unexpected error during tool execution")
-        return ToolExecutionResponse(error="Internal server error")
+@app.post("/register_agent")
+async def register_agent(
+    agent_id: str, credentials: HTTPAuthorizationCredentials = security_dependency
+) -> dict[str, str]:
+    verify_token(credentials)
+
+    ensure_agent_process(agent_id)
+    return {"status": "registered", "agent_id": agent_id}
 
 
 @app.get("/health")
-async def health_check() -> dict[str, str]:
+async def health_check() -> dict[str, Any]:
     return {
         "status": "healthy",
         "sandbox_mode": str(SANDBOX_MODE),
         "environment": "sandbox" if SANDBOX_MODE else "main",
         "auth_configured": "true" if EXPECTED_TOKEN else "false",
+        "active_agents": len(agent_processes),
+        "agents": list(agent_processes.keys()),
     }
+
+
+def cleanup_all_agents() -> None:
+    for agent_id in list(agent_processes.keys()):
+        try:
+            agent_queues[agent_id]["request"].put(None)
+            process = agent_processes[agent_id]["process"]
+
+            process.join(timeout=1)
+
+            if process.is_alive():
+                process.terminate()
+                process.join(timeout=1)
+
+            if process.is_alive():
+                process.kill()
+
+        except (BrokenPipeError, EOFError, OSError):
+            pass
+        except (RuntimeError, ValueError) as e:
+            logging.getLogger(__name__).debug(f"Error during agent cleanup: {e}")
+
+
+def signal_handler(_signum: int, _frame: Any) -> None:
+    signal.signal(signal.SIGPIPE, signal.SIG_IGN) if hasattr(signal, "SIGPIPE") else None
+    cleanup_all_agents()
+    sys.exit(0)
+
+
+if hasattr(signal, "SIGPIPE"):
+    signal.signal(signal.SIGPIPE, signal.SIG_IGN)
+
+signal.signal(signal.SIGTERM, signal_handler)
+signal.signal(signal.SIGINT, signal_handler)
+
+if __name__ == "__main__":
+    try:
+        uvicorn.run(app, host=args.host, port=args.port, log_level="info")
+    finally:
+        cleanup_all_agents()

strix/tools/agents_graph/agents_graph_actions.py

@@ -57,6 +57,10 @@ def _run_agent_in_thread(
         - 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
     </instructions>
 </agent_delegation>"""
 
@@ -192,16 +196,6 @@ def create_agent(
         if prompt_modules:
             module_list = [m.strip() for m in prompt_modules.split(",") if m.strip()]
 
-        if "root_agent" in module_list:
-            return {
-                "success": False,
-                "error": (
-                    "The 'root_agent' module is reserved for the main agent "
-                    "and cannot be used by sub-agents"
-                ),
-                "agent_id": None,
-            }
-
         if len(module_list) > 3:
             return {
                 "success": False,

strix/tools/agents_graph/agents_graph_actions_schema.xml

@@ -59,7 +59,7 @@ Use this tool when:
   <tool name="create_agent">
     <description>Create and spawn a new agent to handle a specific subtask.
 
-MANDATORY REQUIREMENT: You MUST call view_agent_graph FIRST before creating any new agent to check if there is already an agent working on the same or similar task. Only create a new agent if no existing agent is handling the specific task.</description>
+Only create a new agent if no existing agent is handling the specific task.</description>
     <details>The new agent inherits the parent's conversation history and context up to the point
   of creation, then continues with its assigned subtask. This enables decomposition
   of complex penetration testing tasks into specialized sub-agents.
@@ -67,12 +67,6 @@ MANDATORY REQUIREMENT: You MUST call view_agent_graph FIRST before creating any
   The agent runs asynchronously and independently, allowing the parent to continue
   immediately while the new agent executes its task in the background.
 
-  CRITICAL: Before calling this tool, you MUST first use view_agent_graph to:
-  - Examine all existing agents and their current tasks
-  - Verify no agent is already working on the same or similar objective
-  - Avoid duplication of effort and resource waste
-  - Ensure efficient coordination across the multi-agent system
-
   If you as a parent agent don't absolutely have anything to do while your subagents are running, you can use wait_for_message tool. The subagent will continue to run in the background, and update you when it's done.
   </details>
     <parameters>
@@ -86,16 +80,13 @@ MANDATORY REQUIREMENT: You MUST call view_agent_graph FIRST before creating any
         <description>Whether the new agent should inherit parent's conversation history and context</description>
       </parameter>
       <parameter name="prompt_modules" type="string" required="false">
-        <description>Comma-separated list of prompt modules to use for the agent. Most agents should have at least one module in order to be useful. {{DYNAMIC_MODULES_DESCRIPTION}}</description>
+        <description>Comma-separated list of prompt modules to use for the agent (MAXIMUM 3 modules allowed). Most agents should have at least one module in order to be useful. Agents should be highly specialized - use 1-3 related vulnerability modules only. {{DYNAMIC_MODULES_DESCRIPTION}}</description>
       </parameter>
     </parameters>
     <returns type="Dict[str, Any]">
       <description>Response containing: - agent_id: Unique identifier for the created agent - success: Whether the agent was created successfully - message: Status message - agent_info: Details about the created agent</description>
     </returns>
     <examples>
-  # REQUIRED: First check agent graph before creating any new agent
-  <function=view_agent_graph>
-  </function>
   # REQUIRED: Check agent graph again before creating another agent
   <function=view_agent_graph>
   </function>
@@ -108,12 +99,27 @@ MANDATORY REQUIREMENT: You MUST call view_agent_graph FIRST before creating any
   <parameter=prompt_modules>sql_injection</parameter>
   </function>
 
-  # Create specialized authentication testing agent with multiple modules (comma-separated)
   <function=create_agent>
   <parameter=task>Test authentication mechanisms, JWT implementation, and session management
               for security vulnerabilities and bypass techniques.</parameter>
   <parameter=name>Auth Specialist</parameter>
   <parameter=prompt_modules>authentication_jwt, business_logic</parameter>
+  </function>
+
+  # Example of single-module specialization (most focused)
+  <function=create_agent>
+  <parameter=task>Perform comprehensive XSS testing including reflected, stored, and DOM-based
+              variants across all identified input points.</parameter>
+  <parameter=name>XSS Specialist</parameter>
+  <parameter=prompt_modules>xss</parameter>
+  </function>
+
+  # Example of maximum 3 related modules (borderline acceptable)
+  <function=create_agent>
+  <parameter=task>Test for server-side vulnerabilities including SSRF, XXE, and potential
+              RCE vectors in file upload and XML processing endpoints.</parameter>
+  <parameter=name>Server-Side Attack Specialist</parameter>
+  <parameter=prompt_modules>ssrf, xxe, rce</parameter>
   </function>
     </examples>
   </tool>

strix/tools/argument_parser.py

@@ -1,6 +1,7 @@
 import contextlib
 import inspect
 import json
+import types
 from collections.abc import Callable
 from typing import Any, Union, get_args, get_origin
 
@@ -48,7 +49,7 @@ def convert_arguments(func: Callable[..., Any], kwargs: dict[str, Any]) -> dict[
 
 def convert_string_to_type(value: str, param_type: Any) -> Any:
     origin = get_origin(param_type)
-    if origin is Union or origin is type(str | None):
+    if origin is Union or isinstance(param_type, types.UnionType):
         args = get_args(param_type)
         for arg_type in args:
             if arg_type is not type(None):

strix/tools/executor.py

@@ -49,7 +49,10 @@ async def _execute_tool_in_sandbox(tool_name: str, agent_state: Any, **kwargs: A
     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,
     }

strix/tools/terminal/__init__.py

@@ -1,4 +1,4 @@
-from .terminal_actions import terminal_action
+from .terminal_actions import terminal_execute
 
 
-__all__ = ["terminal_action"]
+__all__ = ["terminal_execute"]

strix/tools/terminal/terminal_actions.py

@@ -1,53 +1,35 @@
-from typing import Any, Literal
+from typing import Any
 
 from strix.tools.registry import register_tool
 
 from .terminal_manager import get_terminal_manager
 
 
-TerminalAction = Literal["new_terminal", "send_input", "wait", "close"]
-
-
 @register_tool
-def terminal_action(
-    action: TerminalAction,
-    inputs: list[str] | None = None,
-    time: float | None = None,
+def terminal_execute(
+    command: str,
+    is_input: bool = False,
+    timeout: float | None = None,
     terminal_id: str | None = None,
+    no_enter: bool = False,
 ) -> dict[str, Any]:
-    def _validate_inputs(action_name: str, inputs: list[str] | None) -> None:
-        if not inputs:
-            raise ValueError(f"inputs parameter is required for {action_name} action")
-
-    def _validate_time(time_param: float | None) -> None:
-        if time_param is None:
-            raise ValueError("time parameter is required for wait action")
-
-    def _validate_action(action_name: str) -> None:
-        raise ValueError(f"Unknown action: {action_name}")
-
     manager = get_terminal_manager()
 
     try:
-        match action:
-            case "new_terminal":
-                return manager.create_terminal(terminal_id, inputs)
-
-            case "send_input":
-                _validate_inputs(action, inputs)
-                assert inputs is not None
-                return manager.send_input(terminal_id, inputs)
-
-            case "wait":
-                _validate_time(time)
-                assert time is not None
-                return manager.wait_terminal(terminal_id, time)
-
-            case "close":
-                return manager.close_terminal(terminal_id)
-
-            case _:
-                _validate_action(action)  # type: ignore[unreachable]
-
+        return manager.execute_command(
+            command=command,
+            is_input=is_input,
+            timeout=timeout,
+            terminal_id=terminal_id,
+            no_enter=no_enter,
+        )
     except (ValueError, RuntimeError) as e:
-        return {"error": str(e), "terminal_id": terminal_id, "snapshot": "", "is_running": False}
+        return {
+            "error": str(e),
+            "command": command,
+            "terminal_id": terminal_id or "default",
+            "content": "",
+            "status": "error",
+            "exit_code": None,
+            "working_dir": None,
+        }