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 6ff4550f7e..5549465eb7 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -326,8 +326,9 @@ async def _handle_inner_agent( ) -> AsyncIterable[ResponseStreamEvent | dict[str, Any]]: """Handle a regular agent with Responses-managed MAF session continuity. - Conversation mode reads and writes one MAF session snapshot under - ``conversation_id``. Response chaining reads the snapshot under + Conversation mode reads the latest MAF session snapshot under + ``conversation_id`` and writes each turn under both its immutable + ``response_id`` and the conversation ID. Response chaining reads the snapshot under ``previous_response_id`` and writes the updated session under the current ``response_id``, allowing branches without changing the MAF session's own identifier. Hosted storage uses the request user as its isolation boundary. @@ -372,6 +373,8 @@ async def _handle_inner_agent( try: approval_storage = self._function_approval_storage_provider.get_store(config=self.config) session_storage = self._session_storage_provider.get_store(config=self.config) + if request.get("previous_response_id") is not None and context.conversation_id is not None: + raise RuntimeError("Previous response ID cannot be used in conjunction with conversation ID.") # Agent sessions are either tied to the conversation_id (for multi-turn conversation mode) # or the previous_response_id (for response chaining). If neither is present, a new session # is created for this request and stored under the current response_id. The current response_id @@ -456,7 +459,9 @@ async def _handle_inner_agent( if self._uses_hosted_responses_history: session.state.pop(_HOSTED_RESPONSES_HISTORY_SOURCE_ID, None) try: - await session_storage.set(context.conversation_id or context.response_id, session) + await session_storage.set(context.response_id, session) + if context.conversation_id is not None: + await session_storage.set(context.conversation_id, session) except Exception as save_error: save_failure = save_error if request_interrupted: @@ -506,9 +511,10 @@ async def _handle_inner_workflow( if are_options_set: logger.warning("Workflow agent doesn't support runtime options. They will be ignored.") - if request.get("previous_response_id") is not None and context.conversation_id is not None: + previous_response_id = request.get("previous_response_id") + if previous_response_id is not None and context.conversation_id is not None: raise RuntimeError("Previous response ID cannot be used in conjunction with conversation ID.") - context_id = request.get("previous_response_id") or context.conversation_id + context_id = previous_response_id or context.conversation_id if not isinstance(self._agent, WorkflowAgent): raise RuntimeError("Agent is not a workflow agent.") @@ -539,75 +545,127 @@ async def _handle_inner_workflow( latest_checkpoint = await restore_storage.get_latest(workflow_name=self._agent.workflow.name) if latest_checkpoint is not None: latest_checkpoint_id = latest_checkpoint.checkpoint_id + elif previous_response_id is not None: + raise RuntimeError( + f"Cannot find an existing workflow checkpoint for " + f"previous_response_id={previous_response_id}. " + "Ensure that the previous response was created successfully and that the ID is correct." + ) # Storage that will receive checkpoints written during this turn. - # When the caller chains with previous_response_id, the next turn - # will reference the current response_id as its previous_response_id, - # so new checkpoints must land under the current response_id (or the - # conversation_id when set). When conversation_id is set, this - # matches restore_storage; when only previous_response_id was - # supplied, restore_storage points at the *prior* response's - # directory and write_storage points at the *current* response's. - write_context_id = context.conversation_id or context.response_id + # Every turn writes under its current response_id so a later request + # can branch from that exact state via previous_response_id. When a + # conversation_id is set, its store is updated after the run as a + # latest-state alias while the response snapshot remains unchanged. + write_context_id = context.response_id validate_path_segment(write_context_id, kind="context id") write_storage = self._checkpoint_storage_provider.get_store( config=self.config, context_id=write_context_id, ) - # Multi-turn pattern: when we have a prior checkpoint, restore it - # first (drive the workflow back to idle with prior state intact), - # then make a separate call that delivers the new user input. This - # depends on Workflow.run preserving shared state across calls. The - # restore-only call may yield events from any pending in-flight - # work in the checkpoint; we consume those internally here so they - # don't surface to the response stream as duplicates. - # - # If the restored checkpoint had pending request_info events, the - # restore-only call replays them through - # ``WorkflowAgent._convert_workflow_event_to_agent_response_updates`` - # and populates ``self._agent.pending_requests``. That is the correct - # state: those requests are genuinely outstanding, and the next - # ``run(input_messages, ...)`` call may contain ``function_call_output`` - # items (carried as FunctionResult/FunctionApprovalResponse content) - # that fulfill them via :meth:`WorkflowAgent._process_pending_requests`. - if latest_checkpoint_id is not None: - async for _ in self._agent.run( + request_failure: Exception | None = None + request_interrupted = False + try: + # Multi-turn pattern: when we have a prior checkpoint, restore it + # first (drive the workflow back to idle with prior state intact), + # then make a separate call that delivers the new user input. This + # depends on Workflow.run preserving shared state across calls. The + # restore-only call may yield events from any pending in-flight + # work in the checkpoint; we consume those internally here so they + # don't surface to the response stream as duplicates. + # + # If the restored checkpoint had pending request_info events, the + # restore-only call replays them through + # ``WorkflowAgent._convert_workflow_event_to_agent_response_updates`` + # and populates ``self._agent.pending_requests``. That is the correct + # state: those requests are genuinely outstanding, and the next + # ``run(input_messages, ...)`` call may contain ``function_call_output`` + # items (carried as FunctionResult/FunctionApprovalResponse content) + # that fulfill them via :meth:`WorkflowAgent._process_pending_requests`. + if latest_checkpoint_id is not None: + async for _ in self._agent.run( + stream=True, + checkpoint_id=latest_checkpoint_id, + checkpoint_storage=restore_storage, + ): + pass + + tracker = _OutputItemTracker(response_event_stream) + + # Run the workflow agent in streaming mode with the new user input. + async for update in self._agent.run( + input_messages, stream=True, - checkpoint_id=latest_checkpoint_id, - checkpoint_storage=restore_storage, + checkpoint_storage=write_storage, ): - pass - - tracker = _OutputItemTracker(response_event_stream) - - # Run the workflow agent in streaming mode with the new user input. - async for update in self._agent.run( - input_messages, - stream=True, - checkpoint_storage=write_storage, - ): - for content in update.contents: - for event in tracker.handle(content): - yield event - if tracker.needs_async: - async for item in _to_outputs( - response_event_stream, content, approval_storage=approval_storage - ): - yield item - tracker.needs_async = False - - # Close any remaining active builder - for event in tracker.close(): - yield event - - await self._delete_not_latest_checkpoints(write_storage, self._agent.workflow.name) + for content in update.contents: + for event in tracker.handle(content): + yield event + if tracker.needs_async: + async for item in _to_outputs( + response_event_stream, content, approval_storage=approval_storage + ): + yield item + tracker.needs_async = False + + # Close any remaining active builder + for event in tracker.close(): + yield event + except (asyncio.CancelledError, GeneratorExit): + request_interrupted = True + raise + except Exception as ex: + request_failure = ex + raise + finally: + try: + await self._finalize_workflow_checkpoints( + write_storage, + workflow_name=self._agent.workflow.name, + conversation_id=context.conversation_id, + ) + except Exception as save_error: + if request_interrupted: + logger.error( + "Failed to finalize workflow checkpoints while unwinding an interrupted request", + exc_info=(type(save_error), save_error, save_error.__traceback__), + ) + elif request_failure is not None: + logger.error( + "Failed to finalize workflow checkpoints after a workflow failure", + exc_info=(type(save_error), save_error, save_error.__traceback__), + ) + else: + raise yield response_event_stream.emit_completed() except Exception as ex: logger.exception("Failed to produce response for workflow agent") for event in self._emit_failure(response_event_stream, tracker, ex): yield event + async def _finalize_workflow_checkpoints( + self, + response_storage: CheckpointStorage, + *, + workflow_name: str, + conversation_id: str | None, + ) -> None: + """Keep one response checkpoint and update the conversation's latest-state alias.""" + await self._delete_not_latest_checkpoints(response_storage, workflow_name) + if conversation_id is None: + return + + latest_checkpoint = await response_storage.get_latest(workflow_name=workflow_name) + if latest_checkpoint is None: + return + conversation_storage = self._checkpoint_storage_provider.get_store( + config=self.config, + context_id=conversation_id, + ) + await conversation_storage.save(latest_checkpoint) + await self._delete_not_latest_checkpoints(conversation_storage, workflow_name) + @staticmethod async def _delete_not_latest_checkpoints(checkpoint_storage: CheckpointStorage, workflow_name: str) -> None: """Delete all checkpoints except the latest one. diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index 5b0dfda65c..22e3ce8f28 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -31,6 +31,7 @@ ChatMiddlewareLayer, ChatResponse, ChatResponseUpdate, + CheckpointStorage, Content, FunctionInvocationLayer, HistoryProvider, @@ -43,6 +44,7 @@ SupportsAgentRun, WorkflowAgent, WorkflowBuilder, + WorkflowCheckpoint, WorkflowContext, executor, tool, @@ -457,6 +459,23 @@ async def test_hosted_request_requires_protocol_v2(self) -> None: asyncio.Event(), ) + async def test_previous_response_rejected_with_conversation(self) -> None: + agent = _make_agent() + server = _make_server(agent) + + response = await _post( + server, + previous_response_id="caresp_aaaaaaaaaaaaaaaa00" + "1" * 32, + conversation_id="conversation-1", + ) + + assert response.status_code == 200 + body = response.json() + assert body["status"] == "failed" + assert body["error"]["message"] == "Previous response ID cannot be used in conjunction with conversation ID." + agent.run.assert_not_called() + agent.create_session.assert_not_called() + async def test_previous_response_requires_existing_agent_session(self) -> None: agent = _make_agent() server = _make_server(agent, session_store=SessionStore()) @@ -483,6 +502,85 @@ async def test_previous_response_requires_existing_agent_session(self) -> None: class TestAgentSessionPersistence: + async def test_conversations_are_isolated_and_response_snapshots_are_saved(self) -> None: + seen_counts: list[int] = [] + seen_session_ids: list[str] = [] + + def run_with_state(*args: Any, **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse]: + del args + session = kwargs["session"] + assert isinstance(session, AgentSession) + count = int(session.state.get("turn_count", 0)) + 1 + session.state["turn_count"] = count + seen_counts.append(count) + seen_session_ids.append(session.session_id) + + async def updates() -> AsyncIterator[AgentResponseUpdate]: + yield AgentResponseUpdate(contents=[Content.from_text(f"turn {count}")], role="assistant") + + return ResponseStream(updates(), finalizer=AgentResponse.from_updates) + + agent = _make_agent() + agent.run = MagicMock(side_effect=run_with_state) + store = SessionStore() + server = _make_server(agent, session_store=store) + + first = await _post(server, input_text="first", conversation_id="conversation-1") + second = await _post(server, input_text="other", conversation_id="conversation-2") + third = await _post(server, input_text="continue", conversation_id="conversation-1") + + assert seen_counts == [1, 1, 2] + assert seen_session_ids[0] != seen_session_ids[1] + assert seen_session_ids[2] == seen_session_ids[0] + for snapshot_id in ( + first.json()["id"], + second.json()["id"], + third.json()["id"], + "conversation-1", + "conversation-2", + ): + assert await store.get(snapshot_id) is not None + assert agent.create_session.call_count == 2 + + async def test_conversation_response_snapshots_support_branching(self) -> None: + seen_counts: list[int] = [] + + def run_with_state(*args: Any, **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse]: + del args + session = kwargs["session"] + assert isinstance(session, AgentSession) + count = int(session.state.get("turn_count", 0)) + 1 + session.state["turn_count"] = count + seen_counts.append(count) + + async def updates() -> AsyncIterator[AgentResponseUpdate]: + yield AgentResponseUpdate(contents=[Content.from_text(f"turn {count}")], role="assistant") + + return ResponseStream(updates(), finalizer=AgentResponse.from_updates) + + agent = _make_agent() + agent.run = MagicMock(side_effect=run_with_state) + store = SessionStore() + server = _make_server(agent, session_store=store) + + first = await _post(server, input_text="first", conversation_id="conversation-1") + await _post(server, input_text="second", conversation_id="conversation-1") + branch = await _post(server, input_text="branch", previous_response_id=first.json()["id"]) + + assert first.status_code == 200 + assert branch.status_code == 200 + assert seen_counts == [1, 2, 2] + + conversation_snapshot = await store.get("conversation-1") + first_snapshot = await store.get(first.json()["id"]) + branch_snapshot = await store.get(branch.json()["id"]) + assert conversation_snapshot is not None + assert first_snapshot is not None + assert branch_snapshot is not None + assert conversation_snapshot.state["turn_count"] == 2 + assert first_snapshot.state["turn_count"] == 1 + assert branch_snapshot.state["turn_count"] == 2 + async def test_previous_response_chain_restores_session_state(self) -> None: seen_counts: list[int] = [] seen_session_ids: list[str] = [] @@ -4263,6 +4361,167 @@ async def test_basic_text_response_streaming(self) -> None: text_done = [e for e in events if e["event"] == "response.output_text.done"] assert any(e["data"]["text"] == "hello stream" for e in text_done) + async def test_previous_response_requires_existing_workflow_checkpoint(self) -> None: + workflow_agent = _build_text_workflow_agent("should not run") + server = _make_server(workflow_agent) + missing_response_id = "caresp_aaaaaaaaaaaaaaaa00" + "1" * 32 + + response = await _post(server, previous_response_id=missing_response_id) + + assert response.status_code == 200 + body = response.json() + assert body["status"] == "failed" + assert ( + f"Cannot find an existing workflow checkpoint for previous_response_id={missing_response_id}." + in body["error"]["message"] + ) + assert body["output"] == [] + + @pytest.mark.parametrize("stream", [False, True]) + async def test_conversation_response_checkpoints_support_branching(self, stream: bool) -> None: + @executor + async def count_turns(messages: list[Message], ctx: WorkflowContext[Any, AgentResponse]) -> None: + del messages + turn_count = int(ctx.get_state("turn_count", 0)) + 1 + ctx.set_state("turn_count", turn_count) + await ctx.yield_output( + AgentResponse(messages=[Message("assistant", [Content.from_text(f"turn {turn_count}")])]) + ) + + workflow_agent = WorkflowAgent( + workflow=WorkflowBuilder(start_executor=count_turns).build(), + name="Counting Workflow Agent", + ) + server = _make_server(workflow_agent) + + def response_body(response: httpx.Response) -> dict[str, Any]: + if not stream: + return response.json() + return _parse_sse_events(response.text)[-1]["data"]["response"] + + first = await _post(server, input_text="first", conversation_id="conversation-1", stream=stream) + second = await _post(server, input_text="second", conversation_id="conversation-1", stream=stream) + first_body = response_body(first) + branch = await _post( + server, + input_text="branch", + previous_response_id=first_body["id"], + stream=stream, + ) + + assert first.status_code == 200 + assert second.status_code == 200 + assert branch.status_code == 200 + assert first_body["status"] == "completed" + assert response_body(second)["status"] == "completed" + assert response_body(branch)["status"] == "completed" + branch_text = [ + part["text"] + for item in response_body(branch)["output"] + if item["type"] == "message" + for part in item.get("content", []) + if part["type"] == "output_text" + ] + assert branch_text == ["turn 2"] + + async def test_failed_conversation_workflow_promotes_latest_response_checkpoint(self) -> None: + workflow_agent = _build_text_workflow_agent("ignored") + checkpoint = WorkflowCheckpoint( + workflow_name=workflow_agent.workflow.name, + graph_signature_hash="hash", + ) + + async def updates(checkpoint_storage: CheckpointStorage) -> AsyncIterator[AgentResponseUpdate]: + await checkpoint_storage.save(checkpoint) + raise RuntimeError("workflow failed") + yield # pragma: no cover + + def failing_run(*args: Any, **kwargs: Any) -> AsyncIterator[AgentResponseUpdate]: + del args + return updates(kwargs["checkpoint_storage"]) + + server = _make_server(workflow_agent) + request = CreateResponse(model="m", input="hi") + context = ResponseContext( + response_id="response-1", + conversation_id="conversation-1", + mode_flags=MagicMock(), + ) + + with ( + patch.object(ResponseContext, "get_input_items", new=AsyncMock(return_value=[])), + patch.object(workflow_agent, "run", side_effect=failing_run), + ): + events = [ + event + async for event in server._handle_inner_workflow( # pyright: ignore[reportPrivateUsage] + request, + context, + ) + ] + + conversation_storage = server._checkpoint_storage_provider.get_store( # pyright: ignore[reportPrivateUsage] + config=server.config, + context_id="conversation-1", + ) + latest = await conversation_storage.get_latest(workflow_name=workflow_agent.workflow.name) + assert latest is not None + assert latest.checkpoint_id == checkpoint.checkpoint_id + assert events[-1].get("type") == "response.failed" + + @pytest.mark.parametrize("interruption", ["cancel", "close"]) + async def test_interrupted_conversation_workflow_promotes_latest_response_checkpoint( + self, + interruption: str, + ) -> None: + workflow_agent = _build_text_workflow_agent("ignored") + checkpoint = WorkflowCheckpoint( + workflow_name=workflow_agent.workflow.name, + graph_signature_hash="hash", + ) + + async def updates(checkpoint_storage: CheckpointStorage) -> AsyncIterator[AgentResponseUpdate]: + await checkpoint_storage.save(checkpoint) + yield AgentResponseUpdate(contents=[Content.from_text("started")], role="assistant") + await asyncio.Event().wait() + + def streaming_run(*args: Any, **kwargs: Any) -> AsyncIterator[AgentResponseUpdate]: + del args + return updates(kwargs["checkpoint_storage"]) + + server = _make_server(workflow_agent) + request = CreateResponse(model="m", input="hi", stream=True) + context = ResponseContext( + response_id="response-1", + conversation_id="conversation-1", + mode_flags=MagicMock(), + ) + + with ( + patch.object(ResponseContext, "get_input_items", new=AsyncMock(return_value=[])), + patch.object(workflow_agent, "run", side_effect=streaming_run), + ): + handler = cast( + AsyncGenerator[Any, None], + server._handle_inner_workflow(request, context), # pyright: ignore[reportPrivateUsage] + ) + await anext(handler) + await anext(handler) + await anext(handler) + if interruption == "cancel": + with pytest.raises(asyncio.CancelledError): + await handler.athrow(asyncio.CancelledError()) + else: + await handler.aclose() + + conversation_storage = server._checkpoint_storage_provider.get_store( # pyright: ignore[reportPrivateUsage] + config=server.config, + context_id="conversation-1", + ) + latest = await conversation_storage.get_latest(workflow_name=workflow_agent.workflow.name) + assert latest is not None + assert latest.checkpoint_id == checkpoint.checkpoint_id + async def test_non_streaming_emits_mcp_approval_request_and_persists_to_storage(self) -> None: workflow_agent, mock_agent = _build_approval_workflow_agent(approval_request_id="apr_wf_ns") server = _make_server(workflow_agent)