diff --git a/README.md b/README.md index 390d04de6..217e69766 100644 --- a/README.md +++ b/README.md @@ -185,6 +185,11 @@ strix --target api.your-app.com --instruction-file ./instruction.md strix -n --target ./ --scan-mode quick --scope-mode diff --diff-base origin/main ``` +### Interactive Commands + +While using the default interactive TUI, send `/compact` to the selected agent to force +SDK 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 24a7185f5..a33fde312 100644 --- a/strix/core/agents.py +++ b/strix/core/agents.py @@ -10,15 +10,22 @@ from pathlib import Path from typing import TYPE_CHECKING, Any, Literal, cast +from agents.memory import is_openai_responses_compaction_aware_session + if TYPE_CHECKING: from agents.items import TResponseInputItem - from agents.memory import Session + from agents.memory import OpenAIResponsesCompactionArgs, Session logger = logging.getLogger(__name__) Status = Literal["running", "waiting", "completed", "stopped", "crashed", "failed"] +CompactionResult = Literal["success", "unavailable", "failed"] +COMPACTION_SUCCESS: CompactionResult = "success" +COMPACTION_UNAVAILABLE: CompactionResult = "unavailable" +COMPACTION_FAILED: CompactionResult = "failed" +_FORCED_COMPACTION_ARGS: OpenAIResponsesCompactionArgs = {"force": True} @dataclass(slots=True) @@ -152,6 +159,38 @@ async def send(self, target_agent_id: str, message: dict[str, Any]) -> bool: await self._maybe_snapshot() return True + async def compact_agent_session(self, target_agent_id: str) -> CompactionResult: + """Force the SDK session compaction pipeline for one attached agent.""" + async with self._lock: + if target_agent_id not in self.statuses: + logger.debug("agent.compact dropped unknown target=%s", target_agent_id) + return COMPACTION_UNAVAILABLE + session = self.runtimes.setdefault(target_agent_id, AgentRuntime()).session + + if session is None: + logger.warning( + "agent.compact dropped target=%s because its SDK session is not attached", + target_agent_id, + ) + return COMPACTION_UNAVAILABLE + + if not is_openai_responses_compaction_aware_session(session): + logger.warning( + "agent.compact dropped target=%s because its SDK session does not " + "support compaction", + target_agent_id, + ) + return COMPACTION_UNAVAILABLE + + try: + await session.run_compaction(_FORCED_COMPACTION_ARGS) + except Exception: + logger.exception("agent.compact failed target=%s", target_agent_id) + return COMPACTION_FAILED + + await self._maybe_snapshot() + return COMPACTION_SUCCESS + async def wait_for_message(self, agent_id: str) -> None: while True: async with self._lock: diff --git a/strix/interface/tui/app.py b/strix/interface/tui/app.py index 34965168d..a20995738 100644 --- a/strix/interface/tui/app.py +++ b/strix/interface/tui/app.py @@ -1611,6 +1611,7 @@ def _send_user_message(self, message: str) -> None: live_view=self.live_view, target_agent_id=target_agent_id, message=message, + ui_thread_call=self.call_from_thread, ) if not submitted: self.notify("Scan loop is not ready; message was not sent", severity="warning") diff --git a/strix/interface/tui/live_view.py b/strix/interface/tui/live_view.py index 993074d67..b9ad362ae 100644 --- a/strix/interface/tui/live_view.py +++ b/strix/interface/tui/live_view.py @@ -97,6 +97,17 @@ def record_user_message(self, agent_id: str, content: str) -> None: }, ) + def record_system_message(self, agent_id: str, content: str) -> None: + self._append_event( + agent_id, + "chat", + { + "role": "system", + "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/strix/interface/tui/messages.py b/strix/interface/tui/messages.py index 18cd37c13..80125fb86 100644 --- a/strix/interface/tui/messages.py +++ b/strix/interface/tui/messages.py @@ -4,11 +4,27 @@ import asyncio import logging -from typing import Any +from typing import TYPE_CHECKING, Any + +from strix.core.agents import COMPACTION_FAILED, COMPACTION_SUCCESS, COMPACTION_UNAVAILABLE + + +if TYPE_CHECKING: + from collections.abc import Callable logger = logging.getLogger(__name__) +_COMPACT_COMMANDS = frozenset({"/compact", "/compress"}) +_COMPACT_SUCCESS_MESSAGE = "Context compaction complete." +_COMPACT_UNAVAILABLE_MESSAGE = "Context compaction is unavailable for this agent." +_COMPACT_FAILED_MESSAGE = "Context compaction failed. Please try again." +_COMPACT_RESULT_MESSAGES = { + COMPACTION_SUCCESS: _COMPACT_SUCCESS_MESSAGE, + COMPACTION_UNAVAILABLE: _COMPACT_UNAVAILABLE_MESSAGE, + COMPACTION_FAILED: _COMPACT_FAILED_MESSAGE, +} + def send_user_message_to_agent( *, @@ -17,11 +33,27 @@ def send_user_message_to_agent( live_view: Any, target_agent_id: str, message: str, + ui_thread_call: Callable[..., Any] | None = None, ) -> bool: if loop is None or loop.is_closed(): return False live_view.record_user_message(target_agent_id, message) + if message.strip().lower() in _COMPACT_COMMANDS: + future = asyncio.run_coroutine_threadsafe( + coordinator.compact_agent_session(target_agent_id), + loop, + ) + future.add_done_callback( + lambda compact_future: _record_compaction_result( + compact_future, + live_view=live_view, + target_agent_id=target_agent_id, + ui_thread_call=ui_thread_call, + ), + ) + return True + future = asyncio.run_coroutine_threadsafe( coordinator.send( target_agent_id, @@ -41,3 +73,41 @@ def _log_delivery_failure(future: Any) -> None: return if not delivered: logger.warning("TUI user message was not persisted to the SDK session") + + +def _record_compaction_result( + future: Any, + *, + live_view: Any, + target_agent_id: str, + ui_thread_call: Callable[..., Any] | None = None, +) -> None: + try: + result = future.result() + except Exception: + logger.exception("TUI context compaction failed") + result = COMPACTION_FAILED + + message = _COMPACT_RESULT_MESSAGES.get(result, _COMPACT_FAILED_MESSAGE) + _record_system_message( + live_view=live_view, + target_agent_id=target_agent_id, + message=message, + ui_thread_call=ui_thread_call, + ) + + +def _record_system_message( + *, + live_view: Any, + target_agent_id: str, + message: str, + ui_thread_call: Callable[..., Any] | None, +) -> None: + record_system_message = getattr(live_view, "record_system_message", None) + if not callable(record_system_message): + return + if ui_thread_call is not None: + ui_thread_call(record_system_message, target_agent_id, message) + return + record_system_message(target_agent_id, message) diff --git a/tests/test_manual_compact.py b/tests/test_manual_compact.py new file mode 100644 index 000000000..664ffd5aa --- /dev/null +++ b/tests/test_manual_compact.py @@ -0,0 +1,175 @@ +"""Tests for manual interactive context compaction.""" + +from __future__ import annotations + +import asyncio +import threading +from typing import Any + +import pytest + +from strix.core.agents import AgentCoordinator +from strix.interface.tui.messages import send_user_message_to_agent + + +COMPACT_COMMAND = "/compact" +COMPRESS_COMMAND = "/compress" +TARGET_AGENT_ID = "agent" +WAIT_TIMEOUT_SECONDS = 1.0 +TEST_COMPACTION_SUCCESS = "success" +TEST_COMPACTION_UNAVAILABLE = "unavailable" +TEST_COMPACTION_FAILED = "failed" +COMPACTION_SUCCESS_MESSAGE = "Context compaction complete." +COMPACTION_FAILED_MESSAGE = "Context compaction failed. Please try again." + + +class CompactAwareSession: + def __init__(self) -> None: + self.compaction_args: list[dict[str, Any]] = [] + + async def run_compaction(self, args: dict[str, Any] | None = None) -> None: + self.compaction_args.append(dict(args or {})) + + +class FailingCompactSession: + async def run_compaction(self, _args: dict[str, Any] | None = None) -> None: + raise RuntimeError("compaction backend failed") + + +class RecordingCoordinator: + def __init__( + self, + *, + compact_done: threading.Event, + compact_result: str = TEST_COMPACTION_SUCCESS, + ) -> None: + self.compact_done = compact_done + self.compact_result = compact_result + self.compacted_agents: list[str] = [] + self.sent_messages: list[tuple[str, dict[str, Any]]] = [] + + async def compact_agent_session(self, target_agent_id: str) -> str: + self.compacted_agents.append(target_agent_id) + self.compact_done.set() + return self.compact_result + + async def send(self, target_agent_id: str, message: dict[str, Any]) -> bool: + self.sent_messages.append((target_agent_id, message)) + return True + + +class RecordingLiveView: + def __init__(self) -> None: + self.user_messages: list[tuple[str, str]] = [] + self.system_messages: list[tuple[str, str]] = [] + self.ui_thread_calls: list[tuple[Any, tuple[Any, ...]]] = [] + + def record_user_message(self, target_agent_id: str, message: str) -> None: + self.user_messages.append((target_agent_id, message)) + + def record_system_message(self, target_agent_id: str, message: str) -> None: + self.system_messages.append((target_agent_id, message)) + + def call_from_thread(self, callback: Any, *args: Any) -> None: + self.ui_thread_calls.append((callback, args)) + callback(*args) + + +@pytest.mark.asyncio +async def test_coordinator_forces_compaction_on_attached_session() -> None: + coordinator = AgentCoordinator() + session = CompactAwareSession() + await coordinator.register(TARGET_AGENT_ID, "strix", parent_id=None) + await coordinator.attach_runtime(TARGET_AGENT_ID, session=session) + + compacted = await coordinator.compact_agent_session(TARGET_AGENT_ID) + + assert compacted == TEST_COMPACTION_SUCCESS + assert session.compaction_args == [{"force": True}] + + +@pytest.mark.asyncio +async def test_coordinator_reports_unavailable_when_session_is_not_attached() -> None: + coordinator = AgentCoordinator() + await coordinator.register(TARGET_AGENT_ID, "strix", parent_id=None) + + compacted = await coordinator.compact_agent_session(TARGET_AGENT_ID) + + assert compacted == TEST_COMPACTION_UNAVAILABLE + + +@pytest.mark.asyncio +async def test_coordinator_reports_failed_when_compaction_raises() -> None: + coordinator = AgentCoordinator() + await coordinator.register(TARGET_AGENT_ID, "strix", parent_id=None) + await coordinator.attach_runtime(TARGET_AGENT_ID, session=FailingCompactSession()) + + compacted = await coordinator.compact_agent_session(TARGET_AGENT_ID) + + assert compacted == TEST_COMPACTION_FAILED + + +@pytest.mark.parametrize("command", [COMPACT_COMMAND, COMPRESS_COMMAND]) +def test_compact_command_runs_compaction_without_sending_agent_message(command: str) -> None: + compact_done = threading.Event() + coordinator = RecordingCoordinator(compact_done=compact_done) + live_view = RecordingLiveView() + + loop = asyncio.new_event_loop() + thread = threading.Thread(target=loop.run_forever) + thread.start() + try: + submitted = send_user_message_to_agent( + coordinator=coordinator, + loop=loop, + live_view=live_view, + target_agent_id=TARGET_AGENT_ID, + message=command, + ui_thread_call=live_view.call_from_thread, + ) + + assert submitted is True + assert compact_done.wait(WAIT_TIMEOUT_SECONDS) + finally: + loop.call_soon_threadsafe(loop.stop) + thread.join(timeout=WAIT_TIMEOUT_SECONDS) + loop.close() + + assert coordinator.compacted_agents == [TARGET_AGENT_ID] + assert coordinator.sent_messages == [] + assert live_view.user_messages == [(TARGET_AGENT_ID, command)] + assert live_view.system_messages == [(TARGET_AGENT_ID, COMPACTION_SUCCESS_MESSAGE)] + assert live_view.ui_thread_calls == [ + (live_view.record_system_message, (TARGET_AGENT_ID, COMPACTION_SUCCESS_MESSAGE)) + ] + + +def test_compact_command_reports_failed_compaction() -> None: + compact_done = threading.Event() + coordinator = RecordingCoordinator( + compact_done=compact_done, + compact_result=TEST_COMPACTION_FAILED, + ) + live_view = RecordingLiveView() + + loop = asyncio.new_event_loop() + thread = threading.Thread(target=loop.run_forever) + thread.start() + try: + submitted = send_user_message_to_agent( + coordinator=coordinator, + loop=loop, + live_view=live_view, + target_agent_id=TARGET_AGENT_ID, + message=COMPACT_COMMAND, + ui_thread_call=live_view.call_from_thread, + ) + + assert submitted is True + assert compact_done.wait(WAIT_TIMEOUT_SECONDS) + finally: + loop.call_soon_threadsafe(loop.stop) + thread.join(timeout=WAIT_TIMEOUT_SECONDS) + loop.close() + + assert live_view.system_messages == [(TARGET_AGENT_ID, COMPACTION_FAILED_MESSAGE)]