flowent 0.2.1 → 0.2.3

package/backend/src/flowent/main.py

@@ -1,12 +1,14 @@
 import asyncio
+import copy
 import json
 import logging
 import os
-from collections.abc import AsyncIterator, Mapping, Sequence
+import time
+from collections.abc import AsyncIterator, Awaitable, Mapping, Sequence
 from contextlib import asynccontextmanager, suppress
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Literal
+from typing import Any, Literal
 from uuid import uuid4
 
 from fastapi import FastAPI, HTTPException, Query
@@ -85,10 +87,11 @@ logger = logging.getLogger("flowent.main")
 DEFAULT_STATIC_DIR = Path(__file__).parent / "static"
 COMPACTED_CONTEXT_MARKER = "Context compacted"
 OPTIMIZED_CONTEXT_MARKER = "Context optimized"
-DEFAULT_AUTO_COMPACT_TOKEN_LIMIT = 120_000
+DEFAULT_AUTO_COMPACT_CONTEXT_WINDOW_RATIO = 0.95
 AUTO_COMPACT_RETAINED_MESSAGE_TOKEN_BUDGET = 20_000
 APPROVAL_TRANSCRIPT_MESSAGE_LIMIT = 12
 APPROVAL_TRANSCRIPT_TEXT_LIMIT = 2_000
+WORKSPACE_PROGRESS_FLUSH_INTERVAL_SECONDS = 0.5
 
 
 class ProviderModelsRequest(BaseModel):
@@ -121,6 +124,14 @@ class WorkspaceRunResponse(BaseModel):
     run_id: str
 
 
+class WorkspaceClearResponse(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    active_run_id: str | None = None
+    messages: list[StoredMessage]
+    usage_info: TokenUsageInfo | None = None
+
+
 @dataclass
 class WorkspaceCompactTask:
     task: asyncio.Task[tuple[StoredMessage, TokenUsageInfo]]
@@ -172,10 +183,13 @@ class WritablePathListResponse(BaseModel):
 @dataclass
 class WorkspaceRun:
     condition: asyncio.Condition
+    active_output: Literal["text", "thinking"] | None = None
     discard_on_cancel: bool = False
     events: list[tuple[int, str, dict[str, object]]] = field(default_factory=list)
+    generation: int = 0
     id: str = field(default_factory=lambda: str(uuid4()))
     is_done: bool = False
+    latest_snapshot: StoredMessage | None = None
     task: asyncio.Task[None] | None = None
 
     @property
@@ -183,8 +197,20 @@ class WorkspaceRun:
         return self.events[-1][0] if self.events else 0
 
 
-def stream_event(event: str, data: dict[str, object]) -> str:
-    return f"event: {event}\ndata: {json.dumps(data)}\n\n"
+def stream_event(
+    event: str, data: dict[str, object], event_id: int | None = None
+) -> str:
+    id_line = f"id: {event_id}\n" if event_id is not None else ""
+    return f"{id_line}event: {event}\ndata: {json.dumps(data)}\n\n"
+
+
+def stream_message_data(
+    message: StoredMessage, active_output: Literal["text", "thinking"] | None = None
+) -> dict[str, object]:
+    data = {**message.model_dump(), "status": message.status}
+    if active_output is not None:
+        data["active_output"] = active_output
+    return data
 
 
 def append_or_replace_message(
@@ -196,6 +222,136 @@ def append_or_replace_message(
     ]
 
 
+def run_snapshot_data_at(
+    run: WorkspaceRun, event_index: int
+) -> dict[str, object] | None:
+    snapshot_event_index = 0
+    snapshot: dict[str, object] | None = None
+    for current_event_index, event, data in run.events:
+        if current_event_index > event_index:
+            break
+        if event != "snapshot":
+            if event == "start" and snapshot is None:
+                assistant_id = data.get("id")
+                if isinstance(assistant_id, str):
+                    snapshot_event_index = current_event_index
+                    snapshot = {
+                        "author": "assistant",
+                        "content": "",
+                        "groups": [],
+                        "id": assistant_id,
+                        "status": "running",
+                        "tools": [],
+                    }
+            continue
+        message = data.get("message")
+        if isinstance(message, dict):
+            snapshot_event_index = current_event_index
+            snapshot = copy.deepcopy(message)
+    if snapshot is None:
+        return None
+    for current_event_index, event, data in run.events:
+        if current_event_index <= snapshot_event_index:
+            continue
+        if current_event_index > event_index:
+            break
+        apply_stream_event_to_snapshot(snapshot, event, data)
+    return snapshot
+
+
+def apply_stream_event_to_snapshot(
+    snapshot: dict[str, object], event: str, data: dict[str, object]
+) -> None:
+    if event == "output_start":
+        snapshot.pop("active_output", None)
+        index = data.get("index")
+        if isinstance(index, int):
+            append_snapshot_group(snapshot, index)
+    if event == "delta":
+        append_snapshot_text(snapshot, str(data.get("content") or ""))
+    if event == "thinking_delta":
+        append_snapshot_thinking(snapshot, str(data.get("content") or ""))
+    if event == "output_done":
+        snapshot.pop("active_output", None)
+
+
+def snapshot_groups(snapshot: dict[str, object]) -> list[dict[str, object]]:
+    groups = snapshot.get("groups")
+    if not isinstance(groups, list):
+        groups = []
+        snapshot["groups"] = groups
+    return groups
+
+
+def append_snapshot_group(
+    snapshot: dict[str, object], index: int | None = None
+) -> None:
+    groups = snapshot_groups(snapshot)
+    assistant_id = str(snapshot.get("id") or "assistant")
+    group_index = index if index is not None else len(groups) + 1
+    group_id = f"{assistant_id}-group-{group_index}"
+    if groups and groups[-1].get("id") == group_id:
+        return
+    groups.append({"id": group_id, "items": []})
+
+
+def append_snapshot_text(snapshot: dict[str, object], content: str) -> None:
+    if not content:
+        return
+    snapshot["active_output"] = "text"
+    snapshot["content"] = f"{snapshot.get('content') or ''}{content}"
+    append_snapshot_item_content(snapshot, content, "text")
+
+
+def append_snapshot_thinking(snapshot: dict[str, object], content: str) -> None:
+    if not content:
+        return
+    snapshot["active_output"] = "thinking"
+    snapshot["thinking"] = f"{snapshot.get('thinking') or ''}{content}"
+    append_snapshot_item_content(snapshot, content, "thinking")
+
+
+def append_snapshot_item_content(
+    snapshot: dict[str, object], content: str, item_type: Literal["text", "thinking"]
+) -> None:
+    groups = snapshot_groups(snapshot)
+    if not groups:
+        append_snapshot_group(snapshot)
+    group = groups[-1]
+    items = group.get("items")
+    if not isinstance(items, list):
+        items = []
+        group["items"] = items
+    item = next(
+        (
+            current
+            for current in reversed(items)
+            if isinstance(current, dict) and current.get("type") == item_type
+        ),
+        None,
+    )
+    if item is None:
+        assistant_id = str(snapshot.get("id") or "assistant")
+        snapshot_item_count = 0
+        for current_group in groups:
+            current_items = current_group.get("items")
+            if not isinstance(current_items, list):
+                continue
+            snapshot_item_count += sum(
+                1
+                for current_item in current_items
+                if isinstance(current_item, dict)
+                and current_item.get("type") == item_type
+            )
+        item = {
+            "content": "",
+            "id": f"{assistant_id}-{item_type}-{snapshot_item_count + 1}",
+            "type": item_type,
+        }
+        items.append(item)
+    item["content"] = f"{item.get('content') or ''}{content}"
+
+
 USER_VISIBLE_RUN_ERROR_TITLE = "Request failed"
 USER_VISIBLE_RUN_ERROR_MESSAGE = "Check the model connection settings and try again."
 USER_VISIBLE_CONTEXT_OPTIMIZATION_ERROR_MESSAGE = "Context could not be optimized."
@@ -511,16 +667,22 @@ def is_context_marker(message: StoredMessage) -> bool:
     return message.content in {COMPACTED_CONTEXT_MARKER, OPTIMIZED_CONTEXT_MARKER}
 
 
-def auto_compact_token_limit() -> int:
+def auto_compact_token_limit(context_window: int) -> int:
     raw_limit = os.environ.get("FLOWENT_AUTO_COMPACT_TOKEN_LIMIT", "")
+    if not raw_limit:
+        return max(0, int(context_window * DEFAULT_AUTO_COMPACT_CONTEXT_WINDOW_RATIO))
     try:
         return max(0, int(raw_limit))
     except ValueError:
-        return DEFAULT_AUTO_COMPACT_TOKEN_LIMIT
+        return max(0, int(context_window * DEFAULT_AUTO_COMPACT_CONTEXT_WINDOW_RATIO))
 
 
-def should_auto_compact(messages: list[ChatMessage]) -> bool:
-    token_limit = auto_compact_token_limit()
+def should_auto_compact(
+    messages: list[ChatMessage],
+    *,
+    context_window: int,
+) -> bool:
+    token_limit = auto_compact_token_limit(context_window)
     if token_limit <= 0:
         return False
     return (
@@ -543,19 +705,40 @@ def usage_event_data(usage_info: TokenUsageInfo) -> dict[str, object]:
     return {"usage_info": usage_info.model_dump()}
 
 
+def update_context_usage_for_response(
+    usage_info: TokenUsageInfo | None,
+    *,
+    messages: Sequence[Mapping[str, object]],
+    output_content: str,
+    model_context_window: int,
+) -> TokenUsageInfo:
+    return recompute_context_usage(
+        usage_info,
+        estimated_token_usage_for_messages(
+            model_visible_messages_for_usage(messages),
+            output_content=output_content,
+        ).total_tokens,
+        model_context_window=model_context_window,
+    )
+
+
 def usage_info_for_model(
     usage_info: TokenUsageInfo | None,
-    model_name: str | None,
+    model_context_window: int,
 ) -> TokenUsageInfo | None:
     if usage_info is None:
         return None
-    return usage_info.model_copy(
-        update={"model_context_window": current_model_context_window(model_name)}
-    )
+    return usage_info.model_copy(update={"model_context_window": model_context_window})
+
+
+def context_window_for_settings(settings: StoredSettings) -> int:
+    if settings.context_window_limit is not None:
+        return settings.context_window_limit
+    return current_model_context_window(settings.selected_model)
 
 
 def state_with_current_model_context_window(state: StoredState) -> StoredState:
-    selected_model = state.settings.selected_model
+    model_context_window = context_window_for_settings(state.settings)
     return state.model_copy(
         update={
             "messages": [
@@ -563,7 +746,7 @@ def state_with_current_model_context_window(state: StoredState) -> StoredState:
                     update={
                         "usage_info": usage_info_for_model(
                             message.usage_info,
-                            selected_model,
+                            model_context_window,
                         )
                     }
                 )
@@ -571,7 +754,10 @@ def state_with_current_model_context_window(state: StoredState) -> StoredState:
                 else message
                 for message in state.messages
             ],
-            "usage_info": usage_info_for_model(state.usage_info, selected_model),
+            "usage_info": usage_info_for_model(
+                state.usage_info,
+                model_context_window,
+            ),
         }
     )
 
@@ -671,6 +857,7 @@ def create_app(
     telegram_bot_manager: TelegramBotManager | None = None
     workspace_runs: dict[str, WorkspaceRun] = {}
     active_workspace_run_id: str | None = None
+    workspace_generation = 0
     active_compact_task: WorkspaceCompactTask | None = None
 
     static_dir = frontend_static_directory().resolve(strict=False)
@@ -702,6 +889,7 @@ def create_app(
     async def save_context_checkpoint(
         *,
         connection: ProviderConnection,
+        context_window_limit: int,
         messages: list[StoredMessage],
         model_history: list[ChatMessage],
         marker_content: str,
@@ -723,12 +911,12 @@ def create_app(
             usage_info = append_token_usage(
                 usage_info,
                 compact_result.summary_usage,
-                model_context_window=current_model_context_window(connection.model),
+                model_context_window=context_window_limit,
             )
         usage_info = recompute_context_usage(
             usage_info,
             compact_result.token_after,
-            model_context_window=current_model_context_window(connection.model),
+            model_context_window=context_window_limit,
         )
         store.save_usage_info(usage_info)
         marker = StoredMessage(
@@ -767,16 +955,21 @@ def create_app(
     async def auto_compact_workspace_messages(
         *,
         connection: ProviderConnection,
+        context_window_limit: int,
         messages: list[StoredMessage],
         model_history: list[ChatMessage],
         source_message_id: str | None = None,
     ) -> tuple[StoredMessage, list[dict[str, object]], TokenUsageInfo] | None:
-        if not should_auto_compact(model_history):
+        if not should_auto_compact(
+            model_history,
+            context_window=context_window_limit,
+        ):
             return None
         logger.info("Workspace auto compact requested")
         try:
             return await save_context_checkpoint(
                 connection=connection,
+                context_window_limit=context_window_limit,
                 marker_content=OPTIMIZED_CONTEXT_MARKER,
                 messages=messages,
                 model_history=model_history,
@@ -790,6 +983,7 @@ def create_app(
     async def run_workspace_turn(content: str) -> StoredMessage:
         state = store.read_state()
         connection = selected_connection(state)
+        context_window_limit = context_window_for_settings(state.settings)
         user_message = StoredMessage(
             author="user",
             content=content,
@@ -807,6 +1001,7 @@ def create_app(
         ]
         auto_compaction = await auto_compact_workspace_messages(
             connection=connection,
+            context_window_limit=context_window_limit,
             messages=state.messages,
             model_history=model_history,
             source_message_id=None,
@@ -873,12 +1068,15 @@ def create_app(
             if event.event == "usage":
                 usage_data = event.data.get("usage")
                 if isinstance(usage_data, dict):
-                    usage_info = append_token_usage(
-                        store.read_usage_info(),
-                        TokenUsage.model_validate(usage_data),
-                        model_context_window=current_model_context_window(
-                            connection.model
+                    usage_info = update_context_usage_for_response(
+                        append_token_usage(
+                            store.read_usage_info(),
+                            TokenUsage.model_validate(usage_data),
+                            model_context_window=context_window_limit,
                         ),
+                        messages=request_messages,
+                        output_content=assistant_output.content,
+                        model_context_window=context_window_limit,
                     )
                     store.save_usage_info(usage_info)
                     turn_usage_info = usage_info
@@ -899,15 +1097,20 @@ def create_app(
 
         final_usage_info = turn_usage_info
         if final_usage_info is None:
-            final_usage_info = recompute_context_usage(
+            final_usage_info = update_context_usage_for_response(
                 store.read_usage_info(),
-                estimated_token_usage_for_messages(
-                    model_visible_messages_for_usage(request_messages),
-                    output_content=assistant_output.content,
-                ).total_tokens,
-                model_context_window=current_model_context_window(connection.model),
+                messages=request_messages,
+                output_content=assistant_output.content,
+                model_context_window=context_window_limit,
             )
-            store.save_usage_info(final_usage_info)
+        else:
+            final_usage_info = update_context_usage_for_response(
+                final_usage_info,
+                messages=request_messages,
+                output_content=assistant_output.content,
+                model_context_window=context_window_limit,
+            )
+        store.save_usage_info(final_usage_info)
 
         assistant_message = StoredMessage(
             author="assistant",
@@ -931,6 +1134,58 @@ def create_app(
         telegram_transport=telegram_transport,
     )
 
+    async def gather_shutdown_tasks(
+        label: str, tasks: Sequence[asyncio.Task[Any]]
+    ) -> None:
+        if not tasks:
+            return
+        results = await asyncio.gather(*tasks, return_exceptions=True)
+        for result in results:
+            if result is None or isinstance(result, asyncio.CancelledError):
+                continue
+            if isinstance(result, BaseException):
+                logger.error(
+                    "%s cleanup task failed",
+                    label,
+                    exc_info=(type(result), result, result.__traceback__),
+                )
+
+    async def stop_workspace_runs_for_shutdown() -> None:
+        tasks: list[asyncio.Task[None]] = []
+        for run in workspace_runs.values():
+            if run.task is None or run.task.done():
+                continue
+            run.task.cancel()
+            tasks.append(run.task)
+        await gather_shutdown_tasks("Workspace run", tasks)
+
+    async def stop_workspace_compact_for_shutdown() -> None:
+        nonlocal active_compact_task
+        if active_compact_task is None:
+            store.save_is_compacting(False)
+            return
+        task = active_compact_task.task
+        active_compact_task = None
+        if not task.done():
+            task.cancel()
+        await gather_shutdown_tasks("Workspace compact", [task])
+        store.save_is_compacting(False)
+
+    async def run_shutdown_step(label: str, cleanup: Awaitable[object]) -> None:
+        try:
+            await cleanup
+        except Exception:
+            logger.exception("%s cleanup failed during shutdown", label)
+
+    async def graceful_shutdown() -> None:
+        await run_shutdown_step("Workspace run", stop_workspace_runs_for_shutdown())
+        await run_shutdown_step(
+            "Workspace compact", stop_workspace_compact_for_shutdown()
+        )
+        if telegram_bot_manager is not None:
+            await run_shutdown_step("Telegram", telegram_bot_manager.stop_all())
+        await run_shutdown_step("MCP", mcp_manager.stop_all())
+
     @asynccontextmanager
     async def lifespan(app: FastAPI) -> AsyncIterator[None]:
         app.state.mcp_manager = mcp_manager
@@ -941,9 +1196,7 @@ def create_app(
         try:
             yield
         finally:
-            if telegram_bot_manager is not None:
-                await telegram_bot_manager.stop_all()
-            await mcp_manager.stop_all()
+            await graceful_shutdown()
 
     app = FastAPI(title="Flowent", lifespan=lifespan)
     app.state.mcp_manager = mcp_manager
@@ -1096,14 +1349,23 @@ def create_app(
     async def save_workspace_messages(
         request: WorkspaceMessagesRequest,
     ) -> WorkspaceMessagesRequest:
+        return WorkspaceMessagesRequest(messages=store.save_messages(request.messages))
+
+    @app.post("/api/workspace/clear")
+    async def clear_workspace() -> WorkspaceClearResponse:
         nonlocal active_workspace_run_id
-        if not request.messages:
-            run = active_workspace_run()
-            if run is not None and run.task is not None and not run.task.done():
+        nonlocal workspace_generation
+        workspace_generation += 1
+        for run in workspace_runs.values():
+            run.is_done = True
+            if run.task is not None and not run.task.done():
                 run.discard_on_cancel = True
                 run.task.cancel()
-                active_workspace_run_id = None
-        return WorkspaceMessagesRequest(messages=store.save_messages(request.messages))
+            async with run.condition:
+                run.condition.notify_all()
+        active_workspace_run_id = None
+        messages = store.save_messages([])
+        return WorkspaceClearResponse(messages=messages)
 
     async def append_run_event(
         run: WorkspaceRun, event: str, data: dict[str, object]
@@ -1112,15 +1374,42 @@ def create_app(
             run.events.append((run.latest_event_index + 1, event, data))
             run.condition.notify_all()
 
+    async def append_run_snapshot(run: WorkspaceRun, message: StoredMessage) -> None:
+        if message.author != "assistant":
+            return
+        run.latest_snapshot = message
+        await append_run_event(
+            run,
+            "snapshot",
+            {"message": stream_message_data(message, run.active_output)},
+        )
+
     def active_workspace_run() -> WorkspaceRun | None:
         if active_workspace_run_id is None:
             return None
-        return workspace_runs.get(active_workspace_run_id)
+        run = workspace_runs.get(active_workspace_run_id)
+        if run is None or run.is_done:
+            return None
+        return run
+
+    def has_active_workspace_run() -> bool:
+        return any(
+            not run.is_done and run.task is not None and not run.task.done()
+            for run in workspace_runs.values()
+        )
 
     def create_workspace_run(content: str) -> WorkspaceRun:
         nonlocal active_workspace_run_id
+        if has_active_workspace_run():
+            active_run = active_workspace_run()
+            raise HTTPException(
+                status_code=409,
+                detail="Response in progress",
+                headers={"X-Flowent-Run-Id": active_run.id if active_run else ""},
+            )
         state = store.read_state()
         connection = selected_connection(state)
+        context_window_limit = context_window_for_settings(state.settings)
 
         user_message = StoredMessage(
             author="user",
@@ -1129,7 +1418,10 @@ def create_app(
         )
         next_messages = [*state.messages, user_message]
         store.save_messages(next_messages)
-        run = WorkspaceRun(condition=asyncio.Condition())
+        run = WorkspaceRun(
+            condition=asyncio.Condition(),
+            generation=workspace_generation,
+        )
         workspace_runs[run.id] = run
         active_workspace_run_id = run.id
 
@@ -1143,9 +1435,17 @@ def create_app(
                 status="running",
             )
             assistant_output = AssistantOutputBuilder(assistant_message.id)
+            last_progress_flush_at = 0.0
 
-            def persist_assistant(status: str = "running") -> None:
+            def is_current_generation() -> bool:
+                return run.generation == workspace_generation
+
+            def update_assistant_message(
+                status: str = "running", *, persist: bool
+            ) -> StoredMessage | None:
                 nonlocal next_messages, assistant_message
+                if not is_current_generation() or run.discard_on_cancel:
+                    return None
                 assistant_message = StoredMessage(
                     author="assistant",
                     content=assistant_output.content,
@@ -1159,7 +1459,32 @@ def create_app(
                 next_messages = append_or_replace_message(
                     next_messages, assistant_message
                 )
-                store.upsert_message(assistant_message)
+                if persist:
+                    store.upsert_message(assistant_message)
+                return assistant_message
+
+            def persist_assistant(status: str = "running") -> StoredMessage | None:
+                nonlocal last_progress_flush_at
+                message = update_assistant_message(status, persist=True)
+                if status == "running" and message is not None:
+                    last_progress_flush_at = time.monotonic()
+                return message
+
+            def refresh_assistant(status: str = "running") -> StoredMessage | None:
+                return update_assistant_message(status, persist=False)
+
+            def persist_assistant_progress() -> StoredMessage | None:
+                nonlocal last_progress_flush_at
+                now = time.monotonic()
+                if (
+                    last_progress_flush_at > 0
+                    and now - last_progress_flush_at
+                    < WORKSPACE_PROGRESS_FLUSH_INTERVAL_SECONDS
+                ):
+                    refresh_assistant()
+                    return None
+                last_progress_flush_at = now
+                return update_assistant_message("running", persist=True)
 
             try:
                 current_tool_id: str | None = None
@@ -1176,6 +1501,7 @@ def create_app(
                 )
                 auto_compaction = await auto_compact_workspace_messages(
                     connection=connection,
+                    context_window_limit=context_window_limit,
                     messages=state.messages,
                     model_history=[
                         ChatMessage.model_validate(message)
@@ -1232,6 +1558,8 @@ def create_app(
                     conversation: Sequence[Mapping[str, object]],
                 ) -> AgentContextUpdate | None:
                     nonlocal next_messages
+                    if not is_current_generation() or run.discard_on_cancel:
+                        return None
                     assistant_snapshot = StoredMessage(
                         author="assistant",
                         content=assistant_output.content,
@@ -1267,6 +1595,7 @@ def create_app(
                             )
                     auto_result = await auto_compact_workspace_messages(
                         connection=connection,
+                        context_window_limit=context_window_limit,
                         messages=next_messages,
                         model_history=model_history,
                         source_message_id=assistant_snapshot.id,
@@ -1304,6 +1633,11 @@ def create_app(
                     messages=current_request_messages,
                     tool_runner=tool_runner,
                 ):
+                    if not is_current_generation() or run.discard_on_cancel:
+                        raise asyncio.CancelledError
+                    run_event_data = event.data
+                    should_append_run_event = event.event != "usage"
+                    snapshot_after_event: StoredMessage | None = None
                     if event.event == "start":
                         event_id = event.data.get("id")
                         if isinstance(event_id, str):
@@ -1311,20 +1645,24 @@ def create_app(
                                 update={"id": event_id}
                             )
                             assistant_output.set_assistant_id(event_id)
-                            persist_assistant()
+                            snapshot_after_event = persist_assistant()
                     if event.event == "output_start":
                         index = event.data.get("index")
                         if isinstance(index, int):
+                            run.active_output = None
                             assistant_output.start_group(index)
-                            persist_assistant()
+                            snapshot_after_event = persist_assistant()
+                    if event.event == "output_done":
+                        run.active_output = None
                     if event.event == "tool_start":
                         tool = event.data.get("tool")
                         if isinstance(tool, dict) and isinstance(tool.get("id"), str):
+                            run.active_output = None
                             current_tool_id = tool["id"]
                             assistant_output.start_tool(
                                 StoredToolItem.model_validate(tool)
                             )
-                            persist_assistant()
+                            snapshot_after_event = persist_assistant()
                     if event.event in {"tool_done", "tool_error"}:
                         tool_id = event.data.get("id")
                         if (
@@ -1335,34 +1673,37 @@ def create_app(
                                 None if current_tool_id == tool_id else current_tool_id
                             )
                             assistant_output.update_tool(tool_id, event.data)
-                            persist_assistant()
+                            snapshot_after_event = persist_assistant()
                     if event.event == "delta":
+                        run.active_output = "text"
                         assistant_output.append_text(
                             str(event.data.get("content") or "")
                         )
-                        persist_assistant()
+                        snapshot_after_event = persist_assistant_progress()
                     if event.event == "thinking_delta":
+                        run.active_output = "thinking"
                         assistant_output.append_thinking(
                             str(event.data.get("content") or "")
                         )
-                        persist_assistant()
+                        snapshot_after_event = persist_assistant_progress()
                     if event.event == "usage":
                         usage_data = event.data.get("usage")
                         if isinstance(usage_data, dict):
-                            usage_info = append_token_usage(
-                                store.read_usage_info(),
-                                TokenUsage.model_validate(usage_data),
-                                model_context_window=current_model_context_window(
-                                    connection.model
+                            usage_info = update_context_usage_for_response(
+                                append_token_usage(
+                                    store.read_usage_info(),
+                                    TokenUsage.model_validate(usage_data),
+                                    model_context_window=context_window_limit,
                                 ),
+                                messages=current_request_messages,
+                                output_content=assistant_output.content,
+                                model_context_window=context_window_limit,
                             )
                             store.save_usage_info(usage_info)
                             turn_usage_info = usage_info
-                            await append_run_event(
-                                run,
-                                "usage",
-                                usage_event_data(usage_info),
-                            )
+                            run_event_data = usage_event_data(usage_info)
+                            should_append_run_event = True
+                            snapshot_after_event = persist_assistant()
                     logger.log(
                         TRACE_LEVEL,
                         "Workspace stream event=%s data=%r",
@@ -1372,34 +1713,44 @@ def create_app(
                     if event.event == "done":
                         message = event.data.get("message")
                         if isinstance(message, dict):
+                            run.active_output = None
                             assistant_output.apply_done_message(message)
                             response_usage_info = store.read_usage_info()
                             final_usage_info = turn_usage_info
                             if final_usage_info is None:
-                                final_usage_info = recompute_context_usage(
+                                final_usage_info = update_context_usage_for_response(
                                     response_usage_info,
-                                    estimated_token_usage_for_messages(
-                                        model_visible_messages_for_usage(
-                                            current_request_messages
-                                        ),
-                                        output_content=assistant_output.content,
-                                    ).total_tokens,
-                                    model_context_window=current_model_context_window(
-                                        connection.model
-                                    ),
+                                    messages=current_request_messages,
+                                    output_content=assistant_output.content,
+                                    model_context_window=context_window_limit,
                                 )
-                                store.save_usage_info(final_usage_info)
-                            if final_usage_info == response_usage_info:
-                                assistant_message = assistant_message.model_copy(
-                                    update={"usage_info": final_usage_info}
+                            else:
+                                final_usage_info = update_context_usage_for_response(
+                                    final_usage_info,
+                                    messages=current_request_messages,
+                                    output_content=assistant_output.content,
+                                    model_context_window=context_window_limit,
                                 )
-                            persist_assistant("completed")
-                    if event.event != "usage":
-                        await append_run_event(run, event.event, event.data)
+                            store.save_usage_info(final_usage_info)
+                            snapshot_after_event = persist_assistant("completed")
+                            if snapshot_after_event is not None:
+                                run_event_data = {
+                                    "message": stream_message_data(snapshot_after_event)
+                                }
+                    if event.event == "done" and snapshot_after_event is not None:
+                        await append_run_snapshot(run, snapshot_after_event)
+                        await append_run_event(run, event.event, run_event_data)
+                    else:
+                        if should_append_run_event:
+                            await append_run_event(run, event.event, run_event_data)
+                        if snapshot_after_event is not None:
+                            await append_run_snapshot(run, snapshot_after_event)
             except asyncio.CancelledError:
                 logger.info("Workspace run stopped")
                 if not run.discard_on_cancel:
-                    persist_assistant("interrupted")
+                    interrupted_snapshot = persist_assistant("interrupted")
+                    if interrupted_snapshot is not None:
+                        await append_run_snapshot(run, interrupted_snapshot)
                     await append_run_event(
                         run,
                         "error",
@@ -1423,7 +1774,9 @@ def create_app(
                         str(error) or EMPTY_MODEL_RESPONSE_DETAIL,
                     )
                 )
-                persist_assistant("failed")
+                failed_snapshot = persist_assistant("failed")
+                if failed_snapshot is not None:
+                    await append_run_snapshot(run, failed_snapshot)
                 await append_run_event(run, "error", run_error_event_data(error_item))
             finally:
                 run.is_done = True
@@ -1436,9 +1789,16 @@ def create_app(
         return run
 
     async def workspace_run_stream(
-        run: WorkspaceRun, after: int = 0
+        run: WorkspaceRun, after: int = 0, include_snapshots: bool = True
     ) -> AsyncIterator[str]:
         next_event_index = after + 1
+        reconnect_snapshot = run_snapshot_data_at(run, after) if after > 0 else None
+        if include_snapshots and reconnect_snapshot is not None:
+            yield stream_event(
+                "snapshot",
+                {"message": reconnect_snapshot},
+                event_id=after,
+            )
         while True:
             async with run.condition:
 
@@ -1452,7 +1812,9 @@ def create_app(
 
             for index, event, data in events:
                 next_event_index = index + 1
-                yield stream_event(event, data)
+                if event == "snapshot" and not include_snapshots:
+                    continue
+                yield stream_event(event, data, event_id=index)
                 if event in {"done", "error"}:
                     return
 
@@ -1498,6 +1860,7 @@ def create_app(
             *,
             checkpoint: StoredCompactionCheckpoint | None,
             connection: ProviderConnection,
+            context_window_limit: int,
             state: StoredState,
         ) -> tuple[StoredMessage, TokenUsageInfo]:
             logger.info("Workspace compact requested")
@@ -1513,6 +1876,7 @@ def create_app(
 
                 marker, _, usage_info = await save_context_checkpoint(
                     connection=connection,
+                    context_window_limit=context_window_limit,
                     marker_content=COMPACTED_CONTEXT_MARKER,
                     messages=state.messages,
                     model_history=model_history,
@@ -1551,12 +1915,14 @@ def create_app(
                 )
             state = store.read_state()
             connection = selected_connection(state)
+            context_window_limit = context_window_for_settings(state.settings)
             checkpoint = store.read_active_compaction_checkpoint()
             store.save_is_compacting(True)
             compact_task = asyncio.create_task(
                 run_manual_compact(
                     checkpoint=checkpoint,
                     connection=connection,
+                    context_window_limit=context_window_limit,
                     state=state,
                 )
             )
@@ -1596,7 +1962,7 @@ def create_app(
         logger.log(TRACE_LEVEL, "Workspace user content=%r", request.content)
         run = create_workspace_run(request.content)
         return StreamingResponse(
-            workspace_run_stream(run),
+            workspace_run_stream(run, include_snapshots=False),
             media_type="text/event-stream",
         )