strix-agent 0.4.0__py3-none-any.whl → 0.6.2__py3-none-any.whl

strix/llm/memory_compressor.py

@@ -1,9 +1,10 @@
 import logging
-import os
 from typing import Any
 
 import litellm
 
+from strix.config import Config
+
 
 logger = logging.getLogger(__name__)
 
@@ -85,7 +86,7 @@ def _extract_message_text(msg: dict[str, Any]) -> str:
 def _summarize_messages(
     messages: list[dict[str, Any]],
     model: str,
-    timeout: int = 600,
+    timeout: int = 30,
 ) -> dict[str, Any]:
     if not messages:
         empty_summary = "<context_summary message_count='0'>{text}</context_summary>"
@@ -147,11 +148,11 @@ class MemoryCompressor:
         self,
         max_images: int = 3,
         model_name: str | None = None,
-        timeout: int = 600,
+        timeout: int | None = None,
     ):
         self.max_images = max_images
-        self.model_name = model_name or os.getenv("STRIX_LLM", "openai/gpt-5")
-        self.timeout = timeout
+        self.model_name = model_name or Config.get("strix_llm")
+        self.timeout = timeout or int(Config.get("strix_memory_compressor_timeout") or "30")
 
         if not self.model_name:
             raise ValueError("STRIX_LLM environment variable must be set and not empty")

strix/llm/utils.py

@@ -18,7 +18,7 @@ def _truncate_to_first_function(content: str) -> str:
 
 
 def parse_tool_invocations(content: str) -> list[dict[str, Any]] | None:
-    content = _fix_stopword(content)
+    content = fix_incomplete_tool_call(content)
 
     tool_invocations: list[dict[str, Any]] = []
 
@@ -46,12 +46,15 @@ def parse_tool_invocations(content: str) -> list[dict[str, Any]] | None:
     return tool_invocations if tool_invocations else None
 
 
-def _fix_stopword(content: str) -> str:
-    if "<function=" in content and content.count("<function=") == 1:
-        if content.endswith("</"):
-            content = content.rstrip() + "function>"
-        elif not content.rstrip().endswith("</function>"):
-            content = content + "\n</function>"
+def fix_incomplete_tool_call(content: str) -> str:
+    """Fix incomplete tool calls by adding missing </function> tag."""
+    if (
+        "<function=" in content
+        and content.count("<function=") == 1
+        and "</function>" not in content
+    ):
+        content = content.rstrip()
+        content = content + "function>" if content.endswith("</") else content + "\n</function>"
     return content
 
 
@@ -70,11 +73,17 @@ def clean_content(content: str) -> str:
     if not content:
         return ""
 
-    content = _fix_stopword(content)
+    content = fix_incomplete_tool_call(content)
 
     tool_pattern = r"<function=[^>]+>.*?</function>"
     cleaned = re.sub(tool_pattern, "", content, flags=re.DOTALL)
 
+    incomplete_tool_pattern = r"<function=[^>]+>.*$"
+    cleaned = re.sub(incomplete_tool_pattern, "", cleaned, flags=re.DOTALL)
+
+    partial_tag_pattern = r"<f(?:u(?:n(?:c(?:t(?:i(?:o(?:n(?:=(?:[^>]*)?)?)?)?)?)?)?)?)?$"
+    cleaned = re.sub(partial_tag_pattern, "", cleaned)
+
     hidden_xml_patterns = [
         r"<inter_agent_message>.*?</inter_agent_message>",
         r"<agent_completion_report>.*?</agent_completion_report>",

strix/runtime/__init__.py

@@ -1,10 +1,19 @@
-import os
+from strix.config import Config
 
 from .runtime import AbstractRuntime
 
 
+class SandboxInitializationError(Exception):
+    """Raised when sandbox initialization fails (e.g., Docker issues)."""
+
+    def __init__(self, message: str, details: str | None = None):
+        super().__init__(message)
+        self.message = message
+        self.details = details
+
+
 def get_runtime() -> AbstractRuntime:
-    runtime_backend = os.getenv("STRIX_RUNTIME_BACKEND", "docker")
+    runtime_backend = Config.get("strix_runtime_backend")
 
     if runtime_backend == "docker":
         from .docker_runtime import DockerRuntime
@@ -16,4 +25,4 @@ def get_runtime() -> AbstractRuntime:
     )
 
 
-__all__ = ["AbstractRuntime", "get_runtime"]
+__all__ = ["AbstractRuntime", "SandboxInitializationError", "get_runtime"]

strix/runtime/docker_runtime.py

@@ -1,5 +1,4 @@
 import contextlib
-import logging
 import os
 import secrets
 import socket
@@ -8,31 +7,37 @@ 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
 
 
-STRIX_IMAGE = os.getenv("STRIX_IMAGE", "ghcr.io/usestrix/strix-sandbox:0.1.10")
-logger = logging.getLogger(__name__)
+HOST_GATEWAY_HOSTNAME = "host.docker.internal"
+DOCKER_TIMEOUT = 60
+CONTAINER_TOOL_SERVER_PORT = 48081
 
 
 class DockerRuntime(AbstractRuntime):
     def __init__(self) -> None:
         try:
-            self.client = docker.from_env()
-        except DockerException as e:
-            logger.exception("Failed to connect to Docker daemon")
-            raise RuntimeError("Docker is not available or not configured correctly.") from e
+            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
 
-    def _generate_sandbox_token(self) -> str:
-        return secrets.token_urlsafe(32)
-
     def _find_available_port(self) -> int:
         with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
             s.bind(("", 0))
@@ -45,113 +50,118 @@ class DockerRuntime(AbstractRuntime):
             tracer = get_global_tracer()
             if tracer and tracer.scan_config:
                 return str(tracer.scan_config.get("scan_id", "default-scan"))
-        except ImportError:
-            logger.debug("Failed to import tracer, using fallback scan ID")
-        except AttributeError:
-            logger.debug("Tracer missing scan_config, using fallback scan ID")
-
+        except (ImportError, AttributeError):
+            pass
         return f"scan-{agent_id.split('-')[0]}"
 
     def _verify_image_available(self, image_name: str, max_retries: int = 3) -> None:
-        def _validate_image(image: docker.models.images.Image) -> None:
-            if not image.id or not image.attrs:
-                raise ImageNotFound(f"Image {image_name} metadata incomplete")
-
         for attempt in range(max_retries):
             try:
                 image = self.client.images.get(image_name)
-                _validate_image(image)
-            except ImageNotFound:
+                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:
-                    logger.exception(f"Image {image_name} not found after {max_retries} attempts")
                     raise
-                logger.warning(f"Image {image_name} not ready, attempt {attempt + 1}/{max_retries}")
-                time.sleep(2**attempt)
-            except DockerException:
-                if attempt == max_retries - 1:
-                    logger.exception(f"Failed to verify image {image_name}")
-                    raise
-                logger.warning(f"Docker error verifying image, attempt {attempt + 1}/{max_retries}")
                 time.sleep(2**attempt)
             else:
-                logger.debug(f"Image {image_name} verified as available")
                 return
 
-    def _create_container_with_retry(self, scan_id: str, max_retries: int = 3) -> Container:
-        last_exception = None
-        container_name = f"strix-scan-{scan_id}"
+    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"])
+
+    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:
-                self._verify_image_available(STRIX_IMAGE)
+                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.",
+        )
 
-                try:
-                    existing_container = self.client.containers.get(container_name)
-                    logger.warning(f"Container {container_name} already exists, removing it")
+    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_container.stop(timeout=5)
-                    existing_container.remove(force=True)
+                        existing.stop(timeout=5)
+                    existing.remove(force=True)
                     time.sleep(1)
-                except NotFound:
-                    pass
-                except DockerException as e:
-                    logger.warning(f"Error checking/removing existing container: {e}")
-
-                caido_port = self._find_available_port()
-                tool_server_port = self._find_available_port()
-                tool_server_token = self._generate_sandbox_token()
 
-                self._tool_server_port = tool_server_port
-                self._tool_server_token = tool_server_token
+                self._tool_server_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(
-                    STRIX_IMAGE,
+                    image_name,
                     command="sleep infinity",
                     detach=True,
                     name=container_name,
-                    hostname=f"strix-scan-{scan_id}",
-                    ports={
-                        f"{caido_port}/tcp": caido_port,
-                        f"{tool_server_port}/tcp": tool_server_port,
-                    },
+                    hostname=container_name,
+                    ports={f"{CONTAINER_TOOL_SERVER_PORT}/tcp": self._tool_server_port},
                     cap_add=["NET_ADMIN", "NET_RAW"],
                     labels={"strix-scan-id": scan_id},
                     environment={
                         "PYTHONUNBUFFERED": "1",
-                        "CAIDO_PORT": str(caido_port),
-                        "TOOL_SERVER_PORT": str(tool_server_port),
-                        "TOOL_SERVER_TOKEN": tool_server_token,
+                        "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
-                logger.info("Created container %s for scan %s", container.id, scan_id)
-
-                self._initialize_container(
-                    container, caido_port, tool_server_port, tool_server_token
-                )
-            except DockerException as e:
-                last_exception = e
-                if attempt == max_retries - 1:
-                    logger.exception(f"Failed to create container after {max_retries} attempts")
-                    break
-
-                logger.warning(f"Container creation attempt {attempt + 1}/{max_retries} failed")
-
-                self._tool_server_port = None
-                self._tool_server_token = None
-
-                sleep_time = (2**attempt) + (0.1 * attempt)
-                time.sleep(sleep_time)
+                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
+                    time.sleep(2**attempt)
             else:
                 return container
 
-        raise RuntimeError(
-            f"Failed to create Docker container after {max_retries} attempts: {last_exception}"
-        ) from last_exception
+        raise SandboxInitializationError(
+            "Failed to create container",
+            f"Container creation failed after {max_retries + 1} attempts: {last_error}",
+        ) from last_error
 
-    def _get_or_create_scan_container(self, scan_id: str) -> Container:  # noqa: PLR0912
+    def _get_or_create_container(self, scan_id: str) -> Container:
         container_name = f"strix-scan-{scan_id}"
 
         if self._scan_container:
@@ -168,33 +178,14 @@ class DockerRuntime(AbstractRuntime):
             container = self.client.containers.get(container_name)
             container.reload()
 
-            if (
-                "strix-scan-id" not in container.labels
-                or container.labels["strix-scan-id"] != scan_id
-            ):
-                logger.warning(
-                    f"Container {container_name} exists but missing/wrong label, updating"
-                )
-
             if container.status != "running":
-                logger.info(f"Starting existing container {container_name}")
                 container.start()
                 time.sleep(2)
 
             self._scan_container = container
-
-            for env_var in container.attrs["Config"]["Env"]:
-                if env_var.startswith("TOOL_SERVER_PORT="):
-                    self._tool_server_port = int(env_var.split("=")[1])
-                elif env_var.startswith("TOOL_SERVER_TOKEN="):
-                    self._tool_server_token = env_var.split("=")[1]
-
-            logger.info(f"Reusing existing container {container_name}")
-
+            self._recover_container_state(container)
         except NotFound:
             pass
-        except DockerException as e:
-            logger.warning(f"Failed to get container by name {container_name}: {e}")
         else:
             return container
 
@@ -203,52 +194,18 @@ class DockerRuntime(AbstractRuntime):
                 all=True, filters={"label": f"strix-scan-id={scan_id}"}
             )
             if containers:
-                container = cast("Container", containers[0])
+                container = containers[0]
                 if container.status != "running":
                     container.start()
                     time.sleep(2)
-                self._scan_container = container
-
-                for env_var in container.attrs["Config"]["Env"]:
-                    if env_var.startswith("TOOL_SERVER_PORT="):
-                        self._tool_server_port = int(env_var.split("=")[1])
-                    elif env_var.startswith("TOOL_SERVER_TOKEN="):
-                        self._tool_server_token = env_var.split("=")[1]
 
-                logger.info(f"Found existing container by label for scan {scan_id}")
+                self._scan_container = container
+                self._recover_container_state(container)
                 return container
-        except DockerException as e:
-            logger.warning("Failed to find existing container by label for scan %s: %s", scan_id, e)
-
-        logger.info("Creating new Docker container for scan %s", scan_id)
-        return self._create_container_with_retry(scan_id)
-
-    def _initialize_container(
-        self, container: Container, caido_port: int, tool_server_port: int, tool_server_token: str
-    ) -> None:
-        logger.info("Initializing Caido proxy on port %s", caido_port)
-        result = container.exec_run(
-            f"bash -c 'export CAIDO_PORT={caido_port} && /usr/local/bin/docker-entrypoint.sh true'",
-            detach=False,
-        )
-
-        time.sleep(5)
-
-        result = container.exec_run(
-            "bash -c 'source /etc/profile.d/proxy.sh && echo $CAIDO_API_TOKEN'", user="pentester"
-        )
-        caido_token = result.output.decode().strip() if result.exit_code == 0 else ""
-
-        container.exec_run(
-            f"bash -c 'source /etc/profile.d/proxy.sh && cd /app && "
-            f"STRIX_SANDBOX_MODE=true CAIDO_API_TOKEN={caido_token} CAIDO_PORT={caido_port} "
-            f"poetry run python strix/runtime/tool_server.py --token {tool_server_token} "
-            f"--host 0.0.0.0 --port {tool_server_port} &'",
-            detach=True,
-            user="pentester",
-        )
+        except DockerException:
+            pass
 
-        time.sleep(5)
+        return self._create_container(scan_id)
 
     def _copy_local_directory_to_container(
         self, container: Container, local_path: str, target_name: str | None = None
@@ -259,17 +216,8 @@ class DockerRuntime(AbstractRuntime):
         try:
             local_path_obj = Path(local_path).resolve()
             if not local_path_obj.exists() or not local_path_obj.is_dir():
-                logger.warning(f"Local path does not exist or is not directory: {local_path_obj}")
                 return
 
-            if target_name:
-                logger.info(
-                    f"Copying local directory {local_path_obj} to container at "
-                    f"/workspace/{target_name}"
-                )
-            else:
-                logger.info(f"Copying local directory {local_path_obj} to container")
-
             tar_buffer = BytesIO()
             with tarfile.open(fileobj=tar_buffer, mode="w") as tar:
                 for item in local_path_obj.rglob("*"):
@@ -280,16 +228,12 @@ class DockerRuntime(AbstractRuntime):
 
             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",
             )
-
-            logger.info("Successfully copied local directory to /workspace")
-
         except (OSError, DockerException):
-            logger.exception("Failed to copy local directory to container")
+            pass
 
     async def create_sandbox(
         self,
@@ -298,7 +242,7 @@ class DockerRuntime(AbstractRuntime):
         local_sources: list[dict[str, str]] | None = None,
     ) -> SandboxInfo:
         scan_id = self._get_scan_id(agent_id)
-        container = self._get_or_create_scan_container(scan_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):
@@ -306,40 +250,33 @@ class DockerRuntime(AbstractRuntime):
                 source_path = source.get("source_path")
                 if not source_path:
                     continue
-
-                target_name = source.get("workspace_subdir")
-                if not target_name:
-                    target_name = Path(source_path).name or f"target_{index}"
-
+                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)
 
-        container_id = container.id
-        if container_id is None:
+        if container.id is None:
             raise RuntimeError("Docker container ID is unexpectedly None")
 
-        token = existing_token if existing_token is not None else self._tool_server_token
-
+        token = existing_token or self._tool_server_token
         if self._tool_server_port is None or token is None:
-            raise RuntimeError("Tool server not initialized or no token available")
+            raise RuntimeError("Tool server not initialized")
 
-        api_url = await self.get_sandbox_url(container_id, self._tool_server_port)
+        host = self._resolve_docker_host()
+        api_url = f"http://{host}:{self._tool_server_port}"
 
-        await self._register_agent_with_tool_server(api_url, agent_id, token)
+        await self._register_agent(api_url, agent_id, token)
 
         return {
-            "workspace_id": container_id,
+            "workspace_id": container.id,
             "api_url": api_url,
             "auth_token": token,
             "tool_server_port": self._tool_server_port,
             "agent_id": agent_id,
         }
 
-    async def _register_agent_with_tool_server(
-        self, api_url: str, agent_id: str, token: str
-    ) -> None:
-        import httpx
-
+    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(
@@ -349,51 +286,33 @@ class DockerRuntime(AbstractRuntime):
                     timeout=30,
                 )
                 response.raise_for_status()
-                logger.info(f"Registered agent {agent_id} with tool server")
-        except (httpx.RequestError, httpx.HTTPStatusError) as e:
-            logger.warning(f"Failed to register agent {agent_id}: {e}")
+        except httpx.RequestError:
+            pass
 
     async def get_sandbox_url(self, container_id: str, port: int) -> str:
         try:
-            container = self.client.containers.get(container_id)
-            container.reload()
-
-            host = self._resolve_docker_host()
-
+            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
-        except DockerException as e:
-            raise RuntimeError(f"Failed to get container URL for {container_id}: {e}") from e
-        else:
-            return f"http://{host}:{port}"
 
     def _resolve_docker_host(self) -> str:
         docker_host = os.getenv("DOCKER_HOST", "")
-        if not docker_host:
-            return "127.0.0.1"
-
-        from urllib.parse import urlparse
-
-        parsed = urlparse(docker_host)
-
-        if parsed.scheme in ("tcp", "http", "https") and parsed.hostname:
-            return parsed.hostname
+        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:
-        logger.info("Destroying scan container %s", container_id)
         try:
             container = self.client.containers.get(container_id)
             container.stop()
             container.remove()
-            logger.info("Successfully destroyed container %s", container_id)
-
             self._scan_container = None
             self._tool_server_port = None
             self._tool_server_token = None
-
-        except NotFound:
-            logger.warning("Container %s not found for destruction.", container_id)
-        except DockerException as e:
-            logger.warning("Failed to destroy container %s: %s", container_id, e)
+        except (NotFound, DockerException):
+            pass