From 00dd1fb20302e93e0705b202b6eae6083e480554 Mon Sep 17 00:00:00 2001 From: erauner12 <59383957+erauner12@users.noreply.github.com> Date: Thu, 25 Jun 2026 10:00:27 -0500 Subject: [PATCH 1/2] fix(ui): avoid false stale chat send guard (#2034) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes a false-positive chat send guard that can show `New messages loaded — please review before sending` during normal same-tab usage after an agent response has already completed and is visible. The send guard now compares backend-derived messages against all locally visible comparable messages instead of only persisted/stored messages. This preserves the intended cross-tab stale-message warning while allowing the next same-tab send to proceed. Fixes #2033. ## Changes - Count all locally visible comparable chat messages when deciding whether the backend has unseen messages. - Apply the same comparable-message filtering to backend-derived messages before comparing counts. - Add mounted regression tests covering: - same-tab completed stream does not block the next send - true cross-tab/unseen messages still block - cross-tab/unseen messages still block after a same-tab stream ## Validation Automated: ```bash cd ui npm test -- --runTestsByPath src/components/chat/__tests__/ChatInterface.sendGuard.test.tsx --runInBand ``` Result: ```text Test Suites: 1 passed, 1 total Tests: 3 passed, 3 total ``` Previously also run successfully during branch validation: ```text src/lib/__tests__/messageHandlers.test.ts: 27 passed Full UI Jest suite: 268 passed ``` Manual: - Deployed/tested locally against `kind-kagent`. - Verified the next same-tab message sends without the false `New messages loaded` toast after a completed agent response. --------- Signed-off-by: Evan Rauner Co-authored-by: Eitan Yarmush (cherry picked from commit e1422ab454faed444041603f68ae42f8354487bd) Signed-off-by: Jeremy Alvis --- ui/src/components/chat/ChatInterface.tsx | 29 +- .../ChatInterface.sendGuard.test.tsx | 380 ++++++++++++++++++ ui/src/lib/messageHandlers.ts | 78 ++++ 3 files changed, 475 insertions(+), 12 deletions(-) create mode 100644 ui/src/components/chat/__tests__/ChatInterface.sendGuard.test.tsx diff --git a/ui/src/components/chat/ChatInterface.tsx b/ui/src/components/chat/ChatInterface.tsx index 7783b4b38..98b828a54 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 000000000..b8efc1906 --- /dev/null +++ b/ui/src/components/chat/__tests__/ChatInterface.sendGuard.test.tsx @@ -0,0 +1,380 @@ +/** + * @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("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 15ac895e0..df0be1659 100644 --- a/ui/src/lib/messageHandlers.ts +++ b/ui/src/lib/messageHandlers.ts @@ -448,6 +448,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 From 999bb40a8dc412842af834e97f29a82b87ca8162 Mon Sep 17 00:00:00 2001 From: Jeremy Alvis <3587901+iplay88keys@users.noreply.github.com> Date: Tue, 30 Jun 2026 11:26:02 -0700 Subject: [PATCH 2/2] Stop false stale-send guard after tool-call turns (#2115) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description The chat send guard falsely blocked the next message with "New messages loaded — please review before sending" on essentially every turn after the first — sending twice was needed each time. It surfaced first on a tool-call turn, but the root cause affects every agent response. ## Root cause A2A marks per-message `contextId`/`taskId` as optional — the `Task` is the canonical carrier, so persisted history messages routinely store them as empty strings. After a turn completes the UI keeps the locally-streamed message copies (it does not reload from the backend), and those carry the task's real `contextId`/`taskId`, so the send guard keys them as `["task", contextId, taskId, contentSignature]`. The backend-extracted copies, with empty ids, fall back to `["message", ]` instead — and for converted tool messages that id is even regenerated (`uuidv4`) on every extraction. So a persisted agent message never matches its locally-streamed counterpart, the guard counts the backend as ahead of the local view, and blocks the send. Every turn has an agent response, so every send after the first tripped it. ## Fix Carry `contextId`/`taskId` on agent messages so the streamed and persisted copies key identically. Two layers: - **UI** (`ui/src/lib/messageHandlers.ts`): when flattening `task.history`, backfill `contextId`/`taskId` from the task for **every** extracted agent message (text and tool), treating `""` as absent (`||`, not `??`). This is the complete fix on its own — it repairs already-persisted sessions and works regardless of which runtime produced them. - **Runtime** (source, defense-in-depth, helps new sessions): stamp `contextId`/`taskId` onto agent messages at emission time, in both producers — Python `kagent-adk` (`convert_event_to_a2a_message`) and the Go ADK executor, now the default declarative runtime (`go/adk/pkg/a2a/executor.go`, via a `newAgentStatusEvent` seam). A2A allows omission, and remote agents are out of our control, so the UI backfill stays regardless. Scope is limited to agent messages; user messages are keyed by `messageId` and already match. ## Reproduction image (7) 1. Send a message that triggers a tool call. 2. Send another message → blocked by the stale-send banner with no actual change. --------- Signed-off-by: Jeremy Alvis (cherry picked from commit df828c0d7b5c0f804f3fd3fd8c0d8d90318dd71b) --- go/adk/pkg/a2a/executor.go | 45 +++++++++--- go/adk/pkg/a2a/executor_test.go | 67 +++++++++++++++++ .../kagent/adk/converters/event_converter.py | 25 ++++++- .../converters/test_event_converter.py | 18 +++++ .../ChatInterface.sendGuard.test.tsx | 73 +++++++++++++++++++ ui/src/lib/messageHandlers.ts | 28 +++++-- 6 files changed, 235 insertions(+), 21 deletions(-) create mode 100644 go/adk/pkg/a2a/executor_test.go diff --git a/go/adk/pkg/a2a/executor.go b/go/adk/pkg/a2a/executor.go index b9879988b..e92fc7df1 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 000000000..91af7fe16 --- /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 92c4b664d..fadd1e574 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 6040a7e04..62611648c 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/__tests__/ChatInterface.sendGuard.test.tsx b/ui/src/components/chat/__tests__/ChatInterface.sendGuard.test.tsx index b8efc1906..5a301b447 100644 --- a/ui/src/components/chat/__tests__/ChatInterface.sendGuard.test.tsx +++ b/ui/src/components/chat/__tests__/ChatInterface.sendGuard.test.tsx @@ -337,6 +337,79 @@ describe("ChatInterface send guard", () => { 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)]); diff --git a/ui/src/lib/messageHandlers.ts b/ui/src/lib/messageHandlers.ts index df0be1659..30e2dee11 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 } } : {}), + }); } } }