diff --git a/go/adk/pkg/a2a/executor.go b/go/adk/pkg/a2a/executor.go index b9879988bb..e92fc7df1c 100644 --- a/go/adk/pkg/a2a/executor.go +++ b/go/adk/pkg/a2a/executor.go @@ -100,6 +100,33 @@ func (u *userIDInterceptor) Before(ctx context.Context, callCtx *a2asrv.CallCont return ctx, nil } +// newAgentMessage builds an agent message stamped with the request's context +// and task ids. A2A allows omitting these (the task is the canonical carrier), +// but stamping them lets consumers that flatten task.history into standalone +// messages key each message to its task without backfilling. Mirrors the Python +// kagent-adk event converter. +func newAgentMessage(reqCtx *a2asrv.RequestContext, parts ...a2atype.Part) *a2atype.Message { + msg := a2atype.NewMessage(a2atype.MessageRoleAgent, parts...) + msg.ContextID = reqCtx.ContextID + msg.TaskID = reqCtx.TaskID + return msg +} + +// newAgentStatusEvent builds a working TaskStatusUpdateEvent whose agent message +// carries the given parts, the given metadata, and the request's context/task +// ids (via newAgentMessage). The message and event share the same metadata map, +// matching the executor's emission paths. This is the per-event seam where a +// streamed agent message is turned into an emitted (and persisted) event, so the +// id stamping here is what the send guard relies on when it later flattens +// task.history. +func newAgentStatusEvent(reqCtx *a2asrv.RequestContext, parts a2atype.ContentParts, meta map[string]any) *a2atype.TaskStatusUpdateEvent { + msg := newAgentMessage(reqCtx, parts...) + msg.Metadata = meta + statusEv := a2atype.NewStatusUpdateEvent(reqCtx, a2atype.TaskStateWorking, msg) + statusEv.Metadata = meta + return statusEv +} + // Execute implements a2asrv.AgentExecutor. // It follows the Python _handle_request pattern: set up session, handle HITL, // convert inbound message, run the agent loop, and emit A2A events. @@ -266,7 +293,7 @@ func (e *KAgentExecutor) Execute(ctx context.Context, reqCtx *a2asrv.RequestCont // Events with no content carry metadata only; still track invocationID/usage. // Check for LLM error. if adkEvent.ErrorCode != "" { - errMsg := a2atype.NewMessage(a2atype.MessageRoleAgent, + errMsg := newAgentMessage(reqCtx, a2atype.TextPart{Text: fmt.Sprintf("LLM error: %s %s", adkEvent.ErrorCode, adkEvent.ErrorMessage)}) failed := a2atype.NewStatusUpdateEvent(reqCtx, a2atype.TaskStateFailed, errMsg) failed.Final = true @@ -278,7 +305,7 @@ func (e *KAgentExecutor) Execute(ctx context.Context, reqCtx *a2asrv.RequestCont // Check for LLM error (even with content present). if adkEvent.ErrorCode != "" { - errMsg := a2atype.NewMessage(a2atype.MessageRoleAgent, + errMsg := newAgentMessage(reqCtx, a2atype.TextPart{Text: fmt.Sprintf("LLM error: %s %s", adkEvent.ErrorCode, adkEvent.ErrorMessage)}) failed := a2atype.NewStatusUpdateEvent(reqCtx, a2atype.TaskStateFailed, errMsg) failed.Final = true @@ -326,10 +353,7 @@ func (e *KAgentExecutor) Execute(ctx context.Context, reqCtx *a2asrv.RequestCont if len(textOnly) > 0 { mirrorMeta := maps.Clone(eventMeta) mirrorMeta[adka2a.ToA2AMetaKey("partial")] = true - msg := a2atype.NewMessage(a2atype.MessageRoleAgent, textOnly...) - msg.Metadata = mirrorMeta - statusEv := a2atype.NewStatusUpdateEvent(reqCtx, a2atype.TaskStateWorking, msg) - statusEv.Metadata = mirrorMeta + statusEv := newAgentStatusEvent(reqCtx, textOnly, mirrorMeta) if err := queue.Write(ctx, statusEv); err != nil { return fmt.Errorf("failed to write partial status event: %w", err) } @@ -338,10 +362,7 @@ func (e *KAgentExecutor) Execute(ctx context.Context, reqCtx *a2asrv.RequestCont mirrorParts := a2aParts if len(hitlParts) == 0 { // Only mirror when not accumulating HITL parts (those go into input_required). - msg := a2atype.NewMessage(a2atype.MessageRoleAgent, mirrorParts...) - msg.Metadata = maps.Clone(eventMeta) - statusEv := a2atype.NewStatusUpdateEvent(reqCtx, a2atype.TaskStateWorking, msg) - statusEv.Metadata = maps.Clone(eventMeta) + statusEv := newAgentStatusEvent(reqCtx, mirrorParts, maps.Clone(eventMeta)) if err := queue.Write(ctx, statusEv); err != nil { return fmt.Errorf("failed to write mirror status event: %w", err) } @@ -362,7 +383,7 @@ func (e *KAgentExecutor) Execute(ctx context.Context, reqCtx *a2asrv.RequestCont } if runErr != nil { - errMsg := a2atype.NewMessage(a2atype.MessageRoleAgent, a2atype.TextPart{Text: runErr.Error()}) + errMsg := newAgentMessage(reqCtx, a2atype.TextPart{Text: runErr.Error()}) failed := a2atype.NewStatusUpdateEvent(reqCtx, a2atype.TaskStateFailed, errMsg) failed.Final = true failed.Metadata = finalMeta @@ -371,7 +392,7 @@ func (e *KAgentExecutor) Execute(ctx context.Context, reqCtx *a2asrv.RequestCont if len(hitlParts) > 0 { // input_required: the agent is waiting for HITL decisions. - hitlMsg := a2atype.NewMessage(a2atype.MessageRoleAgent, hitlParts...) + hitlMsg := newAgentMessage(reqCtx, hitlParts...) inputRequired := a2atype.NewStatusUpdateEvent(reqCtx, a2atype.TaskStateInputRequired, hitlMsg) inputRequired.Final = true inputRequired.Metadata = finalMeta diff --git a/go/adk/pkg/a2a/executor_test.go b/go/adk/pkg/a2a/executor_test.go new file mode 100644 index 0000000000..91af7fe160 --- /dev/null +++ b/go/adk/pkg/a2a/executor_test.go @@ -0,0 +1,67 @@ +package a2a + +import ( + "testing" + + a2atype "github.com/a2aproject/a2a-go/a2a" + "github.com/a2aproject/a2a-go/a2asrv" +) + +// TestNewAgentMessage_StampsContextAndTaskID verifies agent messages carry the +// request's context and task ids. A2A allows omitting them (the task is the +// canonical carrier), but stamping them lets consumers that flatten task.history +// key each message to its task without backfilling. +func TestNewAgentMessage_StampsContextAndTaskID(t *testing.T) { + reqCtx := &a2asrv.RequestContext{ + ContextID: "ctx-xyz", + TaskID: a2atype.TaskID("task-xyz"), + } + + msg := newAgentMessage(reqCtx, a2atype.TextPart{Text: "hello"}) + + if msg.ContextID != "ctx-xyz" { + t.Errorf("ContextID = %q, want %q", msg.ContextID, "ctx-xyz") + } + if msg.TaskID != a2atype.TaskID("task-xyz") { + t.Errorf("TaskID = %q, want %q", msg.TaskID, a2atype.TaskID("task-xyz")) + } + if msg.Role != a2atype.MessageRoleAgent { + t.Errorf("Role = %q, want %q", msg.Role, a2atype.MessageRoleAgent) + } +} + +// TestNewAgentStatusEvent_MessageCarriesIDs verifies the per-event emission seam: +// the working status event the executor writes (and that is persisted into +// task.history) carries an agent message stamped with the request's context/task +// ids. This is the property the send guard depends on — without it the persisted +// message keys differently from its locally-streamed counterpart and falsely +// blocks the next send. Mirrors the Python converter test. +func TestNewAgentStatusEvent_MessageCarriesIDs(t *testing.T) { + reqCtx := &a2asrv.RequestContext{ + ContextID: "ctx-xyz", + TaskID: a2atype.TaskID("task-xyz"), + } + meta := map[string]any{"k": "v"} + + ev := newAgentStatusEvent(reqCtx, a2atype.ContentParts{a2atype.TextPart{Text: "hi"}}, meta) + + if ev.Status.State != a2atype.TaskStateWorking { + t.Errorf("State = %q, want %q", ev.Status.State, a2atype.TaskStateWorking) + } + if ev.Status.Message == nil { + t.Fatal("status message is nil") + } + if ev.Status.Message.ContextID != "ctx-xyz" { + t.Errorf("message ContextID = %q, want %q", ev.Status.Message.ContextID, "ctx-xyz") + } + if ev.Status.Message.TaskID != a2atype.TaskID("task-xyz") { + t.Errorf("message TaskID = %q, want %q", ev.Status.Message.TaskID, a2atype.TaskID("task-xyz")) + } + // The event itself also carries the ids (from reqCtx), matching the message. + if ev.ContextID != "ctx-xyz" { + t.Errorf("event ContextID = %q, want %q", ev.ContextID, "ctx-xyz") + } + if ev.TaskID != a2atype.TaskID("task-xyz") { + t.Errorf("event TaskID = %q, want %q", ev.TaskID, a2atype.TaskID("task-xyz")) + } +} diff --git a/python/packages/kagent-adk/src/kagent/adk/converters/event_converter.py b/python/packages/kagent-adk/src/kagent/adk/converters/event_converter.py index 92c4b664d6..fadd1e5740 100644 --- a/python/packages/kagent-adk/src/kagent/adk/converters/event_converter.py +++ b/python/packages/kagent-adk/src/kagent/adk/converters/event_converter.py @@ -161,6 +161,8 @@ def convert_event_to_a2a_message( invocation_context: InvocationContext, role: Role = Role.agent, subagent_session_ids: Optional[Dict[str, str]] = None, + task_id: Optional[str] = None, + context_id: Optional[str] = None, ) -> Optional[Message]: """Converts an ADK event to an A2A message. @@ -171,6 +173,12 @@ def convert_event_to_a2a_message( subagent_session_ids: Optional mapping of tool name to pre-generated subagent session ID. When provided, function_call DataParts for matching tools will have the session ID stamped into their metadata. + task_id: Optional task ID stamped onto the message so it carries the + same identity as the enclosing task. A2A allows these to be omitted + (the task is the canonical carrier), but stamping them lets consumers + that flatten task.history into standalone messages key each message to + its task without backfilling. + context_id: Optional context ID stamped onto the message, as task_id. Returns: An A2A Message if the event has content, None otherwise. @@ -198,7 +206,14 @@ def convert_event_to_a2a_message( if a2a_parts: message_metadata = _get_context_metadata(event, invocation_context) - return Message(message_id=str(uuid.uuid4()), role=role, parts=a2a_parts, metadata=message_metadata) + return Message( + message_id=str(uuid.uuid4()), + role=role, + parts=a2a_parts, + metadata=message_metadata, + task_id=task_id, + context_id=context_id, + ) except Exception as e: logger.error("Failed to convert event to status message: %s", e) @@ -342,7 +357,13 @@ def convert_event_to_a2a_events( a2a_events.append(error_event) # Handle regular message content - message = convert_event_to_a2a_message(event, invocation_context, subagent_session_ids=subagent_session_ids) + message = convert_event_to_a2a_message( + event, + invocation_context, + subagent_session_ids=subagent_session_ids, + task_id=task_id, + context_id=context_id, + ) if message: running_event = _create_status_update_event(message, invocation_context, event, task_id, context_id) a2a_events.append(running_event) diff --git a/python/packages/kagent-adk/tests/unittests/converters/test_event_converter.py b/python/packages/kagent-adk/tests/unittests/converters/test_event_converter.py index 6040a7e047..62611648c8 100644 --- a/python/packages/kagent-adk/tests/unittests/converters/test_event_converter.py +++ b/python/packages/kagent-adk/tests/unittests/converters/test_event_converter.py @@ -117,3 +117,21 @@ def test_convert_event_to_a2a_events(self): error_code_key = get_kagent_metadata_key("error_code") assert error_code_key in error_event.metadata assert error_event.metadata[error_code_key] == str(genai_types.FinishReason.MALFORMED_FUNCTION_CALL) + + def test_message_carries_task_and_context_ids(self): + """The converted message stamps task_id/context_id so consumers that + flatten task.history can key it to its task without backfilling.""" + invocation_context = _create_mock_invocation_context() + content = genai_types.Content(parts=[genai_types.Part(text="hello world")]) + event = _create_mock_event(content=content, invocation_id="test_invocation_ids") + + result = convert_event_to_a2a_events(event, invocation_context, task_id="task-xyz", context_id="ctx-xyz") + + working_events = [ + e for e in result if isinstance(e, TaskStatusUpdateEvent) and e.status.state == TaskState.working + ] + assert len(working_events) == 1 + message = working_events[0].status.message + assert message is not None + assert message.task_id == "task-xyz" + assert message.context_id == "ctx-xyz" diff --git a/ui/src/components/chat/ChatInterface.tsx b/ui/src/components/chat/ChatInterface.tsx index 7783b4b384..98b828a54d 100644 --- a/ui/src/components/chat/ChatInterface.tsx +++ b/ui/src/components/chat/ChatInterface.tsx @@ -26,7 +26,7 @@ import { getUiRuntimeConfig } from "@/app/actions/config"; import { DEFAULT_STREAM_TIMEOUT_MS } from "@/lib/constants"; import { toast } from "sonner"; import { useRouter } from "next/navigation"; -import { createMessageHandlers, extractMessagesFromTasks, extractApprovalMessagesFromTasks, extractTokenStatsFromTasks, createMessage, ADKMetadata, ProcessedToolCallData } from "@/lib/messageHandlers"; +import { createMessageHandlers, extractMessagesFromTasks, extractApprovalMessagesFromTasks, extractTokenStatsFromTasks, createMessage, countSendGuardComparableMessages, countBackendBackedComparableMessages, ADKMetadata, ProcessedToolCallData } from "@/lib/messageHandlers"; import { kagentA2AClient } from "@/lib/a2aClient"; import { formatA2AClientError } from "@/lib/a2aErrors"; import { useChatRunInSandbox, useChatSubstrateSandbox } from "@/components/chat/ChatAgentContext"; @@ -39,6 +39,7 @@ const RESUBSCRIBE_TASK_STATES: TaskState[] = ["submitted", "working"]; // Task states that mean the session is busy (used by the cross-tab send guard). const ACTIVE_TASK_STATES: TaskState[] = ["submitted", "working", "input-required"]; + interface ChatInterfaceProps { selectedAgentName: string; selectedNamespace: string; @@ -243,14 +244,11 @@ export default function ChatInterface({ selectedAgentName, selectedNamespace, se // the full context before their next message goes out. const guardSessionId = session?.id || sessionId; if (guardSessionId) { - // Compare only non-approval messages to avoid false negatives when - // storedMessages includes appended ToolApprovalRequest / AskUserRequest entries. - const localMessageCount = storedMessages.filter(m => { - const meta = m.metadata as ADKMetadata | undefined; - return meta?.originalType !== "ToolApprovalRequest" && meta?.originalType !== "AskUserRequest"; - }).length; + // Compare visible messages against the backend snapshot. Completed + // same-tab stream messages remain in streamingMessages until the next + // send promotes them, but only backend-backed matches count as local. const guardResult = await checkAndSyncSessionBeforeAction(guardSessionId, { - localMessageCount, + localMessages: allMessages, messages: { inFlight: "This session is already being processed — reconnecting to live updates", inputRequired: "Session is awaiting your input — please review before sending", @@ -270,15 +268,18 @@ export default function ChatInterface({ selectedAgentName, selectedNamespace, se pendingRejectionReasonsRef.current = {}; pendingTurnStatsRef.current = undefined; + const messageId = uuidv4(); + // For new sessions or when no stored messages exist, show the user message immediately const userMessage: Message = { kind: "message", - messageId: uuidv4(), + messageId, role: "user", parts: [{ kind: "text", text: userMessageText }], + contextId: guardSessionId, metadata: { timestamp: Date.now() } @@ -378,7 +379,6 @@ export default function ChatInterface({ selectedAgentName, selectedNamespace, se } } - const messageId = uuidv4(); const a2aMessage = createMessage(userMessageText, "user", { messageId, contextId: currentSessionId, @@ -597,6 +597,7 @@ export default function ChatInterface({ selectedAgentName, selectedNamespace, se guardSessionId: string, opts: { expectedTaskId?: string; + localMessages?: Message[]; localMessageCount?: number; messages: { inFlight: string; @@ -648,9 +649,13 @@ export default function ChatInterface({ selectedAgentName, selectedNamespace, se return "blocked"; } - if (opts.localMessageCount !== undefined) { + if (opts.localMessages !== undefined || opts.localMessageCount !== undefined) { const dbMessages = extractMessagesFromTasks(tasksCheck.data); - if (dbMessages.length > opts.localMessageCount) { + const dbMessageCount = countSendGuardComparableMessages(dbMessages); + const localMessageCount = opts.localMessages !== undefined + ? countBackendBackedComparableMessages(opts.localMessages, dbMessages) + : opts.localMessageCount; + if (localMessageCount !== undefined && dbMessageCount > localMessageCount) { await reloadSessionFromDB(); toast.info(opts.messages.staleOrChanged); return "blocked"; diff --git a/ui/src/components/chat/__tests__/ChatInterface.sendGuard.test.tsx b/ui/src/components/chat/__tests__/ChatInterface.sendGuard.test.tsx new file mode 100644 index 0000000000..5a301b447c --- /dev/null +++ b/ui/src/components/chat/__tests__/ChatInterface.sendGuard.test.tsx @@ -0,0 +1,453 @@ +/** + * @jest-environment jsdom + */ +import { fireEvent, render, screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import type { DataPart, Message, Task, TaskStatusUpdateEvent } from "@a2a-js/sdk"; +import { checkSessionExists, createSession, getSessionTasks } from "@/app/actions/sessions"; +import { kagentA2AClient } from "@/lib/a2aClient"; +import { toast } from "sonner"; +import ChatInterface from "@/components/chat/ChatInterface"; +import type { Session } from "@/types"; + +jest.mock("@/app/actions/sessions", () => ({ + checkSessionExists: jest.fn(), + createSession: jest.fn(), + getSessionTasks: jest.fn(), +})); + +jest.mock("@/app/actions/agents", () => ({ + getAgentWithResolvedKind: jest.fn(), + waitForSandboxAgentReady: jest.fn(), +})); + +jest.mock("@/lib/a2aClient", () => ({ + kagentA2AClient: { + sendMessageStream: jest.fn(), + resubscribeStream: jest.fn(), + }, +})); + +jest.mock("sonner", () => ({ + toast: { + info: jest.fn(), + error: jest.fn(), + loading: jest.fn(), + dismiss: jest.fn(), + }, +})); + +jest.mock("@/hooks/useSpeechRecognition", () => ({ + useSpeechRecognition: () => ({ + isListening: false, + isSupported: false, + startListening: jest.fn(), + stopListening: jest.fn(), + error: null, + }), +})); + +jest.mock("@/components/chat/ChatAgentContext", () => ({ + useChatRunInSandbox: () => false, + useChatSubstrateSandbox: () => false, +})); + +jest.mock("@/components/chat/ChatMessage", () => ({ + __esModule: true, + default: ({ message }: { message: Message }) => ( +
+ {message.parts + ?.map((part) => part.kind === "text" ? part.text : JSON.stringify(part)) + .join("")} +
+ ), +})); + +jest.mock("@/components/chat/StreamingMessage", () => ({ + __esModule: true, + default: ({ content }: { content: string }) =>
{content}
, +})); + +const mockCheckSessionExists = checkSessionExists as jest.MockedFunction; +const mockCreateSession = createSession as jest.MockedFunction; +const mockGetSessionTasks = getSessionTasks as jest.MockedFunction; +const mockSendMessageStream = kagentA2AClient.sendMessageStream as jest.MockedFunction; +const mockToastInfo = toast.info as jest.MockedFunction; + +const staleToastMessage = "New messages loaded — please review before sending"; + +function mockBackendTasks(tasks: Task[]) { + mockGetSessionTasks.mockResolvedValue({ data: tasks }); +} + +function textMessage(messageId: string, role: "user" | "agent", text: string, contextId = "session-1", taskId = "task-1"): Message { + return { + kind: "message", + messageId, + role, + contextId, + taskId, + parts: [{ kind: "text", text }], + metadata: { timestamp: Date.now() }, + } as Message; +} + +function completedTask(taskId: string, history: Message[], contextId = "session-1"): Task { + return { + id: taskId, + contextId, + status: { + state: "completed", + timestamp: new Date().toISOString(), + }, + history, + } as Task; +} + +function completedStatusEvent(text: string, contextId = "session-1", taskId = "task-streamed"): TaskStatusUpdateEvent { + return { + kind: "status-update", + contextId, + taskId, + final: true, + status: { + state: "completed", + timestamp: new Date().toISOString(), + message: textMessage(`assistant-${taskId}`, "agent", text, contextId, taskId), + }, + } as TaskStatusUpdateEvent; +} + +function toolCallMessage(messageId: string, contextId = "session-1", taskId = "task-tool", callId = "shared-call"): Message { + return { + kind: "message", + messageId, + role: "agent", + contextId, + taskId, + parts: [ + { + kind: "data", + data: { id: callId, name: "kubectl_get_pods", args: { namespace: "default" } }, + metadata: { adk_type: "function_call" }, + } as DataPart, + ], + metadata: { timestamp: Date.now() }, + } as Message; +} + +function completedToolCallStatusEvent(contextId = "session-1", taskId = "task-streamed", callId = "shared-call"): TaskStatusUpdateEvent { + return { + kind: "status-update", + contextId, + taskId, + final: true, + status: { + state: "completed", + timestamp: new Date().toISOString(), + message: toolCallMessage(`tool-${taskId}`, contextId, taskId, callId), + }, + } as TaskStatusUpdateEvent; +} + +async function* streamOf(...events: unknown[]): AsyncIterable { + for (const event of events) { + yield event; + } +} + +function sessionFixture(overrides: Partial = {}): Session { + return { + id: "session-1", + name: "Existing chat", + agent_id: "kagent__NS__test-agent", + user_id: "user-1", + created_at: "2026-03-07T10:00:00Z", + updated_at: "2026-03-07T10:05:00Z", + deleted_at: "", + ...overrides, + }; +} + +function renderExistingSession() { + return render( + , + ); +} + +async function sendText(text: string) { + const user = userEvent.setup(); + const textbox = screen.getByRole("textbox"); + await waitFor(() => expect(textbox).not.toBeDisabled()); + await user.clear(textbox); + await user.type(textbox, text); + await user.click(screen.getByRole("button", { name: /send/i })); +} + +function sentMessage(callIndex = 0): Message { + return (mockSendMessageStream.mock.calls[callIndex][2] as { message: Message }).message; +} + +describe("ChatInterface send guard", () => { + const initialTurn = [ + textMessage("initial-user", "user", "initial user", "session-1", "task-initial"), + textMessage("initial-agent", "agent", "initial answer", "session-1", "task-initial"), + ]; + const sameTabTurn = [ + textMessage("same-tab-user", "user", "same tab question", "session-1", "task-streamed"), + textMessage("same-tab-agent", "agent", "same tab answer", "session-1", "task-streamed"), + ]; + const externalTurn = [ + textMessage("external-user", "user", "external user", "session-1", "task-external"), + textMessage("external-agent", "agent", "external answer", "session-1", "task-external"), + ]; + + beforeEach(() => { + jest.clearAllMocks(); + mockCheckSessionExists.mockResolvedValue({ data: true }); + mockCreateSession.mockResolvedValue({ error: "unexpected createSession call" }); + }); + + it("does not block the next send when completed same-tab stream messages are already visible", async () => { + mockBackendTasks([completedTask("task-initial", initialTurn)]); + mockSendMessageStream + .mockResolvedValueOnce(streamOf(completedStatusEvent("same tab answer"))) + .mockResolvedValueOnce(streamOf(completedStatusEvent("next answer", "session-1", "task-next"))); + + renderExistingSession(); + + expect(await screen.findByText("initial answer")).toBeInTheDocument(); + + await sendText("same tab question"); + await waitFor(() => expect(mockSendMessageStream).toHaveBeenCalledTimes(1)); + expect(await screen.findByText("same tab answer")).toBeInTheDocument(); + + mockBackendTasks([ + completedTask("task-initial", initialTurn), + completedTask("task-streamed", [ + textMessage(sentMessage().messageId, "user", "same tab question", "session-1", "task-streamed"), + textMessage("same-tab-agent", "agent", "same tab answer", "session-1", "task-streamed"), + ]), + ]); + await sendText("next question"); + + await waitFor(() => expect(mockSendMessageStream).toHaveBeenCalledTimes(2)); + expect(mockToastInfo).not.toHaveBeenCalledWith(staleToastMessage); + }); + + it("still blocks after a same-tab stream when the backend also has an unseen cross-tab message", async () => { + mockBackendTasks([completedTask("task-initial", initialTurn)]); + mockSendMessageStream + .mockResolvedValueOnce(streamOf(completedStatusEvent("same tab answer"))); + + renderExistingSession(); + + expect(await screen.findByText("initial answer")).toBeInTheDocument(); + + await sendText("same tab question"); + await waitFor(() => expect(mockSendMessageStream).toHaveBeenCalledTimes(1)); + expect(await screen.findByText("same tab answer")).toBeInTheDocument(); + + mockBackendTasks([ + completedTask("task-initial", initialTurn), + completedTask("task-streamed", sameTabTurn), + completedTask("task-external", externalTurn), + ]); + await sendText("should review cross-tab first"); + + await waitFor(() => expect(mockToastInfo).toHaveBeenCalledWith(staleToastMessage)); + expect(mockSendMessageStream).toHaveBeenCalledTimes(1); + }); + + it("does not let local-only streaming messages mask an unseen backend turn", async () => { + mockBackendTasks([completedTask("task-initial", initialTurn)]); + mockSendMessageStream + .mockResolvedValueOnce(streamOf(completedStatusEvent("local-only answer"))); + + renderExistingSession(); + + expect(await screen.findByText("initial answer")).toBeInTheDocument(); + + await sendText("local optimistic question"); + await waitFor(() => expect(mockSendMessageStream).toHaveBeenCalledTimes(1)); + expect(await screen.findByText("local-only answer")).toBeInTheDocument(); + + mockBackendTasks([ + completedTask("task-initial", initialTurn), + completedTask("task-external", externalTurn), + ]); + await sendText("should review backend first"); + + await waitFor(() => expect(mockToastInfo).toHaveBeenCalledWith(staleToastMessage)); + expect(mockSendMessageStream).toHaveBeenCalledTimes(1); + }); + + it("blocks when local-only streaming messages share text and role with an unseen backend turn", async () => { + mockBackendTasks([completedTask("task-initial", initialTurn)]); + mockSendMessageStream + .mockResolvedValueOnce(streamOf(completedStatusEvent("duplicate answer", "session-1", "task-streamed"))); + + renderExistingSession(); + + expect(await screen.findByText("initial answer")).toBeInTheDocument(); + + await sendText("duplicate question"); + await waitFor(() => expect(mockSendMessageStream).toHaveBeenCalledTimes(1)); + expect(await screen.findByText("duplicate answer")).toBeInTheDocument(); + + mockBackendTasks([ + completedTask("task-initial", initialTurn), + completedTask("task-unseen", [ + textMessage("unseen-user", "user", "duplicate question", "session-1", "task-unseen"), + textMessage("unseen-agent", "agent", "duplicate answer", "session-1", "task-unseen"), + ]), + ]); + await sendText("should review duplicate backend first"); + + await waitFor(() => expect(mockToastInfo).toHaveBeenCalledWith(staleToastMessage)); + expect(mockSendMessageStream).toHaveBeenCalledTimes(1); + }); + + it("blocks when a local-only empty tool stream collides with an unseen backend data turn", async () => { + mockBackendTasks([completedTask("task-initial", initialTurn)]); + mockSendMessageStream + .mockResolvedValueOnce(streamOf(completedToolCallStatusEvent("session-1", "task-streamed", "shared-call"))); + + renderExistingSession(); + + expect(await screen.findByText("initial answer")).toBeInTheDocument(); + + await sendText("local tool question"); + await waitFor(() => expect(mockSendMessageStream).toHaveBeenCalledTimes(1)); + + mockBackendTasks([ + completedTask("task-initial", initialTurn), + completedTask("task-unseen-tool", [ + toolCallMessage("backend-tool", "session-1", "task-unseen-tool", "shared-call"), + ]), + ]); + await sendText("should review tool turn first"); + + await waitFor(() => expect(mockToastInfo).toHaveBeenCalledWith(staleToastMessage)); + expect(mockSendMessageStream).toHaveBeenCalledTimes(1); + }); + + it("does not block the next send after a same-tab tool-call turn is persisted", async () => { + // Reproduces the false positive seen with tool calls: a same-tab turn that + // makes a tool call streams a ToolCallRequestEvent locally (keyed by its + // real contextId/taskId), but the backend persists agent messages with + // EMPTY contextId/taskId. extractMessagesFromTasks leaves those empty + // (`"" ?? task.contextId` is still "") and rebuilds the converted tool + // message with a fresh uuidv4 messageId, so it keys on ["message", ] + // — a different key on every extraction. The guard therefore can never match + // the persisted tool message against the local one, counts the backend as + // ahead, and falsely blocks the next send. + mockBackendTasks([completedTask("task-initial", initialTurn)]); + mockSendMessageStream + .mockResolvedValueOnce(streamOf(completedToolCallStatusEvent("session-1", "task-streamed", "shared-call"))) + .mockResolvedValueOnce(streamOf(completedStatusEvent("next answer", "session-1", "task-next"))); + + renderExistingSession(); + + expect(await screen.findByText("initial answer")).toBeInTheDocument(); + + await sendText("tool question"); + await waitFor(() => expect(mockSendMessageStream).toHaveBeenCalledTimes(1)); + + // Backend now reflects exactly that same-tab tool turn. The persisted agent + // tool message carries empty contextId/taskId, as the real backend stores it. + mockBackendTasks([ + completedTask("task-initial", initialTurn), + completedTask("task-streamed", [ + textMessage(sentMessage().messageId, "user", "tool question", "session-1", "task-streamed"), + toolCallMessage("backend-tool", "", "", "shared-call"), + ]), + ]); + await sendText("next question"); + + await waitFor(() => expect(mockSendMessageStream).toHaveBeenCalledTimes(2)); + expect(mockToastInfo).not.toHaveBeenCalledWith(staleToastMessage); + }); + + it("does not block the next send after a same-tab text turn persisted with empty agent ids", async () => { + // The general case behind "every message needs sending twice": the backend + // persists AGENT messages with empty contextId/taskId (A2A optional fields; + // the task is the canonical carrier). The locally-streamed agent message + // carries the task's real ids, so it keys on ["task", ...] while the + // backend-extracted copy — pushed as-is with empty ids — keys on + // ["message", ]. They never match, so every turn (each has an + // agent text response) counts the backend as ahead and falsely blocks. + mockBackendTasks([completedTask("task-initial", initialTurn)]); + mockSendMessageStream + .mockResolvedValueOnce(streamOf(completedStatusEvent("same tab answer", "session-1", "task-streamed"))) + .mockResolvedValueOnce(streamOf(completedStatusEvent("next answer", "session-1", "task-next"))); + + renderExistingSession(); + + expect(await screen.findByText("initial answer")).toBeInTheDocument(); + + await sendText("same tab question"); + await waitFor(() => expect(mockSendMessageStream).toHaveBeenCalledTimes(1)); + expect(await screen.findByText("same tab answer")).toBeInTheDocument(); + + // Backend reflects the turn, with the agent message persisted with empty + // contextId/taskId (the user message keeps its contextId, as the client stamps it). + mockBackendTasks([ + completedTask("task-initial", initialTurn), + completedTask("task-streamed", [ + textMessage(sentMessage().messageId, "user", "same tab question", "session-1", ""), + textMessage("same-tab-agent", "agent", "same tab answer", "", ""), + ]), + ]); + await sendText("next question"); + + await waitFor(() => expect(mockSendMessageStream).toHaveBeenCalledTimes(2)); + expect(mockToastInfo).not.toHaveBeenCalledWith(staleToastMessage); + }); + + it("still blocks when the backend has a cross-tab message not visible locally", async () => { + mockBackendTasks([completedTask("task-initial", initialTurn)]); + + renderExistingSession(); + + expect(await screen.findByText("initial answer")).toBeInTheDocument(); + + mockBackendTasks([ + completedTask("task-initial", initialTurn), + completedTask("task-external", externalTurn), + ]); + await sendText("should review first"); + + await waitFor(() => expect(mockToastInfo).toHaveBeenCalledWith(staleToastMessage)); + expect(mockSendMessageStream).not.toHaveBeenCalled(); + }); + + it.each([ + ["Cmd+Enter", { metaKey: true }], + ["Ctrl+Enter", { ctrlKey: true }], + ])("applies the stale-message send guard for %s", async (_shortcut, modifier) => { + const user = userEvent.setup(); + mockBackendTasks([completedTask("task-initial", initialTurn)]); + + renderExistingSession(); + + expect(await screen.findByText("initial answer")).toBeInTheDocument(); + + mockBackendTasks([ + completedTask("task-initial", initialTurn), + completedTask("task-external", externalTurn), + ]); + + const textbox = screen.getByRole("textbox"); + await user.type(textbox, "should review first"); + fireEvent.keyDown(textbox, { key: "Enter", code: "Enter", ...modifier }); + + await waitFor(() => expect(mockToastInfo).toHaveBeenCalledWith(staleToastMessage)); + expect(mockSendMessageStream).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/src/lib/messageHandlers.ts b/ui/src/lib/messageHandlers.ts index 15ac895e0d..30e2dee116 100644 --- a/ui/src/lib/messageHandlers.ts +++ b/ui/src/lib/messageHandlers.ts @@ -60,8 +60,18 @@ export function extractMessagesFromTasks(tasks: Task[]): Message[] { // Agent messages: convert function_call / function_response DataParts to // the same ToolCallRequestEvent / ToolCallExecutionEvent format produced // by the live-stream handlers so the rendering component can display them. - const msgContextId = historyItem.contextId ?? task.contextId; - const msgTaskId = historyItem.taskId ?? task.id; + // + // Backfill contextId/taskId from the task when the history item omits them. + // Persisted agent messages (both tool and text) frequently carry empty + // strings here, and `??` would keep the empty string ("" is not nullish). + // The locally-streamed copies of these messages carry the task's real ids, + // so the send guard keys them on (contextId, taskId); without the backfill + // the backend-extracted copies fall back to messageId and never match the + // local ones, counting the backend as ahead and falsely blocking the next + // send on every turn. Treat "" as absent so both copies get the task's + // stable ids. Applied to every extracted agent message below. + const msgContextId = historyItem.contextId || task.contextId; + const msgTaskId = historyItem.taskId || task.id; const source = getSourceFromMetadata(historyItem.metadata as ADKMetadata | undefined, "assistant"); const msgStats = getMessageTokenStats(historyItem.metadata as Record); @@ -149,12 +159,16 @@ export function extractMessagesFromTasks(tasks: Task[]): Message[] { } } - // Text messages (or any message without data parts): push with tokenStats. + // Text messages (or any message without data parts): push with tokenStats + // and the backfilled contextId/taskId so they key the same way the + // locally-streamed copy does. if (!hasConvertedParts) { - messages.push(msgStats - ? { ...historyItem, metadata: { ...(historyItem.metadata as object || {}), tokenStats: msgStats } } - : historyItem - ); + messages.push({ + ...historyItem, + contextId: msgContextId, + taskId: msgTaskId, + ...(msgStats ? { metadata: { ...(historyItem.metadata as object || {}), tokenStats: msgStats } } : {}), + }); } } } @@ -448,6 +462,84 @@ export interface ADKMetadata { [key: string]: unknown; // Allow for additional metadata fields } +const SEND_GUARD_EXCLUDED_ORIGINAL_TYPES = new Set([ + "ToolApprovalRequest", // HITL approval UI duplicates the pending backend task state. + "AskUserRequest", // ask_user UI duplicates the pending backend task state. + "ToolCallSummaryMessage", // UI-only marker that closes tool calls; not a backend chat turn. +]); +const SEND_GUARD_KEY_DELIMITER = "\u0000"; +const SEND_GUARD_PART_DELIMITER = "\u0001"; + +function isSendGuardComparableMessage(message: Message): boolean { + const meta = message.metadata as ADKMetadata | undefined; + return !meta?.originalType || !SEND_GUARD_EXCLUDED_ORIGINAL_TYPES.has(meta.originalType); +} + +export function countSendGuardComparableMessages(messages: Message[]): number { + return messages.filter(isSendGuardComparableMessage).length; +} + +function getSendGuardMessageContentSignature(message: Message): string { + const meta = message.metadata as ADKMetadata | undefined; + const originalType = meta?.originalType === "TextMessage" ? "" : (meta?.originalType ?? ""); + const partsSignature = message.parts + ?.map(part => part.kind === "text" ? `text:${part.text ?? ""}` : `${part.kind}:${JSON.stringify(part)}`) + .join(SEND_GUARD_PART_DELIMITER) ?? ""; + + return [ + message.role, + originalType, + JSON.stringify(meta?.toolCallData ?? null), + JSON.stringify(meta?.toolResultData ?? null), + partsSignature, + ].join(SEND_GUARD_KEY_DELIMITER); +} + +function getSendGuardMessageKey(message: Message): string | undefined { + const contentSignature = getSendGuardMessageContentSignature(message); + + if (message.role === "user" && message.messageId) { + return ["message", message.messageId].join(SEND_GUARD_KEY_DELIMITER); + } + + // Same-tab streamed agent display messages are locally re-created, so their + // messageId may differ from the backend history item. The task/context pair + // is the stable backend identity for those messages; the richer signature + // avoids counting unrelated messages in the same task. + if (message.contextId && message.taskId) { + return ["task", message.contextId, message.taskId, contentSignature].join(SEND_GUARD_KEY_DELIMITER); + } + + if (message.messageId) { + return ["message", message.messageId].join(SEND_GUARD_KEY_DELIMITER); + } + + return undefined; +} + +export function countBackendBackedComparableMessages(localMessages: Message[], backendMessages: Message[]): number { + const backendCounts = new Map(); + for (const message of backendMessages) { + if (!isSendGuardComparableMessage(message)) continue; + const key = getSendGuardMessageKey(message); + if (!key) continue; + backendCounts.set(key, (backendCounts.get(key) ?? 0) + 1); + } + + let count = 0; + for (const message of localMessages) { + if (!isSendGuardComparableMessage(message)) continue; + const key = getSendGuardMessageKey(message); + if (!key) continue; + const remaining = backendCounts.get(key) ?? 0; + if (remaining > 0) { + count += 1; + backendCounts.set(key, remaining - 1); + } + } + return count; +} + /** * Read a metadata value checking `adk_` first, then `kagent_`. * Allows interoperability with upstream ADK (adk_ prefix) while preserving