Enhance model validation and tool call management

- Updated the build_chat_request function to validate model names against available models, raising an HTTPException for unknown models.
- Enhanced UIMessageStreamEncoder to manage open tool call IDs, ensuring proper closure of tool calls that do not complete, and added error handling for unfinished tool calls.
- Modified the GoogleSearchActivity class to accurately track activity completion per cycle, preventing duplicate events and ensuring correct state management.
- Updated tests to validate the new model validation and tool call management functionalities, ensuring robustness in error handling and state transitions.

app/api.py

@@ -148,7 +148,9 @@ def build_chat_request(payload: ApiChatRequest) -> ChatRequest:
         )
         or DEFAULT_SELECTED_SOURCE_KEYS
     )
-    model_name = (payload.model or DEFAULT_MODEL_NAME).strip()
+    model_name = (payload.model or DEFAULT_MODEL_NAME).strip() or DEFAULT_MODEL_NAME
+    if model_name not in {model["id"] for model in AVAILABLE_MODELS}:
+        raise HTTPException(status_code=422, detail="Unknown model")
     if payload.enabledTools is None:
         enabled_tools = tuple(
             tool["key"]
@@ -185,6 +187,7 @@ class UIMessageStreamEncoder:
         self.text_block_count = 0
         self.active_reasoning_id = ""
         self.source_matches_by_call_id: dict[str, list[dict[str, Any]]] = {}
+        self.open_tool_call_ids: list[str] = []
         self.closed = False
 
     def close_reasoning_block(self) -> list[dict[str, Any]]:
@@ -270,6 +273,8 @@ class UIMessageStreamEncoder:
             parts.extend(self.close_text_block())
             call_id = str(event.data.get("call_id", uuid4().hex))
             tool_name = str(event.data.get("tool_name", "tool"))
+            if call_id not in self.open_tool_call_ids:
+                self.open_tool_call_ids.append(call_id)
             parts.append(
                 {
                     "type": "tool-input-start",
@@ -352,6 +357,8 @@ class UIMessageStreamEncoder:
                 "text": str(event.data.get("output_text", "")),
                 "matches": self.source_matches_by_call_id.get(call_id, []),
             }
+            if call_id in self.open_tool_call_ids:
+                self.open_tool_call_ids.remove(call_id)
             parts.append(
                 {
                     "type": "tool-output-available",
@@ -389,6 +396,17 @@ class UIMessageStreamEncoder:
             parts.extend(self.close_reasoning_block())
             if self.text_block_id:
                 parts.append({"type": "text-end", "id": self.text_block_id})
+            # Close tool calls the turn never finished, or the UI shows
+            # them as running forever.
+            for call_id in self.open_tool_call_ids:
+                parts.append(
+                    {
+                        "type": "tool-output-error",
+                        "toolCallId": call_id,
+                        "errorText": "The response ended before this tool call finished.",
+                    }
+                )
+            self.open_tool_call_ids = []
             if self.message_id:
                 parts.append({"type": "finish-step"})
                 parts.append({"type": "finish"})

app/chat_service.py

@@ -1051,7 +1051,7 @@ async def stream_chat(request: ChatRequest) -> AsyncIterator[ChatEvent]:
                 if metadata.get("lc_source") == "summarization":
                     continue
 
-                step = str(metadata.get("langgraph_step", ""))
+                step = str(metadata.get("langgraph_node", ""))
                 if include_reasoning:
                     thought_text = "\n\n".join(extract_thought_summaries(token.content))
                     if thought_text:
@@ -1171,6 +1171,11 @@ async def stream_chat(request: ChatRequest) -> AsyncIterator[ChatEvent]:
                 )
                 if search_started:
                     yield search_started
+                # This model step is complete, so any search it ran is done:
+                # close the activity now instead of when the whole turn ends.
+                search_step_completed = google_search.completed_event()
+                if search_step_completed:
+                    yield search_step_completed
 
                 if is_anthropic_model(request.model_name):
                     anthropic_updates, anthropic_tool_uses = (

app/provider_events.py

@@ -112,11 +112,17 @@ class GoogleSearchActivity:
         self._message_id = message_id
         self._web_evidence = web_evidence
         self._call_id = ""
+        self._completed = False
         self._queries: list[str] = []
         self._match_count = 0
 
     def observe(self, response_metadata: Any) -> ChatEvent | None:
-        """Record metadata; return a tool_call_started event on first activity."""
+        """Record metadata; return a tool_call_started event on new activity.
+
+        Activity after a completed cycle starts a fresh synthetic call (a
+        later model step ran another search), so completion timing per step
+        stays accurate.
+        """
         new_queries = [
             q
             for q in extract_web_search_queries(response_metadata)
@@ -127,8 +133,11 @@ class GoogleSearchActivity:
             self._web_evidence,
         )
         started: ChatEvent | None = None
-        if (new_queries or new_grounding) and not self._call_id:
+        if (new_queries or new_grounding) and (not self._call_id or self._completed):
             self._call_id = uuid4().hex
+            self._completed = False
+            self._queries = []
+            self._match_count = 0
             joined = "; ".join(new_queries)
             started = ChatEvent(
                 "tool_call_started",
@@ -145,8 +154,10 @@ class GoogleSearchActivity:
         return started
 
     def completed_event(self) -> ChatEvent | None:
-        if not self._call_id:
+        """Close the active cycle; idempotent until new activity arrives."""
+        if not self._call_id or self._completed:
             return None
+        self._completed = True
         joined = "; ".join(self._queries)
         if self._match_count == 0:
             output_text = "Google search ran but returned no grounding results."

frontend/components/chat-shell.tsx

@@ -304,6 +304,7 @@ export function ChatShell() {
       messageId,
       body: {
         sourceKeys: selectedSourceKeys,
+        enabledTools: enabledToolKeys,
         includeReasoning: true,
         threadId,
         model: selectedModel,

tests/test_api.py

@@ -17,7 +17,7 @@ import pytest
 import uvicorn
 from fastapi.testclient import TestClient
 
-from app.api import app
+from app.api import UIMessageStreamEncoder, app
 from app.chat_types import ChatEvent
 
 
@@ -273,6 +273,65 @@ class ApiTestCase(unittest.TestCase):
             part_types.index("tool-input-start"), part_types.index("text-start")
         )
 
+    def test_finish_error_closes_open_tool_calls(self) -> None:
+        encoder = UIMessageStreamEncoder()
+        encoder.encode(ChatEvent("message_started", {"message_id": "m1"}))
+        encoder.encode(
+            ChatEvent(
+                "tool_call_started",
+                {
+                    "message_id": "m1",
+                    "call_id": "call_open",
+                    "tool_name": "google_search",
+                    "args": {"query": "x"},
+                    "args_text": "x",
+                },
+            )
+        )
+        encoder.encode(
+            ChatEvent(
+                "tool_call_started",
+                {
+                    "message_id": "m1",
+                    "call_id": "call_done",
+                    "tool_name": "run_kb_command",
+                    "args": {"command": "ls"},
+                    "args_text": "ls",
+                },
+            )
+        )
+        encoder.encode(
+            ChatEvent(
+                "tool_call_completed",
+                {"message_id": "m1", "call_id": "call_done", "output_text": "ok"},
+            )
+        )
+
+        parts = encoder.finish_error("boom")
+
+        tool_errors = [p for p in parts if p["type"] == "tool-output-error"]
+        self.assertEqual(len(tool_errors), 1)
+        self.assertEqual(tool_errors[0]["toolCallId"], "call_open")
+        # Stream still terminates cleanly after closing the tool part.
+        part_types = [p["type"] for p in parts]
+        self.assertLess(
+            part_types.index("tool-output-error"), part_types.index("finish")
+        )
+
+    def test_chat_rejects_unknown_model(self) -> None:
+        with TestClient(app) as client:
+            response = client.post(
+                "/api/chat",
+                json={
+                    "messages": [
+                        {"role": "user", "parts": [{"type": "text", "text": "hi"}]}
+                    ],
+                    "model": "anthropic:claude-opus-4-8",
+                },
+            )
+
+        self.assertEqual(response.status_code, 422)
+
 
 LIVE_API_E2E = pytest.mark.skipif(
     os.getenv("RUN_LIVE_API_E2E") != "1",

tests/test_chat_service.py

@@ -69,7 +69,7 @@ class FakeStreamingAgent(FakeAgent):
                         }
                     ],
                 ),
-                {"langgraph_step": "model"},
+                {"langgraph_node": "model"},
             ),
         }
         yield {

tests/test_provider_events.py

@@ -1,6 +1,52 @@
 from __future__ import annotations
 
-from app.provider_events import extract_anthropic_source_matches
+from app.provider_events import (
+    GoogleSearchActivity,
+    extract_anthropic_source_matches,
+)
+
+
+def _grounding_metadata(query: str, uri: str, title: str) -> dict:
+    return {
+        "grounding_metadata": {
+            "web_search_queries": [query],
+            "grounding_chunks": [{"web": {"uri": uri, "title": title}}],
+            "grounding_supports": [],
+        }
+    }
+
+
+def test_google_search_activity_completes_once_per_cycle() -> None:
+    evidence: dict = {}
+    activity = GoogleSearchActivity("msg_1", evidence)
+
+    started = activity.observe(_grounding_metadata("q1", "https://a.example", "A"))
+    assert started is not None
+    assert started.data["tool_name"] == "google_search"
+
+    completed = activity.completed_event()
+    assert completed is not None
+    assert "1 web result" in completed.data["output_text"]
+    assert completed.data["call_id"] == started.data["call_id"]
+    # Idempotent: the post-turn fallback must not duplicate the event.
+    assert activity.completed_event() is None
+
+
+def test_google_search_activity_restarts_after_completion() -> None:
+    # A later model step running another search gets its own activity cycle
+    # instead of being swallowed by the completed one.
+    activity = GoogleSearchActivity("msg_1", {})
+    first = activity.observe(_grounding_metadata("q1", "https://a.example", "A"))
+    assert activity.completed_event() is not None
+
+    second = activity.observe(_grounding_metadata("q2", "https://b.example", "B"))
+    assert second is not None
+    assert second.data["call_id"] != first.data["call_id"]
+    assert second.data["args_text"] == "q2"
+
+    completed = activity.completed_event()
+    assert completed is not None
+    assert "1 web result" in completed.data["output_text"]
 
 
 def test_web_fetch_result_object_recorded_as_evidence() -> None: