flowent 0.3.1 → 0.3.2

package/backend/src/flowent/tools.py

@@ -11,7 +11,7 @@ from dataclasses import dataclass
 from pathlib import Path
 from uuid import uuid4
 
-from pydantic import BaseModel, ConfigDict
+from pydantic import BaseModel, ConfigDict, Field
 
 from flowent.network import flowent_user_agent
 from flowent.patch import affected_paths
@@ -23,8 +23,7 @@ from flowent.system_tools import ensure_ripgrep_available
 class ToolResult(BaseModel):
     model_config = ConfigDict(extra="forbid")
 
-    content: str
-    data: dict[str, object] = {}
+    result: dict[str, object] = Field(default_factory=dict)
     ok: bool = True
     title: str
 
@@ -35,6 +34,46 @@ class ToolContext:
     web_searcher: Callable[[str], Sequence[dict[str, str]]] | None = None
 
 
+def text_tool_result(text: str, **metadata: object) -> dict[str, object]:
+    return {"type": "text", "text": text, **metadata}
+
+
+def command_tool_result(
+    *,
+    command: str,
+    exit_code: int,
+    stderr: str,
+    stdout: str,
+) -> dict[str, object]:
+    return {
+        "type": "command",
+        "command": command,
+        "exit_code": exit_code,
+        "stderr": stderr,
+        "stdout": stdout,
+        "output": stdout or stderr,
+    }
+
+
+def tool_result_model_content(result: ToolResult | dict[str, object]) -> str:
+    payload = result.result if isinstance(result, ToolResult) else result
+    result_type = payload.get("type")
+    if result_type == "command":
+        output = str(payload.get("output") or "")
+        metadata: dict[str, object] = {}
+        if "exit_code" in payload:
+            metadata["exit_code"] = payload["exit_code"]
+        return json.dumps(
+            {"output": output, "metadata": metadata},
+            ensure_ascii=False,
+        )
+    for key in ("text", "output"):
+        value = payload.get(key)
+        if value is not None:
+            return str(value)
+    return json.dumps(payload, ensure_ascii=False)
+
+
 def tool_specs() -> list[dict[str, object]]:
     return [
         {
@@ -222,7 +261,7 @@ def run_tool(
         title = (
             "Edit failed" if name == "apply_patch" else tool_call_title(name, arguments)
         )
-        return ToolResult(content=str(error), ok=False, title=title)
+        return ToolResult(result=text_tool_result(str(error)), ok=False, title=title)
 
 
 async def run_tool_async(
@@ -238,7 +277,7 @@ async def run_tool_async(
         title = (
             "Edit failed" if name == "apply_patch" else tool_call_title(name, arguments)
         )
-        return ToolResult(content=str(error), ok=False, title=title)
+        return ToolResult(result=text_tool_result(str(error)), ok=False, title=title)
 
 
 def integer_argument(arguments: dict[str, object], name: str, default: int) -> int:
@@ -266,7 +305,10 @@ def read_file(arguments: dict[str, object], context: ToolContext) -> ToolResult:
     lines = path.read_text(errors="replace").splitlines()
     selected = lines[offset : offset + limit]
     content = "\n".join(selected)
-    return ToolResult(content=content, data={"path": str(path)}, title=f"Read {path}")
+    return ToolResult(
+        result=text_tool_result(content, path=str(path)),
+        title=f"Read {path}",
+    )
 
 
 def list_dir(arguments: dict[str, object], context: ToolContext) -> ToolResult:
@@ -279,7 +321,8 @@ def list_dir(arguments: dict[str, object], context: ToolContext) -> ToolResult:
         f"{entry.name}/" if entry.is_dir() else entry.name for entry in entries[:limit]
     ]
     return ToolResult(
-        content="\n".join(rendered), data={"path": str(path)}, title=f"Listed {path}"
+        result=text_tool_result("\n".join(rendered), path=str(path)),
+        title=f"Listed {path}",
     )
 
 
@@ -297,8 +340,7 @@ def grep_files(arguments: dict[str, object], context: ToolContext) -> ToolResult
     )
     output = completed.stdout or completed.stderr
     return ToolResult(
-        content=output[:20000],
-        data={"path": str(path), "pattern": pattern},
+        result=text_tool_result(output[:20000], path=str(path), pattern=pattern),
         title=f"Searched {pattern}",
     )
 
@@ -317,8 +359,11 @@ def apply_patch_tool(arguments: dict[str, object], context: ToolContext) -> Tool
         raise SandboxError(tool_failure_content(result))
     data = json.loads(result.stdout or "{}")
     return ToolResult(
-        content=result.stdout,
-        data=data if isinstance(data, dict) else {},
+        result={
+            "type": "patch",
+            "output": result.stdout,
+            **(data if isinstance(data, dict) else {}),
+        },
         title=patch_title_from_result(data),
     )
 
@@ -339,8 +384,11 @@ async def apply_patch_tool_async(
         raise SandboxError(tool_failure_content(result))
     data = json.loads(result.stdout or "{}")
     return ToolResult(
-        content=result.stdout,
-        data=data if isinstance(data, dict) else {},
+        result={
+            "type": "patch",
+            "output": result.stdout,
+            **(data if isinstance(data, dict) else {}),
+        },
         title=patch_title_from_result(data),
     )
 
@@ -388,15 +436,13 @@ def shell_command(arguments: dict[str, object], context: ToolContext) -> ToolRes
         invocation.args, env=invocation.env, timeout_seconds=timeout_seconds
     )
     ok = result.exit_code == 0
-    content = result.stdout or result.stderr
     return ToolResult(
-        content=content,
-        data={
-            "command": command,
-            "exit_code": result.exit_code,
-            "stderr": result.stderr,
-            "stdout": result.stdout,
-        },
+        result=command_tool_result(
+            command=command,
+            exit_code=result.exit_code,
+            stderr=result.stderr,
+            stdout=result.stdout,
+        ),
         ok=ok,
         title=f"Ran {command}",
     )
@@ -412,15 +458,13 @@ async def shell_command_async(
         invocation.args, env=invocation.env, timeout_seconds=timeout_seconds
     )
     ok = result.exit_code == 0
-    content = result.stdout or result.stderr
     return ToolResult(
-        content=content,
-        data={
-            "command": command,
-            "exit_code": result.exit_code,
-            "stderr": result.stderr,
-            "stdout": result.stdout,
-        },
+        result=command_tool_result(
+            command=command,
+            exit_code=result.exit_code,
+            stderr=result.stderr,
+            stdout=result.stdout,
+        ),
         ok=ok,
         title=f"Ran {command}",
     )
@@ -430,8 +474,11 @@ def update_plan(arguments: dict[str, object]) -> ToolResult:
     items = arguments.get("items", [])
     content = json.dumps(items, ensure_ascii=False)
     return ToolResult(
-        content=content,
-        data={"items": items if isinstance(items, list) else []},
+        result={
+            "type": "plan",
+            "items": items if isinstance(items, list) else [],
+            "output": content,
+        },
         title="Updated plan",
     )
 
@@ -473,8 +520,12 @@ def web_search(arguments: dict[str, object], context: ToolContext) -> ToolResult
         for result in results
     )
     return ToolResult(
-        content=content or "No results.",
-        data={"query": query, "results": results},
+        result={
+            "type": "web_search",
+            "output": content or "No results.",
+            "query": query,
+            "results": results,
+        },
         title=f"Searched web for {query}",
     )
 

package/backend/src/flowent/usage.py

@@ -148,6 +148,9 @@ def current_model_context_window(model_name: str | None = None) -> int:
 
 def model_context_window_for(model_name: str | None = None) -> int:
     candidates = normalized_model_name_candidates(model_name)
+    metadata_context_window = litellm_input_context_window_for(candidates)
+    if metadata_context_window is not None:
+        return metadata_context_window
     for candidate in candidates:
         context_window = MODEL_CONTEXT_WINDOWS.get(candidate)
         if context_window is not None:
@@ -159,6 +162,22 @@ def model_context_window_for(model_name: str | None = None) -> int:
     return DEFAULT_MODEL_CONTEXT_WINDOW
 
 
+def litellm_input_context_window_for(candidates: Sequence[str]) -> int | None:
+    try:
+        from litellm import model_cost
+    except Exception:
+        return None
+
+    for candidate in candidates:
+        metadata = model_cost.get(candidate)
+        if metadata is None:
+            continue
+        context_window = first_int_value(value_at(metadata, "max_input_tokens"))
+        if context_window is not None and context_window > 0:
+            return context_window
+    return None
+
+
 def normalized_model_name_candidates(model_name: str | None) -> tuple[str, ...]:
     if model_name is None:
         return ()
@@ -261,6 +280,53 @@ def estimated_token_usage_for_messages(
     )
 
 
+def estimated_token_usage_for_request(
+    messages: Sequence[Mapping[str, object]],
+    *,
+    output_content: str = "",
+    tools: Sequence[Mapping[str, object]] = (),
+) -> TokenUsage:
+    message_usage = estimated_token_usage_for_messages(
+        messages,
+        output_content=output_content,
+    )
+    tool_tokens = sum(
+        approximate_token_count(json.dumps(tool, ensure_ascii=False)) for tool in tools
+    )
+    input_tokens = message_usage.input_tokens + tool_tokens
+    return TokenUsage(
+        input_tokens=input_tokens,
+        output_tokens=message_usage.output_tokens,
+        total_tokens=input_tokens + message_usage.output_tokens,
+    )
+
+
+def full_context_usage(
+    usage_info: TokenUsageInfo | None,
+    *,
+    model_context_window: int,
+) -> TokenUsageInfo:
+    info = usage_info or TokenUsageInfo(model_context_window=model_context_window)
+    return TokenUsageInfo(
+        total_token_usage=info.total_token_usage,
+        last_token_usage=TokenUsage(total_tokens=max(0, model_context_window)),
+        model_context_window=model_context_window,
+    )
+
+
+def is_context_window_error(error: BaseException) -> bool:
+    message = str(error).lower()
+    return any(
+        marker in message
+        for marker in (
+            "context window",
+            "context_length_exceeded",
+            "maximum context length",
+            "too many tokens",
+        )
+    )
+
+
 def estimate_mapping_message_tokens(message: Mapping[str, object]) -> int:
     total = approximate_token_count(string_content(message.get("content")))
     tool_calls = message.get("tool_calls")

package/backend/src/flowent/workspace/context.py

@@ -12,10 +12,11 @@ from flowent.storage import (
     StoredSettings,
     StoredState,
 )
+from flowent.tools import tool_result_model_content
 from flowent.usage import (
     TokenUsageInfo,
     current_model_context_window,
-    estimated_token_usage_for_messages,
+    estimated_token_usage_for_request,
     recompute_context_usage,
 )
 from flowent.workspace.output import error_context_summary, message_error_items
@@ -51,13 +52,15 @@ def should_auto_compact(
     messages: Sequence[ChatMessage | Mapping[str, object]],
     *,
     context_window: int,
+    tools: Sequence[Mapping[str, object]] = (),
 ) -> bool:
     token_limit = auto_compact_token_limit(context_window)
     if token_limit <= 0:
         return False
     return (
-        estimated_token_usage_for_messages(
-            model_request_messages_data(messages)
+        estimated_token_usage_for_request(
+            model_request_messages_data(messages),
+            tools=tools,
         ).total_tokens
         >= token_limit
     )
@@ -83,17 +86,19 @@ def update_context_usage_for_response(
     messages: Sequence[Mapping[str, object]],
     output_content: str,
     output_tools: Sequence[Mapping[str, object]] = (),
+    request_tools: Sequence[Mapping[str, object]] = (),
     model_context_window: int,
 ) -> TokenUsageInfo:
     return recompute_context_usage(
         usage_info,
-        estimated_token_usage_for_messages(
+        estimated_token_usage_for_request(
             [
                 *model_visible_messages_for_usage(messages),
                 *model_visible_response_messages_for_usage(
                     output_content, output_tools
                 ),
             ],
+            tools=request_tools,
         ).total_tokens,
         model_context_window=model_context_window,
     )
@@ -107,6 +112,8 @@ def model_visible_response_messages_for_usage(
     for index, tool in enumerate(output_tools):
         tool_id = str(tool.get("id") or f"call_{index}")
         arguments = tool.get("arguments")
+        result_payload = tool.get("result")
+        tool_result = result_payload if isinstance(result_payload, dict) else {}
         visible_messages.append(
             {
                 "role": "assistant",
@@ -130,7 +137,7 @@ def model_visible_response_messages_for_usage(
             {
                 "role": "tool",
                 "tool_call_id": tool_id,
-                "content": str(tool.get("content") or ""),
+                "content": tool_result_model_content(tool_result),
             }
         )
     if output_content:
@@ -175,7 +182,7 @@ def model_visible_assistant_output_messages(
             {
                 "role": "tool",
                 "tool_call_id": tool.id,
-                "content": tool.content or "",
+                "content": tool_result_model_content(tool.result or {}),
             }
             for tool in group_tools
             if tool.status != "running"
@@ -260,7 +267,7 @@ def usage_info_for_model(
     model_context_window: int,
 ) -> TokenUsageInfo | None:
     if usage_info is None:
-        return None
+        return TokenUsageInfo(model_context_window=model_context_window)
     return usage_info.model_copy(update={"model_context_window": model_context_window})
 
 

package/backend/src/flowent/workspace/output.py

@@ -13,6 +13,7 @@ from flowent.storage import (
     StoredToolItem,
     StoredToolOutputItem,
 )
+from flowent.tools import tool_result_model_content
 
 APPROVAL_TRANSCRIPT_MESSAGE_LIMIT = 12
 APPROVAL_TRANSCRIPT_TEXT_LIMIT = 2_000
@@ -84,7 +85,9 @@ def approval_transcript(
             if content:
                 entries.append(ApprovalTranscriptEntry(role=role, content=content))
             for tool in message.tools:
-                tool_content = approval_transcript_text(tool.content)
+                tool_content = approval_transcript_text(
+                    tool_result_model_content(tool.result or {})
+                )
                 if tool_content:
                     entries.append(
                         ApprovalTranscriptEntry(

package/backend/src/flowent/workspace/runtime.py

@@ -27,11 +27,13 @@ from flowent.storage import (
     StoredState,
     StoredToolItem,
 )
-from flowent.tools import ToolContext
+from flowent.tools import ToolContext, text_tool_result, tool_specs
 from flowent.usage import (
     TokenUsage,
     TokenUsageInfo,
     append_token_usage,
+    full_context_usage,
+    is_context_window_error,
     recompute_context_usage,
 )
 from flowent.workspace.context import (
@@ -67,6 +69,7 @@ logger = logging.getLogger("flowent.workspace.runtime")
 
 AUTO_COMPACT_RETAINED_MESSAGE_TOKEN_BUDGET = 20_000
 WORKSPACE_PROGRESS_FLUSH_INTERVAL_SECONDS = 0.5
+USER_VISIBLE_MANUAL_COMPACT_ERROR_MESSAGE = "Context could not be compacted."
 
 
 @dataclass
@@ -188,13 +191,16 @@ class WorkspaceRuntime:
         *,
         connection: ProviderConnection,
         context_window_limit: int,
+        budget_messages: Sequence[ChatMessage | Mapping[str, object]] | None = None,
         messages: list[StoredMessage],
         model_history: Sequence[ChatMessage | Mapping[str, object]],
         source_message_id: str | None = None,
+        tools: Sequence[Mapping[str, object]] = (),
     ) -> tuple[StoredMessage, list[dict[str, object]], TokenUsageInfo] | None:
         if not should_auto_compact(
-            model_history,
+            budget_messages or model_history,
             context_window=context_window_limit,
+            tools=tools,
         ):
             return None
         logger.info("Workspace auto compact requested")
@@ -223,6 +229,10 @@ class WorkspaceRuntime:
         )
         next_messages = [*state.messages, user_message]
         self.store.save_messages(next_messages)
+        model_tool_specs = [
+            *tool_specs(),
+            *list(self.mcp_manager.tool_specs()),
+        ]
         model_history: list[ChatMessage | Mapping[str, object]] = [
             *runtime_context_messages(self.cwd, state.settings.agent_prompt),
             *workspace_chat_messages(
@@ -234,9 +244,13 @@ class WorkspaceRuntime:
         auto_compaction = await self.auto_compact_messages(
             connection=connection,
             context_window_limit=context_window_limit,
+            budget_messages=self.request_messages_for_content(
+                state, next_messages, content
+            ),
             messages=state.messages,
             model_history=model_history,
             source_message_id=None,
+            tools=model_tool_specs,
         )
         if auto_compaction is not None:
             marker, _, _ = auto_compaction
@@ -341,6 +355,7 @@ class WorkspaceRuntime:
                     tool.model_dump(exclude_none=True)
                     for tool in assistant_output.tools.values()
                 ],
+                request_tools=model_tool_specs,
                 model_context_window=context_window_limit,
             )
         self.store.save_usage_info(final_usage_info)
@@ -709,6 +724,10 @@ class WorkspaceRuntime:
                 turn_usage_info: TokenUsageInfo | None = None
                 current_output_index = 0
                 latest_usage_output_index: int | None = None
+                model_tool_specs = [
+                    *tool_specs(),
+                    *list(self.mcp_manager.tool_specs()),
+                ]
                 if request_messages is None:
                     current_request_messages = self.request_messages_for_content(
                         state,
@@ -723,9 +742,11 @@ class WorkspaceRuntime:
                     auto_compaction = await self.auto_compact_messages(
                         connection=connection,
                         context_window_limit=context_window_limit,
+                        budget_messages=current_request_messages,
                         messages=state.messages,
                         model_history=pre_turn_request_messages,
                         source_message_id=None,
+                        tools=model_tool_specs,
                     )
                     if auto_compaction is not None:
                         marker, _, usage_info = auto_compaction
@@ -746,6 +767,42 @@ class WorkspaceRuntime:
                         )
                 else:
                     current_request_messages = request_messages
+                    auto_compaction = await self.auto_compact_messages(
+                        connection=connection,
+                        context_window_limit=context_window_limit,
+                        messages=next_messages,
+                        model_history=compact_prompt_chat_messages(
+                            current_request_messages
+                        ),
+                        source_message_id=assistant_message.id,
+                        tools=model_tool_specs,
+                    )
+                    if auto_compaction is not None:
+                        marker, replacement_history, usage_info = auto_compaction
+                        assistant_message = assistant_message.model_copy(
+                            update={"usage_info": usage_info}
+                        )
+                        next_messages = append_or_replace_message(
+                            [*next_messages, marker], assistant_message
+                        )
+                        self.store.save_messages(next_messages)
+                        await self.append_event(
+                            response,
+                            "context_optimized",
+                            {
+                                "message": marker.model_dump(),
+                                **usage_event_data(usage_info),
+                            },
+                        )
+                        current_request_messages = model_request_messages_data(
+                            [
+                                *runtime_context_messages(
+                                    self.cwd, state.settings.agent_prompt
+                                ),
+                                *explicit_skill_messages(self.cwd, self.store, content),
+                                *replacement_history,
+                            ]
+                        )
                 context_usage_messages = (
                     usage_request_messages
                     if usage_request_messages is not None
@@ -801,6 +858,7 @@ class WorkspaceRuntime:
                         messages=next_messages,
                         model_history=compact_prompt_chat_messages(conversation),
                         source_message_id=assistant_snapshot.id,
+                        tools=model_tool_specs,
                     )
                     if auto_result is None:
                         return None
@@ -944,6 +1002,7 @@ class WorkspaceRuntime:
                                         tool.model_dump(exclude_none=True)
                                         for tool in assistant_output.tools.values()
                                     ],
+                                    request_tools=model_tool_specs,
                                     model_context_window=context_window_limit,
                                 )
                             self.store.save_usage_info(final_usage_info)
@@ -976,6 +1035,12 @@ class WorkspaceRuntime:
                 raise
             except Exception as error:
                 logger.exception("Workspace response failed")
+                if is_context_window_error(error):
+                    usage_info = full_context_usage(
+                        self.store.read_usage_info(),
+                        model_context_window=context_window_limit,
+                    )
+                    self.store.save_usage_info(usage_info)
                 if (
                     current_tool_id is not None
                     and current_tool_id in assistant_output.tools
@@ -983,7 +1048,10 @@ class WorkspaceRuntime:
                 ):
                     assistant_output.update_tool(
                         current_tool_id,
-                        {"content": str(error) or "Tool failed.", "status": "failed"},
+                        {
+                            "result": text_tool_result(str(error) or "Tool failed."),
+                            "status": "failed",
+                        },
                     )
                 error_item = assistant_output.append_error(
                     run_error_output_item(
@@ -1142,10 +1210,31 @@ class WorkspaceRuntime:
         async def compact_events() -> AsyncIterator[str]:
             try:
                 marker, usage_info = await asyncio.shield(compact_task)
-            except Exception:
+            except Exception as error:
+                assistant_id = str(uuid4())
+                assistant_output = AssistantOutputBuilder(assistant_id)
+                error_item = run_error_output_item(assistant_id, str(error)).model_copy(
+                    update={"message": USER_VISIBLE_MANUAL_COMPACT_ERROR_MESSAGE}
+                )
+                assistant_output.append_error(error_item)
+                failed_message = StoredMessage(
+                    author="assistant",
+                    content="",
+                    groups=assistant_output.groups,
+                    id=assistant_id,
+                    status="failed",
+                )
+                self.store.save_messages(
+                    [*self.store.read_state().messages, failed_message]
+                )
+                failed_message_data = stream_message_data(failed_message)
+                yield stream_event("snapshot", {"message": failed_message_data})
                 yield stream_event(
                     "error",
-                    {"message": "Context could not be compacted."},
+                    {
+                        "error": error_item.model_dump(exclude_none=True),
+                        "message": USER_VISIBLE_MANUAL_COMPACT_ERROR_MESSAGE,
+                    },
                 )
                 return
 

package/backend/uv.lock

@@ -701,7 +701,7 @@ wheels = [
 
 [[package]]
 name = "flowent"
-version = "0.3.1"
+version = "0.3.2"
 source = { editable = "." }
 dependencies = [
     { name = "fastapi", extra = ["standard"] },