Merge pull request #1361 from trheyi/main

Implement thread-safe message and block metadata storage

agent/assistant/handlers/stream.go

@@ -200,24 +200,78 @@ func (s *streamState) handleToolCall(data []byte) int {
 	// Track current message type
 	s.currentType = message.TypeToolCall
 
-	// Append to buffer
+	// Append to buffer for message_end event
 	s.buffer = append(s.buffer, data...)
 	s.chunkCount++
 	s.messageSeq++
 
+	// Parse the tool call delta data (JSON array from OpenAI)
+	var toolCallArray []map[string]interface{}
+	if err := jsoniter.Unmarshal(data, &toolCallArray); err != nil {
+		// If parse fails, log and skip this chunk
+		return 0
+	}
+
+	// Extract tool call fields from delta
+	// OpenAI delta typically has one element, but we handle arrays safely
+	var props map[string]interface{}
+	var deltaAction string
+	var deltaPath string
+
+	if len(toolCallArray) == 1 {
+		// Single tool call - flatten to props root level
+		tc := toolCallArray[0]
+		props = map[string]interface{}{}
+
+		// Static fields (only in first chunk): use merge
+		if id, ok := tc["id"].(string); ok {
+			props["id"] = id
+		}
+		if typ, ok := tc["type"].(string); ok {
+			props["type"] = typ
+		}
+		if index, ok := tc["index"].(float64); ok {
+			props["index"] = int(index)
+		}
+		if fn, ok := tc["function"].(map[string]interface{}); ok {
+			if name, ok := fn["name"].(string); ok {
+				props["name"] = name
+			}
+			// Arguments field: use append
+			if args, ok := fn["arguments"].(string); ok {
+				props["arguments"] = args
+				// If this chunk has arguments, use append action for arguments field
+				deltaAction = "append"
+				deltaPath = "arguments"
+			}
+		}
+
+		// If no arguments in this chunk, use merge for other fields
+		if deltaAction == "" {
+			deltaAction = "merge"
+		}
+	} else {
+		// Multiple tool calls in delta (rare) - keep as array
+		props = map[string]interface{}{
+			"calls": toolCallArray,
+		}
+		deltaAction = "merge"
+	}
+
 	// 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)
+	// - DeltaAction: "append" for arguments chunks, "merge" for id/type/name chunks
+	// - DeltaPath: "arguments" when appending arguments field
+	//   OpenAI sends: first chunk has id/type/name, subsequent chunks only have arguments fragments
 	msg := &message.Message{
 		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
+		DeltaAction: deltaAction, // "append" for arguments, "merge" for static fields
-		Props: map[string]interface{}{
+		DeltaPath:   deltaPath,   // "arguments" when appending
-			"raw": string(data), // Raw tool call JSON data
+		Props:       props,       // Flattened tool call fields
-		},
 	}
 
 	if err := s.ctx.Send(msg); err != nil {

agent/context/message.go

@@ -3,8 +3,61 @@ package context
 import (
 	"encoding/json"
 	"fmt"
+	"sync"
 )
 
+// messageMetadataStore provides thread-safe storage for message and block metadata
+type messageMetadataStore struct {
+	messages map[string]*MessageMetadata // Message metadata by MessageID
+	blocks   map[string]*BlockMetadata   // Block metadata by BlockID
+	mu       sync.RWMutex
+}
+
+// newMessageMetadataStore creates a new message metadata store
+func newMessageMetadataStore() *messageMetadataStore {
+	return &messageMetadataStore{
+		messages: make(map[string]*MessageMetadata),
+		blocks:   make(map[string]*BlockMetadata),
+	}
+}
+
+// setMessage stores metadata for a message (thread-safe)
+func (s *messageMetadataStore) setMessage(messageID string, metadata *MessageMetadata) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.messages[messageID] = metadata
+}
+
+// getMessage retrieves metadata for a message (thread-safe)
+func (s *messageMetadataStore) getMessage(messageID string) *MessageMetadata {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	return s.messages[messageID]
+}
+
+// setBlock stores metadata for a block (thread-safe)
+func (s *messageMetadataStore) setBlock(blockID string, metadata *BlockMetadata) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.blocks[blockID] = metadata
+}
+
+// getBlock retrieves metadata for a block (thread-safe)
+func (s *messageMetadataStore) getBlock(blockID string) *BlockMetadata {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	return s.blocks[blockID]
+}
+
+// updateBlock updates block metadata (thread-safe)
+func (s *messageMetadataStore) updateBlock(blockID string, update func(*BlockMetadata)) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if block, exists := s.blocks[blockID]; exists {
+		update(block)
+	}
+}
+
 // UnmarshalJSON custom unmarshaler for Message to handle Content field
 func (m *Message) UnmarshalJSON(data []byte) error {
 	// Define a temporary struct to avoid infinite recursion

agent/context/types.go

@@ -214,58 +214,6 @@ type BlockMetadata struct {
 	MessageCount int       // Number of messages in this block
 }
 
-// messageMetadataStore provides thread-safe storage for message and block metadata
-type messageMetadataStore struct {
-	messages map[string]*MessageMetadata // Message metadata by MessageID
-	blocks   map[string]*BlockMetadata   // Block metadata by BlockID
-	mu       sync.RWMutex
-}
-
-// newMessageMetadataStore creates a new message metadata store
-func newMessageMetadataStore() *messageMetadataStore {
-	return &messageMetadataStore{
-		messages: make(map[string]*MessageMetadata),
-		blocks:   make(map[string]*BlockMetadata),
-	}
-}
-
-// setMessage stores metadata for a message (thread-safe)
-func (s *messageMetadataStore) setMessage(messageID string, metadata *MessageMetadata) {
-	s.mu.Lock()
-	defer s.mu.Unlock()
-	s.messages[messageID] = metadata
-}
-
-// getMessage retrieves metadata for a message (thread-safe)
-func (s *messageMetadataStore) getMessage(messageID string) *MessageMetadata {
-	s.mu.RLock()
-	defer s.mu.RUnlock()
-	return s.messages[messageID]
-}
-
-// setBlock stores metadata for a block (thread-safe)
-func (s *messageMetadataStore) setBlock(blockID string, metadata *BlockMetadata) {
-	s.mu.Lock()
-	defer s.mu.Unlock()
-	s.blocks[blockID] = metadata
-}
-
-// getBlock retrieves metadata for a block (thread-safe)
-func (s *messageMetadataStore) getBlock(blockID string) *BlockMetadata {
-	s.mu.RLock()
-	defer s.mu.RUnlock()
-	return s.blocks[blockID]
-}
-
-// updateBlock updates block metadata (thread-safe)
-func (s *messageMetadataStore) updateBlock(blockID string, update func(*BlockMetadata)) {
-	s.mu.Lock()
-	defer s.mu.Unlock()
-	if block, exists := s.blocks[blockID]; exists {
-		update(block)
-	}
-}
-
 // Context the context
 type Context struct {
 

agent/llm/providers/openai/openai.go

@@ -587,6 +587,8 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess
 				}
 
 				// Notify handler of tool call progress
+				// Send the raw delta from OpenAI (as JSON bytes)
+				// Handler will convert to object for frontend merge
 				if handler != nil {
 					toolCallData, _ := jsoniter.Marshal(delta.ToolCalls)
 					handler(message.ChunkToolCall, toolCallData)

agent/types/types.go

@@ -1,7 +1,6 @@
 package types
 
 import (
-	"github.com/gin-gonic/gin"
 	"github.com/yaoapp/yao/agent/assistant"
 	store "github.com/yaoapp/yao/agent/store/types"
 )
@@ -15,26 +14,15 @@ type DSL struct {
 	StoreSetting store.Setting `json:"store" yaml:"store"`                   // The store setting of the assistant
 	Cache        string        `json:"cache" yaml:"cache"`                   // The cache store of the assistant, if not set, default is "__yao.agent.cache"
 
-	// AuthSetting      *Auth         `json:"auth,omitempty" yaml:"auth,omitempty"`           // Authenticate Settings
-	// UploadSetting    *Upload       `json:"upload,omitempty" yaml:"upload,omitempty"`       // Upload Settings
-	// KnowledgeSetting *Knowledge    `json:"knowledge,omitempty" yaml:"knowledge,omitempty"` // Knowledge base Settings
-
 	// Global External Settings - model capabilities, tools, etc.
 	// ===============================
 	Models map[string]assistant.ModelCapabilities `json:"models,omitempty" yaml:"models,omitempty"` // The model capabilities configuration
 
-	// Agent API Settings
-	// ===============================s
-	// Guard  string   `json:"guard,omitempty" yaml:"guard,omitempty"`   // The guard of the assistant
-	// Allows []string `json:"allows,omitempty" yaml:"allows,omitempty"` // The allowed domains of the assistant
-
 	// Internal
 	// ===============================
 	// ID            string            `json:"-" yaml:"-"` // The id of the instance
 	Assistant assistant.API `json:"-" yaml:"-"` // The default assistant
 	Store     store.Store   `json:"-" yaml:"-"` // The store of the assistant
-	// Vision        *vision.Vision    `json:"-" yaml:"-"`
-	GuardHandlers []gin.HandlerFunc `json:"-" yaml:"-"`
 }
 
 // Uses the default assistant settings
@@ -49,34 +37,6 @@ type Uses struct {
 	Fetch   string `json:"fetch,omitempty" yaml:"fetch,omitempty"`     // The assistant for fetching the http/https/ftp/sftp/etc. file, and return the file's content. if not set, use the http process to fetch the file.
 }
 
-// Auth Authenticate Settings
-// ===============================
-type Auth struct {
-	Models        *AuthModels        `json:"models,omitempty" yaml:"models,omitempty"`                 // The models of the user, it is used to handle the user, and the user is a user in the database. (Guest and User model must have the id and permission fields)
-	Fields        *AuthFields        `json:"fields,omitempty" yaml:"fields,omitempty"`                 // The fields of the user model, it is used to handle the user, and the user is a user in the database. (Guest and User model must have the id and permission fields)
-	SessionFields *AuthSessionFields `json:"session_fields,omitempty" yaml:"session_fields,omitempty"` // The session fields of the user, it is used to handle the user, and the user is a user in the database.
-}
-
-// AuthModels the auth model
-type AuthModels struct {
-	User  string `json:"user,omitempty" yaml:"user,omitempty"`   // default is admin.user, The user model is a special model, it is used to handle the user, and the user is a user in the database.
-	Guest string `json:"guest,omitempty" yaml:"guest,omitempty"` // The guest model is a special model, it is used to handle the guest user, and the guest user is not a user in the database.
-}
-
-// AuthSessionFields the auth session field
-type AuthSessionFields struct {
-	ID    string `json:"id,omitempty" yaml:"id,omitempty"`       // the field name of the user id, default is user_id
-	Roles string `json:"roles,omitempty" yaml:"roles,omitempty"` // the field name of the user roles, default is user_roles. the value must be an JSON array string.
-	Guest string `json:"guest,omitempty" yaml:"guest,omitempty"` // the field name of the guest user, default is guest_id
-}
-
-// AuthFields the auth field
-type AuthFields struct {
-	ID         string `json:"id,omitempty" yaml:"id,omitempty"`                 // the field name of the user id, default is id
-	Roles      string `json:"roles,omitempty" yaml:"roles,omitempty"`           // the field name of the user roles, default is roles, it must be an JSON field.
-	Permission string `json:"permission,omitempty" yaml:"permission,omitempty"` // the field name of the user permission, default is permission
-}
-
 // Mention Structure
 // ===============================
 type Mention struct {