fix: fix revert and edit CTA

This commit is contained in:
Abhishek Kumar 2026-05-21 13:15:03 +05:30
parent 08a2435ba5
commit 5d9ae9da6d
3 changed files with 227 additions and 99 deletions

View file

@ -103,7 +103,7 @@ async def initialize_text_chat_session(
actual_revision=e.actual_revision,
) from e
return await _reload_text_chat_session(run_id, text_session)
return await _reload_text_chat_session(run_id)
async def append_text_chat_user_message(
@ -138,7 +138,7 @@ async def append_text_chat_user_message(
actual_revision=e.actual_revision,
) from e
return await _reload_text_chat_session(run_id, text_session)
return await _reload_text_chat_session(run_id)
async def rewind_text_chat_session_state(
@ -175,7 +175,7 @@ async def rewind_text_chat_session_state(
},
)
return await _reload_text_chat_session(run_id, text_session)
return await _reload_text_chat_session(run_id)
async def execute_pending_text_chat_turn(
@ -262,7 +262,7 @@ async def execute_pending_text_chat_turn(
if cost_info is not None:
await db_client.update_workflow_run(run_id, cost_info=cost_info)
return await _reload_text_chat_session(run_id, text_session)
return await _reload_text_chat_session(run_id)
def validate_text_chat_turn_cursor(
@ -361,11 +361,12 @@ async def _mark_pending_turn_failed(
return
async def _reload_text_chat_session(
run_id: int,
text_session: WorkflowRunTextSessionModel,
) -> WorkflowRunTextSessionModel:
organization_id = text_session.workflow_run.workflow.organization_id
async def _reload_text_chat_session(run_id: int) -> WorkflowRunTextSessionModel:
organization_id = await db_client.get_organization_id_by_workflow_run_id(run_id)
if organization_id is None:
raise TextChatSessionExecutionError(
"Workflow run organization not found after update"
)
updated_text_session = await db_client.get_workflow_run_text_session(
run_id,
organization_id=organization_id,

View file

@ -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)