diff --git a/api/controllers/console/app/completion.py b/api/controllers/console/app/completion.py index 3fe721def62..beca00d621b 100644 --- a/api/controllers/console/app/completion.py +++ b/api/controllers/console/app/completion.py @@ -351,24 +351,31 @@ def _resolve_current_user_agent_debug_conversation_id( app_model: App, agent_id: str | None, draft_type: AgentConfigDraftType, + start_new: bool = False, ) -> str: - """Resolve the current editor's conversation without crossing draft surfaces.""" + """Resolve or rotate the current editor's conversation within one draft surface. + + ``start_new`` rotates the scoped mapping through ``AgentRosterService`` so + the old runtime session is retired before the new conversation is used. + Continuations and Build chat keep resolving the existing mapping. + """ roster_service = AgentRosterService(session) - if agent_id: - return roster_service.get_or_create_agent_app_debug_conversation_id( - tenant_id=current_tenant_id, - agent_id=agent_id, - account_id=current_user.id, - draft_type=draft_type, - ) + resolved_agent_id = agent_id + if not resolved_agent_id: + agent = roster_service.get_app_backing_agent(tenant_id=current_tenant_id, app_id=str(app_model.id)) + if agent is None: + raise AgentNotFoundError() + resolved_agent_id = agent.id - agent = roster_service.get_app_backing_agent(tenant_id=current_tenant_id, app_id=str(app_model.id)) - if agent is None: - raise AgentNotFoundError() - return roster_service.get_or_create_agent_app_debug_conversation_id( + resolve_conversation = ( + roster_service.refresh_agent_app_debug_conversation_id + if start_new + else roster_service.get_or_create_agent_app_debug_conversation_id + ) + return resolve_conversation( tenant_id=current_tenant_id, - agent_id=agent.id, + agent_id=resolved_agent_id, account_id=current_user.id, draft_type=draft_type, ) @@ -387,13 +394,17 @@ def _create_chat_message( args = args_model.model_dump(exclude_none=True, by_alias=True) if AppMode.value_of(app_model.mode) == AppMode.AGENT: + draft_type = AgentConfigDraftType(args_model.draft_type) + # Preview follows the normal chat contract: an omitted/empty conversation ID starts a new + # conversation. Build chat keeps its stable mapping so build drafts and finalization stay continuous. debug_conversation_id = _resolve_current_user_agent_debug_conversation_id( session=session, current_tenant_id=current_tenant_id or app_model.tenant_id, current_user=current_user, app_model=app_model, agent_id=agent_id, - draft_type=AgentConfigDraftType(args_model.draft_type), + draft_type=draft_type, + start_new=draft_type == AgentConfigDraftType.DRAFT and not args_model.conversation_id, ) if args_model.conversation_id and args_model.conversation_id != debug_conversation_id: raise NotFound("Conversation Not Exists.") diff --git a/api/services/agent/roster_service.py b/api/services/agent/roster_service.py index 979cffafc64..8e8758890c9 100644 --- a/api/services/agent/roster_service.py +++ b/api/services/agent/roster_service.py @@ -660,15 +660,18 @@ class AgentRosterService: agent_id: str, account_id: str, draft_type: AgentConfigDraftType = AgentConfigDraftType.DEBUG_BUILD, - commit: bool = True, ) -> str: """Start a new scoped console conversation for the current Agent App editor. If this account already has a mapping for the requested draft surface, the previous - conversation is abandoned first: any ACTIVE conversation-owned Agent - runtime sessions for that old conversation are sent through best-effort - backend cleanup and then retired locally even when enqueueing fails. The - other draft surface is left untouched. + conversation is abandoned after the replacement mapping is committed: any ACTIVE + conversation-owned Agent runtime sessions for that old conversation are sent through + best-effort backend cleanup and then retired locally even when enqueueing fails. This + order prevents a failed database commit from retiring the still-current runtime session. + The other draft surface is left untouched. + + A user and draft surface own one current mapping. If new-conversation requests overlap, + the last committed rotation becomes current and earlier response IDs cannot be continued. """ agent = self._session.scalar( @@ -691,6 +694,7 @@ class AgentRosterService: app_id=backing_app_id, account_id=account_id, ) + previous_conversation: tuple[str, str] | None = None mapping = self._session.scalar( select(AgentDebugConversation).where( AgentDebugConversation.tenant_id == tenant_id, @@ -714,19 +718,22 @@ class AgentRosterService: previous_app_id = mapping.app_id previous_conversation_id = mapping.conversation_id if previous_conversation_id: - self._cleanup_debug_conversation_runtime_sessions( - tenant_id=tenant_id, - agent_id=agent_id, - account_id=account_id, - draft_type=draft_type, - app_id=previous_app_id or backing_app_id, - conversation_id=previous_conversation_id, - ) + previous_conversation = (previous_app_id or backing_app_id, previous_conversation_id) mapping.app_id = backing_app_id mapping.conversation_id = conversation_id self._session.flush() - if commit: - self._session.commit() + self._session.commit() + + if previous_conversation: + previous_app_id, previous_conversation_id = previous_conversation + self._cleanup_debug_conversation_runtime_sessions( + tenant_id=tenant_id, + agent_id=agent_id, + account_id=account_id, + draft_type=draft_type, + app_id=previous_app_id, + conversation_id=previous_conversation_id, + ) return conversation_id def _cleanup_debug_conversation_runtime_sessions( @@ -739,8 +746,8 @@ class AgentRosterService: app_id: str, conversation_id: str, ) -> None: - session_store = AgentAppRuntimeSessionStore() try: + session_store = AgentAppRuntimeSessionStore() stored_sessions = session_store.list_active_sessions_for_conversation( tenant_id=tenant_id, app_id=app_id, diff --git a/api/tests/unit_tests/controllers/console/agent/test_agent_controllers.py b/api/tests/unit_tests/controllers/console/agent/test_agent_controllers.py index ee13618fe1b..8f395f5549d 100644 --- a/api/tests/unit_tests/controllers/console/agent/test_agent_controllers.py +++ b/api/tests/unit_tests/controllers/console/agent/test_agent_controllers.py @@ -1530,8 +1530,35 @@ def test_drain_streaming_generate_response_raises_when_stream_ends_early() -> No completion_controller._drain_streaming_generate_response(response) -def test_agent_chat_helper_forces_agent_streaming_and_external_trace( - app: Flask, monkeypatch: pytest.MonkeyPatch, account_id: str +@pytest.mark.parametrize( + ("payload_extra", "expected_draft_type", "expected_start_new"), + [ + ({}, AgentConfigDraftType.DRAFT, True), + ({"conversation_id": None}, AgentConfigDraftType.DRAFT, True), + ({"conversation_id": ""}, AgentConfigDraftType.DRAFT, True), + ( + {"conversation_id": "00000000-0000-0000-0000-000000000001"}, + AgentConfigDraftType.DRAFT, + False, + ), + ({"draft_type": "debug_build"}, AgentConfigDraftType.DEBUG_BUILD, False), + ( + { + "draft_type": "debug_build", + "conversation_id": "00000000-0000-0000-0000-000000000001", + }, + AgentConfigDraftType.DEBUG_BUILD, + False, + ), + ], +) +def test_agent_chat_helper_resolves_scoped_conversation_and_forces_streaming( + app: Flask, + monkeypatch: pytest.MonkeyPatch, + account_id: str, + payload_extra: dict[str, str | None], + expected_draft_type: AgentConfigDraftType, + expected_start_new: bool, ) -> None: app_model = SimpleNamespace(id="app-1", tenant_id="tenant-1", mode="agent") current_user = SimpleNamespace(id=account_id) @@ -1543,7 +1570,7 @@ def test_agent_chat_helper_forces_agent_streaming_and_external_trace( def resolve_debug_conversation(**kwargs: object) -> str: captured["resolve_debug_conversation"] = kwargs - return "debug-conversation-1" + return "00000000-0000-0000-0000-000000000001" monkeypatch.setattr(completion_controller.AppGenerateService, "generate", generate) monkeypatch.setattr( @@ -1555,7 +1582,8 @@ def test_agent_chat_helper_forces_agent_streaming_and_external_trace( completion_controller.helper, "compact_generate_response", lambda response: {"response": response} ) with app.test_request_context( - json={"inputs": {}, "query": "hello", "response_mode": "streaming"}, headers={"X-Trace-Id": "trace-1"} + json={"inputs": {}, "query": "hello", "response_mode": "streaming", **payload_extra}, + headers={"X-Trace-Id": "trace-1"}, ): result = completion_controller._create_chat_message( current_user=current_user, app_model=app_model, session=Mock() @@ -1566,10 +1594,12 @@ def test_agent_chat_helper_forces_agent_streaming_and_external_trace( assert captured["streaming"] is True args = cast(dict[str, object], captured["args"]) assert args["response_mode"] == "streaming" - assert args["conversation_id"] == "debug-conversation-1" + assert args["conversation_id"] == "00000000-0000-0000-0000-000000000001" assert args["auto_generate_name"] is False assert args["external_trace_id"] == "trace-1" - assert cast(dict[str, object], captured["resolve_debug_conversation"])["draft_type"] == AgentConfigDraftType.DRAFT + resolve_call = cast(dict[str, object], captured["resolve_debug_conversation"]) + assert resolve_call["draft_type"] == expected_draft_type + assert resolve_call["start_new"] is expected_start_new def test_agent_chat_helper_ignores_private_exit_intent_payload_key( @@ -1617,14 +1647,28 @@ def test_agent_chat_helper_ignores_private_exit_intent_payload_key( assert completion_controller.AGENT_RUNTIME_EXIT_INTENT_ARG not in args -def test_agent_chat_helper_rejects_foreign_debug_conversation( - app: Flask, monkeypatch: pytest.MonkeyPatch, account_id: str +@pytest.mark.parametrize( + ("payload_extra", "expected_draft_type"), + [ + ({}, AgentConfigDraftType.DRAFT), + ({"draft_type": "debug_build"}, AgentConfigDraftType.DEBUG_BUILD), + ], +) +def test_agent_chat_helper_rejects_foreign_debug_conversation_before_generation( + app: Flask, + monkeypatch: pytest.MonkeyPatch, + account_id: str, + payload_extra: dict[str, str], + expected_draft_type: AgentConfigDraftType, ) -> None: app_model = SimpleNamespace(id="app-1", tenant_id="tenant-1", mode="agent") + generate = MagicMock() + resolve_debug_conversation = MagicMock(return_value="owned-conversation") + monkeypatch.setattr(completion_controller.AppGenerateService, "generate", generate) monkeypatch.setattr( completion_controller, "_resolve_current_user_agent_debug_conversation_id", - lambda **kwargs: "owned-conversation", + resolve_debug_conversation, ) with app.test_request_context( json={ @@ -1632,6 +1676,7 @@ def test_agent_chat_helper_rejects_foreign_debug_conversation( "query": "hello", "response_mode": "streaming", "conversation_id": "00000000-0000-0000-0000-000000000001", + **payload_extra, } ): with pytest.raises(NotFound): @@ -1643,6 +1688,12 @@ def test_agent_chat_helper_rejects_foreign_debug_conversation( session=Mock(), ) + resolve_debug_conversation.assert_called_once() + resolve_call = resolve_debug_conversation.call_args.kwargs + assert resolve_call["draft_type"] == expected_draft_type + assert resolve_call["start_new"] is False + generate.assert_not_called() + def test_resolve_current_user_agent_debug_conversation_uses_agent_or_backing_app( monkeypatch: pytest.MonkeyPatch, @@ -1657,6 +1708,10 @@ def test_resolve_current_user_agent_debug_conversation_uses_agent_or_backing_app calls.append({"get_or_create": kwargs}) return f"debug-{kwargs['agent_id']}" + def refresh_agent_app_debug_conversation_id(self, **kwargs: object) -> str: + calls.append({"refresh": kwargs}) + return f"new-{kwargs['agent_id']}" + def get_app_backing_agent(self, **kwargs: object) -> object: calls.append({"get_app_backing_agent": kwargs}) return SimpleNamespace(id="backing-agent") @@ -1669,6 +1724,7 @@ def test_resolve_current_user_agent_debug_conversation_uses_agent_or_backing_app app_model=SimpleNamespace(id="app-1"), agent_id="agent-1", draft_type=AgentConfigDraftType.DRAFT, + start_new=True, ) fallback_id = completion_controller._resolve_current_user_agent_debug_conversation_id( session="session-1", # type: ignore[arg-type] @@ -1678,10 +1734,20 @@ def test_resolve_current_user_agent_debug_conversation_uses_agent_or_backing_app agent_id=None, draft_type=AgentConfigDraftType.DEBUG_BUILD, ) - assert explicit_id == "debug-agent-1" + fallback_preview_id = completion_controller._resolve_current_user_agent_debug_conversation_id( + session="session-1", # type: ignore[arg-type] + current_tenant_id="tenant-1", + current_user=SimpleNamespace(id="account-1"), + app_model=SimpleNamespace(id="app-1"), + agent_id=None, + draft_type=AgentConfigDraftType.DRAFT, + start_new=True, + ) + assert explicit_id == "new-agent-1" assert fallback_id == "debug-backing-agent" + assert fallback_preview_id == "new-backing-agent" assert calls[1] == { - "get_or_create": { + "refresh": { "tenant_id": "tenant-1", "agent_id": "agent-1", "account_id": "account-1", @@ -1697,6 +1763,15 @@ def test_resolve_current_user_agent_debug_conversation_uses_agent_or_backing_app "draft_type": AgentConfigDraftType.DEBUG_BUILD, } } + assert calls[6] == {"get_app_backing_agent": {"tenant_id": "tenant-1", "app_id": "app-1"}} + assert calls[7] == { + "refresh": { + "tenant_id": "tenant-1", + "agent_id": "backing-agent", + "account_id": "account-1", + "draft_type": AgentConfigDraftType.DRAFT, + } + } @pytest.mark.parametrize( diff --git a/api/tests/unit_tests/services/agent/test_agent_services.py b/api/tests/unit_tests/services/agent/test_agent_services.py index 0ff248feb0f..6ac297125ae 100644 --- a/api/tests/unit_tests/services/agent/test_agent_services.py +++ b/api/tests/unit_tests/services/agent/test_agent_services.py @@ -3496,8 +3496,61 @@ class TestAgentAppBackingAgent: assert session.deleted == [] assert session.commits == 1 - def test_refresh_agent_app_debug_conversation_enqueues_cleanup_for_old_runtime_sessions( + def test_refresh_agent_app_debug_conversation_rotates_preview_mapping_each_time( self, monkeypatch: pytest.MonkeyPatch + ): + agent = Agent( + id="agent-1", + tenant_id="tenant-1", + name="Iris", + description="", + agent_kind=AgentKind.DIFY_AGENT, + scope=AgentScope.ROSTER, + source=AgentSource.AGENT_APP, + status=AgentStatus.ACTIVE, + app_id="app-1", + ) + session = FakeSession(scalar=[agent, None]) + service = AgentRosterService(session) + cleanup = MagicMock() + monkeypatch.setattr(service, "_cleanup_debug_conversation_runtime_sessions", cleanup) + + first_conversation_id = service.refresh_agent_app_debug_conversation_id( + tenant_id="tenant-1", + agent_id="agent-1", + account_id="account-1", + draft_type=AgentConfigDraftType.DRAFT, + ) + mapping = next(value for value in session.added if isinstance(value, AgentDebugConversation)) + session._scalar.extend([agent, mapping]) + second_conversation_id = service.refresh_agent_app_debug_conversation_id( + tenant_id="tenant-1", + agent_id="agent-1", + account_id="account-1", + draft_type=AgentConfigDraftType.DRAFT, + ) + + assert first_conversation_id != second_conversation_id + assert mapping.draft_type == AgentConfigDraftType.DRAFT + assert mapping.conversation_id == second_conversation_id + assert len([value for value in session.added if isinstance(value, Conversation)]) == 2 + cleanup.assert_called_once_with( + tenant_id="tenant-1", + agent_id="agent-1", + account_id="account-1", + draft_type=AgentConfigDraftType.DRAFT, + app_id="app-1", + conversation_id=first_conversation_id, + ) + + @pytest.mark.parametrize( + "draft_type", + [AgentConfigDraftType.DRAFT, AgentConfigDraftType.DEBUG_BUILD], + ) + def test_refresh_agent_app_debug_conversation_enqueues_cleanup_for_old_runtime_sessions( + self, + monkeypatch: pytest.MonkeyPatch, + draft_type: AgentConfigDraftType, ): agent = Agent( id="agent-1", @@ -3525,9 +3578,21 @@ class TestAgentAppBackingAgent: ) session = FakeSession(scalar=[agent, mapping]) service = AgentRosterService(session) + events: list[str] = [] + original_commit = session.commit + + def record_commit() -> None: + events.append("commit") + original_commit() + + def list_active_sessions(**kwargs: object) -> list[object]: + events.append("cleanup") + return [stored_session] + cleanup_delay = MagicMock() cleanup_store = MagicMock() - cleanup_store.list_active_sessions_for_conversation.return_value = [stored_session] + cleanup_store.list_active_sessions_for_conversation.side_effect = list_active_sessions + monkeypatch.setattr(session, "commit", record_commit) monkeypatch.setattr(roster_service, "AgentAppRuntimeSessionStore", lambda: cleanup_store) monkeypatch.setattr(roster_service.cleanup_conversation_agent_runtime_session, "delay", cleanup_delay) @@ -3535,6 +3600,7 @@ class TestAgentAppBackingAgent: tenant_id="tenant-1", agent_id="agent-1", account_id="account-1", + draft_type=draft_type, ) cleanup_store.list_active_sessions_for_conversation.assert_called_once_with( @@ -3546,10 +3612,10 @@ class TestAgentAppBackingAgent: payload = cleanup_delay.call_args.args[0] assert payload["metadata"]["conversation_id"] == "old-conversation" assert payload["metadata"]["agent_id"] == "agent-9" - assert payload["metadata"]["draft_type"] == "debug_build" + assert payload["metadata"]["draft_type"] == draft_type.value assert ( payload["idempotency_key"] - == "tenant-1:agent-1:account-1:debug_build:old-conversation:debug-session-cleanup:" + == f"tenant-1:agent-1:account-1:{draft_type.value}:old-conversation:debug-session-cleanup:" "agent-9:snap-9:run-old" ) cleanup_store.mark_cleaned.assert_called_once_with( @@ -3558,6 +3624,41 @@ class TestAgentAppBackingAgent: ) assert mapping.app_id == "app-1" assert mapping.conversation_id == conversation_id + assert events == ["commit", "cleanup"] + + def test_refresh_agent_app_debug_conversation_does_not_cleanup_when_commit_fails( + self, monkeypatch: pytest.MonkeyPatch + ): + agent = Agent( + id="agent-1", + tenant_id="tenant-1", + name="Iris", + description="", + agent_kind=AgentKind.DIFY_AGENT, + scope=AgentScope.ROSTER, + source=AgentSource.AGENT_APP, + status=AgentStatus.ACTIVE, + app_id="app-1", + ) + mapping = SimpleNamespace(app_id="old-app", conversation_id="old-conversation") + session = FakeSession(scalar=[agent, mapping]) + service = AgentRosterService(session) + runtime_store_factory = MagicMock() + cleanup_delay = MagicMock() + monkeypatch.setattr(session, "commit", MagicMock(side_effect=RuntimeError("database unavailable"))) + monkeypatch.setattr(roster_service, "AgentAppRuntimeSessionStore", runtime_store_factory) + monkeypatch.setattr(roster_service.cleanup_conversation_agent_runtime_session, "delay", cleanup_delay) + + with pytest.raises(RuntimeError, match="database unavailable"): + service.refresh_agent_app_debug_conversation_id( + tenant_id="tenant-1", + agent_id="agent-1", + account_id="account-1", + draft_type=AgentConfigDraftType.DRAFT, + ) + + runtime_store_factory.assert_not_called() + cleanup_delay.assert_not_called() def test_refresh_agent_app_debug_conversation_marks_old_runtime_sessions_clean_when_enqueue_fails( self, monkeypatch: pytest.MonkeyPatch