Make AgentPool worker callbacks spawn-safe and add pickle regression test

[mahimairaja]

Motivation

• The AgentPool previously registered a nested universal_session closure and used a bound self._prewarm method as worker callbacks, which captures self and is not pickleable under multiprocessing spawn, causing macOS dev-mode failures.
• Worker bootstrap must be serialization-safe so workers can be spawned without accidentally pickling the entire AgentPool instance.
• Add a regression test to prevent reintroducing this class of bug.

Description

• Introduce a small serializable state carrier @dataclass _PoolRuntimeState that holds the shared agents mapping.
• Add module-level spawn-safe entrypoints def _prewarm_worker(...) and async def _run_universal_session(...) that operate on _PoolRuntimeState instead of closing over AgentPool instances.
• Bind the entrypoints to the livekit hooks using functools.partial so AgentServer.setup_fnc and the RTC session handler reference only module-level callables plus the serializable runtime state.
• Move routing and dependency-loading helpers to module-level helpers (_resolve_agent_config, _load_shared_runtime_dependencies, and associated metadata helpers) and make AgentPool delegate to them, preserving external behavior while removing instance-bound callbacks.
• Remove the nested closure and the instance-bound _prewarm method; keep runtime behavior (prewarm, session creation, greeting generation) unchanged.
• Add a regression test test_worker_callbacks_are_pickleable_and_keep_registered_agents in tests/test_pool.py that pickles/unpickles the registered callbacks, executes the setup callback to populate proc.userdata, and runs the session callback against a fake context with a fake AgentSession to verify behavior.

Testing

• Ran pytest tests/test_pool.py tests/test_routing.py and all tests passed (25 passed).
• Ran ruff check src/openrtc/pool.py tests/test_pool.py tests/test_routing.py and linting passed (All checks passed).

---

Codex Task

Summary by CodeRabbit

• Refactor

  • Reorganized worker and session handling internals to centralize shared runtime state and plugin loading.

• Tests

  • Added tests ensuring worker callbacks remain serializable and session handlers behave correctly after serialization.

Note: Internal improvements only; no changes to user-facing functionality or APIs.

[coderabbitai]

Caution

Review failed

Pull request was closed or merged during review

📝 Walkthrough

Walkthrough

Extracts worker lifecycle and session handling from AgentPool into module-level functions using a shared _PoolRuntimeState dataclass; binds runtime state into worker callbacks via functools.partial; moves shared plugin loading to module scope; adds tests ensuring worker callbacks remain pickleable and preserve registered-agent behavior.

Changes

Cohort / File(s)	Summary
Core Worker Lifecycle Refactor src/openrtc/pool.py	Moved worker setup from AgentPool._prewarm to module-level _prewarm_worker(runtime_state, proc) that validates registration and populates proc.userdata with vad and turn_detection. Moved session handling from AgentPool._handle_session to async module-level _run_universal_session(runtime_state, ctx) which resolves agent, constructs and starts AgentSession using ctx.proc.userdata, connects context, and may generate greetings. Introduced _PoolRuntimeState dataclass and used functools.partial to bind runtime state into callbacks. Replaced instance routing helpers with module-level _resolve_agent_config, _agent_name_from_metadata, _agent_name_from_mapping, and _get_registered_agent; AgentPool._resolve_agent now delegates to _resolve_agent_config. Also lifted _load_shared_runtime_dependencies to module level.
Tests: Pickleability & Roundtrip tests/test_pool.py	Added test_worker_callbacks_are_pickleable_and_keep_registered_agents which captures and pickles the pool setup callback and session handler, unpickles and runs them with monkeypatched _load_shared_runtime_dependencies and a fake AgentSession, and asserts proc.userdata is populated and the fake session starts with the registered agent in a non-null room. Added imports (asyncio, pickle, SimpleNamespace, openrtc.pool as pool_module).

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~50 minutes

Possibly related PRs

• Introduce AgentPool implementation, examples, tests, and CI workflows #2 — Prior refactor touching src/openrtc/pool.py and worker/session setup; this change restructures and extracts callback/state binding introduced there.

Poem

🐰 I hopped through code to bind the state,

Partial stitched callbacks, never late,
Workers packed, then boxed to send,
Pickled friends rejoin and blend,
Agents wake in rooms at fate. 🥕

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 36.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately summarizes the main change: making AgentPool worker callbacks spawn-safe by refactoring them to use module-level functions and a serializable dataclass, plus adding a pickle regression test.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch feat/fix-spawn-safety-in-agentpool-callbacks

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

src/openrtc/pool.py (1)

1-1: ⚠️ Potential issue | 🟡 Minor

Fix formatting to pass the linting pipeline.

The CI pipeline indicates that ruff format --check failed for this file. Run the formatter to resolve.

#!/bin/bash
# Run ruff format to see what changes are needed
ruff format --diff src/openrtc/pool.py

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/openrtc/pool.py` at line 1, The file fails ruff formatting; run and apply
ruff format to src/openrtc/pool.py (the current top-level "from __future__
import annotations" import line indicates formatting changes are needed) so the
file conforms to ruff/PEP8 style; simply run `ruff format src/openrtc/pool.py`
(or `ruff format --diff` to preview), accept the changes, and commit the updated
formatted file.

🧹 Nitpick comments (2)

tests/test_pool.py (2)

166-174: Consider using the monkeypatch fixture for cleaner test setup.

Creating pytest.MonkeyPatch() instances manually works, but using the pytest fixture is more idiomatic and ensures automatic cleanup even if exceptions occur.

♻️ Refactor to use monkeypatch fixture

Change the test signature to accept the fixture:

def test_worker_callbacks_are_pickleable_and_keep_registered_agents(
    monkeypatch: pytest.MonkeyPatch,
) -> None:

Then replace manual MonkeyPatch() instantiation and try/finally blocks with direct monkeypatch.setattr() calls—pytest handles cleanup automatically.

Also applies to: 203-208

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/test_pool.py` around lines 166 - 174, Refactor the test to use pytest's
monkeypatch fixture instead of creating pytest.MonkeyPatch() manually: update
the test function signature (e.g.,
test_worker_callbacks_are_pickleable_and_keep_registered_agents) to accept
monkeypatch: pytest.MonkeyPatch, then replace the manual instantiation and
try/finally around
monkeypatch.setattr("openrtc.pool._load_shared_runtime_dependencies", lambda:
(FakeSilero, lambda: "turn")) with a direct monkeypatch.setattr(...) call so
that pytest performs automatic cleanup; apply the same change for the other
occurrences around lines referencing the same pattern (the block using
FakeSilero and the _load_shared_runtime_dependencies replacement).

---

188-194: Class-level instances list may cause test pollution across test runs.

FakeSession.instances is a class variable that persists across tests. If this test file is extended or run in certain pytest configurations, stale instances could affect assertions.

♻️ Clear instances list before use

     class FakeSession:
         instances: list[FakeSession] = []

         def __init__(self, **kwargs: object) -> None:
             self.kwargs = kwargs
             self.started = False
             FakeSession.instances.append(self)
+
+    FakeSession.instances.clear()

Add the clear call after the class definition and before ctx = FakeJobContext().

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/test_pool.py` around lines 188 - 194, The class-level list
FakeSession.instances can leak state between tests; reset it before creating the
test context by calling FakeSession.instances.clear() immediately after the
FakeSession class definition and before instantiating ctx = FakeJobContext(),
ensuring each test run starts with an empty instances list.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Outside diff comments:
In `@src/openrtc/pool.py`:
- Line 1: The file fails ruff formatting; run and apply ruff format to
src/openrtc/pool.py (the current top-level "from __future__ import annotations"
import line indicates formatting changes are needed) so the file conforms to
ruff/PEP8 style; simply run `ruff format src/openrtc/pool.py` (or `ruff format
--diff` to preview), accept the changes, and commit the updated formatted file.

---

Nitpick comments:
In `@tests/test_pool.py`:
- Around line 166-174: Refactor the test to use pytest's monkeypatch fixture
instead of creating pytest.MonkeyPatch() manually: update the test function
signature (e.g.,
test_worker_callbacks_are_pickleable_and_keep_registered_agents) to accept
monkeypatch: pytest.MonkeyPatch, then replace the manual instantiation and
try/finally around
monkeypatch.setattr("openrtc.pool._load_shared_runtime_dependencies", lambda:
(FakeSilero, lambda: "turn")) with a direct monkeypatch.setattr(...) call so
that pytest performs automatic cleanup; apply the same change for the other
occurrences around lines referencing the same pattern (the block using
FakeSilero and the _load_shared_runtime_dependencies replacement).
- Around line 188-194: The class-level list FakeSession.instances can leak state
between tests; reset it before creating the test context by calling
FakeSession.instances.clear() immediately after the FakeSession class definition
and before instantiating ctx = FakeJobContext(), ensuring each test run starts
with an empty instances list.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 953a812e-6950-480c-bd35-2ccd63f65c79

📥 Commits

Reviewing files that changed from the base of the PR and between 87895c3 and d84190a.

📒 Files selected for processing (2)

• src/openrtc/pool.py
• tests/test_pool.py