diff --git a/coworker/conversations.py b/coworker/conversations.py index fd2131bf..e67ef2fb 100644 --- a/coworker/conversations.py +++ b/coworker/conversations.py @@ -95,6 +95,7 @@ class ConversationStore: "ALTER TABLE sessions ADD COLUMN auto_title TEXT", "ALTER TABLE sessions ADD COLUMN renamed INTEGER DEFAULT 0", "ALTER TABLE sessions ADD COLUMN grants TEXT", + "ALTER TABLE sessions ADD COLUMN compaction TEXT", ): try: self._conn.execute(ddl) @@ -191,13 +192,14 @@ class ConversationStore: title = record.title or title_from(record.messages) self._conn.execute( """ - INSERT INTO sessions (session_id, workspace, model, mode, title, agent, n_msgs, messages, extra_roots, grants, updated_at) - VALUES (?, ?, ?, ?, ?, ?, ?, NULL, ?, ?, CURRENT_TIMESTAMP) + INSERT INTO sessions (session_id, workspace, model, mode, title, agent, n_msgs, messages, extra_roots, grants, compaction, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, NULL, ?, ?, ?, CURRENT_TIMESTAMP) ON CONFLICT(session_id) DO UPDATE SET workspace = excluded.workspace, model = excluded.model, mode = excluded.mode, title = COALESCE(sessions.title, excluded.title), agent = excluded.agent, n_msgs = excluded.n_msgs, messages = NULL, extra_roots = excluded.extra_roots, - grants = excluded.grants, updated_at = CURRENT_TIMESTAMP + grants = excluded.grants, compaction = excluded.compaction, + updated_at = CURRENT_TIMESTAMP """, ( sid, @@ -209,6 +211,7 @@ class ConversationStore: len(record.messages), json.dumps(record.extra_roots or []), json.dumps(record.grants or {}), + json.dumps(record.compaction or {}), ), ) self._conn.commit() @@ -241,6 +244,10 @@ class ConversationStore: row["extra_roots"] if "extra_roots" in row.keys() else None ), grants=_load_grants(row["grants"] if "grants" in row.keys() else None), + # Auto-compaction state (OPE-27) — same defensive parse as grants. + compaction=_load_grants( + row["compaction"] if "compaction" in row.keys() else None + ), pinned=bool(row["pinned"]), archived=bool(row["archived"]), origin=row["origin"], diff --git a/coworker/engine.py b/coworker/engine.py index 3ce77d1e..f3379d61 100644 --- a/coworker/engine.py +++ b/coworker/engine.py @@ -19,6 +19,7 @@ from dataclasses import dataclass from enum import Enum from typing import Any, AsyncIterator, Awaitable, Callable, Optional +from . import compaction as _compaction from .events import Event, EventType from .permissions import Mode, PermissionEngine from .providers import AssistantTurn, ProviderClient, ToolCall @@ -103,6 +104,14 @@ class TurnEngine: # (answerable inline in a live session or from the Inbox when unattended). None on surfaces # that can't ask (the tool then no-ops). self.question_asker = question_asker + # Auto-compaction (OPE-27) — set post-construction by the surface/manager so the + # constructor footprint stays put. `compaction_settings` is a live getter (Settings + # changes apply without a rebuild); `is_attended` gates the failure prompt (None → + # treat as unattended: never park a background run on internal bookkeeping). + self.compaction_state: Optional[_compaction.CompactionState] = None + self.compaction_settings: Optional[Callable[[], dict[str, Any]]] = None + self.is_attended: Optional[Callable[[], bool]] = None + self._last_context_tokens: Optional[int] = None self.audit_context: dict[str, Any] = {} if instructions and not ( self.messages and self.messages[0].get("role") == "system" @@ -302,6 +311,13 @@ class TurnEngine: return iterations += 1 + # Auto-compaction checkpoint (OPE-27): between tool turns and before a new + # turn's first call. Deliberately no "wrap up" warning to the model. + notice = await self._compact_now() + if notice: + self._append_notice("compacted", notice) + yield Event(EventType.COMPACTED, {"text": notice}) + turn: Optional[AssistantTurn] = None streamed: list[str] = [] streamed_reasoning: list[str] = [] @@ -329,6 +345,16 @@ class TurnEngine: if chunk.turn is not None: turn = chunk.turn except Exception as exc: # provider failure + # A raw context-overflow 400 (compaction mispredicted, e.g. the estimate + # path) routes into the compaction policy instead of surfacing. The retry + # is progress-guarded: each pass moves the boundary forward or gives up, + # so a model that keeps overflowing still terminates in the error path. + if _compaction.is_context_overflow(exc) and not self._cancel.is_set(): + notice = await self._compact_now(force=True) + if notice: + self._append_notice("compacted", notice) + yield Event(EventType.COMPACTED, {"text": notice}) + continue # Same contract as the stop path below: the partial the user watched # arrive survives the failure. if streamed or streamed_reasoning: @@ -352,6 +378,10 @@ class TurnEngine: return if turn is None: turn = AssistantTurn() + if turn.usage is not None: + # The trigger signal: the prompt-side total that actually occupied the + # window on this round-trip (estimate fallback when never reported). + self._last_context_tokens = turn.usage.context_tokens self.messages.append(_assistant_message(turn, model=self.model)) payload: dict[str, Any] = { @@ -386,6 +416,97 @@ class TurnEngine: if self._steering: self._inject_steering() + # -- auto-compaction (OPE-27) ------------------------------------------------ + def _compaction_config(self) -> dict[str, Any]: + cfg = dict(self.compaction_settings() or {}) if self.compaction_settings else {} + if not cfg.get("context_window"): + from .providers.matrix import model_context_windows + + cfg["context_window"] = model_context_windows().get(self.model) + cfg.setdefault("threshold_pct", _compaction.DEFAULT_THRESHOLD_PCT) + cfg.setdefault("cap_tokens", _compaction.DEFAULT_CAP_TOKENS) + return cfg + + async def _compact_now(self, *, force: bool = False) -> Optional[str]: + """Run the compaction policy when the trigger fires (or `force`, the overflow + path). Returns the user-facing notice text when the outbound view changed, else + None. Failure policy per spec: retry once (both modes); attended → Retry / Trim + prompt; unattended → auto-trim and continue (never park a run on bookkeeping).""" + cfg = self._compaction_config() + if cfg.get("enabled") is False and not force: + return None + pct = float(cfg["threshold_pct"]) + cap = int(cfg["cap_tokens"]) + window = cfg.get("context_window") + if not force: + signal = self._last_context_tokens or _compaction.estimate_tokens( + self._outbound_messages() + ) + if not _compaction.should_compact( + signal, window, threshold_pct=pct, cap_tokens=cap + ): + return None + keep = int( + _compaction.KEEP_RECENT_FRACTION + * _compaction.trigger_tokens(window, threshold_pct=pct, cap_tokens=cap) + ) + model = str(cfg.get("model") or "") or self.model + + def _build() -> Optional[_compaction.CompactionState]: + return _compaction.build_state( + self.messages, + provider=self.provider, + model=model, + keep_tokens=keep, + prior=self.compaction_state, + ) + + state: Optional[_compaction.CompactionState] = None + failed = False + for _attempt in range(2): # first try + the unconditional single retry + try: + state = await asyncio.to_thread(_build) + failed = False + break + except Exception: + failed = True + if failed and self.question_asker is not None and self.is_attended and self.is_attended(): + while True: + answer = await self._interruptible( + self.question_asker( + { + "question": ( + "Context compaction failed — the summarizer couldn't " + "condense this session's history. How should I proceed?" + ), + "options": ["Retry", "Trim oldest 10%"], + "allow_text": False, + "header": "Compaction", + }, + None, + ), + interrupted=None, + ) + if not answer or answer.get("answer") != "Retry": + break + try: + state = await asyncio.to_thread(_build) + failed = False + break + except Exception: + continue + if state is not None: + self.compaction_state = state + self._last_context_tokens = None # stale once the outbound view shrank + return "Context compacted — earlier turns were summarized" + if failed or force: + trimmed = _compaction.trim_state(self.messages, prior=self.compaction_state) + if trimmed is not None: + self.compaction_state = trimmed + self._last_context_tokens = None + return "Context trimmed — oldest turns dropped (summary unavailable)" + return None + # -- helpers ---------------------------------------------------------------- async def _astream(self): """Bridge the provider's blocking stream generator to the async loop via a @@ -893,13 +1014,19 @@ class TurnEngine: # one. Whole `notice` messages (error/interrupted/model-switch markers) are # display-only too: dropped entirely. _SIDECARS = ("source", "_display", "ts", "reasoning", "usage") + # Auto-compaction (OPE-27): everything before the boundary is represented by the + # compacted block. Outbound-only — the canonical history stays intact — and the + # block+tail are byte-stable between turns, so prompt caching keeps working. + source_messages = _compaction.apply_to_outbound( + self.messages, self.compaction_state + ) out = [ ( {k: v for k, v in msg.items() if k not in _SIDECARS} if any(s in msg for s in _SIDECARS) else msg ) - for msg in self.messages + for msg in source_messages if msg.get("role") != "notice" ] # PDF attachments (stored as `file` parts) are adapted to the ACTIVE model right diff --git a/coworker/events.py b/coworker/events.py index 4cdfc192..0b146485 100644 --- a/coworker/events.py +++ b/coworker/events.py @@ -31,6 +31,7 @@ class EventType(str, Enum): TURN_END = "turn_end" ERROR = "error" INTERRUPTED = "interrupted" + COMPACTED = "compacted" # outbound history was compacted (summary or trim) @dataclass diff --git a/coworker/server/app.py b/coworker/server/app.py index 3b054945..b27f6c7d 100644 --- a/coworker/server/app.py +++ b/coworker/server/app.py @@ -1677,6 +1677,9 @@ def create_app(manager: SessionManager) -> FastAPI: ) await ws.close() return + # Auto-compaction failure prompt (OPE-27): only an ATTENDED session may be asked + # Retry/Trim — unattended runs auto-trim (the policy in engine._compact_now). + engine.is_attended = lambda: _visibility() == VIS_INLINE await ws.send_json( { "type": "ready", diff --git a/coworker/server/manager.py b/coworker/server/manager.py index 048604d7..d43a76c3 100644 --- a/coworker/server/manager.py +++ b/coworker/server/manager.py @@ -460,6 +460,13 @@ class SessionManager: ) if record is not None and record.grants: self._apply_grants(engine, record.grants) + # Auto-compaction (OPE-27): restore the persisted view boundary and wire the live + # Settings getter — post-construction, so build_engine's signature stays put. + if record is not None and record.compaction: + from ..compaction import CompactionState + + engine.compaction_state = CompactionState.from_dict(record.compaction) + engine.compaction_settings = self.compaction_settings self._engines[session_id] = engine if is_new_session: self._emit_session_created(session_id, agent_name) @@ -1860,6 +1867,23 @@ class SessionManager: "pdf_max_mb": max(1, min(mb, 10)), } + def compaction_settings(self) -> dict[str, Any]: + """The live auto-compaction knobs (OPE-27) — read by every engine per check, so a + Settings change applies without a rebuild. Only the two spec'd overrides plus the + summarizer-model pin; absent keys fall back to compaction.py defaults.""" + from ..compaction import DEFAULT_CAP_TOKENS, DEFAULT_THRESHOLD_PCT + + return { + "threshold_pct": float( + self._prefs.get("compaction_threshold_pct") or DEFAULT_THRESHOLD_PCT + ), + "cap_tokens": int( + self._prefs.get("compaction_cap_tokens") or DEFAULT_CAP_TOKENS + ), + # "" → the session's own model (engine falls back to self.model). + "model": str(self._prefs.get("compaction_model") or ""), + } + def set_pdf_settings( self, fallback: Any = None, @@ -3249,6 +3273,11 @@ class SessionManager: agent=getattr(engine, "agent_name", "code"), extra_roots=self._extra_roots_of(engine), grants=_grants_of(engine), + compaction=( + engine.compaction_state.as_dict() + if getattr(engine, "compaction_state", None) + else {} + ), ) ) diff --git a/coworker/sessions.py b/coworker/sessions.py index cc6c4cf5..4bd857c7 100644 --- a/coworker/sessions.py +++ b/coworker/sessions.py @@ -34,3 +34,6 @@ class SessionRecord: # (e.g. origin="slack", origin_label="#general · T0ABCD"). Set once at spawn. origin: Optional[str] = None origin_label: Optional[str] = None + # Auto-compaction state (OPE-27): CompactionState.as_dict(), {} when never compacted. + # Persisted so a reloaded session keeps its compacted outbound view. + compaction: dict[str, Any] = field(default_factory=dict) diff --git a/tests/test_compaction_engine.py b/tests/test_compaction_engine.py new file mode 100644 index 00000000..ff54748f --- /dev/null +++ b/tests/test_compaction_engine.py @@ -0,0 +1,232 @@ +"""OPE-27 engine hook: the mid-run trigger, the outbound view, the usage signal, the +failure policy (attended prompt / unattended auto-trim), raw-overflow routing, and the +session persistence round-trip. Scripted providers, tiny forced windows, no network.""" + +import asyncio + +from coworker.engine import TurnEngine +from coworker.events import EventType +from coworker.permissions import PermissionEngine +from coworker.providers import ( + AssistantTurn, + ModelCapabilities, + ProviderClient, + ToolCall, +) +from coworker.providers.base import TokenUsage +from coworker.tools import ToolRegistry + +SUMMARY = "## Primary request and intent\nkeep building the report" + + +class CompactingProvider(ProviderClient): + """Scripted main turns; summarizer calls (recognized by the compaction system prompt) + are answered out-of-band so they never consume the main script.""" + + def __init__(self, turns, *, summary=SUMMARY, summary_fails=0, main_overflows=0): + self._turns = list(turns) + self.summary = summary + self.summary_fails = summary_fails + self.main_overflows = main_overflows + self.summary_calls = [] + self.main_calls = 0 + + def complete(self, *, model, messages, tools=None, **settings): + if messages and "compacting an AI coworker" in str( + messages[0].get("content", "") + ): + self.summary_calls.append({"model": model, "messages": messages}) + if self.summary_fails > 0: + self.summary_fails -= 1 + raise RuntimeError("summarizer down") + return AssistantTurn(text=self.summary, finish_reason="stop") + self.main_calls += 1 + if self.main_overflows > 0: + self.main_overflows -= 1 + raise RuntimeError( + "Error 400: maximum context length is 100000 tokens, request used more" + ) + return self._turns.pop(0) + + def capabilities(self, model): + return ModelCapabilities() + + +def long_history(turns=8, bulk=1500): + msgs = [{"role": "system", "content": "be helpful"}] + for i in range(turns): + msgs.append({"role": "user", "content": f"request {i}", "ts": 1.0}) + msgs.append( + {"role": "assistant", "content": f"answer {i} " + "x" * bulk, "ts": 1.0} + ) + return msgs + + +def make_engine(tmp_path, provider, *, messages=None, cap=400): + engine = TurnEngine( + provider=provider, + registry=ToolRegistry(), + permissions=PermissionEngine(workspace_root=tmp_path), + model="gpt-5.5", + messages=messages, + ) + engine.compaction_settings = lambda: { + "cap_tokens": cap, + "threshold_pct": 0.8, + "context_window": 100_000, + } + return engine + + +def collect(engine, text="continue"): + async def _run(): + return [e async for e in engine.run(text)] + + return asyncio.run(_run()) + + +def test_compacts_before_the_turn_when_estimate_crosses(tmp_path): + provider = CompactingProvider([AssistantTurn(text="done", finish_reason="stop")]) + engine = make_engine(tmp_path,provider, messages=long_history(), cap=400) + events = collect(engine) + + assert any(e.type == EventType.COMPACTED for e in events) + assert not any(e.type == EventType.ERROR for e in events) + state = engine.compaction_state + assert state is not None and not state.trimmed + assert provider.summary_calls[0]["model"] == "gpt-5.5" # session's own model + + # Outbound view: system survives, the block stands in for the old turns, the + # canonical transcript is untouched, and the persisted notice marks the spot. + out = engine._outbound_messages() + assert out[0]["role"] == "system" + assert "" in out[1]["content"] + assert SUMMARY.splitlines()[-1] in out[1]["content"] + assert "request 0" in out[1]["content"] # mechanical user-message list + assert any("answer 0" in str(m.get("content")) for m in engine.messages) + assert any( + m.get("role") == "notice" and m.get("kind") == "compacted" + for m in engine.messages + ) + + +def test_usage_signal_triggers_between_tool_turns(tmp_path): + # History too small for the estimate path — only the reported usage crosses the + # trigger, after iteration 1's round-trip. The compaction runs before iteration 2. + provider = CompactingProvider( + [ + AssistantTurn( + tool_calls=[ToolCall(id="c1", name="nonexistent_tool", arguments={})], + finish_reason="tool_calls", + usage=TokenUsage(input=90_000, output=10), + ), + AssistantTurn(text="done", finish_reason="stop"), + ] + ) + engine = make_engine(tmp_path,provider, messages=long_history(turns=2, bulk=10), cap=400) + events = collect(engine) + assert any(e.type == EventType.COMPACTED for e in events) + assert provider.summary_calls # driven by usage, not the (tiny) estimate + assert engine._last_context_tokens is None # reset once the view shrank + + +def test_summarizer_failure_unattended_auto_trims(tmp_path): + provider = CompactingProvider( + [AssistantTurn(text="done", finish_reason="stop")], summary_fails=99 + ) + engine = make_engine(tmp_path,provider, messages=long_history(), cap=400) + events = collect(engine) # is_attended is None → unattended policy + + compacted = [e for e in events if e.type == EventType.COMPACTED] + assert compacted and "trimmed" in compacted[0].data["text"].lower() + assert engine.compaction_state is not None and engine.compaction_state.trimmed + assert len(provider.summary_calls) == 2 # the one unconditional retry, then trim + + +def test_summarizer_failure_attended_prompts_retry_then_succeeds(tmp_path): + provider = CompactingProvider( + [AssistantTurn(text="done", finish_reason="stop")], summary_fails=2 + ) + engine = make_engine(tmp_path,provider, messages=long_history(), cap=400) + engine.is_attended = lambda: True + asked = [] + + async def asker(args, tool_call_id=None): + asked.append(args) + return {"answer": "Retry"} + + engine.question_asker = asker + collect(engine) + + assert asked and asked[0]["options"] == ["Retry", "Trim oldest 10%"] + assert engine.compaction_state is not None and not engine.compaction_state.trimmed + + +def test_summarizer_failure_attended_choose_trim(tmp_path): + provider = CompactingProvider( + [AssistantTurn(text="done", finish_reason="stop")], summary_fails=99 + ) + engine = make_engine(tmp_path,provider, messages=long_history(), cap=400) + engine.is_attended = lambda: True + + async def asker(args, tool_call_id=None): + return {"answer": "Trim oldest 10%"} + + engine.question_asker = asker + collect(engine) + assert engine.compaction_state is not None and engine.compaction_state.trimmed + + +def test_raw_overflow_routes_into_compaction_and_retries(tmp_path): + # Trigger never fires (huge cap) — the provider 400 is the only signal. The engine + # must compact (force) and retry the call instead of surfacing the error. + provider = CompactingProvider( + [AssistantTurn(text="recovered", finish_reason="stop")], main_overflows=1 + ) + engine = make_engine(tmp_path,provider, messages=long_history(), cap=1_000_000) + events = collect(engine) + + assert any(e.type == EventType.COMPACTED for e in events) + assert not any(e.type == EventType.ERROR for e in events) + finals = [e for e in events if e.type == EventType.ASSISTANT_MESSAGE] + assert finals and finals[-1].data["text"] == "recovered" + assert provider.main_calls == 2 + + +def test_non_overflow_provider_errors_still_surface(tmp_path): + class FailingProvider(CompactingProvider): + def complete(self, *, model, messages, tools=None, **settings): + raise RuntimeError("rate limit exceeded") + + engine = make_engine(tmp_path,FailingProvider([]), messages=long_history(turns=1), cap=1_000_000) + events = collect(engine) + assert any(e.type == EventType.ERROR for e in events) + assert not any(e.type == EventType.COMPACTED for e in events) + + +def test_compaction_state_survives_save_and_rebuild(tmp_path): + from coworker.compaction import CompactionState + from coworker.server.manager import SessionManager + + class Provider(ProviderClient): + def complete(self, *, model, messages, tools=None, **settings): + return AssistantTurn(text="hi", finish_reason="stop") + + def capabilities(self, model): + return ModelCapabilities() + + mgr = SessionManager(workspace=tmp_path, provider=Provider()) + sid = "compact-persist" + engine = mgr.get_engine(sid, agent="cowork", workspace=str(tmp_path)) + assert callable(engine.compaction_settings) # live Settings getter is wired + assert engine.compaction_settings()["threshold_pct"] == 0.8 + + engine.messages += long_history(turns=3)[1:] + engine.compaction_state = CompactionState( + boundary_index=3, summary_text="the gist", working_state="", user_messages=["u"] + ) + mgr.save(sid, engine) + mgr._engines.pop(sid) + + rebuilt = mgr.get_engine(sid, agent="cowork", workspace=str(tmp_path)) + assert rebuilt.compaction_state == engine.compaction_state