diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py index 54e1323f62..4391da0b70 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -329,8 +329,16 @@ def _is_hosted_responses_history_sentinel(provider: ContextProvider) -> bool: ) -def _reject_busy_conversation(session: AgentSession) -> None: +def _reject_busy_conversation(session: AgentSession, *, agent_server_history: bool) -> None: if session.state.get(_HOSTED_CONVERSATION_CLAIM_KEY) is not None: + if agent_server_history: + raise _PublicResponseFailure( + RuntimeError( + "The conversation already has an in-flight turn. Wait for it to finish before sending " + "another request; if that turn ended abnormally (for example, the host stopped or its " + "cleanup failed), start a new conversation." + ) + ) raise _PublicResponseFailure( RuntimeError( "A service-backed conversation already has an in-flight turn. Wait for it to finish; " @@ -340,6 +348,60 @@ def _reject_busy_conversation(session: AgentSession) -> None: ) +async def _release_conversation_claim( + store: SessionStore, + conversation_id: str, + response_id: str, + *, + persist_state: AgentSession | None = None, +) -> None: + """Remove a claim still owned by ``response_id`` so an uncommitted turn leaves the conversation usable. + + The head is re-read so the conditional write carries the ETag of the claim record. By default the claim + record's unchanged pre-turn state is restored. When ``persist_state`` is provided it is written instead, + so state the turn already consumed (such as an approval whose tool may have run) is not reopened. + A claim owned by another response is left untouched. + """ + head = await store.get(conversation_id) + if head is None or head.state.get(_HOSTED_CONVERSATION_CLAIM_KEY) != response_id: + return + if persist_state is None: + head.state.pop(_HOSTED_CONVERSATION_CLAIM_KEY, None) + await store.set(conversation_id, head) + else: + persist_state.state.pop(_HOSTED_CONVERSATION_CLAIM_KEY, None) + await store.set(conversation_id, persist_state) + + +async def _release_conversation_claim_best_effort( + store: SessionStore, + conversation_id: str, + response_id: str, + *, + persist_state: AgentSession | None, +) -> None: + """Release a claim without letting a cancellation or a storage failure skip or mask the cleanup. + + The release is shielded so it still completes when the surrounding task is cancelled, and the + cancellation itself is left to propagate. A failed release is logged; the claim then stays held, + which blocks later turns rather than risking a repeated approval. + """ + try: + await asyncio.shield( + _release_conversation_claim(store, conversation_id, response_id, persist_state=persist_state) + ) + except Exception as release_error: + _log_response_failure("conversation claim release", response_id, release_error) + + +def _request_carries_approval_response(request: CreateResponse) -> bool: + """Return whether the caller's input resumes a paused tool call with an approval response.""" + items = request.get("input") + return isinstance(items, list) and any( + isinstance(item, Mapping) and item.get("type") == "mcp_approval_response" for item in items + ) + + def _create_response_event_stream(context: ResponseContext) -> ResponseEventStream: """Create a response stream seeded from recovery state when available.""" if context.is_recovery: @@ -1296,7 +1358,7 @@ async def _handle_prepared_response( ) session = agent.create_session() if context.conversation_id is not None: - _reject_busy_conversation(session) + _reject_busy_conversation(session, agent_server_history=configuration.agent_server_history) if previous_response_id is not None and context.conversation_id is None: if session.service_session_id is not None and session.state.get( _HOSTED_SOURCE_CONVERSATION_KEY @@ -1483,9 +1545,11 @@ async def _handle_inner_agent( raise RuntimeError("A non-resumable agent cannot be replayed after a process crash.") stored = request.get("store") is not False + carries_approval_response = _request_carries_approval_response(request) request_messages_task: asyncio.Task[list[Message]] | None = None conversation_already_committed = False + releasable_claim_conversation_id: str | None = None try: if isinstance(agent, RawAgent): try: @@ -1618,7 +1682,7 @@ async def _handle_inner_agent( ) ) else: - _reject_busy_conversation(session) + _reject_busy_conversation(session, agent_server_history=configuration.agent_server_history) if not stored and self._history_source == "agent" and session.service_session_id is not None: raise _PublicResponseFailure( ValueError( @@ -1725,17 +1789,29 @@ async def _handle_inner_agent( session.state.pop(_HOSTED_PROVIDER_STATE_KEY, None) session.state.pop(_HOSTED_CONVERSATION_COMMITTED_KEY, None) + # With steering, AgentServer already serializes same-conversation turns and deliberately lets a newer + # turn supersede an in-flight one, so an agent_server-history turn must not be fenced by a claim. + steering_owns_admission = configuration.agent_server_history and bool( + self._host_options and self._host_options.steerable_conversations + ) if ( stored and context.conversation_id is not None - and not configuration.agent_server_history and not conversation_already_committed + and not steering_owns_admission ): if session_storage is None: - raise RuntimeError("Service history requires agent session storage.") - # The scoped store's conditional write must succeed before the provider can - # mutate its linear thread; a losing turn never reaches agent.run(). + raise RuntimeError("Conversation admission requires agent session storage.") + # The scoped store's conditional write must succeed before the agent can dispatch a tool or + # mutate a provider thread; a losing turn never reaches agent.run(). This also applies to + # agent_server history, where it is the only guard against two concurrent turns each consuming + # the same pending approval from stale session copies. session.state[_HOSTED_CONVERSATION_CLAIM_KEY] = context.response_id + if configuration.agent_server_history: + # No provider thread is mutated, so an uncommitted turn can hand the conversation back. + # Recorded before the write because a cancellation can land after the backend commits the + # claim but before the await returns; the ownership-checked release then still cleans up. + releasable_claim_conversation_id = context.conversation_id try: await _set_agent_session(session_storage, context.conversation_id, session) finally: @@ -1867,6 +1943,7 @@ async def _handle_inner_agent( session.state.pop(_HOSTED_PROVIDER_STATE_KEY, None) session.state[_HOSTED_CONVERSATION_COMMITTED_KEY] = context.response_id await _set_agent_session(session_storage, context.conversation_id, session) + releasable_claim_conversation_id = None except Exception as save_error: save_failure = save_error if request_interrupted: @@ -1876,6 +1953,16 @@ async def _handle_inner_agent( else: message = "Failed to persist the Agent Framework session after a successful request" _log_response_failure(message, context.response_id, save_error) + finally: + if releasable_claim_conversation_id is not None and session_storage is not None: + # An approval turn may already have run its tool, so keep the consumed approval rather than + # reopening it; any other turn restores the pre-turn state. + await _release_conversation_claim_best_effort( + session_storage, + releasable_claim_conversation_id, + context.response_id, + persist_state=session if carries_approval_response else None, + ) if request_failure is not None and save_failure is not None: raise RuntimeError( diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index aa9a5f0ce9..54fe4c67bb 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -103,6 +103,7 @@ _output_item_to_message, # pyright: ignore[reportPrivateUsage] _output_items_to_messages, # pyright: ignore[reportPrivateUsage] _OutputItemTracker, # pyright: ignore[reportPrivateUsage] + _release_conversation_claim, # pyright: ignore[reportPrivateUsage] _SignalledIterator, # pyright: ignore[reportPrivateUsage] _stringify_mcp_output, # pyright: ignore[reportPrivateUsage] consent_url_from_error, @@ -573,16 +574,39 @@ async def set(self, session_id: str, session: AgentSession) -> None: class _ConflictingConversationStore(SessionStore): + """Accept the pre-dispatch claim write but reject any later write to the conversation head.""" + def __init__(self) -> None: super().__init__() - self.fail_conversation = False + self.fail_conversation_commit = False async def set(self, session_id: str, session: AgentSession) -> None: - if self.fail_conversation and session_id == "conversation-head": + is_claim = "_foundry_conversation_claim" in session.state + if self.fail_conversation_commit and session_id == "conversation-head" and not is_claim: raise RuntimeError("Another request advanced this agent session; reload before writing.") await super().set(session_id, session) +class _InterruptingSessionStore(SessionStore): + """Interrupt a chosen write: block it until released, or cancel it after the backend has committed it.""" + + def __init__(self) -> None: + super().__init__() + self.block_key: str | None = None + self.cancel_after_claim_commit_key: str | None = None + self.blocked = asyncio.Event() + self.release = asyncio.Event() + + async def set(self, session_id: str, session: AgentSession) -> None: + if session_id == self.block_key: + self.blocked.set() + await self.release.wait() + await super().set(session_id, session) + if session_id == self.cancel_after_claim_commit_key and "_foundry_conversation_claim" in session.state: + self.cancel_after_claim_commit_key = None + raise asyncio.CancelledError + + _SESSION_STORE_UNSET = object() @@ -3136,10 +3160,318 @@ async def collect(response_context: ResponseContext, cancellation: asyncio.Event with suppress(asyncio.CancelledError): await pending + async def test_agent_server_conversation_claim_runs_replayed_approval_once(self) -> None: + conversation_id = f"conversation-{uuid.uuid4().hex}" + both_loaded = asyncio.Event() + release_run = asyncio.Event() + run_started = asyncio.Event() + reads = 0 + original_get = FoundryAgentSessionStore.get + + async def concurrent_get(store: FoundryAgentSessionStore, key: str) -> AgentSession | None: + nonlocal reads + session = await original_get(store, key) + if key == conversation_id and reads < 2: + reads += 1 + if reads == 2: + both_loaded.set() + await asyncio.wait_for(both_loaded.wait(), timeout=10) + return session + + dispatched: list[str] = [] + agent = _make_agent() + + def run(*args: Any, **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse]: + del args + session: AgentSession = kwargs["session"] + pending: list[str] = session.state.get("pending_approvals", []) + dispatched.extend(pending) + session.state["pending_approvals"] = [] + + async def updates() -> AsyncIterator[AgentResponseUpdate]: + if pending: + run_started.set() + await release_run.wait() + yield AgentResponseUpdate(contents=[Content.from_text("done")], role="assistant") + + return ResponseStream(updates(), finalizer=AgentResponse.from_updates) + + agent.run = MagicMock(side_effect=run) + server = _make_server(agent) + session_store = server._session_storage_provider.get_store( # pyright: ignore[reportPrivateUsage] + config=server.config, platform_context=get_request_context() + ) + pending_session = AgentSession() + pending_session.state["pending_approvals"] = ["approval-1"] + await session_store.set(conversation_id, pending_session) + + async def collect(response_id: str) -> list[Any]: + context = ResponseContext( + response_id=response_id, + conversation_id=conversation_id, + mode_flags=MagicMock(), + ) + return [ + event + async for event in server._handle_response( + CreateResponse(input="approve", store=True), context, asyncio.Event() + ) + ] + + with ( + patch.object(FoundryAgentSessionStore, "get", new=concurrent_get), + patch.object(ResponseContext, "get_input_items", new=AsyncMock(return_value=[])), + patch.object(ResponseContext, "get_history", new=AsyncMock(return_value=[])), + ): + pending = {asyncio.create_task(collect(f"response-{uuid.uuid4().hex}")) for _ in range(2)} + try: + await asyncio.wait_for(run_started.wait(), timeout=10) + done, pending = await asyncio.wait(pending, timeout=10, return_when=asyncio.FIRST_COMPLETED) + assert len(done) == len(pending) == 1 + assert "Another request advanced this agent session" in _failure_message(next(iter(done)).result()) + assert dispatched == ["approval-1"] + + blocked = await collect("blocked-while-turn-in-flight") + assert "in-flight turn" in _failure_message(blocked) + assert dispatched == ["approval-1"] + + release_run.set() + winner = await asyncio.wait_for(next(iter(pending)), timeout=10) + assert winner[-1]["type"] == "response.completed" + finally: + release_run.set() + for task in pending: + task.cancel() + await asyncio.gather(*pending, return_exceptions=True) + + head_store = server._session_storage_provider.get_store( # pyright: ignore[reportPrivateUsage] + config=server.config, platform_context=get_request_context() + ) + head = await head_store.get(conversation_id) + assert head is not None and "_foundry_conversation_claim" not in head.state + assert head.state["pending_approvals"] == [] + + replay = await collect("sequential-replay") + assert replay[-1]["type"] == "response.completed" + assert dispatched == ["approval-1"] + + async def test_failed_agent_server_conversation_turn_releases_claim_and_restores_pre_turn_state(self) -> None: + store = SessionStore() + conversation_id = "failed-agent-server-conversation" + pre_turn = AgentSession() + pre_turn.state["pending_approvals"] = ["approval-1"] + await store.set(conversation_id, pre_turn) + calls = 0 + agent = _make_agent() + + def run(*args: Any, **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse]: + nonlocal calls + del args + calls += 1 + session: AgentSession = kwargs["session"] + if calls == 1: + session.state["pending_approvals"] = [] + return ResponseStream(_raising_updates("model timed out"), finalizer=AgentResponse.from_updates) + return ResponseStream(_single_state_update(session), finalizer=AgentResponse.from_updates) + + agent.run = MagicMock(side_effect=run) + server = _make_server(agent, session_store=store) + + async def collect(response_id: str) -> list[Any]: + context = ResponseContext(response_id=response_id, conversation_id=conversation_id, mode_flags=MagicMock()) + return [ + event + async for event in server._handle_response( + CreateResponse(input="hello", store=True), context, asyncio.Event() + ) + ] + + with ( + patch.object(ResponseContext, "get_input_items", new=AsyncMock(return_value=[])), + patch.object(ResponseContext, "get_history", new=AsyncMock(return_value=[])), + ): + failed = await collect("failed-turn") + head_after_failure = await store.get(conversation_id) + retry = await collect("retry-turn") + + assert _failure_message(failed) == "An internal server error occurred." + assert head_after_failure is not None + assert "_foundry_conversation_claim" not in head_after_failure.state + assert head_after_failure.state["pending_approvals"] == ["approval-1"] + assert retry[-1]["type"] == "response.completed" + assert calls == 2 + + async def test_cancelled_agent_server_conversation_turn_releases_claim(self) -> None: + store = SessionStore() + started = asyncio.Event() + agent = _make_agent() + + async def blocked() -> AsyncIterator[AgentResponseUpdate]: + started.set() + await asyncio.Event().wait() + yield AgentResponseUpdate(contents=[Content.from_text("too late")], role="assistant") + + agent.run = MagicMock(side_effect=lambda **kwargs: ResponseStream(blocked())) + server = _make_server(agent, session_store=store) + conversation_id = "cancelled-agent-server-conversation" + context = ResponseContext(response_id="cancelled-turn", conversation_id=conversation_id, mode_flags=MagicMock()) + signal = asyncio.Event() + + async def collect() -> list[Any]: + return [ + event + async for event in server._handle_response(CreateResponse(input="hello", store=True), context, signal) + ] + + with ( + patch.object(ResponseContext, "get_input_items", new=AsyncMock(return_value=[])), + patch.object(ResponseContext, "get_history", new=AsyncMock(return_value=[])), + ): + pending = asyncio.create_task(collect()) + try: + await asyncio.wait_for(started.wait(), timeout=5) + held = await store.get(conversation_id) + assert held is not None and held.state["_foundry_conversation_claim"] == "cancelled-turn" + context.client_cancelled = True + signal.set() + await asyncio.wait_for(pending, timeout=5) + finally: + pending.cancel() + with suppress(asyncio.CancelledError): + await pending + + head = await store.get(conversation_id) + assert head is not None and "_foundry_conversation_claim" not in head.state + + async def test_failed_approval_turn_keeps_consumed_approval_and_does_not_repeat_the_tool(self) -> None: + store = SessionStore() + conversation_id = "failed-approval-conversation" + pre_turn = AgentSession() + pre_turn.state["pending_approvals"] = ["approval-1"] + await store.set(conversation_id, pre_turn) + dispatched: list[str] = [] + agent = _make_agent() + + def run(*args: Any, **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse]: + del args + session: AgentSession = kwargs["session"] + pending: list[str] = session.state.get("pending_approvals", []) + dispatched.extend(pending) + session.state["pending_approvals"] = [] + if pending: + # The approved tool ran; the following model call fails. + return ResponseStream(_raising_updates("model timed out"), finalizer=AgentResponse.from_updates) + return ResponseStream(_single_state_update(session), finalizer=AgentResponse.from_updates) + + agent.run = MagicMock(side_effect=run) + server = _make_server(agent, session_store=store) + approval: Any = [{"type": "mcp_approval_response", "approval_request_id": "approval-1", "approve": True}] + + async def collect(response_id: str) -> list[Any]: + context = ResponseContext(response_id=response_id, conversation_id=conversation_id, mode_flags=MagicMock()) + return [ + event + async for event in server._handle_response( + CreateResponse(input=approval, store=True), context, asyncio.Event() + ) + ] + + with ( + patch.object(ResponseContext, "get_input_items", new=AsyncMock(return_value=[])), + patch.object(ResponseContext, "get_history", new=AsyncMock(return_value=[])), + ): + failed = await collect("failed-approval-turn") + head_after_failure = await store.get(conversation_id) + retry = await collect("retry-approval-turn") + + assert _failure_message(failed) == "An internal server error occurred." + assert head_after_failure is not None + assert "_foundry_conversation_claim" not in head_after_failure.state + assert head_after_failure.state["pending_approvals"] == [] + assert retry[-1]["type"] == "response.completed" + assert dispatched == ["approval-1"] + + async def test_cancellation_during_response_snapshot_write_releases_claim(self) -> None: + store = _InterruptingSessionStore() + agent = _make_agent() + agent.run = MagicMock( + side_effect=lambda **kwargs: ResponseStream( + _single_state_update(kwargs["session"]), finalizer=AgentResponse.from_updates + ) + ) + server = _make_server(agent, session_store=store) + conversation_id = "cancelled-snapshot-conversation" + store.block_key = "snapshot-response" + context = ResponseContext( + response_id="snapshot-response", conversation_id=conversation_id, mode_flags=MagicMock() + ) + + async def collect() -> list[Any]: + return [ + event + async for event in server._handle_response( + CreateResponse(input="hello", store=True), context, asyncio.Event() + ) + ] + + with ( + patch.object(ResponseContext, "get_input_items", new=AsyncMock(return_value=[])), + patch.object(ResponseContext, "get_history", new=AsyncMock(return_value=[])), + ): + pending = asyncio.create_task(collect()) + try: + await asyncio.wait_for(store.blocked.wait(), timeout=5) + held = await store.get(conversation_id) + assert held is not None and held.state["_foundry_conversation_claim"] == "snapshot-response" + pending.cancel() + with suppress(asyncio.CancelledError): + await pending + finally: + store.release.set() + + head = await store.get(conversation_id) + assert head is not None and "_foundry_conversation_claim" not in head.state + + async def test_cancellation_after_claim_commit_releases_claim(self) -> None: + store = _InterruptingSessionStore() + conversation_id = "cancelled-claim-conversation" + store.cancel_after_claim_commit_key = conversation_id + agent = _make_agent() + agent.run = MagicMock() + server = _make_server(agent, session_store=store) + context = ResponseContext(response_id="claim-response", conversation_id=conversation_id, mode_flags=MagicMock()) + + async def collect() -> None: + async for _ in server._handle_response(CreateResponse(input="hello", store=True), context, asyncio.Event()): + pass + + with ( + patch.object(ResponseContext, "get_input_items", new=AsyncMock(return_value=[])), + patch.object(ResponseContext, "get_history", new=AsyncMock(return_value=[])), + ): + task = asyncio.create_task(collect()) + with suppress(asyncio.CancelledError): + await task + + agent.run.assert_not_called() + head = await store.get(conversation_id) + assert head is not None and "_foundry_conversation_claim" not in head.state + + async def test_release_conversation_claim_leaves_a_claim_owned_by_another_response(self) -> None: + store = SessionStore() + other = AgentSession() + other.state["_foundry_conversation_claim"] = "other-response" + await store.set("conversation", other) + + await _release_conversation_claim(store, "conversation", "this-response") + + head = await store.get("conversation") + assert head is not None and head.state["_foundry_conversation_claim"] == "other-response" + async def test_conversation_write_conflict_keeps_response_snapshot(self) -> None: store = _ConflictingConversationStore() await store.set("conversation-head", AgentSession()) - store.fail_conversation = True + store.fail_conversation_commit = True agent = _make_agent() agent.run = MagicMock( side_effect=lambda **kwargs: ResponseStream(