Preserve native vision output from parallel workers
Queue model-visible history produced by direct parallel workers and promote it only after the parent parallel result is recorded. This keeps native vision images available to Main while preserving provider-safe ordering, background collection, and the delegated sidecar path.
Alessandro committed
Aug 26, 2026 at 00:21 UTC
66cc11096c7ec0cab0cb346ac6933e887366da61
8 files changed
+163
-22
helpers/parallel_tools.py
+41
-6
@@ -65,6 +65,7 @@ class ParallelJob:
65
log_id: str = field(default_factory=lambda: str(uuid.uuid4()))
66
log_item: "LogItem | None" = field(default=None, repr=False)
67
deferred_task: DeferredTask | None = field(default=None, repr=False)
68
+ parent_history: list[tuple[Any, int]] = field(default_factory=list, repr=False)
69
70
def elapsed(self) -> float:
71
end = self.completed_at or time.time()
@@ -170,6 +171,26 @@ def is_parallel_worker(agent: "Agent | None") -> bool:
171
return _parallel_worker_kind(agent) == "tool"
172
173
174
+def queue_parallel_parent_history(
175
+ agent: "Agent",
176
+ *,
177
+ content: Any,
178
+ tokens: int = 0,
179
+) -> bool:
180
+ if not is_parallel_worker(agent):
181
+ return False
182
+ context = agent.context
183
+ parent_context_id = str(
184
+ context.get_data(PARALLEL_WORKER_PARENT_CONTEXT_KEY) or ""
185
+ )
186
+ job_id = str(context.get_data(PARALLEL_WORKER_JOB_KEY) or "")
187
+ job = _get_job(parent_context_id, job_id)
188
+ if not job or job.kind != "tool":
189
+ return False
190
+ job.parent_history.append((content, tokens))
191
+ return True
192
+
193
+
194
def _jobs_for_context(context: "AgentContext") -> dict[str, ParallelJob]:
195
jobs = context.get_data(PARALLEL_JOBS_KEY)
196
if not isinstance(jobs, dict):
@@ -242,7 +263,6 @@ async def await_parallel_jobs(
263
raise ValueError("No `job_ids` were provided to await.")
264
265
deadline = time.time() + timeout
245
- known_job_ids = set(job_ids)
266
wait_timed_out_job_ids: set[str] = set()
267
while True:
268
await refresh_parallel_jobs(agent)
@@ -271,11 +291,7 @@ async def await_parallel_jobs(
291
snapshots.append(snapshot)
292
293
if collect:
274
- for job_id in known_job_ids:
275
- job = _jobs_for_context(agent.context).get(job_id)
276
- if job and job.state in TERMINAL_STATES:
277
- await cleanup_parallel_job(agent, job)
278
- _jobs_for_context(agent.context).pop(job_id, None)
294
+ await collect_parallel_jobs(agent, job_ids)
295
296
return snapshots
297
@@ -326,6 +342,25 @@ async def cleanup_parallel_job(agent: "Agent", job: ParallelJob) -> None:
342
await _remove_context(job.worker_context_id)
343
344
345
+async def collect_parallel_jobs(
346
+ agent: "Agent",
347
+ job_ids: list[str],
348
+ *,
349
+ promote_parent_history: bool = False,
350
+) -> None:
351
+ jobs = _jobs_for_context(agent.context)
352
+ for job_id in dict.fromkeys(job_ids):
353
+ job = jobs.get(job_id)
354
+ if not job or job.state not in TERMINAL_STATES:
355
+ continue
356
+ if promote_parent_history:
357
+ for content, tokens in job.parent_history:
358
+ agent.hist_add_message(False, content=content, tokens=tokens)
359
+ job.parent_history.clear()
360
+ await cleanup_parallel_job(agent, job)
361
+ jobs.pop(job_id, None)
362
+
363
+
364
async def build_parallel_jobs_extras(agent: "Agent") -> str:
365
await refresh_parallel_jobs(agent)
366
jobs = [
helpers/parallel_tools.py.dox.md
+2
@@ -38,6 +38,7 @@
38
- Child tool logs mirror normal tool-call visible args; job ids remain available through wrapper results and prompt extras rather than visible process-step args.
39
- Wrapped tool child logs use each tool's native `get_log_object()` output when available, preserving special log rendering (for example: `code_execution_tool` uses `code_exe`, `wait` uses `progress`, MCP tools use `mcp`, and regular tools use `tool`).
40
- Direct parallel worker execution reuses the parent-visible child log item so tool `before_execution()` cannot create a second generic worker log or lose the native badge type.
41
+- Direct tools may explicitly queue model-visible history for their parent. Terminal collection records the outer `parallel` result first, promotes queued messages in job order, and only then removes disposable worker state; background jobs retain queued history until they are collected.
42
- Job IDs are stable handles for later await, collect, or cancel operations.
43
- Prompt extras must stay bounded and expose only job IDs, tool names, status, and compact result/error summaries.
44
@@ -47,6 +48,7 @@
48
- `wait=True` starts jobs and awaits them before returning until all requested jobs finish or the wait timeout is reached; the timeout stops waiting but does not cancel running jobs.
49
- `collect` returns already-finished job results without waiting; `await` waits for requested job IDs.
50
- Canceled jobs should be marked terminal and should stop their background `DeferredTask` when cancellation is possible.
51
+- `queue_parallel_parent_history(...)` accepts messages only from registered direct tool workers. `collect_parallel_jobs(...)` optionally promotes those messages while collecting terminal jobs; arbitrary worker history is never copied.
52
53
## Work Guidance
54
tests/test_parallel_tool.py
+55
@@ -1046,6 +1046,61 @@ async def test_parallel_tool_keeps_wrapper_out_of_visible_log() -> None:
1046
assert agent.tool_results == [("parallel", "done", {"extra": "value"})]
1047
1048
1049
+@pytest.mark.asyncio
1050
+async def test_parallel_collect_promotes_history_after_wrapper_result() -> None:
1051
+ from tools.parallel import ParallelTool
1052
+
1053
+ class HistoryAgent(_FakeAgent):
1054
+ def __init__(self) -> None:
1055
+ super().__init__()
1056
+ self.history_events = []
1057
+
1058
+ def hist_add_tool_result(self, tool_name, tool_result, **kwargs):
1059
+ self.history_events.append(("tool", tool_name, tool_result, kwargs))
1060
+
1061
+ def hist_add_message(self, ai, content, tokens=0, **kwargs):
1062
+ self.history_events.append(("message", ai, content, tokens, kwargs))
1063
+
1064
+ agent = HistoryAgent()
1065
+ job = parallel_tools.ParallelJob(
1066
+ id="vision-ready",
1067
+ parent_context_id=agent.context.id,
1068
+ index=0,
1069
+ tool_name="vision_load",
1070
+ tool_args={"paths": ["/image.png"]},
1071
+ kind="tool",
1072
+ state="success",
1073
+ result="Loaded images (1)",
1074
+ parent_history=[("raw-image-message", 1500)],
1075
+ )
1076
+ agent.context.set_data(parallel_tools.PARALLEL_JOBS_KEY, {job.id: job})
1077
+ tool = ParallelTool(
1078
+ agent, # type: ignore[arg-type]
1079
+ "parallel",
1080
+ None,
1081
+ {"action": "collect", "job_ids": [job.id]},
1082
+ "",
1083
+ None,
1084
+ )
1085
+
1086
+ response = await tool.execute(**tool.args)
1087
+
1088
+ assert job.id in agent.context.get_data(parallel_tools.PARALLEL_JOBS_KEY)
1089
+ assert agent.history_events == []
1090
+
1091
+ await tool.after_execution(response)
1092
+
1093
+ assert agent.history_events[0][0] == "tool"
1094
+ assert agent.history_events[1] == (
1095
+ "message",
1096
+ False,
1097
+ "raw-image-message",
1098
+ 1500,
1099
+ {},
1100
+ )
1101
+ assert job.id not in agent.context.get_data(parallel_tools.PARALLEL_JOBS_KEY)
1102
+
1103
+
1104
@pytest.mark.asyncio
1105
async def test_parallel_child_contexts_are_chats_not_tasks(monkeypatch) -> None:
1106
from agent import AgentContext
tests/test_vision_load_image_refs.py
+25
-4
@@ -108,7 +108,7 @@ async def test_vision_load_materializes_local_image_to_chat_artifact(monkeypatch
108
messages = []
109
updates = []
110
agent = SimpleNamespace(
111
- context=SimpleNamespace(id="ctx-vision"),
111
+ context=SimpleNamespace(id="ctx-vision", get_data=lambda *_args, **_kwargs: None),
112
agent_name="Agent 0",
113
hist_add_tool_result=lambda *args, **kwargs: tool_results.append((args, kwargs)),
114
hist_add_message=lambda *args, **kwargs: messages.append((args, kwargs)),
@@ -303,6 +303,12 @@ async def test_parallel_worker_consumes_parent_ephemeral_image(monkeypatch, tmp_
303
lambda _agent: {"vision": True, "max_embeds": 10},
304
)
305
monkeypatch.setattr(vision_load_module, "get_vision_model_config", lambda _agent: {})
306
+ queued = []
307
+ monkeypatch.setattr(
308
+ vision_load_module.parallel_tools,
309
+ "queue_parallel_parent_history",
310
+ lambda _agent, **message: queued.append(message) or True,
311
+ )
312
313
ref = vision_load_module.ephemeral_images.put_image_bytes(
314
context_id=parent_id,
@@ -313,10 +319,17 @@ async def test_parallel_worker_consumes_parent_ephemeral_image(monkeypatch, tmp_
319
context = SimpleNamespace(
320
id="parallel-worker",
321
get_data=lambda key: parent_id
316
- if key == vision_load_module.PARALLEL_WORKER_PARENT_CONTEXT_KEY
322
+ if key == vision_load_module.parallel_tools.PARALLEL_WORKER_PARENT_CONTEXT_KEY
323
else None,
324
)
319
- agent = SimpleNamespace(context=context, agent_name="Agent 0")
325
+ tool_results = []
326
+ local_messages = []
327
+ agent = SimpleNamespace(
328
+ context=context,
329
+ agent_name="Agent 0",
330
+ hist_add_tool_result=lambda *args, **kwargs: tool_results.append((args, kwargs)),
331
+ hist_add_message=lambda *args, **kwargs: local_messages.append((args, kwargs)),
332
+ )
333
tool = vision_load_module.VisionLoad(
334
agent=agent,
335
name="vision_load",
@@ -325,14 +338,22 @@ async def test_parallel_worker_consumes_parent_ephemeral_image(monkeypatch, tmp_
338
message="",
339
loop_data=None,
340
)
341
+ tool.log = SimpleNamespace(id="vision-log", update=lambda **kwargs: None)
342
329
- await tool.execute(paths=[ref])
343
+ response = await tool.execute(paths=[ref])
344
+ await tool.after_execution(response)
345
346
assert tool._context_id() == parent_id
347
assert tool.loaded_paths == ["shot.png"]
348
assert vision_load_module.ephemeral_images.get_image(ref, context_id=parent_id) is None
349
stored_ref = tool.images_dict["shot.png"]
350
assert stored_ref.startswith("/a0/usr/chats/parent-vision/images/vision-load/shot-")
351
+ assert local_messages == []
352
+ assert queued[0]["tokens"] == vision_load_module.TOKENS_ESTIMATE
353
+ raw_content = queued[0]["content"]["raw_content"]
354
+ assert raw_content == [
355
+ {"type": "image_url", "image_url": {"url": stored_ref}}
356
+ ]
357
358
359
@pytest.mark.anyio
tools/parallel.py
+10
-2
@@ -14,8 +14,14 @@ class ParallelTool(Tool):
14
text,
15
**(response.additional or {}),
16
)
17
+ await parallel_tools.collect_parallel_jobs(
18
+ self.agent,
19
+ getattr(self, "_collect_job_ids", []),
20
+ promote_parent_history=True,
21
+ )
22
23
async def execute(self, **kwargs) -> Response:
24
+ self._collect_job_ids = []
25
args = {**self.args, **kwargs}
26
action = str(args.get("action") or "").strip().lower()
27
@@ -68,9 +74,10 @@ class ParallelTool(Tool):
74
self.agent,
75
all_job_ids,
76
timeout=timeout,
71
- collect=True,
77
+ collect=False,
78
wait=False,
79
)
80
+ self._collect_job_ids = [result["job_id"] for result in results]
81
return Response(
82
message=parallel_tools.format_parallel_results(results),
83
break_loop=False,
@@ -80,9 +87,10 @@ class ParallelTool(Tool):
87
self.agent,
88
all_job_ids,
89
timeout=timeout,
83
- collect=True,
90
+ collect=False,
91
wait=True,
92
)
93
+ self._collect_job_ids = [result["job_id"] for result in results]
94
return Response(
95
message=parallel_tools.format_parallel_results(results),
96
break_loop=False,
tools/parallel.py.dox.md
+1
@@ -31,6 +31,7 @@
31
- Recursive use of `parallel` from inside a direct background tool worker is blocked before execution; numbered subordinate child chats can use normal child-chat tools, including `parallel`, to create their next-level descendants.
32
- Wrapped `call_subordinate` uses the same lifecycle as a top-level call. `job_id` identifies one parallel invocation, while its returned `context_id` identifies the reusable child agent for later `reset=false` calls.
33
- The wrapper tool does not create its own visible process-step log; each wrapped child call owns the visible log row, and the wrapper result is recorded only in model history.
34
+- Terminal direct jobs are collected after the wrapper result enters model history. Any explicitly queued parent-history messages are appended next, preserving result-before-content ordering for native multimodal tools.
35
36
## Key Concepts
37
tools/vision_load.py
+28
-10
@@ -2,8 +2,15 @@ from mimetypes import guess_type
2
3
from langchain_core.messages import HumanMessage
4
5
-from helpers import chat_media, ephemeral_images, files, history, images, runtime
6
-from helpers.parallel_tools import PARALLEL_WORKER_PARENT_CONTEXT_KEY
5
+from helpers import (
6
+ chat_media,
7
+ ephemeral_images,
8
+ files,
9
+ history,
10
+ images,
11
+ parallel_tools,
12
+ runtime,
13
+)
14
from helpers.tool import Response, Tool
15
from plugins._model_config.helpers.model_config import (
16
build_vision_model,
@@ -96,7 +103,11 @@ class VisionLoad(Tool):
103
def _context_id(self) -> str:
104
context = getattr(self.agent, "context", None)
105
get_data = getattr(context, "get_data", None)
99
- parent_id = get_data(PARALLEL_WORKER_PARENT_CONTEXT_KEY) if get_data else ""
106
+ parent_id = (
107
+ get_data(parallel_tools.PARALLEL_WORKER_PARENT_CONTEXT_KEY)
108
+ if get_data
109
+ else ""
110
+ )
111
return str(parent_id or getattr(context, "id", "") or "").strip()
112
113
async def _call_vision_model(self, image_paths: list[str]) -> str:
@@ -191,11 +202,18 @@ class VisionLoad(Tool):
202
{"type": "image_url", "image_url": {"url": image_path}}
203
for image_path in self.images_dict.values()
204
]
194
- self.agent.hist_add_message(
195
- False,
196
- content=history.RawMessage(
197
- raw_content=content,
198
- preview="<Image attachments loaded by path>",
199
- ),
200
- tokens=TOKENS_ESTIMATE * len(content),
205
+ raw_message = history.RawMessage(
206
+ raw_content=content,
207
+ preview="<Image attachments loaded by path>",
208
)
209
+ tokens = TOKENS_ESTIMATE * len(content)
210
+ if not parallel_tools.queue_parallel_parent_history(
211
+ self.agent,
212
+ content=raw_message,
213
+ tokens=tokens,
214
+ ):
215
+ self.agent.hist_add_message(
216
+ False,
217
+ content=raw_message,
218
+ tokens=tokens,
219
+ )
tools/vision_load.py.dox.md
+1
@@ -24,6 +24,7 @@
24
- The public tool contract is route-agnostic. A Vision Model receives the current user request through `fw.vision_load.md`; direct parallel workers inherit that request from their parent.
25
- Delegation completes during `execute(...)` so native Responses function output contains the real capsule before `after_execution(...)` persists it.
26
- Delegated history contains the text capsule only. Native history contains the tool result followed by one raw message holding all loaded image blocks.
27
+- In a direct parallel worker, native image content is queued for the parent and promoted immediately after the outer `parallel` result; the disposable worker never owns the only copy of model-visible pixels.
28
- Direct parallel workers inherit the parent's model override generically. This tool uses their recorded parent context only to resolve ephemeral refs and durable chat media.
29
- `max_embeds` comes from the model that actually receives the images.
30
- Vision Model calls use the selected model's Advanced `kwargs`; this tool does not impose a separate timeout or output-token limit.