From ea167b35a72948cb2645a676b2358d0d0ccb76fb Mon Sep 17 00:00:00 2001
From: Max <shadow.iqka@qq.com>
Date: Sat, 22 Nov 2025 17:32:51 +0800
Subject: [PATCH] Enhance response structure and stream start data in Assistant
 for improved context handling

- Updated the Response struct to include ContextID, RequestID, ChatID, and AssistantID for better traceability and context management.
- Modified the Stream method to return the new fields in the response, enhancing the output for client interactions.
- Added ContextID and ChatID to the StreamStartData struct to improve event tracking during streaming requests.
---
 agent/assistant/agent.go   | 12 +++++++++---
 agent/context/types.go     | 14 +++++++++-----
 agent/context/types_llm.go |  1 +
 3 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/agent/assistant/agent.go b/agent/assistant/agent.go
index f85fc1f6..cf4b09c6 100644
--- a/agent/assistant/agent.go
+++ b/agent/assistant/agent.go
@@ -249,7 +249,12 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
 		}
 	}
 
-	return &context.Response{Create: createResponse, Done: doneResponse, Completion: completionResponse}, nil
+	return &context.Response{
+		ContextID:   ctx.ID,
+		RequestID:   ctx.RequestID(),
+		ChatID:      ctx.ChatID,
+		AssistantID: ast.ID,
+		Create:      createResponse, Done: doneResponse, Completion: completionResponse}, nil
 }
 
 // GetConnector get the connector object, capabilities, and error with priority: createResponse > ctx > ast
@@ -676,11 +681,12 @@ func (ast *Assistant) sendAgentStreamStart(ctx *context.Context, handler context
 
 	// Build the start data
 	startData := context.StreamStartData{
+		ContextID: ctx.ID,
+		ChatID:    ctx.ChatID,
+		TraceID:   ctx.TraceID(),
 		RequestID: ctx.RequestID(),
 		Timestamp: startTime.UnixMilli(),
 		Assistant: ast.Info(ctx.Locale),
-		ChatID:    ctx.ChatID,
-		TraceID:   ctx.TraceID(),
 	}
 
 	if startJSON, err := jsoniter.Marshal(startData); err == nil {
diff --git a/agent/context/types.go b/agent/context/types.go
index 9ff39fd1..65ad94f1 100644
--- a/agent/context/types.go
+++ b/agent/context/types.go
@@ -263,11 +263,15 @@ type Stack struct {
 // Response the response
 // 100% compatible with the OpenAI API
 type Response struct {
-	Create     *HookCreateResponse   `json:"create,omitempty"`
-	MCP        *ResponseHookMCP      `json:"mcp,omitempty"`
-	Done       *ResponseHookDone     `json:"done,omitempty"`
-	Failback   *ResponseHookFailback `json:"failback,omitempty"`
-	Completion *CompletionResponse   `json:"completion,omitempty"`
+	RequestID   string                `json:"request_id"`   // Request ID for the response
+	ContextID   string                `json:"context_id"`   // Context ID for the response
+	ChatID      string                `json:"chat_id"`      // Chat ID for the response
+	AssistantID string                `json:"assistant_id"` // Assistant ID for the response
+	Create      *HookCreateResponse   `json:"create,omitempty"`
+	MCP         *ResponseHookMCP      `json:"mcp,omitempty"`
+	Done        *ResponseHookDone     `json:"done,omitempty"`
+	Failback    *ResponseHookFailback `json:"failback,omitempty"`
+	Completion  *CompletionResponse   `json:"completion,omitempty"`
 }
 
 // HookCreateResponse the response of the create hook
diff --git a/agent/context/types_llm.go b/agent/context/types_llm.go
index 2b6dd853..57f54cd7 100644
--- a/agent/context/types_llm.go
+++ b/agent/context/types_llm.go
@@ -221,6 +221,7 @@ 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")