flowent 0.2.0 → 0.2.2

package/backend/src/flowent/main.py

@@ -3,7 +3,7 @@ import json
 import logging
 import os
 from collections.abc import AsyncIterator, Mapping, Sequence
-from contextlib import asynccontextmanager
+from contextlib import asynccontextmanager, suppress
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Literal
@@ -70,6 +70,14 @@ from flowent.storage import (
     StoredWritablePath,
 )
 from flowent.tools import ToolContext
+from flowent.usage import (
+    TokenUsage,
+    TokenUsageInfo,
+    append_token_usage,
+    current_model_context_window,
+    estimated_token_usage_for_messages,
+    recompute_context_usage,
+)
 
 logger = logging.getLogger("flowent.main")
 
@@ -77,7 +85,7 @@ 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
@@ -113,10 +121,17 @@ class WorkspaceRunResponse(BaseModel):
     run_id: str
 
 
-class WorkspaceCompactResponse(BaseModel):
+class WorkspaceClearResponse(BaseModel):
     model_config = ConfigDict(extra="forbid")
 
-    message: StoredMessage
+    active_run_id: str | None = None
+    messages: list[StoredMessage]
+    usage_info: TokenUsageInfo | None = None
+
+
+@dataclass
+class WorkspaceCompactTask:
+    task: asyncio.Task[tuple[StoredMessage, TokenUsageInfo]]
 
 
 class AboutResponse(BaseModel):
@@ -167,8 +182,10 @@ class WorkspaceRun:
     condition: asyncio.Condition
     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
@@ -176,8 +193,15 @@ 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) -> dict[str, object]:
+    return {**message.model_dump(), "status": message.status}
 
 
 def append_or_replace_message(
@@ -189,6 +213,18 @@ def append_or_replace_message(
     ]
 
 
+def run_snapshot_data_at(
+    run: WorkspaceRun, event_index: int
+) -> dict[str, object] | None:
+    for current_event_index, event, data in reversed(run.events):
+        if current_event_index > event_index or event != "snapshot":
+            continue
+        message = data.get("message")
+        if isinstance(message, dict):
+            return message
+    return None
+
+
 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."
@@ -504,16 +540,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 (
@@ -522,6 +564,77 @@ def should_auto_compact(messages: list[ChatMessage]) -> bool:
     )
 
 
+def model_visible_messages_for_usage(
+    messages: Sequence[Mapping[str, object]],
+) -> list[dict[str, object]]:
+    return [
+        dict(message)
+        for message in messages
+        if message.get("role") in {"system", "user", "assistant", "tool"}
+    ]
+
+
+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_context_window: int,
+) -> TokenUsageInfo | None:
+    if usage_info is None:
+        return None
+    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:
+    model_context_window = context_window_for_settings(state.settings)
+    return state.model_copy(
+        update={
+            "messages": [
+                message.model_copy(
+                    update={
+                        "usage_info": usage_info_for_model(
+                            message.usage_info,
+                            model_context_window,
+                        )
+                    }
+                )
+                if message.usage_info is not None
+                else message
+                for message in state.messages
+            ],
+            "usage_info": usage_info_for_model(
+                state.usage_info,
+                model_context_window,
+            ),
+        }
+    )
+
+
 def workspace_chat_messages(
     messages: list[StoredMessage],
     compacted_context: str = "",
@@ -617,6 +730,8 @@ 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)
     logger.debug("Flowent app created serve_frontend=%s", serve_frontend)
@@ -647,17 +762,13 @@ def create_app(
     async def save_context_checkpoint(
         *,
         connection: ProviderConnection,
+        context_window_limit: int,
         messages: list[StoredMessage],
         model_history: list[ChatMessage],
         marker_content: str,
         source_message_id: str | None = None,
         trigger: Literal["manual", "auto"],
-    ) -> tuple[StoredMessage, list[dict[str, object]]]:
-        marker = StoredMessage(
-            author="system",
-            content=marker_content,
-            id=str(uuid4()),
-        )
+    ) -> tuple[StoredMessage, list[dict[str, object]], TokenUsageInfo]:
         compact_result = await compact_provider.compact(
             connection,
             CompactInput(
@@ -668,6 +779,25 @@ def create_app(
             ),
             completion=chat_completion,
         )
+        usage_info = store.read_usage_info()
+        if compact_result.summary_usage is not None:
+            usage_info = append_token_usage(
+                usage_info,
+                compact_result.summary_usage,
+                model_context_window=context_window_limit,
+            )
+        usage_info = recompute_context_usage(
+            usage_info,
+            compact_result.token_after,
+            model_context_window=context_window_limit,
+        )
+        store.save_usage_info(usage_info)
+        marker = StoredMessage(
+            author="system",
+            content=marker_content,
+            id=str(uuid4()),
+            usage_info=usage_info,
+        )
         store.save_compaction_checkpoint(
             StoredCompactionCheckpoint(
                 id=str(uuid4()),
@@ -689,23 +819,30 @@ def create_app(
             compact_result.token_after,
         )
         logger.log(TRACE_LEVEL, "Workspace compact summary=%r", compact_result.summary)
-        return marker, [
-            message.model_dump() for message in compact_result.replacement_history
-        ]
+        return (
+            marker,
+            [message.model_dump() for message in compact_result.replacement_history],
+            usage_info,
+        )
 
     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]]] | None:
-        if not should_auto_compact(model_history):
+    ) -> tuple[StoredMessage, list[dict[str, object]], TokenUsageInfo] | None:
+        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,
@@ -719,6 +856,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,
@@ -736,17 +874,19 @@ 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,
         )
         if auto_compaction is not None:
-            marker, _ = auto_compaction
+            marker, _, _ = auto_compaction
             next_messages = [*state.messages, marker, user_message]
             store.save_messages(next_messages)
         request_messages = request_messages_for_content(state, next_messages, content)
         assistant_id = str(uuid4())
         assistant_output = AssistantOutputBuilder(assistant_id)
+        turn_usage_info: TokenUsageInfo | None = None
 
         async def review_tool_approval(request: ApprovalReviewRequest):
             return await review_approval_request(
@@ -798,6 +938,21 @@ def create_app(
                 assistant_output.append_text(str(event.data.get("content") or ""))
             if event.event == "thinking_delta":
                 assistant_output.append_thinking(str(event.data.get("content") or ""))
+            if event.event == "usage":
+                usage_data = event.data.get("usage")
+                if isinstance(usage_data, dict):
+                    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
             if event.event == "tool_start":
                 tool = event.data.get("tool")
                 if isinstance(tool, dict) and isinstance(tool.get("id"), str):
@@ -813,6 +968,23 @@ def create_app(
                     assistant_output.set_assistant_id(assistant_id)
                     assistant_output.apply_done_message(message)
 
+        final_usage_info = turn_usage_info
+        if final_usage_info is None:
+            final_usage_info = update_context_usage_for_response(
+                store.read_usage_info(),
+                messages=request_messages,
+                output_content=assistant_output.content,
+                model_context_window=context_window_limit,
+            )
+        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",
             content=assistant_output.content,
@@ -821,6 +993,7 @@ def create_app(
             status="completed",
             thinking=assistant_output.thinking,
             tools=list(assistant_output.tools.values()),
+            usage_info=final_usage_info,
         )
         store.save_messages([*next_messages, assistant_message])
         return assistant_message
@@ -858,7 +1031,7 @@ def create_app(
 
     @app.get("/api/state")
     async def app_state() -> StoredState:
-        state = store.read_state()
+        state = state_with_current_model_context_window(store.read_state())
         active_run = (
             workspace_runs.get(active_workspace_run_id)
             if active_workspace_run_id
@@ -999,14 +1172,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]
@@ -1015,15 +1197,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)},
+        )
+
     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",
@@ -1032,7 +1241,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
 
@@ -1047,8 +1259,13 @@ def create_app(
             )
             assistant_output = AssistantOutputBuilder(assistant_message.id)
 
-            def persist_assistant(status: str = "running") -> None:
+            def is_current_generation() -> bool:
+                return run.generation == workspace_generation
+
+            def persist_assistant(status: str = "running") -> 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,
@@ -1057,14 +1274,17 @@ def create_app(
                     status=status,
                     thinking=assistant_output.thinking,
                     tools=list(assistant_output.tools.values()),
+                    usage_info=store.read_usage_info(),
                 )
                 next_messages = append_or_replace_message(
                     next_messages, assistant_message
                 )
                 store.upsert_message(assistant_message)
+                return assistant_message
 
             try:
                 current_tool_id: str | None = None
+                turn_usage_info: TokenUsageInfo | None = None
                 current_request_messages = request_messages_for_content(
                     state,
                     next_messages,
@@ -1077,6 +1297,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)
@@ -1085,13 +1306,16 @@ def create_app(
                     source_message_id=None,
                 )
                 if auto_compaction is not None:
-                    marker, _ = auto_compaction
+                    marker, _, usage_info = auto_compaction
                     next_messages = [*state.messages, marker, user_message]
                     store.save_messages(next_messages)
                     await append_run_event(
                         run,
                         "context_optimized",
-                        {"message": marker.model_dump()},
+                        {
+                            "message": marker.model_dump(),
+                            **usage_event_data(usage_info),
+                        },
                     )
                     current_request_messages = request_messages_for_content(
                         state,
@@ -1130,6 +1354,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,
@@ -1138,6 +1364,7 @@ def create_app(
                         status="running",
                         thinking=assistant_output.thinking,
                         tools=list(assistant_output.tools.values()),
+                        usage_info=store.read_usage_info(),
                     )
                     model_history: list[ChatMessage] = []
                     for message in conversation:
@@ -1164,13 +1391,17 @@ 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,
                     )
                     if auto_result is None:
                         return None
-                    marker, replacement_history = auto_result
+                    marker, replacement_history, usage_info = auto_result
+                    assistant_snapshot = assistant_snapshot.model_copy(
+                        update={"usage_info": usage_info}
+                    )
                     next_messages = append_or_replace_message(
                         [*next_messages, marker], assistant_snapshot
                     )
@@ -1181,7 +1412,10 @@ def create_app(
                     ]
                     return AgentContextUpdate(
                         conversation=compacted_conversation,
-                        message=marker.model_dump(),
+                        message={
+                            **marker.model_dump(),
+                            "usage_info": usage_info.model_dump(),
+                        },
                     )
 
                 async for event in run_agent_stream(
@@ -1195,6 +1429,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):
@@ -1202,12 +1441,12 @@ 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):
                             assistant_output.start_group(index)
-                            persist_assistant()
+                            snapshot_after_event = persist_assistant()
                     if event.event == "tool_start":
                         tool = event.data.get("tool")
                         if isinstance(tool, dict) and isinstance(tool.get("id"), str):
@@ -1215,7 +1454,7 @@ def create_app(
                             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 (
@@ -1226,17 +1465,35 @@ 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":
                         assistant_output.append_text(
                             str(event.data.get("content") or "")
                         )
-                        persist_assistant()
+                        snapshot_after_event = persist_assistant()
                     if event.event == "thinking_delta":
                         assistant_output.append_thinking(
                             str(event.data.get("content") or "")
                         )
-                        persist_assistant()
+                        snapshot_after_event = persist_assistant()
+                    if event.event == "usage":
+                        usage_data = event.data.get("usage")
+                        if isinstance(usage_data, dict):
+                            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
+                            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",
@@ -1247,12 +1504,42 @@ def create_app(
                         message = event.data.get("message")
                         if isinstance(message, dict):
                             assistant_output.apply_done_message(message)
-                            persist_assistant("completed")
-                    await append_run_event(run, event.event, event.data)
+                            response_usage_info = store.read_usage_info()
+                            final_usage_info = turn_usage_info
+                            if final_usage_info is None:
+                                final_usage_info = update_context_usage_for_response(
+                                    response_usage_info,
+                                    messages=current_request_messages,
+                                    output_content=assistant_output.content,
+                                    model_context_window=context_window_limit,
+                                )
+                            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,
+                                )
+                            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",
@@ -1276,7 +1563,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
@@ -1289,9 +1578,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:
 
@@ -1305,7 +1601,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
 
@@ -1343,47 +1641,105 @@ def create_app(
             run.task.cancel()
         return {"ok": True}
 
-    @app.post("/api/workspace/compact")
-    async def compact_workspace() -> WorkspaceCompactResponse:
-        if active_workspace_run() is not None:
-            raise HTTPException(
-                status_code=409,
-                detail="Compact is unavailable while Flowent is responding.",
+    @app.post("/api/workspace/compact", response_class=StreamingResponse)
+    async def compact_workspace() -> StreamingResponse:
+        nonlocal active_compact_task
+
+        async def run_manual_compact(
+            *,
+            checkpoint: StoredCompactionCheckpoint | None,
+            connection: ProviderConnection,
+            context_window_limit: int,
+            state: StoredState,
+        ) -> tuple[StoredMessage, TokenUsageInfo]:
+            logger.info("Workspace compact requested")
+            try:
+                model_history = [
+                    *runtime_context_messages(cwd, state.settings.agent_prompt),
+                    *workspace_chat_messages(
+                        state.messages,
+                        store.read_compacted_context(),
+                        checkpoint,
+                    ),
+                ]
+
+                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,
+                    source_message_id=None,
+                    trigger="manual",
+                )
+                store.save_messages([*state.messages, marker])
+                logger.info("Workspace compact completed")
+                return marker, usage_info
+            except Exception:
+                logger.exception("Workspace compact failed")
+                raise
+            finally:
+                store.save_is_compacting(False)
+
+        def clear_active_compact_task(
+            task: asyncio.Task[tuple[StoredMessage, TokenUsageInfo]],
+        ) -> None:
+            nonlocal active_compact_task
+            if active_compact_task is not None and active_compact_task.task is task:
+                active_compact_task = None
+            with suppress(asyncio.CancelledError):
+                task.exception()
+
+        if active_compact_task is not None:
+            if not active_compact_task.task.done():
+                compact_task = active_compact_task.task
+            else:
+                active_compact_task = None
+
+        if active_compact_task is None:
+            if active_workspace_run() is not None:
+                raise HTTPException(
+                    status_code=409,
+                    detail="Compact is unavailable while Flowent is responding.",
+                )
+            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,
+                )
             )
-        logger.info("Workspace compact requested")
-        state = store.read_state()
-        connection = selected_connection(state)
-        checkpoint = store.read_active_compaction_checkpoint()
-        model_history = [
-            *runtime_context_messages(cwd, state.settings.agent_prompt),
-            *workspace_chat_messages(
-                state.messages,
-                store.read_compacted_context(),
-                checkpoint,
-            ),
-        ]
+            compact_task.add_done_callback(clear_active_compact_task)
+            active_compact_task = WorkspaceCompactTask(task=compact_task)
 
-        try:
-            marker, _ = await save_context_checkpoint(
-                connection=connection,
-                marker_content=COMPACTED_CONTEXT_MARKER,
-                messages=state.messages,
-                model_history=model_history,
-                source_message_id=None,
-                trigger="manual",
+        async def compact_workspace_stream() -> AsyncIterator[str]:
+            try:
+                marker, usage_info = await asyncio.shield(compact_task)
+            except Exception:
+                yield stream_event(
+                    "error",
+                    {"message": "Context could not be compacted."},
+                )
+                return
+
+            marker_data = marker.model_dump()
+            yield stream_event("usage", usage_event_data(usage_info))
+            yield stream_event(
+                "context_optimized",
+                {"message": marker_data, **usage_event_data(usage_info)},
             )
-        except HTTPException:
-            raise
-        except Exception as error:
-            logger.exception("Workspace compact failed")
-            raise HTTPException(
-                status_code=500,
-                detail="Context could not be compacted.",
-            ) from error
+            yield stream_event("done", {"message": marker_data})
 
-        store.save_messages([*state.messages, marker])
-        logger.info("Workspace compact completed")
-        return WorkspaceCompactResponse(message=marker)
+        return StreamingResponse(
+            compact_workspace_stream(),
+            media_type="text/event-stream",
+        )
 
     @app.post("/api/workspace/respond")
     async def respond_to_workspace(
@@ -1395,7 +1751,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",
         )