From 130d765ec7fb46a0388475482b1abbf720ea2fe3 Mon Sep 17 00:00:00 2001
From: Max <shadow.iqka@qq.com>
Date: Tue, 25 Nov 2025 11:11:43 +0800
Subject: [PATCH] Enhance stream data structures in Assistant for improved
 context and metadata handling

- Added Metadata field to StreamStartData and StreamEndData structs to facilitate passing additional context information.
- Updated sendAgentStreamStart and sendAgentStreamEnd methods to include new fields for better traceability and context management during streaming events.
- Enhanced the overall structure of stream data to support more comprehensive event tracking and context sharing.
---
 agent/assistant/agent.go   |  4 ++++
 agent/context/types_llm.go | 28 ++++++++++++++++------------
 2 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/agent/assistant/agent.go b/agent/assistant/agent.go
index ca508885..feb329ed 100644
--- a/agent/assistant/agent.go
+++ b/agent/assistant/agent.go
@@ -710,6 +710,7 @@ func (ast *Assistant) sendAgentStreamStart(ctx *context.Context, handler context
 		RequestID: ctx.RequestID(),
 		Timestamp: startTime.UnixMilli(),
 		Assistant: ast.Info(ctx.Locale),
+		Metadata:  ctx.Metadata,
 	}
 
 	if startJSON, err := jsoniter.Marshal(startData); err == nil {
@@ -731,9 +732,12 @@ func (ast *Assistant) sendAgentStreamEnd(ctx *context.Context, handler context.S
 
 	endData := &context.StreamEndData{
 		RequestID:  ctx.RequestID(),
+		ContextID:  ctx.ID,
 		Timestamp:  time.Now().UnixMilli(),
 		DurationMs: time.Since(startTime).Milliseconds(),
 		Status:     status,
+		TraceID:    ctx.TraceID(),
+		Metadata:   ctx.Metadata,
 	}
 
 	if err != nil {
diff --git a/agent/context/types_llm.go b/agent/context/types_llm.go
index 57f54cd7..918f5a67 100644
--- a/agent/context/types_llm.go
+++ b/agent/context/types_llm.go
@@ -221,23 +221,27 @@ type JSONSchema struct {
 // StreamStartData represents the data for stream_start event
 // Sent when a streaming request begins
 type StreamStartData struct {
-	ContextID string         `json:"context_id"`          // Context ID for the response
-	RequestID string         `json:"request_id"`          // Unique identifier for this request
-	Timestamp int64          `json:"timestamp"`           // Unix timestamp when stream started
-	ChatID    string         `json:"chat_id"`             // Chat ID being used (e.g., "chat-123")
-	TraceID   string         `json:"trace_id"`            // Trace ID being used (e.g., "trace-123")
-	Assistant *AssistantInfo `json:"assistant,omitempty"` // Assistant information
+	ContextID string                 `json:"context_id"`          // Context ID for the response
+	RequestID string                 `json:"request_id"`          // Unique identifier for this request
+	Timestamp int64                  `json:"timestamp"`           // Unix timestamp when stream started
+	ChatID    string                 `json:"chat_id"`             // Chat ID being used (e.g., "chat-123")
+	TraceID   string                 `json:"trace_id"`            // Trace ID being used (e.g., "trace-123")
+	Assistant *AssistantInfo         `json:"assistant,omitempty"` // Assistant information
+	Metadata  map[string]interface{} `json:"metadata,omitempty"`  // Metadata to pass to the page for CUI context
 }
 
 // StreamEndData represents the data for stream_end event
 // Sent when a streaming request completes (successfully or with error)
 type StreamEndData struct {
-	RequestID  string     `json:"request_id"`      // Corresponding request ID
-	Timestamp  int64      `json:"timestamp"`       // Unix timestamp when stream ended
-	DurationMs int64      `json:"duration_ms"`     // Total duration in milliseconds
-	Status     string     `json:"status"`          // "completed" | "error" | "cancelled"
-	Error      string     `json:"error,omitempty"` // Error message if status is "error"
-	Usage      *UsageInfo `json:"usage,omitempty"` // Token usage statistics
+	RequestID  string                 `json:"request_id"`         // Corresponding request ID
+	ContextID  string                 `json:"context_id"`         // Context ID for the response
+	TraceID    string                 `json:"trace_id"`           // Trace ID being used (e.g., "trace-123")
+	Timestamp  int64                  `json:"timestamp"`          // Unix timestamp when stream ended
+	DurationMs int64                  `json:"duration_ms"`        // Total duration in milliseconds
+	Status     string                 `json:"status"`             // "completed" | "error" | "cancelled"
+	Error      string                 `json:"error,omitempty"`    // Error message if status is "error"
+	Usage      *UsageInfo             `json:"usage,omitempty"`    // Token usage statistics
+	Metadata   map[string]interface{} `json:"metadata,omitempty"` // Metadata to pass to the page for CUI context
 }
 
 // GroupStartData represents the data for group_start event