Skip to content

Commit d9ce0da

Browse files
Kowserclaude
andcommitted
Make tool workers spawn-safe: ToolWorkerEntry replaces the make_tool_worker closure
Group A of the idea-5 spawn-safety design. Previously make_tool_worker returned a nested tool_worker closure with a spoofed identity (__name__/__qualname__ reassigned, __module__ left as _dispatch), which could not be pickled at Process.start() under the default 'spawn' start method — the direct cause of the agent-e2e pickle failures ("attribute lookup <tool> on ..._dispatch failed"). - make_tool_worker now returns a picklable ToolWorkerEntry (module-level class instance): tool function travels as FunctionRef (module-level fns, incl. the @tool wrapper's __wrapped__ hop) or by value (picklable callable instances); guardrails, credential names and the framework- callable marker are carried as instance state because parent-populated _dispatch registries are empty in spawn children. - The tool execution body moves to module-level run_tool_task(); the identity spoof is deleted. Credential-name priority is resolved at registration; the workflow-level fallback stays a runtime lookup. - Skill workers: the make_script_func / make_read_func closures become picklable ScriptRunner / SkillFileReader classes (plain-data attrs). - All three registration paths (native @tool via ToolRegistry, skill workers, framework-extracted) now run probe_spawn_safety, failing fast at registration with a named offender instead of an opaque PicklingError at worker start. Probe group "tools" enabled. - test fixture: the langchain integration fixture's tool impl moves to module level (extracted <locals> functions are legitimately rejected by the probe under spawn). New test proves a full ToolWorkerEntry pickles, crosses a REAL spawn boundary, and executes a Task with zero parent registry state. tests/unit: 2510 passed, 1 skipped (default and explicit CONDUCTOR_MP_START_METHOD=spawn). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent b570b49 commit d9ce0da

8 files changed

Lines changed: 273 additions & 80 deletions

File tree

src/conductor/ai/agents/runtime/_dispatch.py

Lines changed: 72 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -267,29 +267,69 @@ def _needs_context(func):
267267

268268

269269
def make_tool_worker(tool_func, tool_name, guardrails=None, tool_def=None, credential_names=None):
270-
"""Create a Conductor worker wrapper for a @tool function.
271-
272-
The wrapper accepts a ``Task`` object so it can extract metadata
273-
(execution ID) for ``ToolContext`` injection, then maps the task's
274-
``inputParameters`` to the tool function's arguments.
275-
On failure the exception propagates so Conductor marks the task FAILED.
276-
277-
If *guardrails* are provided, they wrap the tool execution:
278-
- Pre-execution guardrails check the input parameters.
279-
- Post-execution guardrails check the tool result.
280-
281-
If *credential_names* are provided (e.g. for framework-extracted tools),
282-
they are captured in the closure and used as the primary source of
283-
credential names at invocation time.
270+
"""Create a spawn-safe Conductor worker for a @tool function.
271+
272+
Returns a picklable ``ToolWorkerEntry`` (module-level class instance) —
273+
NOT a closure. Under the ``spawn`` start method every worker is pickled
274+
at ``Process.start()``, so the entry carries everything the tool needs
275+
(function reference, guardrails, credential names) as instance state;
276+
parent-populated module registries are empty in spawn children.
277+
278+
Credential-name priority is resolved here, at registration:
279+
explicit *credential_names* (framework-extracted tools) > *tool_def*
280+
credentials > the function's ``_tool_def`` attribute. The workflow-level
281+
fallback stays a runtime lookup in :func:`run_tool_task`.
284282
"""
285-
# Capture credential names in closure for framework-extracted tools
286-
_closure_cred_names = list(credential_names) if credential_names else []
287-
# Store tool_def in module-level registry so it's accessible across
288-
# spawn-mode multiprocessing boundaries (closures are not picklable).
289283
if tool_def is not None:
284+
# Parent-side registry kept for in-process readers; spawn children
285+
# rely on the entry's carried state instead.
290286
_tool_def_registry[tool_name] = tool_def
287+
# Resolve PEP 563 string annotations eagerly (documented factory behavior;
288+
# run_tool_task repeats this in the child, where re-imported functions
289+
# start over with string annotations).
290+
import typing
291+
292+
try:
293+
tool_func.__annotations__ = typing.get_type_hints(tool_func)
294+
except Exception:
295+
pass
296+
creds = list(credential_names) if credential_names else []
297+
if not creds and tool_def is not None:
298+
creds = [c for c in (getattr(tool_def, "credentials", None) or []) if isinstance(c, str)]
299+
if not creds:
300+
creds = [c for c in _get_credential_names_from_tool(tool_func) if isinstance(c, str)]
301+
302+
from conductor.ai.agents.runtime._worker_entries import ToolWorkerEntry
303+
304+
# No probe here: this factory is also used for in-process execution (and
305+
# directly by tests). The spawn probe runs at the worker_task registration
306+
# sites, where crossing a process boundary becomes real.
307+
return ToolWorkerEntry.for_callable(
308+
tool_func, tool_name, guardrails=guardrails, credential_names=creds
309+
)
310+
311+
312+
def run_tool_task(task, *, tool_name, tool_func, guardrails=None, credential_names=None,
313+
framework_callable=False):
314+
"""Execute one tool task — the module-level body behind ``ToolWorkerEntry``.
315+
316+
Maps the task's ``inputParameters`` to the tool function's arguments,
317+
injects ``ToolContext``/credentials, and applies guardrails. On failure
318+
the returned ``TaskResult`` is marked FAILED (terminal for
319+
``TerminalToolError``). Runs in the worker child process; must not rely
320+
on parent-populated registries beyond best-effort fallbacks.
321+
"""
322+
if framework_callable:
323+
# Re-apply the parent-side marker: a spawn child's re-imported
324+
# function copy does not carry attributes set in the parent.
325+
try:
326+
tool_func._agentspan_framework_callable = True
327+
except Exception:
328+
pass
329+
_closure_cred_names = list(credential_names) if credential_names else []
291330
# Resolve PEP 563 string annotations (from __future__ import annotations)
292-
# to real types so downstream code can use isinstance().
331+
# to real types so downstream code can use isinstance(). Idempotent —
332+
# done per call because the resolution must happen in the child process.
293333
import typing
294334

295335
try:
@@ -406,7 +446,9 @@ def tool_worker(task: Task) -> TaskResult:
406446
if _closure_cred_names:
407447
credential_names = list(_closure_cred_names)
408448
else:
409-
_td = _tool_def_registry.get(tool_name) or tool_def
449+
# Best-effort parent-side registry (empty in spawn children —
450+
# the entry-carried credential_names above is the real path).
451+
_td = _tool_def_registry.get(tool_name)
410452
raw_secrets = (
411453
list(getattr(_td, "credentials", []))
412454
if _td
@@ -442,7 +484,11 @@ def tool_worker(task: Task) -> TaskResult:
442484
continue
443485
if param_name in task.input_data:
444486
raw_value = task.input_data[param_name]
445-
ann = tool_func.__annotations__.get(param_name, inspect.Parameter.empty)
487+
# getattr: callable-instance workers (e.g. skill entries)
488+
# have no __annotations__ of their own.
489+
ann = getattr(tool_func, "__annotations__", {}).get(
490+
param_name, inspect.Parameter.empty
491+
)
446492
fn_kwargs[param_name] = _coerce_value(raw_value, ann)
447493
elif sig.parameters[param_name].default is not inspect.Parameter.empty:
448494
fn_kwargs[param_name] = sig.parameters[param_name].default
@@ -503,10 +549,11 @@ def _invoke_with_context():
503549
task_result.reason_for_incompletion = str(e)
504550
return task_result
505551

506-
tool_worker.__name__ = tool_func.__name__
507-
tool_worker.__qualname__ = tool_func.__qualname__
508-
tool_worker.__doc__ = tool_func.__doc__
509-
return tool_worker
552+
# The nested helpers above exist only for the duration of this call — the
553+
# picklable unit is ToolWorkerEntry; nothing here crosses a process
554+
# boundary (the identity spoof that used to live here broke pickling:
555+
# idea-5 spawn-vs-fork analysis).
556+
return tool_worker(task)
510557

511558

512559
# ── Native function calling workers ─────────────────────────────────────

src/conductor/ai/agents/runtime/_worker_entries.py

Lines changed: 73 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -144,12 +144,84 @@ def resolve(self) -> Callable:
144144
return obj
145145

146146

147+
# ── Worker entries ───────────────────────────────────────────────────────
148+
149+
150+
class ToolWorkerEntry:
151+
"""Spawn-safe replacement for ``make_tool_worker``'s nested ``tool_worker``.
152+
153+
Pickles by value (module-level class + picklable attrs) — no closure, no
154+
identity spoof. The tool callable travels either as a :class:`FunctionRef`
155+
(module-level functions) or directly (``fn_direct``) for picklable
156+
callable instances and — under ``fork``, where nothing is pickled — legacy
157+
closures. Credential names are resolved at registration and carried here,
158+
because the parent's ``_dispatch`` registries are empty in spawn children.
159+
"""
160+
161+
def __init__(self, tool_name, fn_ref=None, fn_direct=None, guardrails=None,
162+
credential_names=None, framework_callable=False):
163+
if (fn_ref is None) == (fn_direct is None):
164+
raise ValueError("exactly one of fn_ref/fn_direct is required")
165+
self.tool_name = tool_name
166+
self.fn_ref = fn_ref
167+
self.fn_direct = fn_direct
168+
self.guardrails = guardrails
169+
self.credential_names = list(credential_names) if credential_names else []
170+
# Carried explicitly: the parent sets _agentspan_framework_callable on
171+
# the function object, which a spawn child's re-imported copy lacks.
172+
self.framework_callable = framework_callable
173+
174+
@classmethod
175+
def for_callable(cls, fn, tool_name, guardrails=None, credential_names=None):
176+
"""Build an entry for *fn*, preferring by-reference transport.
177+
178+
Falls back to direct transport for callable instances (picklable by
179+
value) and for closures — the latter only work under ``fork``; the
180+
registration probe polices that under ``spawn``.
181+
"""
182+
framework = bool(getattr(fn, "_agentspan_framework_callable", False))
183+
try:
184+
ref = FunctionRef.of(fn)
185+
except SpawnSafetyError:
186+
entry = cls(tool_name, fn_direct=fn, guardrails=guardrails,
187+
credential_names=credential_names, framework_callable=framework)
188+
else:
189+
entry = cls(tool_name, fn_ref=ref, guardrails=guardrails,
190+
credential_names=credential_names, framework_callable=framework)
191+
# Introspection compatibility (logging etc.). Plain string INSTANCE
192+
# attrs — unlike the old wrapper's reassigned function identity, these
193+
# don't participate in pickling-by-reference, so they can't break it.
194+
entry.__name__ = getattr(fn, "__name__", tool_name)
195+
entry.__qualname__ = getattr(fn, "__qualname__", tool_name)
196+
entry.__doc__ = getattr(fn, "__doc__", None)
197+
return entry
198+
199+
def _target(self) -> Callable:
200+
return self.fn_ref.resolve() if self.fn_ref is not None else self.fn_direct
201+
202+
def __call__(self, task):
203+
# Late import: _dispatch owns the execution helpers; importing it here
204+
# (not at module top) avoids an import cycle and works in the child.
205+
from conductor.ai.agents.runtime._dispatch import run_tool_task
206+
207+
return run_tool_task(
208+
task,
209+
tool_name=self.tool_name,
210+
tool_func=self._target(),
211+
guardrails=self.guardrails,
212+
credential_names=self.credential_names,
213+
framework_callable=self.framework_callable,
214+
)
215+
216+
147217
# ── Registration-time spawn probe ────────────────────────────────────────
148218
#
149219
# Groups are enabled stage-by-stage as each worker family is converted to a
150220
# spawn-safe form (idea-5 implementation plan): probing an unconverted group
151221
# would fail every registration immediately.
152-
_ENABLED_PROBE_GROUPS: FrozenSet[str] = frozenset()
222+
# - "tools": make_tool_worker/ToolWorkerEntry (native @tool, skill workers,
223+
# framework-extracted) — converted in Stage 2.
224+
_ENABLED_PROBE_GROUPS: FrozenSet[str] = frozenset({"tools"})
153225

154226

155227
def _spawn_probe_active(group: str) -> bool:

src/conductor/ai/agents/runtime/runtime.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -872,11 +872,13 @@ def prepare(self, agent: Any) -> None:
872872
# TaskHandler is created — avoids incremental fork() on macOS.
873873
from conductor.ai.agents.frameworks.serializer import serialize_agent
874874
from conductor.ai.agents.runtime._dispatch import make_tool_worker
875+
from conductor.ai.agents.runtime._worker_entries import probe_spawn_safety
875876
from conductor.client.worker.worker_task import worker_task
876877

877878
_, workers = serialize_agent(agent)
878879
for w in workers:
879880
wrapper = make_tool_worker(w.func, w.name)
881+
probe_spawn_safety(wrapper, w.name, group="tools")
880882
worker_task(
881883
task_definition_name=w.name,
882884
task_def=_default_task_def(w.name),
@@ -1335,6 +1337,7 @@ def _register_and_start_skill_workers(
13351337
def _register_skill_workers(self, agent: Agent, domain: "Optional[str]" = None) -> None:
13361338
"""Register skill workers (scripts + read_skill_file) for a skill-based agent."""
13371339
from conductor.ai.agents.runtime._dispatch import make_tool_worker
1340+
from conductor.ai.agents.runtime._worker_entries import probe_spawn_safety
13381341
from conductor.ai.agents.skill import create_skill_workers
13391342
from conductor.client.worker.worker_task import worker_task
13401343

@@ -1344,6 +1347,7 @@ def _register_skill_workers(self, agent: Agent, domain: "Optional[str]" = None)
13441347

13451348
for sw in skill_workers:
13461349
wrapper = make_tool_worker(sw.func, sw.name)
1350+
probe_spawn_safety(wrapper, sw.name, group="tools")
13471351
worker_task(
13481352
task_definition_name=sw.name,
13491353
task_def=_default_task_def(sw.name),
@@ -3240,6 +3244,7 @@ def _register_framework_workers(
32403244
return
32413245

32423246
from conductor.ai.agents.runtime._dispatch import make_tool_worker
3247+
from conductor.ai.agents.runtime._worker_entries import probe_spawn_safety
32433248
from conductor.client.worker.worker_task import worker_task
32443249

32453250
for w in workers:
@@ -3248,6 +3253,7 @@ def _register_framework_workers(
32483253
except Exception:
32493254
pass
32503255
wrapper = make_tool_worker(w.func, w.name, credential_names=credentials)
3256+
probe_spawn_safety(wrapper, w.name, group="tools")
32513257
worker_task(
32523258
task_definition_name=w.name,
32533259
task_def=_default_task_def(w.name),

src/conductor/ai/agents/runtime/tool_registry.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,10 +67,15 @@ def register_tool_workers(
6767
_tool_approval_flags[td.name] = True
6868
logger.info("Tool '%s' registered with approval_required=True", td.name)
6969

70+
from conductor.ai.agents.runtime._worker_entries import probe_spawn_safety
71+
7072
for td in tool_defs:
7173
if td.func is not None and td.tool_type in ("worker", "cli"):
7274
guardrails = td.guardrails if td.guardrails else None
7375
wrapper = make_tool_worker(td.func, td.name, guardrails=guardrails, tool_def=td)
76+
# Fail fast (with the offender named) if this worker can't
77+
# cross the spawn process boundary — see idea-5 spawn safety.
78+
probe_spawn_safety(wrapper, td.name, group="tools")
7479
worker_task(
7580
task_definition_name=td.name,
7681
task_def=_default_task_def(

0 commit comments

Comments
 (0)