From a868b4d7416e6ccbda6347dca0376fbd3e2a1c3b Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 30 Nov 2025 16:40:45 +0800 Subject: [PATCH] Refactor context initialization and enhance tool call handling - Updated context creation methods to remove unnecessary payload parameter, simplifying initialization. - Enhanced handleToolCall method to manage message sequencing and delta actions more effectively, ensuring proper message tracking. - Improved error handling in message sending to ensure robustness in tool call processing. --- agent/assistant/handlers/stream.go | 35 +++++++++--- agent/context/authorized_test.go | 84 ++++++++++++++++++++++++++++ agent/context/context.go | 38 +++++++------ agent/context/jsapi_output_test.go | 3 +- agent/context/message_events_test.go | 2 +- agent/context/openapi.go | 3 +- 6 files changed, 134 insertions(+), 31 deletions(-) create mode 100644 agent/context/authorized_test.go diff --git a/agent/assistant/handlers/stream.go b/agent/assistant/handlers/stream.go index 179dd963..45aced84 100644 --- a/agent/assistant/handlers/stream.go +++ b/agent/assistant/handlers/stream.go @@ -193,20 +193,37 @@ func (s *streamState) handleThinking(data []byte) int { // handleToolCall handles tool call chunks func (s *streamState) handleToolCall(data []byte) int { - // Tool calls are usually complete JSON objects - // Parse and send as tool_call message + if len(data) == 0 { + return 0 + } + + // Track current message type + s.currentType = message.TypeToolCall + + // Append to buffer + s.buffer = append(s.buffer, data...) + s.chunkCount++ + s.messageSeq++ + + // Send delta message + // - ChunkID: Unique chunk ID (C1, C2, C3...) for this fragment + // - MessageID: Same for all chunks of this logical message (frontend merges by message_id) + // - DeltaAction: "replace" for tool call raw data (each chunk contains complete state, not incremental) msg := &message.Message{ - ChunkID: s.ctx.IDGenerator.GenerateChunkID(), - MessageID: s.ctx.IDGenerator.GenerateMessageID(), // Tool call is a new message - Type: message.TypeToolCall, - Delta: true, + ChunkID: s.ctx.IDGenerator.GenerateChunkID(), // Unique chunk ID + MessageID: s.currentGroupID, // Message ID for merging (all chunks share this) + Type: message.TypeToolCall, + Delta: true, + DeltaAction: "replace", // Replace entire raw field with latest state Props: map[string]interface{}{ - // TODO: Parse tool call data - "raw": string(data), + "raw": string(data), // Raw tool call JSON data }, } - s.ctx.Send(msg) + if err := s.ctx.Send(msg); err != nil { + return 0 + } + return 0 // Continue } diff --git a/agent/context/authorized_test.go b/agent/context/authorized_test.go new file mode 100644 index 00000000..47eeeea2 --- /dev/null +++ b/agent/context/authorized_test.go @@ -0,0 +1,84 @@ +package context + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/yaoapp/yao/config" + "github.com/yaoapp/yao/openapi/oauth/types" + "github.com/yaoapp/yao/test" + "github.com/yaoapp/yao/trace" +) + +func TestContextNew_PreservesAuthorizedInfo(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() + + // Create authorized info + authInfo := &types.AuthorizedInfo{ + UserID: "716942074991", + TeamID: "565955042879", + TenantID: "tenant-001", + } + + // Create context using New() + ctx := New(context.Background(), authInfo, "test-chat-123") + + // Verify authorized info is preserved + assert.NotNil(t, ctx) + assert.NotNil(t, ctx.Authorized) + assert.Equal(t, "716942074991", ctx.Authorized.UserID) + assert.Equal(t, "565955042879", ctx.Authorized.TeamID) + assert.Equal(t, "tenant-001", ctx.Authorized.TenantID) + assert.Equal(t, "test-chat-123", ctx.ChatID) +} + +func TestContextTrace_SavesAuthorizedInfo(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() + + // Create authorized info + authInfo := &types.AuthorizedInfo{ + UserID: "716942074991", + TeamID: "565955042879", + TenantID: "tenant-001", + } + + // Create context + ctx := New(context.Background(), authInfo, "test-chat-456") + ctx.AssistantID = "test-assistant" + + // Initialize trace + manager, err := ctx.Trace() + assert.NoError(t, err) + assert.NotNil(t, manager) + + // Get trace info + info, err := manager.GetTraceInfo() + assert.NoError(t, err) + assert.NotNil(t, info) + + // Verify auth info is saved in trace + assert.Equal(t, "716942074991", info.CreatedBy) + assert.Equal(t, "565955042879", info.TeamID) + assert.Equal(t, "tenant-001", info.TenantID) + + // Clean up + if ctx.Stack != nil && ctx.Stack.TraceID != "" { + trace.Release(ctx.Stack.TraceID) + trace.Remove(context.Background(), trace.Local, ctx.Stack.TraceID) + } +} + +func TestContextNew_NilAuthorized(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() + + // Create context with nil authorized info (should not panic) + ctx := New(context.Background(), nil, "test-chat-789") + + assert.NotNil(t, ctx) + assert.Nil(t, ctx.Authorized) + assert.Equal(t, "test-chat-789", ctx.ChatID) +} diff --git a/agent/context/context.go b/agent/context/context.go index b250fdb2..3af8f4da 100644 --- a/agent/context/context.go +++ b/agent/context/context.go @@ -21,56 +21,60 @@ var ( contextRegistry = &sync.Map{} // map[contextID]*Context ) -// New create a new context -func New(parent context.Context, authorized *types.AuthorizedInfo, chatID, payload string) Context { - +// New create a new context with basic initialization +func New(parent context.Context, authorized *types.AuthorizedInfo, chatID string) *Context { if parent == nil { parent = context.Background() } - // Validate the client type - ctx := Context{ + ctx := &Context{ Context: parent, ID: generateContextID(), // Generate unique ID for the context + Authorized: authorized, // Set authorized info Space: plan.NewMemorySharedSpace(), ChatID: chatID, IDGenerator: message.NewIDGenerator(), // Initialize ID generator for this context messageMetadata: newMessageMetadataStore(), // Initialize message metadata store } - if payload == "" { - return ctx - } + return ctx +} - err := jsoniter.Unmarshal([]byte(payload), &ctx) - if err != nil { - log.Error("%s", err.Error()) +// NewWithPayload create a new context and unmarshal from payload +func NewWithPayload(parent context.Context, authorized *types.AuthorizedInfo, chatID, payload string) *Context { + ctx := New(parent, authorized, chatID) + + if payload != "" { + err := jsoniter.Unmarshal([]byte(payload), ctx) + if err != nil { + log.Error("%s", err.Error()) + } } return ctx } // NewWithCancel create a new context with cancel -func NewWithCancel(parent context.Context, authorized *types.AuthorizedInfo, chatID, payload string) (Context, context.CancelFunc) { - ctx := New(parent, authorized, chatID, payload) +func NewWithCancel(parent context.Context, authorized *types.AuthorizedInfo, chatID string) (*Context, context.CancelFunc) { + ctx := New(parent, authorized, chatID) return WithCancel(ctx) } // NewWithTimeout create a new context with timeout -func NewWithTimeout(parent context.Context, authorized *types.AuthorizedInfo, chatID, payload string, timeout time.Duration) (Context, context.CancelFunc) { - ctx := New(parent, authorized, chatID, payload) +func NewWithTimeout(parent context.Context, authorized *types.AuthorizedInfo, chatID string, timeout time.Duration) (*Context, context.CancelFunc) { + ctx := New(parent, authorized, chatID) return WithTimeout(ctx, timeout) } // WithCancel create a new context -func WithCancel(parent Context) (Context, context.CancelFunc) { +func WithCancel(parent *Context) (*Context, context.CancelFunc) { new, cancel := context.WithCancel(parent.Context) parent.Context = new return parent, cancel } // WithTimeout create a new context -func WithTimeout(parent Context, timeout time.Duration) (Context, context.CancelFunc) { +func WithTimeout(parent *Context, timeout time.Duration) (*Context, context.CancelFunc) { new, cancel := context.WithTimeout(parent.Context, timeout) parent.Context = new return parent, cancel diff --git a/agent/context/jsapi_output_test.go b/agent/context/jsapi_output_test.go index 1ebe1ab7..2539a0e4 100644 --- a/agent/context/jsapi_output_test.go +++ b/agent/context/jsapi_output_test.go @@ -784,8 +784,7 @@ func TestJsValueEndBlock(t *testing.T) { mockWriter := newMockResponseWriter() // Use New() to properly initialize messageMetadata - ctxValue := New(context.Background(), nil, "test-chat-id", "") - cxt := &ctxValue + cxt := New(context.Background(), nil, "test-chat-id") cxt.AssistantID = "test-assistant-id" cxt.Accept = AcceptWebCUI cxt.Locale = "en" diff --git a/agent/context/message_events_test.go b/agent/context/message_events_test.go index 71f7fd75..b06ae1f6 100644 --- a/agent/context/message_events_test.go +++ b/agent/context/message_events_test.go @@ -21,7 +21,7 @@ func TestMessageLifecycleEvents(t *testing.T) { } // Create context using New() to ensure proper initialization - ctx := context.New(stdContext.Background(), nil, "test-chat", "") + ctx := context.New(stdContext.Background(), nil, "test-chat") ctx.Accept = context.AcceptWebCUI ctx.Writer = mockWriter ctx.AssistantID = "test-assistant" diff --git a/agent/context/openapi.go b/agent/context/openapi.go index 38b0e81d..d35677dc 100644 --- a/agent/context/openapi.go +++ b/agent/context/openapi.go @@ -44,8 +44,7 @@ func GetCompletionRequest(c *gin.Context, cache store.Store) (*CompletionRequest clientIP := c.ClientIP() // Create context with unique ID using New() to ensure proper initialization - ctxValue := New(c.Request.Context(), authInfo, chatID, "") - ctx := &ctxValue + ctx := New(c.Request.Context(), authInfo, chatID) // Set additional fields ctx.Cache = cache