fix: run sync sandbox tools off the event loop so traced tools work in agent code mode (#1523)

## Problem

With `Agent(code_mode=True)`, live runs showed **every** tool call
failing with:

```
Deadlock detected: blocking call used in syncify thread flyte_syncify when calling
function <function _fetch_action_outputs ...>, use .aio() if in an async call.
```

Root cause: when the agent is entered through a syncified call (e.g. a
sync `agent.run()`), the whole Monty sandbox loop runs on the
`flyte_syncify` background loop. `ExternalFunctionBridge.execute_monty`
invoked external functions **inline**, so a sync `@flyte.trace` tool's
wrapper made its blocking syncify call (`_fetch_action_outputs`) from
the syncify thread itself — a guaranteed deadlock that the SDK detects
and aborts. The JSON tool-calling path was unaffected because
`_make_callable_tool` already dispatches sync tools via
`asyncio.to_thread`.

Downstream workaround that motivated this fix: unionai/unionai-agents#86
(had to disable code mode entirely).

## Fix

`ExternalFunctionBridge` now routes every non-coroutine-function
external through `asyncio.to_thread` (new `_call_external` helper, used
by both `execute_monty` and `_map_callable`):

- `asyncio.to_thread` copies contextvars, so the Flyte task context —
and therefore tracing — works exactly as it does on the JSON tool path.
- The blocking syncify calls inside sync trace wrappers now happen on a
worker thread while the syncify loop stays free to serve them.
- Side benefit: slow sync tools no longer stall the event loop driving
the sandbox (previously they blocked it outright), including under
`flyte_map`.
- Coroutine functions (`TaskTemplate.aio`, async tools, call-handler
wrappers) keep their inline path, and the existing unwrap loop still
drains coroutines returned by non-coroutine callables (local-mode
`TaskTemplate.aio()`).

## Also in this PR: remove a flaky Python 3.14 test

CI on this PR surfaced a pre-existing flaky failure, unrelated to the
sandbox change:
`test_task_call_sequence.py::test_interleaved_allocations_prevent_id_reuse`
asserted that CPython never hands a freed override object's address to a
later allocation when a same-size "noise" object is allocated in
between. That's unspecified allocator behavior, and the 3.14 runner
reused an address once (`seqs [1,1,1,1,1,1,1,2,1,1]`).

The `id()`-keyed sequencing that test demonstrates is also historical —
the remote controller has keyed sequences on a call key (task identity +
inputs hash + group) via `TaskCallSequencer` since #607, with
deterministic coverage in `test_action_name_stability_gaps.py`. The test
is removed and the file's docstring now notes that the remaining
demonstrations (which hold every object alive, so ids are guaranteed
distinct) are deterministic and describe pre-#607 behavior.

## Verification

- New regression tests drive `orchestrate_local` on the syncify loop
with a sync tool that makes a blocking syncify call (direct call and via
`flyte_map`) — they reproduce the exact production error before the fix
and pass after.
- Repro script confirmed: unfixed bridge → `Deadlock detected: blocking
call used in syncify thread flyte_syncify ...`; fixed bridge → returns
the tool result.
- `tests/flyte/sandbox/` + `tests/flyte/ai/`: 694 passed; sequencing
suites (`test_task_call_sequence.py`,
`test_action_name_stability_gaps.py`): 22 passed. Lint/mypy/ty clean.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

https://claude.ai/code/session_01TqG84zzd1AFjBXPrMnWDxh

---------

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

Author: cosmicBboy

src/flyte/sandbox/_bridge.py

@@ -29,6 +29,30 @@ def _to_monty(value: Any) -> Any:
     return value
 
 
+async def _call_external(fn: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
+    """Call an external ref from the bridge's event loop, resolving to a concrete value.
+
+    Coroutine functions are called inline and awaited. Everything else runs in a
+    worker thread via `asyncio.to_thread` (which propagates contextvars, so the
+    Flyte task context survives): a sync external may block — in particular a sync
+    `@flyte.trace` wrapper makes blocking syncify calls, and when this loop *is*
+    the syncify loop (e.g. agent code mode entered through a syncified call), an
+    inline call would trip syncify's deadlock detection and kill the tool call.
+    Offloading also keeps slow sync tools from stalling the loop.
+
+    The trailing loop unwraps coroutines returned by non-coroutine-function
+    callables — e.g. `TaskTemplate.aio()` in local mode may return an unawaited
+    coroutine from `forward()`.
+    """
+    if inspect.iscoroutinefunction(fn):
+        result = fn(*args, **kwargs)
+    else:
+        result = await asyncio.to_thread(fn, *args, **kwargs)
+    while inspect.iscoroutine(result):
+        result = await result
+    return result
+
+
 def _from_monty(value: Any) -> Any:
     """Unmarshal a tagged dict back to a flyte.io type."""
     if isinstance(value, dict) and _IO_TYPE_KEY in value:
@@ -140,10 +164,7 @@ async def _map_callable(
 
         async def run_row(row: tuple[Any, ...]) -> Any:
             try:
-                result = fn(*row)
-                while inspect.iscoroutine(result):
-                    result = await result
-                return result
+                return await _call_external(fn, *row)
             except Exception as exc:
                 if return_exceptions:
                     return exc
@@ -202,12 +223,7 @@ async def execute_monty(self, monty_cls: Any, code: str, input_names: list[str],
                 args = [_from_monty(a) for a in progress.args]
                 kwargs = {k: _from_monty(v) for k, v in progress.kwargs.items()}
 
-                # Call the external function and await if async.
-                # Loop because TaskTemplate.aio() in local mode may return
-                # an unawaited coroutine from forward() for async functions.
-                result = fn(*args, **kwargs)
-                while inspect.iscoroutine(result):
-                    result = await result
+                result = await _call_external(fn, *args, **kwargs)
 
                 progress = progress.resume({"return_value": _to_monty(result)})
             else:

tests/flyte/internal/runtime/test_task_call_sequence.py

@@ -10,6 +10,14 @@
 Combined with a stale completion event in ActionCache (remove() does not clean
 up _completion_events), duplicate action names cause
 "Task X did not return an output path, but the task has outputs defined."
+
+Note: the id()-based logic replicated below is the *historical* behavior these
+tests were written against; the remote controller now keys sequences on a call
+key (task identity + inputs hash + group) via TaskCallSequencer, covered in
+test_action_name_stability_gaps.py. The demonstrations here deliberately hold
+every override object alive so their id() values are guaranteed distinct —
+assertions about *freed* addresses would depend on unspecified CPython
+allocator behavior (and did flake on Python 3.14).
 """
 
 import hashlib
@@ -135,28 +143,6 @@ def test_sequence_increments_for_same_object_identity(self):
 
         assert seqs == list(range(1, 11))
 
-    def test_interleaved_allocations_prevent_id_reuse(self):
-        """
-        Simulate what happens in real async code: allocations between loop
-        iterations (logging, protobuf serialization, other coroutines) prevent
-        CPython from reusing the same memory address for the next override
-        object.
-        """
-        sequencer: dict[int, int] = defaultdict(int)
-        noise = []  # hold references to prevent address reuse
-
-        seqs = []
-        for i in range(10):
-            override = parse_entity_extraction.override(short_name="same_name")
-            seq = generate_task_call_sequence(override, sequencer)
-            seqs.append(seq)
-            del override
-            # Allocate an object of the same type/size to grab the freed address
-            noise.append(parse_entity_extraction.override(short_name=f"noise-{i}"))
-
-        # With interleaved allocations, id() is not reused -> seq stays at 1
-        assert all(s == 1 for s in seqs), f"Expected all sequences to be 1 (id not reused), got {seqs}"
-
     def test_action_name_determinism(self):
         """
         Verify that ActionID.new_sub_action_from produces identical names when

tests/flyte/sandbox/test_bridge.py

@@ -1,6 +1,13 @@
 """Tests for the ExternalFunctionBridge."""
 
-from flyte.sandbox._bridge import ExternalFunctionBridge
+import asyncio
+import threading
+
+import pytest
+
+import flyte.sandbox
+from flyte.sandbox._bridge import ExternalFunctionBridge, _call_external
+from flyte.syncify import syncify
 
 
 class TestExternalFunctionBridge:
@@ -28,3 +35,106 @@ def test_init_empty_refs(self):
             durable_refs={},
         )
         assert bridge._all_refs == {}
+
+
+class TestCallExternal:
+    """`_call_external` dispatches sync work off the loop thread and awaits async work inline."""
+
+    @pytest.mark.asyncio
+    async def test_sync_fn_runs_off_loop_thread(self):
+        seen: dict[str, threading.Thread] = {}
+
+        def sync_fn(x: int) -> int:
+            seen["thread"] = threading.current_thread()
+            return x + 1
+
+        assert await _call_external(sync_fn, 1) == 2
+        assert seen["thread"] is not threading.current_thread()
+
+    @pytest.mark.asyncio
+    async def test_async_fn_runs_on_loop(self):
+        seen: dict[str, threading.Thread] = {}
+
+        async def async_fn(x: int) -> int:
+            seen["thread"] = threading.current_thread()
+            return x + 1
+
+        assert await _call_external(async_fn, 1) == 2
+        assert seen["thread"] is threading.current_thread()
+
+    @pytest.mark.asyncio
+    async def test_returned_coroutine_is_awaited(self):
+        # TaskTemplate.aio() in local mode may hand back an unawaited coroutine
+        # from forward(); _call_external must drain it to a concrete value.
+        async def inner() -> str:
+            return "done"
+
+        def sync_returning_coro():
+            return inner()
+
+        assert await _call_external(sync_returning_coro) == "done"
+
+
+# --- Regression: agent code mode with sync @flyte.trace tools -----------------
+#
+# With `code_mode=True` the whole sandbox loop can end up running on the
+# `flyte_syncify` background loop (any syncified entry point). Sync
+# `@flyte.trace` wrappers make *blocking* syncify calls
+# (`_fetch_action_outputs` / `_record_trace_action`); if the bridge invokes
+# them inline on that same thread, syncify's deadlock detection aborts every
+# tool call. The bridge must run sync externals in a worker thread instead.
+
+
+@syncify
+async def _syncified_helper() -> str:
+    return "ok"
+
+
+def blocking_tool(x: int) -> int:
+    """A sync tool that blocks on syncify, exactly like a sync @flyte.trace wrapper."""
+    assert _syncified_helper() == "ok"
+    return x + 1
+
+
+@syncify
+async def _orchestrate_on_syncify_loop(code: str, tools: list, inputs: dict):
+    return await flyte.sandbox.orchestrate_local(code, inputs=inputs, tasks=tools)
+
+
+class TestSyncToolOnSyncifyLoop:
+    def test_blocking_sync_tool_does_not_deadlock(self):
+        # Before the fix this raised: "Deadlock detected: blocking call used in
+        # syncify thread flyte_syncify ... use .aio() if in an async call."
+        result = _orchestrate_on_syncify_loop("blocking_tool(x)", [blocking_tool], {"x": 1})
+        assert result == 2
+
+    def test_flyte_map_over_blocking_sync_tool_does_not_deadlock(self):
+        result = _orchestrate_on_syncify_loop(
+            "flyte_map('blocking_tool', xs)",
+            [blocking_tool],
+            {"xs": [1, 2, 3]},
+        )
+        assert result == [2, 3, 4]
+
+    def test_event_loop_stays_responsive_while_sync_tool_blocks(self):
+        # While a sync tool sleeps in its worker thread, the loop driving the
+        # bridge must still be able to run other coroutines.
+        import time
+
+        async def heartbeat() -> int:
+            ticks = 0
+            for _ in range(5):
+                await asyncio.sleep(0.01)
+                ticks += 1
+            return ticks
+
+        def parked(x: int) -> int:
+            time.sleep(0.05)
+            return x
+
+        async def drive() -> tuple[int, int]:
+            return await asyncio.gather(_call_external(parked, 7), heartbeat())
+
+        tool_result, ticks = asyncio.run(drive())
+        assert tool_result == 7
+        assert ticks == 5