Keep Tool calls stable across provider history and async polls

backend/app/services/agent_runtime/node_executor.py

@@ -778,6 +778,8 @@ async def _model(
                 repair_limit = (
                     WRITE_FILE_PROTOCOL_REPAIR_LIMIT
                     if is_write_file_repair
+                    else 10
+                    if repair_code == "invalid_tool_call"
                     else 1
                 )
                 repair_counter_key = (

backend/app/services/agent_runtime/tool_execution.py

@@ -44,7 +44,7 @@
     "reconcile",
 ]
 ToolSideEffectState = Literal["none", "confirmed", "possible", "unknown"]
-SAFE_READ_MAX_ATTEMPTS = 3
+SAFE_READ_MAX_ATTEMPTS = 10
 
 # These tools dispatch an external image-generation request and can therefore
 # leave the provider outcome uncertain after a response timeout.  Direct Chat

backend/app/services/agent_runtime/tool_repair_budget.py

@@ -10,7 +10,7 @@
 from app.services.agent_runtime.state import JsonObject
 
 SAME_FINGERPRINT_FAILURE_LIMIT = 10
-TOOL_EPISODE_FAILURE_LIMIT = 20
+TOOL_EPISODE_FAILURE_LIMIT = 10
 _REPAIRABLE_MODEL_ACTIONS = frozenset(
     {"repair_arguments", "choose_other_tool"}
 )

backend/app/services/agent_runtime/tool_step_service.py

@@ -541,6 +541,7 @@ def _async_pending_step_result(
     run_id: uuid.UUID,
     execution_id: uuid.UUID,
     call_id: str,
+    origin_call_id: str,
     tool_name: str,
     outcome: ToolExecutionOutcome,
     prior_messages: Sequence[JsonObject],
@@ -597,6 +598,7 @@ def _async_pending_step_result(
         "tool_calls": [poll_call],
         "runtime_intent": "async_poll",
         "runtime_run_id": str(run_id),
+        "runtime_origin_tool_call_id": origin_call_id,
     }
     return ToolStepResult(
         messages=(
@@ -1678,6 +1680,7 @@ async def execute_pending(
         tool_calls: tuple[JsonObject, ...],
     ) -> ToolStepResult:
         step_context_update: JsonObject | None = None
+        async_origin_call_id: str | None = None
         try:
             tenant_id = uuid.UUID(context.tenant_id)
             run_id = uuid.UUID(context.run_id)
@@ -1708,7 +1711,18 @@ async def execute_pending(
                 )
             except ToolContractError as exc:
                 raise ToolExecutionError("tool_context_corrupt", str(exc)) from exc
-            if (
+            is_async_poll = assistant_message.get("runtime_intent") == "async_poll"
+            if is_async_poll:
+                raw_origin_call_id = assistant_message.get(
+                    "runtime_origin_tool_call_id"
+                )
+                if not isinstance(raw_origin_call_id, str) or not raw_origin_call_id:
+                    raise ToolExecutionError(
+                        "tool_context_corrupt",
+                        "async poll is missing its origin Tool Call ID",
+                    )
+                async_origin_call_id = raw_origin_call_id
+            elif (
                 step_context is not None
                 and step_context.assistant_message_id != assistant_message_id
             ):
@@ -1746,6 +1760,7 @@ async def execute_pending(
                     tools=legacy_tools,
                 )
                 step_context_update = step_context.to_json()
+                async_origin_call_id = None
                 logger.warning(
                     "[RuntimeToolCompatibility] event=legacy_tool_context_resolved "
                     "run_id={} assistant_message_id={} accepted_call_count={} "
@@ -1770,15 +1785,27 @@ async def execute_pending(
                         step_tool_context=step_context_update,
                     )
                 call_id, tool_name, arguments = _call_fields(call)
-                accepted = (
-                    _accepted_call(
+                if async_origin_call_id is not None:
+                    origin_call = _accepted_call(
                         step_context,
-                        call_id=call_id,
+                        call_id=async_origin_call_id,
                         tool_name=tool_name,
                     )
-                    if step_context is not None
-                    else None
-                )
+                    accepted = AcceptedToolCall(
+                        call_instance_id=call_id,
+                        provider_call_id=None,
+                        entry=origin_call.entry,
+                    )
+                else:
+                    accepted = (
+                        _accepted_call(
+                            step_context,
+                            call_id=call_id,
+                            tool_name=tool_name,
+                        )
+                        if step_context is not None
+                        else None
+                    )
                 if accepted is None:  # pragma: no cover - new contexts are mandatory here
                     raise ToolExecutionError(
                         "tool_context_corrupt",
@@ -1933,6 +1960,7 @@ async def execute_pending(
                             run_id=run_id,
                             execution_id=reservation.execution.id,
                             call_id=call_id,
+                            origin_call_id=async_origin_call_id or call_id,
                             tool_name=tool_name,
                             outcome=reservation.reusable_result,
                             prior_messages=messages,
@@ -2482,6 +2510,7 @@ async def execute_pending(
                         run_id=run_id,
                         execution_id=reservation.execution.id,
                         call_id=call_id,
+                        origin_call_id=async_origin_call_id or call_id,
                         tool_name=tool_name,
                         outcome=outcome,
                         prior_messages=messages,

backend/app/services/llm/caller.py

@@ -66,7 +66,7 @@ async def execute_tool(*args, **kwargs):
     "send_message_to_agent", "send_feishu_message", "send_email"
 })
 
-WRITE_FILE_PROTOCOL_REPAIR_LIMIT = 3
+WRITE_FILE_PROTOCOL_REPAIR_LIMIT = 10
 WRITE_FILE_PROTOCOL_REPAIR_COUNTER_KEY = "invalid_tool_call:write_file"
 WRITE_FILE_PROTOCOL_REPAIR_INSTRUCTION = (
     "Your previous `write_file` call was not executed because `function.arguments` "
@@ -788,7 +788,7 @@ async def _buffer_chunk(_text: str) -> None:
             repair_limit = (
                 WRITE_FILE_PROTOCOL_REPAIR_LIMIT
                 if retry_tool_name == "write_file"
-                else 1
+                else 10
             )
             repair_counter_key = (
                 WRITE_FILE_PROTOCOL_REPAIR_COUNTER_KEY

backend/app/services/llm/client.py

@@ -1564,19 +1564,6 @@ def _content_to_gemini_parts(self, content: Any) -> list[dict[str, Any]]:
 
         return [{"text": str(content)}]
 
-    def _extract_tool_name_map(self, messages: list[LLMMessage]) -> dict[str, str]:
-        """Build tool_call_id -> function_name map from assistant messages."""
-        out: dict[str, str] = {}
-        for msg in messages:
-            if msg.role != "assistant" or not msg.tool_calls:
-                continue
-            for tc in msg.tool_calls:
-                tc_id = tc.get("id")
-                tc_name = tc.get("function", {}).get("name")
-                if tc_id and tc_name:
-                    out[tc_id] = tc_name
-        return out
-
     def _convert_tools(self, tools: list[dict] | None) -> tuple[list[dict[str, Any]] | None, dict[str, Any] | None]:
         """Convert OpenAI-style tools to Gemini function declarations."""
         if not tools:
@@ -1617,7 +1604,7 @@ def _build_payload(
         messages = normalize_provider_messages(messages)
         system_blocks: list[str] = []
         contents: list[dict[str, Any]] = []
-        tool_name_map = self._extract_tool_name_map(messages)
+        pending_tool_names: dict[str, str] = {}
 
         for msg in messages:
             if msg.role == "system":
@@ -1630,16 +1617,22 @@ def _build_payload(
                 continue
 
             if msg.role == "user":
+                pending_tool_names = {}
                 parts = self._content_to_gemini_parts(msg.content)
                 if parts:
                     contents.append({"role": "user", "parts": parts})
                 continue
 
             if msg.role == "assistant":
+                pending_tool_names = {}
                 parts = self._content_to_gemini_parts(msg.content)
                 if msg.tool_calls:
                     for tc in msg.tool_calls:
                         fn = tc.get("function", {})
+                        tc_id = tc.get("id")
+                        tc_name = fn.get("name")
+                        if tc_id and tc_name:
+                            pending_tool_names[tc_id] = tc_name
                         args = fn.get("arguments", "{}")
                         if isinstance(args, str):
                             try:
@@ -1666,7 +1659,7 @@ def _build_payload(
                 continue
 
             if msg.role == "tool":
-                name = tool_name_map.get(msg.tool_call_id or "", msg.tool_call_id or "tool_result")
+                name = pending_tool_names.get(msg.tool_call_id or "", msg.tool_call_id or "tool_result")
                 response_content = msg.content or ""
                 if isinstance(response_content, str):
                     try:

backend/tests/test_agent_runtime_model_step_service.py

@@ -598,7 +598,7 @@ async def complete(model_arg, _messages, **_kwargs):
 
 
 @pytest.mark.asyncio
-async def test_invalid_write_file_arguments_request_three_protocol_repairs() -> None:
+async def test_invalid_write_file_arguments_request_ten_protocol_repairs() -> None:
     tenant_id = uuid.uuid4()
     model = _model(tenant_id)
     agent = _agent(tenant_id)

backend/tests/test_agent_runtime_node_executor.py

@@ -1351,15 +1351,16 @@ async def test_empty_output_is_repaired_once_then_fails_explicitly() -> None:
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
-    ("repair_code", "instruction"),
+    ("repair_code", "instruction", "repair_limit"),
     [
-        ("invalid_finish", "Retry finish with valid content."),
-        ("invalid_tool_call", "Retry with valid JSON tool arguments."),
+        ("invalid_finish", "Retry finish with valid content.", 1),
+        ("invalid_tool_call", "Retry with valid JSON tool arguments.", 10),
     ],
 )
 async def test_repeated_model_tool_protocol_repair_code_fails_explicitly(
     repair_code: str,
     instruction: str,
+    repair_limit: int,
 ) -> None:
     run_id = uuid.uuid4()
     repair = ModelStepResult(
@@ -1368,7 +1369,7 @@ async def test_repeated_model_tool_protocol_repair_code_fails_explicitly(
         repair_instruction=instruction,
         repair_code=repair_code,
     )
-    model = ModelService(repair, repair)
+    model = ModelService(*([repair] * (repair_limit + 1)))
     executor = _executor(model)
 
     result = await _invoke(run_id, executor, model_turn_limit=50)
@@ -1377,13 +1378,13 @@ async def test_repeated_model_tool_protocol_repair_code_fails_explicitly(
     assert lifecycle["status"] == "failed"
     assert lifecycle["reason"] == "model_tool_protocol_violation"
     assert lifecycle["error"]["code"] == "model_tool_protocol_violation"
-    assert lifecycle["model_protocol_repairs"] == {repair_code: 1}
-    assert lifecycle["model_step_count"] == 2
-    assert model.calls == 2
+    assert lifecycle["model_protocol_repairs"] == {repair_code: repair_limit}
+    assert lifecycle["model_step_count"] == repair_limit + 1
+    assert model.calls == repair_limit + 1
 
 
 @pytest.mark.asyncio
-async def test_write_file_protocol_repair_uses_three_attempts_then_guides_user() -> None:
+async def test_write_file_protocol_repair_uses_ten_attempts_then_guides_user() -> None:
     run_id = uuid.uuid4()
     repair = ModelStepResult(
         intent="text",
@@ -1392,7 +1393,7 @@ async def test_write_file_protocol_repair_uses_three_attempts_then_guides_user()
         repair_code="invalid_tool_call",
         repair_tool_name="write_file",
     )
-    model = ModelService(repair, repair, repair, repair)
+    model = ModelService(*([repair] * 11))
     executor = _executor(model)
 
     result = await _invoke(run_id, executor, model_turn_limit=50)
@@ -1408,14 +1409,14 @@ async def test_write_file_protocol_repair_uses_three_attempts_then_guides_user()
         ),
     }
     assert lifecycle["model_protocol_repairs"] == {
-        "invalid_tool_call:write_file": 3,
+        "invalid_tool_call:write_file": 10,
     }
-    assert lifecycle["model_step_count"] == 4
-    assert model.calls == 4
+    assert lifecycle["model_step_count"] == 11
+    assert model.calls == 11
 
 
 @pytest.mark.asyncio
-async def test_write_file_protocol_can_recover_on_the_third_repair() -> None:
+async def test_write_file_protocol_can_recover_on_the_tenth_repair() -> None:
     run_id = uuid.uuid4()
     repair = ModelStepResult(
         intent="text",
@@ -1424,9 +1425,7 @@ async def test_write_file_protocol_can_recover_on_the_third_repair() -> None:
         repair_tool_name="write_file",
     )
     model = ModelService(
-        repair,
-        repair,
-        repair,
+        *([repair] * 10),
         ModelStepResult(intent="finish", finish_content="Recovered"),
     )
     executor = _executor(model)
@@ -1435,9 +1434,9 @@ async def test_write_file_protocol_can_recover_on_the_third_repair() -> None:
 
     assert result["lifecycle"]["status"] == "completed"
     assert result["lifecycle"]["model_protocol_repairs"] == {
-        "invalid_tool_call:write_file": 3,
+        "invalid_tool_call:write_file": 10,
     }
-    assert model.calls == 4
+    assert model.calls == 11
 
 
 @pytest.mark.asyncio

backend/tests/test_agent_runtime_tool_repair_budget.py

@@ -50,7 +50,7 @@ def test_tenth_consecutive_fingerprint_pauses_without_off_by_one() -> None:
     assert _episode(state)["total_failures"] == 10
 
 
-def test_twentieth_tool_failure_pauses_even_when_fingerprint_changes() -> None:
+def test_tenth_tool_failure_pauses_even_when_fingerprint_changes() -> None:
     state: dict = {}
     transition = None
     for model_step in range(1, TOOL_EPISODE_FAILURE_LIMIT + 1):
@@ -63,7 +63,7 @@ def test_twentieth_tool_failure_pauses_even_when_fingerprint_changes() -> None:
 
     assert transition is not None
     assert transition.pause_reason == "tool_repair_episode_limit_reached"
-    assert _episode(state)["total_failures"] == 20
+    assert _episode(state)["total_failures"] == 10
     assert _episode(state)["same_fingerprint_failures"] == 1