Improve Context Doctor output handling
Repair completed non-tool output into compact thoughts JSON, preserve schema-guided selection of the most complete repaired tool call, refresh log fields, and create a response log only for response tool calls using text. Suppress XML-shaped malformed responses as configured. Cover raw fallback, repaired-call selection, and response-log handling.
linkliti committed
Aug 24, 2026 at 21:43 UTC
e85842c5dc81b828d875bdf24ad4b1b1c50deb9a
4 files changed
+195
-71
plugins/_context_doctor/AGENTS.md
+8
-7
@@ -2,24 +2,25 @@
2
3
## Purpose
4
5
-- Repair malformed Agent Zero tool-call JSON and persist compact repaired output.
5
+- Repair malformed Agent Zero tool-call JSON and preserve raw output as compact thoughts JSON when repair cannot produce a tool call.
6
7
## Ownership
8
9
-- `helpers/context_doctor.py` validates and repairs tool-call JSON.
9
+- `helpers/context_doctor.py` transforms output and refreshes log fields.
10
- `extensions/python/message_loop_result/` normalizes completed model output before default processing.
11
-- `webui/config.html` exposes XML-output suppression only.
11
+- `webui/config.html` exposes XML suppression and log-detail settings.
12
13
## Local Contracts
14
15
-- Repaired tool-call JSON is always minified.
16
-- Invalid non-tool output is unchanged; native processing retains ownership.
17
-- Do not write settings that alter repair mode or log content.
15
+- Repaired and fallback JSON is always minified.
16
+- Nonempty non-tool output becomes `{"thoughts":[raw]}`; XML-like output becomes `{}` only when suppression is enabled.
17
+- Log kvps and heading always reflect transformed output; `update_log` controls only View Details content.
18
+- A repaired `response` tool call refreshes the response log item when streaming did not create it.
19
20
## Work Guidance
21
22
- Keep repair scoped to complete tool-call JSON.
22
-- Use framework-installed `json_repair`; apply the plugin-local parser patch before repair. Do not vendor dependencies.
23
+- Use framework-installed `json_repair`; apply plugin-local parser patch before repair. Do not vendor dependencies.
24
25
## Verification
26
plugins/_context_doctor/extensions/python/message_loop_result/_10_context_doctor.py
+38
-10
@@ -1,12 +1,13 @@
1
-"""Repair and minify model tool-call JSON before default processing."""
1
+"""Repair model output and refresh its log details before default processing."""
2
3
from __future__ import annotations
4
5
+import json
6
from typing import Any, override
7
8
from helpers.extension import Extension
9
from helpers.plugins import get_plugin_config
9
-from plugins._context_doctor.helpers.context_doctor import repair_and_minify, update_log_item
10
+from plugins._context_doctor.helpers.context_doctor import transform_response, update_log_item
11
12
13
class ContextDoctor(Extension):
@@ -17,18 +18,45 @@ class ContextDoctor(Extension):
18
19
llm_result = result_data.get("llm_result")
20
response = getattr(llm_result, "response", None)
20
- if not isinstance(response, str):
21
+ if not isinstance(response, str) or not response:
22
return
23
24
config = get_plugin_config("_context_doctor", agent=self.agent) or {}
24
- repaired = repair_and_minify(
25
+ transformed = transform_response(
26
response, suppress_xml=config.get("suppress_xml", True)
27
)
27
- if repaired is None:
28
- return
28
+ llm_result.response = transformed
29
30
- llm_result.response = repaired
30
params = getattr(getattr(self.agent, "loop_data", None), "params_temporary", None)
32
- log_item = params.get("log_item_generating") if isinstance(params, dict) else None
33
- if log_item is not None and config.get("update_log", False):
34
- update_log_item(self.agent, log_item, repaired)
31
+ if not isinstance(params, dict) or transformed == "{}":
32
+ return
33
+
34
+ log_item = params.get("log_item_generating")
35
+ if log_item is not None:
36
+ update_log_item(
37
+ self.agent,
38
+ log_item,
39
+ transformed,
40
+ update_log=config.get("update_log", False),
41
+ raw_response=response,
42
+ )
43
+
44
+ try:
45
+ parsed = json.loads(transformed)
46
+ tool_args = parsed.get("tool_args", {}) if isinstance(parsed, dict) else {}
47
+ response_text = tool_args.get("text") if parsed.get("tool_name") == "response" else None
48
+ except (AttributeError, TypeError, ValueError):
49
+ response_text = None
50
+ if not isinstance(response_text, str) or not response_text.strip():
51
+ return
52
+
53
+ response_item = params.get("log_item_response")
54
+ if response_item is None and log_item is not None:
55
+ response_item = self.agent.context.log.log(
56
+ type="response",
57
+ heading=f"icon://chat {self.agent.agent_name}: Responding",
58
+ id=getattr(log_item, "id", ""),
59
+ )
60
+ params["log_item_response"] = response_item
61
+ if response_item is not None:
62
+ response_item.update(content=response_text)
plugins/_context_doctor/helpers/context_doctor.py
+76
-26
@@ -1,4 +1,4 @@
1
-"""Repair complete Agent Zero tool-call JSON into compact JSON."""
1
+"""Repair completed model output into compact JSON."""
2
3
from __future__ import annotations
4
@@ -22,38 +22,88 @@ def _is_tool_call(value: Any) -> bool:
22
)
23
24
25
-def repair_and_minify(response: str, *, suppress_xml: bool) -> str | None:
26
- """Return compact tool-call JSON, XML fallback, or ``None`` for other output."""
27
- if not response:
28
- return None
25
+_A0_SALVAGE_SCHEMA = {
26
+ "type": "object",
27
+ "properties": {
28
+ "thoughts": {"type": "array", "items": {"type": "string"}},
29
+ "headline": {"type": "string"},
30
+ "tool_name": {"type": "string"},
31
+ "tool_args": {"type": "object"},
32
+ },
33
+ "required": [],
34
+}
35
30
- try:
31
- from plugins._context_doctor.helpers.json_repair_patch import apply_patch
32
- from json_repair import repair_json
36
34
- apply_patch()
35
- repaired = repair_json(response, return_objects=True)
36
- except Exception:
37
- return "{}" if suppress_xml and "<" in response and ">" in response else None
37
+def _a0_completeness_score(value: dict[str, Any]) -> int:
38
+ return sum(field in value for field in ("thoughts", "headline", "tool_name", "tool_args")) + bool(
39
+ value.get("tool_args")
40
+ )
41
+
42
39
- if isinstance(repaired, list):
40
- repaired = next((item for item in repaired if _is_tool_call(item)), None)
41
- if _is_tool_call(repaired):
42
- return json.dumps(repaired, ensure_ascii=False, separators=(",", ":"))
43
- return "{}" if suppress_xml and "<" in response and ">" in response else None
43
+def transform_response(response: str, *, suppress_xml: bool) -> str:
44
+ """Return compact repaired tool JSON or a compact raw-text fallback."""
45
+ if response:
46
+ try:
47
+ from plugins._context_doctor.helpers.json_repair_patch import apply_patch
48
+ from json_repair import repair_json
49
50
+ apply_patch()
51
+ try:
52
+ repaired = repair_json(
53
+ response,
54
+ return_objects=True,
55
+ schema=_A0_SALVAGE_SCHEMA,
56
+ schema_repair_mode="salvage",
57
+ )
58
+ except Exception:
59
+ repaired = repair_json(response, return_objects=True)
60
+
61
+ candidates = repaired if isinstance(repaired, list) else [repaired]
62
+ try:
63
+ no_schema = repair_json(response, return_objects=True)
64
+ except Exception:
65
+ no_schema = None
66
+ if isinstance(no_schema, list):
67
+ valid = [item for item in no_schema if _is_tool_call(item)]
68
+ if len(valid) > 1 or (len(valid) == 1 and not _is_tool_call(repaired)):
69
+ candidates = no_schema
70
+ repaired = max(
71
+ (item for item in candidates if _is_tool_call(item)),
72
+ key=_a0_completeness_score,
73
+ default=None,
74
+ )
75
+ except Exception:
76
+ repaired = None
77
46
-def update_log_item(agent: Any, log_item: Any, response: str) -> None:
47
- """Replace final log details with repaired JSON and derived display fields."""
78
+ if _is_tool_call(repaired):
79
+ return json.dumps(repaired, ensure_ascii=False, separators=(",", ":"))
80
+
81
+ if suppress_xml and "<" in response and ">" in response:
82
+ return "{}"
83
+ return json.dumps({"thoughts": [response]}, ensure_ascii=False, separators=(",", ":"))
84
+
85
+
86
+def update_log_item(
87
+ agent: Any,
88
+ log_item: Any,
89
+ response: str,
90
+ *,
91
+ update_log: bool,
92
+ raw_response: str,
93
+) -> None:
94
+ """Refresh log fields from transformed JSON; optionally replace raw details."""
95
try:
96
parsed = json.loads(response)
50
- if not _is_tool_call(parsed):
97
+ if not isinstance(parsed, dict):
98
return
52
- heading = parsed.get("headline") or f"Using {parsed['tool_name']}"
53
- log_item.update(
54
- content=response,
55
- kvps=parsed,
56
- heading=f"{getattr(agent, 'agent_name', 'A0')}: {heading}",
57
- )
99
+ heading = parsed.get("headline")
100
+ if not isinstance(heading, str) or not heading:
101
+ tool_name = parsed.get("tool_name")
102
+ heading = f"Using {tool_name}" if isinstance(tool_name, str) else ""
103
+ kwargs: dict[str, Any] = {"kvps": parsed}
104
+ if heading:
105
+ kwargs["heading"] = f"{getattr(agent, 'agent_name', 'A0')}: {heading}"
106
+ kwargs["content"] = response if update_log else raw_response
107
+ log_item.update(**kwargs)
108
except (AttributeError, TypeError, ValueError):
109
pass
plugins/_context_doctor/tests/test_context_doctor.py
+73
-28
@@ -1,78 +1,123 @@
1
from types import SimpleNamespace
2
3
-from plugins._context_doctor.helpers.context_doctor import repair_and_minify, update_log_item
3
from plugins._context_doctor.extensions.python.message_loop_result._10_context_doctor import (
4
ContextDoctor,
5
)
6
+from plugins._context_doctor.helpers.context_doctor import transform_response, update_log_item
7
8
9
def test_repairs_and_minifies_tool_call():
10
response = '{"tool_name":"response","tool_args":{"text":"ok",},}'
11
12
- assert repair_and_minify(response, suppress_xml=True) == (
12
+ assert transform_response(response, suppress_xml=True) == (
13
'{"tool_name":"response","tool_args":{"text":"ok"}}'
14
)
15
16
17
-def test_ignores_non_tool_json():
18
- assert repair_and_minify('{"message":"ok"}', suppress_xml=True) is None
19
-
17
+def test_chooses_most_complete_tool_call():
18
+ response = (
19
+ '{"tool_name":"first","tool_args":{}} '
20
+ '{"thoughts":["x"],"headline":"Second","tool_name":"second","tool_args":{"x":1}}'
21
+ )
22
21
-def test_suppresses_xml_when_enabled():
22
- assert repair_and_minify('<tool>response</tool>', suppress_xml=True) == "{}"
23
- assert repair_and_minify('<tool>response</tool>', suppress_xml=False) is None
23
+ assert transform_response(response, suppress_xml=True) == (
24
+ '{"thoughts":["x"],"headline":"Second","tool_name":"second","tool_args":{"x":1}}'
25
+ )
26
27
26
-def test_extension_replaces_completed_result(monkeypatch):
27
- monkeypatch.setattr(
28
- "plugins._context_doctor.extensions.python.message_loop_result._10_context_doctor.get_plugin_config",
29
- lambda *args, **kwargs: {"suppress_xml": True},
28
+def test_wraps_raw_text_in_thoughts():
29
+ assert transform_response("plain text", suppress_xml=True) == (
30
+ '{"thoughts":["plain text"]}'
31
)
31
- llm_result = SimpleNamespace(
32
- response='{"tool_name":"response","tool_args":{"text":"ok",},}'
33
- )
34
- agent = SimpleNamespace(loop_data=SimpleNamespace(params_temporary={}))
32
36
- ContextDoctor(agent).execute({"llm_result": llm_result})
33
38
- assert llm_result.response == '{"tool_name":"response","tool_args":{"text":"ok"}}'
34
+def test_suppresses_xml_when_enabled():
35
+ assert transform_response("<tool>response</tool>", suppress_xml=True) == "{}"
36
+ assert transform_response("<tool>response</tool>", suppress_xml=False) == (
37
+ '{"thoughts":["<tool>response</tool>"]}'
38
+ )
39
40
41
-def test_updates_log_with_repaired_tool_call():
41
+def test_updates_log_kvps_and_heading_while_preserving_raw_content():
42
log_item = SimpleNamespace(update=lambda **kwargs: setattr(log_item, "data", kwargs))
43
+ raw = '{"tool_name":"response","tool_args":{"text":"ok",},}'
44
+ repaired = '{"headline":"Done","tool_name":"response","tool_args":{"text":"ok"}}'
45
46
update_log_item(
47
SimpleNamespace(agent_name="A0"),
48
log_item,
47
- '{"headline":"Done","tool_name":"response","tool_args":{"text":"ok"}}',
49
+ repaired,
50
+ update_log=False,
51
+ raw_response=raw,
52
)
53
50
- assert log_item.data["content"] == (
51
- '{"headline":"Done","tool_name":"response","tool_args":{"text":"ok"}}'
52
- )
54
+ assert log_item.data["content"] == raw
55
+ assert log_item.data["kvps"]["tool_name"] == "response"
56
assert log_item.data["heading"] == "A0: Done"
57
58
56
-def test_extension_updates_log_only_when_enabled(monkeypatch):
59
+def test_extension_replaces_result_refreshes_log_and_response_item(monkeypatch):
60
+ monkeypatch.setattr(
61
+ "plugins._context_doctor.extensions.python.message_loop_result._10_context_doctor.get_plugin_config",
62
+ lambda *args, **kwargs: {"suppress_xml": True, "update_log": False},
63
+ )
64
llm_result = SimpleNamespace(
65
response='{"tool_name":"response","tool_args":{"text":"ok",},}'
66
)
60
- log_item = SimpleNamespace(update=lambda **kwargs: setattr(log_item, "data", kwargs))
67
+ log_item = SimpleNamespace(
68
+ id="generating", update=lambda **kwargs: setattr(log_item, "data", kwargs)
69
+ )
70
+ response_item = SimpleNamespace(update=lambda **kwargs: setattr(response_item, "data", kwargs))
71
+ context = SimpleNamespace(log=SimpleNamespace(log=lambda **kwargs: response_item))
72
agent = SimpleNamespace(
73
agent_name="A0",
74
+ context=context,
75
loop_data=SimpleNamespace(params_temporary={"log_item_generating": log_item}),
76
)
77
78
+ ContextDoctor(agent).execute({"llm_result": llm_result})
79
+
80
+ assert llm_result.response == '{"tool_name":"response","tool_args":{"text":"ok"}}'
81
+ assert log_item.data["content"] != llm_result.response
82
+ assert log_item.data["kvps"]["tool_name"] == "response"
83
+ assert response_item.data == {"content": "ok"}
84
+
85
+
86
+def test_extension_does_not_create_response_item_for_other_tools(monkeypatch):
87
monkeypatch.setattr(
88
"plugins._context_doctor.extensions.python.message_loop_result._10_context_doctor.get_plugin_config",
89
lambda *args, **kwargs: {"suppress_xml": True, "update_log": False},
90
)
91
+ log_item = SimpleNamespace(id="generating", update=lambda **kwargs: None)
92
+ agent = SimpleNamespace(
93
+ agent_name="A0",
94
+ context=SimpleNamespace(log=SimpleNamespace(log=lambda **kwargs: None)),
95
+ loop_data=SimpleNamespace(params_temporary={"log_item_generating": log_item}),
96
+ )
97
+ llm_result = SimpleNamespace(
98
+ response='{"tool_name":"notify_user","tool_args":{"message":"not final"}}'
99
+ )
100
+
101
ContextDoctor(agent).execute({"llm_result": llm_result})
71
- assert not hasattr(log_item, "data")
102
103
+ assert "log_item_response" not in agent.loop_data.params_temporary
104
+
105
+
106
+def test_extension_does_not_use_legacy_response_message_key(monkeypatch):
107
monkeypatch.setattr(
108
"plugins._context_doctor.extensions.python.message_loop_result._10_context_doctor.get_plugin_config",
75
- lambda *args, **kwargs: {"suppress_xml": True, "update_log": True},
109
+ lambda *args, **kwargs: {"suppress_xml": True, "update_log": False},
110
+ )
111
+ log_item = SimpleNamespace(id="generating", update=lambda **kwargs: None)
112
+ agent = SimpleNamespace(
113
+ agent_name="A0",
114
+ context=SimpleNamespace(log=SimpleNamespace(log=lambda **kwargs: None)),
115
+ loop_data=SimpleNamespace(params_temporary={"log_item_generating": log_item}),
116
+ )
117
+ llm_result = SimpleNamespace(
118
+ response='{"tool_name":"response","tool_args":{"message":"not final"}}'
119
)
120
+
121
ContextDoctor(agent).execute({"llm_result": llm_result})
78
- assert log_item.data["content"] == llm_result.response
122
+
123
+ assert "log_item_response" not in agent.loop_data.params_temporary