@samitouri / QOSami_Agents_Zero-2 / commits / 6b7302f6

fix(tools): harden malformed text tool handling

Require complete standalone text requests for execution while accepting supported text-tool formats. Route fenced and structurally broken tool intent through the existing misformat repair path so streamed chats do not remain in a thinking state.

Alessandro committed Jul 23, 2026 at 11:34 UTC 6b7302f6642e9f4e68f8bb8afbfc8b2550511c06

10 files changed +480 -73

agent.py

+24 -28

@@ -447,24 +447,15 @@ class Agent:

447 printer.print("Response: ") # start of response

448 # Pass chunk and full data to extensions for processing

449 stream_data = {"chunk": chunk, "full": full}

450 - stop_response: str | None = None

451 -

452 - snapshot = extract_tools.extract_json_root_string(full)

453 - if snapshot:

454 - parsed_snapshot = extract_tools.json_parse_dirty(snapshot)

455 - if parsed_snapshot is not None:

456 - try:

457 - await self.validate_tool_request(parsed_snapshot)

458 - except Exception:

459 - pass

460 - else:

461 - previous_full = last_response_stream_full

462 - stream_data["full"] = snapshot

463 - if snapshot.startswith(previous_full):

464 - stream_data["chunk"] = snapshot[len(previous_full) :]

465 - else:

466 - stream_data["chunk"] = snapshot

467 - stop_response = snapshot

450 + tool_request = extract_tools.extract_tool_request(full)

451 + if tool_request is not None:

452 + try:

453 + await self.validate_tool_request(tool_request)

454 + except Exception:

455 + pass

456 + else:

457 + await self.handle_response_stream(full)

458 + return full.strip()

459

460 await extension.call_extensions_async(

461 "response_stream_chunk",

@@ -478,8 +469,6 @@ class Agent:

469 # Use the potentially modified full text for downstream processing

470 await self.handle_response_stream(stream_data["full"])

471 last_response_stream_full = stream_data["full"]

481 - if stop_response is not None:

482 - return stop_response

472

473 # call main LLM

474 llm_result = await self.call_chat_model_turn(

@@ -1123,13 +1112,18 @@ class Agent:

1112 return None

1113 if llm_result.builtin_items and not llm_result.response:

1114 return None

1126 - if (

1127 - llm_result.mode == "responses"

1128 - and llm_result.response

1129 - and extract_tools.json_parse_dirty(llm_result.response) is None

1130 - ):

1131 - return llm_result.response

1132 - return await self.process_tools(llm_result.response)

1115 + message = llm_result.response

1116 + if not message and llm_result.reasoning:

1117 + if (

1118 + extract_tools.extract_tool_request(llm_result.reasoning) is not None

1119 + or extract_tools.is_misformatted_tool_request(llm_result.reasoning)

1120 + ):

1121 + message = llm_result.reasoning

1122 + if extract_tools.extract_tool_request(message) is None:

1123 + if extract_tools.is_misformatted_tool_request(message):

1124 + return await self.process_tools(message)

1125 + return message

1126 + return await self.process_tools(message)

1127

1128 async def _execute_tool_request(

1129 self,

@@ -1408,7 +1402,7 @@ class Agent:

1402 @extension.extensible

1403 async def process_tools(self, msg: str):

1404 # search for tool usage requests in agent message

1411 - tool_request = extract_tools.json_parse_dirty(msg)

1405 + tool_request = extract_tools.extract_tool_request(msg)

1406

1407 raw_tool_name = ""

1408 tool_args = {}

@@ -1417,6 +1411,7 @@ class Agent:

1411 # block was found - the misformat warning path below handles that.

1412 if tool_request is not None:

1413 try:

1414 + await self.validate_tool_request(tool_request)

1415 raw_tool_name, tool_args = extract_tools.normalize_tool_request(

1416 tool_request

1417 )

@@ -1458,6 +1453,7 @@ class Agent:

1453 )

1454

1455 if tool:

1456 + tool.args = tool_args

1457 self.loop_data.current_tool = tool # type: ignore

1458 try:

1459 await self.handle_intervention()

extensions/python/_functions/agent/Agent/hist_add_ai_response/end/_10_log_plain_responses.py

+1 -1

@@ -23,7 +23,7 @@ class LogPlainResponses(Extension):

23 message = call_args[1]

24 if not isinstance(message, str) or not message:

25 return

26 - if extract_tools.json_parse_dirty(message) is not None:

26 + if extract_tools.extract_tool_request(message) is not None:

27 return

29 params = getattr(getattr(self.agent, "loop_data", None), "params_temporary", None)

extensions/python/response_stream/_20_live_response.py

+4 -8

@@ -1,4 +1,5 @@

1 from helpers import persist_chat, tokens

2 +from helpers import extract_tools

3 from helpers.extension import Extension

4 from agent import LoopData

5 import asyncio

@@ -19,13 +20,8 @@ class LiveResponse(Extension):

20 return

22 try:

22 - if (

23 - not "tool_name" in parsed

24 - or parsed["tool_name"] != "response"

25 - or "tool_args" not in parsed

26 - or "text" not in parsed["tool_args"]

27 - or not parsed["tool_args"]["text"]

28 - ):

23 + tool_name, tool_args = extract_tools.normalize_tool_request(parsed)

24 + if tool_name != "response" or not tool_args.get("text"):

25 return # not a response

27 # create log message and store it in loop data temporary params

@@ -43,6 +39,6 @@ class LiveResponse(Extension):

40 # update log message

41 log_item = loop_data.params_temporary["log_item_response"]

46 - log_item.update(content=parsed["tool_args"]["text"])

42 + log_item.update(content=tool_args["text"])

43 except Exception as e:

44 pass

helpers/extract_tools.py

+71

@@ -20,17 +20,88 @@ def json_parse_dirty(json: str) -> dict[str, Any] | None:

20 return first_data

23 +def extract_tool_request(content: str) -> dict[str, Any] | None:

24 + if not content or not isinstance(content, str):

25 + return None

26 +

27 + content = content.strip()

28 + root = extract_json_root_string(content)

29 + if root != content:

30 + return None

31 +

32 + request = _parse_json_root_object(root)

33 + return request if request is not None and _is_tool_request(request) else None

34 +

35 +

36 +def is_misformatted_tool_request(content: str) -> bool:

37 + if not content or not isinstance(content, str):

38 + return False

39 +

40 + content = content.strip()

41 + for fenced_content in re.findall(

42 + r"```(?:json)?\s*(.*?)```", content, flags=re.IGNORECASE | re.DOTALL

43 + ):

44 + request = json_parse_dirty(fenced_content)

45 + if isinstance(request, dict) and _is_tool_request(request):

46 + return True

47 +

48 + if (

49 + not content.endswith("}")

50 + or re.match(r'^\{\s*"thoughts"\s*:', content) is None

51 + ):

52 + return False

53 +

54 + request = json_parse_dirty(content)

55 + thoughts = request.get("thoughts") if isinstance(request, dict) else None

56 + thoughts_text = (

57 + "\n".join(thought for thought in thoughts if isinstance(thought, str))

58 + if isinstance(thoughts, list)

59 + else ""

60 + )

61 + return (

62 + isinstance(thoughts, list)

63 + and all(

64 + f'{field}\":' in thoughts_text

65 + for field in ("headline", "tool_name", "tool_args")

66 + )

67 + )

68 +

69 +

70 def normalize_tool_request(tool_request: Any) -> tuple[str, dict]:

71 if not isinstance(tool_request, dict):

72 raise ValueError("Tool request must be a dictionary")

73 + if (

74 + not tool_request.get("tool_name")

75 + and not tool_request.get("tool")

76 + and "actions" in tool_request

77 + ):

78 + actions = tool_request["actions"]

79 + # Text tool calls allow one request per turn; do not silently discard extras.

80 + if (

81 + not isinstance(actions, list)

82 + or len(actions) != 1

83 + or not isinstance(actions[0], dict)

84 + ):

85 + raise ValueError(

86 + "Tool request actions wrapper must contain exactly one dictionary"

87 + )

88 + tool_request = actions[0]

89 +

90 tool_name = tool_request.get("tool_name")

91 if not tool_name or not isinstance(tool_name, str):

92 tool_name = tool_request.get("tool")

93 + if (

94 + (not tool_name or not isinstance(tool_name, str))

95 + and tool_request.get("type") == "function"

96 + ):

97 + tool_name = tool_request.get("name")

98 if not tool_name or not isinstance(tool_name, str):

99 raise ValueError("Tool request must have a tool_name (type string) field")

100 tool_args = tool_request.get("tool_args")

101 if not isinstance(tool_args, dict):

102 tool_args = tool_request.get("args")

103 + if not isinstance(tool_args, dict) and tool_request.get("type") == "function":

104 + tool_args = tool_request.get("parameters")

105 if not isinstance(tool_args, dict):

106 raise ValueError("Tool request must have a tool_args (type dictionary) field")

107 tool_args = dict(tool_args)

helpers/extract_tools.py.dox.md

+7 -2

@@ -12,6 +12,8 @@

12 - `extract_tools.py.dox.md` owns durable notes about responsibilities, contracts, side effects, and verification for that implementation.

13 - Top-level functions:

14 - `json_parse_dirty(json: str) -> dict[str, Any] | None`

15 +- `extract_tool_request(content: str) -> dict[str, Any] | None`

16 +- `is_misformatted_tool_request(content: str) -> bool`

17 - `normalize_tool_request(tool_request: Any) -> tuple[str, dict]`

18 - `extract_json_root_string(content: str) -> str | None`

19 - `extract_json_root_strings(content: str) -> list[str]`

@@ -24,8 +26,11 @@

26 - Helper modules own reusable framework APIs and must preserve public callers unless all callers, tests, and docs are updated together.

27 - Update this file whenever public functions, classes, persistence behavior, path/security assumptions, side effects, or cross-module contracts change.

28 - Observed side-effect areas: settings/state persistence.

27 -- Dirty parsing scans complete JSON object roots in prose and prefers the first object that normalizes as a valid tool request, so a leading text preamble or incidental non-tool object does not force a misformat warning when a valid tool call follows.

28 -- Streaming tool snapshots use the same valid-tool preference through `extract_json_root_string`, while preserving the first complete object fallback when no valid tool-call object is present.

29 +- Dirty parsing scans complete JSON object roots in prose and prefers the first object that normalizes as a valid tool request for permissive repair and legacy callers.

30 + Normalization accepts canonical `tool_name`/`tool_args`, legacy `tool`/`args`, native `type="function"` `name`/`parameters`, and a single-item `actions` wrapper; malformed or multi-action wrappers are rejected.

31 +- `extract_tool_request` is the execution boundary: it accepts a request only when the complete trimmed content is one valid tool object. Plain text, ordinary JSON, and tool-shaped JSON embedded in prose remain final text.

32 +- `is_misformatted_tool_request` identifies either a tool request wrapped in a JSON code fence or a complete Agent Zero envelope that starts with `thoughts` and whose dirty parser has absorbed `headline`, `tool_name`, and `tool_args` into that list. It routes that output to the existing repair prompt without executing it.

33 +- Streaming tool snapshots use `extract_tool_request`; the permissive root helpers remain available for repair and legacy callers, not tool execution.

34 - Root extraction ignores objects nested inside an open parent object, so streamed wrapper tools such as `parallel` cannot stop early on the first nested `tool_calls` item.

35 - Imported dependency areas include: `dirty_json`, `helpers.modules`, `re`, `regex`, `typing`.

models.py

+1 -6

@@ -738,14 +738,9 @@ class LiteLLMChatWrapper(SimpleChatModel):

738 output["response_delta"]

739 )

740 )

741 - if (

742 - stop_response is not None

743 - and not transport.policy.using_responses

744 - ):

745 - result.response = stop_response

746 - break

741 if stop_response is not None:

742 result.response = stop_response

743 + break

744 else:

745 parsed = await transport.acomplete()

746 output = result.add_chunk(parsed)

tests/test_plain_response_logging.py

+41

@@ -1,6 +1,8 @@

1 +import pytest

2 from types import SimpleNamespace

4 from helpers.log import Log

5 +from extensions.python.response_stream._20_live_response import LiveResponse

6 from extensions.python._functions.agent.Agent.hist_add_ai_response.end._10_log_plain_responses import (

7 LogPlainResponses,

8 )

@@ -50,6 +52,20 @@ def test_responses_tool_json_keeps_generating_log_as_agent_step():

52 assert "log_item_response" not in agent.loop_data.params_temporary

55 +def test_responses_plain_json_completion_finishes_generating_log_as_response():

56 + agent, item = _agent_with_generating_log()

57 + data = {

58 + "args": (agent, '{"status":"ok"}'),

59 + "kwargs": {"id": "msg-1", "llm_result": SimpleNamespace(mode="responses")},

60 + }

61 +

62 + LogPlainResponses(agent=agent).execute(data=data)

63 +

64 + assert item.type == "response"

65 + assert item.content == '{"status":"ok"}'

66 + assert agent.loop_data.params_temporary["log_item_response"] is item

67 +

68 +

69 def test_responses_plain_text_completion_does_not_replace_live_response_log():

70 agent, item = _agent_with_generating_log()

71 live_response = Log().log(type="response", content="Already live")

@@ -64,3 +80,28 @@ def test_responses_plain_text_completion_does_not_replace_live_response_log():

80 assert item.type == "agent"

81 assert item.content == ""

82 assert agent.loop_data.params_temporary["log_item_response"] is live_response

83 +

84 +

85 +@pytest.mark.asyncio

86 +async def test_live_response_renders_single_action_wrapper():

87 + log = Log()

88 + generating = log.log(type="agent", id="msg-1")

89 + loop_data = SimpleNamespace(params_temporary={"log_item_generating": generating})

90 + agent = SimpleNamespace(

91 + context=SimpleNamespace(log=log),

92 + agent_name="A0",

93 + )

94 +

95 + await LiveResponse(agent=agent).execute(

96 + loop_data=loop_data,

97 + parsed={

98 + "actions": [

99 + {"tool_name": "response", "tool_args": {"text": "wrapper works"}}

100 + ]

101 + },

102 + )

103 +

104 + response = loop_data.params_temporary["log_item_response"]

105 + assert response.type == "response"

106 + assert response.content == "wrapper works"

107 + assert response.id == "msg-1"

tests/test_responses_architecture.py

+200 -3

@@ -1,3 +1,4 @@

1 +import asyncio

2 import sys

3 from pathlib import Path

@@ -11,7 +12,7 @@ if str(PROJECT_ROOT) not in sys.path:

13 import models

14 from agent import Agent, AgentConfig, AgentContextType, LoopData

14 -from helpers import history, litellm_transport

15 +from helpers import extract_tools, history, litellm_transport

16 from helpers.log import Log

17 from helpers.llm_result import LLMResult, result_from_metadata

18 from helpers.persist_chat import _collect_response_ids

@@ -43,6 +44,28 @@ class _AsyncEventStream:

44 self.closed = True

47 +class _StallingAsyncEventStream:

48 + def __init__(self, event: dict):

49 + self.event = event

50 + self.sent = False

51 + self.closed = False

52 + self._stalled = asyncio.Event()

53 +

54 + def __aiter__(self):

55 + return self

56 +

57 + async def __anext__(self):

58 + if not self.sent:

59 + self.sent = True

60 + return self.event

61 + await self._stalled.wait()

62 + raise StopAsyncIteration

63 +

64 + async def aclose(self):

65 + self.closed = True

66 + self._stalled.set()

67 +

68 +

69 def test_llm_result_persists_only_durable_responses_metadata():

70 result = LLMResult.from_response(

71 {

@@ -248,7 +271,7 @@ async def test_transport_downgrades_unsupported_builtin_tools(monkeypatch):

271

272

273 @pytest.mark.asyncio

251 -async def test_unified_turn_keeps_stream_open_to_capture_response_id(monkeypatch):

274 +async def test_unified_turn_captures_response_id_without_stop_request(monkeypatch):

275 stream = _AsyncEventStream(

276 [

277 {

@@ -303,7 +326,7 @@ async def test_unified_turn_keeps_stream_open_to_capture_response_id(monkeypatch

326 )

327

328 async def response_callback(chunk: str, full: str):

306 - return full

329 + return None

330

331 result = await wrapper.unified_turn(

332 messages=[HumanMessage(content="hi")],

@@ -316,6 +339,46 @@ async def test_unified_turn_keeps_stream_open_to_capture_response_id(monkeypatch

339 assert result.function_calls[0].call_id == "call_1"

340

341

342 +@pytest.mark.asyncio

343 +async def test_unified_turn_stops_responses_stream_after_callback_stop(monkeypatch):

344 + message = (

345 + '{"thoughts":["test"],"actions":['

346 + '{"tool_name":"response","tool_args":{"text":"ok"}}]}'

347 + )

348 + stream = _StallingAsyncEventStream(

349 + {"type": "response.output_text.delta", "delta": message}

350 + )

351 +

352 + async def fake_aresponses(*args, **kwargs):

353 + return stream

354 +

355 + async def fake_rate_limiter(*args, **kwargs):

356 + return None

357 +

358 + monkeypatch.setattr(litellm_transport, "aresponses", fake_aresponses)

359 + monkeypatch.setattr(models, "apply_rate_limiter", fake_rate_limiter)

360 +

361 + wrapper = models.LiteLLMChatWrapper(

362 + model="test-model",

363 + provider="openai",

364 + model_config=None,

365 + )

366 +

367 + async def response_callback(chunk: str, full: str):

368 + return full if extract_tools.extract_tool_request(full) else None

369 +

370 + result = await asyncio.wait_for(

371 + wrapper.unified_turn(

372 + messages=[HumanMessage(content="hi")],

373 + response_callback=response_callback,

374 + ),

375 + timeout=1,

376 + )

377 +

378 + assert result.response == message

379 + assert stream.closed is True

380 +

381 +

382 def test_collect_response_ids_from_agent_state_and_history_metadata():

383 payload = {

384 "agents": [

@@ -411,3 +474,137 @@ async def test_agent_executes_native_responses_function_call_and_records_output(

474 "output": "done:a0",

475 }

476 ]

477 +

478 +

479 +@pytest.mark.asyncio

480 +async def test_agent_routes_only_complete_text_tool_requests() -> None:

481 + agent = object.__new__(Agent)

482 + processed: list[str] = []

483 +

484 + async def log_builtin_items(result):

485 + return None

486 +

487 + async def process_tools(message):

488 + processed.append(message)

489 + return message

490 +

491 + agent._log_response_builtin_items = log_builtin_items

492 + agent.process_tools = process_tools

493 +

494 + tool_request = '{"type":"function","name":"response","parameters":{"text":"ok"}}'

495 + for message in (

496 + "Plain final answer.",

497 + '{"status":"planning"}',

498 + f"Example tool JSON: {tool_request}",

499 + ):

500 + assert await Agent.process_llm_result_tools(

501 + agent, LLMResult.from_chat(response=message)

502 + ) == message

503 + assert processed == []

504 +

505 + assert await Agent.process_llm_result_tools(

506 + agent, LLMResult.from_chat(response=tool_request)

507 + ) == tool_request

508 + assert processed == [tool_request]

509 +

510 + processed.clear()

511 + assert await Agent.process_llm_result_tools(

512 + agent, LLMResult(response="", reasoning=tool_request)

513 + ) == tool_request

514 + assert processed == [tool_request]

515 +

516 + processed.clear()

517 + assert await Agent.process_llm_result_tools(

518 + agent, LLMResult(response="", reasoning='{"status":"planning"}')

519 + ) == ""

520 + assert processed == []

521 +

522 +

523 +@pytest.mark.asyncio

524 +async def test_agent_routes_misformatted_tool_intent_to_repair() -> None:

525 + agent = object.__new__(Agent)

526 + processed: list[str] = []

527 +

528 + async def log_builtin_items(result):

529 + return None

530 +

531 + async def process_tools(message):

532 + processed.append(message)

533 + return None

534 +

535 + agent._log_response_builtin_items = log_builtin_items

536 + agent.process_tools = process_tools

537 +

538 + malformed = (

539 + '{"thoughts":["Plan the work", "Run the tools", '

540 + '"headline":"Save results", "tool_name":"parallel", '

541 + '"tool_args":{"tool_calls":[{"tool_name":"memory_save",'

542 + '"tool_args":{"text":"ok"}}],"wait":true}}'

543 + )

544 +

545 + assert await Agent.process_llm_result_tools(

546 + agent, LLMResult.from_chat(response=malformed)

547 + ) is None

548 + assert processed == [malformed]

549 +

550 + processed.clear()

551 + assert await Agent.process_llm_result_tools(

552 + agent, LLMResult(response="", reasoning=malformed)

553 + ) is None

554 + assert processed == [malformed]

555 +

556 + fenced = (

557 + "I will call the tool.\n\n```json\n"

558 + '{"tool_name":"response","tool_args":{"text":"ok"}}\n```'

559 + )

560 + processed.clear()

561 + assert await Agent.process_llm_result_tools(

562 + agent, LLMResult.from_chat(response=fenced)

563 + ) is None

564 + assert processed == [fenced]

565 +

566 +

567 +@pytest.mark.asyncio

568 +async def test_text_tool_execution_uses_normalized_tool_args(monkeypatch) -> None:

569 + class DummyMCPConfig:

570 + def get_tool(self, agent, tool_name):

571 + return None

572 +

573 + class DummyTool:

574 + def __init__(self):

575 + self.args = {}

576 +

577 + async def before_execution(self, **kwargs):

578 + assert self.args == {"text": "ok"}

579 +

580 + async def execute(self, **kwargs):

581 + assert kwargs == {"text": "ok"}

582 + return Response(message=self.args["text"], break_loop=True)

583 +

584 + async def after_execution(self, response):

585 + return None

586 +

587 + async def no_extension(*args, **kwargs):

588 + return None

589 +

590 + async def no_intervention(*args, **kwargs):

591 + return None

592 +

593 + import agent as agent_module

594 + from helpers import mcp_handler

595 +

596 + monkeypatch.setattr(

597 + mcp_handler.MCPConfig, "get_instance", lambda: DummyMCPConfig()

598 + )

599 + monkeypatch.setattr(agent_module.extension, "call_extensions_async", no_extension)

600 +

601 + tool = DummyTool()

602 + agent = object.__new__(Agent)

603 + agent.data = {}

604 + agent.loop_data = LoopData()

605 + agent.handle_intervention = no_intervention

606 + agent.get_tool = lambda **kwargs: tool

607 +

608 + assert await Agent.process_tools(

609 + agent, '{"actions":[{"tool_name":"response","tool_args":{"text":"ok"}}]}'

610 + ) == "ok"

tests/test_stream_tool_early_stop.py

+50 -24

@@ -239,9 +239,7 @@ async def test_unified_call_stops_after_canonical_root_snapshot(monkeypatch):

239 stream = _AsyncChunkStream(

240 [

241 {"type": "response.created"},

242 - _response_event(

243 - '{"tool_name":"response","tool_args":{"text":"hello"}} trailing text'

244 - ),

242 + _response_event('{"tool_name":"response","tool_args":{"text":"hello"}}'),

243 _response_event(" unreachable"),

244 ]

245 )

@@ -273,10 +271,7 @@ async def test_unified_call_stops_after_canonical_root_snapshot(monkeypatch):

271

272 async def response_callback(chunk: str, full: str):

273 seen.append((chunk, full))

276 - snapshot = extract_tools.extract_json_root_string(full)

277 - if snapshot:

278 - return snapshot

279 - return None

274 + return full.strip() if extract_tools.extract_tool_request(full) else None

275

276 response, reasoning = await wrapper.unified_call(

277 messages=[],

@@ -288,11 +283,11 @@ async def test_unified_call_stops_after_canonical_root_snapshot(monkeypatch):

283 assert stream.index == 2

284 assert stream.closed is True

285 assert len(seen) == 1

291 - assert seen[0][1] == '{"tool_name":"response","tool_args":{"text":"hello"}} trailing text'

286 + assert seen[0][1] == '{"tool_name":"response","tool_args":{"text":"hello"}}'

287

288

289 @pytest.mark.asyncio

295 -async def test_unified_call_stops_after_tool_root_with_incidental_json(monkeypatch):

290 +async def test_unified_call_does_not_stop_for_embedded_tool_json(monkeypatch):

291 stream = _AsyncChunkStream(

292 [

293 {"type": "response.created"},

@@ -331,28 +326,21 @@ async def test_unified_call_stops_after_tool_root_with_incidental_json(monkeypat

326

327 async def response_callback(chunk: str, full: str):

328 seen.append((chunk, full))

334 - snapshot = extract_tools.extract_json_root_string(full)

335 - if not snapshot:

336 - return None

337 - parsed_snapshot = extract_tools.json_parse_dirty(snapshot)

338 - if parsed_snapshot is None:

339 - return None

340 - try:

341 - extract_tools.normalize_tool_request(parsed_snapshot)

342 - except ValueError:

343 - return None

344 - return snapshot

329 + return full.strip() if extract_tools.extract_tool_request(full) else None

330

331 response, reasoning = await wrapper.unified_call(

332 messages=[],

333 response_callback=response_callback,

334 )

335

351 - assert response == '{"tool_name":"response","tool_args":{"text":"ok"}}'

336 + assert response == (

337 + 'Preamble {"note":"not the tool"}.\n'

338 + '{"tool_name":"response","tool_args":{"text":"ok"}} trailing text unreachable'

339 + )

340 assert reasoning == ""

353 - assert stream.index == 3

354 - assert stream.closed is True

355 - assert len(seen) == 2

341 + assert stream.index == 4

342 + assert stream.closed is False

343 + assert len(seen) == 3

344 assert seen[0][1] == 'Preamble {"note":"not the tool"}.\n'

345 assert (

346 seen[1][1]

@@ -437,6 +425,44 @@ async def test_chat_completions_escape_hatch_still_uses_acompletion(monkeypatch)

425 assert calls == ["chat"]

426

427

428 +@pytest.mark.asyncio

429 +async def test_unified_turn_stops_chat_stream_after_text_tool_request(monkeypatch):

430 + message = (

431 + '{"thoughts":["test"],"actions":['

432 + '{"tool_name":"response","tool_args":{"text":"ok"}}]}'

433 + )

434 + stream = _AsyncChunkStream([_chunk(message), _chunk(" unreachable")])

435 +

436 + async def fake_acompletion(*args, **kwargs):

437 + assert kwargs["stream"] is True

438 + return stream

439 +

440 + async def fake_rate_limiter(*args, **kwargs):

441 + return None

442 +

443 + monkeypatch.setattr(litellm_transport, "acompletion", fake_acompletion)

444 + monkeypatch.setattr(models, "apply_rate_limiter", fake_rate_limiter)

445 +

446 + wrapper = models.LiteLLMChatWrapper(

447 + model="test-model",

448 + provider="openai",

449 + model_config=None,

450 + a0_api_mode="chat",

451 + )

452 +

453 + async def response_callback(chunk: str, full: str):

454 + return full if extract_tools.extract_tool_request(full) else None

455 +

456 + result = await wrapper.unified_turn(

457 + messages=[],

458 + response_callback=response_callback,

459 + )

460 +

461 + assert result.response == message

462 + assert stream.index == 1

463 + assert stream.closed is True

464 +

465 +

466 @pytest.mark.asyncio

467 async def test_unified_call_retries_responses_with_high_reasoning(monkeypatch):

468 validation_error = ValueError(

tests/test_tool_request_normalization.py

+81 -1

@@ -9,7 +9,12 @@ PROJECT_ROOT = Path(__file__).resolve().parents[1]

9 if str(PROJECT_ROOT) not in sys.path:

10 sys.path.insert(0, str(PROJECT_ROOT))

12 -from helpers.extract_tools import normalize_tool_request

12 +from helpers.extract_tools import (

13 + extract_tool_request,

14 + is_misformatted_tool_request,

15 + json_parse_dirty,

16 + normalize_tool_request,

17 +)

18 from helpers import parallel_tools

@@ -68,6 +73,81 @@ def test_normalize_tool_request_rejects_missing_args() -> None:

73 normalize_tool_request({"tool_name": "response"})

76 +def test_normalize_tool_request_accepts_native_function_format() -> None:

77 + request = {

78 + "type": "function",

79 + "name": "search_engine",

80 + "parameters": {"query": "latest Agent Zero release"},

81 + }

82 +

83 + assert json_parse_dirty(str(request)) == request

84 + assert normalize_tool_request(request) == (

85 + "search_engine",

86 + {"query": "latest Agent Zero release"},

87 + )

88 +

89 +

90 +def test_normalize_tool_request_accepts_single_action_wrapper() -> None:

91 + request = {

92 + "thoughts": ["Read the requested file."],

93 + "actions": [

94 + {

95 + "tool_name": "text_editor",

96 + "tool_args": {"action": "read", "path": "README.md"},

97 + }

98 + ],

99 + }

100 +

101 + assert json_parse_dirty(str(request)) == request

102 + assert normalize_tool_request(request) == (

103 + "text_editor",

104 + {"action": "read", "path": "README.md"},

105 + )

106 +

107 +

108 +def test_normalize_tool_request_rejects_multiple_wrapped_actions() -> None:

109 + with pytest.raises(ValueError, match="exactly one"):

110 + normalize_tool_request(

111 + {

112 + "actions": [

113 + {"tool_name": "response", "tool_args": {"text": "first"}},

114 + {"tool_name": "response", "tool_args": {"text": "second"}},

115 + ]

116 + }

117 + )

118 +

119 +

120 +def test_extract_tool_request_requires_a_complete_tool_message() -> None:

121 + request = '{"tool_name":"response","tool_args":{"text":"ok"}}'

122 +

123 + assert extract_tool_request(request) == {

124 + "tool_name": "response",

125 + "tool_args": {"text": "ok"},

126 + }

127 + assert extract_tool_request('{"status":"ok"}') is None

128 + assert extract_tool_request(f"Example: {request}") is None

129 + assert extract_tool_request(f"{request} trailing text") is None

130 +

131 +

132 +def test_is_misformatted_tool_request_requires_agent_tool_envelope() -> None:

133 + request = '{"tool_name":"response","tool_args":{"text":"ok"}}'

134 + malformed = (

135 + '{"thoughts":["Plan the work", "Run the tools", '

136 + '"headline":"Save results", "tool_name":"parallel", '

137 + '"tool_args":{"tool_calls":[{"tool_name":"memory_save",'

138 + '"tool_args":{"text":"ok"}}],"wait":true}}'

139 + )

140 +

141 + assert extract_tool_request(malformed) is None

142 + assert is_misformatted_tool_request(malformed) is True

143 + assert is_misformatted_tool_request(f"Intro\n```json\n{request}\n```") is True

144 + assert is_misformatted_tool_request('{"status":"planning"}') is False

145 + assert is_misformatted_tool_request(f"Example: {request}") is False

146 + assert is_misformatted_tool_request(

147 + malformed.replace('{"thoughts"', '{"status":"planning","thoughts"')

148 + ) is False

149 +

150 +

151 def test_normalize_parallel_tool_calls_accepts_full_agent_reply_shape() -> None:

152 calls = parallel_tools.normalize_parallel_tool_calls(

153 [