From 56dd8b58855335bcf193719fc7937fb30d73ccb3 Mon Sep 17 00:00:00 2001
From: viettranx <viettranx@gmail.com>
Date: Sat, 7 Mar 2026 23:41:14 +0700
Subject: [PATCH] fix(ws-media): plumb subagent media paths through announce to
 inbound message
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Root cause: SubagentTask had no Media field. Tool results with media
(e.g. image generation) captured result.Media but it was discarded.
AnnounceQueueItem also had no Media field, so PublishInbound for
subagent announces always had empty Media — ContentSuffix never triggered.

Changes:
- Add Media []string to SubagentTask (subagent.go)
- Capture result.Media from tool executions in subagent_exec.go
- Add Media []string to AnnounceQueueItem (announce_queue.go)
- Pass task.Media through direct publish and batched announce drain
- Collect batch media in gateway.go announce queue drain callback
---
 cmd/gateway.go                   | 6 ++++++
 internal/tools/announce_queue.go | 1 +
 internal/tools/subagent.go       | 1 +
 internal/tools/subagent_exec.go  | 9 +++++++++
 4 files changed, 17 insertions(+)

diff --git a/cmd/gateway.go b/cmd/gateway.go
index f337ff34..9b9d51df 100644
--- a/cmd/gateway.go
+++ b/cmd/gateway.go
@@ -245,6 +245,11 @@ func runGateway() {
 				if meta.OriginSessionKey != "" {
 					batchMeta["origin_session_key"] = meta.OriginSessionKey
 				}
+				// Collect media from all items in the batch.
+				var batchMedia []string
+				for _, item := range items {
+					batchMedia = append(batchMedia, item.Media...)
+				}
 				msgBus.PublishInbound(bus.InboundMessage{
 					Channel:  "system",
 					SenderID: senderID,
@@ -252,6 +257,7 @@ func runGateway() {
 					Content:  content,
 					UserID:   meta.OriginUserID,
 					Metadata: batchMeta,
+					Media:    batchMedia,
 				})
 			},
 			func(parentID string) int {
diff --git a/internal/tools/announce_queue.go b/internal/tools/announce_queue.go
index 5bc95704..0ec262ea 100644
--- a/internal/tools/announce_queue.go
+++ b/internal/tools/announce_queue.go
@@ -13,6 +13,7 @@ type AnnounceQueueItem struct {
 	Label      string
 	Status     string // "completed", "failed", "cancelled"
 	Result     string
+	Media      []string // media file paths from tool results
 	Runtime    time.Duration
 	Iterations int
 }
diff --git a/internal/tools/subagent.go b/internal/tools/subagent.go
index 3b691cec..1b0e707b 100644
--- a/internal/tools/subagent.go
+++ b/internal/tools/subagent.go
@@ -59,6 +59,7 @@ type SubagentTask struct {
 	OriginSessionKey string `json:"originSessionKey,omitempty"` // exact parent session key for announce routing (WS uses non-standard format)
 	CreatedAt        int64  `json:"createdAt"`
 	CompletedAt      int64  `json:"completedAt,omitempty"`
+	Media            []string  `json:"-"` // media file paths from tool results
 	OriginTraceID    uuid.UUID `json:"-"` // parent trace for announce linking
 	OriginRootSpanID uuid.UUID `json:"-"` // parent agent's root span ID
 	cancelFunc       context.CancelFunc `json:"-"` // per-task context cancel
diff --git a/internal/tools/subagent_exec.go b/internal/tools/subagent_exec.go
index 2c4eba9a..67763c1e 100644
--- a/internal/tools/subagent_exec.go
+++ b/internal/tools/subagent_exec.go
@@ -28,6 +28,7 @@ func (sm *SubagentManager) runTask(ctx context.Context, task *SubagentTask, call
 			Label:      task.Label,
 			Status:     task.Status,
 			Result:     task.Result,
+			Media:      task.Media,
 			Runtime:    elapsed,
 			Iterations: iterations,
 		}
@@ -74,6 +75,7 @@ func (sm *SubagentManager) runTask(ctx context.Context, task *SubagentTask, call
 				Content:  announceContent,
 				UserID:   task.OriginUserID,
 				Metadata: announceMeta,
+				Media:    task.Media,
 			})
 		}
 	}
@@ -163,6 +165,7 @@ func (sm *SubagentManager) executeTask(ctx context.Context, task *SubagentTask)
 	}
 
 	// Run LLM iteration loop (similar to agent loop but simplified)
+	var mediaPaths []string
 	maxIterations := 20
 
 	for iteration < maxIterations {
@@ -223,6 +226,11 @@ func (sm *SubagentManager) executeTask(ctx context.Context, task *SubagentTask)
 			argsJSON, _ := json.Marshal(tc.Arguments)
 			sm.emitToolSpan(subTraceCtx, toolStart, tc.Name, tc.ID, string(argsJSON), result.ForLLM, result.IsError)
 
+			// Capture media file paths from tool results (e.g. image generation).
+			if len(result.Media) > 0 {
+				mediaPaths = append(mediaPaths, result.Media...)
+			}
+
 			messages = append(messages, providers.Message{
 				Role:       "tool",
 				Content:    result.ForLLM,
@@ -237,6 +245,7 @@ func (sm *SubagentManager) executeTask(ctx context.Context, task *SubagentTask)
 	}
 	task.Status = TaskStatusCompleted
 	task.Result = finalContent
+	task.Media = mediaPaths
 	sm.mu.Unlock()
 
 	slog.Info("subagent completed", "id", task.ID, "iterations", iteration)