diff --git a/agent/store/CHAT_STORAGE_DESIGN.md b/agent/store/CHAT_STORAGE_DESIGN.md index cec06b1e..ca22e1ee 100644 --- a/agent/store/CHAT_STORAGE_DESIGN.md +++ b/agent/store/CHAT_STORAGE_DESIGN.md @@ -881,8 +881,11 @@ const ( ```go // ChatFilter for listing chats type ChatFilter struct { - UserID string `json:"user_id,omitempty"` - TeamID string `json:"team_id,omitempty"` + // Permission filters (direct filtering on Yao permission fields) + UserID string `json:"user_id,omitempty"` // Filter by __yao_created_by + TeamID string `json:"team_id,omitempty"` // Filter by __yao_team_id + + // Business filters AssistantID string `json:"assistant_id,omitempty"` Status string `json:"status,omitempty"` Keywords string `json:"keywords,omitempty"` @@ -903,8 +906,9 @@ type ChatFilter struct { Page int `json:"page,omitempty"` PageSize int `json:"pagesize,omitempty"` - // Permission filter (not serialized) - QueryFilter func(query.Query) `json:"-"` // Custom query function for permission filtering + // Advanced permission filter (not serialized) + // Use for complex conditions like: (created_by = user OR team_id = team) + QueryFilter func(query.Query) `json:"-"` } // MessageFilter for listing messages @@ -1162,9 +1166,9 @@ Multimedia content storage: ### 5. Load Chat History ```go -// Example 1: Flat list (default) +// Example 1: Filter by user (simple permission check) chats, _ := chatStore.ListChats(ChatFilter{ - UserID: "user123", + UserID: "user123", // Filters by __yao_created_by Status: "active", OrderBy: "last_message_at", Order: "desc", @@ -1173,7 +1177,35 @@ chats, _ := chatStore.ListChats(ChatFilter{ }) // Response: chats.Data = [...], chats.Groups = nil -// Example 2: Grouped by time +// Example 2: Filter by team +chats, _ := chatStore.ListChats(ChatFilter{ + TeamID: "team456", // Filters by __yao_team_id + Status: "active", + Page: 1, + PageSize: 20, +}) + +// Example 3: Filter by user AND team (both must match) +chats, _ := chatStore.ListChats(ChatFilter{ + UserID: "user123", + TeamID: "team456", + Page: 1, + PageSize: 20, +}) + +// Example 4: Complex permission filter (user OR team) using QueryFilter +chats, _ := chatStore.ListChats(ChatFilter{ + Page: 1, + PageSize: 20, + QueryFilter: func(qb query.Query) { + qb.Where(func(sub query.Query) { + sub.Where("__yao_created_by", "user123"). + OrWhere("__yao_team_id", "team456") + }) + }, +}) + +// Example 5: Grouped by time chats, _ := chatStore.ListChats(ChatFilter{ UserID: "user123", GroupBy: "time", // Enable time-based grouping @@ -1191,7 +1223,7 @@ chats, _ := chatStore.ListChats(ChatFilter{ // { Key: "earlier", Label: "Earlier", Chats: [...], Count: 0 }, // ] -// Example 3: Filter by time range +// Example 6: Filter by time range startTime := time.Now().AddDate(0, 0, -7) // Last 7 days chats, _ := chatStore.ListChats(ChatFilter{ UserID: "user123", @@ -1201,7 +1233,7 @@ chats, _ := chatStore.ListChats(ChatFilter{ Order: "desc", }) -// Example 4: Filter specific date range +// Example 7: Filter specific date range start := time.Date(2024, 12, 1, 0, 0, 0, 0, time.Local) end := time.Date(2024, 12, 31, 23, 59, 59, 0, time.Local) chats, _ := chatStore.ListChats(ChatFilter{ @@ -1211,6 +1243,17 @@ chats, _ := chatStore.ListChats(ChatFilter{ TimeField: "created_at", // Filter by creation time }) +// Example 8: Combine permission with business filters +chats, _ := chatStore.ListChats(ChatFilter{ + UserID: "user123", + TeamID: "team456", + AssistantID: "weather_assistant", + Status: "active", + Keywords: "weather", + Page: 1, + PageSize: 20, +}) + // Get messages for a chat messages, _ := chatStore.GetMessages("chat_123", MessageFilter{ Limit: 100, diff --git a/agent/store/xun/chat.go b/agent/store/xun/chat.go index 4a8545f0..e850fc11 100644 --- a/agent/store/xun/chat.go +++ b/agent/store/xun/chat.go @@ -7,7 +7,6 @@ import ( "github.com/google/uuid" jsoniter "github.com/json-iterator/go" - "github.com/yaoapp/xun/dbal/query" "github.com/yaoapp/yao/agent/store/types" ) @@ -224,7 +223,15 @@ func (store *Xun) ListChats(filter types.ChatFilter) (*types.ChatList, error) { // Build base query qb := store.newQueryChat().WhereNull("deleted_at") - // Apply filters + // Apply permission filters (UserID and TeamID) + if filter.UserID != "" { + qb.Where("__yao_created_by", filter.UserID) + } + if filter.TeamID != "" { + qb.Where("__yao_team_id", filter.TeamID) + } + + // Apply business filters if filter.AssistantID != "" { qb.Where("assistant_id", filter.AssistantID) } @@ -243,7 +250,8 @@ func (store *Xun) ListChats(filter types.ChatFilter) (*types.ChatList, error) { qb.Where(filter.TimeField, "<=", *filter.EndTime) } - // Apply custom query filter (for permission filtering) + // Apply custom query filter (for advanced permission filtering) + // This allows flexible combinations like: (created_by = user OR team_id = team) if filter.QueryFilter != nil { qb.Where(filter.QueryFilter) } @@ -439,14 +447,3 @@ func (store *Xun) UpdateChatLastMessageAt(chatID string, timestamp time.Time) er return err } - -// newQueryChatWithPermission creates a new query builder with permission filtering -func (store *Xun) newQueryChatWithPermission(filter types.ChatFilter) query.Query { - qb := store.newQueryChat().WhereNull("deleted_at") - - if filter.QueryFilter != nil { - qb.Where(filter.QueryFilter) - } - - return qb -} diff --git a/agent/store/xun/chat_test.go b/agent/store/xun/chat_test.go index 0a790de7..0b439c02 100644 --- a/agent/store/xun/chat_test.go +++ b/agent/store/xun/chat_test.go @@ -5,6 +5,7 @@ import ( "testing" "time" + goumodel "github.com/yaoapp/gou/model" "github.com/yaoapp/xun/dbal/query" "github.com/yaoapp/yao/agent/store/types" "github.com/yaoapp/yao/agent/store/xun" @@ -786,6 +787,224 @@ func TestListChats(t *testing.T) { }) } +// TestListChatsByUserAndTeam tests filtering chats by UserID and TeamID +func TestListChatsByUserAndTeam(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() + + store, err := xun.NewXun(types.Setting{ + Connector: "default", + }) + if err != nil { + t.Fatalf("Failed to create store: %v", err) + } + + // Create chats with different user/team combinations + // Note: __yao_created_by and __yao_team_id are managed by Yao's permission system + // For testing, we'll create chats and then update these fields directly via raw query + + chat1 := &types.Chat{AssistantID: "test_assistant", Title: "User1 Team1 Chat"} + chat2 := &types.Chat{AssistantID: "test_assistant", Title: "User1 Team2 Chat"} + chat3 := &types.Chat{AssistantID: "test_assistant", Title: "User2 Team1 Chat"} + chat4 := &types.Chat{AssistantID: "test_assistant", Title: "User2 Team2 Chat"} + + for _, chat := range []*types.Chat{chat1, chat2, chat3, chat4} { + err := store.CreateChat(chat) + if err != nil { + t.Fatalf("Failed to create chat: %v", err) + } + } + defer func() { + store.DeleteChat(chat1.ChatID) + store.DeleteChat(chat2.ChatID) + store.DeleteChat(chat3.ChatID) + store.DeleteChat(chat4.ChatID) + }() + + // Update permission fields directly for testing + // In production, these would be set by Yao's permission middleware + updatePermissionFields := func(chatID, userID, teamID string) error { + // Use Yao model to update permission fields + m := goumodel.Select("__yao.agent.chat") + if m == nil { + return fmt.Errorf("model __yao.agent.chat not found") + } + _, err := m.UpdateWhere( + goumodel.QueryParam{Wheres: []goumodel.QueryWhere{{Column: "chat_id", Value: chatID}}}, + map[string]interface{}{ + "__yao_created_by": userID, + "__yao_team_id": teamID, + }, + ) + return err + } + + // Set up permission fields + updatePermissionFields(chat1.ChatID, "user1", "team1") + updatePermissionFields(chat2.ChatID, "user1", "team2") + updatePermissionFields(chat3.ChatID, "user2", "team1") + updatePermissionFields(chat4.ChatID, "user2", "team2") + + t.Run("FilterByUserID", func(t *testing.T) { + result, err := store.ListChats(types.ChatFilter{ + UserID: "user1", + Page: 1, + PageSize: 20, + }) + if err != nil { + t.Fatalf("Failed to list chats by user: %v", err) + } + + if len(result.Data) != 2 { + t.Errorf("Expected 2 chats for user1, got %d", len(result.Data)) + } + + // Verify all returned chats belong to user1 + for _, chat := range result.Data { + if chat.Title != "User1 Team1 Chat" && chat.Title != "User1 Team2 Chat" { + t.Errorf("Unexpected chat title: %s", chat.Title) + } + } + }) + + t.Run("FilterByTeamID", func(t *testing.T) { + result, err := store.ListChats(types.ChatFilter{ + TeamID: "team1", + Page: 1, + PageSize: 20, + }) + if err != nil { + t.Fatalf("Failed to list chats by team: %v", err) + } + + if len(result.Data) != 2 { + t.Errorf("Expected 2 chats for team1, got %d", len(result.Data)) + } + + // Verify all returned chats belong to team1 + for _, chat := range result.Data { + if chat.Title != "User1 Team1 Chat" && chat.Title != "User2 Team1 Chat" { + t.Errorf("Unexpected chat title: %s", chat.Title) + } + } + }) + + t.Run("FilterByUserIDAndTeamID", func(t *testing.T) { + result, err := store.ListChats(types.ChatFilter{ + UserID: "user1", + TeamID: "team1", + Page: 1, + PageSize: 20, + }) + if err != nil { + t.Fatalf("Failed to list chats by user and team: %v", err) + } + + if len(result.Data) != 1 { + t.Errorf("Expected 1 chat for user1+team1, got %d", len(result.Data)) + } + + if len(result.Data) > 0 && result.Data[0].Title != "User1 Team1 Chat" { + t.Errorf("Expected 'User1 Team1 Chat', got '%s'", result.Data[0].Title) + } + }) + + t.Run("FilterByUserIDWithOtherFilters", func(t *testing.T) { + // Combine UserID with Status filter + result, err := store.ListChats(types.ChatFilter{ + UserID: "user1", + Status: "active", + Page: 1, + PageSize: 20, + }) + if err != nil { + t.Fatalf("Failed to list chats: %v", err) + } + + // All user1's chats should be active (default status) + if len(result.Data) != 2 { + t.Errorf("Expected 2 active chats for user1, got %d", len(result.Data)) + } + }) + + t.Run("FilterByTeamIDWithQueryFilter", func(t *testing.T) { + // Combine TeamID with custom QueryFilter + result, err := store.ListChats(types.ChatFilter{ + TeamID: "team2", + Page: 1, + PageSize: 20, + QueryFilter: func(qb query.Query) { + // Additional filter: only chats with "User1" in title + qb.Where("title", "like", "%User1%") + }, + }) + if err != nil { + t.Fatalf("Failed to list chats: %v", err) + } + + if len(result.Data) != 1 { + t.Errorf("Expected 1 chat (User1 in team2), got %d", len(result.Data)) + } + + if len(result.Data) > 0 && result.Data[0].Title != "User1 Team2 Chat" { + t.Errorf("Expected 'User1 Team2 Chat', got '%s'", result.Data[0].Title) + } + }) + + t.Run("FilterByNonExistentUser", func(t *testing.T) { + result, err := store.ListChats(types.ChatFilter{ + UserID: "nonexistent_user", + Page: 1, + PageSize: 20, + }) + if err != nil { + t.Fatalf("Failed to list chats: %v", err) + } + + if len(result.Data) != 0 { + t.Errorf("Expected 0 chats for nonexistent user, got %d", len(result.Data)) + } + }) + + t.Run("FilterByNonExistentTeam", func(t *testing.T) { + result, err := store.ListChats(types.ChatFilter{ + TeamID: "nonexistent_team", + Page: 1, + PageSize: 20, + }) + if err != nil { + t.Fatalf("Failed to list chats: %v", err) + } + + if len(result.Data) != 0 { + t.Errorf("Expected 0 chats for nonexistent team, got %d", len(result.Data)) + } + }) + + t.Run("QueryFilterForOrCondition", func(t *testing.T) { + // Use QueryFilter for complex OR condition: + // Get chats where user is user1 OR team is team2 + result, err := store.ListChats(types.ChatFilter{ + Page: 1, + PageSize: 20, + QueryFilter: func(qb query.Query) { + qb.Where(func(sub query.Query) { + sub.Where("__yao_created_by", "user1"). + OrWhere("__yao_team_id", "team2") + }) + }, + }) + if err != nil { + t.Fatalf("Failed to list chats with OR condition: %v", err) + } + + // Should return: user1+team1, user1+team2, user2+team2 = 3 chats + if len(result.Data) != 3 { + t.Errorf("Expected 3 chats (user1 OR team2), got %d", len(result.Data)) + } + }) +} + // TestChatCompleteWorkflow tests a complete chat workflow func TestChatCompleteWorkflow(t *testing.T) { test.Prepare(t, config.Conf) diff --git a/agent/store/xun/message.go b/agent/store/xun/message.go index 09bfd8fd..b196a42b 100644 --- a/agent/store/xun/message.go +++ b/agent/store/xun/message.go @@ -1,6 +1,11 @@ package xun import ( + "fmt" + "time" + + "github.com/google/uuid" + jsoniter "github.com/json-iterator/go" "github.com/yaoapp/yao/agent/store/types" ) @@ -8,28 +13,354 @@ import ( // Message Management // ============================================================================= -// SaveMessages batch saves messages for a chat +// SaveMessages batch saves messages for a chat using a single database call // This is the primary write method - messages are buffered during execution // and batch-written at the end of a request func (store *Xun) SaveMessages(chatID string, messages []*types.Message) error { - // TODO: implement - return nil + if chatID == "" { + return fmt.Errorf("chat_id is required") + } + if len(messages) == 0 { + return nil // Nothing to save + } + + // Prepare batch insert data + now := time.Now() + rows := make([]map[string]interface{}, 0, len(messages)) + + for _, msg := range messages { + if msg == nil { + continue + } + + // Generate message_id if not provided + messageID := msg.MessageID + if messageID == "" { + messageID = uuid.New().String() + } + + // Validate required fields + if msg.Role == "" { + return fmt.Errorf("message role is required") + } + if msg.Type == "" { + return fmt.Errorf("message type is required") + } + if msg.Props == nil { + return fmt.Errorf("message props is required") + } + + // Serialize JSON fields + propsJSON, err := jsoniter.MarshalToString(msg.Props) + if err != nil { + return fmt.Errorf("failed to marshal props: %w", err) + } + + // Build row with all fields (including nullable ones for consistent batch insert) + row := map[string]interface{}{ + "message_id": messageID, + "chat_id": chatID, + "role": msg.Role, + "type": msg.Type, + "props": propsJSON, + "sequence": msg.Sequence, + "request_id": nil, + "block_id": nil, + "thread_id": nil, + "assistant_id": nil, + "metadata": nil, + "created_at": now, + "updated_at": now, + } + + // Set nullable fields if they have values + if msg.RequestID != "" { + row["request_id"] = msg.RequestID + } + if msg.BlockID != "" { + row["block_id"] = msg.BlockID + } + if msg.ThreadID != "" { + row["thread_id"] = msg.ThreadID + } + if msg.AssistantID != "" { + row["assistant_id"] = msg.AssistantID + } + if msg.Metadata != nil { + metadataJSON, err := jsoniter.MarshalToString(msg.Metadata) + if err != nil { + return fmt.Errorf("failed to marshal metadata: %w", err) + } + row["metadata"] = metadataJSON + } + + rows = append(rows, row) + } + + if len(rows) == 0 { + return nil + } + + // Single batch insert - one database call for all messages + return store.newQueryMessage().Insert(rows) } // GetMessages retrieves messages for a chat with filtering func (store *Xun) GetMessages(chatID string, filter types.MessageFilter) ([]*types.Message, error) { - // TODO: implement - return nil, nil + if chatID == "" { + return nil, fmt.Errorf("chat_id is required") + } + + qb := store.newQueryMessage(). + Where("chat_id", chatID). + WhereNull("deleted_at") + + // Apply filters + if filter.RequestID != "" { + qb.Where("request_id", filter.RequestID) + } + if filter.Role != "" { + qb.Where("role", filter.Role) + } + if filter.BlockID != "" { + qb.Where("block_id", filter.BlockID) + } + if filter.ThreadID != "" { + qb.Where("thread_id", filter.ThreadID) + } + if filter.Type != "" { + qb.Where("type", filter.Type) + } + + // Apply pagination (MySQL requires LIMIT when using OFFSET) + if filter.Limit > 0 { + qb.Limit(filter.Limit) + if filter.Offset > 0 { + qb.Offset(filter.Offset) + } + } else if filter.Offset > 0 { + // If only offset is specified, use a large limit + qb.Limit(1000000).Offset(filter.Offset) + } + + // Order by sequence + qb.OrderBy("sequence", "asc") + + rows, err := qb.Get() + if err != nil { + return nil, err + } + + messages := make([]*types.Message, 0, len(rows)) + for _, row := range rows { + data := row.ToMap() + if data == nil || data["message_id"] == nil { + continue + } + + msg, err := store.rowToMessage(data) + if err != nil { + continue + } + messages = append(messages, msg) + } + + return messages, nil } // UpdateMessage updates a single message func (store *Xun) UpdateMessage(messageID string, updates map[string]interface{}) error { - // TODO: implement - return nil + if messageID == "" { + return fmt.Errorf("message_id is required") + } + if len(updates) == 0 { + return fmt.Errorf("no fields to update") + } + + // Check if message exists + exists, err := store.newQueryMessage(). + Where("message_id", messageID). + WhereNull("deleted_at"). + Exists() + if err != nil { + return err + } + if !exists { + return fmt.Errorf("message %s not found", messageID) + } + + // Prepare update data + data := make(map[string]interface{}) + + for key, value := range updates { + // Skip system fields + if key == "message_id" || key == "chat_id" || key == "created_at" { + continue + } + + // Handle JSON fields + if key == "props" || key == "metadata" { + if value != nil { + jsonStr, err := jsoniter.MarshalToString(value) + if err != nil { + return fmt.Errorf("failed to marshal %s: %w", key, err) + } + data[key] = jsonStr + } else { + data[key] = nil + } + continue + } + + data[key] = value + } + + // Always update updated_at + data["updated_at"] = time.Now() + + if len(data) == 0 { + return fmt.Errorf("no valid fields to update") + } + + _, err = store.newQueryMessage(). + Where("message_id", messageID). + Update(data) + + return err } -// DeleteMessages deletes specific messages from a chat +// DeleteMessages soft deletes specific messages from a chat func (store *Xun) DeleteMessages(chatID string, messageIDs []string) error { - // TODO: implement - return nil + if chatID == "" { + return fmt.Errorf("chat_id is required") + } + if len(messageIDs) == 0 { + return nil // Nothing to delete + } + + // Soft delete all specified messages in one query + _, err := store.newQueryMessage(). + Where("chat_id", chatID). + WhereIn("message_id", messageIDs). + WhereNull("deleted_at"). + Update(map[string]interface{}{ + "deleted_at": time.Now(), + "updated_at": time.Now(), + }) + + return err +} + +// GetMessageByID retrieves a single message by ID +func (store *Xun) GetMessageByID(messageID string) (*types.Message, error) { + if messageID == "" { + return nil, fmt.Errorf("message_id is required") + } + + row, err := store.newQueryMessage(). + Where("message_id", messageID). + WhereNull("deleted_at"). + First() + if err != nil { + return nil, err + } + + if row == nil { + return nil, fmt.Errorf("message %s not found", messageID) + } + + data := row.ToMap() + if len(data) == 0 || data["message_id"] == nil { + return nil, fmt.Errorf("message %s not found", messageID) + } + + return store.rowToMessage(data) +} + +// GetMessageCount returns the count of messages for a chat +func (store *Xun) GetMessageCount(chatID string) (int64, error) { + if chatID == "" { + return 0, fmt.Errorf("chat_id is required") + } + + return store.newQueryMessage(). + Where("chat_id", chatID). + WhereNull("deleted_at"). + Count() +} + +// GetLastSequence returns the last sequence number for a chat +func (store *Xun) GetLastSequence(chatID string) (int, error) { + if chatID == "" { + return 0, fmt.Errorf("chat_id is required") + } + + row, err := store.newQueryMessage(). + Where("chat_id", chatID). + WhereNull("deleted_at"). + OrderBy("sequence", "desc"). + First() + if err != nil { + return 0, err + } + + if row == nil { + return 0, nil + } + + data := row.ToMap() + return getInt(data, "sequence"), nil +} + +// ============================================================================= +// Helper Functions +// ============================================================================= + +// rowToMessage converts a database row to a Message struct +func (store *Xun) rowToMessage(data map[string]interface{}) (*types.Message, error) { + msg := &types.Message{ + MessageID: getString(data, "message_id"), + ChatID: getString(data, "chat_id"), + RequestID: getString(data, "request_id"), + Role: getString(data, "role"), + Type: getString(data, "type"), + BlockID: getString(data, "block_id"), + ThreadID: getString(data, "thread_id"), + AssistantID: getString(data, "assistant_id"), + Sequence: getInt(data, "sequence"), + } + + // Handle timestamps + if createdAt := getTime(data, "created_at"); createdAt != nil { + msg.CreatedAt = *createdAt + } + if updatedAt := getTime(data, "updated_at"); updatedAt != nil { + msg.UpdatedAt = *updatedAt + } + + // Handle props (required) + if props := data["props"]; props != nil { + if propsStr, ok := props.(string); ok && propsStr != "" { + var propsMap map[string]interface{} + if err := jsoniter.UnmarshalFromString(propsStr, &propsMap); err == nil { + msg.Props = propsMap + } + } else if propsMap, ok := props.(map[string]interface{}); ok { + msg.Props = propsMap + } + } + + // Handle metadata (optional) + if metadata := data["metadata"]; metadata != nil { + if metaStr, ok := metadata.(string); ok && metaStr != "" { + var metaMap map[string]interface{} + if err := jsoniter.UnmarshalFromString(metaStr, &metaMap); err == nil { + msg.Metadata = metaMap + } + } else if metaMap, ok := metadata.(map[string]interface{}); ok { + msg.Metadata = metaMap + } + } + + return msg, nil } diff --git a/agent/store/xun/message_test.go b/agent/store/xun/message_test.go new file mode 100644 index 00000000..c589c2fe --- /dev/null +++ b/agent/store/xun/message_test.go @@ -0,0 +1,896 @@ +package xun_test + +import ( + "fmt" + "testing" + "time" + + "github.com/yaoapp/yao/agent/store/types" + "github.com/yaoapp/yao/agent/store/xun" + "github.com/yaoapp/yao/config" + "github.com/yaoapp/yao/test" +) + +// TestSaveMessages tests batch saving messages +func TestSaveMessages(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() + + store, err := xun.NewXun(types.Setting{ + Connector: "default", + }) + if err != nil { + t.Fatalf("Failed to create store: %v", err) + } + + // Create a chat first + chat := &types.Chat{ + AssistantID: "test_assistant", + Title: "Message Test Chat", + } + err = store.CreateChat(chat) + if err != nil { + t.Fatalf("Failed to create chat: %v", err) + } + defer store.DeleteChat(chat.ChatID) + + t.Run("SaveSingleMessage", func(t *testing.T) { + messages := []*types.Message{ + { + Role: "user", + Type: "text", + Props: map[string]interface{}{"content": "Hello, world!"}, + Sequence: 1, + }, + } + + err := store.SaveMessages(chat.ChatID, messages) + if err != nil { + t.Fatalf("Failed to save message: %v", err) + } + + // Verify + retrieved, err := store.GetMessages(chat.ChatID, types.MessageFilter{}) + if err != nil { + t.Fatalf("Failed to get messages: %v", err) + } + + if len(retrieved) < 1 { + t.Fatal("Expected at least 1 message") + } + + // Find the message we just saved + var found *types.Message + for _, msg := range retrieved { + if msg.Sequence == 1 && msg.Type == "text" { + found = msg + break + } + } + + if found == nil { + t.Fatal("Could not find saved message") + } + + if found.Role != "user" { + t.Errorf("Expected role 'user', got '%s'", found.Role) + } + if found.Props["content"] != "Hello, world!" { + t.Errorf("Expected content 'Hello, world!', got '%v'", found.Props["content"]) + } + }) + + t.Run("SaveBatchMessages", func(t *testing.T) { + // Create a new chat for this test + batchChat := &types.Chat{ + AssistantID: "test_assistant", + Title: "Batch Message Test", + } + err := store.CreateChat(batchChat) + if err != nil { + t.Fatalf("Failed to create chat: %v", err) + } + defer store.DeleteChat(batchChat.ChatID) + + // Save multiple messages in one batch + messages := []*types.Message{ + { + Role: "user", + Type: "user_input", + Props: map[string]interface{}{"content": "What's the weather?"}, + Sequence: 1, + RequestID: "req_001", + AssistantID: "weather_assistant", + }, + { + Role: "assistant", + Type: "loading", + Props: map[string]interface{}{"message": "Checking weather..."}, + Sequence: 2, + RequestID: "req_001", + BlockID: "B1", + AssistantID: "weather_assistant", + }, + { + Role: "assistant", + Type: "text", + Props: map[string]interface{}{"content": "The weather is sunny, 25°C."}, + Sequence: 3, + RequestID: "req_001", + BlockID: "B1", + AssistantID: "weather_assistant", + }, + } + + err = store.SaveMessages(batchChat.ChatID, messages) + if err != nil { + t.Fatalf("Failed to save batch messages: %v", err) + } + + // Verify all messages saved + retrieved, err := store.GetMessages(batchChat.ChatID, types.MessageFilter{}) + if err != nil { + t.Fatalf("Failed to get messages: %v", err) + } + + if len(retrieved) != 3 { + t.Errorf("Expected 3 messages, got %d", len(retrieved)) + } + + // Verify order (should be by sequence) + if len(retrieved) >= 3 { + if retrieved[0].Sequence != 1 { + t.Errorf("Expected first message sequence 1, got %d", retrieved[0].Sequence) + } + if retrieved[2].Sequence != 3 { + t.Errorf("Expected last message sequence 3, got %d", retrieved[2].Sequence) + } + } + + t.Logf("Saved %d messages in single batch call", len(messages)) + }) + + t.Run("SaveMessageWithAllFields", func(t *testing.T) { + fullChat := &types.Chat{ + AssistantID: "test_assistant", + } + err := store.CreateChat(fullChat) + if err != nil { + t.Fatalf("Failed to create chat: %v", err) + } + defer store.DeleteChat(fullChat.ChatID) + + messages := []*types.Message{ + { + Role: "assistant", + Type: "tool_call", + Props: map[string]interface{}{"id": "call_123", "name": "get_weather", "arguments": `{"location":"SF"}`}, + Sequence: 1, + RequestID: "req_full", + BlockID: "B1", + ThreadID: "T1", + AssistantID: "weather_assistant", + Metadata: map[string]interface{}{"tool_call_id": "call_123", "is_tool_result": false}, + }, + } + + err = store.SaveMessages(fullChat.ChatID, messages) + if err != nil { + t.Fatalf("Failed to save message: %v", err) + } + + retrieved, err := store.GetMessages(fullChat.ChatID, types.MessageFilter{}) + if err != nil { + t.Fatalf("Failed to get messages: %v", err) + } + + if len(retrieved) != 1 { + t.Fatalf("Expected 1 message, got %d", len(retrieved)) + } + + msg := retrieved[0] + if msg.RequestID != "req_full" { + t.Errorf("Expected request_id 'req_full', got '%s'", msg.RequestID) + } + if msg.BlockID != "B1" { + t.Errorf("Expected block_id 'B1', got '%s'", msg.BlockID) + } + if msg.ThreadID != "T1" { + t.Errorf("Expected thread_id 'T1', got '%s'", msg.ThreadID) + } + if msg.AssistantID != "weather_assistant" { + t.Errorf("Expected assistant_id 'weather_assistant', got '%s'", msg.AssistantID) + } + if msg.Metadata == nil { + t.Error("Expected metadata to be set") + } else if msg.Metadata["tool_call_id"] != "call_123" { + t.Errorf("Expected metadata tool_call_id 'call_123', got '%v'", msg.Metadata["tool_call_id"]) + } + }) + + t.Run("SaveEmptyMessages", func(t *testing.T) { + err := store.SaveMessages(chat.ChatID, []*types.Message{}) + if err != nil { + t.Errorf("Expected no error for empty messages, got: %v", err) + } + }) + + t.Run("SaveMessagesWithoutChatID", func(t *testing.T) { + messages := []*types.Message{{Role: "user", Type: "text", Props: map[string]interface{}{"content": "test"}}} + err := store.SaveMessages("", messages) + if err == nil { + t.Error("Expected error when saving without chat_id") + } + }) + + t.Run("SaveMessageWithoutRole", func(t *testing.T) { + messages := []*types.Message{{Type: "text", Props: map[string]interface{}{"content": "test"}, Sequence: 1}} + err := store.SaveMessages(chat.ChatID, messages) + if err == nil { + t.Error("Expected error when saving message without role") + } + }) + + t.Run("SaveMessageWithoutType", func(t *testing.T) { + messages := []*types.Message{{Role: "user", Props: map[string]interface{}{"content": "test"}, Sequence: 1}} + err := store.SaveMessages(chat.ChatID, messages) + if err == nil { + t.Error("Expected error when saving message without type") + } + }) + + t.Run("SaveMessageWithoutProps", func(t *testing.T) { + messages := []*types.Message{{Role: "user", Type: "text", Sequence: 1}} + err := store.SaveMessages(chat.ChatID, messages) + if err == nil { + t.Error("Expected error when saving message without props") + } + }) +} + +// TestGetMessages tests retrieving messages with filters +func TestGetMessages(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() + + store, err := xun.NewXun(types.Setting{ + Connector: "default", + }) + if err != nil { + t.Fatalf("Failed to create store: %v", err) + } + + // Create chat and messages + chat := &types.Chat{ + AssistantID: "test_assistant", + } + err = store.CreateChat(chat) + if err != nil { + t.Fatalf("Failed to create chat: %v", err) + } + defer store.DeleteChat(chat.ChatID) + + // Save test messages + messages := []*types.Message{ + {Role: "user", Type: "user_input", Props: map[string]interface{}{"content": "Hello"}, Sequence: 1, RequestID: "req_001"}, + {Role: "assistant", Type: "text", Props: map[string]interface{}{"content": "Hi there!"}, Sequence: 2, RequestID: "req_001", BlockID: "B1"}, + {Role: "user", Type: "user_input", Props: map[string]interface{}{"content": "Weather?"}, Sequence: 3, RequestID: "req_002"}, + {Role: "assistant", Type: "loading", Props: map[string]interface{}{"message": "Checking..."}, Sequence: 4, RequestID: "req_002", BlockID: "B2"}, + {Role: "assistant", Type: "text", Props: map[string]interface{}{"content": "Sunny!"}, Sequence: 5, RequestID: "req_002", BlockID: "B2", ThreadID: "T1"}, + } + err = store.SaveMessages(chat.ChatID, messages) + if err != nil { + t.Fatalf("Failed to save messages: %v", err) + } + + t.Run("GetAllMessages", func(t *testing.T) { + retrieved, err := store.GetMessages(chat.ChatID, types.MessageFilter{}) + if err != nil { + t.Fatalf("Failed to get messages: %v", err) + } + + if len(retrieved) != 5 { + t.Errorf("Expected 5 messages, got %d", len(retrieved)) + } + + // Verify order by sequence + for i := 1; i < len(retrieved); i++ { + if retrieved[i].Sequence < retrieved[i-1].Sequence { + t.Error("Messages not ordered by sequence") + } + } + }) + + t.Run("FilterByRole", func(t *testing.T) { + retrieved, err := store.GetMessages(chat.ChatID, types.MessageFilter{Role: "user"}) + if err != nil { + t.Fatalf("Failed to get messages: %v", err) + } + + if len(retrieved) != 2 { + t.Errorf("Expected 2 user messages, got %d", len(retrieved)) + } + + for _, msg := range retrieved { + if msg.Role != "user" { + t.Errorf("Expected role 'user', got '%s'", msg.Role) + } + } + }) + + t.Run("FilterByRequestID", func(t *testing.T) { + retrieved, err := store.GetMessages(chat.ChatID, types.MessageFilter{RequestID: "req_002"}) + if err != nil { + t.Fatalf("Failed to get messages: %v", err) + } + + if len(retrieved) != 3 { + t.Errorf("Expected 3 messages for req_002, got %d", len(retrieved)) + } + }) + + t.Run("FilterByBlockID", func(t *testing.T) { + retrieved, err := store.GetMessages(chat.ChatID, types.MessageFilter{BlockID: "B2"}) + if err != nil { + t.Fatalf("Failed to get messages: %v", err) + } + + if len(retrieved) != 2 { + t.Errorf("Expected 2 messages in block B2, got %d", len(retrieved)) + } + }) + + t.Run("FilterByThreadID", func(t *testing.T) { + retrieved, err := store.GetMessages(chat.ChatID, types.MessageFilter{ThreadID: "T1"}) + if err != nil { + t.Fatalf("Failed to get messages: %v", err) + } + + if len(retrieved) != 1 { + t.Errorf("Expected 1 message in thread T1, got %d", len(retrieved)) + } + }) + + t.Run("FilterByType", func(t *testing.T) { + retrieved, err := store.GetMessages(chat.ChatID, types.MessageFilter{Type: "loading"}) + if err != nil { + t.Fatalf("Failed to get messages: %v", err) + } + + if len(retrieved) != 1 { + t.Errorf("Expected 1 loading message, got %d", len(retrieved)) + } + }) + + t.Run("FilterWithLimit", func(t *testing.T) { + retrieved, err := store.GetMessages(chat.ChatID, types.MessageFilter{Limit: 2}) + if err != nil { + t.Fatalf("Failed to get messages: %v", err) + } + + if len(retrieved) != 2 { + t.Errorf("Expected 2 messages with limit, got %d", len(retrieved)) + } + }) + + t.Run("FilterWithOffset", func(t *testing.T) { + retrieved, err := store.GetMessages(chat.ChatID, types.MessageFilter{Offset: 3}) + if err != nil { + t.Fatalf("Failed to get messages: %v", err) + } + + if len(retrieved) != 2 { + t.Errorf("Expected 2 messages with offset 3, got %d", len(retrieved)) + } + }) + + t.Run("FilterWithLimitAndOffset", func(t *testing.T) { + retrieved, err := store.GetMessages(chat.ChatID, types.MessageFilter{Limit: 2, Offset: 1}) + if err != nil { + t.Fatalf("Failed to get messages: %v", err) + } + + if len(retrieved) != 2 { + t.Errorf("Expected 2 messages, got %d", len(retrieved)) + } + + // Should be sequence 2 and 3 + if len(retrieved) >= 2 { + if retrieved[0].Sequence != 2 { + t.Errorf("Expected first message sequence 2, got %d", retrieved[0].Sequence) + } + } + }) + + t.Run("GetMessagesWithEmptyChatID", func(t *testing.T) { + _, err := store.GetMessages("", types.MessageFilter{}) + if err == nil { + t.Error("Expected error when getting messages without chat_id") + } + }) + + t.Run("GetMessagesFromNonExistentChat", func(t *testing.T) { + retrieved, err := store.GetMessages("nonexistent_chat", types.MessageFilter{}) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + if len(retrieved) != 0 { + t.Errorf("Expected 0 messages from non-existent chat, got %d", len(retrieved)) + } + }) +} + +// TestUpdateMessage tests updating messages +func TestUpdateMessage(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() + + store, err := xun.NewXun(types.Setting{ + Connector: "default", + }) + if err != nil { + t.Fatalf("Failed to create store: %v", err) + } + + // Create chat and message + chat := &types.Chat{ + AssistantID: "test_assistant", + } + err = store.CreateChat(chat) + if err != nil { + t.Fatalf("Failed to create chat: %v", err) + } + defer store.DeleteChat(chat.ChatID) + + messages := []*types.Message{ + { + MessageID: fmt.Sprintf("msg_%d", time.Now().UnixNano()), + Role: "assistant", + Type: "loading", + Props: map[string]interface{}{"message": "Loading..."}, + Sequence: 1, + }, + } + err = store.SaveMessages(chat.ChatID, messages) + if err != nil { + t.Fatalf("Failed to save message: %v", err) + } + + messageID := messages[0].MessageID + + t.Run("UpdateProps", func(t *testing.T) { + err := store.UpdateMessage(messageID, map[string]interface{}{ + "props": map[string]interface{}{"content": "Updated content"}, + }) + if err != nil { + t.Fatalf("Failed to update message: %v", err) + } + + retrieved, err := store.GetMessages(chat.ChatID, types.MessageFilter{}) + if err != nil { + t.Fatalf("Failed to get messages: %v", err) + } + + var found *types.Message + for _, msg := range retrieved { + if msg.MessageID == messageID { + found = msg + break + } + } + + if found == nil { + t.Fatal("Could not find updated message") + } + + if found.Props["content"] != "Updated content" { + t.Errorf("Expected props content 'Updated content', got '%v'", found.Props["content"]) + } + }) + + t.Run("UpdateType", func(t *testing.T) { + err := store.UpdateMessage(messageID, map[string]interface{}{ + "type": "text", + }) + if err != nil { + t.Fatalf("Failed to update message: %v", err) + } + + retrieved, err := store.GetMessages(chat.ChatID, types.MessageFilter{}) + if err != nil { + t.Fatalf("Failed to get messages: %v", err) + } + + var found *types.Message + for _, msg := range retrieved { + if msg.MessageID == messageID { + found = msg + break + } + } + + if found == nil { + t.Fatal("Could not find updated message") + } + + if found.Type != "text" { + t.Errorf("Expected type 'text', got '%s'", found.Type) + } + }) + + t.Run("UpdateMetadata", func(t *testing.T) { + err := store.UpdateMessage(messageID, map[string]interface{}{ + "metadata": map[string]interface{}{"updated": true}, + }) + if err != nil { + t.Fatalf("Failed to update metadata: %v", err) + } + + retrieved, err := store.GetMessages(chat.ChatID, types.MessageFilter{}) + if err != nil { + t.Fatalf("Failed to get messages: %v", err) + } + + var found *types.Message + for _, msg := range retrieved { + if msg.MessageID == messageID { + found = msg + break + } + } + + if found == nil { + t.Fatal("Could not find updated message") + } + + if found.Metadata == nil || found.Metadata["updated"] != true { + t.Errorf("Expected metadata updated=true, got %v", found.Metadata) + } + }) + + t.Run("UpdateNonExistentMessage", func(t *testing.T) { + err := store.UpdateMessage("nonexistent_msg", map[string]interface{}{ + "type": "text", + }) + if err == nil { + t.Error("Expected error when updating non-existent message") + } + }) + + t.Run("UpdateWithEmptyID", func(t *testing.T) { + err := store.UpdateMessage("", map[string]interface{}{ + "type": "text", + }) + if err == nil { + t.Error("Expected error when updating with empty ID") + } + }) + + t.Run("UpdateWithEmptyFields", func(t *testing.T) { + err := store.UpdateMessage(messageID, map[string]interface{}{}) + if err == nil { + t.Error("Expected error when updating with empty fields") + } + }) +} + +// TestDeleteMessages tests deleting messages +func TestDeleteMessages(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() + + store, err := xun.NewXun(types.Setting{ + Connector: "default", + }) + if err != nil { + t.Fatalf("Failed to create store: %v", err) + } + + t.Run("DeleteSingleMessage", func(t *testing.T) { + chat := &types.Chat{AssistantID: "test_assistant"} + err := store.CreateChat(chat) + if err != nil { + t.Fatalf("Failed to create chat: %v", err) + } + defer store.DeleteChat(chat.ChatID) + + msgID := fmt.Sprintf("msg_del_%d", time.Now().UnixNano()) + messages := []*types.Message{ + {MessageID: msgID, Role: "user", Type: "text", Props: map[string]interface{}{"content": "test"}, Sequence: 1}, + } + err = store.SaveMessages(chat.ChatID, messages) + if err != nil { + t.Fatalf("Failed to save message: %v", err) + } + + err = store.DeleteMessages(chat.ChatID, []string{msgID}) + if err != nil { + t.Fatalf("Failed to delete message: %v", err) + } + + // Verify deleted + retrieved, err := store.GetMessages(chat.ChatID, types.MessageFilter{}) + if err != nil { + t.Fatalf("Failed to get messages: %v", err) + } + + for _, msg := range retrieved { + if msg.MessageID == msgID { + t.Error("Message should have been deleted") + } + } + }) + + t.Run("DeleteMultipleMessages", func(t *testing.T) { + chat := &types.Chat{AssistantID: "test_assistant"} + err := store.CreateChat(chat) + if err != nil { + t.Fatalf("Failed to create chat: %v", err) + } + defer store.DeleteChat(chat.ChatID) + + msgID1 := fmt.Sprintf("msg_del1_%d", time.Now().UnixNano()) + msgID2 := fmt.Sprintf("msg_del2_%d", time.Now().UnixNano()) + msgID3 := fmt.Sprintf("msg_del3_%d", time.Now().UnixNano()) + + messages := []*types.Message{ + {MessageID: msgID1, Role: "user", Type: "text", Props: map[string]interface{}{"content": "1"}, Sequence: 1}, + {MessageID: msgID2, Role: "assistant", Type: "text", Props: map[string]interface{}{"content": "2"}, Sequence: 2}, + {MessageID: msgID3, Role: "user", Type: "text", Props: map[string]interface{}{"content": "3"}, Sequence: 3}, + } + err = store.SaveMessages(chat.ChatID, messages) + if err != nil { + t.Fatalf("Failed to save messages: %v", err) + } + + // Delete first two + err = store.DeleteMessages(chat.ChatID, []string{msgID1, msgID2}) + if err != nil { + t.Fatalf("Failed to delete messages: %v", err) + } + + // Verify + retrieved, err := store.GetMessages(chat.ChatID, types.MessageFilter{}) + if err != nil { + t.Fatalf("Failed to get messages: %v", err) + } + + if len(retrieved) != 1 { + t.Errorf("Expected 1 remaining message, got %d", len(retrieved)) + } + + if len(retrieved) > 0 && retrieved[0].MessageID != msgID3 { + t.Errorf("Expected remaining message to be %s, got %s", msgID3, retrieved[0].MessageID) + } + }) + + t.Run("DeleteEmptyList", func(t *testing.T) { + chat := &types.Chat{AssistantID: "test_assistant"} + err := store.CreateChat(chat) + if err != nil { + t.Fatalf("Failed to create chat: %v", err) + } + defer store.DeleteChat(chat.ChatID) + + err = store.DeleteMessages(chat.ChatID, []string{}) + if err != nil { + t.Errorf("Expected no error for empty delete list, got: %v", err) + } + }) + + t.Run("DeleteWithEmptyChatID", func(t *testing.T) { + err := store.DeleteMessages("", []string{"msg_123"}) + if err == nil { + t.Error("Expected error when deleting with empty chat_id") + } + }) +} + +// TestMessageCompleteWorkflow tests a complete message workflow +func TestMessageCompleteWorkflow(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() + + store, err := xun.NewXun(types.Setting{ + Connector: "default", + }) + if err != nil { + t.Fatalf("Failed to create store: %v", err) + } + + t.Run("CompleteWorkflow", func(t *testing.T) { + // 1. Create chat + chat := &types.Chat{ + AssistantID: "workflow_assistant", + Title: "Message Workflow Test", + } + err := store.CreateChat(chat) + if err != nil { + t.Fatalf("Failed to create chat: %v", err) + } + defer store.DeleteChat(chat.ChatID) + + // 2. Save batch messages (simulating a request) + requestID := fmt.Sprintf("req_%d", time.Now().UnixNano()) + messages := []*types.Message{ + { + Role: "user", + Type: "user_input", + Props: map[string]interface{}{"content": "What's the weather in SF?"}, + Sequence: 1, + RequestID: requestID, + AssistantID: "workflow_assistant", + }, + { + Role: "assistant", + Type: "loading", + Props: map[string]interface{}{"message": "Checking weather..."}, + Sequence: 2, + RequestID: requestID, + BlockID: "B1", + AssistantID: "workflow_assistant", + }, + { + Role: "assistant", + Type: "tool_call", + Props: map[string]interface{}{"id": "call_weather", "name": "get_weather", "arguments": `{"location":"SF"}`}, + Sequence: 3, + RequestID: requestID, + BlockID: "B1", + AssistantID: "workflow_assistant", + }, + { + Role: "assistant", + Type: "text", + Props: map[string]interface{}{"content": "The weather in San Francisco is 18°C and sunny."}, + Sequence: 4, + RequestID: requestID, + BlockID: "B1", + AssistantID: "workflow_assistant", + Metadata: map[string]interface{}{"tool_call_id": "call_weather"}, + }, + } + + err = store.SaveMessages(chat.ChatID, messages) + if err != nil { + t.Fatalf("Failed to save messages: %v", err) + } + t.Logf("Saved %d messages in single batch", len(messages)) + + // 3. Get all messages + retrieved, err := store.GetMessages(chat.ChatID, types.MessageFilter{}) + if err != nil { + t.Fatalf("Failed to get messages: %v", err) + } + + if len(retrieved) != 4 { + t.Errorf("Expected 4 messages, got %d", len(retrieved)) + } + + // 4. Filter by request + byRequest, err := store.GetMessages(chat.ChatID, types.MessageFilter{RequestID: requestID}) + if err != nil { + t.Fatalf("Failed to filter by request: %v", err) + } + + if len(byRequest) != 4 { + t.Errorf("Expected 4 messages for request, got %d", len(byRequest)) + } + + // 5. Filter by block + byBlock, err := store.GetMessages(chat.ChatID, types.MessageFilter{BlockID: "B1"}) + if err != nil { + t.Fatalf("Failed to filter by block: %v", err) + } + + if len(byBlock) != 3 { + t.Errorf("Expected 3 messages in block B1, got %d", len(byBlock)) + } + + // 6. Update loading message to text (simulating stream completion) + var loadingMsgID string + for _, msg := range retrieved { + if msg.Type == "loading" { + loadingMsgID = msg.MessageID + break + } + } + + if loadingMsgID != "" { + err = store.UpdateMessage(loadingMsgID, map[string]interface{}{ + "type": "text", + "props": map[string]interface{}{"content": "Weather check complete."}, + }) + if err != nil { + t.Fatalf("Failed to update message: %v", err) + } + } + + // 7. Delete a message + if len(retrieved) > 0 { + err = store.DeleteMessages(chat.ChatID, []string{retrieved[0].MessageID}) + if err != nil { + t.Fatalf("Failed to delete message: %v", err) + } + } + + // 8. Verify final state + final, err := store.GetMessages(chat.ChatID, types.MessageFilter{}) + if err != nil { + t.Fatalf("Failed to get final messages: %v", err) + } + + if len(final) != 3 { + t.Errorf("Expected 3 messages after delete, got %d", len(final)) + } + + t.Log("Complete message workflow passed!") + }) +} + +// TestConcurrentMessages tests concurrent message storage +func TestConcurrentMessages(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() + + store, err := xun.NewXun(types.Setting{ + Connector: "default", + }) + if err != nil { + t.Fatalf("Failed to create store: %v", err) + } + + t.Run("ConcurrentThreadMessages", func(t *testing.T) { + chat := &types.Chat{AssistantID: "test_assistant"} + err := store.CreateChat(chat) + if err != nil { + t.Fatalf("Failed to create chat: %v", err) + } + defer store.DeleteChat(chat.ChatID) + + // Simulate concurrent operations with different threads + messages := []*types.Message{ + {Role: "assistant", Type: "text", Props: map[string]interface{}{"content": "Weather result"}, Sequence: 1, BlockID: "B1", ThreadID: "T1"}, + {Role: "assistant", Type: "text", Props: map[string]interface{}{"content": "News result"}, Sequence: 2, BlockID: "B1", ThreadID: "T2"}, + {Role: "assistant", Type: "text", Props: map[string]interface{}{"content": "Stock result"}, Sequence: 3, BlockID: "B1", ThreadID: "T3"}, + {Role: "assistant", Type: "text", Props: map[string]interface{}{"content": "Summary"}, Sequence: 4, BlockID: "B2"}, + } + + err = store.SaveMessages(chat.ChatID, messages) + if err != nil { + t.Fatalf("Failed to save concurrent messages: %v", err) + } + + // Verify all saved + all, err := store.GetMessages(chat.ChatID, types.MessageFilter{}) + if err != nil { + t.Fatalf("Failed to get messages: %v", err) + } + + if len(all) != 4 { + t.Errorf("Expected 4 messages, got %d", len(all)) + } + + // Filter by thread + t1Messages, err := store.GetMessages(chat.ChatID, types.MessageFilter{ThreadID: "T1"}) + if err != nil { + t.Fatalf("Failed to filter by thread: %v", err) + } + + if len(t1Messages) != 1 { + t.Errorf("Expected 1 message in thread T1, got %d", len(t1Messages)) + } + + // Filter by block + b1Messages, err := store.GetMessages(chat.ChatID, types.MessageFilter{BlockID: "B1"}) + if err != nil { + t.Fatalf("Failed to filter by block: %v", err) + } + + if len(b1Messages) != 3 { + t.Errorf("Expected 3 messages in block B1, got %d", len(b1Messages)) + } + + t.Log("Concurrent thread messages test passed!") + }) +} diff --git a/agent/store/xun/resume.go b/agent/store/xun/resume.go index 0f39ef04..29315c6d 100644 --- a/agent/store/xun/resume.go +++ b/agent/store/xun/resume.go @@ -1,6 +1,11 @@ package xun import ( + "fmt" + "time" + + "github.com/google/uuid" + jsoniter "github.com/json-iterator/go" "github.com/yaoapp/yao/agent/store/types" ) @@ -8,41 +13,367 @@ import ( // Resume Management (only called on failure/interrupt) // ============================================================================= -// SaveResume batch saves resume records +// SaveResume batch saves resume records using a single database call // Only called when request is interrupted or failed func (store *Xun) SaveResume(records []*types.Resume) error { - // TODO: implement - return nil + if len(records) == 0 { + return nil // Nothing to save + } + + // Prepare batch insert data + now := time.Now() + rows := make([]map[string]interface{}, 0, len(records)) + + for _, record := range records { + if record == nil { + continue + } + + // Generate resume_id if not provided + resumeID := record.ResumeID + if resumeID == "" { + resumeID = uuid.New().String() + } + + // Validate required fields + if record.ChatID == "" { + return fmt.Errorf("chat_id is required") + } + if record.RequestID == "" { + return fmt.Errorf("request_id is required") + } + if record.AssistantID == "" { + return fmt.Errorf("assistant_id is required") + } + if record.StackID == "" { + return fmt.Errorf("stack_id is required") + } + if record.Type == "" { + return fmt.Errorf("type is required") + } + if record.Status == "" { + return fmt.Errorf("status is required") + } + + // Build row with all fields (including nullable ones for consistent batch insert) + row := map[string]interface{}{ + "resume_id": resumeID, + "chat_id": record.ChatID, + "request_id": record.RequestID, + "assistant_id": record.AssistantID, + "stack_id": record.StackID, + "stack_parent_id": nil, + "stack_depth": record.StackDepth, + "type": record.Type, + "status": record.Status, + "input": nil, + "output": nil, + "space_snapshot": nil, + "error": nil, + "sequence": record.Sequence, + "metadata": nil, + "created_at": now, + "updated_at": now, + } + + // Set nullable fields if they have values + if record.StackParentID != "" { + row["stack_parent_id"] = record.StackParentID + } + if record.Input != nil { + inputJSON, err := jsoniter.MarshalToString(record.Input) + if err != nil { + return fmt.Errorf("failed to marshal input: %w", err) + } + row["input"] = inputJSON + } + if record.Output != nil { + outputJSON, err := jsoniter.MarshalToString(record.Output) + if err != nil { + return fmt.Errorf("failed to marshal output: %w", err) + } + row["output"] = outputJSON + } + if record.SpaceSnapshot != nil { + snapshotJSON, err := jsoniter.MarshalToString(record.SpaceSnapshot) + if err != nil { + return fmt.Errorf("failed to marshal space_snapshot: %w", err) + } + row["space_snapshot"] = snapshotJSON + } + if record.Error != "" { + row["error"] = record.Error + } + if record.Metadata != nil { + metadataJSON, err := jsoniter.MarshalToString(record.Metadata) + if err != nil { + return fmt.Errorf("failed to marshal metadata: %w", err) + } + row["metadata"] = metadataJSON + } + + rows = append(rows, row) + } + + if len(rows) == 0 { + return nil + } + + // Single batch insert - one database call for all records + return store.newQueryResume().Insert(rows) } // GetResume retrieves all resume records for a chat func (store *Xun) GetResume(chatID string) ([]*types.Resume, error) { - // TODO: implement - return nil, nil + if chatID == "" { + return nil, fmt.Errorf("chat_id is required") + } + + rows, err := store.newQueryResume(). + Where("chat_id", chatID). + WhereNull("deleted_at"). + OrderBy("sequence", "asc"). + Get() + if err != nil { + return nil, err + } + + records := make([]*types.Resume, 0, len(rows)) + for _, row := range rows { + data := row.ToMap() + if data == nil || data["resume_id"] == nil { + continue + } + + record, err := store.rowToResume(data) + if err != nil { + continue + } + records = append(records, record) + } + + return records, nil } // GetLastResume retrieves the last (most recent) resume record for a chat func (store *Xun) GetLastResume(chatID string) (*types.Resume, error) { - // TODO: implement - return nil, nil + if chatID == "" { + return nil, fmt.Errorf("chat_id is required") + } + + row, err := store.newQueryResume(). + Where("chat_id", chatID). + WhereNull("deleted_at"). + OrderBy("sequence", "desc"). + First() + if err != nil { + return nil, err + } + + if row == nil { + return nil, nil // No resume records found + } + + data := row.ToMap() + if len(data) == 0 || data["resume_id"] == nil { + return nil, nil + } + + return store.rowToResume(data) } // GetResumeByStackID retrieves resume records for a specific stack func (store *Xun) GetResumeByStackID(stackID string) ([]*types.Resume, error) { - // TODO: implement - return nil, nil + if stackID == "" { + return nil, fmt.Errorf("stack_id is required") + } + + rows, err := store.newQueryResume(). + Where("stack_id", stackID). + WhereNull("deleted_at"). + OrderBy("sequence", "asc"). + Get() + if err != nil { + return nil, err + } + + records := make([]*types.Resume, 0, len(rows)) + for _, row := range rows { + data := row.ToMap() + if data == nil || data["resume_id"] == nil { + continue + } + + record, err := store.rowToResume(data) + if err != nil { + continue + } + records = append(records, record) + } + + return records, nil } // GetStackPath returns the stack path from root to the given stack // Returns: [root_stack_id, ..., current_stack_id] func (store *Xun) GetStackPath(stackID string) ([]string, error) { - // TODO: implement - return nil, nil + if stackID == "" { + return nil, fmt.Errorf("stack_id is required") + } + + path := []string{stackID} + currentStackID := stackID + + // Walk up the stack tree by following stack_parent_id + for { + row, err := store.newQueryResume(). + Where("stack_id", currentStackID). + WhereNull("deleted_at"). + First() + if err != nil { + return nil, err + } + + if row == nil { + break + } + + data := row.ToMap() + parentID := getString(data, "stack_parent_id") + if parentID == "" { + break // Reached root + } + + // Prepend parent to path + path = append([]string{parentID}, path...) + currentStackID = parentID + } + + return path, nil } -// DeleteResume deletes all resume records for a chat +// DeleteResume soft deletes all resume records for a chat // Called after successful resume to clean up func (store *Xun) DeleteResume(chatID string) error { - // TODO: implement - return nil + if chatID == "" { + return fmt.Errorf("chat_id is required") + } + + _, err := store.newQueryResume(). + Where("chat_id", chatID). + WhereNull("deleted_at"). + Update(map[string]interface{}{ + "deleted_at": time.Now(), + "updated_at": time.Now(), + }) + + return err +} + +// GetResumeByRequestID retrieves resume records for a specific request +func (store *Xun) GetResumeByRequestID(requestID string) ([]*types.Resume, error) { + if requestID == "" { + return nil, fmt.Errorf("request_id is required") + } + + rows, err := store.newQueryResume(). + Where("request_id", requestID). + WhereNull("deleted_at"). + OrderBy("sequence", "asc"). + Get() + if err != nil { + return nil, err + } + + records := make([]*types.Resume, 0, len(rows)) + for _, row := range rows { + data := row.ToMap() + if data == nil || data["resume_id"] == nil { + continue + } + + record, err := store.rowToResume(data) + if err != nil { + continue + } + records = append(records, record) + } + + return records, nil +} + +// ============================================================================= +// Helper Functions +// ============================================================================= + +// rowToResume converts a database row to a Resume struct +func (store *Xun) rowToResume(data map[string]interface{}) (*types.Resume, error) { + record := &types.Resume{ + ResumeID: getString(data, "resume_id"), + ChatID: getString(data, "chat_id"), + RequestID: getString(data, "request_id"), + AssistantID: getString(data, "assistant_id"), + StackID: getString(data, "stack_id"), + StackParentID: getString(data, "stack_parent_id"), + StackDepth: getInt(data, "stack_depth"), + Type: getString(data, "type"), + Status: getString(data, "status"), + Error: getString(data, "error"), + Sequence: getInt(data, "sequence"), + } + + // Handle timestamps + if createdAt := getTime(data, "created_at"); createdAt != nil { + record.CreatedAt = *createdAt + } + if updatedAt := getTime(data, "updated_at"); updatedAt != nil { + record.UpdatedAt = *updatedAt + } + + // Handle JSON fields + if input := data["input"]; input != nil { + if inputStr, ok := input.(string); ok && inputStr != "" { + var inputMap map[string]interface{} + if err := jsoniter.UnmarshalFromString(inputStr, &inputMap); err == nil { + record.Input = inputMap + } + } else if inputMap, ok := input.(map[string]interface{}); ok { + record.Input = inputMap + } + } + + if output := data["output"]; output != nil { + if outputStr, ok := output.(string); ok && outputStr != "" { + var outputMap map[string]interface{} + if err := jsoniter.UnmarshalFromString(outputStr, &outputMap); err == nil { + record.Output = outputMap + } + } else if outputMap, ok := output.(map[string]interface{}); ok { + record.Output = outputMap + } + } + + if snapshot := data["space_snapshot"]; snapshot != nil { + if snapshotStr, ok := snapshot.(string); ok && snapshotStr != "" { + var snapshotMap map[string]interface{} + if err := jsoniter.UnmarshalFromString(snapshotStr, &snapshotMap); err == nil { + record.SpaceSnapshot = snapshotMap + } + } else if snapshotMap, ok := snapshot.(map[string]interface{}); ok { + record.SpaceSnapshot = snapshotMap + } + } + + if metadata := data["metadata"]; metadata != nil { + if metaStr, ok := metadata.(string); ok && metaStr != "" { + var metaMap map[string]interface{} + if err := jsoniter.UnmarshalFromString(metaStr, &metaMap); err == nil { + record.Metadata = metaMap + } + } else if metaMap, ok := metadata.(map[string]interface{}); ok { + record.Metadata = metaMap + } + } + + return record, nil } diff --git a/agent/store/xun/resume_test.go b/agent/store/xun/resume_test.go new file mode 100644 index 00000000..f0866a7b --- /dev/null +++ b/agent/store/xun/resume_test.go @@ -0,0 +1,839 @@ +package xun_test + +import ( + "fmt" + "testing" + "time" + + "github.com/yaoapp/yao/agent/store/types" + "github.com/yaoapp/yao/agent/store/xun" + "github.com/yaoapp/yao/config" + "github.com/yaoapp/yao/test" +) + +// TestSaveResume tests batch saving resume records +func TestSaveResume(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() + + store, err := xun.NewXun(types.Setting{ + Connector: "default", + }) + if err != nil { + t.Fatalf("Failed to create store: %v", err) + } + + // Create a chat first + chat := &types.Chat{ + AssistantID: "test_assistant", + Title: "Resume Test Chat", + } + err = store.CreateChat(chat) + if err != nil { + t.Fatalf("Failed to create chat: %v", err) + } + defer store.DeleteChat(chat.ChatID) + + t.Run("SaveSingleRecord", func(t *testing.T) { + requestID := fmt.Sprintf("req_%d", time.Now().UnixNano()) + records := []*types.Resume{ + { + ChatID: chat.ChatID, + RequestID: requestID, + AssistantID: "test_assistant", + StackID: "stack_001", + StackDepth: 0, + Type: types.ResumeTypeLLM, + Status: types.ResumeStatusInterrupted, + Sequence: 1, + }, + } + + err := store.SaveResume(records) + if err != nil { + t.Fatalf("Failed to save resume record: %v", err) + } + + // Verify + retrieved, err := store.GetResume(chat.ChatID) + if err != nil { + t.Fatalf("Failed to get resume records: %v", err) + } + + found := false + for _, r := range retrieved { + if r.RequestID == requestID { + found = true + if r.Type != types.ResumeTypeLLM { + t.Errorf("Expected type '%s', got '%s'", types.ResumeTypeLLM, r.Type) + } + if r.Status != types.ResumeStatusInterrupted { + t.Errorf("Expected status '%s', got '%s'", types.ResumeStatusInterrupted, r.Status) + } + break + } + } + + if !found { + t.Error("Could not find saved resume record") + } + + // Clean up + store.DeleteResume(chat.ChatID) + }) + + t.Run("SaveBatchRecords", func(t *testing.T) { + // Create a new chat for this test + batchChat := &types.Chat{ + AssistantID: "test_assistant", + } + err := store.CreateChat(batchChat) + if err != nil { + t.Fatalf("Failed to create chat: %v", err) + } + defer store.DeleteChat(batchChat.ChatID) + + requestID := fmt.Sprintf("req_%d", time.Now().UnixNano()) + records := []*types.Resume{ + { + ChatID: batchChat.ChatID, + RequestID: requestID, + AssistantID: "test_assistant", + StackID: "stack_001", + StackDepth: 0, + Type: types.ResumeTypeInput, + Status: types.ResumeStatusInterrupted, + Sequence: 1, + }, + { + ChatID: batchChat.ChatID, + RequestID: requestID, + AssistantID: "test_assistant", + StackID: "stack_001", + StackDepth: 0, + Type: types.ResumeTypeHookCreate, + Status: types.ResumeStatusInterrupted, + Sequence: 2, + }, + { + ChatID: batchChat.ChatID, + RequestID: requestID, + AssistantID: "test_assistant", + StackID: "stack_001", + StackDepth: 0, + Type: types.ResumeTypeLLM, + Status: types.ResumeStatusFailed, + Sequence: 3, + Error: "Connection timeout", + }, + } + + err = store.SaveResume(records) + if err != nil { + t.Fatalf("Failed to save batch resume records: %v", err) + } + + // Verify all records saved + retrieved, err := store.GetResume(batchChat.ChatID) + if err != nil { + t.Fatalf("Failed to get resume records: %v", err) + } + + if len(retrieved) != 3 { + t.Errorf("Expected 3 records, got %d", len(retrieved)) + } + + // Verify order (should be by sequence) + if len(retrieved) >= 3 { + if retrieved[0].Sequence != 1 { + t.Errorf("Expected first record sequence 1, got %d", retrieved[0].Sequence) + } + if retrieved[2].Sequence != 3 { + t.Errorf("Expected last record sequence 3, got %d", retrieved[2].Sequence) + } + if retrieved[2].Error != "Connection timeout" { + t.Errorf("Expected error 'Connection timeout', got '%s'", retrieved[2].Error) + } + } + + t.Logf("Saved %d resume records in single batch call", len(records)) + }) + + t.Run("SaveRecordWithAllFields", func(t *testing.T) { + fullChat := &types.Chat{ + AssistantID: "test_assistant", + } + err := store.CreateChat(fullChat) + if err != nil { + t.Fatalf("Failed to create chat: %v", err) + } + defer store.DeleteChat(fullChat.ChatID) + + requestID := fmt.Sprintf("req_%d", time.Now().UnixNano()) + records := []*types.Resume{ + { + ChatID: fullChat.ChatID, + RequestID: requestID, + AssistantID: "test_assistant", + StackID: "stack_001", + StackParentID: "stack_000", + StackDepth: 1, + Type: types.ResumeTypeDelegate, + Status: types.ResumeStatusInterrupted, + Input: map[string]interface{}{"agent_id": "sub_agent", "messages": []interface{}{}}, + Output: map[string]interface{}{"partial": true}, + SpaceSnapshot: map[string]interface{}{"key1": "value1", "key2": 123}, + Error: "User cancelled", + Sequence: 1, + Metadata: map[string]interface{}{"retry_count": 0}, + }, + } + + err = store.SaveResume(records) + if err != nil { + t.Fatalf("Failed to save record: %v", err) + } + + retrieved, err := store.GetResume(fullChat.ChatID) + if err != nil { + t.Fatalf("Failed to get records: %v", err) + } + + if len(retrieved) != 1 { + t.Fatalf("Expected 1 record, got %d", len(retrieved)) + } + + r := retrieved[0] + if r.StackParentID != "stack_000" { + t.Errorf("Expected stack_parent_id 'stack_000', got '%s'", r.StackParentID) + } + if r.StackDepth != 1 { + t.Errorf("Expected stack_depth 1, got %d", r.StackDepth) + } + if r.Input == nil { + t.Error("Expected input to be set") + } + if r.Output == nil { + t.Error("Expected output to be set") + } + if r.SpaceSnapshot == nil { + t.Error("Expected space_snapshot to be set") + } else if r.SpaceSnapshot["key1"] != "value1" { + t.Errorf("Expected space_snapshot key1='value1', got '%v'", r.SpaceSnapshot["key1"]) + } + if r.Metadata == nil { + t.Error("Expected metadata to be set") + } + }) + + t.Run("SaveEmptyRecords", func(t *testing.T) { + err := store.SaveResume([]*types.Resume{}) + if err != nil { + t.Errorf("Expected no error for empty records, got: %v", err) + } + }) + + t.Run("SaveRecordWithoutChatID", func(t *testing.T) { + records := []*types.Resume{{RequestID: "req", AssistantID: "ast", StackID: "stk", Type: "llm", Status: "failed", Sequence: 1}} + err := store.SaveResume(records) + if err == nil { + t.Error("Expected error when saving without chat_id") + } + }) + + t.Run("SaveRecordWithoutRequestID", func(t *testing.T) { + records := []*types.Resume{{ChatID: chat.ChatID, AssistantID: "ast", StackID: "stk", Type: "llm", Status: "failed", Sequence: 1}} + err := store.SaveResume(records) + if err == nil { + t.Error("Expected error when saving without request_id") + } + }) + + t.Run("SaveRecordWithoutAssistantID", func(t *testing.T) { + records := []*types.Resume{{ChatID: chat.ChatID, RequestID: "req", StackID: "stk", Type: "llm", Status: "failed", Sequence: 1}} + err := store.SaveResume(records) + if err == nil { + t.Error("Expected error when saving without assistant_id") + } + }) + + t.Run("SaveRecordWithoutStackID", func(t *testing.T) { + records := []*types.Resume{{ChatID: chat.ChatID, RequestID: "req", AssistantID: "ast", Type: "llm", Status: "failed", Sequence: 1}} + err := store.SaveResume(records) + if err == nil { + t.Error("Expected error when saving without stack_id") + } + }) + + t.Run("SaveRecordWithoutType", func(t *testing.T) { + records := []*types.Resume{{ChatID: chat.ChatID, RequestID: "req", AssistantID: "ast", StackID: "stk", Status: "failed", Sequence: 1}} + err := store.SaveResume(records) + if err == nil { + t.Error("Expected error when saving without type") + } + }) + + t.Run("SaveRecordWithoutStatus", func(t *testing.T) { + records := []*types.Resume{{ChatID: chat.ChatID, RequestID: "req", AssistantID: "ast", StackID: "stk", Type: "llm", Sequence: 1}} + err := store.SaveResume(records) + if err == nil { + t.Error("Expected error when saving without status") + } + }) +} + +// TestGetResume tests retrieving resume records +func TestGetResume(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() + + store, err := xun.NewXun(types.Setting{ + Connector: "default", + }) + if err != nil { + t.Fatalf("Failed to create store: %v", err) + } + + // Create chat and resume records + chat := &types.Chat{ + AssistantID: "test_assistant", + } + err = store.CreateChat(chat) + if err != nil { + t.Fatalf("Failed to create chat: %v", err) + } + defer store.DeleteChat(chat.ChatID) + + requestID := fmt.Sprintf("req_%d", time.Now().UnixNano()) + records := []*types.Resume{ + {ChatID: chat.ChatID, RequestID: requestID, AssistantID: "ast1", StackID: "stk1", Type: types.ResumeTypeInput, Status: types.ResumeStatusInterrupted, Sequence: 1}, + {ChatID: chat.ChatID, RequestID: requestID, AssistantID: "ast1", StackID: "stk1", Type: types.ResumeTypeHookCreate, Status: types.ResumeStatusInterrupted, Sequence: 2}, + {ChatID: chat.ChatID, RequestID: requestID, AssistantID: "ast1", StackID: "stk1", Type: types.ResumeTypeLLM, Status: types.ResumeStatusFailed, Sequence: 3}, + } + err = store.SaveResume(records) + if err != nil { + t.Fatalf("Failed to save records: %v", err) + } + defer store.DeleteResume(chat.ChatID) + + t.Run("GetAllRecords", func(t *testing.T) { + retrieved, err := store.GetResume(chat.ChatID) + if err != nil { + t.Fatalf("Failed to get records: %v", err) + } + + if len(retrieved) != 3 { + t.Errorf("Expected 3 records, got %d", len(retrieved)) + } + + // Verify order by sequence + for i := 1; i < len(retrieved); i++ { + if retrieved[i].Sequence < retrieved[i-1].Sequence { + t.Error("Records not ordered by sequence") + } + } + }) + + t.Run("GetRecordsWithEmptyChatID", func(t *testing.T) { + _, err := store.GetResume("") + if err == nil { + t.Error("Expected error when getting records without chat_id") + } + }) + + t.Run("GetRecordsFromNonExistentChat", func(t *testing.T) { + retrieved, err := store.GetResume("nonexistent_chat") + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + if len(retrieved) != 0 { + t.Errorf("Expected 0 records from non-existent chat, got %d", len(retrieved)) + } + }) +} + +// TestGetLastResume tests retrieving the last resume record +func TestGetLastResume(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() + + store, err := xun.NewXun(types.Setting{ + Connector: "default", + }) + if err != nil { + t.Fatalf("Failed to create store: %v", err) + } + + chat := &types.Chat{ + AssistantID: "test_assistant", + } + err = store.CreateChat(chat) + if err != nil { + t.Fatalf("Failed to create chat: %v", err) + } + defer store.DeleteChat(chat.ChatID) + + t.Run("GetLastRecordFromMultiple", func(t *testing.T) { + requestID := fmt.Sprintf("req_%d", time.Now().UnixNano()) + records := []*types.Resume{ + {ChatID: chat.ChatID, RequestID: requestID, AssistantID: "ast", StackID: "stk", Type: types.ResumeTypeInput, Status: types.ResumeStatusInterrupted, Sequence: 1}, + {ChatID: chat.ChatID, RequestID: requestID, AssistantID: "ast", StackID: "stk", Type: types.ResumeTypeHookCreate, Status: types.ResumeStatusInterrupted, Sequence: 2}, + {ChatID: chat.ChatID, RequestID: requestID, AssistantID: "ast", StackID: "stk", Type: types.ResumeTypeLLM, Status: types.ResumeStatusFailed, Sequence: 3, Error: "Last error"}, + } + err := store.SaveResume(records) + if err != nil { + t.Fatalf("Failed to save records: %v", err) + } + defer store.DeleteResume(chat.ChatID) + + last, err := store.GetLastResume(chat.ChatID) + if err != nil { + t.Fatalf("Failed to get last record: %v", err) + } + + if last == nil { + t.Fatal("Expected last record, got nil") + } + + if last.Sequence != 3 { + t.Errorf("Expected sequence 3, got %d", last.Sequence) + } + if last.Type != types.ResumeTypeLLM { + t.Errorf("Expected type '%s', got '%s'", types.ResumeTypeLLM, last.Type) + } + if last.Error != "Last error" { + t.Errorf("Expected error 'Last error', got '%s'", last.Error) + } + }) + + t.Run("GetLastRecordFromEmpty", func(t *testing.T) { + emptyChat := &types.Chat{AssistantID: "test_assistant"} + err := store.CreateChat(emptyChat) + if err != nil { + t.Fatalf("Failed to create chat: %v", err) + } + defer store.DeleteChat(emptyChat.ChatID) + + last, err := store.GetLastResume(emptyChat.ChatID) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + + if last != nil { + t.Error("Expected nil for empty chat, got record") + } + }) + + t.Run("GetLastRecordWithEmptyChatID", func(t *testing.T) { + _, err := store.GetLastResume("") + if err == nil { + t.Error("Expected error when getting last record without chat_id") + } + }) +} + +// TestGetResumeByStackID tests retrieving records by stack ID +func TestGetResumeByStackID(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() + + store, err := xun.NewXun(types.Setting{ + Connector: "default", + }) + if err != nil { + t.Fatalf("Failed to create store: %v", err) + } + + chat := &types.Chat{ + AssistantID: "test_assistant", + } + err = store.CreateChat(chat) + if err != nil { + t.Fatalf("Failed to create chat: %v", err) + } + defer store.DeleteChat(chat.ChatID) + + requestID := fmt.Sprintf("req_%d", time.Now().UnixNano()) + records := []*types.Resume{ + {ChatID: chat.ChatID, RequestID: requestID, AssistantID: "ast1", StackID: "stack_A", Type: types.ResumeTypeInput, Status: types.ResumeStatusInterrupted, Sequence: 1}, + {ChatID: chat.ChatID, RequestID: requestID, AssistantID: "ast1", StackID: "stack_A", Type: types.ResumeTypeLLM, Status: types.ResumeStatusInterrupted, Sequence: 2}, + {ChatID: chat.ChatID, RequestID: requestID, AssistantID: "ast2", StackID: "stack_B", StackParentID: "stack_A", StackDepth: 1, Type: types.ResumeTypeDelegate, Status: types.ResumeStatusFailed, Sequence: 3}, + } + err = store.SaveResume(records) + if err != nil { + t.Fatalf("Failed to save records: %v", err) + } + defer store.DeleteResume(chat.ChatID) + + t.Run("GetRecordsByStackA", func(t *testing.T) { + retrieved, err := store.GetResumeByStackID("stack_A") + if err != nil { + t.Fatalf("Failed to get records: %v", err) + } + + if len(retrieved) != 2 { + t.Errorf("Expected 2 records for stack_A, got %d", len(retrieved)) + } + }) + + t.Run("GetRecordsByStackB", func(t *testing.T) { + retrieved, err := store.GetResumeByStackID("stack_B") + if err != nil { + t.Fatalf("Failed to get records: %v", err) + } + + if len(retrieved) != 1 { + t.Errorf("Expected 1 record for stack_B, got %d", len(retrieved)) + } + + if len(retrieved) > 0 { + if retrieved[0].StackParentID != "stack_A" { + t.Errorf("Expected stack_parent_id 'stack_A', got '%s'", retrieved[0].StackParentID) + } + if retrieved[0].StackDepth != 1 { + t.Errorf("Expected stack_depth 1, got %d", retrieved[0].StackDepth) + } + } + }) + + t.Run("GetRecordsByNonExistentStack", func(t *testing.T) { + retrieved, err := store.GetResumeByStackID("nonexistent_stack") + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + if len(retrieved) != 0 { + t.Errorf("Expected 0 records, got %d", len(retrieved)) + } + }) + + t.Run("GetRecordsByEmptyStackID", func(t *testing.T) { + _, err := store.GetResumeByStackID("") + if err == nil { + t.Error("Expected error when getting records without stack_id") + } + }) +} + +// TestGetStackPath tests retrieving the stack path +func TestGetStackPath(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() + + store, err := xun.NewXun(types.Setting{ + Connector: "default", + }) + if err != nil { + t.Fatalf("Failed to create store: %v", err) + } + + chat := &types.Chat{ + AssistantID: "test_assistant", + } + err = store.CreateChat(chat) + if err != nil { + t.Fatalf("Failed to create chat: %v", err) + } + defer store.DeleteChat(chat.ChatID) + + // Create a nested stack structure: root -> child -> grandchild + requestID := fmt.Sprintf("req_%d", time.Now().UnixNano()) + records := []*types.Resume{ + {ChatID: chat.ChatID, RequestID: requestID, AssistantID: "ast1", StackID: "root_stack", Type: types.ResumeTypeInput, Status: types.ResumeStatusInterrupted, Sequence: 1}, + {ChatID: chat.ChatID, RequestID: requestID, AssistantID: "ast2", StackID: "child_stack", StackParentID: "root_stack", StackDepth: 1, Type: types.ResumeTypeDelegate, Status: types.ResumeStatusInterrupted, Sequence: 2}, + {ChatID: chat.ChatID, RequestID: requestID, AssistantID: "ast3", StackID: "grandchild_stack", StackParentID: "child_stack", StackDepth: 2, Type: types.ResumeTypeLLM, Status: types.ResumeStatusFailed, Sequence: 3}, + } + err = store.SaveResume(records) + if err != nil { + t.Fatalf("Failed to save records: %v", err) + } + defer store.DeleteResume(chat.ChatID) + + t.Run("GetPathFromGrandchild", func(t *testing.T) { + path, err := store.GetStackPath("grandchild_stack") + if err != nil { + t.Fatalf("Failed to get stack path: %v", err) + } + + if len(path) != 3 { + t.Errorf("Expected path length 3, got %d", len(path)) + } + + if len(path) >= 3 { + if path[0] != "root_stack" { + t.Errorf("Expected first element 'root_stack', got '%s'", path[0]) + } + if path[1] != "child_stack" { + t.Errorf("Expected second element 'child_stack', got '%s'", path[1]) + } + if path[2] != "grandchild_stack" { + t.Errorf("Expected third element 'grandchild_stack', got '%s'", path[2]) + } + } + + t.Logf("Stack path: %v", path) + }) + + t.Run("GetPathFromChild", func(t *testing.T) { + path, err := store.GetStackPath("child_stack") + if err != nil { + t.Fatalf("Failed to get stack path: %v", err) + } + + if len(path) != 2 { + t.Errorf("Expected path length 2, got %d", len(path)) + } + + if len(path) >= 2 { + if path[0] != "root_stack" { + t.Errorf("Expected first element 'root_stack', got '%s'", path[0]) + } + if path[1] != "child_stack" { + t.Errorf("Expected second element 'child_stack', got '%s'", path[1]) + } + } + }) + + t.Run("GetPathFromRoot", func(t *testing.T) { + path, err := store.GetStackPath("root_stack") + if err != nil { + t.Fatalf("Failed to get stack path: %v", err) + } + + if len(path) != 1 { + t.Errorf("Expected path length 1, got %d", len(path)) + } + + if len(path) >= 1 && path[0] != "root_stack" { + t.Errorf("Expected 'root_stack', got '%s'", path[0]) + } + }) + + t.Run("GetPathWithEmptyStackID", func(t *testing.T) { + _, err := store.GetStackPath("") + if err == nil { + t.Error("Expected error when getting path without stack_id") + } + }) +} + +// TestDeleteResume tests deleting resume records +func TestDeleteResume(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() + + store, err := xun.NewXun(types.Setting{ + Connector: "default", + }) + if err != nil { + t.Fatalf("Failed to create store: %v", err) + } + + t.Run("DeleteExistingRecords", func(t *testing.T) { + chat := &types.Chat{AssistantID: "test_assistant"} + err := store.CreateChat(chat) + if err != nil { + t.Fatalf("Failed to create chat: %v", err) + } + defer store.DeleteChat(chat.ChatID) + + requestID := fmt.Sprintf("req_%d", time.Now().UnixNano()) + records := []*types.Resume{ + {ChatID: chat.ChatID, RequestID: requestID, AssistantID: "ast", StackID: "stk", Type: types.ResumeTypeLLM, Status: types.ResumeStatusFailed, Sequence: 1}, + {ChatID: chat.ChatID, RequestID: requestID, AssistantID: "ast", StackID: "stk", Type: types.ResumeTypeTool, Status: types.ResumeStatusFailed, Sequence: 2}, + } + err = store.SaveResume(records) + if err != nil { + t.Fatalf("Failed to save records: %v", err) + } + + // Delete + err = store.DeleteResume(chat.ChatID) + if err != nil { + t.Fatalf("Failed to delete records: %v", err) + } + + // Verify deleted + retrieved, err := store.GetResume(chat.ChatID) + if err != nil { + t.Fatalf("Failed to get records: %v", err) + } + + if len(retrieved) != 0 { + t.Errorf("Expected 0 records after delete, got %d", len(retrieved)) + } + }) + + t.Run("DeleteFromEmptyChat", func(t *testing.T) { + chat := &types.Chat{AssistantID: "test_assistant"} + err := store.CreateChat(chat) + if err != nil { + t.Fatalf("Failed to create chat: %v", err) + } + defer store.DeleteChat(chat.ChatID) + + // Delete from chat with no records - should not error + err = store.DeleteResume(chat.ChatID) + if err != nil { + t.Errorf("Expected no error when deleting from empty chat, got: %v", err) + } + }) + + t.Run("DeleteWithEmptyChatID", func(t *testing.T) { + err := store.DeleteResume("") + if err == nil { + t.Error("Expected error when deleting with empty chat_id") + } + }) +} + +// TestResumeCompleteWorkflow tests a complete resume/retry workflow +func TestResumeCompleteWorkflow(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() + + store, err := xun.NewXun(types.Setting{ + Connector: "default", + }) + if err != nil { + t.Fatalf("Failed to create store: %v", err) + } + + t.Run("CompleteA2AWorkflow", func(t *testing.T) { + // Create chat + chat := &types.Chat{ + AssistantID: "main_assistant", + Title: "A2A Workflow Test", + } + err := store.CreateChat(chat) + if err != nil { + t.Fatalf("Failed to create chat: %v", err) + } + defer store.DeleteChat(chat.ChatID) + + // Simulate A2A call that gets interrupted + // Main assistant -> Sub assistant (interrupted during LLM call) + requestID := fmt.Sprintf("req_%d", time.Now().UnixNano()) + records := []*types.Resume{ + // Main assistant steps + { + ChatID: chat.ChatID, + RequestID: requestID, + AssistantID: "main_assistant", + StackID: "main_stack", + StackDepth: 0, + Type: types.ResumeTypeInput, + Status: types.ResumeStatusInterrupted, + Input: map[string]interface{}{"messages": []interface{}{map[string]interface{}{"role": "user", "content": "Analyze this"}}}, + Sequence: 1, + }, + { + ChatID: chat.ChatID, + RequestID: requestID, + AssistantID: "main_assistant", + StackID: "main_stack", + StackDepth: 0, + Type: types.ResumeTypeDelegate, + Status: types.ResumeStatusInterrupted, + SpaceSnapshot: map[string]interface{}{"task": "analyze", "data_id": "123"}, + Sequence: 2, + }, + // Sub assistant steps + { + ChatID: chat.ChatID, + RequestID: requestID, + AssistantID: "sub_assistant", + StackID: "sub_stack", + StackParentID: "main_stack", + StackDepth: 1, + Type: types.ResumeTypeInput, + Status: types.ResumeStatusInterrupted, + Sequence: 3, + }, + { + ChatID: chat.ChatID, + RequestID: requestID, + AssistantID: "sub_assistant", + StackID: "sub_stack", + StackParentID: "main_stack", + StackDepth: 1, + Type: types.ResumeTypeLLM, + Status: types.ResumeStatusInterrupted, + Input: map[string]interface{}{"messages": []interface{}{}}, + Output: map[string]interface{}{"partial_content": "The analysis shows..."}, + SpaceSnapshot: map[string]interface{}{"task": "analyze", "data_id": "123"}, + Sequence: 4, + }, + } + + err = store.SaveResume(records) + if err != nil { + t.Fatalf("Failed to save resume records: %v", err) + } + t.Logf("Saved %d resume records for A2A workflow", len(records)) + + // 1. Get last resume record (should be the interrupted LLM call) + last, err := store.GetLastResume(chat.ChatID) + if err != nil { + t.Fatalf("Failed to get last resume: %v", err) + } + + if last == nil { + t.Fatal("Expected last resume record") + } + + if last.Type != types.ResumeTypeLLM { + t.Errorf("Expected type '%s', got '%s'", types.ResumeTypeLLM, last.Type) + } + if last.StackDepth != 1 { + t.Errorf("Expected stack_depth 1, got %d", last.StackDepth) + } + + // 2. Get stack path to understand the call hierarchy + path, err := store.GetStackPath(last.StackID) + if err != nil { + t.Fatalf("Failed to get stack path: %v", err) + } + + if len(path) != 2 { + t.Errorf("Expected path length 2, got %d", len(path)) + } + t.Logf("Stack path: %v", path) + + // 3. Get all records for the sub stack + subRecords, err := store.GetResumeByStackID("sub_stack") + if err != nil { + t.Fatalf("Failed to get sub stack records: %v", err) + } + + if len(subRecords) != 2 { + t.Errorf("Expected 2 records for sub_stack, got %d", len(subRecords)) + } + + // 4. Verify space snapshot is preserved + if last.SpaceSnapshot == nil { + t.Error("Expected space_snapshot to be set") + } else { + if last.SpaceSnapshot["task"] != "analyze" { + t.Errorf("Expected task='analyze', got '%v'", last.SpaceSnapshot["task"]) + } + } + + // 5. Clean up after successful resume + err = store.DeleteResume(chat.ChatID) + if err != nil { + t.Fatalf("Failed to delete resume records: %v", err) + } + + // 6. Verify cleanup + remaining, err := store.GetResume(chat.ChatID) + if err != nil { + t.Fatalf("Failed to get remaining records: %v", err) + } + + if len(remaining) != 0 { + t.Errorf("Expected 0 records after cleanup, got %d", len(remaining)) + } + + t.Log("Complete A2A workflow test passed!") + }) +} diff --git a/agent/store/xun/xun_test.go b/agent/store/xun/xun_test.go deleted file mode 100644 index c0711016..00000000 --- a/agent/store/xun/xun_test.go +++ /dev/null @@ -1,1870 +0,0 @@ -package xun - -// import ( -// "fmt" -// "testing" -// "time" - -// "github.com/stretchr/testify/assert" -// "github.com/yaoapp/gou/connector" -// "github.com/yaoapp/xun/capsule" -// "github.com/yaoapp/yao/config" -// "github.com/yaoapp/yao/test" -// ) - -// func TestNewXunDefault(t *testing.T) { -// test.Prepare(t, config.Conf) -// defer test.Clean() -// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_history") -// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat") -// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant") - -// err := capsule.Schema().DropTableIfExists("__unit_test_conversation_history") -// if err != nil { -// t.Fatal(err) -// } - -// err = capsule.Schema().DropTableIfExists("__unit_test_conversation_chat") -// if err != nil { -// t.Fatal(err) -// } - -// err = capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant") -// if err != nil { -// t.Fatal(err) -// } - -// // Add a small delay to ensure table is created -// time.Sleep(100 * time.Millisecond) - -// store, err := NewXun(Setting{ -// Connector: "default", -// Prefix: "__unit_test_conversation_", -// }) - -// if err != nil { -// t.Error(err) -// return -// } - -// // Check history table -// has, err := capsule.Schema().HasTable("__unit_test_conversation_history") -// if err != nil { -// t.Fatal(err) -// } -// assert.Equal(t, true, has) - -// // Check chat table -// has, err = capsule.Schema().HasTable("__unit_test_conversation_chat") -// if err != nil { -// t.Fatal(err) -// } -// assert.Equal(t, true, has) - -// // Check assistant table -// has, err = capsule.Schema().HasTable("__unit_test_conversation_assistant") -// if err != nil { -// t.Fatal(err) -// } -// assert.Equal(t, true, has) - -// // Validate table structure by attempting operations -// // Test history operations -// messages := []map[string]interface{}{ -// {"role": "user", "content": "test message"}, -// } -// err = store.SaveHistory("test_user", messages, "test_chat", nil) -// assert.Nil(t, err) - -// history, err := store.GetHistory("test_user", "test_chat") -// assert.Nil(t, err) -// assert.NotEmpty(t, history) - -// // Test chat operations -// err = store.UpdateChatTitle("test_user", "test_chat", "Test Chat") -// assert.Nil(t, err) - -// chat, err := store.GetChat("test_user", "test_chat") -// assert.Nil(t, err) -// assert.NotNil(t, chat) - -// // Test assistant operations -// assistant := map[string]interface{}{ -// "name": "Test Assistant", -// "type": "assistant", -// "connector": "test", -// "description": "Test Description", -// "tags": []string{"test"}, -// "mentionable": true, -// "automated": true, -// } - -// id, err := store.SaveAssistant(assistant) -// assert.Nil(t, err) -// assert.NotNil(t, id) - -// // Clean up test data -// err = store.DeleteChat("test_user", "test_chat") -// assert.Nil(t, err) - -// err = store.DeleteAssistant(id.(string)) -// assert.Nil(t, err) -// } - -// func TestNewXunConnector(t *testing.T) { -// test.Prepare(t, config.Conf) -// defer test.Clean() - -// conn, err := connector.Select("mysql") -// if err != nil { -// t.Fatal(err) -// } - -// sch, err := conn.Schema() -// if err != nil { -// t.Fatal(err) -// } - -// defer sch.DropTableIfExists("__unit_test_conversation_history") -// defer sch.DropTableIfExists("__unit_test_conversation_chat") -// defer sch.DropTableIfExists("__unit_test_conversation_assistant") -// defer sch.DropTableIfExists("__unit_test_conversation_knowledge") -// defer sch.DropTableIfExists("__unit_test_conversation_attachment") - -// sch.DropTableIfExists("__unit_test_conversation_history") -// sch.DropTableIfExists("__unit_test_conversation_chat") -// sch.DropTableIfExists("__unit_test_conversation_assistant") -// sch.DropTableIfExists("__unit_test_conversation_knowledge") -// sch.DropTableIfExists("__unit_test_conversation_attachment") - -// // Add a small delay to ensure table is created -// time.Sleep(100 * time.Millisecond) - -// store, err := NewXun(Setting{ -// Connector: "mysql", -// Prefix: "__unit_test_conversation_", -// }) - -// if err != nil { -// t.Error(err) -// return -// } - -// // Check history table -// has, err := sch.HasTable("__unit_test_conversation_history") -// if err != nil { -// t.Fatal(err) -// } -// assert.Equal(t, true, has) - -// // Check chat table -// has, err = sch.HasTable("__unit_test_conversation_chat") -// if err != nil { -// t.Fatal(err) -// } -// assert.Equal(t, true, has) - -// // Check assistant table -// has, err = sch.HasTable("__unit_test_conversation_assistant") -// if err != nil { -// t.Fatal(err) -// } -// assert.Equal(t, true, has) - -// // Test basic operations -// messages := []map[string]interface{}{ -// {"role": "user", "content": "test message"}, -// } -// err = store.SaveHistory("test_user", messages, "test_chat", nil) -// assert.Nil(t, err) - -// history, err := store.GetHistory("test_user", "test_chat") -// assert.Nil(t, err) -// assert.NotEmpty(t, history) - -// err = store.DeleteChat("test_user", "test_chat") -// assert.Nil(t, err) -// } - -// func TestXunSaveAndGetHistory(t *testing.T) { -// test.Prepare(t, config.Conf) -// defer test.Clean() -// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_history") -// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat") - -// err := capsule.Schema().DropTableIfExists("__unit_test_conversation_history") -// if err != nil { -// t.Fatal(err) -// } - -// err = capsule.Schema().DropTableIfExists("__unit_test_conversation_chat") -// if err != nil { -// t.Fatal(err) -// } - -// store, err := NewXun(Setting{ -// Connector: "default", -// Prefix: "__unit_test_conversation_", -// TTL: 3600, -// }) - -// // save the history -// cid := "123456" -// err = store.SaveHistory("123456", []map[string]interface{}{ -// {"role": "user", "name": "user1", "content": "hello"}, -// {"role": "assistant", "name": "user1", "content": "Hello there, how"}, -// }, cid, nil) -// assert.Nil(t, err) - -// // get the history -// data, err := store.GetHistory("123456", cid) -// if err != nil { -// t.Fatal(err) -// } -// assert.Equal(t, 2, len(data)) -// } - -// func TestXunSaveAndGetHistoryWithCID(t *testing.T) { -// test.Prepare(t, config.Conf) -// defer test.Clean() -// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_history") -// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat") - -// err := capsule.Schema().DropTableIfExists("__unit_test_conversation_history") -// if err != nil { -// t.Fatal(err) -// } - -// err = capsule.Schema().DropTableIfExists("__unit_test_conversation_chat") -// if err != nil { -// t.Fatal(err) -// } - -// store, err := NewXun(Setting{ -// Connector: "default", -// Prefix: "__unit_test_conversation_", -// TTL: 3600, -// }) - -// // save the history with specific cid -// sid := "123456" -// cid := "789012" -// assistantID := "test-assistant-1" -// messages := []map[string]interface{}{ -// {"role": "user", "name": "user1", "content": "hello"}, -// {"role": "assistant", "name": "assistant1", "content": "Hi! How can I help you?"}, -// } -// context := map[string]interface{}{ -// "assistant_id": assistantID, -// } -// err = store.SaveHistory(sid, messages, cid, context) -// assert.Nil(t, err) - -// // get the history for specific cid -// data, err := store.GetHistory(sid, cid) -// if err != nil { -// t.Fatal(err) -// } -// assert.Equal(t, 2, len(data)) - -// // Verify assistant_id is saved in chat -// chat, err := store.GetChat(sid, cid) -// assert.Nil(t, err) -// assert.Equal(t, assistantID, chat.Chat["assistant_id"]) - -// // save another message with different cid and assistant -// anotherCID := "345678" -// anotherAssistantID := "test-assistant-2" -// moreMessages := []map[string]interface{}{ -// {"role": "user", "name": "user1", "content": "another message"}, -// {"role": "assistant", "name": "assistant2", "content": "Hello!"}, -// } -// anotherContext := map[string]interface{}{ -// "assistant_id": anotherAssistantID, -// } -// err = store.SaveHistory(sid, moreMessages, anotherCID, anotherContext) -// assert.Nil(t, err) - -// // Verify second chat's assistant_id -// chat2, err := store.GetChat(sid, anotherCID) -// assert.Nil(t, err) -// assert.Equal(t, anotherAssistantID, chat2.Chat["assistant_id"]) - -// // get history for the first cid - should still be 2 messages -// data, err = store.GetHistory(sid, cid) -// if err != nil { -// t.Fatal(err) -// } -// assert.Equal(t, 2, len(data)) - -// // get history for the second cid - should be 2 messages -// data, err = store.GetHistory(sid, anotherCID) -// if err != nil { -// t.Fatal(err) -// } -// assert.Equal(t, 2, len(data)) - -// // get all history for the sid without specifying cid -// allData, err := store.GetHistory(sid, cid) -// if err != nil { -// t.Fatal(err) -// } -// assert.Equal(t, 2, len(allData)) -// } - -// func TestXunGetChats(t *testing.T) { -// test.Prepare(t, config.Conf) -// defer test.Clean() -// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_history") -// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat") -// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant") - -// // Drop tables before test -// err := capsule.Schema().DropTableIfExists("__unit_test_conversation_history") -// if err != nil { -// t.Fatal(err) -// } -// err = capsule.Schema().DropTableIfExists("__unit_test_conversation_chat") -// if err != nil { -// t.Fatal(err) -// } -// err = capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant") -// if err != nil { -// t.Fatal(err) -// } - -// store, err := NewXun(Setting{ -// Connector: "default", -// Prefix: "__unit_test_conversation_", -// }) -// if err != nil { -// t.Fatal(err) -// } - -// // Create test assistants first -// assistant1 := map[string]interface{}{ -// "assistant_id": "test-assistant-1", -// "name": "Test Assistant 1", -// "avatar": "avatar1.png", -// "type": "assistant", -// "connector": "test", -// } -// assistant2 := map[string]interface{}{ -// "assistant_id": "test-assistant-2", -// "name": "Test Assistant 2", -// "avatar": "avatar2.png", -// "type": "assistant", -// "connector": "test", -// } -// _, err = store.SaveAssistant(assistant1) -// assert.Nil(t, err) -// _, err = store.SaveAssistant(assistant2) -// assert.Nil(t, err) - -// // Save some test chats -// sid := "test_user" -// messages := []map[string]interface{}{ -// {"role": "user", "content": "test message"}, -// } - -// // Create chats with different dates and assistants -// for i := 0; i < 5; i++ { -// chatID := fmt.Sprintf("chat_%d", i) -// title := fmt.Sprintf("Test Chat %d", i) -// var context map[string]interface{} - -// // Alternate between having assistant and no assistant -// if i%2 == 0 { -// context = map[string]interface{}{ -// "assistant_id": "test-assistant-1", -// } -// } else if i%3 == 0 { -// context = map[string]interface{}{ -// "assistant_id": "test-assistant-2", -// } -// } - -// // Save history first to create the chat -// err = store.SaveHistory(sid, messages, chatID, context) -// assert.Nil(t, err) - -// // Update the chat title -// err = store.UpdateChatTitle(sid, chatID, title) -// assert.Nil(t, err) - -// // Verify chat was created with correct assistant info -// chat, err := store.GetChat(sid, chatID) -// assert.Nil(t, err) -// assert.NotNil(t, chat) -// assert.Equal(t, chatID, chat.Chat["chat_id"]) -// assert.Equal(t, title, chat.Chat["title"]) - -// if i%2 == 0 { -// assert.Equal(t, "test-assistant-1", chat.Chat["assistant_id"]) -// assert.Equal(t, "Test Assistant 1", chat.Chat["assistant_name"]) -// assert.Equal(t, "avatar1.png", chat.Chat["assistant_avatar"]) -// } else if i%3 == 0 { -// assert.Equal(t, "test-assistant-2", chat.Chat["assistant_id"]) -// assert.Equal(t, "Test Assistant 2", chat.Chat["assistant_name"]) -// assert.Equal(t, "avatar2.png", chat.Chat["assistant_avatar"]) -// } else { -// assert.Nil(t, chat.Chat["assistant_id"]) -// assert.Nil(t, chat.Chat["assistant_name"]) -// assert.Nil(t, chat.Chat["assistant_avatar"]) -// } -// } - -// // Test GetChats -// filter := ChatFilter{ -// PageSize: 10, -// Order: "desc", -// } -// groups, err := store.GetChats(sid, filter) -// assert.Nil(t, err) -// assert.NotNil(t, groups) -// assert.Greater(t, len(groups.Groups), 0) - -// // Verify assistant information in chat list -// for _, group := range groups.Groups { -// for _, chat := range group.Chats { -// if assistantID, ok := chat["assistant_id"].(string); ok && assistantID != "" { -// if assistantID == "test-assistant-1" { -// assert.Equal(t, "Test Assistant 1", chat["assistant_name"]) -// assert.Equal(t, "avatar1.png", chat["assistant_avatar"]) -// } else if assistantID == "test-assistant-2" { -// assert.Equal(t, "Test Assistant 2", chat["assistant_name"]) -// assert.Equal(t, "avatar2.png", chat["assistant_avatar"]) -// } -// } else { -// assert.Nil(t, chat["assistant_name"]) -// assert.Nil(t, chat["assistant_avatar"]) -// } -// } -// } - -// // Test with keywords -// filter.Keywords = "test" -// groups, err = store.GetChats(sid, filter) -// assert.Nil(t, err) -// assert.Greater(t, len(groups.Groups), 0) -// } - -// func TestXunDeleteChat(t *testing.T) { -// test.Prepare(t, config.Conf) -// defer test.Clean() -// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_history") -// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat") - -// store, err := NewXun(Setting{ -// Connector: "default", -// Prefix: "__unit_test_conversation_", -// }) -// if err != nil { -// t.Fatal(err) -// } - -// // Create a test chat -// sid := "test_user" -// cid := "test_chat" -// messages := []map[string]interface{}{ -// {"role": "user", "content": "test message"}, -// } - -// // Save the chat and history -// err = store.SaveHistory(sid, messages, cid, nil) -// assert.Nil(t, err) - -// // Verify chat exists -// chat, err := store.GetChat(sid, cid) -// assert.Nil(t, err) -// assert.NotNil(t, chat) - -// // Delete the chat -// err = store.DeleteChat(sid, cid) -// assert.Nil(t, err) - -// // Verify chat is deleted -// chat, err = store.GetChat(sid, cid) -// assert.Nil(t, err) -// assert.Equal(t, (*ChatInfo)(nil), chat) -// } - -// func TestXunDeleteAllChats(t *testing.T) { -// test.Prepare(t, config.Conf) -// defer test.Clean() -// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_history") -// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat") - -// store, err := NewXun(Setting{ -// Connector: "default", -// Prefix: "__unit_test_conversation_", -// }) -// if err != nil { -// t.Fatal(err) -// } - -// // Create multiple test chats -// sid := "test_user" -// messages := []map[string]interface{}{ -// {"role": "user", "content": "test message"}, -// } - -// // Save multiple chats -// for i := 0; i < 3; i++ { -// cid := fmt.Sprintf("test_chat_%d", i) -// err = store.SaveHistory(sid, messages, cid, nil) -// assert.Nil(t, err) -// } - -// // Verify chats exist -// response, err := store.GetChats(sid, ChatFilter{}) -// assert.Nil(t, err) -// assert.Greater(t, response.Total, int64(0)) - -// // Delete all chats -// err = store.DeleteAllChats(sid) -// assert.Nil(t, err) - -// // Verify all chats are deleted -// response, err = store.GetChats(sid, ChatFilter{}) -// assert.Nil(t, err) -// assert.Equal(t, int64(0), response.Total) -// } - -// func TestXunAssistantCRUD(t *testing.T) { -// test.Prepare(t, config.Conf) -// defer test.Clean() -// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant") - -// // Drop assistant table before test -// err := capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant") -// if err != nil { -// t.Fatal(err) -// } - -// // Add a small delay to ensure table is created -// time.Sleep(100 * time.Millisecond) - -// store, err := NewXun(Setting{ -// Connector: "default", -// Prefix: "__unit_test_conversation_", -// }) -// if err != nil { -// t.Fatal(err) -// } - -// // Clean up any existing data -// _, err = store.DeleteAssistants(AssistantFilter{}) -// assert.Nil(t, err) - -// // Test case 1: JSON fields as strings -// tagsJSON := `["tag1", "tag2", "tag3"]` -// optionsJSON := `{"model": "gpt-4"}` -// placeholderJSON := `{"title": "Test Title", "description": "Test Description", "prompts": ["prompt1", "prompt2"]}` -// assistant := map[string]interface{}{ -// "name": "Test Assistant", -// "type": "assistant", -// "avatar": "https://example.com/avatar.png", -// "connector": "openai", -// "description": "Test Description", -// "path": "/assistants/test", -// "sort": 100, -// "built_in": true, -// "tags": tagsJSON, -// "options": optionsJSON, -// "placeholder": placeholderJSON, -// "mentionable": true, -// "automated": true, -// } - -// // Test SaveAssistant (Create) with string JSON -// v, err := store.SaveAssistant(assistant) -// assert.Nil(t, err) -// assistantID := v.(string) -// assert.NotEmpty(t, assistantID) - -// // Test GetAssistant for the first assistant -// assistantData, err := store.GetAssistant(assistantID) -// assert.Nil(t, err) -// assert.NotNil(t, assistantData) -// assert.Equal(t, "Test Assistant", assistantData["name"]) -// assert.Equal(t, "assistant", assistantData["type"]) -// assert.Equal(t, "https://example.com/avatar.png", assistantData["avatar"]) -// assert.Equal(t, "openai", assistantData["connector"]) -// assert.Equal(t, "Test Description", assistantData["description"]) -// assert.Equal(t, "/assistants/test", assistantData["path"]) -// assert.Equal(t, int64(100), assistantData["sort"]) -// assert.Equal(t, int64(1), assistantData["built_in"]) -// assert.Equal(t, []interface{}{"tag1", "tag2", "tag3"}, assistantData["tags"]) -// assert.Equal(t, map[string]interface{}{"model": "gpt-4"}, assistantData["options"]) -// assert.Equal(t, map[string]interface{}{ -// "title": "Test Title", -// "description": "Test Description", -// "prompts": []interface{}{"prompt1", "prompt2"}, -// }, assistantData["placeholder"]) -// assert.Equal(t, int64(1), assistantData["mentionable"]) -// assert.Equal(t, int64(1), assistantData["automated"]) - -// // Test case 2: JSON fields as native types -// assistant2 := map[string]interface{}{ -// "name": "Test Assistant 2", -// "type": "assistant", -// "avatar": "https://example.com/avatar2.png", -// "connector": "openai", -// "description": "Test Description 2", -// "path": "/assistants/test2", -// "sort": 200, -// "built_in": false, -// "tags": []string{"tag1", "tag2", "tag3"}, -// "options": map[string]interface{}{"model": "gpt-4"}, -// "prompts": []string{"prompt1", "prompt2"}, -// "workflow": []string{"flow1", "flow2"}, -// "knowledge": []string{"file1", "file2"}, -// "tools": []map[string]interface{}{{"name": "tool1"}, {"name": "tool2"}}, -// "permissions": map[string]interface{}{"read": true, "write": true}, -// "placeholder": map[string]interface{}{ -// "title": "Test Title 2", -// "description": "Test Description 2", -// "prompts": []string{"prompt3", "prompt4"}, -// }, -// "mentionable": true, -// "automated": true, -// } - -// // Test SaveAssistant (Create) with native types -// v, err = store.SaveAssistant(assistant2) -// assert.Nil(t, err) -// assistant2ID := v.(string) -// assert.NotEmpty(t, assistant2ID) - -// // Test case 3: Test with nil JSON fields -// assistant3 := map[string]interface{}{ -// "name": "Test Assistant 3", -// "type": "assistant", -// "connector": "openai", -// "description": "Test Description 3", -// "path": nil, -// "sort": 9999, -// "built_in": false, -// "tags": nil, -// "options": nil, -// "prompts": nil, -// "workflow": nil, -// "knowledge": nil, -// "tools": nil, -// "permissions": nil, -// "placeholder": nil, -// "mentionable": true, -// "automated": true, -// } - -// // Test SaveAssistant (Create) with nil fields -// v, err = store.SaveAssistant(assistant3) -// assert.Nil(t, err) -// assistant3ID := v.(string) -// assert.NotEmpty(t, assistant3ID) - -// // Test GetAssistant for the third assistant -// assistant3Data, err := store.GetAssistant(assistant3ID) -// assert.Nil(t, err) -// assert.NotNil(t, assistant3Data) -// assert.Equal(t, "Test Assistant 3", assistant3Data["name"]) -// assert.Nil(t, assistant3Data["tags"]) -// assert.Nil(t, assistant3Data["options"]) -// assert.Nil(t, assistant3Data["prompts"]) -// assert.Nil(t, assistant3Data["workflow"]) -// assert.Nil(t, assistant3Data["knowledge"]) -// assert.Nil(t, assistant3Data["tools"]) -// assert.Nil(t, assistant3Data["permissions"]) -// assert.Nil(t, assistant3Data["placeholder"]) -// assert.Equal(t, int64(1), assistant3Data["mentionable"]) -// assert.Equal(t, int64(1), assistant3Data["automated"]) - -// // Test GetAssistant with non-existent ID -// nonExistentData, err := store.GetAssistant("non-existent-id") -// assert.Error(t, err) -// assert.Nil(t, nonExistentData) -// assert.Contains(t, err.Error(), "is empty") - -// // Test GetAssistants to verify JSON fields are properly stored -// resp, err := store.GetAssistants(AssistantFilter{}) -// assert.Nil(t, err) -// assert.Equal(t, 3, len(resp.Data)) - -// // Clean up all test data -// _, err = store.DeleteAssistants(AssistantFilter{}) -// assert.Nil(t, err) - -// // Verify cleanup -// resp, err = store.GetAssistants(AssistantFilter{}) -// assert.Nil(t, err) -// assert.Equal(t, 0, len(resp.Data)) -// } - -// func TestXunAssistantPagination(t *testing.T) { -// test.Prepare(t, config.Conf) -// defer test.Clean() -// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_history") -// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat") -// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant") - -// // Drop assistant table before test -// err := capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant") -// if err != nil { -// t.Fatal(err) -// } - -// // Add a small delay to ensure table is created -// time.Sleep(100 * time.Millisecond) - -// store, err := NewXun(Setting{ -// Connector: "default", -// Prefix: "__unit_test_conversation_", -// }) -// if err != nil { -// t.Fatal(err) -// } - -// // Create test data for filtering tests -// testAssistants := []map[string]interface{}{} -// for i := 0; i < 25; i++ { -// assistant := map[string]interface{}{ -// "name": fmt.Sprintf("Filter Test Assistant %d", i), -// "type": "assistant", -// "connector": fmt.Sprintf("connector%d", i%3), -// "description": fmt.Sprintf("Filter Test Description %d", i), -// "tags": []string{fmt.Sprintf("tag%d", i%5)}, -// "built_in": i%2 == 0, -// "mentionable": i%2 == 0, -// "automated": i%3 == 0, -// "sort": 9999 - i, -// } -// id, err := store.SaveAssistant(assistant) -// assert.Nil(t, err) -// assistant["assistant_id"] = id -// testAssistants = append(testAssistants, assistant) -// } - -// // Get first assistant ID for later use -// firstAssistantID := testAssistants[0]["assistant_id"].(string) - -// // Test filtering with assistantIDs -// assistantIDs := []string{firstAssistantID} -// if len(testAssistants) > 1 { -// assistantIDs = append(assistantIDs, testAssistants[1]["assistant_id"].(string)) -// } - -// // Test multiple assistant_ids -// resp, err := store.GetAssistants(AssistantFilter{ -// AssistantIDs: assistantIDs, -// Page: 1, -// PageSize: 10, -// }) -// assert.Nil(t, err) -// assert.Equal(t, len(assistantIDs), len(resp.Data)) -// for _, assistant := range resp.Data { -// found := false -// for _, id := range assistantIDs { -// if assistant["assistant_id"] == id { -// found = true -// break -// } -// } -// assert.True(t, found, "Assistant ID should be in the requested list") -// } - -// // Test assistantIDs with other filters -// resp, err = store.GetAssistants(AssistantFilter{ -// AssistantIDs: assistantIDs, -// Select: []string{"name", "assistant_id", "description"}, -// Page: 1, -// PageSize: 10, -// }) -// assert.Nil(t, err) -// assert.Equal(t, len(assistantIDs), len(resp.Data)) -// // Verify only selected fields are returned -// for _, item := range resp.Data { -// assert.Contains(t, item, "name") -// assert.Contains(t, item, "assistant_id") -// assert.Contains(t, item, "description") -// assert.NotContains(t, item, "tags") -// assert.NotContains(t, item, "options") -// } - -// // Test filtering with select fields -// resp, err = store.GetAssistants(AssistantFilter{ -// Select: []string{"name", "description", "tags"}, -// Page: 1, -// PageSize: 10, -// }) -// assert.Nil(t, err) -// assert.Equal(t, 10, len(resp.Data)) - -// // Test filtering with select fields and other filters combined -// resp, err = store.GetAssistants(AssistantFilter{ -// Tags: []string{"tag0"}, -// Keywords: "Filter Test", -// Select: []string{"name", "tags"}, -// Page: 1, -// PageSize: 10, -// }) -// assert.Nil(t, err) -// assert.Greater(t, len(resp.Data), 0) - -// // Test combined filters -// mentionableTrue := true -// automatedTrue := true -// resp, err = store.GetAssistants(AssistantFilter{ -// Tags: []string{"tag0"}, -// Keywords: "Filter Test", -// Connector: "connector0", -// Mentionable: &mentionableTrue, -// Automated: &automatedTrue, -// Page: 1, -// PageSize: 10, -// }) -// assert.Nil(t, err) - -// // Now test the delete operations -// // Test delete by connector -// var count int64 -// count, err = store.DeleteAssistants(AssistantFilter{ -// Connector: "connector0", -// }) -// assert.Nil(t, err) -// assert.Greater(t, count, int64(0)) - -// // Verify deletion -// resp, err = store.GetAssistants(AssistantFilter{ -// Connector: "connector0", -// }) -// assert.Nil(t, err) -// assert.Equal(t, 0, len(resp.Data)) - -// // Test delete by built_in status -// builtInTrue := true -// count, err = store.DeleteAssistants(AssistantFilter{ -// BuiltIn: &builtInTrue, -// }) -// assert.Nil(t, err) -// assert.Greater(t, count, int64(0)) - -// // Verify deletion -// resp, err = store.GetAssistants(AssistantFilter{ -// BuiltIn: &builtInTrue, -// }) -// assert.Nil(t, err) -// assert.Equal(t, 0, len(resp.Data)) - -// // Test delete by tags -// count, err = store.DeleteAssistants(AssistantFilter{ -// Tags: []string{"tag1"}, -// }) -// assert.Nil(t, err) -// assert.Greater(t, count, int64(0)) - -// // Verify deletion -// resp, err = store.GetAssistants(AssistantFilter{ -// Tags: []string{"tag1"}, -// }) -// assert.Nil(t, err) -// assert.Equal(t, 0, len(resp.Data)) - -// // Test delete by keywords -// count, err = store.DeleteAssistants(AssistantFilter{ -// Keywords: "Filter Test", -// }) -// assert.Nil(t, err) -// assert.Greater(t, count, int64(0)) - -// // Verify all assistants are deleted -// resp, err = store.GetAssistants(AssistantFilter{}) -// assert.Nil(t, err) -// assert.Equal(t, 0, len(resp.Data)) - -// // Test delete by assistantIDs -// // First create some test assistants -// testIDs := []string{} -// for i := 0; i < 3; i++ { -// assistant := map[string]interface{}{ -// "name": fmt.Sprintf("AssistantIDs Test Assistant %d", i), -// "type": "assistant", -// "connector": "test", -// "description": fmt.Sprintf("AssistantIDs Test Description %d", i), -// "tags": []string{"test-tag"}, -// "built_in": false, -// "mentionable": true, -// "automated": true, -// } -// id, err := store.SaveAssistant(assistant) -// assert.Nil(t, err) -// testIDs = append(testIDs, id.(string)) -// } - -// // Delete by assistantIDs -// count, err = store.DeleteAssistants(AssistantFilter{ -// AssistantIDs: testIDs, -// }) -// assert.Nil(t, err) -// assert.Equal(t, int64(len(testIDs)), count) - -// // Verify deletion -// resp, err = store.GetAssistants(AssistantFilter{ -// AssistantIDs: testIDs, -// }) -// assert.Nil(t, err) -// assert.Equal(t, 0, len(resp.Data)) - -// // Verify all assistants are deleted -// resp, err = store.GetAssistants(AssistantFilter{}) -// assert.Nil(t, err) -// assert.Equal(t, 0, len(resp.Data)) -// } - -// func TestGetAssistantTags(t *testing.T) { - -// test.Prepare(t, config.Conf) -// defer test.Clean() -// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant") - -// store, err := NewXun(Setting{ -// Connector: "default", -// Prefix: "__unit_test_conversation_", -// }) -// if err != nil { -// t.Fatal(err) -// } - -// // Create test assistants with tags -// assistants := []map[string]interface{}{ -// { -// "assistant_id": "test-assistant-1", -// "type": "assistant", -// "connector": "test", -// "tags": []string{"tag1", "tag2"}, -// "name": "Test Assistant 1", -// }, -// { -// "assistant_id": "test-assistant-2", -// "type": "assistant", -// "connector": "test", -// "tags": []string{"tag2", "tag3"}, -// "name": "Test Assistant 2", -// }, -// { -// "assistant_id": "test-assistant-3", -// "type": "assistant", -// "connector": "test", -// "tags": []string{"tag1", "tag3", "tag4"}, -// "name": "Test Assistant 3", -// }, -// } - -// // Save test assistants -// for _, assistant := range assistants { -// _, err := store.SaveAssistant(assistant) -// if err != nil { -// t.Fatal(err) -// } -// } - -// // Get tags -// tags, err := store.GetAssistantTags() -// if err != nil { -// t.Fatal(err) -// } - -// // Verify results -// expectedTags := map[string]bool{ -// "tag1": true, -// "tag2": true, -// "tag3": true, -// "tag4": true, -// } - -// if len(tags) != len(expectedTags) { -// t.Errorf("Expected %d tags, got %d", len(expectedTags), len(tags)) -// } - -// for _, tag := range tags { -// value := tag.Value -// if !expectedTags[value] { -// t.Errorf("Unexpected tag found: %s", tag) -// } -// } - -// // Cleanup -// for _, assistant := range assistants { -// err := store.DeleteAssistant(assistant["assistant_id"].(string)) -// if err != nil { -// t.Fatal(err) -// } -// } -// } - -// func TestXunSaveAndGetHistoryWithSilent(t *testing.T) { -// test.Prepare(t, config.Conf) -// defer test.Clean() -// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_history") -// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat") - -// err := capsule.Schema().DropTableIfExists("__unit_test_conversation_history") -// if err != nil { -// t.Fatal(err) -// } - -// err = capsule.Schema().DropTableIfExists("__unit_test_conversation_chat") -// if err != nil { -// t.Fatal(err) -// } - -// store, err := NewXun(Setting{ -// Connector: "default", -// Prefix: "__unit_test_conversation_", -// TTL: 3600, -// }) - -// // save the history with silent messages -// sid := "123456" -// cid := "silent_test" - -// // First save regular messages -// messages := []map[string]interface{}{ -// {"role": "user", "name": "user1", "content": "hello"}, -// {"role": "assistant", "name": "assistant1", "content": "Hi! How can I help you?"}, -// } -// context := map[string]interface{}{ -// "assistant_id": "test-assistant-1", -// } -// err = store.SaveHistory(sid, messages, cid, context) -// assert.Nil(t, err) - -// // Then save silent messages -// silentMessages := []map[string]interface{}{ -// {"role": "user", "name": "user1", "content": "silent message"}, -// {"role": "assistant", "name": "assistant1", "content": "This is a silent response"}, -// } -// silentContext := map[string]interface{}{ -// "assistant_id": "test-assistant-1", -// "silent": true, -// } -// err = store.SaveHistory(sid, silentMessages, cid, silentContext) -// assert.Nil(t, err) - -// // Get history without filter (should only return non-silent messages) -// data, err := store.GetHistory(sid, cid) -// if err != nil { -// t.Fatal(err) -// } -// assert.Equal(t, 2, len(data)) -// for _, msg := range data { -// // Check if silent is false, handling different types -// isSilent := false -// switch v := msg["silent"].(type) { -// case bool: -// isSilent = v -// case int: -// isSilent = v != 0 -// case int64: -// isSilent = v != 0 -// case float64: -// isSilent = v != 0 -// } -// assert.False(t, isSilent, "message should not be silent") -// } - -// // Get history with silent=true filter (should return all messages) -// silentTrue := true -// filter := ChatFilter{ -// Silent: &silentTrue, -// } -// allData, err := store.GetHistoryWithFilter(sid, cid, filter) -// if err != nil { -// t.Fatal(err) -// } -// assert.Equal(t, 4, len(allData)) - -// // Count silent messages -// silentCount := 0 -// for _, msg := range allData { -// // Check if silent is true, handling different types -// isSilent := false -// switch v := msg["silent"].(type) { -// case bool: -// isSilent = v -// case int: -// isSilent = v != 0 -// case int64: -// isSilent = v != 0 -// case float64: -// isSilent = v != 0 -// } -// if isSilent { -// silentCount++ -// } -// } -// assert.Equal(t, 2, silentCount) - -// // Get chat with filter (should include silent messages) -// chat, err := store.GetChatWithFilter(sid, cid, filter) -// assert.Nil(t, err) -// assert.Equal(t, 4, len(chat.History)) - -// // Get chat without filter (should exclude silent messages) -// chatNoSilent, err := store.GetChat(sid, cid) -// assert.Nil(t, err) -// assert.Equal(t, 2, len(chatNoSilent.History)) -// } - -// func TestXunGetChatsWithSilent(t *testing.T) { -// test.Prepare(t, config.Conf) -// defer test.Clean() -// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_history") -// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat") -// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant") - -// // Drop tables before test -// err := capsule.Schema().DropTableIfExists("__unit_test_conversation_history") -// if err != nil { -// t.Fatal(err) -// } -// err = capsule.Schema().DropTableIfExists("__unit_test_conversation_chat") -// if err != nil { -// t.Fatal(err) -// } -// err = capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant") -// if err != nil { -// t.Fatal(err) -// } - -// store, err := NewXun(Setting{ -// Connector: "default", -// Prefix: "__unit_test_conversation_", -// }) -// if err != nil { -// t.Fatal(err) -// } - -// // Create test assistant -// assistant := map[string]interface{}{ -// "assistant_id": "test-assistant-1", -// "name": "Test Assistant 1", -// "avatar": "avatar1.png", -// "type": "assistant", -// "connector": "test", -// } -// _, err = store.SaveAssistant(assistant) -// assert.Nil(t, err) - -// // Save some test chats -// sid := "test_user" -// messages := []map[string]interface{}{ -// {"role": "user", "content": "test message"}, -// } - -// // Create regular chats -// for i := 0; i < 3; i++ { -// chatID := fmt.Sprintf("regular_chat_%d", i) -// title := fmt.Sprintf("Regular Chat %d", i) -// context := map[string]interface{}{ -// "assistant_id": "test-assistant-1", -// "silent": false, -// } - -// // Save history to create the chat -// err = store.SaveHistory(sid, messages, chatID, context) -// assert.Nil(t, err) - -// // Update the chat title -// err = store.UpdateChatTitle(sid, chatID, title) -// assert.Nil(t, err) -// } - -// // Create silent chats -// for i := 0; i < 2; i++ { -// chatID := fmt.Sprintf("silent_chat_%d", i) -// title := fmt.Sprintf("Silent Chat %d", i) -// context := map[string]interface{}{ -// "assistant_id": "test-assistant-1", -// "silent": true, -// } - -// // Save history to create the chat -// err = store.SaveHistory(sid, messages, chatID, context) -// assert.Nil(t, err) - -// // Update the chat title -// err = store.UpdateChatTitle(sid, chatID, title) -// assert.Nil(t, err) -// } - -// // Test GetChats with default filter (should exclude silent chats) -// defaultFilter := ChatFilter{ -// PageSize: 10, -// Order: "desc", -// } -// defaultGroups, err := store.GetChats(sid, defaultFilter) -// assert.Nil(t, err) -// assert.NotNil(t, defaultGroups) - -// // Count total chats in all groups -// totalDefaultChats := 0 -// for _, group := range defaultGroups.Groups { -// totalDefaultChats += len(group.Chats) -// } -// assert.Equal(t, 3, totalDefaultChats, "Default filter should only return non-silent chats") - -// // Test GetChats with silent=true filter (should include all chats) -// silentTrue := true -// silentFilter := ChatFilter{ -// PageSize: 10, -// Order: "desc", -// Silent: &silentTrue, -// } -// silentGroups, err := store.GetChats(sid, silentFilter) -// assert.Nil(t, err) -// assert.NotNil(t, silentGroups) - -// // Count total chats in all groups -// totalSilentChats := 0 -// for _, group := range silentGroups.Groups { -// totalSilentChats += len(group.Chats) -// } -// assert.Equal(t, 5, totalSilentChats, "Silent filter should return all chats") - -// // Test GetChats with silent=false filter (should only include non-silent chats) -// silentFalse := false -// nonSilentFilter := ChatFilter{ -// PageSize: 10, -// Order: "desc", -// Silent: &silentFalse, -// } -// nonSilentGroups, err := store.GetChats(sid, nonSilentFilter) -// assert.Nil(t, err) -// assert.NotNil(t, nonSilentGroups) - -// // Count total chats in all groups -// totalNonSilentChats := 0 -// for _, group := range nonSilentGroups.Groups { -// totalNonSilentChats += len(group.Chats) -// } -// assert.Equal(t, 3, totalNonSilentChats, "Non-silent filter should only return non-silent chats") -// } - -// func TestXunAttachmentCRUD(t *testing.T) { -// test.Prepare(t, config.Conf) -// defer test.Clean() -// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_attachment") - -// // Drop attachment table before test -// err := capsule.Schema().DropTableIfExists("__unit_test_conversation_attachment") -// if err != nil { -// t.Fatal(err) -// } - -// // Add a small delay to ensure table is created -// time.Sleep(100 * time.Millisecond) - -// store, err := NewXun(Setting{ -// Connector: "default", -// Prefix: "__unit_test_conversation_", -// }) -// if err != nil { -// t.Fatal(err) -// } - -// // Clean up any existing data -// _, err = store.DeleteAttachments(AttachmentFilter{}) -// assert.Nil(t, err) - -// // Test SaveAttachment (Create) -// attachment := map[string]interface{}{ -// "file_id": "test-file-123", -// "uid": "user-123", -// "manager": "local", -// "content_type": "image/jpeg", -// "name": "test-image.jpg", -// "guest": false, -// "public": true, -// "gzip": false, -// "bytes": 102400, -// "scope": []string{"user", "admin"}, -// "status": "uploaded", -// "progress": "100%", -// "error": nil, -// } - -// v, err := store.SaveAttachment(attachment) -// assert.Nil(t, err) -// fileID := v.(string) -// assert.Equal(t, "test-file-123", fileID) - -// // Test GetAttachment -// attachmentData, err := store.GetAttachment(fileID) -// assert.Nil(t, err) -// assert.NotNil(t, attachmentData) -// assert.Equal(t, "test-file-123", attachmentData["file_id"]) -// assert.Equal(t, "user-123", attachmentData["uid"]) -// assert.Equal(t, "local", attachmentData["manager"]) -// assert.Equal(t, "image/jpeg", attachmentData["content_type"]) -// assert.Equal(t, "test-image.jpg", attachmentData["name"]) -// assert.Equal(t, int64(1), attachmentData["public"]) -// assert.Equal(t, []interface{}{"user", "admin"}, attachmentData["scope"]) -// assert.Equal(t, "uploaded", attachmentData["status"]) -// assert.Equal(t, "100%", attachmentData["progress"]) -// assert.Nil(t, attachmentData["error"]) - -// // Test SaveAttachment (Update) -// attachment["name"] = "updated-image.jpg" -// attachment["bytes"] = 204800 -// attachment["status"] = "indexing" -// attachment["progress"] = "Processing..." -// attachment["error"] = "Connection timeout" -// v, err = store.SaveAttachment(attachment) -// assert.Nil(t, err) -// assert.Equal(t, "test-file-123", v.(string)) - -// // Verify update -// attachmentData, err = store.GetAttachment(fileID) -// assert.Nil(t, err) -// assert.Equal(t, "updated-image.jpg", attachmentData["name"]) -// assert.Equal(t, int64(204800), attachmentData["bytes"]) -// assert.Equal(t, "indexing", attachmentData["status"]) -// assert.Equal(t, "Processing...", attachmentData["progress"]) -// assert.Equal(t, "Connection timeout", attachmentData["error"]) - -// // Test GetAttachments with filters -// resp, err := store.GetAttachments(AttachmentFilter{ -// UID: "user-123", -// Manager: "local", -// Page: 1, -// PageSize: 10, -// }) -// assert.Nil(t, err) -// assert.Equal(t, 1, len(resp.Data)) -// assert.Equal(t, "test-file-123", resp.Data[0]["file_id"]) - -// // Test with non-existent file -// nonExistentData, err := store.GetAttachment("non-existent-file") -// assert.Error(t, err) -// assert.Nil(t, nonExistentData) -// assert.Contains(t, err.Error(), "is empty") - -// // Test DeleteAttachment -// err = store.DeleteAttachment(fileID) -// assert.Nil(t, err) - -// // Verify deletion -// _, err = store.GetAttachment(fileID) -// assert.Error(t, err) -// } - -// func TestXunKnowledgeCRUD(t *testing.T) { -// test.Prepare(t, config.Conf) -// defer test.Clean() -// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_knowledge") - -// // Drop knowledge table before test -// err := capsule.Schema().DropTableIfExists("__unit_test_conversation_knowledge") -// if err != nil { -// t.Fatal(err) -// } - -// // Add a small delay to ensure table is created -// time.Sleep(100 * time.Millisecond) - -// store, err := NewXun(Setting{ -// Connector: "default", -// Prefix: "__unit_test_conversation_", -// }) -// if err != nil { -// t.Fatal(err) -// } - -// // Clean up any existing data -// _, err = store.DeleteKnowledges(KnowledgeFilter{}) -// assert.Nil(t, err) - -// // Test SaveKnowledge (Create) -// knowledge := map[string]interface{}{ -// "collection_id": "test-collection-123", -// "name": "Test Knowledge Collection", -// "description": "A test knowledge collection for unit tests", -// "uid": "user-123", -// "public": true, -// "readonly": false, -// "system": false, -// "sort": 100, -// "cover": "cover-image.jpg", -// "scope": []string{"user", "admin"}, -// "option": map[string]interface{}{"embedding": "openai", "chunk_size": 1000}, -// } - -// v, err := store.SaveKnowledge(knowledge) -// assert.Nil(t, err) -// collectionID := v.(string) -// assert.Equal(t, "test-collection-123", collectionID) - -// // Test GetKnowledge -// knowledgeData, err := store.GetKnowledge(collectionID) -// assert.Nil(t, err) -// assert.NotNil(t, knowledgeData) -// assert.Equal(t, "test-collection-123", knowledgeData["collection_id"]) -// assert.Equal(t, "Test Knowledge Collection", knowledgeData["name"]) -// assert.Equal(t, "A test knowledge collection for unit tests", knowledgeData["description"]) -// assert.Equal(t, "user-123", knowledgeData["uid"]) -// assert.Equal(t, int64(1), knowledgeData["public"]) -// assert.Equal(t, int64(100), knowledgeData["sort"]) -// assert.Equal(t, []interface{}{"user", "admin"}, knowledgeData["scope"]) -// assert.Equal(t, map[string]interface{}{"embedding": "openai", "chunk_size": float64(1000)}, knowledgeData["option"]) - -// // Test SaveKnowledge (Update) -// knowledge["name"] = "Updated Knowledge Collection" -// knowledge["description"] = "Updated description" -// knowledge["sort"] = 200 -// v, err = store.SaveKnowledge(knowledge) -// assert.Nil(t, err) -// assert.Equal(t, "test-collection-123", v.(string)) - -// // Verify update -// knowledgeData, err = store.GetKnowledge(collectionID) -// assert.Nil(t, err) -// assert.Equal(t, "Updated Knowledge Collection", knowledgeData["name"]) -// assert.Equal(t, "Updated description", knowledgeData["description"]) -// assert.Equal(t, int64(200), knowledgeData["sort"]) - -// // Test knowledge without sort field (should get default value 9999) -// knowledgeWithoutSort := map[string]interface{}{ -// "collection_id": "test-collection-456", -// "name": "Test Knowledge Without Sort", -// "description": "Test knowledge without explicit sort value", -// "uid": "user-123", -// } -// v2, err := store.SaveKnowledge(knowledgeWithoutSort) -// assert.Nil(t, err) - -// // Verify default sort value -// knowledgeData2, err := store.GetKnowledge(v2.(string)) -// assert.Nil(t, err) -// assert.Equal(t, int64(9999), knowledgeData2["sort"]) - -// // Test GetKnowledges with filters -// resp, err := store.GetKnowledges(KnowledgeFilter{ -// UID: "user-123", -// Keywords: "Updated", -// Page: 1, -// PageSize: 10, -// }) -// assert.Nil(t, err) -// assert.Equal(t, 1, len(resp.Data)) -// assert.Equal(t, "test-collection-123", resp.Data[0]["collection_id"]) - -// // Test with non-existent collection -// nonExistentData, err := store.GetKnowledge("non-existent-collection") -// assert.Error(t, err) -// assert.Nil(t, nonExistentData) -// assert.Contains(t, err.Error(), "is empty") - -// // Test DeleteKnowledge -// err = store.DeleteKnowledge(collectionID) -// assert.Nil(t, err) -// err = store.DeleteKnowledge(v2.(string)) -// assert.Nil(t, err) - -// // Verify deletion -// _, err = store.GetKnowledge(collectionID) -// assert.Error(t, err) -// } - -// func TestXunKnowledgeFiltering(t *testing.T) { -// test.Prepare(t, config.Conf) -// defer test.Clean() -// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_knowledge") - -// // Drop knowledge table before test -// err := capsule.Schema().DropTableIfExists("__unit_test_conversation_knowledge") -// if err != nil { -// t.Fatal(err) -// } - -// // Add a small delay to ensure table is created -// time.Sleep(100 * time.Millisecond) - -// store, err := NewXun(Setting{ -// Connector: "default", -// Prefix: "__unit_test_conversation_", -// }) -// if err != nil { -// t.Fatal(err) -// } - -// // Create test data for filtering tests -// testKnowledges := []map[string]interface{}{} -// for i := 0; i < 15; i++ { -// knowledge := map[string]interface{}{ -// "collection_id": fmt.Sprintf("test-collection-%d", i), -// "name": fmt.Sprintf("Collection %d", i), -// "description": fmt.Sprintf("Description for collection %d", i), -// "uid": fmt.Sprintf("user-%d", i%3), -// "public": i%2 == 0, -// "readonly": i%3 == 0, -// "system": i%4 == 0, -// "sort": 100 + i*10, // Different sort values for testing ordering -// "cover": fmt.Sprintf("cover%d.jpg", i), -// } -// id, err := store.SaveKnowledge(knowledge) -// assert.Nil(t, err) -// knowledge["collection_id"] = id -// testKnowledges = append(testKnowledges, knowledge) -// } - -// // Test sorting functionality - should return results ordered by sort ASC then created_at DESC -// respAll, err := store.GetKnowledges(KnowledgeFilter{ -// Page: 1, -// PageSize: 15, -// }) -// assert.Nil(t, err) -// assert.Equal(t, 15, len(respAll.Data)) - -// // Verify sort order - first item should have the smallest sort value -// firstSort := respAll.Data[0]["sort"].(int64) -// lastSort := respAll.Data[len(respAll.Data)-1]["sort"].(int64) -// assert.LessOrEqual(t, firstSort, lastSort, "Results should be ordered by sort ASC") - -// // More specific sort order verification -// for i := 1; i < len(respAll.Data); i++ { -// prevSort := respAll.Data[i-1]["sort"].(int64) -// currSort := respAll.Data[i]["sort"].(int64) -// assert.LessOrEqual(t, prevSort, currSort, "Sort order should be ascending") -// } - -// // Test filtering by UID -// resp, err := store.GetKnowledges(KnowledgeFilter{ -// UID: "user-0", -// Page: 1, -// PageSize: 10, -// }) -// assert.Nil(t, err) -// assert.Greater(t, len(resp.Data), 0) - -// // Test filtering by public status -// publicTrue := true -// resp, err = store.GetKnowledges(KnowledgeFilter{ -// Public: &publicTrue, -// Page: 1, -// PageSize: 10, -// }) -// assert.Nil(t, err) -// assert.Greater(t, len(resp.Data), 0) - -// // Test filtering by readonly status -// readonlyTrue := true -// resp, err = store.GetKnowledges(KnowledgeFilter{ -// Readonly: &readonlyTrue, -// Page: 1, -// PageSize: 10, -// }) -// assert.Nil(t, err) -// assert.Greater(t, len(resp.Data), 0) - -// // Test filtering by system status -// systemTrue := true -// resp, err = store.GetKnowledges(KnowledgeFilter{ -// System: &systemTrue, -// Page: 1, -// PageSize: 10, -// }) -// assert.Nil(t, err) -// assert.Greater(t, len(resp.Data), 0) - -// // Test filtering by keywords -// resp, err = store.GetKnowledges(KnowledgeFilter{ -// Keywords: "Collection 1", -// Page: 1, -// PageSize: 10, -// }) -// assert.Nil(t, err) -// assert.Greater(t, len(resp.Data), 0) - -// // Test DeleteKnowledges with filter -// count, err := store.DeleteKnowledges(KnowledgeFilter{ -// UID: "user-0", -// }) -// assert.Nil(t, err) -// assert.Greater(t, count, int64(0)) - -// // Verify deletion -// resp, err = store.GetKnowledges(KnowledgeFilter{ -// UID: "user-0", -// }) -// assert.Nil(t, err) -// assert.Equal(t, 0, len(resp.Data)) - -// // Clean up all test data -// _, err = store.DeleteKnowledges(KnowledgeFilter{}) -// assert.Nil(t, err) -// } - -// func TestXunAttachmentFiltering(t *testing.T) { -// test.Prepare(t, config.Conf) -// defer test.Clean() -// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_attachment") - -// // Drop attachment table before test -// err := capsule.Schema().DropTableIfExists("__unit_test_conversation_attachment") -// if err != nil { -// t.Fatal(err) -// } - -// // Add a small delay to ensure table is created -// time.Sleep(100 * time.Millisecond) - -// store, err := NewXun(Setting{ -// Connector: "default", -// Prefix: "__unit_test_conversation_", -// }) -// if err != nil { -// t.Fatal(err) -// } - -// // Create test data for filtering tests -// testAttachments := []map[string]interface{}{} -// for i := 0; i < 15; i++ { -// attachment := map[string]interface{}{ -// "file_id": fmt.Sprintf("test-file-%d", i), -// "uid": fmt.Sprintf("user-%d", i%3), -// "manager": fmt.Sprintf("manager%d", i%2), -// "content_type": fmt.Sprintf("type/%d", i%4), -// "name": fmt.Sprintf("file%d.txt", i), -// "guest": i%2 == 0, -// "public": i%3 == 0, -// "gzip": i%4 == 0, -// "bytes": 1024 * (i + 1), -// "collection_id": fmt.Sprintf("collection-%d", i%5), -// } -// id, err := store.SaveAttachment(attachment) -// assert.Nil(t, err) -// attachment["file_id"] = id -// testAttachments = append(testAttachments, attachment) -// } - -// // Test filtering by UID -// resp, err := store.GetAttachments(AttachmentFilter{ -// UID: "user-0", -// Page: 1, -// PageSize: 10, -// }) -// assert.Nil(t, err) -// assert.Greater(t, len(resp.Data), 0) - -// // Test filtering by manager -// resp, err = store.GetAttachments(AttachmentFilter{ -// Manager: "manager0", -// Page: 1, -// PageSize: 10, -// }) -// assert.Nil(t, err) -// assert.Greater(t, len(resp.Data), 0) - -// // Test filtering by content_type -// resp, err = store.GetAttachments(AttachmentFilter{ -// ContentType: "type/0", -// Page: 1, -// PageSize: 10, -// }) -// assert.Nil(t, err) -// assert.Greater(t, len(resp.Data), 0) - -// // Test filtering by guest status -// guestTrue := true -// resp, err = store.GetAttachments(AttachmentFilter{ -// Guest: &guestTrue, -// Page: 1, -// PageSize: 10, -// }) -// assert.Nil(t, err) -// assert.Greater(t, len(resp.Data), 0) - -// // Test filtering by public status -// publicTrue := true -// resp, err = store.GetAttachments(AttachmentFilter{ -// Public: &publicTrue, -// Page: 1, -// PageSize: 10, -// }) -// assert.Nil(t, err) -// assert.Greater(t, len(resp.Data), 0) - -// // Test filtering by keywords -// resp, err = store.GetAttachments(AttachmentFilter{ -// Keywords: "file1", -// Page: 1, -// PageSize: 10, -// }) -// assert.Nil(t, err) -// assert.Greater(t, len(resp.Data), 0) - -// // Test DeleteAttachments with filter -// count, err := store.DeleteAttachments(AttachmentFilter{ -// Manager: "manager0", -// }) -// assert.Nil(t, err) -// assert.Greater(t, count, int64(0)) - -// // Verify deletion -// resp, err = store.GetAttachments(AttachmentFilter{ -// Manager: "manager0", -// }) -// assert.Nil(t, err) -// assert.Equal(t, 0, len(resp.Data)) - -// // Clean up all test data -// _, err = store.DeleteAttachments(AttachmentFilter{}) -// assert.Nil(t, err) -// } - -// func TestXunAttachmentStatusFields(t *testing.T) { -// test.Prepare(t, config.Conf) -// defer test.Clean() -// defer capsule.Schema().DropTableIfExists("__unit_test_conversation_attachment") - -// // Drop attachment table before test -// err := capsule.Schema().DropTableIfExists("__unit_test_conversation_attachment") -// if err != nil { -// t.Fatal(err) -// } - -// // Add a small delay to ensure table is created -// time.Sleep(100 * time.Millisecond) - -// store, err := NewXun(Setting{ -// Connector: "default", -// Prefix: "__unit_test_conversation_", -// }) -// if err != nil { -// t.Fatal(err) -// } - -// // Clean up any existing data -// _, err = store.DeleteAttachments(AttachmentFilter{}) -// assert.Nil(t, err) - -// // Test all possible enum status values -// statusValues := []string{"uploading", "uploaded", "indexing", "indexed", "upload_failed", "index_failed"} - -// for i, status := range statusValues { -// // Create attachment with specific status -// attachment := map[string]interface{}{ -// "file_id": fmt.Sprintf("test-file-%s-%d", status, i), -// "uid": "user-123", -// "manager": "local", -// "content_type": "image/jpeg", -// "name": fmt.Sprintf("test-%s.jpg", status), -// "guest": false, -// "public": true, -// "gzip": false, -// "bytes": 102400, -// "status": status, -// "progress": fmt.Sprintf("%s in progress", status), -// "error": nil, -// } - -// // Set error message for failed statuses -// if status == "upload_failed" || status == "index_failed" { -// attachment["error"] = fmt.Sprintf("%s error occurred", status) -// } - -// v, err := store.SaveAttachment(attachment) -// assert.Nil(t, err) -// fileID := v.(string) - -// // Verify the attachment was saved with correct status -// attachmentData, err := store.GetAttachment(fileID) -// assert.Nil(t, err) -// assert.Equal(t, status, attachmentData["status"]) -// assert.Equal(t, fmt.Sprintf("%s in progress", status), attachmentData["progress"]) - -// if status == "upload_failed" || status == "index_failed" { -// assert.Equal(t, fmt.Sprintf("%s error occurred", status), attachmentData["error"]) -// } else { -// assert.Nil(t, attachmentData["error"]) -// } -// } - -// // Test default status value (should be "uploading") -// attachmentWithoutStatus := map[string]interface{}{ -// "file_id": "test-file-default", -// "uid": "user-123", -// "manager": "local", -// "content_type": "image/jpeg", -// "name": "test-default.jpg", -// "guest": false, -// "public": true, -// "gzip": false, -// "bytes": 102400, -// // status not specified - should use default -// } - -// v, err := store.SaveAttachment(attachmentWithoutStatus) -// assert.Nil(t, err) -// fileID := v.(string) - -// // Verify default status -// attachmentData, err := store.GetAttachment(fileID) -// assert.Nil(t, err) -// assert.Equal(t, "uploading", attachmentData["status"]) // Should be default value -// assert.Nil(t, attachmentData["progress"]) // Should be null -// assert.Nil(t, attachmentData["error"]) // Should be null - -// // Test updating status workflow: uploading -> uploaded -> indexing -> indexed -// workflowAttachment := map[string]interface{}{ -// "file_id": "test-file-workflow", -// "uid": "user-123", -// "manager": "local", -// "content_type": "text/plain", -// "name": "workflow-test.txt", -// "status": "uploading", -// "progress": "Starting upload...", -// } - -// v, err = store.SaveAttachment(workflowAttachment) -// assert.Nil(t, err) -// workflowFileID := v.(string) - -// // Update to uploaded -// workflowAttachment["status"] = "uploaded" -// workflowAttachment["progress"] = "Upload completed, starting indexing..." -// _, err = store.SaveAttachment(workflowAttachment) -// assert.Nil(t, err) - -// attachmentData, err = store.GetAttachment(workflowFileID) -// assert.Nil(t, err) -// assert.Equal(t, "uploaded", attachmentData["status"]) -// assert.Equal(t, "Upload completed, starting indexing...", attachmentData["progress"]) - -// // Update to indexing -// workflowAttachment["status"] = "indexing" -// workflowAttachment["progress"] = "Indexing in progress..." -// _, err = store.SaveAttachment(workflowAttachment) -// assert.Nil(t, err) - -// attachmentData, err = store.GetAttachment(workflowFileID) -// assert.Nil(t, err) -// assert.Equal(t, "indexing", attachmentData["status"]) -// assert.Equal(t, "Indexing in progress...", attachmentData["progress"]) - -// // Update to indexed (final state) -// workflowAttachment["status"] = "indexed" -// workflowAttachment["progress"] = "Indexing completed" -// _, err = store.SaveAttachment(workflowAttachment) -// assert.Nil(t, err) - -// attachmentData, err = store.GetAttachment(workflowFileID) -// assert.Nil(t, err) -// assert.Equal(t, "indexed", attachmentData["status"]) -// assert.Equal(t, "Indexing completed", attachmentData["progress"]) - -// // Clean up test data -// _, err = store.DeleteAttachments(AttachmentFilter{}) -// assert.Nil(t, err) -// }