fix: fix revert and edit CTA

api/tests/test_text_chat_session_service.py

@@ -1,7 +1,13 @@
+from unittest.mock import AsyncMock
+
 import pytest
 
+import api.services.workflow.text_chat_session_service as text_chat_session_service
+from api.db.models import WorkflowRunTextSessionModel
 from api.services.workflow.text_chat_session_service import (
+    TextChatSessionExecutionError,
     TextChatTurnNotFoundError,
+    _reload_text_chat_session,
     build_pending_text_chat_turn,
     truncate_text_chat_future_turns,
     validate_text_chat_turn_cursor,
@@ -43,3 +49,43 @@ def test_validate_text_chat_turn_cursor_raises_for_missing_turn():
             {"turns": [{"id": "turn-1"}]},
             "turn-404",
         )
+
+
+@pytest.mark.asyncio
+async def test_reload_text_chat_session_uses_run_id_to_resolve_organization(
+    monkeypatch,
+):
+    reloaded_session = WorkflowRunTextSessionModel(workflow_run_id=123)
+    get_org_id = AsyncMock(return_value=77)
+    get_text_session = AsyncMock(return_value=reloaded_session)
+
+    monkeypatch.setattr(
+        text_chat_session_service.db_client,
+        "get_organization_id_by_workflow_run_id",
+        get_org_id,
+    )
+    monkeypatch.setattr(
+        text_chat_session_service.db_client,
+        "get_workflow_run_text_session",
+        get_text_session,
+    )
+
+    result = await _reload_text_chat_session(123)
+
+    assert result is reloaded_session
+    get_org_id.assert_awaited_once_with(123)
+    get_text_session.assert_awaited_once_with(123, organization_id=77)
+
+
+@pytest.mark.asyncio
+async def test_reload_text_chat_session_raises_when_run_organization_is_missing(
+    monkeypatch,
+):
+    monkeypatch.setattr(
+        text_chat_session_service.db_client,
+        "get_organization_id_by_workflow_run_id",
+        AsyncMock(return_value=None),
+    )
+
+    with pytest.raises(TextChatSessionExecutionError, match="organization not found"):
+        await _reload_text_chat_session(123)