fix: Accept Stream instances in as_tool() and subagent_tool() stream parameter

[Oxygen56]

What

Fixes #3021 — MemoryStream passed to as_tool() does not capture sub-agent events.

Root Cause

Agent.as_tool(stream=...) only accepted StreamFactory (a callable), not a Stream instance. When users naturally passed MemoryStream(), the code tried to call it as a function, the sub-agent silently fell back to creating an ephemeral stream in run_task(), and the user's stream stayed empty.

Fix

1. agent.py: Changed stream parameter type from StreamFactory | None to Stream | StreamFactory | None in as_tool()
2. subagent_tool.py: Added auto-wrapping logic — if a Stream instance is passed, it's wrapped in a lambda factory. Raises TypeError for invalid types.
3. background.py: Same fix for background_agent_tool() which had the same pattern
4. Tests: 4 new tests in test/beta/tools/test_task.py::TestAsToolStream covering:
   • Stream instance (the reported bug)
   • StreamFactory callable (backward compat)
   • None (default behavior)
   • Invalid type (TypeError)

Verification

$ pytest test/beta/tools/test_task.py -xvs
37 passed  # all existing tests + 4 new tests

Backward Compatibility

✅ Fully backward compatible — StreamFactory callables and None continue to work.

[Lancetnik]

Hi @Oxygen56 — this PR now has merge conflicts with main and has been moved to status:changes-requested.

Please rebase onto the current main and resolve the conflicts. Once you push the updated branch, the PR will return to the review queue automatically.

⏳ If there's no update within 30 days, this PR will be marked stale and closed 7 days after that.

[genisis0x]

This is a clean fix for #3021 — widening stream to Stream | StreamFactory and auto-wrapping an instance in a factory is the right shape, and I like that you added the TypeError guard, the docstrings, and test coverage. Confirmed the discriminator is safe: the Stream protocol has no __call__ and MemoryStream doesn't define one, so not callable(stream) correctly separates an instance from a factory, and the isinstance(stream, Stream) check backs it up.

Two things worth thinking about before merge:

1. Shared-instance semantics under background_agent_tool. The auto-wrap returns the same stream object on every factory call (lambda _agent, _ctx: _captured_stream). For subagent_tool/as_tool that's fine and is exactly what #3021 wants — the user's stream captures the sub-agent's events. But background_agent_tool is fire-and-forget, so a parent can kick off several background sub-agents that now all share one MemoryStream. MemoryStream carries a per-instance _ag2_turn_lock, so those concurrent background runs will serialize on that lock and interleave their turns into the one event log. That's arguably the desired "observe everything in one place" behavior, but it's a meaningful semantic change for the background case (independent tasks now contend on a shared lock and share history) and I'd suggest calling it out in the docstring, plus a test that spins up two concurrent background runs against one shared stream to pin down the interleaving contract.

2. Minor DRY: the normalize block (if stream is not None and not callable(stream): ...) is now duplicated verbatim in background.py and subagent_tool.py. Since as_tool already delegates to subagent_tool, pulling this into a small _as_stream_factory(stream) -> StreamFactory | None helper (raising the same TypeError) would keep the two call sites from drifting.

Neither is a correctness blocker for the non-background path — the core fix looks right.

[Oxygen56]

Thanks for the review and for flagging the merge conflict. I rebased the branch onto the current main and resolved the package-path migration conflict.

I also addressed both follow-up suggestions:

• centralized stream normalization in _as_stream_factory so the background and foreground paths cannot drift;
• documented the shared-instance behavior for concurrent background calls and added a concurrent test confirming both calls use the same stream/history.

After the package move, the affected paths are now under ag2/... and test/tools/....

Verification:

• uv run --group test-core pytest test/tools/test_task.py test/tools/test_background.py -q → 47 passed
• Ruff format/check on the changed modules and tests → passed
• Mypy on the changed subagent modules → passed

The PR is now mergeable against main. The repository workflows currently show action_required without jobs, so they appear to be awaiting maintainer approval to run. Thanks!

[genisis0x]

Following up on my earlier review now that _as_stream_factory is in — thanks for pulling the normalize logic into one helper. One thing on its ordering: it checks callable(stream) before isinstance(stream, Stream), which is correct today only because no Stream type defines __call__. If a Stream subclass ever becomes callable (a stream that doubles as its own factory isn't far-fetched), a real instance would take the callable branch and get invoked as stream(agent, ctx) — the exact failure this PR fixes. Checking isinstance(stream, Stream) first, then callable, then None would be robust regardless. Minor and non-blocking; the rest still looks good.