From 78ffc57a8ecb621c94a4b246d739adb5f576d0c7 Mon Sep 17 00:00:00 2001
From: viettranx <viettranx@gmail.com>
Date: Thu, 19 Mar 2026 19:01:10 +0700
Subject: [PATCH] fix(ws): propagate origin_session_key for team task announces

Team task announces rebuilt session keys via BuildScopedSessionKey(),
producing different format than WS sessions (ws:direct:uuid vs
ws-userId-ts). Now stores and forwards origin_session_key through
task metadata and dispatch, matching subagent announce pattern.

Also adds "ws" to InternalChannels to suppress spurious warnings.
---
 cmd/gateway_consumer_handlers.go       |  9 +++++++--
 internal/channels/channel.go           |  1 +
 internal/tools/team_tasks_mutations.go |  5 +++++
 internal/tools/team_tool_dispatch.go   | 11 +++++++++++
 4 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/cmd/gateway_consumer_handlers.go b/cmd/gateway_consumer_handlers.go
index 8490d1d7..ced28f18 100644
--- a/cmd/gateway_consumer_handlers.go
+++ b/cmd/gateway_consumer_handlers.go
@@ -455,8 +455,13 @@ func handleTeammateMessage(
 			origPeerKind = string(sessions.PeerDirect)
 		}
 		origLocalKey := inMeta["origin_local_key"]
-		leadSessionKey := sessions.BuildScopedSessionKey(leadAgent, origCh, sessions.PeerKind(origPeerKind), origChatID, cfg.Sessions.Scope, cfg.Sessions.DmScope, cfg.Sessions.MainKey)
-		leadSessionKey = overrideSessionKeyFromLocalKey(leadSessionKey, origLocalKey, leadAgent, origCh, origChatID, origPeerKind)
+		// Use exact origin session key if available (WS uses non-standard format).
+		leadSessionKey := inMeta["origin_session_key"]
+		if leadSessionKey == "" {
+			// Fallback: rebuild session key from origin metadata (works for Telegram, Discord, etc.)
+			leadSessionKey = sessions.BuildScopedSessionKey(leadAgent, origCh, sessions.PeerKind(origPeerKind), origChatID, cfg.Sessions.Scope, cfg.Sessions.DmScope, cfg.Sessions.MainKey)
+			leadSessionKey = overrideSessionKeyFromLocalKey(leadSessionKey, origLocalKey, leadAgent, origCh, origChatID, origPeerKind)
+		}
 
 		// Extract trace context for announce linking.
 		var parentTraceID, parentRootSpanID uuid.UUID
diff --git a/internal/channels/channel.go b/internal/channels/channel.go
index dcfb3807..5d4d26e8 100644
--- a/internal/channels/channel.go
+++ b/internal/channels/channel.go
@@ -25,6 +25,7 @@ var InternalChannels = map[string]bool{
 	"system":   true,
 	"subagent": true,
 	"browser":  true,
+	"ws":       true, // WebSocket — responses delivered via events/RPC, not outbound dispatch
 }
 
 // IsInternalChannel checks if a channel name is internal.
diff --git a/internal/tools/team_tasks_mutations.go b/internal/tools/team_tasks_mutations.go
index 5dbd7d3e..b2320e76 100644
--- a/internal/tools/team_tasks_mutations.go
+++ b/internal/tools/team_tasks_mutations.go
@@ -162,6 +162,11 @@ func (t *TeamTasksTool) executeCreate(ctx context.Context, args map[string]any)
 	if lk := ToolLocalKeyFromCtx(ctx); lk != "" {
 		taskMeta["local_key"] = lk
 	}
+	// Store origin session key so deferred dispatches route announces correctly.
+	// WS sessions use non-standard key format that BuildScopedSessionKey() cannot reproduce.
+	if sk := ToolSessionKeyFromCtx(ctx); sk != "" {
+		taskMeta["origin_session_key"] = sk
+	}
 	// Store leader's trace context so unblocked dispatch links back to the leader's trace.
 	if traceID := tracing.TraceIDFromContext(ctx); traceID != uuid.Nil {
 		taskMeta["origin_trace_id"] = traceID.String()
diff --git a/internal/tools/team_tool_dispatch.go b/internal/tools/team_tool_dispatch.go
index 110c5f10..558fd1e6 100644
--- a/internal/tools/team_tool_dispatch.go
+++ b/internal/tools/team_tool_dispatch.go
@@ -136,6 +136,17 @@ func (m *TeamToolManager) dispatchTaskToAgent(ctx context.Context, task *store.T
 	if localKey != "" {
 		meta["origin_local_key"] = localKey
 	}
+	// Resolve origin session key from context; fallback to task metadata for deferred dispatches.
+	// WS sessions use non-standard key format that BuildScopedSessionKey() cannot reproduce.
+	originSessionKey := ToolSessionKeyFromCtx(ctx)
+	if originSessionKey == "" {
+		if sk, ok := task.Metadata["origin_session_key"].(string); ok {
+			originSessionKey = sk
+		}
+	}
+	if originSessionKey != "" {
+		meta["origin_session_key"] = originSessionKey
+	}
 	// Pass the team workspace dir so member agents write files to the shared folder.
 	if ws := taskTeamWorkspace(task); ws != "" {
 		meta["team_workspace"] = ws