From 5d9ae9da6da878f96e4593d9f023967c76b4597e Mon Sep 17 00:00:00 2001 From: Abhishek Kumar Date: Thu, 21 May 2026 13:15:03 +0530 Subject: [PATCH] fix: fix revert and edit CTA --- .../workflow/text_chat_session_service.py | 19 +- api/tests/test_text_chat_session_service.py | 46 +++ .../components/WorkflowTesterPanel.tsx | 261 ++++++++++++------ 3 files changed, 227 insertions(+), 99 deletions(-) diff --git a/api/services/workflow/text_chat_session_service.py b/api/services/workflow/text_chat_session_service.py index 6090ab7..886c9e3 100644 --- a/api/services/workflow/text_chat_session_service.py +++ b/api/services/workflow/text_chat_session_service.py @@ -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, diff --git a/api/tests/test_text_chat_session_service.py b/api/tests/test_text_chat_session_service.py index 6a91043..abbba0e 100644 --- a/api/tests/test_text_chat_session_service.py +++ b/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) diff --git a/ui/src/app/workflow/[workflowId]/components/WorkflowTesterPanel.tsx b/ui/src/app/workflow/[workflowId]/components/WorkflowTesterPanel.tsx index dc39382..a0df70b 100644 --- a/ui/src/app/workflow/[workflowId]/components/WorkflowTesterPanel.tsx +++ b/ui/src/app/workflow/[workflowId]/components/WorkflowTesterPanel.tsx @@ -1,6 +1,6 @@ "use client"; -import { AlertCircle, ArrowUpRight, Loader2, MessageSquareText, Mic, Phone, RefreshCw, RotateCcw, Sparkles, X } from "lucide-react"; +import { AlertCircle, Loader2, MessageSquareText, Mic, Pencil, Phone, RefreshCw, RotateCcw, Sparkles, X } from "lucide-react"; import { useRouter } from "next/navigation"; import { type ReactNode, useCallback, useEffect, useRef, useState } from "react"; import { toast } from "sonner"; @@ -80,6 +80,13 @@ interface TextChatToolEvent { resultText?: string; } +interface TurnActionState { + turnId: string; + type: "rewind" | "edit"; +} + +const EMPTY_TEXT_CHAT_TURNS: TextChatTurn[] = []; + function toTextChatSession(response: WorkflowRunTextSessionResponse): TextChatSession { return { ...response, @@ -229,6 +236,14 @@ function extractToolEvents(events: Array>): TextChatTool }, []); } +function getReplayCursorTurnId(turns: TextChatTurn[], turnId: string): string | null { + const turnIndex = turns.findIndex((turn) => turn.id === turnId); + if (turnIndex < 0) { + throw new Error("Turn not found"); + } + return turns[turnIndex - 1]?.id ?? null; +} + function ToolEventBubble({ event }: { event: TextChatToolEvent }) { return (
@@ -320,30 +335,6 @@ function EmbeddedVoiceTester({ return ( <>
-
-
-
- - Run {workflowRunId} - - - Browser voice test - -
-

- The call starts as soon as this test run is created. -

-
- -
-
(null); + const [editingTurnId, setEditingTurnId] = useState(null); + const [activeTurnAction, setActiveTurnAction] = useState(null); const scrollEndRef = useRef(null); - const turns = session?.session_data.turns ?? []; + const turns = session?.session_data.turns ?? EMPTY_TEXT_CHAT_TURNS; + const editingTurn = editingTurnId + ? turns.find((turn) => turn.id === editingTurnId) ?? null + : null; + const composerId = `workflow-tester-compose-${workflowId}`; const createSession = useCallback(async () => { if (disabled) return; @@ -482,15 +478,43 @@ function ManualTextChat({ onActiveChange?.(started); }, [onActiveChange, started]); - const sendMessage = useCallback(async () => { - if (!session || !draft.trim() || disabled) return; + const submitMessage = useCallback(async ( + messageText: string, + replayOptions?: TurnActionState, + ) => { + const trimmedText = messageText.trim(); + if (!session || !trimmedText || disabled) return; setSendingMessage(true); + if (replayOptions) { + setActiveTurnAction(replayOptions); + } try { + let activeSession = session; + + if (replayOptions) { + const rewindResponse = await rewindTextChatSessionApiV1WorkflowWorkflowIdTextChatSessionsRunIdRewindPost({ + path: { workflow_id: workflowId, run_id: activeSession.workflow_run_id }, + body: { + cursor_turn_id: getReplayCursorTurnId( + activeSession.session_data.turns, + replayOptions.turnId, + ), + expected_revision: activeSession.revision, + }, + }); + if (rewindResponse.error || !rewindResponse.data) { + throw new Error(extractSdkErrorMessage(rewindResponse.error, "Failed to rewind session")); + } + + activeSession = toTextChatSession(rewindResponse.data); + setSession(activeSession); + } + const response = await appendTextChatMessageApiV1WorkflowWorkflowIdTextChatSessionsRunIdMessagesPost({ - path: { workflow_id: workflowId, run_id: session.workflow_run_id }, + path: { workflow_id: workflowId, run_id: activeSession.workflow_run_id }, body: { - text: draft.trim(), - expected_revision: session.revision, + text: trimmedText, + expected_revision: activeSession.revision, }, }); if (response.error || !response.data) { @@ -498,38 +522,59 @@ function ManualTextChat({ } setSession(toTextChatSession(response.data)); setDraft(""); + setEditingTurnId(null); } catch (error) { toast.error(getErrorMessage(error)); } finally { setSendingMessage(false); - } - }, [disabled, draft, session, workflowId]); - - const rewindToTurn = useCallback(async (turnId: string) => { - if (!session || disabled) return; - setRewindingTurnId(turnId); - try { - const response = await rewindTextChatSessionApiV1WorkflowWorkflowIdTextChatSessionsRunIdRewindPost({ - path: { workflow_id: workflowId, run_id: session.workflow_run_id }, - body: { - cursor_turn_id: turnId, - expected_revision: session.revision, - }, - }); - if (response.error || !response.data) { - throw new Error(extractSdkErrorMessage(response.error, "Failed to rewind session")); - } - setSession(toTextChatSession(response.data)); - } catch (error) { - toast.error(getErrorMessage(error)); - } finally { - setRewindingTurnId(null); + setActiveTurnAction(null); } }, [disabled, session, workflowId]); + const rewindTurn = useCallback(async (turn: TextChatTurn) => { + if (!turn.user_message) return; + await submitMessage(turn.user_message.text, { turnId: turn.id, type: "rewind" }); + }, [submitMessage]); + + const startEditingTurn = useCallback((turn: TextChatTurn) => { + if (!turn.user_message) return; + const nextText = turn.user_message.text; + + setEditingTurnId(turn.id); + setDraft(nextText); + requestAnimationFrame(() => { + const textarea = document.getElementById(composerId) as HTMLTextAreaElement | null; + textarea?.focus(); + textarea?.setSelectionRange(nextText.length, nextText.length); + }); + }, [composerId]); + + const cancelEditingTurn = useCallback(() => { + setEditingTurnId(null); + setDraft(""); + }, []); + + const submitComposer = useCallback(async () => { + if (editingTurnId) { + await submitMessage(draft, { turnId: editingTurnId, type: "edit" }); + return; + } + await submitMessage(draft); + }, [draft, editingTurnId, submitMessage]); + + useEffect(() => { + if (!editingTurnId) { + return; + } + if (!turns.some((turn) => turn.id === editingTurnId)) { + setEditingTurnId(null); + setDraft(""); + } + }, [editingTurnId, turns]); + useEffect(() => { scrollEndRef.current?.scrollIntoView({ behavior: "smooth" }); - }, [turns.length, sendingMessage]); + }, [session?.revision, sendingMessage, turns.length]); const inputDisabled = disabled || !session; @@ -578,38 +623,60 @@ function ManualTextChat({
{turns.map((turn) => { const toolEvents = extractToolEvents(turn.events); + const rewindingThisTurn = activeTurnAction?.turnId === turn.id && activeTurnAction.type === "rewind"; + const rerunningEditedTurn = activeTurnAction?.turnId === turn.id && activeTurnAction.type === "edit"; return ( -
- {turn.user_message ? ( - - ) : null} - {toolEvents.map((event, index) => ( - - ))} - {turn.assistant_message ? ( - - ) : turn.status === "failed" ? ( - - ) : null} -
- +
+ {turn.user_message ? ( +
+ +
+ + +
+
+ ) : null} + {toolEvents.map((event, index) => ( + + ))} + {turn.assistant_message ? ( + + ) : turn.status === "failed" ? ( + + ) : null}
-
); })} {sendingMessage ? : null} @@ -619,11 +686,25 @@ function ManualTextChat({
+ {editingTurn ? ( +
+ Edit the selected user message, then press Enter to rerun from that point. + +
+ ) : null}