From b2df7dc01b6bf2ff52b1e0f54e8231d3cd8f0a6f Mon Sep 17 00:00:00 2001 From: Aly Raffauf Date: Mon, 29 Jun 2026 13:51:56 -0400 Subject: [PATCH] context: add ledger and deterministic compaction --- AGENTS.md | 1 + README.md | 1 + lib/agents.nix | 9 + tartarus/agent_loop.py | 23 ++- tartarus/cli.py | 163 +++++++++++++++-- tartarus/config.py | 14 ++ tartarus/context.py | 322 ++++++++++++++++++++++++++++++++++ tartarus/manifest.py | 14 +- tartarus/session.py | 16 +- tests/test_agent_loop.py | 85 +++++++++ tests/test_cli.py | 34 +++- tests/test_context.py | 189 ++++++++++++++++++++ tests/test_jail.py | 2 +- tests/test_manifest.py | 26 +++ tests/test_manifest_loader.py | 20 +++ 15 files changed, 898 insertions(+), 21 deletions(-) create mode 100644 tartarus/context.py create mode 100644 tests/test_context.py diff --git a/AGENTS.md b/AGENTS.md index 75113b0..811edf1 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -105,6 +105,7 @@ From `tartarus/jail.py` and `PLAN.md §8`: output is forced (reading `config` is free), mirroring `system.build.toplevel`. - Capabilities are keyed attrsets under `capabilities.`. Do not put `name` in the capability body; the attrset key is the identity. + `context_status` and `context_read` are reserved internal tool names. - The agent's `shell` is the baseline PATH baked into the manifest; keep it minimal. Tool-specific programs go in that capability's `grants.packages`. `shell.env` adds env vars to every call (reserved names rejected at build diff --git a/README.md b/README.md index a4951f4..974b6c7 100644 --- a/README.md +++ b/README.md @@ -190,6 +190,7 @@ example agent is in `agent.nix`. A capability declares: - the attrset key as its name, plus `description` and model-facing `params` + (`context_status` and `context_read` are reserved for internal context tools) - `policy`: `auto`, `ask-once`, `ask-always`, or `deny` - `grants.packages`: package binaries available only to that tool - `grants.network.allowedHosts`: proxy-allowed HTTP(S) hosts diff --git a/lib/agents.nix b/lib/agents.nix index 6b76010..78dc939 100644 --- a/lib/agents.nix +++ b/lib/agents.nix @@ -125,6 +125,11 @@ let || grant.writable != [ ] || grant.unrestricted; + reservedCapabilityNames = [ + "context_status" + "context_read" + ]; + capabilityType = { options = { description = lib.mkOption { @@ -197,6 +202,10 @@ let capabilityAssertions = name: capability: [ + { + assertion = !(lib.elem name reservedCapabilityNames); + message = "Tartarus capability '${name}' uses a reserved internal tool name."; + } { assertion = !(capability.grants.unrestricted && capability.policy == "auto"); message = "Tartarus capability '${name}' cannot combine unrestricted = true with policy = \"auto\"."; diff --git a/tartarus/agent_loop.py b/tartarus/agent_loop.py index 6576d76..ac8cfec 100644 --- a/tartarus/agent_loop.py +++ b/tartarus/agent_loop.py @@ -15,6 +15,7 @@ import asyncio from dataclasses import dataclass from tartarus.broker import Broker +from tartarus.context import CONTEXT_TOOL_NAMES, CONTEXT_TOOLS, ContextManager from tartarus.manifest import Manifest from tartarus.models import ( TextDelta, @@ -48,11 +49,13 @@ class AgentLoop: broker: Broker, manifest: Manifest, system_prompt: str, + context_manager: ContextManager | None = None, ): self._provider = provider self._broker = broker self._manifest = manifest self._system_prompt = system_prompt + self._context_manager = context_manager or ContextManager() async def run_turn(self, messages: list[dict]): """Drive one human turn to completion, yielding UI events as they happen. @@ -65,8 +68,9 @@ class AgentLoop: while True: text_parts: list[str] = [] turn = None + effective_messages = self._context_manager.effective_messages(messages) async for event in self._provider.stream( - self._system_prompt, messages, self._manifest.tools + self._system_prompt, effective_messages, self._tools() ): if isinstance(event, TextDelta): text_parts.append(event.text) @@ -86,7 +90,7 @@ class AgentLoop: for call in turn.tool_calls: yield ToolStarted(call) result = None - async for tool_event in self._run_tool(call): + async for tool_event in self._run_tool(call, messages): if isinstance(tool_event, ToolOutputDelta): yield tool_event else: @@ -100,7 +104,17 @@ class AgentLoop: messages.append(assistant_message) messages.extend(self._provider.tool_result_messages(results)) - async def _run_tool(self, call: ToolCall): + async def _run_tool(self, call: ToolCall, messages: list[dict]): + # Context tools only inspect local context state — no host reach to gate — + # so they answer here directly, bypassing the broker/jail/policy/audit path. + if call.name in CONTEXT_TOOL_NAMES: + yield ToolResult( + call.id, + self._context_manager.handle_tool(call.name, call.arguments, messages), + is_error=False, + ) + return + output_queue: asyncio.Queue[str] = asyncio.Queue() # The broker runs on this loop, so streamed output lands on the queue @@ -142,3 +156,6 @@ class AgentLoop: finally: if pending_get is not None: pending_get.cancel() + + def _tools(self) -> list[dict]: + return [*self._manifest.tools, *CONTEXT_TOOLS] diff --git a/tartarus/cli.py b/tartarus/cli.py index 3dd5fdf..5a7867c 100644 --- a/tartarus/cli.py +++ b/tartarus/cli.py @@ -8,6 +8,7 @@ killing the process. """ import asyncio +import os import signal import sys from dataclasses import dataclass @@ -16,15 +17,17 @@ from tartarus.agent_loop import AgentLoop, ToolFinished, ToolStarted from tartarus.audit import FileAuditLog from tartarus.background import BackgroundRegistry, Notice from tartarus.broker import Broker +from tartarus.bundle import BundleError, base_env_from, load_bundle, resolve_bundle from tartarus.config import ( Config, ConfigError, ResolvedRuntime, + context_dir_from_env, load_config, resolve_runtime, session_dir_from_env, ) -from tartarus.bundle import BundleError, base_env_from, load_bundle, resolve_bundle +from tartarus.context import ContextError, ContextLedger, ContextLimits, ContextManager from tartarus.jail import JailBuilder from tartarus.manifest_loader import host_system from tartarus.models import TextDelta, ToolOutputDelta @@ -41,6 +44,8 @@ class SessionFlags: continue_latest: bool = False # --continue: reopen the most recent disabled: bool = False # --no-session: don't persist list_sessions: bool = False # --list-sessions: print and exit + context_status: bool = False # --context-status: print status and exit + compact_context: bool = False # --compact-context: compact current session and exit def _parse_session_flags(argv: list[str]) -> tuple[SessionFlags, list[str]]: @@ -60,6 +65,10 @@ def _parse_session_flags(argv: list[str]) -> tuple[SessionFlags, list[str]]: flags.disabled = True elif arg == "--list-sessions": flags.list_sessions = True + elif arg == "--context-status": + flags.context_status = True + elif arg == "--compact-context": + flags.compact_context = True elif arg == "--resume": if i + 1 >= len(argv): raise ConfigError("--resume requires a session id") @@ -157,14 +166,25 @@ async def _send(loop: AgentLoop, messages: list[dict], user_text: str) -> bool: pass -def _persist(store: SessionStore | None, messages: list[dict]) -> None: +def _persist( + store: SessionStore | None, + ledger: ContextLedger | None, + messages: list[dict], +) -> None: """Flush newly committed messages, warning (not failing) on write errors.""" if store is None: return try: - store.append(messages) + start_index = store.append(messages) except SessionError as error: print(f"warning: could not save session: {error}", file=sys.stderr) + return + if start_index is None or ledger is None: + return + try: + ledger.append_message_events(messages, start_index) + except ContextError as error: + print(f"warning: could not save context ledger: {error}", file=sys.stderr) async def _run_one_shot( @@ -172,17 +192,18 @@ async def _run_one_shot( prompt: str, messages: list[dict], store: SessionStore | None, + ledger: ContextLedger | None, registry: BackgroundRegistry, notices: "asyncio.Queue[Notice]", ) -> int: try: if await _send(loop, messages, prompt): - _persist(store, messages) + _persist(store, ledger, messages) # A one-shot run that launched background work waits it out, reacting to # each completion, so the task is not killed the instant the turn ends. failed = False while registry.has_running or not notices.empty(): - if not await _drain_notice(loop, messages, store, notices): + if not await _drain_notice(loop, messages, store, ledger, notices): failed = True return 1 if failed else 0 except ProviderError as error: @@ -194,6 +215,7 @@ async def _run_repl( loop: AgentLoop, messages: list[dict], store: SessionStore | None, + ledger: ContextLedger | None, registry: BackgroundRegistry, notices: "asyncio.Queue[Notice]", ) -> int: @@ -218,7 +240,7 @@ async def _run_repl( continue if notice_task in done: - await _react_to_notice(loop, messages, store, notice_task.result()) + await _react_to_notice(loop, messages, store, ledger, notice_task.result()) continue # notice_task did not win, so the input task is the one that completed. @@ -234,7 +256,7 @@ async def _run_repl( continue try: if await _send(loop, messages, user_text): - _persist(store, messages) + _persist(store, ledger, messages) except ProviderError as error: print(f"provider error: {error}", file=sys.stderr) @@ -247,15 +269,17 @@ async def _drain_notice( loop: AgentLoop, messages: list[dict], store: SessionStore | None, + ledger: ContextLedger | None, notices: "asyncio.Queue[Notice]", ) -> bool: - return await _react_to_notice(loop, messages, store, await notices.get()) + return await _react_to_notice(loop, messages, store, ledger, await notices.get()) async def _react_to_notice( loop: AgentLoop, messages: list[dict], store: SessionStore | None, + ledger: ContextLedger | None, notice: Notice, ) -> bool: """Turn one background completion into a transcript message + follow-up turn. @@ -276,7 +300,7 @@ async def _react_to_notice( print(f"\n [background] {notice.task_id} finished (exit {notice.exit_code})") try: if await _send(loop, messages, text): - _persist(store, messages) + _persist(store, ledger, messages) return True return False except ProviderError as error: @@ -308,6 +332,106 @@ def _print_session_list(session_dir: str) -> None: print(f"{session_id} {preview}") +def _context_limits(max_chars: int | None, recent_turns: int | None) -> ContextLimits: + """Validate context limits, falling back to defaults for unset (None) values. + + The single resolution point for both the env-only inspection path and the + config-driven live run, so they cannot validate differently. + """ + defaults = ContextLimits() + resolved_max_chars = max_chars if max_chars is not None else defaults.max_chars + resolved_recent_turns = ( + recent_turns if recent_turns is not None else defaults.recent_turns + ) + if resolved_max_chars < 0: + raise ConfigError("TARTARUS_CONTEXT_MAX_CHARS must be non-negative") + if resolved_recent_turns < 0: + raise ConfigError("TARTARUS_CONTEXT_RECENT_TURNS must be non-negative") + return ContextLimits( + max_chars=resolved_max_chars, + recent_turns=resolved_recent_turns, + ) + + +def _context_limits_from_env() -> ContextLimits: + return _context_limits( + _int_from_env("TARTARUS_CONTEXT_MAX_CHARS", None), + _int_from_env("TARTARUS_CONTEXT_RECENT_TURNS", None), + ) + + +def _context_limits_from_config(config: Config) -> ContextLimits: + return _context_limits(config.context_max_chars, config.context_recent_turns) + + +def _int_from_env(name: str, default: int | None) -> int | None: + """Parse an integer env var, or return the default when unset. + + Range validation lives in _context_limits, the single resolution point; this + only turns the raw string into an int. + """ + value = os.environ.get(name) + if value is None or value == "": + return default + try: + return int(value) + except ValueError as error: + raise ConfigError(f"{name} must be an integer") from error + + +def _resolve_read_only_session(flags: SessionFlags) -> tuple[SessionStore, list[dict]]: + session_dir = session_dir_from_env() + if flags.disabled: + raise SessionError("--no-session cannot be combined with context inspection") + if flags.resume is not None: + session_id = SessionStore.resolve(session_dir, flags.resume) + else: + session_id = SessionStore.latest(session_dir) + if session_id is None: + raise SessionError(f"no sessions in {session_dir}") + store = SessionStore(session_dir, session_id) + return store, store.load() + + +def _print_context_status(flags: SessionFlags) -> int: + try: + store, messages = _resolve_read_only_session(flags) + ledger = ContextLedger(context_dir_from_env(), store.session_id) + manager = ContextManager(ledger, _context_limits_from_env()) + status = manager.status(messages) + except (ConfigError, ContextError, SessionError) as error: + print(f"configuration error: {error}", file=sys.stderr) + return 1 + print(f"session: {store.session_id}") + print(f"messages: {status.message_count}") + print(f"estimated context chars: {status.estimated_chars}") + print(f"effective messages: {status.effective_message_count}") + print(f"effective estimated chars: {status.effective_estimated_chars}") + print(f"ledger events: {status.ledger_event_count}") + print(f"ledger: {status.ledger_path}") + return 0 + + +def _compact_context(flags: SessionFlags) -> int: + try: + store, messages = _resolve_read_only_session(flags) + ledger = ContextLedger(context_dir_from_env(), store.session_id) + event = ContextManager(ledger, _context_limits_from_env()).compact(messages) + except (ConfigError, ContextError, SessionError) as error: + print(f"configuration error: {error}", file=sys.stderr) + return 1 + if event is None: + print(f"session: {store.session_id}") + print("compaction: nothing to compact") + print(f"ledger: {ledger.path}") + return 0 + covered = event["covered"] + print(f"session: {store.session_id}") + print(f"compacted messages: {covered['start']}-{covered['end']}") + print(f"ledger: {ledger.path}") + return 0 + + def _open_session( config: Config, flags: SessionFlags ) -> tuple[SessionStore | None, list[dict]]: @@ -345,6 +469,10 @@ async def _async_main(argv: list[str]) -> int: if session_flags.list_sessions: _print_session_list(session_dir_from_env()) return 0 + if session_flags.context_status: + return _print_context_status(session_flags) + if session_flags.compact_context: + return _compact_context(session_flags) try: config = load_config() @@ -356,6 +484,13 @@ async def _async_main(argv: list[str]) -> int: print(f"configuration error: {error}", file=sys.stderr) return 1 + ledger = ( + ContextLedger(config.context_dir, store.session_id) + if store is not None + else None + ) + context_manager = ContextManager(ledger, _context_limits_from_config(config)) + print("loading agent bundle...", file=sys.stderr) try: bundle_path = resolve_bundle(config) @@ -406,7 +541,7 @@ async def _async_main(argv: list[str]) -> int: ) # The agent's Nix definition owns its persona; the config default is a fallback. system_prompt = manifest.system_prompt or config.system_prompt - loop = AgentLoop(provider, broker, manifest, system_prompt) + loop = AgentLoop(provider, broker, manifest, system_prompt, context_manager) tool_names = ", ".join(tool["name"] for tool in manifest.tools) mode = "headless" if config.headless else "interactive" @@ -421,12 +556,16 @@ async def _async_main(argv: list[str]) -> int: print(f"audit log: {config.audit_path}", file=sys.stderr) if store is not None: print(f"session: {store.session_id} ({store.path})", file=sys.stderr) + if ledger is not None: + print(f"context ledger: {ledger.path}", file=sys.stderr) prompt = " ".join(argv).strip() try: if prompt: - return await _run_one_shot(loop, prompt, messages, store, registry, notices) - return await _run_repl(loop, messages, store, registry, notices) + return await _run_one_shot( + loop, prompt, messages, store, ledger, registry, notices + ) + return await _run_repl(loop, messages, store, ledger, registry, notices) finally: registry.shutdown_all() diff --git a/tartarus/config.py b/tartarus/config.py index e343020..dc023f5 100644 --- a/tartarus/config.py +++ b/tartarus/config.py @@ -32,6 +32,7 @@ DEFAULT_STATE_DIR = ".tartarus" # Leaf names under /.tartarus, shared by every path-deriving call site. AUDIT_LOG_LEAF = "audit.jsonl" SESSIONS_LEAF = "sessions" +CONTEXT_LEAF = "context" # `path:` copies the directory regardless of git tracking, which keeps local # capability edits visible before they are committed. DEFAULT_FLAKE_REF = "path:." @@ -92,6 +93,10 @@ class Config(BaseSettings): audit_path: str = "" # Directory holding per-conversation transcript files (.jsonl). session_dir: str = Field("", validation_alias=AliasChoices("TARTARUS_SESSIONS_DIR")) + # Directory holding append-only per-session context ledgers. + context_dir: str = Field("", validation_alias=AliasChoices("TARTARUS_CONTEXT_DIR")) + context_max_chars: int | None = None + context_recent_turns: int | None = None output_truncate: int = DEFAULT_OUTPUT_TRUNCATE_CHARS @model_validator(mode="after") @@ -102,6 +107,8 @@ class Config(BaseSettings): self.audit_path = _default_state_path(self.work_tree, AUDIT_LOG_LEAF) if not self.session_dir: self.session_dir = _default_state_path(self.work_tree, SESSIONS_LEAF) + if not self.context_dir: + self.context_dir = _default_state_path(self.work_tree, CONTEXT_LEAF) return self @@ -117,6 +124,13 @@ def session_dir_from_env() -> str: ) +def context_dir_from_env() -> str: + work_tree = os.environ.get("TARTARUS_WORK_TREE") or os.getcwd() + return os.environ.get("TARTARUS_CONTEXT_DIR") or _default_state_path( + work_tree, CONTEXT_LEAF + ) + + def _default_state_path(work_tree: str, leaf: str) -> str: return os.path.join(work_tree, DEFAULT_STATE_DIR, leaf) diff --git a/tartarus/context.py b/tartarus/context.py new file mode 100644 index 0000000..6172aac --- /dev/null +++ b/tartarus/context.py @@ -0,0 +1,322 @@ +"""Context ledger, selection, and deterministic compaction. + +Sessions remain the provider-native transcript of record. The context layer is a +derived, append-only audit trail plus a selector that decides which messages are +sent to the provider for the next round-trip. +""" + +from __future__ import annotations + +import json +import os +import time +from dataclasses import asdict, dataclass +from typing import Any + +CONTEXT_SUFFIX = ".jsonl" +DEFAULT_CONTEXT_MAX_CHARS = 120_000 +DEFAULT_CONTEXT_RECENT_TURNS = 20 +SUMMARY_ROLE = "system" +DEFAULT_LEDGER_READ_LIMIT = 20 +MAX_LEDGER_READ_LIMIT = 100 +SUMMARY_LINE_MAX_CHARS = 240 +ELLIPSIS = "..." +CONTEXT_TOOL_NAMES = {"context_status", "context_read"} +CONTEXT_TOOLS = [ + { + "name": "context_status", + "description": "Read the current conversation context status.", + "parameters": { + "type": "object", + "properties": {}, + "additionalProperties": False, + }, + }, + { + "name": "context_read", + "description": "Read recent append-only context ledger events.", + "parameters": { + "type": "object", + "properties": { + "limit": { + "type": "integer", + "minimum": 1, + "maximum": MAX_LEDGER_READ_LIMIT, + "description": "Maximum number of latest ledger events to return.", + } + }, + "additionalProperties": False, + }, + }, +] + + +class ContextError(Exception): + """Raised when context state cannot be read or written.""" + + +@dataclass(frozen=True) +class ContextLimits: + max_chars: int = DEFAULT_CONTEXT_MAX_CHARS + recent_turns: int = DEFAULT_CONTEXT_RECENT_TURNS + + +@dataclass(frozen=True) +class ContextStatus: + message_count: int + estimated_chars: int + ledger_event_count: int + effective_message_count: int + effective_estimated_chars: int + ledger_path: str | None + + +class ContextLedger: + def __init__(self, context_dir: str, session_id: str): + self._dir = os.path.abspath(context_dir) + self.session_id = session_id + self._path = os.path.join(self._dir, session_id + CONTEXT_SUFFIX) + + @property + def path(self) -> str: + return self._path + + def append_event(self, event: dict[str, Any]) -> None: + if self._dir: + os.makedirs(self._dir, exist_ok=True) + enriched = {"created_at": _timestamp(), **event} + try: + with open(self._path, "a", encoding="utf-8") as context_file: + context_file.write(json.dumps(enriched) + "\n") + except OSError as error: + raise ContextError( + f"cannot write context ledger {self._path}: {error}" + ) from error + + def append_message_events(self, messages: list[dict], start_index: int) -> None: + for offset, message in enumerate(messages[start_index:], start=start_index): + self.append_event(message_event(offset, message)) + + def load_events(self) -> list[dict[str, Any]]: + try: + with open(self._path, encoding="utf-8") as context_file: + return [json.loads(line) for line in context_file if line.strip()] + except FileNotFoundError: + return [] + except (OSError, json.JSONDecodeError) as error: + raise ContextError( + f"cannot read context ledger {self._path}: {error}" + ) from error + + +class ContextManager: + def __init__( + self, + ledger: ContextLedger | None = None, + limits: ContextLimits | None = None, + ): + self._ledger = ledger + self._limits = limits or ContextLimits() + + @property + def ledger_path(self) -> str | None: + return self._ledger.path if self._ledger is not None else None + + def effective_messages(self, messages: list[dict]) -> list[dict]: + events = self._load_events() + summary = latest_summary(events) + if summary is None: + return list(messages) + + # Keep every message after the summarized range; nothing between the + # summary and the recent turns is dropped. fit_to_limit trims oldest + # whole turns only when the result exceeds max_chars. + covered_end = int(summary["covered"]["end"]) + suffix_start = valid_boundary_start(messages, covered_end) + suffix = messages[suffix_start:] + summary_message = { + "role": SUMMARY_ROLE, + "content": "Context summary from earlier transcript:\n\n" + + str(summary["summary"]), + } + return fit_to_limit([summary_message], suffix, self._limits) + + def status(self, messages: list[dict]) -> ContextStatus: + events = self._load_events() + effective = self.effective_messages(messages) + return ContextStatus( + message_count=len(messages), + estimated_chars=estimate_messages(messages), + ledger_event_count=len(events), + effective_message_count=len(effective), + effective_estimated_chars=estimate_messages(effective), + ledger_path=self.ledger_path, + ) + + def read_events( + self, limit: int = DEFAULT_LEDGER_READ_LIMIT + ) -> list[dict[str, Any]]: + events = self._load_events() + bounded_limit = max(1, min(limit, MAX_LEDGER_READ_LIMIT)) + return events[-bounded_limit:] + + def handle_tool( + self, name: str, arguments: dict[str, Any], messages: list[dict] + ) -> str: + if name == "context_status": + return json.dumps(asdict(self.status(messages)), sort_keys=True) + if name == "context_read": + limit = arguments.get("limit", DEFAULT_LEDGER_READ_LIMIT) + if not isinstance(limit, int): + limit = DEFAULT_LEDGER_READ_LIMIT + return json.dumps(self.read_events(limit), sort_keys=True) + raise ContextError(f"unknown context tool '{name}'") + + def compact(self, messages: list[dict]) -> dict[str, Any] | None: + if self._ledger is None: + raise ContextError("cannot compact without a context ledger") + if not messages: + return None + + suffix_start = recent_turn_start(messages, self._limits.recent_turns) + suffix_start = valid_boundary_start(messages, suffix_start) + if suffix_start <= 0: + return None + + summary_text = deterministic_summary(messages[:suffix_start]) + event = { + "type": "context_summary", + "covered": {"start": 0, "end": suffix_start}, + "summary": summary_text, + "source": "deterministic-local", + "estimated_chars": len(summary_text), + } + self._ledger.append_event(event) + return event + + def _load_events(self) -> list[dict[str, Any]]: + if self._ledger is None: + return [] + return self._ledger.load_events() + + +def message_event(index: int, message: dict) -> dict[str, Any]: + role = str(message.get("role", "unknown")) + event_type = { + "user": "user_turn", + "assistant": "assistant_turn", + "tool": "tool_result", + }.get(role, "transcript_message") + if role == "user" and _is_background_notice(message): + event_type = "background_notice" + return { + "type": event_type, + "message_index": index, + "role": role, + "message": message, + "estimated_chars": estimate_message(message), + } + + +def latest_summary(events: list[dict[str, Any]]) -> dict[str, Any] | None: + summaries = [event for event in events if event.get("type") == "context_summary"] + return summaries[-1] if summaries else None + + +def recent_turn_start(messages: list[dict], recent_turns: int) -> int: + if recent_turns <= 0: + return len(messages) + + seen = 0 + for index in range(len(messages) - 1, -1, -1): + if messages[index].get("role") == "user": + seen += 1 + if seen == recent_turns: + return index + return 0 + + +def valid_boundary_start(messages: list[dict], start: int) -> int: + start = max(0, min(start, len(messages))) + while start < len(messages) and messages[start].get("role") == "tool": + start += 1 + return start + + +def fit_to_limit( + prefix: list[dict], + suffix: list[dict], + limits: ContextLimits, +) -> list[dict]: + candidate = list(prefix) + list(suffix) + if estimate_messages(candidate) <= limits.max_chars: + return candidate + + # Over the configured ceiling: drop oldest whole turns from the suffix front + # (each cut starts at a user message), never mid-turn. This is the explicit + # max_chars boundary, not silent loss. + for index, message in enumerate(suffix): + if message.get("role") != "user": + continue + candidate = list(prefix) + suffix[index:] + if estimate_messages(candidate) <= limits.max_chars: + return candidate + return candidate + + +def estimate_message(message: dict) -> int: + return len(json.dumps(message, sort_keys=True)) + + +def estimate_messages(messages: list[dict]) -> int: + return sum(estimate_message(message) for message in messages) + + +def deterministic_summary(messages: list[dict]) -> str: + lines = [ + "# Context Summary", + "", + f"Covered transcript messages: 0-{len(messages)}", + "", + ] + for index, message in enumerate(messages): + role = message.get("role", "unknown") + content = _message_content(message) + if role == "assistant" and message.get("tool_calls"): + lines.append(f"- assistant message {index}: requested tools") + for tool_call in message["tool_calls"]: + function = tool_call.get("function", {}) + name = function.get("name") or tool_call.get("name") or "unknown" + lines.append(f" - tool call: {name}") + continue + if role == "tool": + tool_call_id = message.get("tool_call_id", "unknown") + lines.append( + f"- tool result {index} ({tool_call_id}): {_one_line(content)}" + ) + continue + lines.append(f"- {role} message {index}: {_one_line(content)}") + return "\n".join(lines) + + +def _message_content(message: dict) -> str: + content = message.get("content", "") + if isinstance(content, str): + return content + return json.dumps(content, sort_keys=True) + + +def _one_line(text: str) -> str: + compact = " ".join(text.split()) + if len(compact) > SUMMARY_LINE_MAX_CHARS: + return compact[: SUMMARY_LINE_MAX_CHARS - len(ELLIPSIS)] + ELLIPSIS + return compact + + +def _is_background_notice(message: dict) -> bool: + content = message.get("content") + return isinstance(content, str) and content.startswith("[background] ") + + +def _timestamp() -> str: + return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) diff --git a/tartarus/manifest.py b/tartarus/manifest.py index 5d8a281..6a21402 100644 --- a/tartarus/manifest.py +++ b/tartarus/manifest.py @@ -14,6 +14,7 @@ from typing import Any, Literal from pydantic import BaseModel, Field, field_validator, model_validator from tartarus.constants import CERT_ENV_VARS, STRICT_CONFIG +from tartarus.context import CONTEXT_TOOL_NAMES from typing_extensions import Self @@ -140,6 +141,13 @@ class Capability(BaseModel): # one of "status" | "output" | "stop". None for every other kind. control: Literal["status", "output", "stop"] | None = None + @field_validator("name") + @classmethod + def _reject_reserved_tool_name(cls, v: str) -> str: + if v in CONTEXT_TOOL_NAMES: + raise ValueError(f"capability name '{v}' is reserved") + return v + @field_validator("timeout", mode="before") @classmethod def _reject_bool_timeout(cls, v: object) -> int | None: @@ -332,7 +340,9 @@ class Manifest(BaseModel): ) upper = key.upper() if upper in _RESERVED_SHELL_ENV_NAMES: - raise ValueError(f"shellEnv key '{key}' is reserved and cannot be overridden") + raise ValueError( + f"shellEnv key '{key}' is reserved and cannot be overridden" + ) if upper.endswith("_PROXY"): raise ValueError( f"shellEnv key '{key}' matches the reserved *_PROXY suffix" @@ -358,6 +368,8 @@ class Manifest(BaseModel): name = tool.get("name") if not isinstance(name, str): raise ValueError("tool 'name' must be a string") + if name in CONTEXT_TOOL_NAMES: + raise ValueError(f"tool name '{name}' is reserved") capability = self.capabilities.get(name) if capability is None: raise ValueError(f"tool '{name}' has no matching capability") diff --git a/tartarus/session.py b/tartarus/session.py index c230af9..82b5c44 100644 --- a/tartarus/session.py +++ b/tartarus/session.py @@ -98,11 +98,20 @@ class SessionStore: self._flushed = len(messages) return messages - def append(self, messages: list[dict]) -> None: - """Persist any messages added since the last flush.""" + @property + def flushed_count(self) -> int: + return self._flushed + + def append(self, messages: list[dict]) -> int | None: + """Persist any messages added since the last flush. + + Returns the first appended message index, or None when there was no new + tail to write. + """ tail = messages[self._flushed :] if not tail: - return + return None + start_index = self._flushed if self._dir: os.makedirs(self._dir, exist_ok=True) try: @@ -112,6 +121,7 @@ class SessionStore: except OSError as error: raise SessionError(f"cannot write session {self._path}: {error}") from error self._flushed = len(messages) + return start_index def first_user_message(self) -> str | None: """First user-authored text in the session, for listing previews.""" diff --git a/tests/test_agent_loop.py b/tests/test_agent_loop.py index ce2e0df..1707f03 100644 --- a/tests/test_agent_loop.py +++ b/tests/test_agent_loop.py @@ -6,6 +6,7 @@ from typing import cast from tartarus.agent_loop import AgentLoop, ToolFinished, ToolStarted from tartarus.broker import Broker +from tartarus.context import ContextLedger, ContextLimits, ContextManager from tartarus.jail import ExecResult, JailBuilder from tartarus.models import ( AssistantTurn, @@ -46,6 +47,7 @@ class ScriptedProvider: def __init__(self, turns: list[AssistantTurn]): self._turns = list(turns) self.received_results: list = [] + self.received_messages: list[list[dict]] = [] async def complete( self, system: str, messages: list[dict], tools: list[dict] @@ -55,6 +57,7 @@ class ScriptedProvider: async def stream( self, system: str, messages: list[dict], tools: list[dict] ) -> AsyncIterator[StreamEvent]: + self.received_messages.append(list(messages)) turn = self._turns.pop(0) if turn.text: yield TextDelta(turn.text) @@ -370,6 +373,88 @@ def test_loop_survives_unknown_tool_call(): assert "unknown tool" in finished[0].result.output +def test_loop_sends_effective_messages_without_mutating_raw_transcript(tmp_path): + manifest = echo_manifest() + ledger = ContextLedger(str(tmp_path), "s1") + ledger.append_event( + { + "type": "context_summary", + "covered": {"start": 0, "end": 2}, + "summary": "Earlier work was completed.", + "source": "deterministic-local", + "estimated_chars": 29, + } + ) + provider = ScriptedProvider( + [ + AssistantTurn( + text="done", + tool_calls=[], + raw={"role": "assistant"}, + stop_reason="end", + ) + ] + ) + loop = AgentLoop( + provider, + Broker(manifest, cast(JailBuilder, LocalJail()), PolicyEngine()), + manifest, + "system", + ContextManager(ledger, ContextLimits(max_chars=10_000, recent_turns=1)), + ) + messages = [ + {"role": "user", "content": "old"}, + {"role": "assistant", "content": "old reply"}, + {"role": "user", "content": "new"}, + ] + + asyncio.run(_drain(loop, messages)) + + assert provider.received_messages[0][0]["role"] == "system" + assert "Earlier work" in provider.received_messages[0][0]["content"] + assert provider.received_messages[0][1:] == [{"role": "user", "content": "new"}] + assert messages == [ + {"role": "user", "content": "old"}, + {"role": "assistant", "content": "old reply"}, + {"role": "user", "content": "new"}, + {"role": "assistant", "content": "done"}, + ] + + +def test_loop_handles_context_status_as_internal_tool(tmp_path): + manifest = echo_manifest() + provider = ScriptedProvider( + [ + AssistantTurn( + text=None, + tool_calls=[ToolCall("call-1", "context_status", {})], + raw={"role": "assistant"}, + stop_reason="tool_calls", + ), + AssistantTurn( + text="I checked context.", + tool_calls=[], + raw={"role": "assistant"}, + stop_reason="end", + ), + ] + ) + loop = AgentLoop( + provider, + Broker(manifest, cast(JailBuilder, LocalJail()), PolicyEngine()), + manifest, + "system", + ContextManager(ContextLedger(str(tmp_path), "s1")), + ) + messages = [{"role": "user", "content": "status"}] + + events = asyncio.run(_drain(loop, messages)) + + assert _text(events) == "I checked context." + assert len(provider.received_results) == 1 + assert '"message_count": 1' in provider.received_results[0].output + + def test_loop_brokers_multiple_parallel_tool_calls(): """One assistant turn with two tool calls brokers both, aggregates results.""" manifest = echo_manifest() diff --git a/tests/test_cli.py b/tests/test_cli.py index e031d31..69c7710 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -11,6 +11,7 @@ from tartarus.cli import ( _bundle_manifest_source, _parse_agent_selector, _parse_session_flags, + _print_context_status, _run_one_shot, ) from tartarus.config import ConfigError @@ -76,6 +77,13 @@ def test_parse_session_flags_no_session_and_list(): assert flags.list_sessions is True +def test_parse_session_flags_context_commands(): + flags, _ = _parse_session_flags(["--context-status"]) + assert flags.context_status is True + flags, _ = _parse_session_flags(["--compact-context"]) + assert flags.compact_context is True + + def test_parse_session_flags_resume_without_id_errors(): with pytest.raises(ConfigError, match="--resume requires"): _parse_session_flags(["--resume"]) @@ -123,6 +131,29 @@ def test_print_session_list_reports_unreadable_dir(tmp_path, capsys): assert "warning: could not list sessions" in captured.err +def test_print_context_status_uses_latest_session_without_api_key( + tmp_path, monkeypatch, capsys +): + from tartarus.session import SessionStore + + monkeypatch.setenv("TARTARUS_WORK_TREE", str(tmp_path)) + session_dir = tmp_path / ".tartarus" / "sessions" + SessionStore(str(session_dir), "s1").append( + [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + ] + ) + + result = _print_context_status(SessionFlags()) + + captured = capsys.readouterr() + assert result == 0 + assert "session: s1" in captured.out + assert "messages: 2" in captured.out + assert "ledger events: 0" in captured.out + + def test_run_one_shot_returns_one_when_background_reaction_fails(monkeypatch): """A provider-level failure while reacting to a completion yields exit code 1.""" @@ -136,7 +167,7 @@ def test_run_one_shot_returns_one_when_background_reaction_fails(monkeypatch): async def fake_send(_loop, _messages, _text): return True - async def fake_drain(_loop, _messages, _store, _notices): + async def fake_drain(_loop, _messages, _store, _ledger, _notices): state["running"] = False return False @@ -156,6 +187,7 @@ def test_run_one_shot_returns_one_when_background_reaction_fails(monkeypatch): "prompt", [], None, + None, cast(BackgroundRegistry, FakeRegistry()), asyncio.Queue(), ) diff --git a/tests/test_context.py b/tests/test_context.py new file mode 100644 index 0000000..f4dde7b --- /dev/null +++ b/tests/test_context.py @@ -0,0 +1,189 @@ +import pytest + +from tartarus.context import ( + ContextError, + ContextLedger, + ContextLimits, + ContextManager, + deterministic_summary, + estimate_messages, + message_event, + valid_boundary_start, +) +from tartarus.session import SessionStore + + +def test_ledger_append_and_load_round_trips_events(tmp_path): + ledger = ContextLedger(str(tmp_path), "s1") + ledger.append_event({"type": "user_turn", "message_index": 0}) + ledger.append_message_events( + [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + ], + 0, + ) + + events = ContextLedger(str(tmp_path), "s1").load_events() + + assert [event["type"] for event in events] == [ + "user_turn", + "user_turn", + "assistant_turn", + ] + assert events[1]["message"]["content"] == "hi" + + +def test_ledger_rejects_corrupt_json(tmp_path): + path = tmp_path / "s1.jsonl" + path.write_text("not json\n") + + with pytest.raises(ContextError, match="cannot read context ledger"): + ContextLedger(str(tmp_path), "s1").load_events() + + +def test_message_event_classifies_tool_and_background_messages(): + assert message_event(0, {"role": "tool", "content": "ok"})["type"] == "tool_result" + event = message_event(1, {"role": "user", "content": "[background] t1 done"}) + assert event["type"] == "background_notice" + + +def test_effective_messages_are_identity_without_summary(tmp_path): + manager = ContextManager(ContextLedger(str(tmp_path), "s1")) + messages = [{"role": "user", "content": "hi"}] + + assert manager.effective_messages(messages) == messages + assert manager.effective_messages(messages) is not messages + + +def test_effective_messages_keep_all_messages_after_the_summary(tmp_path): + # The transcript has grown well past the summarized range; everything after + # covered_end must survive even though it is older than the recent window. + ledger = ContextLedger(str(tmp_path), "s1") + ledger.append_event( + { + "type": "context_summary", + "covered": {"start": 0, "end": 2}, + "summary": "Earlier work was completed.", + "source": "deterministic-local", + "estimated_chars": 27, + } + ) + messages = [ + {"role": "user", "content": "covered 0"}, + {"role": "assistant", "content": "covered 1"}, + {"role": "user", "content": "gap 2"}, + {"role": "assistant", "content": "gap 3"}, + {"role": "user", "content": "recent 4"}, + {"role": "assistant", "content": "recent 5"}, + ] + + effective = ContextManager( + ledger, + ContextLimits(max_chars=10_000, recent_turns=1), + ).effective_messages(messages) + + assert effective[0]["role"] == "system" + assert effective[1:] == messages[2:] + assert messages == [ + {"role": "user", "content": "covered 0"}, + {"role": "assistant", "content": "covered 1"}, + {"role": "user", "content": "gap 2"}, + {"role": "assistant", "content": "gap 3"}, + {"role": "user", "content": "recent 4"}, + {"role": "assistant", "content": "recent 5"}, + ] + + +def test_effective_messages_use_summary_plus_valid_recent_suffix(tmp_path): + ledger = ContextLedger(str(tmp_path), "s1") + ledger.append_event( + { + "type": "context_summary", + "covered": {"start": 0, "end": 2}, + "summary": "Earlier: user asked for setup.", + "source": "deterministic-local", + "estimated_chars": 31, + } + ) + messages = [ + {"role": "user", "content": "old"}, + {"role": "assistant", "content": "old reply"}, + {"role": "user", "content": "new"}, + {"role": "assistant", "content": "new reply"}, + ] + + effective = ContextManager( + ledger, + ContextLimits(max_chars=10_000, recent_turns=1), + ).effective_messages(messages) + + assert effective[0]["role"] == "system" + assert "Earlier: user asked" in effective[0]["content"] + assert effective[1:] == messages[2:] + + +def test_valid_boundary_selection_skips_orphan_tool_results(): + messages = [ + {"role": "user", "content": "run"}, + {"role": "assistant", "tool_calls": [{"id": "call-1"}]}, + {"role": "tool", "tool_call_id": "call-1", "content": "ok"}, + {"role": "assistant", "content": "done"}, + ] + + assert valid_boundary_start(messages, 2) == 3 + + +def test_deterministic_compaction_appends_summary_and_preserves_session(tmp_path): + session = SessionStore(str(tmp_path / "sessions"), "s1") + messages = [ + {"role": "user", "content": "old"}, + {"role": "assistant", "content": "old reply"}, + {"role": "user", "content": "new"}, + {"role": "assistant", "content": "new reply"}, + ] + session.append(messages) + ledger = ContextLedger(str(tmp_path / "context"), "s1") + + event = ContextManager( + ledger, + ContextLimits(max_chars=10_000, recent_turns=1), + ).compact(messages) + + assert event is not None + assert event["covered"] == {"start": 0, "end": 2} + assert "user message 0: old" in event["summary"] + assert SessionStore(str(tmp_path / "sessions"), "s1").load() == messages + persisted_event = ContextLedger(str(tmp_path / "context"), "s1").load_events()[-1] + assert persisted_event["type"] == "context_summary" + assert persisted_event["covered"] == {"start": 0, "end": 2} + assert persisted_event["summary"] == event["summary"] + + +def test_status_counts_raw_and_effective_messages(tmp_path): + ledger = ContextLedger(str(tmp_path), "s1") + messages = [{"role": "user", "content": "hi"}] + + status = ContextManager(ledger).status(messages) + + assert status.message_count == 1 + assert status.estimated_chars == estimate_messages(messages) + assert status.ledger_event_count == 0 + + +def test_deterministic_summary_mentions_tool_calls_and_results(): + summary = deterministic_summary( + [ + {"role": "user", "content": "run it"}, + { + "role": "assistant", + "tool_calls": [ + {"id": "call-1", "function": {"name": "bash"}}, + ], + }, + {"role": "tool", "tool_call_id": "call-1", "content": "ok"}, + ] + ) + + assert "tool call: bash" in summary + assert "tool result 2 (call-1): ok" in summary diff --git a/tests/test_jail.py b/tests/test_jail.py index 22cb3ed..aa17a0e 100644 --- a/tests/test_jail.py +++ b/tests/test_jail.py @@ -111,7 +111,7 @@ def test_bwrap_argv_wraps_command_without_hook(tmp_path): @_NEEDS_SANDBOX def test_shell_hook_runs_before_unrestricted_command(tmp_path, shell_path): hook = tmp_path / "hook" - hook.write_text('export HOOK_FLAG=ran\n') + hook.write_text("export HOOK_FLAG=ran\n") jail = JailBuilder(str(tmp_path), shell_path) spec = JailSpec( work_tree=str(tmp_path), diff --git a/tests/test_manifest.py b/tests/test_manifest.py index 7946206..888d288 100644 --- a/tests/test_manifest.py +++ b/tests/test_manifest.py @@ -1,6 +1,9 @@ +import pytest + from tartarus.manifest import ( Capability, Grant, + Manifest, Param, _RESERVED_SHELL_ENV_NAMES, build_manifest, @@ -86,3 +89,26 @@ def test_deny_capabilities_are_not_projected_into_tools(): tool_names = [tool["name"] for tool in manifest.tools] assert tool_names == ["open"] assert "locked" in manifest.capabilities + + +@pytest.mark.parametrize("reserved_name", ["context_status", "context_read"]) +def test_reserved_context_capability_names_are_rejected(reserved_name): + with pytest.raises(ValueError, match=f"'{reserved_name}' is reserved"): + Capability( + name=reserved_name, + description="reserved", + policy="auto", + params={}, + grants=Grant(), + runner="true", + ) + + +def test_reserved_context_tool_names_are_rejected_without_capability(): + with pytest.raises(ValueError, match="'context_status' is reserved"): + Manifest( + tools=[{"name": "context_status", "description": "", "parameters": {}}], + capabilities={}, + ca_bundle_file="/nix/store/cacert/etc/ssl/certs/ca-bundle.crt", + shell_closure_file="/nix/store/shell-closure/store-paths", + ) diff --git a/tests/test_manifest_loader.py b/tests/test_manifest_loader.py index a3fd448..edccaa1 100644 --- a/tests/test_manifest_loader.py +++ b/tests/test_manifest_loader.py @@ -278,6 +278,26 @@ def test_tool_without_capability_is_rejected(): build_manifest_from_raw(raw) +@pytest.mark.parametrize("reserved_name", ["context_status", "context_read"]) +def test_reserved_context_capability_name_is_rejected(reserved_name): + raw = _valid_raw() + raw["capabilities"][reserved_name] = raw["capabilities"].pop("echo") + raw["tools"][0]["name"] = reserved_name + + with pytest.raises(ManifestError, match=f"'{reserved_name}' is reserved"): + build_manifest_from_raw(raw) + + +def test_reserved_context_tool_name_without_capability_is_rejected(): + raw = _valid_raw() + raw["tools"].append( + {"name": "context_status", "description": "", "parameters": {}} + ) + + with pytest.raises(ManifestError, match="'context_status' is reserved"): + build_manifest_from_raw(raw) + + def test_invalid_policy_literal_is_rejected(): raw = _valid_raw() raw["capabilities"]["echo"]["policy"] = "sometimes" -- 2.51.2