Merge pull request Stack-Cairn#239 from Stack-Cairn/fix/edit-resend-message-ref

fix(chat): propagate user message refs so edit-resend can anchor remotely

Author: su-fen

crates/agent-gateway/internal/session/conversation_ingress.go

@@ -55,14 +55,23 @@ func (m *Manager) ingestChatEvent(agentID, requestID string, event *gatewayv2.Ch
 	if event.GetType() == gatewayv2.ChatEvent_USER_MESSAGE {
 		if record := s.runs[agentScopedKey(agentID, runID)]; record != nil && record.userMessageSeeded {
 			messageID, _ := payload["message_id"].(string)
+			if strings.TrimSpace(messageID) == "" {
+				// The full stable ref (edit-resend rebase anchoring) carries
+				// the same id; either field proves the echo has new identity.
+				if ref, ok := payload["message_ref"].(map[string]any); ok {
+					refMessageID, _ := ref["message_id"].(string)
+					messageID = refMessageID
+				}
+			}
 			if strings.TrimSpace(messageID) == "" || record.userMessageIdentityForwarded {
 				// The gateway already appended this run's user_message at accept
 				// time. A plain or replayed agent echo adds no new identity.
 				return
 			}
 			// Forward one authoritative desktop echo carrying the stable message
-			// id. WebUI upserts it into the run's single user slot, so this enriches
-			// identity without creating a second bubble.
+			// identity (message_id, plus message_ref so a follow-up edit-resend
+			// can anchor its rebase). WebUI upserts it into the run's single user
+			// slot, so this enriches identity without creating a second bubble.
 			record.userMessageIdentityForwarded = true
 		}
 	}
@@ -418,6 +427,7 @@ func (s *conversationStreamStore) bindPendingRunLocked(
 		stream, pending.runID, pending.clientRequestID, pending.seeded, now,
 	)
 	record.userMessageSeeded = seededPayloadsIncludeUserMessage(pending.seeded)
+	record.rebaseSeeded = seededPayloadsIncludeRebased(pending.seeded)
 	s.updateChatCommandDedupeLocked(
 		pending.agentID,
 		pending.clientRequestID,

crates/agent-gateway/internal/session/conversation_stream.go

@@ -186,9 +186,11 @@ type chatRunRecord struct {
 	// queuedInGUI marks commands the desktop app parked in its prompt queue;
 	// the startup watchdog must leave them alone.
 	queuedInGUI bool
-	// rebaseSeeded marks runs whose rebased event was already appended from
-	// the agent's ref-bearing user_message, so a reconnect replay of the same
-	// event cannot seed a second truncation.
+	// rebaseSeeded marks runs whose rebased event was already appended — from
+	// the agent's ref-bearing user_message (GUI-local edits) or from the
+	// gateway-seeded payloads of a webui edit_resend command — so neither a
+	// reconnect replay nor the identity-forwarded desktop echo can seed a
+	// second truncation.
 	rebaseSeeded bool
 	// lostInferred marks a run whose terminal was inferred from missing
 	// liveness reports (desktop_run_lost and friends) rather than delivered by
@@ -1035,6 +1037,7 @@ func (m *Manager) StartChatCommand(
 	s.markRunQueuedLocked(stream, runID, clientRequestID, now)
 	acceptedSeq := s.appendSeededPayloadsLocked(stream, runID, clientRequestID, seededPayloads, now)
 	record.userMessageSeeded = seededPayloadsIncludeUserMessage(seededPayloads)
+	record.rebaseSeeded = seededPayloadsIncludeRebased(seededPayloads)
 	start := ChatCommandStart{
 		AgentID:        agentID,
 		RunID:          runID,
@@ -1063,6 +1066,7 @@ func (s *conversationStreamStore) flushDeferredSeedsLocked(
 	record.deferredSeeds = nil
 	s.appendSeededPayloadsLocked(stream, runID, record.clientRequestID, seeds, now)
 	record.userMessageSeeded = seededPayloadsIncludeUserMessage(seeds)
+	record.rebaseSeeded = seededPayloadsIncludeRebased(seeds)
 }
 
 func (s *conversationStreamStore) appendSeededPayloadsLocked(
@@ -1097,6 +1101,19 @@ func (s *conversationStreamStore) appendSeededPayloadsLocked(
 	return acceptedSeq
 }
 
+// seededPayloadsIncludeRebased mirrors seededPayloadsIncludeUserMessage for
+// the webui edit_resend truncation seed: marking rebaseSeeded at accept time
+// keeps the identity-forwarded desktop echo (which still carries the same
+// base_message_ref) from appending a second rebased to the log.
+func seededPayloadsIncludeRebased(seededPayloads []map[string]any) bool {
+	for _, payload := range seededPayloads {
+		if eventType, _ := payload["type"].(string); eventType == StreamEventRebased {
+			return true
+		}
+	}
+	return false
+}
+
 func seededPayloadsIncludeUserMessage(seededPayloads []map[string]any) bool {
 	for _, payload := range seededPayloads {
 		if eventType, _ := payload["type"].(string); eventType == "user_message" {

crates/agent-gateway/internal/session/conversation_stream_test.go

@@ -1063,3 +1063,140 @@ func TestWebuiEditResendEchoSeedsNoSecondRebased(t *testing.T) {
 		t.Fatalf("user_message count = %d (types %v), want 1 (echo swallowed)", got, eventTypes(sub.Events))
 	}
 }
+
+func userMessageEventWithRef(conversationID string, message string, ref map[string]any) *gatewayv2.ChatEvent {
+	payload := map[string]any{"message": message}
+	if ref != nil {
+		payload["message_ref"] = ref
+	}
+	data, _ := json.Marshal(payload)
+	return &gatewayv2.ChatEvent{
+		Type:           gatewayv2.ChatEvent_USER_MESSAGE,
+		ConversationId: conversationID,
+		Data:           string(data),
+	}
+}
+
+func testNewMessageRef() map[string]any {
+	return map[string]any{
+		"segment_index": 0,
+		"message_index": 4,
+		"segment_id":    "seg-1",
+		"message_id":    "msg-9",
+		"role":          "user",
+		"content_hash":  "hash-9",
+	}
+}
+
+// The gateway-seeded user_message cannot carry the message's persisted
+// identity (ids are minted at desktop persist time). An echo whose identity
+// arrives only through message_ref (no bare message_id) is forwarded exactly
+// once, replay-safe, so subscribers can anchor a later edit-resend rebase.
+func TestSeededUserMessageForwardsRefIdentityOnce(t *testing.T) {
+	m := NewManager()
+	m.StartChatCommand(conversationTestAgentID, "run-1", "conv-1", "", "client-1", []map[string]any{
+		{"type": "user_message", "message": "prompt"},
+	})
+	m.ingestChatControl(conversationTestAgentID, "run-1", startedControl("run-1", "conv-1"))
+	m.ingestChatEvent(conversationTestAgentID, "run-1", userMessageEventWithRef("conv-1", "prompt", testNewMessageRef()))
+	// Reconnect replay redelivers the same echo: no third bubble.
+	m.ingestChatEvent(conversationTestAgentID, "run-1", userMessageEventWithRef("conv-1", "prompt", testNewMessageRef()))
+
+	sub := m.SubscribeConversationStream(conversationTestAgentID, "conv-1", 0, "")
+	defer sub.Cleanup()
+	if got := countEventType(sub.Events, "user_message"); got != 2 {
+		t.Fatalf("user_message count = %d (types %v), want 2 (seed + one forwarded echo)", got, eventTypes(sub.Events))
+	}
+	forwarded := sub.Events[len(sub.Events)-1]
+	if forwarded.Type != "user_message" {
+		t.Fatalf("last event = %s, want forwarded user_message", forwarded.Type)
+	}
+	ref, ok := forwarded.Payload["message_ref"].(map[string]any)
+	if !ok || ref["message_id"] != "msg-9" || ref["content_hash"] != "hash-9" {
+		t.Fatalf("forwarded message_ref = %#v", forwarded.Payload["message_ref"])
+	}
+}
+
+// Echoes without usable identity (absent, null, or blank-id message_ref —
+// old desktop versions) swallow silently, exactly as before.
+func TestSeededEchoWithoutIdentitySwallowed(t *testing.T) {
+	m := NewManager()
+	m.StartChatCommand(conversationTestAgentID, "run-1", "conv-1", "", "client-1", []map[string]any{
+		{"type": "user_message", "message": "prompt"},
+	})
+	m.ingestChatControl(conversationTestAgentID, "run-1", startedControl("run-1", "conv-1"))
+	m.ingestChatEvent(conversationTestAgentID, "run-1", userMessageEventWithRef("conv-1", "prompt", nil))
+
+	nullRef, _ := json.Marshal(map[string]any{"message": "prompt", "message_ref": nil})
+	m.ingestChatEvent(conversationTestAgentID, "run-1", &gatewayv2.ChatEvent{
+		Type: gatewayv2.ChatEvent_USER_MESSAGE, ConversationId: "conv-1", Data: string(nullRef),
+	})
+	blankRef, _ := json.Marshal(map[string]any{
+		"message":     "prompt",
+		"message_ref": map[string]any{"message_id": "  ", "content_hash": "hash-9"},
+	})
+	m.ingestChatEvent(conversationTestAgentID, "run-1", &gatewayv2.ChatEvent{
+		Type: gatewayv2.ChatEvent_USER_MESSAGE, ConversationId: "conv-1", Data: string(blankRef),
+	})
+
+	sub := m.SubscribeConversationStream(conversationTestAgentID, "conv-1", 0, "")
+	defer sub.Cleanup()
+	if got := countEventType(sub.Events, "user_message"); got != 1 {
+		t.Fatalf("user_message count = %d (types %v), want 1", got, eventTypes(sub.Events))
+	}
+}
+
+// A GUI-local send is never seeded, so the ref rides inside the single
+// user_message itself.
+func TestGUILocalUserMessageKeepsInlineRef(t *testing.T) {
+	m := NewManager()
+	m.ingestChatControl(conversationTestAgentID, "run-1", startedControl("run-1", "conv-1"))
+	m.ingestChatEvent(conversationTestAgentID, "run-1", userMessageEventWithRef("conv-1", "prompt", testNewMessageRef()))
+
+	sub := m.SubscribeConversationStream(conversationTestAgentID, "conv-1", 0, "")
+	defer sub.Cleanup()
+	if got := countEventType(sub.Events, "user_message"); got != 1 {
+		t.Fatalf("user_message count = %d (types %v), want 1", got, eventTypes(sub.Events))
+	}
+	for _, event := range sub.Events {
+		if event.Type != "user_message" {
+			continue
+		}
+		ref, ok := event.Payload["message_ref"].(map[string]any)
+		if !ok || ref["message_id"] != "msg-9" {
+			t.Fatalf("user_message message_ref = %#v, want inline ref", event.Payload["message_ref"])
+		}
+	}
+}
+
+// A webui edit_resend echo that carries identity is forwarded (so the ref
+// binds), but its base_message_ref must not seed a second rebased — the
+// truncation was already seeded from the command's accept-time payloads.
+func TestWebuiEditResendIdentityEchoSeedsNoSecondRebased(t *testing.T) {
+	m := NewManager()
+	ref := testBaseMessageRef()
+	m.StartChatCommand(conversationTestAgentID, "run-1", "conv-1", "/workspace", "client-1", []map[string]any{
+		{"type": StreamEventRebased, "base_message_ref": ref, "reason": "edit_resend"},
+		{"type": "user_message", "message": "edited prompt", "base_message_ref": ref, "reason": "edit_resend"},
+	})
+	m.ingestChatControl(conversationTestAgentID, "run-1", startedControl("run-1", "conv-1"))
+	identityEcho, _ := json.Marshal(map[string]any{
+		"message":          "edited prompt",
+		"message_ref":      testNewMessageRef(),
+		"base_message_ref": ref,
+		"reason":           "edit_resend",
+	})
+	m.ingestChatEvent(conversationTestAgentID, "run-1", &gatewayv2.ChatEvent{
+		Type: gatewayv2.ChatEvent_USER_MESSAGE, ConversationId: "conv-1", Data: string(identityEcho),
+	})
+	m.ingestChatEvent(conversationTestAgentID, "run-1", tokenEvent("conv-1", "reply"))
+
+	sub := m.SubscribeConversationStream(conversationTestAgentID, "conv-1", 0, "")
+	defer sub.Cleanup()
+	if got := countEventType(sub.Events, StreamEventRebased); got != 1 {
+		t.Fatalf("rebased count = %d (types %v), want 1", got, eventTypes(sub.Events))
+	}
+	if got := countEventType(sub.Events, "user_message"); got != 2 {
+		t.Fatalf("user_message count = %d (types %v), want 2 (seed + forwarded identity echo)", got, eventTypes(sub.Events))
+	}
+}

crates/agent-gateway/web/src/lib/chat/transcript/historyAlignment.ts

@@ -335,7 +335,11 @@ function replaceAll(entries: ChatEntry[], turns: Turn[], historyTurns: HistoryTu
 // "Covered" requires the reply, not just the prompt: a settled turn whose
 // window twin is still user-only (the desktop's post-run flush races the
 // fetch) keeps rendering its streamed content (protectLaggedSettledTurns).
-function alignReplace(params: { turns: Turn[]; entries: ChatEntry[] }): AlignResult {
+function alignReplace(params: {
+  turns: Turn[];
+  entries: ChatEntry[];
+  rebaseReconcile?: boolean;
+}): AlignResult {
   let historyTurns = groupHistoryEntriesIntoTurns(params.entries);
 
   // Trim persisted echoes of active exchanges. Ref-anchored first (covers
@@ -419,7 +423,14 @@ function alignReplace(params: { turns: Turn[]; entries: ChatEntry[] }): AlignRes
     reflessSettled.length,
     Math.max(0, historyUserCount - refMatchedCount),
   );
-  const keptRefless = new Set(reflessSettled.slice(coveredRefless));
+  // Rebase reconciliation: the server authoritatively deleted a settled
+  // suffix (edit-resend truncation whose anchor this client missed), so
+  // uncovered ref-less settled turns are stale edit versions, not
+  // persistence-lagged exchanges — none survive. Genuinely lagged replies
+  // were already captured by protectLaggedSettledTurns above.
+  const keptRefless = new Set(
+    params.rebaseReconcile === true ? [] : reflessSettled.slice(coveredRefless),
+  );
 
   const turns = params.turns.flatMap((turn) => {
     const protectedTurn = protectedTurns.get(turn);
@@ -557,9 +568,18 @@ export function alignHistory(params: {
   turns: Turn[];
   entries: ChatEntry[];
   mode: HistoryApplyMode;
+  // Set after a rebased event whose truncation anchor was missing locally:
+  // the server deleted a settled suffix this client still renders, so the
+  // replace must not keep uncovered ref-less settled turns (only meaningful
+  // with mode "replace").
+  rebaseReconcile?: boolean;
 }): AlignResult {
   if (params.mode === "replace") {
-    return alignReplace({ turns: params.turns, entries: params.entries });
+    return alignReplace({
+      turns: params.turns,
+      entries: params.entries,
+      rebaseReconcile: params.rebaseReconcile,
+    });
   }
   return alignEnrich(params);
 }