opencode/yao
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Enhance chat filtering and management in Xun store

Browse source 
- Updated the ChatFilter structure to include advanced permission filters for UserID and TeamID, allowing for more granular chat retrieval.
- Added examples in the documentation to demonstrate new filtering capabilities, including combinations of user and team filters, as well as complex conditions using QueryFilter.
- Implemented batch saving and retrieval functionalities for messages and resumes, improving data management efficiency.
- Revised related tests to validate the new filtering features and ensure robust functionality across chat management operations.
This commit is contained in:
Max 2025-12-09 11:18:55 +08:00
parent f2e0312e61
commit 76bbfa0927
8 changed files with 2703 additions and 1917 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

61
agent/store/CHAT_STORAGE_DESIGN.md
View file

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

25
agent/store/xun/chat.go
View file

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

219
agent/store/xun/chat_test.go
View file

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

351
agent/store/xun/message.go
View file

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

896
agent/store/xun/message_test.go Normal file
View file

@ -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!")
})
}

359
agent/store/xun/resume.go
View file

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

839
agent/store/xun/resume_test.go Normal file
View file

@ -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!")
})
}

1870
agent/store/xun/xun_test.go
View file

File diff suppressed because it is too large Load diff