Backport: stop false stale-send guard (#2115 + prerequisite #2034) to release/v0.9.x

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

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"))
+	}
+}

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)

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"

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";