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

@@ -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
 }
 

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)
+}

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

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"

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"

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