diff --git a/agent/relay_llm.py b/agent/relay_llm.py index 85286c47dae7..3e3f54d4dfd7 100644 --- a/agent/relay_llm.py +++ b/agent/relay_llm.py @@ -344,6 +344,11 @@ class ManagedLlmStream(Iterator[Any]): self.output_modified = False callback_context = contextvars.copy_context() + def run_callback(callback: Callable[..., Any], *args: Any) -> Any: + # Relay can invoke stream surfaces while another callback still + # owns the captured Context. A fresh copy is safe to enter. + return callback_context.copy().run(callback, *args) + runtime, session, parent = relay_runtime.resolve_execution_context(session_id) if ( runtime is None @@ -378,7 +383,7 @@ class ManagedLlmStream(Iterator[Any]): async def provider_stream(next_request: Any): raw_stream = None try: - raw_stream = callback_context.run( + raw_stream = run_callback( stream_factory, _provider_request( request, @@ -390,7 +395,7 @@ class ManagedLlmStream(Iterator[Any]): ) if ( completed_response_predicate is not None - and callback_context.run( + and run_callback( completed_response_predicate, raw_stream, ) @@ -399,14 +404,14 @@ class ManagedLlmStream(Iterator[Any]): self._provider_completed = True return if on_stream_created is not None: - callback_context.run(on_stream_created, raw_stream) - raw_iterator = callback_context.run(iter, raw_stream) + run_callback(on_stream_created, raw_stream) + raw_iterator = run_callback(iter, raw_stream) while True: try: - chunk = callback_context.run(next, raw_iterator) + chunk = run_callback(next, raw_iterator) except StopIteration: break - if self._accept_chunk is not None and not callback_context.run( + if self._accept_chunk is not None and not run_callback( self._accept_chunk, chunk, ): @@ -421,17 +426,17 @@ class ManagedLlmStream(Iterator[Any]): finally: close = getattr(raw_stream, "close", None) if callable(close): - callback_context.run(close) + run_callback(close) def observe_chunk(chunk: Any) -> None: if self._on_chunk is not None: - callback_context.run(self._on_chunk, _jsonable(chunk)) + run_callback(self._on_chunk, _jsonable(chunk)) def relay_finalizer() -> Any: try: if self.final_response is not None: return _jsonable(self.final_response) - return _jsonable(callback_context.run(finalizer)) + return _jsonable(run_callback(finalizer)) except BaseException as exc: self._callback_error = exc raise diff --git a/tests/agent/test_relay_llm.py b/tests/agent/test_relay_llm.py index 0266d21bd2f2..bca9b6d1648d 100644 --- a/tests/agent/test_relay_llm.py +++ b/tests/agent/test_relay_llm.py @@ -5,6 +5,7 @@ from __future__ import annotations import asyncio import contextvars import json +import threading from types import SimpleNamespace import pytest @@ -436,6 +437,97 @@ def test_stream_provider_callbacks_preserve_caller_context(relay_turn): ] +def test_anthropic_stream_callbacks_do_not_reenter_captured_context( + relay_turn, + monkeypatch, +): + del relay_turn + caller_value = contextvars.ContextVar( + "anthropic_stream_caller_value", + default="default", + ) + caller_value.set("caller") + callback_context = contextvars.copy_context() + real_copy_context = contextvars.copy_context + copy_count = 0 + + def capture_callback_context(): + nonlocal copy_count + copy_count += 1 + if copy_count == 1: + return callback_context + return real_copy_context() + + monkeypatch.setattr( + relay_llm.contextvars, + "copy_context", + capture_callback_context, + ) + observed = [] + accumulator = relay_llm.AnthropicStreamAccumulator() + + def observe_chunk(chunk): + observed.append(caller_value.get()) + accumulator.observe(chunk) + + chunks = [ + { + "type": "message_start", + "message": { + "id": "message-1", + "type": "message", + "role": "assistant", + "model": "claude-test", + "usage": {"input_tokens": 1, "output_tokens": 0}, + }, + }, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 1}, + }, + ] + stream = relay_llm.stream( + { + "model": "claude-test", + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}], + }, + lambda _request: iter(chunks), + session_id="session-1", + name="anthropic", + model_name="claude-test", + finalizer=accumulator.finalize, + on_chunk=observe_chunk, + metadata={ + "api_mode": "anthropic_messages", + "api_request_id": "request-anthropic-context-reentry", + }, + ) + + entered = threading.Event() + release = threading.Event() + + def hold_callback_context() -> None: + def wait() -> None: + entered.set() + assert release.wait(timeout=5) + + callback_context.run(wait) + + holder = threading.Thread(target=hold_callback_context) + holder.start() + assert entered.wait(timeout=1) + try: + assert list(stream) == chunks + finally: + release.set() + holder.join(timeout=1) + + assert holder.is_alive() is False + assert observed == ["caller", "caller"] + + def test_non_stream_does_not_forward_relay_session_headers(relay_turn): _relay, _turn = relay_turn captured_requests = [] diff --git a/tests/run_agent/test_streaming.py b/tests/run_agent/test_streaming.py index 0f2fc27a5158..199d1ae355e1 100644 --- a/tests/run_agent/test_streaming.py +++ b/tests/run_agent/test_streaming.py @@ -1313,6 +1313,95 @@ class TestAnthropicStreamCallbacks: assert touch_calls.count("receiving stream response") == len(events) mock_stream.close.assert_called_once() + @pytest.mark.filterwarnings("ignore:Pydantic serializer warnings:UserWarning") + def test_anthropic_sdk_stream_runs_through_relay_managed_execution( + self, + tmp_path, + monkeypatch, + ): + anthropic = pytest.importorskip("anthropic") + httpx = pytest.importorskip("httpx") + pytest.importorskip("nemo_relay") + from agent import relay_runtime + from run_agent import AIAgent + + monkeypatch.setenv("HERMES_HOME", str(tmp_path / "hermes-home")) + response_body = b"""event: message_start +data: {"type":"message_start","message":{"id":"msg_test","type":"message","role":"assistant","content":[],"model":"claude-test","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":1,"output_tokens":0}}} + +event: content_block_start +data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}} + +event: content_block_delta +data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello"}} + +event: content_block_stop +data: {"type":"content_block_stop","index":0} + +event: message_delta +data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":1}} + +event: message_stop +data: {"type":"message_stop"} + +""" + + def respond(request): + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=response_body, + request=request, + ) + + client = anthropic.Anthropic( + api_key="test-key", + http_client=httpx.Client(transport=httpx.MockTransport(respond)), + ) + relay_runtime._reset_for_tests() + agent = AIAgent( + api_key="test-key", + base_url="https://api.anthropic.com", + provider="anthropic", + model="claude-test", + quiet_mode=True, + skip_context_files=True, + skip_memory=True, + ) + agent.api_mode = "anthropic_messages" + agent.session_id = "anthropic-relay-session" + agent._interrupt_requested = False + agent._create_request_anthropic_client = lambda *args, **kwargs: client + lease = relay_runtime.SESSION_COORDINATOR.acquire_conversation( + profile_key=relay_runtime.current_profile_key(), + session_id=agent.session_id, + platform="cli", + ) + turn = relay_runtime.SESSION_COORDINATOR.begin_turn( + lease, + turn_id="anthropic-relay-turn", + task_id="anthropic-relay-task", + ) + lease.host.retain_managed_execution("test.anthropic_relay") + + try: + result = agent._interruptible_streaming_api_call( + { + "model": "claude-test", + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}], + } + ) + finally: + lease.host.release_managed_execution("test.anthropic_relay") + relay_runtime.SESSION_COORDINATOR.end_turn(turn, outcome="success") + relay_runtime.SESSION_COORDINATOR.release_conversation(lease) + relay_runtime._reset_for_tests() + client.close() + + assert result.content[0].text == "hello" + assert result.stop_reason == "end_turn" + @patch("run_agent.AIAgent._rebuild_anthropic_client") @patch("run_agent.AIAgent._replace_primary_openai_client") def test_anthropic_stream_parser_valueerror_retries_before_delivery(