diff --git a/README.md b/README.md index 351b96e7..0ae3e8bf 100644 --- a/README.md +++ b/README.md @@ -240,6 +240,11 @@ strix --target-list ./targets.txt See the [CLI reference](https://docs.strix.ai/usage/cli) for every option, including scan modes, diff scope, instruction files, and budgets. +### Interactive Commands + +While using the default interactive TUI, send `/compact` to the selected agent to force +provider-agnostic session compaction for long-running scans. `/compress` is supported as an alias. + ### Headless Mode Run Strix programmatically without interactive UI using the `-n/--non-interactive` flag - perfect for servers and automated jobs. The CLI prints real-time vulnerability findings and the final report before exiting. Exits with non-zero code when vulnerabilities are found. diff --git a/strix/core/agents.py b/strix/core/agents.py index 4d3d65cc..05031bd8 100644 --- a/strix/core/agents.py +++ b/strix/core/agents.py @@ -6,6 +6,7 @@ import asyncio import json import logging import tempfile +from collections.abc import Awaitable, Callable from dataclasses import dataclass, field from pathlib import Path from typing import TYPE_CHECKING, Any, Literal, cast @@ -14,8 +15,6 @@ from strix.core.sessions import session_write_lock if TYPE_CHECKING: - from collections.abc import Callable - from agents.items import TResponseInputItem from agents.memory import Session @@ -30,6 +29,11 @@ TERMINAL_STATUSES: frozenset[str] = frozenset({"completed", "stopped", "crashed" # position in the tree - decides whether waiting is bounded: only an agent waiting # on other agents is re-checked on a timer. WaitKind = Literal["user", "agents", "stalled"] +CompactionResult = Literal["success", "unavailable", "failed"] +COMPACTION_SUCCESS: CompactionResult = "success" +COMPACTION_UNAVAILABLE: CompactionResult = "unavailable" +COMPACTION_FAILED: CompactionResult = "failed" +CompactCallback = Callable[[], Awaitable[bool]] @dataclass(slots=True) @@ -45,6 +49,7 @@ class AgentRuntime: wake: asyncio.Event = field(default_factory=asyncio.Event) mailbox: list[dict[str, Any]] = field(default_factory=list) user_wake_required: bool = False + compact: CompactCallback | None = None class AgentCoordinator: @@ -181,6 +186,7 @@ class AgentCoordinator: session: Session | None = None, task: asyncio.Task[Any] | None = None, interrupt_on_message: bool | None = None, + compact: CompactCallback | None = None, resumable: bool | None = None, ) -> None: async with self._lock: @@ -191,6 +197,8 @@ class AgentCoordinator: runtime.task = task if interrupt_on_message is not None: runtime.interrupt_on_message = interrupt_on_message + if compact is not None: + runtime.compact = compact if resumable is not None: runtime.resumable = resumable @@ -340,6 +348,24 @@ class AgentCoordinator: await self._maybe_snapshot() return True + async def compact_agent_session(self, target_agent_id: str) -> CompactionResult: + """Run the compaction pipeline attached to one agent runtime.""" + async with self._lock: + runtime = self.runtimes.get(target_agent_id) + compact = runtime.compact if runtime is not None else None + if compact is None: + logger.warning( + "agent.compact dropped target=%s because its runtime is not attached", + target_agent_id, + ) + return COMPACTION_UNAVAILABLE + try: + compacted = await compact() + except Exception: + logger.exception("agent.compact failed target=%s", target_agent_id) + return COMPACTION_FAILED + return COMPACTION_SUCCESS if compacted else COMPACTION_UNAVAILABLE + async def wait_for_message(self, agent_id: str, *, timeout: float | None = None) -> bool: """Wait until a message is ready for ``agent_id``; False on ``timeout``.""" while True: diff --git a/strix/core/execution.py b/strix/core/execution.py index dfcd39fa..f5e04409 100644 --- a/strix/core/execution.py +++ b/strix/core/execution.py @@ -7,7 +7,7 @@ import contextlib import logging import uuid from collections.abc import Callable -from functools import cache +from functools import cache, partial from typing import TYPE_CHECKING, Any, cast from agents import RunConfig, Runner @@ -119,6 +119,16 @@ async def _compact_session( ) +async def _force_compact_session( + agent: Any, + session: Session | None, + run_config: RunConfig, +) -> bool: + if session is None: + return False + return await _compact_session(agent, session, run_config, force=True) + + _MAX_TRANSIENT_MODEL_RETRIES = 5 _TRANSIENT_MODEL_RETRY_BASE_DELAY_S = 2.0 _TRANSIENT_MODEL_RETRY_MAX_DELAY_S = 90.0 @@ -202,6 +212,7 @@ async def run_agent_loop( agent_id, session=session, interrupt_on_message=interactive, + compact=partial(_force_compact_session, agent, session, run_config), resumable=interactive, ) result: RunResultBase | None = None diff --git a/strix/interface/tui/backend/controller.py b/strix/interface/tui/backend/controller.py index 6f3b3fb3..07bdb28a 100644 --- a/strix/interface/tui/backend/controller.py +++ b/strix/interface/tui/backend/controller.py @@ -6,13 +6,14 @@ import asyncio import contextlib import math import webbrowser -from collections.abc import Awaitable, Callable +from collections.abc import Awaitable, Callable, Coroutine from pathlib import Path from typing import TYPE_CHECKING, Any from strix.config import load_settings from strix.config.models import is_recommended_or_frontier_model from strix.config.settings import DEFAULT_MAX_TURNS +from strix.core.agents import COMPACTION_FAILED, COMPACTION_SUCCESS from strix.interface.tui.backend.live_view import TuiLiveView from strix.interface.tui.backend.projection import ( MAX_TERMINAL_EVENTS, @@ -34,6 +35,14 @@ if TYPE_CHECKING: _STOPPABLE_AGENT_STATUSES = frozenset({"running", "waiting", "budget_paused"}) +_COMPACT_COMMANDS = frozenset({"/compact", "/compress"}) +_COMPACTION_SUCCESS_MESSAGE = "Context compaction complete." +_COMPACTION_UNAVAILABLE_MESSAGE = "Context compaction could not be completed for this agent." +_COMPACTION_FAILED_MESSAGE = "Context compaction failed. Please try again." +_COMPACTION_RESULT_MESSAGES = { + COMPACTION_SUCCESS: _COMPACTION_SUCCESS_MESSAGE, + COMPACTION_FAILED: _COMPACTION_FAILED_MESSAGE, +} ChangeCallback = Callable[[], None] StartCallback = Callable[[], Awaitable[None]] @@ -408,6 +417,16 @@ class TuiController: if self.scan_loop is None or self.scan_loop.is_closed(): raise RuntimeError("Scan loop is not ready") self.live_view.record_user_message(agent_id, message) + if message.strip().lower() in _COMPACT_COMMANDS: + compacted = await self._run_on_scan_loop( + self.coordinator.compact_agent_session(agent_id) + ) + feedback = _COMPACTION_RESULT_MESSAGES.get( + compacted, + _COMPACTION_UNAVAILABLE_MESSAGE, + ) + self.live_view.record_system_message(agent_id, feedback) + return {"compacted": compacted == COMPACTION_SUCCESS} if self.scan_loop is asyncio.get_running_loop(): delivered = await self.coordinator.send( agent_id, @@ -427,6 +446,17 @@ class TuiController: self.live_view.upsert_agent(agent_id, status="waiting", error_message=None) return {"sent": True} + async def _run_on_scan_loop(self, coroutine: Coroutine[Any, Any, Any]) -> Any: + scan_loop = self.scan_loop + if scan_loop is None: + raise RuntimeError("Scan loop is not ready") + if scan_loop is asyncio.get_running_loop(): + return await coroutine + future: asyncio.Future[Any] = asyncio.wrap_future( + asyncio.run_coroutine_threadsafe(coroutine, scan_loop) + ) + return await future + async def _stop_agent(self, payload: dict[str, Any]) -> dict[str, Any]: agent_id = self._required_string(payload, "agent_id") agent = self.live_view.agents.get(agent_id) diff --git a/strix/interface/tui/live_view.py b/strix/interface/tui/live_view.py index dbcf4028..16c8c768 100644 --- a/strix/interface/tui/live_view.py +++ b/strix/interface/tui/live_view.py @@ -207,6 +207,17 @@ class TuiLiveView: }, ) + def record_system_message(self, agent_id: str, content: str) -> None: + self._append_event( + agent_id, + "chat", + { + "role": "assistant", + "content": content, + "metadata": {"source": "tui_system"}, + }, + ) + def ingest_sdk_event(self, agent_id: str, event: Any) -> None: event_type = getattr(event, "type", "") if event_type == "raw_response_event": diff --git a/tests/test_execution.py b/tests/test_execution.py index 8fbf18ff..2953b53d 100644 --- a/tests/test_execution.py +++ b/tests/test_execution.py @@ -16,7 +16,12 @@ from agents.tool_context import ToolContext from openai.types.responses import ResponseOutputMessage, ResponseOutputRefusal from strix.core import execution -from strix.core.agents import AgentCoordinator +from strix.core.agents import ( + COMPACTION_FAILED, + COMPACTION_SUCCESS, + COMPACTION_UNAVAILABLE, + AgentCoordinator, +) from strix.core.execution import ( _notify_root_on_budget_reserve, notify_parent_on_terminal, @@ -1308,3 +1313,66 @@ async def test_autonomous_nudge_does_not_offer_the_user() -> None: ) assert "respond_to_user" not in items[0]["content"] + + +COMPACTION_AGENT_ID = "compact-agent" + + +@pytest.mark.asyncio +async def test_coordinator_runs_attached_manual_compactor() -> None: + compacted_agents: list[str] = [] + + async def compact() -> bool: + compacted_agents.append(COMPACTION_AGENT_ID) + return True + + coordinator = AgentCoordinator() + await coordinator.register(COMPACTION_AGENT_ID, "Compactor", parent_id=None) + await coordinator.attach_runtime(COMPACTION_AGENT_ID, compact=compact) + + assert await coordinator.compact_agent_session(COMPACTION_AGENT_ID) == COMPACTION_SUCCESS + assert compacted_agents == [COMPACTION_AGENT_ID] + + +@pytest.mark.asyncio +async def test_coordinator_rejects_manual_compaction_without_runtime() -> None: + coordinator = AgentCoordinator() + await coordinator.register(COMPACTION_AGENT_ID, "Compactor", parent_id=None) + + assert await coordinator.compact_agent_session(COMPACTION_AGENT_ID) == COMPACTION_UNAVAILABLE + + +@pytest.mark.asyncio +async def test_coordinator_reports_failed_manual_compactor() -> None: + async def compact() -> bool: + raise RuntimeError("compaction failed") + + coordinator = AgentCoordinator() + await coordinator.register(COMPACTION_AGENT_ID, "Compactor", parent_id=None) + await coordinator.attach_runtime(COMPACTION_AGENT_ID, compact=compact) + + assert await coordinator.compact_agent_session(COMPACTION_AGENT_ID) == COMPACTION_FAILED + + +@pytest.mark.asyncio +async def test_force_compact_session_rejects_missing_session() -> None: + assert await execution._force_compact_session(MagicMock(), None, MagicMock()) is False + + +@pytest.mark.asyncio +async def test_force_compact_session_uses_provider_pipeline( + monkeypatch: pytest.MonkeyPatch, +) -> None: + session = MagicMock() + calls: list[tuple[Any, Any, Any, bool]] = [] + + async def compact(agent: Any, attached: Any, config: Any, *, force: bool) -> bool: + calls.append((agent, attached, config, force)) + return True + + agent = MagicMock() + run_config = MagicMock() + monkeypatch.setattr(execution, "_compact_session", compact) + + assert await execution._force_compact_session(agent, session, run_config) is True + assert calls == [(agent, session, run_config, True)] diff --git a/tests/test_tui_backend_controller.py b/tests/test_tui_backend_controller.py index 3a28d020..1fd92db3 100644 --- a/tests/test_tui_backend_controller.py +++ b/tests/test_tui_backend_controller.py @@ -3,15 +3,26 @@ from __future__ import annotations import argparse import asyncio import os +import threading from pathlib import Path import pytest from strix.config import apply_config_override, loader from strix.config.settings import DEFAULT_MAX_TURNS +from strix.core.agents import COMPACTION_FAILED, COMPACTION_SUCCESS, COMPACTION_UNAVAILABLE from strix.interface.tui.backend.controller import TuiController +COMPACT_COMMAND = "/compact" +COMPRESS_COMMAND = "/compress" +COMPACTION_AGENT_ID = "compact-agent" +COMPACTION_SUCCESS_MESSAGE = "Context compaction complete." +COMPACTION_UNAVAILABLE_MESSAGE = "Context compaction could not be completed for this agent." +COMPACTION_FAILED_MESSAGE = "Context compaction failed. Please try again." +THREAD_JOIN_TIMEOUT_SECONDS = 1.0 + + class _SendingCoordinator: def __init__(self, delivered: bool = True) -> None: self.delivered = delivered @@ -521,6 +532,107 @@ async def test_stop_handles_coordinator_rejection_after_stale_active_projection( await controller.handle("agent.stop", {"agent_id": "agent-1"}) +@pytest.mark.asyncio +@pytest.mark.parametrize("command", [COMPACT_COMMAND, COMPRESS_COMMAND]) +async def test_compact_command_compacts_without_sending_message(command: str) -> None: + class Coordinator: + def __init__(self) -> None: + self.compacted: list[str] = [] + self.sent: list[tuple[str, dict[str, str]]] = [] + + async def compact_agent_session(self, agent_id: str) -> str: + self.compacted.append(agent_id) + return COMPACTION_SUCCESS + + async def send(self, agent_id: str, message: dict[str, str]) -> bool: + self.sent.append((agent_id, message)) + return True + + coordinator = Coordinator() + controller = TuiController(args(), coordinator=coordinator) + controller.set_runtime(scan_loop=asyncio.get_running_loop()) + + result = await controller.handle( + "agent.send_message", + {"agent_id": COMPACTION_AGENT_ID, "message": command}, + ) + + assert result == {"compacted": True} + assert coordinator.compacted == [COMPACTION_AGENT_ID] + assert coordinator.sent == [] + events = controller.live_view.events_for_agent(COMPACTION_AGENT_ID) + assert [event["data"]["content"] for event in events] == [ + command, + COMPACTION_SUCCESS_MESSAGE, + ] + + +@pytest.mark.asyncio +async def test_compact_command_reports_unavailable_compaction() -> None: + class Coordinator: + async def compact_agent_session(self, _agent_id: str) -> str: + return COMPACTION_UNAVAILABLE + + controller = TuiController(args(), coordinator=Coordinator()) + controller.set_runtime(scan_loop=asyncio.get_running_loop()) + + result = await controller.handle( + "agent.send_message", + {"agent_id": COMPACTION_AGENT_ID, "message": COMPACT_COMMAND}, + ) + + assert result == {"compacted": False} + events = controller.live_view.events_for_agent(COMPACTION_AGENT_ID) + assert events[-1]["data"]["content"] == COMPACTION_UNAVAILABLE_MESSAGE + + +@pytest.mark.asyncio +async def test_compact_command_reports_failed_compaction() -> None: + class Coordinator: + async def compact_agent_session(self, _agent_id: str) -> str: + return COMPACTION_FAILED + + controller = TuiController(args(), coordinator=Coordinator()) + controller.set_runtime(scan_loop=asyncio.get_running_loop()) + + result = await controller.handle( + "agent.send_message", + {"agent_id": COMPACTION_AGENT_ID, "message": COMPACT_COMMAND}, + ) + + assert result == {"compacted": False} + events = controller.live_view.events_for_agent(COMPACTION_AGENT_ID) + assert events[-1]["data"]["content"] == COMPACTION_FAILED_MESSAGE + + +@pytest.mark.asyncio +async def test_scan_loop_runner_rejects_missing_loop() -> None: + controller = TuiController(args()) + coroutine = asyncio.sleep(0) + try: + with pytest.raises(RuntimeError, match="Scan loop is not ready"): + await controller._run_on_scan_loop(coroutine) + finally: + coroutine.close() + + +@pytest.mark.asyncio +async def test_scan_loop_runner_dispatches_to_background_loop() -> None: + controller = TuiController(args()) + scan_loop = asyncio.new_event_loop() + thread = threading.Thread(target=scan_loop.run_forever) + thread.start() + try: + controller.set_runtime(scan_loop=scan_loop) + result = await controller._run_on_scan_loop(asyncio.sleep(0, result=COMPACTION_SUCCESS)) + finally: + scan_loop.call_soon_threadsafe(scan_loop.stop) + thread.join(timeout=THREAD_JOIN_TIMEOUT_SECONDS) + scan_loop.close() + + assert result == COMPACTION_SUCCESS + + @pytest.mark.asyncio async def test_unknown_command_is_rejected() -> None: controller = TuiController(args()) diff --git a/tests/test_tui_backend_server.py b/tests/test_tui_backend_server.py index 957b9b5c..23179703 100644 --- a/tests/test_tui_backend_server.py +++ b/tests/test_tui_backend_server.py @@ -12,6 +12,7 @@ import pytest from agents.tool import ToolOutputImage from strix.config.settings import DEFAULT_MAX_TURNS +from strix.core.agents import COMPACTION_SUCCESS from strix.interface.tui.backend.controller import TuiController from strix.interface.tui.backend.projection import bounded_state_projection, terminal_projection from strix.interface.tui.backend.protocol import ( @@ -25,6 +26,12 @@ from strix.interface.tui.backend.server import TuiBackendServer from strix.interface.tui.live_view import TuiLiveView +COMPACTION_AGENT_ID = "compact-agent" +COMPACT_COMMAND = "/compact" +COMPACTION_REQUEST_ID = "compact-request" +COMPACTION_SUCCESS_MESSAGE = "Context compaction complete." + + def args() -> argparse.Namespace: return argparse.Namespace( needs_setup=True, @@ -203,6 +210,56 @@ async def test_server_command_round_trip_over_inherited_socket() -> None: await server.close() +@pytest.mark.asyncio +async def test_compact_command_round_trip_over_tui_socket() -> None: + class Coordinator: + def __init__(self) -> None: + self.compacted: list[str] = [] + + async def compact_agent_session(self, agent_id: str) -> str: + self.compacted.append(agent_id) + return COMPACTION_SUCCESS + + coordinator = Coordinator() + controller = TuiController(args(), coordinator=coordinator) + controller.set_runtime(scan_loop=asyncio.get_running_loop()) + backend, child = socket.socketpair() + child.setblocking(False) # noqa: FBT003 + server = TuiBackendServer(controller) + await start_server(server, backend, child) + try: + await receive_initial_state(child) + await send_message( + child, + { + "version": PROTOCOL_VERSION, + "type": "agent.send_message", + "request_id": COMPACTION_REQUEST_ID, + "payload": { + "agent_id": COMPACTION_AGENT_ID, + "message": COMPACT_COMMAND, + }, + }, + ) + + result = await receive_until( + child, + "command_result", + request_id=COMPACTION_REQUEST_ID, + ) + assert result["payload"]["ok"] is True + assert result["payload"]["result"] == {"compacted": True} + assert coordinator.compacted == [COMPACTION_AGENT_ID] + events = controller.live_view.events_for_agent(COMPACTION_AGENT_ID) + assert [event["data"]["content"] for event in events] == [ + COMPACT_COMMAND, + COMPACTION_SUCCESS_MESSAGE, + ] + finally: + child.close() + await server.close() + + def test_unicode_heavy_setup_state_stays_within_control_frame_limit() -> None: controller = TuiController(args()) controller.instruction = "🔒" * 10_000