Agent Zero AI framework
stop tool dispatch at first completed json object
Tool execution no longer waits for the full streamed assistant text. We now detect the first explicitly closed top-level JSON object, freeze that snapshot as the canonical tool request, and stop the model stream there for dispatch. To make that safe, DirtyJson completion semantics are tightened so completed=true only means the root object was explicitly closed, not that parsing hit end of file. I also restricted the new extraction path to object roots only, since tool calls are always brace-delimited objects, and added tests for parser completion and early stream stop.
Alessandro committed
Apr 3, 2026 at 16:56 UTC
5a2223596a4aaea93d24e628d15c9d2d30b8ec35
6 files changed
+256
-40
agent.py
+25
-1
@@ -388,6 +388,7 @@ class Agent:
388
self.context.streaming_agent = self # mark self as current streamer
389
self.loop_data.iteration += 1
390
self.loop_data.params_temporary = {} # clear temporary params
391
+ last_response_stream_full = ""
392
393
# call message_loop_start extensions
394
await extension.call_extensions_async(
@@ -425,12 +426,32 @@ class Agent:
426
await self.handle_reasoning_stream(stream_data["full"])
427
428
async def stream_callback(chunk: str, full: str):
429
+ nonlocal last_response_stream_full
430
await self.handle_intervention()
431
# output the agent response stream
432
if chunk == full:
433
printer.print("Response: ") # start of response
434
# Pass chunk and full data to extensions for processing
435
stream_data = {"chunk": chunk, "full": full}
436
+ stop_response: str | None = None
437
+
438
+ snapshot = extract_tools.extract_json_root_string(full)
439
+ if snapshot:
440
+ parsed_snapshot = extract_tools.json_parse_dirty(snapshot)
441
+ if parsed_snapshot is not None:
442
+ try:
443
+ await self.validate_tool_request(parsed_snapshot)
444
+ except Exception:
445
+ pass
446
+ else:
447
+ previous_full = last_response_stream_full
448
+ stream_data["full"] = snapshot
449
+ if snapshot.startswith(previous_full):
450
+ stream_data["chunk"] = snapshot[len(previous_full) :]
451
+ else:
452
+ stream_data["chunk"] = snapshot
453
+ stop_response = snapshot
454
+
455
await extension.call_extensions_async(
456
"response_stream_chunk",
457
self,
@@ -442,6 +463,9 @@ class Agent:
463
printer.stream(stream_data["chunk"])
464
# Use the potentially modified full text for downstream processing
465
await self.handle_response_stream(stream_data["full"])
466
+ last_response_stream_full = stream_data["full"]
467
+ if stop_response is not None:
468
+ return stop_response
469
470
# call main LLM
471
agent_response, _reasoning = await self.call_chat_model(
@@ -770,7 +794,7 @@ class Agent:
794
async def call_chat_model(
795
self,
796
messages: list[BaseMessage],
773
- response_callback: Callable[[str, str], Awaitable[None]] | None = None,
797
+ response_callback: Callable[[str, str], Awaitable[str | None]] | None = None,
798
reasoning_callback: Callable[[str, str], Awaitable[None]] | None = None,
799
background: bool = False,
800
explicit_caching: bool = True,
helpers/dirty_json.py
+16
-8
@@ -28,10 +28,10 @@ class DirtyJson:
28
self.completed = False
29
self._parsing_started = False
30
31
- def _pop_stack(self):
32
- """Pop from the parsing stack and mark completed if root structure is closed."""
31
+ def _pop_stack(self, root_closed: bool = False):
32
+ """Pop from the parsing stack and mark completed only on an explicit root close."""
33
self.stack.pop()
34
- if self._parsing_started and not self.stack:
34
+ if root_closed and self._parsing_started and not self.stack:
35
self.completed = True
36
37
@staticmethod
@@ -103,6 +103,8 @@ class DirtyJson:
103
self._advance()
104
105
def _parse(self):
106
+ if self.completed and not self.stack:
107
+ return
108
if self.result is None:
109
self.result = self._parse_value()
110
else:
@@ -110,6 +112,8 @@ class DirtyJson:
112
113
def _continue_parsing(self):
114
while self.current_char is not None:
115
+ if self.completed and not self.stack:
116
+ return
117
if isinstance(self.result, dict):
118
self._parse_object_content()
119
elif isinstance(self.result, list):
@@ -122,7 +126,9 @@ class DirtyJson:
126
def _parse_value(self):
127
self._skip_whitespace()
128
if self.current_char == "{":
125
- if self._peek(1) == "{": # Handle {{
129
+ # Only treat doubled braces as a wrapper at the root; nested objects
130
+ # must keep their closing braces paired correctly.
131
+ if not self.stack and self._peek(1) == "{": # Handle {{
132
self._advance(2)
133
return self._parse_object()
134
elif self.current_char == "[":
@@ -169,11 +175,13 @@ class DirtyJson:
175
while self.current_char is not None:
176
self._skip_whitespace()
177
if self.current_char == "}":
172
- if self._peek(1) == "}": # Handle }}
178
+ # Root-level wrapper outputs may end in "}}"; nested objects must
179
+ # still close one brace at a time.
180
+ if len(self.stack) == 1 and self._peek(1) == "}": # Handle }}
181
self._advance(2)
182
else:
183
self._advance()
176
- self._pop_stack()
184
+ self._pop_stack(root_closed=True)
185
return
186
if self.current_char is None:
187
self._pop_stack()
@@ -234,7 +242,7 @@ class DirtyJson:
242
self._skip_whitespace()
243
if self.current_char == "]":
244
self._advance()
237
- self._pop_stack()
245
+ self._pop_stack(root_closed=True)
246
return
247
value = self._parse_value()
248
self.stack[-1].append(value)
@@ -246,7 +254,7 @@ class DirtyJson:
254
if self.current_char is None or self.current_char == "]":
255
if self.current_char == "]":
256
self._advance()
249
- self._pop_stack()
257
+ self._pop_stack(root_closed=True)
258
return
259
elif self.current_char != "]":
260
self._pop_stack()
helpers/extract_tools.py
+22
@@ -19,6 +19,28 @@ def json_parse_dirty(json: str) -> dict[str, Any] | None:
19
return None
20
return None
21
22
+def extract_json_root_string(content: str) -> str | None:
23
+ if not content or not isinstance(content, str):
24
+ return None
25
+
26
+ start = content.find("{")
27
+ if start == -1:
28
+ return None
29
+ first_array = content.find("[")
30
+ if first_array != -1 and first_array < start:
31
+ return None
32
+
33
+ parser = DirtyJson()
34
+ try:
35
+ parser.parse(content[start:])
36
+ except Exception:
37
+ return None
38
+
39
+ if not parser.completed:
40
+ return None
41
+
42
+ return content[start : start + parser.index]
43
+
44
45
def extract_json_object_string(content):
46
start = content.find("{")
models.py
+41
-31
@@ -475,7 +475,7 @@ class LiteLLMChatWrapper(SimpleChatModel):
475
system_message="",
476
user_message="",
477
messages: List[BaseMessage] | None = None,
478
- response_callback: Callable[[str, str], Awaitable[None]] | None = None,
478
+ response_callback: Callable[[str, str], Awaitable[str | None]] | None = None,
479
reasoning_callback: Callable[[str, str], Awaitable[None]] | None = None,
480
tokens_callback: Callable[[str, int], Awaitable[None]] | None = None,
481
rate_limiter_callback: (
@@ -526,36 +526,46 @@ class LiteLLMChatWrapper(SimpleChatModel):
526
527
if stream:
528
# iterate over chunks
529
- async for chunk in _completion: # type: ignore
530
- got_any_chunk = True
531
- # parse chunk
532
- parsed = _parse_chunk(chunk)
533
- output = result.add_chunk(parsed)
534
-
535
- # collect reasoning delta and call callbacks
536
- if output["reasoning_delta"]:
537
- if reasoning_callback:
538
- await reasoning_callback(output["reasoning_delta"], result.reasoning)
539
- if tokens_callback:
540
- await tokens_callback(
541
- output["reasoning_delta"],
542
- approximate_tokens(output["reasoning_delta"]),
543
- )
544
- # Add output tokens to rate limiter if configured
545
- if limiter:
546
- limiter.add(output=approximate_tokens(output["reasoning_delta"]))
547
- # collect response delta and call callbacks
548
- if output["response_delta"]:
549
- if response_callback:
550
- await response_callback(output["response_delta"], result.response)
551
- if tokens_callback:
552
- await tokens_callback(
553
- output["response_delta"],
554
- approximate_tokens(output["response_delta"]),
555
- )
556
- # Add output tokens to rate limiter if configured
557
- if limiter:
558
- limiter.add(output=approximate_tokens(output["response_delta"]))
529
+ stop_response: str | None = None
530
+ try:
531
+ async for chunk in _completion: # type: ignore
532
+ got_any_chunk = True
533
+ # parse chunk
534
+ parsed = _parse_chunk(chunk)
535
+ output = result.add_chunk(parsed)
536
+
537
+ # collect reasoning delta and call callbacks
538
+ if output["reasoning_delta"]:
539
+ if reasoning_callback:
540
+ await reasoning_callback(output["reasoning_delta"], result.reasoning)
541
+ if tokens_callback:
542
+ await tokens_callback(
543
+ output["reasoning_delta"],
544
+ approximate_tokens(output["reasoning_delta"]),
545
+ )
546
+ # Add output tokens to rate limiter if configured
547
+ if limiter:
548
+ limiter.add(output=approximate_tokens(output["reasoning_delta"]))
549
+ # collect response delta and call callbacks
550
+ if output["response_delta"]:
551
+ if response_callback:
552
+ stop_response = await response_callback(
553
+ output["response_delta"], result.response
554
+ )
555
+ if tokens_callback:
556
+ await tokens_callback(
557
+ output["response_delta"],
558
+ approximate_tokens(output["response_delta"]),
559
+ )
560
+ # Add output tokens to rate limiter if configured
561
+ if limiter:
562
+ limiter.add(output=approximate_tokens(output["response_delta"]))
563
+ if stop_response is not None:
564
+ result.response = stop_response
565
+ break
566
+ finally:
567
+ if stop_response is not None and hasattr(_completion, "aclose"):
568
+ await _completion.aclose() # type: ignore[attr-defined]
569
570
# non-stream response
571
else:
tests/test_dirty_json.py
new
+56
@@ -0,0 +1,56 @@
1
+from __future__ import annotations
2
+
3
+import sys
4
+from pathlib import Path
5
+
6
+import pytest
7
+
8
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
9
+if str(PROJECT_ROOT) not in sys.path:
10
+ sys.path.insert(0, str(PROJECT_ROOT))
11
+
12
+from helpers.dirty_json import DirtyJson
13
+
14
+
15
+@pytest.mark.parametrize(
16
+ ("payload", "expected"),
17
+ [
18
+ (
19
+ '{"tool_name":"x","tool_args":{}}',
20
+ {"tool_name": "x", "tool_args": {}},
21
+ ),
22
+ ("[1, 2, 3]", [1, 2, 3]),
23
+ ],
24
+)
25
+def test_completed_true_when_root_is_explicitly_closed(payload, expected) -> None:
26
+ parser = DirtyJson()
27
+
28
+ assert parser.parse(payload) == expected
29
+ assert parser.completed is True
30
+
31
+
32
+def test_completed_false_when_root_hits_eof_before_closing() -> None:
33
+ parser = DirtyJson()
34
+
35
+ assert parser.parse('{"tool_name":"x","tool_args":{}') == {
36
+ "tool_name": "x",
37
+ "tool_args": {},
38
+ }
39
+ assert parser.completed is False
40
+
41
+
42
+def test_completed_remains_true_after_trailing_content() -> None:
43
+ parser = DirtyJson()
44
+
45
+ assert parser.feed('{"tool_name":"x","tool_args":{}}') == {
46
+ "tool_name": "x",
47
+ "tool_args": {},
48
+ }
49
+ assert parser.completed is True
50
+
51
+ assert parser.feed(" trailing noise") == {
52
+ "tool_name": "x",
53
+ "tool_args": {},
54
+ }
55
+
56
+ assert parser.completed is True
tests/test_stream_tool_early_stop.py
new
+96
@@ -0,0 +1,96 @@
1
+import sys
2
+from pathlib import Path
3
+
4
+import pytest
5
+
6
+
7
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
8
+if str(PROJECT_ROOT) not in sys.path:
9
+ sys.path.insert(0, str(PROJECT_ROOT))
10
+
11
+import models
12
+from helpers import extract_tools
13
+
14
+
15
+def _chunk(content: str) -> dict:
16
+ return {"choices": [{"delta": {"content": content}, "message": {}}]}
17
+
18
+
19
+class _AsyncChunkStream:
20
+ def __init__(self, chunks: list[dict]):
21
+ self._chunks = chunks
22
+ self.index = 0
23
+
24
+ def __aiter__(self):
25
+ return self
26
+
27
+ async def __anext__(self):
28
+ if self.index >= len(self._chunks):
29
+ raise StopAsyncIteration
30
+ chunk = self._chunks[self.index]
31
+ self.index += 1
32
+ return chunk
33
+
34
+
35
+def test_extract_json_root_string_returns_canonical_snapshot():
36
+ text = (
37
+ 'prefix {"tool_name":"response","tool_args":{"text":"brace } inside"}} '
38
+ "trailing noise"
39
+ )
40
+
41
+ root = extract_tools.extract_json_root_string(text)
42
+
43
+ assert root == '{"tool_name":"response","tool_args":{"text":"brace } inside"}}'
44
+ assert extract_tools.json_parse_dirty(root)["tool_args"]["text"] == "brace } inside"
45
+ assert extract_tools.extract_json_root_string(
46
+ '{"tool_name":"response","tool_args":{"text":"missing"'
47
+ ) is None
48
+ assert extract_tools.extract_json_root_string('[{"tool_name":"response"}]') is None
49
+
50
+
51
+@pytest.mark.asyncio
52
+async def test_unified_call_stops_after_canonical_root_snapshot(monkeypatch):
53
+ stream = _AsyncChunkStream(
54
+ [
55
+ _chunk(
56
+ '{"tool_name":"response","tool_args":{"text":"hello"}} trailing text'
57
+ ),
58
+ _chunk(" unreachable"),
59
+ ]
60
+ )
61
+
62
+ async def fake_acompletion(*args, **kwargs):
63
+ assert kwargs["stream"] is True
64
+ return stream
65
+
66
+ async def fake_rate_limiter(*args, **kwargs):
67
+ return None
68
+
69
+ monkeypatch.setattr(models, "acompletion", fake_acompletion)
70
+ monkeypatch.setattr(models, "apply_rate_limiter", fake_rate_limiter)
71
+
72
+ wrapper = models.LiteLLMChatWrapper(
73
+ model="test-model",
74
+ provider="openai",
75
+ model_config=None,
76
+ )
77
+
78
+ seen: list[tuple[str, str]] = []
79
+
80
+ async def response_callback(chunk: str, full: str):
81
+ seen.append((chunk, full))
82
+ snapshot = extract_tools.extract_json_root_string(full)
83
+ if snapshot:
84
+ return snapshot
85
+ return None
86
+
87
+ response, reasoning = await wrapper.unified_call(
88
+ messages=[],
89
+ response_callback=response_callback,
90
+ )
91
+
92
+ assert response == '{"tool_name":"response","tool_args":{"text":"hello"}}'
93
+ assert reasoning == ""
94
+ assert stream.index == 1
95
+ assert len(seen) == 1
96
+ assert seen[0][1] == '{"tool_name":"response","tool_args":{"text":"hello"}} trailing text'