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

Merge pull request #1373 from trheyi/main

Browse source 
Implement chat management functionalities
This commit is contained in:
Max 2025-12-09 18:59:30 +08:00 • committed by GitHub
commit 74fb693753
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
37 changed files with 11298 additions and 4328 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

119
agent/assistant/agent.go
View file

@ -57,6 +57,38 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
_, _, done := context.EnterStack(ctx, ast.ID, opts)
defer done()
// ================================================
// Initialize Chat Buffer (for root stack only)
// Buffer is flushed in defer block at the end
// ================================================
ast.InitBuffer(ctx)
// Track final status for buffer flush
var finalStatus = context.StepStatusCompleted
var finalError error
// Defer buffer flush - always executes on exit (success, error, interrupt, panic)
defer func() {
// Handle panic recovery for status tracking
if r := recover(); r != nil {
finalStatus = context.ResumeStatusFailed
if e, ok := r.(error); ok {
finalError = e
} else {
finalError = fmt.Errorf("panic: %v", r)
}
log.Error("[AGENT] Panic recovered in Stream: %v", r)
// Re-panic after flush to preserve original behavior
defer panic(r)
}
// Flush buffer to database
ast.FlushBuffer(ctx, finalStatus, finalError)
}()
// Buffer user input messages
ast.BufferUserInput(ctx, inputMessages)
// Determine stream handler
streamHandler := ast.getStreamHandler(ctx, opts)
@ -64,6 +96,8 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
// so that output adapters can use them when converting stream_start event
err = ast.initializeCapabilities(ctx, opts)
if err != nil {
finalStatus = context.ResumeStatusFailed
finalError = err
ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
return nil, err
}
@ -76,6 +110,9 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
// Use async version to not block the main flow
ast.InitializeConversationAsync(ctx, opts)
// Ensure chat session exists
ast.EnsureChat(ctx)
// Initialize agent trace node
agentNode := ast.initAgentTraceNode(ctx, inputMessages)
@ -95,15 +132,27 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
// Request Create hook ( Optional )
var createResponse *context.HookCreateResponse
if ast.HookScript != nil {
// Begin step tracking for hook_create
ast.BeginStep(ctx, context.StepTypeHookCreate, map[string]interface{}{
"messages": fullMessages,
})
var err error
createResponse, opts, err = ast.HookScript.Create(ctx, fullMessages, opts)
if err != nil {
finalStatus = context.ResumeStatusFailed
finalError = err
ast.traceAgentFail(agentNode, err)
// Send error stream_end for root stack
ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
return nil, err
}
// Complete step
ast.CompleteStep(ctx, map[string]interface{}{
"response": createResponse,
})
// Log the create response
ast.traceCreateHook(agentNode, createResponse)
}
@ -119,6 +168,8 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
// Build the LLM request first
completionMessages, completionOptions, err = ast.BuildRequest(ctx, inputMessages, createResponse)
if err != nil {
finalStatus = context.ResumeStatusFailed
finalError = err
ast.traceAgentFail(agentNode, err)
// Send error stream_end for root stack
ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
@ -128,19 +179,34 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
// Build content - convert extended types (file, data) to standard LLM types (text, image_url, input_audio)
completionMessages, err = ast.BuildContent(ctx, completionMessages, completionOptions, opts)
if err != nil {
finalStatus = context.ResumeStatusFailed
finalError = err
ast.traceAgentFail(agentNode, err)
ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
return nil, err
}
// Begin step tracking for LLM call
ast.BeginStep(ctx, context.StepTypeLLM, map[string]interface{}{
"messages": completionMessages,
})
// Execute the LLM streaming call
completionResponse, err = ast.executeLLMStream(ctx, completionMessages, completionOptions, agentNode, streamHandler, opts)
if err != nil {
finalStatus = context.ResumeStatusFailed
finalError = err
ast.traceAgentFail(agentNode, err)
// Send error stream_end for root stack
ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
return nil, err
}
// Complete LLM step
ast.CompleteStep(ctx, map[string]interface{}{
"content": completionResponse.Content,
"tool_calls": completionResponse.ToolCalls,
})
}
// ================================================
@ -155,6 +221,12 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
for attempt := 0; attempt < maxToolRetries; attempt++ {
// Begin step tracking for tool calls
ast.BeginStep(ctx, context.StepTypeTool, map[string]interface{}{
"tool_calls": currentResponse.ToolCalls,
"attempt": attempt,
})
// Execute all tool calls
toolResults, hasErrors := ast.executeToolCalls(ctx, currentResponse.ToolCalls, attempt)
@ -175,8 +247,11 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
}
}
// If all successful, break out
// If all successful, complete step and break out
if !hasErrors {
ast.CompleteStep(ctx, map[string]interface{}{
"results": toolCallResponses,
})
log.Trace("[AGENT] All tool calls succeeded (attempt %d)", attempt)
break
}
@ -193,6 +268,8 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
// If no retryable errors, don't retry (MCP internal issues)
if !hasRetryableErrors {
err := fmt.Errorf("tool calls failed with non-retryable errors (MCP internal issues)")
finalStatus = context.ResumeStatusFailed
finalError = err
log.Error("[AGENT] %v", err)
ast.traceAgentFail(agentNode, err)
ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
@ -202,19 +279,35 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
// If it's the last attempt, return error
if attempt == maxToolRetries-1 {
err := fmt.Errorf("tool calls failed after %d attempts", maxToolRetries)
finalStatus = context.ResumeStatusFailed
finalError = err
log.Error("[AGENT] %v", err)
ast.traceAgentFail(agentNode, err)
ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
return nil, err
}
// Complete current step (with partial results)
ast.CompleteStep(ctx, map[string]interface{}{
"results": toolCallResponses,
"has_errors": true,
})
// Build retry messages with tool call results (including errors)
retryMessages := ast.buildToolRetryMessages(currentMessages, currentResponse, toolResults)
// Begin LLM retry step
ast.BeginStep(ctx, context.StepTypeLLM, map[string]interface{}{
"messages": retryMessages,
"retry_attempt": attempt + 1,
})
// Retry LLM call (streaming to keep user informed)
log.Trace("[AGENT] Retrying LLM for tool call correction (attempt %d/%d)", attempt+1, maxToolRetries-1)
currentResponse, err = ast.executeLLMForToolRetry(ctx, retryMessages, completionOptions, agentNode, streamHandler, opts)
if err != nil {
finalStatus = context.ResumeStatusFailed
finalError = err
log.Error("[AGENT] LLM retry failed: %v", err)
ast.traceAgentFail(agentNode, err)
ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
@ -224,12 +317,20 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
// If LLM didn't return tool calls, it might have given up
if currentResponse.ToolCalls == nil {
err := fmt.Errorf("LLM did not return tool calls in retry attempt %d", attempt+1)
finalStatus = context.ResumeStatusFailed
finalError = err
log.Error("[AGENT] %v", err)
ast.traceAgentFail(agentNode, err)
ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
return nil, err
}
// Complete LLM retry step
ast.CompleteStep(ctx, map[string]interface{}{
"content": currentResponse.Content,
"tool_calls": currentResponse.ToolCalls,
})
// Update messages for next iteration
currentMessages = retryMessages
}
@ -245,6 +346,13 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
var nextResponse *context.NextHookResponse = nil
if ast.HookScript != nil {
// Begin step tracking for hook_next
ast.BeginStep(ctx, context.StepTypeHookNext, map[string]interface{}{
"messages": fullMessages,
"completion": completionResponse,
"tools": toolCallResponses,
})
var err error
nextResponse, opts, err = ast.HookScript.Next(ctx, &context.NextHookPayload{
Messages: fullMessages,
@ -252,11 +360,18 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
Tools: toolCallResponses,
}, opts)
if err != nil {
finalStatus = context.ResumeStatusFailed
finalError = err
ast.traceAgentFail(agentNode, err)
ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
return nil, err
}
// Complete hook_next step
ast.CompleteStep(ctx, map[string]interface{}{
"response": nextResponse,
})
// Process Next hook response
finalResponse, err = ast.processNextResponse(&NextProcessContext{
Context: ctx,
@ -268,6 +383,8 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
CreateResponse: createResponse,
})
if err != nil {
finalStatus = context.ResumeStatusFailed
finalError = err
ast.traceAgentFail(agentNode, err)
ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
return nil, err

281
agent/assistant/chat.go
View file

@ -4,9 +4,13 @@ import (
"fmt"
"strings"
"sync"
"time"
"github.com/google/uuid"
"github.com/yaoapp/kun/log"
agentcontext "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/i18n"
storetypes "github.com/yaoapp/yao/agent/store/types"
"github.com/yaoapp/yao/kb"
kbapi "github.com/yaoapp/yao/kb/api"
"github.com/yaoapp/yao/trace/types"
@ -210,6 +214,283 @@ func mergeChatMetadata(defaultMetadata map[string]interface{}, ctx *agentcontext
return metadata
}
// =============================================================================
// Chat Buffer Integration
// =============================================================================
// InitBuffer initializes the chat buffer for the context
// Should be called at the start of Stream() for root stack only
func (ast *Assistant) InitBuffer(ctx *agentcontext.Context) {
// Only initialize for root stack
if ctx.Stack == nil || !ctx.Stack.IsRoot() {
return
}
// Skip if buffer already exists
if ctx.Buffer != nil {
return
}
// Skip if History is disabled in options
if ctx.Stack.Options != nil && ctx.Stack.Options.Skip != nil && ctx.Stack.Options.Skip.History {
log.Trace("[CHAT] Buffer skipped: Skip.History is true")
return
}
// Generate request ID if not set
requestID := ctx.RequestID()
if requestID == "" {
requestID = uuid.New().String()
}
// Get connector from options
connector := ""
if ctx.Stack.Options != nil {
connector = ctx.Stack.Options.Connector
}
ctx.Buffer = agentcontext.NewChatBuffer(ctx.ChatID, requestID, ast.ID, connector)
log.Trace("[CHAT] Buffer initialized: chatID=%s, requestID=%s, assistantID=%s, connector=%s", ctx.ChatID, requestID, ast.ID, connector)
}
// BufferUserInput adds user input messages to the buffer
// Should be called after InitBuffer
func (ast *Assistant) BufferUserInput(ctx *agentcontext.Context, inputMessages []agentcontext.Message) {
if ctx.Buffer == nil {
return
}
// Convert input messages to buffer format
for _, msg := range inputMessages {
// Extract content from message
var content interface{}
var name string
content = msg.Content
if msg.Name != nil {
name = *msg.Name
}
ctx.Buffer.AddUserInput(content, name)
}
}
// UpdateSpaceSnapshot updates the space snapshot in the buffer
// Should be called when space data changes
func (ast *Assistant) UpdateSpaceSnapshot(ctx *agentcontext.Context) {
if ctx.Buffer == nil || ctx.Space == nil {
return
}
snapshot := ctx.Space.Snapshot()
ctx.Buffer.SetSpaceSnapshot(snapshot)
}
// BeginStep starts tracking an execution step
// Returns the step for further updates
func (ast *Assistant) BeginStep(ctx *agentcontext.Context, stepType string, input map[string]interface{}) *agentcontext.BufferedStep {
if ctx.Buffer == nil {
return nil
}
// Update space snapshot before beginning step
ast.UpdateSpaceSnapshot(ctx)
return ctx.Buffer.BeginStep(stepType, input, ctx.Stack)
}
// CompleteStep marks the current step as completed
func (ast *Assistant) CompleteStep(ctx *agentcontext.Context, output map[string]interface{}) {
if ctx.Buffer == nil {
return
}
ctx.Buffer.CompleteStep(output)
}
// FlushBuffer saves all buffered data to the database
// Should be called in defer block at the end of Stream()
func (ast *Assistant) FlushBuffer(ctx *agentcontext.Context, finalStatus string, err error) {
if ctx.Buffer == nil {
return
}
// Only flush for root stack
if ctx.Stack == nil || !ctx.Stack.IsRoot() {
return
}
// Get chat store
chatStore := GetChatStore()
if chatStore == nil {
log.Error("[CHAT] Chat store not available, cannot flush buffer")
return
}
// Mark current step as failed/interrupted if needed
if finalStatus != agentcontext.StepStatusCompleted && err != nil {
ctx.Buffer.FailCurrentStep(finalStatus, err)
}
// 1. Save all messages (user input + assistant responses)
messages := ast.convertBufferedMessages(ctx.Buffer.GetMessages())
if len(messages) > 0 {
if saveErr := chatStore.SaveMessages(ctx.ChatID, messages); saveErr != nil {
log.Error("[CHAT] Failed to save messages: %v", saveErr)
} else {
log.Trace("[CHAT] Saved %d messages for chat=%s", len(messages), ctx.ChatID)
}
}
// 2. Update chat last_message_at and last_connector
if len(messages) > 0 {
now := time.Now()
updates := map[string]interface{}{
"last_message_at": now,
}
// Also update last_connector if available
if connector := ctx.Buffer.Connector(); connector != "" {
updates["last_connector"] = connector
}
if updateErr := chatStore.UpdateChat(ctx.ChatID, updates); updateErr != nil {
log.Trace("[CHAT] Failed to update chat: %v", updateErr)
}
}
// 3. Only save resume steps on error/interrupt (not on success)
if finalStatus != agentcontext.StepStatusCompleted {
steps := ast.convertBufferedSteps(ctx.Buffer.GetStepsForResume(finalStatus))
if len(steps) > 0 {
if saveErr := chatStore.SaveResume(steps); saveErr != nil {
log.Error("[CHAT] Failed to save resume steps: %v", saveErr)
} else {
log.Trace("[CHAT] Saved %d resume steps for chat=%s (status=%s)", len(steps), ctx.ChatID, finalStatus)
}
}
}
}
// convertBufferedMessages converts BufferedMessage slice to store Message slice
func (ast *Assistant) convertBufferedMessages(buffered []*agentcontext.BufferedMessage) []*storetypes.Message {
if len(buffered) == 0 {
return nil
}
messages := make([]*storetypes.Message, len(buffered))
for i, msg := range buffered {
messages[i] = &storetypes.Message{
MessageID: msg.MessageID,
ChatID: msg.ChatID,
RequestID: msg.RequestID,
Role: msg.Role,
Type: msg.Type,
Props: msg.Props,
BlockID: msg.BlockID,
ThreadID: msg.ThreadID,
AssistantID: msg.AssistantID,
Connector: msg.Connector,
Sequence: msg.Sequence,
Metadata: msg.Metadata,
CreatedAt: msg.CreatedAt,
UpdatedAt: msg.CreatedAt,
}
}
return messages
}
// convertBufferedSteps converts BufferedStep slice to store Resume slice
func (ast *Assistant) convertBufferedSteps(buffered []*agentcontext.BufferedStep) []*storetypes.Resume {
if len(buffered) == 0 {
return nil
}
steps := make([]*storetypes.Resume, len(buffered))
for i, step := range buffered {
steps[i] = &storetypes.Resume{
ResumeID: step.ResumeID,
ChatID: step.ChatID,
RequestID: step.RequestID,
AssistantID: step.AssistantID,
StackID: step.StackID,
StackParentID: step.StackParentID,
StackDepth: step.StackDepth,
Type: step.Type,
Status: step.Status,
Input: step.Input,
Output: step.Output,
SpaceSnapshot: step.SpaceSnapshot,
Error: step.Error,
Sequence: step.Sequence,
Metadata: step.Metadata,
CreatedAt: step.CreatedAt,
UpdatedAt: step.CreatedAt,
}
}
return steps
}
// EnsureChat ensures a chat session exists, creates if not
func (ast *Assistant) EnsureChat(ctx *agentcontext.Context) error {
if ctx.ChatID == "" {
return nil // No chat ID, skip
}
// Skip if history is disabled
if ctx.Stack != nil && ctx.Stack.Options != nil && ctx.Stack.Options.Skip != nil && ctx.Stack.Options.Skip.History {
return nil // Skip.History is true, don't create chat session
}
chatStore := GetChatStore()
if chatStore == nil {
return nil // No store, skip
}
// Check if chat exists
_, err := chatStore.GetChat(ctx.ChatID)
if err == nil {
return nil // Chat exists
}
// Create new chat with permission fields
chat := &storetypes.Chat{
ChatID: ctx.ChatID,
AssistantID: ast.ID,
Mode: "chat",
Status: "active",
Share: "private",
Sort: 0,
CreatedAt: time.Now(),
UpdatedAt: time.Now(),
}
// Set last_connector from options (user selected connector)
if ctx.Stack != nil && ctx.Stack.Options != nil && ctx.Stack.Options.Connector != "" {
chat.LastConnector = ctx.Stack.Options.Connector
}
// Set permission fields from authorized info
if ctx.Authorized != nil {
chat.CreatedBy = ctx.Authorized.UserID
chat.UpdatedBy = ctx.Authorized.UserID
chat.TeamID = ctx.Authorized.TeamID
chat.TenantID = ctx.Authorized.TenantID
}
return chatStore.CreateChat(chat)
}
// GetChatStore returns the chat store instance
// Returns nil if storage is not configured
func GetChatStore() storetypes.ChatStore {
if storage == nil {
return nil
}
return storage
}
// =============================================================================
// Deprecated methods (kept for compatibility)
// =============================================================================
func (ast *Assistant) saveChat(ctx *agentcontext.Context, input []agentcontext.Message, opts *agentcontext.Options) error {
_ = ctx
_ = input

595
agent/assistant/chat_test.go
View file

@ -7,10 +7,13 @@ import (
"testing"
"time"
"github.com/google/uuid"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/yaoapp/gou/plan"
"github.com/yaoapp/yao/agent/assistant"
agentcontext "github.com/yaoapp/yao/agent/context"
storetypes "github.com/yaoapp/yao/agent/store/types"
"github.com/yaoapp/yao/agent/testutils"
"github.com/yaoapp/yao/kb"
oauthtypes "github.com/yaoapp/yao/openapi/oauth/types"
@ -300,3 +303,595 @@ func TestInitializeConversation(t *testing.T) {
t.Logf("✓ Correctly skipped with history flag")
})
}
// =============================================================================
// Buffer Integration Tests
// =============================================================================
func TestBufferInitialization(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
ast, err := assistant.Get("mohe")
require.NoError(t, err)
require.NotNil(t, ast)
t.Run("InitBufferForRootStack", func(t *testing.T) {
ctx := agentcontext.New(context.Background(), nil, "test_chat_buffer_001")
// Enter stack to simulate root stack
_, _, done := agentcontext.EnterStack(ctx, ast.ID, nil)
defer done()
// Initialize buffer
ast.InitBuffer(ctx)
// Verify buffer was created
assert.NotNil(t, ctx.Buffer, "Buffer should be initialized for root stack")
assert.Equal(t, "test_chat_buffer_001", ctx.Buffer.ChatID())
assert.Equal(t, ast.ID, ctx.Buffer.AssistantID())
t.Logf("✓ Buffer initialized: chatID=%s, assistantID=%s", ctx.Buffer.ChatID(), ctx.Buffer.AssistantID())
})
t.Run("SkipBufferForNestedStack", func(t *testing.T) {
ctx := agentcontext.New(context.Background(), nil, "test_chat_buffer_nested")
// Enter root stack
_, _, doneRoot := agentcontext.EnterStack(ctx, "root_assistant", nil)
defer doneRoot()
// Enter nested stack
_, _, doneNested := agentcontext.EnterStack(ctx, "nested_assistant", nil)
defer doneNested()
// Try to initialize buffer (should be skipped for nested stack)
ast.InitBuffer(ctx)
// Buffer should be nil because we're not at root
assert.Nil(t, ctx.Buffer, "Buffer should not be initialized for nested stack")
t.Logf("✓ Buffer correctly skipped for nested stack")
})
t.Run("IdempotentBufferInit", func(t *testing.T) {
ctx := agentcontext.New(context.Background(), nil, "test_chat_buffer_idem")
// Enter stack
_, _, done := agentcontext.EnterStack(ctx, ast.ID, nil)
defer done()
// Initialize buffer twice
ast.InitBuffer(ctx)
firstBuffer := ctx.Buffer
ast.InitBuffer(ctx)
secondBuffer := ctx.Buffer
// Should be the same buffer instance
assert.Same(t, firstBuffer, secondBuffer, "Buffer should be idempotent")
t.Logf("✓ Buffer initialization is idempotent")
})
}
func TestBufferUserInput(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
ast, err := assistant.Get("mohe")
require.NoError(t, err)
t.Run("BufferSimpleTextInput", func(t *testing.T) {
ctx := agentcontext.New(context.Background(), nil, "test_chat_input_001")
// Enter stack and init buffer
_, _, done := agentcontext.EnterStack(ctx, ast.ID, nil)
defer done()
ast.InitBuffer(ctx)
// Create input messages
inputMessages := []agentcontext.Message{
{
Role: agentcontext.RoleUser,
Content: "Hello, how are you?",
},
}
// Buffer user input
ast.BufferUserInput(ctx, inputMessages)
// Verify buffer contains the message
messages := ctx.Buffer.GetMessages()
assert.Len(t, messages, 1, "Should have 1 buffered message")
assert.Equal(t, "user", messages[0].Role)
assert.Equal(t, "user_input", messages[0].Type)
assert.Equal(t, "Hello, how are you?", messages[0].Props["content"])
t.Logf("✓ User input buffered: %v", messages[0].Props)
})
t.Run("BufferMultipleMessages", func(t *testing.T) {
ctx := agentcontext.New(context.Background(), nil, "test_chat_input_multi")
// Enter stack and init buffer
_, _, done := agentcontext.EnterStack(ctx, ast.ID, nil)
defer done()
ast.InitBuffer(ctx)
// Create multiple input messages
inputMessages := []agentcontext.Message{
{Role: agentcontext.RoleUser, Content: "First message"},
{Role: agentcontext.RoleUser, Content: "Second message"},
}
// Buffer user input
ast.BufferUserInput(ctx, inputMessages)
// Verify buffer contains all messages
messages := ctx.Buffer.GetMessages()
assert.Len(t, messages, 2, "Should have 2 buffered messages")
assert.Equal(t, 1, messages[0].Sequence)
assert.Equal(t, 2, messages[1].Sequence)
t.Logf("✓ Multiple messages buffered with correct sequence")
})
t.Run("BufferWithNilBuffer", func(t *testing.T) {
ctx := agentcontext.New(context.Background(), nil, "test_chat_input_nil")
// Don't initialize buffer
inputMessages := []agentcontext.Message{
{Role: agentcontext.RoleUser, Content: "Test"},
}
// Should not panic
ast.BufferUserInput(ctx, inputMessages)
t.Logf("✓ BufferUserInput handles nil buffer gracefully")
})
}
func TestBufferStepTracking(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
ast, err := assistant.Get("mohe")
require.NoError(t, err)
t.Run("BeginAndCompleteStep", func(t *testing.T) {
ctx := agentcontext.New(context.Background(), nil, "test_chat_step_001")
ctx.Space = plan.NewMemorySharedSpace()
// Enter stack and init buffer
_, _, done := agentcontext.EnterStack(ctx, ast.ID, nil)
defer done()
ast.InitBuffer(ctx)
// Set some space data
ctx.Space.Set("test_key", "test_value")
// Begin a step
step := ast.BeginStep(ctx, agentcontext.StepTypeLLM, map[string]interface{}{
"messages": []string{"Hello"},
})
assert.NotNil(t, step, "Step should be created")
assert.Equal(t, agentcontext.StepTypeLLM, step.Type)
assert.Equal(t, agentcontext.StepStatusRunning, step.Status)
assert.NotEmpty(t, step.StackID)
// Complete the step
ast.CompleteStep(ctx, map[string]interface{}{
"content": "Response",
})
// Verify step is completed
steps := ctx.Buffer.GetAllSteps()
assert.Len(t, steps, 1)
assert.Equal(t, agentcontext.StepStatusCompleted, steps[0].Status)
assert.Equal(t, "Response", steps[0].Output["content"])
t.Logf("✓ Step tracking works correctly")
})
t.Run("SpaceSnapshotCapture", func(t *testing.T) {
ctx := agentcontext.New(context.Background(), nil, "test_chat_space_001")
ctx.Space = plan.NewMemorySharedSpace()
// Enter stack and init buffer
_, _, done := agentcontext.EnterStack(ctx, ast.ID, nil)
defer done()
ast.InitBuffer(ctx)
// Set space data before step
ctx.Space.Set("key1", "value1")
ctx.Space.Set("key2", 123)
// Begin step (should capture space snapshot)
ast.BeginStep(ctx, agentcontext.StepTypeHookCreate, nil)
// Verify space snapshot was captured
steps := ctx.Buffer.GetAllSteps()
require.Len(t, steps, 1)
assert.NotNil(t, steps[0].SpaceSnapshot)
assert.Equal(t, "value1", steps[0].SpaceSnapshot["key1"])
assert.Equal(t, 123, steps[0].SpaceSnapshot["key2"])
t.Logf("✓ Space snapshot captured: %v", steps[0].SpaceSnapshot)
})
t.Run("MultipleSteps", func(t *testing.T) {
ctx := agentcontext.New(context.Background(), nil, "test_chat_multi_step")
ctx.Space = plan.NewMemorySharedSpace()
// Enter stack and init buffer
_, _, done := agentcontext.EnterStack(ctx, ast.ID, nil)
defer done()
ast.InitBuffer(ctx)
// Step 1: hook_create
ast.BeginStep(ctx, agentcontext.StepTypeHookCreate, map[string]interface{}{"phase": "create"})
ast.CompleteStep(ctx, map[string]interface{}{"result": "created"})
// Step 2: llm
ast.BeginStep(ctx, agentcontext.StepTypeLLM, map[string]interface{}{"phase": "llm"})
ast.CompleteStep(ctx, map[string]interface{}{"result": "completed"})
// Step 3: hook_next
ast.BeginStep(ctx, agentcontext.StepTypeHookNext, map[string]interface{}{"phase": "next"})
ast.CompleteStep(ctx, map[string]interface{}{"result": "done"})
// Verify all steps
steps := ctx.Buffer.GetAllSteps()
assert.Len(t, steps, 3)
assert.Equal(t, agentcontext.StepTypeHookCreate, steps[0].Type)
assert.Equal(t, agentcontext.StepTypeLLM, steps[1].Type)
assert.Equal(t, agentcontext.StepTypeHookNext, steps[2].Type)
t.Logf("✓ Multiple steps tracked correctly")
})
}
func TestFlushBuffer(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
ast, err := assistant.Get("mohe")
require.NoError(t, err)
// Skip if chat store not available
chatStore := assistant.GetChatStore()
if chatStore == nil {
t.Skip("Chat store not configured, skipping flush tests")
}
t.Run("FlushOnSuccess", func(t *testing.T) {
chatID := fmt.Sprintf("test_flush_success_%s", uuid.New().String()[:8])
ctx := agentcontext.New(context.Background(), nil, chatID)
ctx.Space = plan.NewMemorySharedSpace()
// Enter stack and init buffer
_, _, done := agentcontext.EnterStack(ctx, ast.ID, nil)
defer done()
ast.InitBuffer(ctx)
// Ensure chat exists
err := chatStore.CreateChat(&storetypes.Chat{
ChatID: chatID,
AssistantID: ast.ID,
Mode: "chat",
Status: "active",
Share: "private",
CreatedAt: time.Now(),
UpdatedAt: time.Now(),
})
require.NoError(t, err)
// Add some messages to buffer
require.NotNil(t, ctx.Buffer, "Buffer should be initialized")
ctx.Buffer.AddUserInput("Test question", "")
ctx.Buffer.AddAssistantMessage("M1", "text", map[string]interface{}{"content": "Test answer"}, "", "", ast.ID, nil)
// Add a step
ast.BeginStep(ctx, agentcontext.StepTypeLLM, nil)
ast.CompleteStep(ctx, nil)
// Flush buffer (success case)
ast.FlushBuffer(ctx, agentcontext.StepStatusCompleted, nil)
// Verify messages were saved
messages, err := chatStore.GetMessages(chatID, storetypes.MessageFilter{})
assert.NoError(t, err)
assert.Len(t, messages, 2, "Should have 2 messages saved")
// Verify no resume records (success case)
resumes, err := chatStore.GetResume(chatID)
assert.NoError(t, err)
assert.Len(t, resumes, 0, "Should have no resume records on success")
// Cleanup
chatStore.DeleteChat(chatID)
t.Logf("✓ Buffer flushed on success: %d messages saved, no resume records", len(messages))
})
t.Run("FlushOnFailure", func(t *testing.T) {
chatID := fmt.Sprintf("test_flush_fail_%s", uuid.New().String()[:8])
ctx := agentcontext.New(context.Background(), nil, chatID)
ctx.Space = plan.NewMemorySharedSpace()
// Enter stack and init buffer
_, _, done := agentcontext.EnterStack(ctx, ast.ID, nil)
defer done()
ast.InitBuffer(ctx)
// Ensure chat exists
err := chatStore.CreateChat(&storetypes.Chat{
ChatID: chatID,
AssistantID: ast.ID,
Mode: "chat",
Status: "active",
Share: "private",
CreatedAt: time.Now(),
UpdatedAt: time.Now(),
})
require.NoError(t, err)
// Add messages
ctx.Buffer.AddUserInput("Test question", "")
// Add a step that will "fail"
ast.BeginStep(ctx, agentcontext.StepTypeLLM, map[string]interface{}{"test": "data"})
// Don't complete - simulate failure
// Flush buffer (failure case)
testErr := fmt.Errorf("simulated error")
ast.FlushBuffer(ctx, agentcontext.ResumeStatusFailed, testErr)
// Verify messages were saved
messages, err := chatStore.GetMessages(chatID, storetypes.MessageFilter{})
assert.NoError(t, err)
assert.Len(t, messages, 1, "Should have 1 message saved")
// Verify resume records were saved
resumes, err := chatStore.GetResume(chatID)
assert.NoError(t, err)
assert.Len(t, resumes, 1, "Should have 1 resume record on failure")
assert.Equal(t, agentcontext.ResumeStatusFailed, resumes[0].Status)
// Cleanup
chatStore.DeleteResume(chatID)
chatStore.DeleteChat(chatID)
t.Logf("✓ Buffer flushed on failure: messages and resume records saved")
})
t.Run("FlushOnInterrupt", func(t *testing.T) {
chatID := fmt.Sprintf("test_flush_interrupt_%s", uuid.New().String()[:8])
ctx := agentcontext.New(context.Background(), nil, chatID)
ctx.Space = plan.NewMemorySharedSpace()
// Enter stack and init buffer
_, _, done := agentcontext.EnterStack(ctx, ast.ID, nil)
defer done()
ast.InitBuffer(ctx)
// Ensure chat exists
err := chatStore.CreateChat(&storetypes.Chat{
ChatID: chatID,
AssistantID: ast.ID,
Mode: "chat",
Status: "active",
Share: "private",
CreatedAt: time.Now(),
UpdatedAt: time.Now(),
})
require.NoError(t, err)
// Add messages and steps
ctx.Buffer.AddUserInput("Test question", "")
ast.BeginStep(ctx, agentcontext.StepTypeLLM, nil)
// Flush buffer (interrupt case)
ast.FlushBuffer(ctx, agentcontext.ResumeStatusInterrupted, nil)
// Verify resume records were saved with interrupted status
resumes, err := chatStore.GetResume(chatID)
assert.NoError(t, err)
assert.Len(t, resumes, 1, "Should have 1 resume record on interrupt")
assert.Equal(t, agentcontext.ResumeStatusInterrupted, resumes[0].Status)
// Cleanup
chatStore.DeleteResume(chatID)
chatStore.DeleteChat(chatID)
t.Logf("✓ Buffer flushed on interrupt: resume records saved with interrupted status")
})
}
func TestEnsureChat(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
ast, err := assistant.Get("mohe")
require.NoError(t, err)
// Skip if chat store not available
chatStore := assistant.GetChatStore()
if chatStore == nil {
t.Skip("Chat store not configured, skipping EnsureChat tests")
}
t.Run("CreateNewChat", func(t *testing.T) {
chatID := fmt.Sprintf("test_ensure_new_%s", uuid.New().String()[:8])
ctx := agentcontext.New(context.Background(), nil, chatID)
// Ensure chat creates it
err := ast.EnsureChat(ctx)
assert.NoError(t, err)
// Verify chat was created
chat, err := chatStore.GetChat(chatID)
assert.NoError(t, err)
assert.NotNil(t, chat)
assert.Equal(t, chatID, chat.ChatID)
assert.Equal(t, ast.ID, chat.AssistantID)
assert.Equal(t, "active", chat.Status)
// Cleanup
chatStore.DeleteChat(chatID)
t.Logf("✓ New chat created: %s", chatID)
})
t.Run("SkipExistingChat", func(t *testing.T) {
chatID := fmt.Sprintf("test_ensure_exist_%s", uuid.New().String()[:8])
ctx := agentcontext.New(context.Background(), nil, chatID)
// Create chat first
err := chatStore.CreateChat(&storetypes.Chat{
ChatID: chatID,
AssistantID: ast.ID,
Title: "Existing Chat",
Mode: "chat",
Status: "active",
Share: "private",
CreatedAt: time.Now(),
UpdatedAt: time.Now(),
})
require.NoError(t, err)
// EnsureChat should not error
err = ast.EnsureChat(ctx)
assert.NoError(t, err)
// Verify chat still has original title
chat, err := chatStore.GetChat(chatID)
assert.NoError(t, err)
assert.Equal(t, "Existing Chat", chat.Title)
// Cleanup
chatStore.DeleteChat(chatID)
t.Logf("✓ Existing chat preserved")
})
t.Run("SkipEmptyChatID", func(t *testing.T) {
ctx := agentcontext.New(context.Background(), nil, "")
// Should not error with empty chat ID
err := ast.EnsureChat(ctx)
assert.NoError(t, err)
t.Logf("✓ Empty chat ID handled gracefully")
})
t.Run("CreateChatWithPermissions", func(t *testing.T) {
chatID := fmt.Sprintf("test_ensure_perm_%s", uuid.New().String()[:8])
// Create context with authorized info
ctx := agentcontext.New(context.Background(), &oauthtypes.AuthorizedInfo{
UserID: "test_user_001",
TeamID: "test_team_001",
TenantID: "test_tenant_001",
}, chatID)
// EnsureChat should create with permission fields
err := ast.EnsureChat(ctx)
assert.NoError(t, err)
// Verify permission fields were saved
chat, err := chatStore.GetChat(chatID)
assert.NoError(t, err)
assert.NotNil(t, chat)
assert.Equal(t, "test_user_001", chat.CreatedBy, "CreatedBy should be set")
assert.Equal(t, "test_user_001", chat.UpdatedBy, "UpdatedBy should be set")
assert.Equal(t, "test_team_001", chat.TeamID, "TeamID should be set")
assert.Equal(t, "test_tenant_001", chat.TenantID, "TenantID should be set")
// Cleanup
chatStore.DeleteChat(chatID)
t.Logf("✓ Chat created with permission fields: user=%s, team=%s, tenant=%s",
chat.CreatedBy, chat.TeamID, chat.TenantID)
})
t.Run("SkipHistoryEnabled", func(t *testing.T) {
chatID := fmt.Sprintf("test_ensure_skip_%s", uuid.New().String()[:8])
// Create context
ctx := agentcontext.New(context.Background(), nil, chatID)
// Set up stack with Skip.History = true
ctx.Stack = &agentcontext.Stack{
ID: "test_stack",
AssistantID: ast.ID,
Depth: 0,
Options: &agentcontext.Options{
Skip: &agentcontext.Skip{
History: true,
},
},
}
// EnsureChat should NOT create chat when Skip.History is true
err := ast.EnsureChat(ctx)
assert.NoError(t, err)
// Verify chat was NOT created
_, err = chatStore.GetChat(chatID)
assert.Error(t, err, "Chat should not be created when Skip.History is true")
t.Logf("✓ Chat not created when Skip.History is true")
})
}
func TestConvertBufferedTypes(t *testing.T) {
t.Run("ConvertBufferedMessages", func(t *testing.T) {
// Create buffered messages
buffered := []*agentcontext.BufferedMessage{
{
MessageID: "msg_001",
ChatID: "chat_001",
RequestID: "req_001",
Role: "user",
Type: "user_input",
Props: map[string]interface{}{"content": "Hello"},
Sequence: 1,
CreatedAt: time.Now(),
},
{
MessageID: "msg_002",
ChatID: "chat_001",
RequestID: "req_001",
Role: "assistant",
Type: "text",
Props: map[string]interface{}{"content": "Hi there!"},
BlockID: "block_001",
AssistantID: "test_assistant",
Sequence: 2,
CreatedAt: time.Now(),
},
}
// Verify structure matches store types
assert.Len(t, buffered, 2)
assert.Equal(t, "user", buffered[0].Role)
assert.Equal(t, "assistant", buffered[1].Role)
assert.Equal(t, "block_001", buffered[1].BlockID)
t.Logf("✓ Buffered messages have correct structure")
})
t.Run("ConvertBufferedSteps", func(t *testing.T) {
// Create buffered steps
buffered := []*agentcontext.BufferedStep{
{
ResumeID: "resume_001",
ChatID: "chat_001",
RequestID: "req_001",
AssistantID: "test_assistant",
StackID: "stack_001",
StackDepth: 0,
Type: agentcontext.StepTypeLLM,
Status: agentcontext.ResumeStatusFailed,
Input: map[string]interface{}{"messages": []string{"Hello"}},
SpaceSnapshot: map[string]interface{}{"key": "value"},
Error: "Test error",
Sequence: 1,
CreatedAt: time.Now(),
},
}
// Verify structure
assert.Len(t, buffered, 1)
assert.Equal(t, agentcontext.StepTypeLLM, buffered[0].Type)
assert.Equal(t, agentcontext.ResumeStatusFailed, buffered[0].Status)
assert.Equal(t, "Test error", buffered[0].Error)
assert.Equal(t, "value", buffered[0].SpaceSnapshot["key"])
t.Logf("✓ Buffered steps have correct structure")
})
}

52
agent/assistant/handlers/stream.go
View file

@ -323,6 +323,58 @@ func (s *streamState) handleMessageEnd(data []byte) int {
threadID = s.ctx.Stack.ID
}
// Get BlockID from metadata if available
var blockID string
if s.ctx != nil {
if metadata := s.ctx.GetMessageMetadata(s.currentGroupID); metadata != nil {
blockID = metadata.BlockID
}
}
// Buffer the complete LLM message for storage
// Delta chunks are not stored, but we need to save the final complete content
// Skip if History is disabled in options
shouldSkipHistory := s.ctx.Stack != nil && s.ctx.Stack.Options != nil &&
s.ctx.Stack.Options.Skip != nil && s.ctx.Stack.Options.Skip.History
if s.ctx.Buffer != nil && len(s.buffer) > 0 && !shouldSkipHistory {
assistantID := ""
if s.ctx.Stack != nil {
assistantID = s.ctx.Stack.AssistantID
}
// Build props based on message type
var props map[string]interface{}
if msgType == message.TypeToolCall {
// For tool calls, try to parse the accumulated buffer as JSON
var toolCallData interface{}
if err := jsoniter.Unmarshal(s.buffer, &toolCallData); err == nil {
props = map[string]interface{}{
"calls": toolCallData,
}
} else {
props = map[string]interface{}{
"content": string(s.buffer),
}
}
} else {
// For text/thinking, content is the accumulated text
props = map[string]interface{}{
"content": string(s.buffer),
}
}
s.ctx.Buffer.AddAssistantMessage(
s.currentGroupID, // Use the message ID
msgType,
props,
blockID,
threadID,
assistantID,
nil,
)
}
// Build EventMessageEndData with complete content
endData := message.EventMessageEndData{
MessageID: s.currentGroupID, // Use the message ID

914
agent/context/JSAPI.md
View file

File diff suppressed because it is too large Load diff

886
agent/context/JSAPI_OUTPUT.md
View file

@ -1,886 +0,0 @@
# Context Output JS API
The Context object provides `Send`, `SendGroup`, `SendGroupStart`, and `SendGroupEnd` methods for sending messages to clients from JavaScript within Agent Hook functions.
## Hook Functions Overview
Agent Hook functions are lifecycle callbacks that allow you to customize the behavior of AI assistants. The Context object passed to these hooks includes output methods for real-time communication with clients.
### Available Hooks
- `Create(ctx, messages)` - Called before the assistant processes messages
- `Before(ctx, messages, response)` - Called before sending LLM response
- `After(ctx, messages, response)` - Called after receiving LLM response
- `Done(ctx, messages, response)` - Called after assistant completes
- `Error(ctx, messages, error)` - Called when an error occurs
## Quick Start
### Basic Usage in Create Hook
```javascript
/**
* Create hook - send initial messages to client
*/
function Create(ctx, messages) {
// Send welcome message (string shorthand, auto-flushes)
ctx.Send("Welcome! Let me help you with that...");
// Send loading indicator (auto-flushes)
ctx.Send({
type: "loading",
props: { message: "Analyzing your request..." },
});
// Continue with normal processing
return { messages };
}
```
### Streaming Updates Example
```javascript
/**
* Create hook - demonstrate streaming updates
*/
function Create(ctx, messages) {
// Send initial message
ctx.Send({
type: "text",
props: { content: "Processing" },
id: "status_msg",
});
ctx.Flush();
time.Sleep(500); // Simulate work
// Append to message (delta update)
ctx.Send({
type: "text",
props: { content: "..." },
id: "status_msg",
delta: true,
delta_path: "content",
delta_action: "append",
});
ctx.Flush();
time.Sleep(500); // More work
// Complete the message
ctx.Send({
type: "text",
props: { content: " Done!" },
id: "status_msg",
delta: true,
delta_path: "content",
delta_action: "append",
});
ctx.Flush();
return { messages };
}
```
## API Reference
### ctx.Send(message)
Send a single message to the client.
**String Shorthand:**
```javascript
ctx.Send("Hello World");
```
**Object Format:**
```javascript
// Text message
ctx.Send({
type: "text",
props: { content: "Hello from JavaScript" },
});
// Loading indicator
ctx.Send({
type: "loading",
props: { message: "Processing..." },
});
// Error message
ctx.Send({
type: "error",
props: { message: "Something went wrong", code: "ERR_500" },
});
```
**Complete Message Object:**
```javascript
ctx.Send({
type: "text",
props: { content: "Hello" },
id: "msg_123", // Optional: message ID for delta updates
delta: true, // Optional: incremental update flag
done: false, // Optional: completion flag
delta_path: "content", // Optional: update path
delta_action: "append", // Optional: append, replace, merge, set
group_id: "grp_1", // Optional: message group ID
metadata: {
// Optional: custom metadata
timestamp: Date.now(),
sequence: 1,
trace_id: "trace_123",
},
});
```
### ctx.SendGroup(group)
Send a group of related messages together.
```javascript
ctx.SendGroup({
id: "group_123",
messages: [
{ type: "text", props: { content: "First message" } },
{ type: "text", props: { content: "Second message" } },
],
metadata: { timestamp: Date.now() },
});
```
### ctx.SendGroupStart(type?, id?)
Start a message group and return the group ID. Messages sent after this should include the returned `group_id`.
**Parameters:**
- `type` (optional): Group type (`"text"`, `"thinking"`, `"tool_call"`, `"mixed"`), defaults to `"mixed"`
- `id` (optional): Custom group ID, auto-generates if not provided
**Returns:** Group ID (string)
```javascript
// Auto-generate ID with default type
const groupId = ctx.SendGroupStart();
// Specify type, auto-generate ID
const groupId = ctx.SendGroupStart("text");
// Specify both type and custom ID
const groupId = ctx.SendGroupStart("thinking", "my-group-123");
```
### ctx.SendGroupEnd(id, chunkCount?)
End a message group.
**Parameters:**
- `id` (required): Group ID returned from `SendGroupStart`
- `chunkCount` (optional): Number of messages in the group
```javascript
// Basic usage
ctx.SendGroupEnd(groupId);
// With chunk count
ctx.SendGroupEnd(groupId, 5);
```
## Complete Hook Examples
### 1. Create Hook - Welcome Message
```javascript
/**
* Send welcome message when conversation starts
*/
function Create(ctx, messages) {
// Send welcome message (auto-flushes)
ctx.Send("Welcome to AI Assistant! How can I help you today?");
// Return messages to continue processing
return { messages };
}
```
### 2. Create Hook - Progress Updates
```javascript
/**
* Show progress indicators during preprocessing
*/
function Create(ctx, messages) {
// Step 1: Analyzing
ctx.Send({
type: "loading",
props: { message: "Analyzing your request..." },
});
ctx.Flush();
// Perform analysis...
const userIntent = analyzeIntent(messages);
// Step 2: Searching
ctx.Send({
type: "loading",
props: { message: "Searching knowledge base..." },
});
ctx.Flush();
// Search knowledge base...
const context = searchKnowledgeBase(userIntent);
// Add context to messages
if (context) {
messages.unshift({
role: "system",
content: `Context: ${context}`,
});
}
return { messages };
}
```
### 3. Before Hook - Show Thinking Process
```javascript
/**
* Display model's reasoning before sending response
*/
function Before(ctx, messages, response) {
// If response includes thinking/reasoning
if (response.thinking) {
ctx.Send({
type: "thinking",
props: { content: response.thinking },
});
ctx.Flush();
}
return { response };
}
```
### 4. After Hook - Process Tool Calls
```javascript
/**
* Handle tool calls and send results
*/
function After(ctx, messages, response) {
// Process tool calls
if (response.tool_calls && response.tool_calls.length > 0) {
response.tool_calls.forEach((toolCall) => {
// Show tool being called
ctx.Send({
type: "tool_call",
props: {
id: toolCall.id,
name: toolCall.function.name,
arguments: toolCall.function.arguments,
},
});
ctx.Flush();
// Execute tool and send result
const result = executeTool(toolCall);
ctx.Send({
type: "text",
props: { content: `Tool result: ${result}` },
});
ctx.Flush();
});
}
return { response };
}
```
### 5. Done Hook - Completion Message
```javascript
/**
* Send completion message and cleanup
*/
function Done(ctx, messages, response) {
// Send completion indicator
ctx.Send({
type: "text",
props: { content: "\n✅ Task completed successfully!" },
});
ctx.Flush();
// Log metrics
console.log("Conversation completed:", {
chat_id: ctx.chat_id,
message_count: messages.length,
tokens_used: response.usage?.total_tokens,
});
return {};
}
```
### 6. Error Hook - Handle Errors Gracefully
```javascript
/**
* Send user-friendly error messages
*/
function Error(ctx, messages, error) {
console.error("Assistant error:", error);
// Send error message to user
ctx.Send({
type: "error",
props: {
message: "I encountered an issue while processing your request.",
code: error.code || "UNKNOWN_ERROR",
details:
process.env.YAO_ENV === "development" ? error.message : undefined,
},
});
ctx.Flush();
// Return error to be logged
return { error };
}
```
### 7. Multi-Step Process with Progress
```javascript
/**
* Complex processing with multiple steps
*/
function Create(ctx, messages) {
const steps = [
{ name: "Validating input", duration: 500 },
{ name: "Loading context", duration: 1000 },
{ name: "Preparing response", duration: 800 },
];
// Create progress message
const progressId = "progress_" + Date.now();
steps.forEach((step, index) => {
// Update progress
ctx.Send({
type: "loading",
props: {
message: `${step.name}... (${index + 1}/${steps.length})`,
},
id: progressId,
delta: index > 0,
});
ctx.Flush();
// Simulate work
time.Sleep(step.duration);
});
// Clear progress indicator
ctx.Send({
type: "loading",
props: { message: "" },
id: progressId,
done: true,
});
ctx.Flush();
return { messages };
}
```
### 8. Real-time Streaming Updates
```javascript
/**
* Send streaming updates as processing progresses
*/
function Create(ctx, messages) {
const messageId = "stream_" + Date.now();
// Start message
ctx.Send({
type: "text",
props: { content: "Processing" },
id: messageId,
});
ctx.Flush();
// Simulate incremental processing
const updates = [".", ".", ".", " analyzing", ".", ".", ".", " complete!"];
updates.forEach((update) => {
time.Sleep(200);
ctx.Send({
type: "text",
props: { content: update },
id: messageId,
delta: true,
delta_path: "content",
delta_action: "append",
});
ctx.Flush();
});
return { messages };
}
```
### 9. Message Groups for Related Content (High-level API)
```javascript
/**
* Send groups of related messages together using SendGroup (auto-handles events)
*/
function Before(ctx, messages, response) {
// SendGroup automatically sends group_start and group_end events
ctx.SendGroup({
messages: [
{
type: "text",
props: { content: "**Context Information:**" },
},
{
type: "text",
props: { content: `User: ${ctx.authorized?.user_id || "Anonymous"}` },
},
{
type: "text",
props: { content: `Session: ${ctx.chat_id}` },
},
{
type: "text",
props: { content: `Locale: ${ctx.locale}` },
},
],
metadata: {
timestamp: Date.now(),
type: "context",
},
});
return { response };
}
```
### 10. Manual Group Control (Low-level API)
```javascript
/**
* Manually control group boundaries with SendGroupStart and SendGroupEnd
*/
function Create(ctx, messages) {
// Start a text group
const groupId = ctx.SendGroupStart("text");
// Send messages with group_id
ctx.Send({
type: "text",
props: { content: "First message in group" },
group_id: groupId,
});
ctx.Send({
type: "text",
props: { content: "Second message in group" },
group_id: groupId,
});
// End the group
ctx.SendGroupEnd(groupId, 2);
return { messages };
}
```
### 11. Streaming with Groups
```javascript
/**
* Stream delta updates within a group
*/
function Create(ctx, messages) {
// Start thinking group
const thinkingId = ctx.SendGroupStart("thinking");
// Stream thinking process
const steps = ["Analyzing", "Processing", "Generating"];
const msgId = "thinking_msg";
steps.forEach((step, i) => {
if (i === 0) {
// First message
ctx.Send({
type: "thinking",
props: { content: step },
id: msgId,
group_id: thinkingId,
delta: false,
});
} else {
// Delta updates
ctx.Send({
type: "thinking",
props: { content: ` → ${step}` },
id: msgId,
group_id: thinkingId,
delta: true,
delta_path: "content",
delta_action: "append",
});
}
});
// End thinking group
ctx.SendGroupEnd(thinkingId, steps.length);
return { messages };
}
```
## Message Types
Built-in message types supported:
- `user_input` - User input (display only)
- `text` - Text content (supports Markdown)
- `thinking` - Reasoning/thinking process
- `loading` - Loading indicator
- `tool_call` - Tool/function call
- `error` - Error message
- `image` - Image content
- `audio` - Audio content
- `video` - Video content
- `action` - System action (silent in OpenAI clients)
- `event` - Lifecycle event (CUI only)
## Message Props by Type
### Text Message
```javascript
{
type: "text",
props: {
content: "Text content (supports Markdown)"
}
}
```
### Thinking Message
```javascript
{
type: "thinking",
props: {
content: "Reasoning process..."
}
}
```
### Loading Message
```javascript
{
type: "loading",
props: {
message: "Loading message..."
}
}
```
### Tool Call Message
```javascript
{
type: "tool_call",
props: {
id: "call_123",
name: "function_name",
arguments: '{"key": "value"}'
}
}
```
### Error Message
```javascript
{
type: "error",
props: {
message: "Error message",
code: "ERROR_CODE",
details: "Additional details"
}
}
```
### Image Message
```javascript
{
type: "image",
props: {
url: "https://example.com/image.jpg",
alt: "Image description",
width: 800,
height: 600
}
}
```
### Audio Message
```javascript
{
type: "audio",
props: {
url: "https://example.com/audio.mp3",
format: "mp3",
duration: 120.5,
transcript: "Audio transcript...",
autoplay: false,
controls: true
}
}
```
### Video Message
```javascript
{
type: "video",
props: {
url: "https://example.com/video.mp4",
format: "mp4",
thumbnail: "https://example.com/thumb.jpg",
width: 1920,
height: 1080,
autoplay: false,
controls: true
}
}
```
## Delta Updates
Use delta updates for streaming scenarios:
```javascript
// Initial message
ctx.Send({
type: "text",
props: { content: "Hello" },
id: "msg_1",
delta: false,
});
// Append to content
ctx.Send({
type: "text",
props: { content: " World" },
id: "msg_1",
delta: true,
delta_path: "content",
delta_action: "append",
});
// Mark as complete
ctx.Send({
type: "text",
props: {},
id: "msg_1",
done: true,
});
```
**Delta Actions:**
- `append` - Append to string or array
- `replace` - Replace value
- `merge` - Merge objects
- `set` - Set new field
## Hook Function Patterns
### Pattern 1: Fire-and-Forget Notifications
```javascript
function Create(ctx, messages) {
ctx.Send("Starting processing..."); // Auto-flushes
// Continue processing immediately
return { messages };
}
```
### Pattern 2: Progress Tracking
```javascript
function Create(ctx, messages) {
const stages = ["validate", "analyze", "prepare"];
stages.forEach((stage) => {
ctx.Send({ type: "loading", props: { message: `${stage}...` } }); // Auto-flushes
performStage(stage);
});
return { messages };
}
```
### Pattern 3: Conditional Messaging
```javascript
function Before(ctx, messages, response) {
// Only show reasoning for complex queries
if (messages[messages.length - 1].content.length > 100) {
ctx.Send({
type: "thinking",
props: { content: "Analyzing complex query..." },
}); // Auto-flushes
}
return { response };
}
```
### Pattern 4: Error Recovery
```javascript
function Error(ctx, messages, error) {
if (error.code === "RATE_LIMIT") {
ctx.Send("Service is busy, retrying..."); // Auto-flushes
time.Sleep(1000);
return { retry: true };
}
ctx.Send({
type: "error",
props: { message: "Sorry, something went wrong.", code: error.code },
}); // Auto-flushes
return { error };
}
```
## Important Notes
### 1. Hook Function Signatures
Each hook receives different parameters:
- `Create(ctx, messages)` - Context and input messages
- `Before(ctx, messages, response)` - Context, messages, and LLM response
- `After(ctx, messages, response)` - Context, messages, and LLM response
- `Done(ctx, messages, response)` - Context, messages, and final response
- `Error(ctx, messages, error)` - Context, messages, and error object
### 2. Messages Auto-Flush for Real-time Updates
```javascript
// Messages are automatically flushed after each Send
ctx.Send("Processing..."); // Sent immediately to client
// Multiple sends work seamlessly
ctx.Send("Step 1"); // Flushed
ctx.Send("Step 2"); // Flushed
ctx.Send("Step 3"); // Flushed
```
### 3. Delta Updates Require Unique IDs
```javascript
// Initial message
ctx.Send({ type: "text", props: { content: "Step 1" }, id: "progress" });
// Update same message
ctx.Send({
type: "text",
props: { content: ", Step 2" },
id: "progress",
delta: true,
delta_path: "content",
delta_action: "append",
});
```
### 4. Message Types and Client Support
- **OpenAI Client** (`ctx.accept === "standard"`): Supports `text`, `thinking`, `tool_call`, `image`, `audio`, `video`
- **CUI Client** (`ctx.accept === "cui-web"` etc.): Supports all types including `loading`, `error`, `action`, `event`
### 5. Performance Considerations
- Messages auto-flush after each `Send()` for real-time delivery
- Batch related messages with `SendGroup()` when possible for better performance
- Avoid sending too many small updates (combine them when feasible)
- Use `SendGroupStart`/`SendGroupEnd` for fine-grained control over grouping
### 6. Context Information Available
The `ctx` object provides access to:
```javascript
ctx.chat_id; // Chat session ID
ctx.assistant_id; // Assistant ID
ctx.locale; // User locale (e.g., "en", "zh-cn")
ctx.authorized; // User authorization info
ctx.metadata; // Custom metadata
ctx.client; // Client information (type, user_agent, ip)
```
## Migration Guide
### From Old Output API
**Before (Deprecated):**
```javascript
function Create(ctx, messages) {
const output = new Output(ctx);
output.Send("Hello");
output.SendGroup({ id: "grp1", messages: [...] });
}
```
**After (Current):**
```javascript
function Create(ctx, messages) {
ctx.Send("Hello"); // Auto-flushes
ctx.SendGroup({ messages: [...] }); // Auto-handles events and flushing
return { messages };
}
```
## Best Practices
1. **Use String Shorthand**: `ctx.Send("Hello")` is simpler than `ctx.Send({ type: "text", props: { content: "Hello" } })`
2. **Messages Auto-Flush**: Each `Send()` automatically flushes for real-time delivery - no manual flushing needed
3. **Choose the Right API Level**:
- **High-level**: Use `SendGroup()` for simple grouped messages (auto-handles events)
- **Low-level**: Use `SendGroupStart()`/`SendGroupEnd()` for fine-grained control
4. **Handle Errors Gracefully**: Always provide user-friendly error messages
5. **Show Progress for Long Operations**: Use loading indicators for better UX
6. **Return Hook Results**: Always return required objects from hooks:
- `Create`: `{ messages }`
- `Before/After`: `{ response }`
- `Done`: `{}` or `{ response }`
- `Error`: `{ error }` or `{ retry: true }`
7. **Test with Different Clients**: Verify behavior with both OpenAI and CUI clients
8. **Group Related Messages**: Use groups to organize related content for better frontend rendering

464
agent/context/buffer.go Normal file
View file

@ -0,0 +1,464 @@
package context
import (
"sync"
"time"
"github.com/google/uuid"
)
// =============================================================================
// Chat Buffer - Buffers messages and steps during execution for batch saving
// =============================================================================
// ChatBuffer buffers messages and resume steps during agent execution
// All data is held in memory and batch-written at the end of Stream()
type ChatBuffer struct {
// Identity
chatID string
requestID string
assistantID string
connector string // Current connector ID (for data analysis)
// Message buffer
messages []*BufferedMessage
msgSequence int
// Step buffer (for Resume)
steps []*BufferedStep
currentStep *BufferedStep
stepSequence int
// Space snapshot (captured when step starts, for recovery)
spaceSnapshot map[string]interface{}
mu sync.Mutex
}
// BufferedMessage represents a message waiting to be saved
type BufferedMessage struct {
MessageID string `json:"message_id"`
ChatID string `json:"chat_id"`
RequestID string `json:"request_id,omitempty"`
Role string `json:"role"` // "user" or "assistant"
Type string `json:"type"` // "text", "image", "loading", "tool_call", "retrieval", etc.
Props map[string]interface{} `json:"props"`
BlockID string `json:"block_id,omitempty"`
ThreadID string `json:"thread_id,omitempty"`
AssistantID string `json:"assistant_id,omitempty"`
Connector string `json:"connector,omitempty"` // Connector ID used for this message
Sequence int `json:"sequence"`
Metadata map[string]interface{} `json:"metadata,omitempty"`
CreatedAt time.Time `json:"created_at"`
IsStreaming bool `json:"-"` // Internal flag: true if message is still streaming (not saved until End)
}
// BufferedStep represents an execution step waiting to be saved (for Resume)
// Only saved when request is interrupted or failed
type BufferedStep struct {
ResumeID string `json:"resume_id"`
ChatID string `json:"chat_id"`
RequestID string `json:"request_id"`
AssistantID string `json:"assistant_id"`
StackID string `json:"stack_id"`
StackParentID string `json:"stack_parent_id,omitempty"`
StackDepth int `json:"stack_depth"`
Type string `json:"type"` // "input", "hook_create", "llm", "tool", "hook_next", "delegate"
Status string `json:"status"` // "running", "completed", "failed", "interrupted"
Input map[string]interface{} `json:"input,omitempty"`
Output map[string]interface{} `json:"output,omitempty"`
SpaceSnapshot map[string]interface{} `json:"space_snapshot,omitempty"`
Error string `json:"error,omitempty"`
Sequence int `json:"sequence"`
Metadata map[string]interface{} `json:"metadata,omitempty"`
CreatedAt time.Time `json:"created_at"`
}
// Step status constants (internal use only, not stored in database)
const (
StepStatusRunning = "running"
StepStatusCompleted = "completed"
)
// Step type constants
const (
StepTypeInput = "input"
StepTypeHookCreate = "hook_create"
StepTypeLLM = "llm"
StepTypeTool = "tool"
StepTypeHookNext = "hook_next"
StepTypeDelegate = "delegate"
)
// Resume status constants (for database storage)
const (
ResumeStatusFailed = "failed"
ResumeStatusInterrupted = "interrupted"
)
// NewChatBuffer creates a new chat buffer
func NewChatBuffer(chatID, requestID, assistantID, connector string) *ChatBuffer {
return &ChatBuffer{
chatID: chatID,
requestID: requestID,
assistantID: assistantID,
connector: connector,
messages: make([]*BufferedMessage, 0),
steps: make([]*BufferedStep, 0),
}
}
// =============================================================================
// Message Buffer Methods
// =============================================================================
// AddMessage adds a message to the buffer
func (b *ChatBuffer) AddMessage(msg *BufferedMessage) {
if msg == nil {
return
}
b.mu.Lock()
defer b.mu.Unlock()
// Auto-generate IDs if not provided
if msg.MessageID == "" {
msg.MessageID = uuid.New().String()
}
if msg.ChatID == "" {
msg.ChatID = b.chatID
}
if msg.RequestID == "" {
msg.RequestID = b.requestID
}
if msg.CreatedAt.IsZero() {
msg.CreatedAt = time.Now()
}
// Auto-increment sequence
b.msgSequence++
msg.Sequence = b.msgSequence
b.messages = append(b.messages, msg)
}
// AddUserInput adds user input message to the buffer
func (b *ChatBuffer) AddUserInput(content interface{}, name string) {
props := map[string]interface{}{
"content": content,
"role": "user",
}
if name != "" {
props["name"] = name
}
b.AddMessage(&BufferedMessage{
Role: "user",
Type: "user_input",
Props: props,
})
}
// AddAssistantMessage adds an assistant message to the buffer
// This is called by ctx.Send() to buffer messages for batch saving
func (b *ChatBuffer) AddAssistantMessage(messageID, msgType string, props map[string]interface{}, blockID, threadID, assistantID string, metadata map[string]interface{}) {
// Skip event type messages (transient, not stored)
if msgType == "event" {
return
}
b.AddMessage(&BufferedMessage{
MessageID: messageID, // Use the same MessageID as sent to client
Role: "assistant",
Type: msgType,
Props: props,
BlockID: blockID,
ThreadID: threadID,
AssistantID: assistantID,
Connector: b.connector, // Use current connector
Metadata: metadata,
})
}
// AddStreamingMessage adds a streaming message to the buffer
// Streaming messages are not saved until CompleteStreamingMessage is called
// This is called by ctx.SendStream() to start a streaming message
func (b *ChatBuffer) AddStreamingMessage(messageID, msgType string, props map[string]interface{}, blockID, threadID, assistantID string, metadata map[string]interface{}) {
// Skip event type messages (transient, not stored)
if msgType == "event" {
return
}
// Deep copy props to avoid mutation issues
propsCopy := make(map[string]interface{})
for k, v := range props {
propsCopy[k] = v
}
b.AddMessage(&BufferedMessage{
MessageID: messageID, // Use provided message ID
Role: "assistant",
Type: msgType,
Props: propsCopy,
BlockID: blockID,
ThreadID: threadID,
AssistantID: assistantID,
Connector: b.connector,
Metadata: metadata,
IsStreaming: true, // Mark as streaming
})
}
// AppendMessageContent appends content to a streaming message
// This is called by ctx.Append() to accumulate content
func (b *ChatBuffer) AppendMessageContent(messageID string, content string) bool {
b.mu.Lock()
defer b.mu.Unlock()
// Find the message by ID
for _, msg := range b.messages {
if msg.MessageID == messageID && msg.IsStreaming {
// Append to existing content
if msg.Props == nil {
msg.Props = make(map[string]interface{})
}
if existing, ok := msg.Props["content"].(string); ok {
msg.Props["content"] = existing + content
} else {
msg.Props["content"] = content
}
return true
}
}
return false
}
// CompleteStreamingMessage marks a streaming message as complete
// This is called by ctx.End() to finalize the message
// Returns the complete content for the message_end event
func (b *ChatBuffer) CompleteStreamingMessage(messageID string) (string, bool) {
b.mu.Lock()
defer b.mu.Unlock()
// Find the message by ID
for _, msg := range b.messages {
if msg.MessageID == messageID && msg.IsStreaming {
msg.IsStreaming = false
// Return the accumulated content
if content, ok := msg.Props["content"].(string); ok {
return content, true
}
return "", true
}
}
return "", false
}
// GetStreamingMessage returns a streaming message by ID
func (b *ChatBuffer) GetStreamingMessage(messageID string) *BufferedMessage {
b.mu.Lock()
defer b.mu.Unlock()
for _, msg := range b.messages {
if msg.MessageID == messageID && msg.IsStreaming {
return msg
}
}
return nil
}
// GetMessages returns all buffered messages
func (b *ChatBuffer) GetMessages() []*BufferedMessage {
b.mu.Lock()
defer b.mu.Unlock()
result := make([]*BufferedMessage, len(b.messages))
copy(result, b.messages)
return result
}
// GetMessageCount returns the number of buffered messages
func (b *ChatBuffer) GetMessageCount() int {
b.mu.Lock()
defer b.mu.Unlock()
return len(b.messages)
}
// =============================================================================
// Step Buffer Methods (for Resume)
// =============================================================================
// BeginStep starts tracking a new execution step
// Returns the step for further updates
func (b *ChatBuffer) BeginStep(stepType string, input map[string]interface{}, stack *Stack) *BufferedStep {
b.mu.Lock()
defer b.mu.Unlock()
b.stepSequence++
step := &BufferedStep{
ResumeID: uuid.New().String(),
ChatID: b.chatID,
RequestID: b.requestID,
AssistantID: b.assistantID,
Type: stepType,
Status: StepStatusRunning,
Input: input,
Sequence: b.stepSequence,
CreatedAt: time.Now(),
}
// Set stack information if available
if stack != nil {
step.StackID = stack.ID
step.StackParentID = stack.ParentID
step.StackDepth = stack.Depth
}
// Capture current space snapshot
if b.spaceSnapshot != nil {
step.SpaceSnapshot = copyMap(b.spaceSnapshot)
}
b.steps = append(b.steps, step)
b.currentStep = step
return step
}
// CompleteStep marks the current step as completed
func (b *ChatBuffer) CompleteStep(output map[string]interface{}) {
b.mu.Lock()
defer b.mu.Unlock()
if b.currentStep != nil {
b.currentStep.Output = output
b.currentStep.Status = StepStatusCompleted
b.currentStep = nil
}
}
// FailCurrentStep marks the current step as failed or interrupted
func (b *ChatBuffer) FailCurrentStep(status string, err error) {
b.mu.Lock()
defer b.mu.Unlock()
if b.currentStep != nil && b.currentStep.Status == StepStatusRunning {
b.currentStep.Status = status
if err != nil {
b.currentStep.Error = err.Error()
}
}
}
// GetCurrentStep returns the current running step
func (b *ChatBuffer) GetCurrentStep() *BufferedStep {
b.mu.Lock()
defer b.mu.Unlock()
return b.currentStep
}
// GetStepsForResume returns steps that need to be saved for resume
// Only returns steps with failed or interrupted status
func (b *ChatBuffer) GetStepsForResume(finalStatus string) []*BufferedStep {
b.mu.Lock()
defer b.mu.Unlock()
// If completed successfully, no steps need to be saved
if finalStatus == StepStatusCompleted {
return nil
}
// Mark current running step with final status
if b.currentStep != nil && b.currentStep.Status == StepStatusRunning {
b.currentStep.Status = finalStatus
}
// Return all steps (they will all have the context for recovery)
result := make([]*BufferedStep, len(b.steps))
copy(result, b.steps)
return result
}
// GetAllSteps returns all buffered steps (for debugging/testing)
func (b *ChatBuffer) GetAllSteps() []*BufferedStep {
b.mu.Lock()
defer b.mu.Unlock()
result := make([]*BufferedStep, len(b.steps))
copy(result, b.steps)
return result
}
// =============================================================================
// Space Snapshot Methods
// =============================================================================
// SetSpaceSnapshot sets the space snapshot for recovery
// Should be called when space data changes
func (b *ChatBuffer) SetSpaceSnapshot(snapshot map[string]interface{}) {
b.mu.Lock()
defer b.mu.Unlock()
b.spaceSnapshot = copyMap(snapshot)
}
// GetSpaceSnapshot returns the current space snapshot
func (b *ChatBuffer) GetSpaceSnapshot() map[string]interface{} {
b.mu.Lock()
defer b.mu.Unlock()
return copyMap(b.spaceSnapshot)
}
// =============================================================================
// Identity Methods
// =============================================================================
// ChatID returns the chat ID
func (b *ChatBuffer) ChatID() string {
return b.chatID
}
// RequestID returns the request ID
func (b *ChatBuffer) RequestID() string {
return b.requestID
}
// AssistantID returns the assistant ID
func (b *ChatBuffer) AssistantID() string {
return b.assistantID
}
// SetAssistantID updates the assistant ID (for A2A calls)
func (b *ChatBuffer) SetAssistantID(assistantID string) {
b.mu.Lock()
defer b.mu.Unlock()
b.assistantID = assistantID
}
// Connector returns the current connector ID
func (b *ChatBuffer) Connector() string {
return b.connector
}
// SetConnector updates the connector ID (when user switches connector)
func (b *ChatBuffer) SetConnector(connector string) {
b.mu.Lock()
defer b.mu.Unlock()
b.connector = connector
}
// =============================================================================
// Helper Functions
// =============================================================================
// copyMap creates a shallow copy of a map
func copyMap(src map[string]interface{}) map[string]interface{} {
if src == nil {
return nil
}
dst := make(map[string]interface{}, len(src))
for k, v := range src {
dst[k] = v
}
return dst
}

1417
agent/context/buffer_test.go Normal file
View file

File diff suppressed because it is too large Load diff

90
agent/context/context.go
View file

@ -387,3 +387,93 @@ func (ctx *Context) getMessageMetadata(messageID string) *MessageMetadata {
}
return ctx.messageMetadata.getMessage(messageID)
}
// GetMessageMetadata returns metadata for a message (public version)
func (ctx *Context) GetMessageMetadata(messageID string) *MessageMetadata {
return ctx.getMessageMetadata(messageID)
}
// =============================================================================
// Chat Buffer Methods
// =============================================================================
// InitBuffer initializes the chat buffer for this context
// Should be called at the start of Stream() to begin buffering messages and steps
func (ctx *Context) InitBuffer(assistantID, connector string) *ChatBuffer {
ctx.Buffer = NewChatBuffer(ctx.ChatID, ctx.RequestID(), assistantID, connector)
return ctx.Buffer
}
// HasBuffer returns true if the buffer is initialized
func (ctx *Context) HasBuffer() bool {
return ctx.Buffer != nil
}
// BufferUserInput adds user input to the buffer
// Should be called at the start of Stream() to buffer the user's input message
func (ctx *Context) BufferUserInput(messages []Message) {
if ctx.Buffer == nil {
return
}
for _, msg := range messages {
if msg.Role == RoleUser {
// Get name if available
var name string
if msg.Name != nil {
name = *msg.Name
}
ctx.Buffer.AddUserInput(msg.Content, name)
}
}
}
// BufferAssistantMessage adds an assistant message to the buffer
// Called by ctx.Send() to buffer messages for batch saving
func (ctx *Context) BufferAssistantMessage(messageID, msgType string, props map[string]interface{}, blockID, threadID string, metadata map[string]interface{}) {
if ctx.Buffer == nil {
return
}
ctx.Buffer.AddAssistantMessage(messageID, msgType, props, blockID, threadID, ctx.AssistantID, metadata)
}
// BeginStep starts tracking a new execution step
// Returns the step for further updates
func (ctx *Context) BeginStep(stepType string, input map[string]interface{}) *BufferedStep {
if ctx.Buffer == nil {
return nil
}
// Update space snapshot before starting step
if ctx.Space != nil {
ctx.Buffer.SetSpaceSnapshot(ctx.Space.Snapshot())
}
return ctx.Buffer.BeginStep(stepType, input, ctx.Stack)
}
// CompleteStep marks the current step as completed
func (ctx *Context) CompleteStep(output map[string]interface{}) {
if ctx.Buffer == nil {
return
}
ctx.Buffer.CompleteStep(output)
}
// FailCurrentStep marks the current step as failed or interrupted
func (ctx *Context) FailCurrentStep(status string, err error) {
if ctx.Buffer == nil {
return
}
ctx.Buffer.FailCurrentStep(status, err)
}
// shouldSkipHistory checks if history saving should be skipped
// Returns true if Skip.History is set in the current stack options
func (ctx *Context) shouldSkipHistory() bool {
if ctx.Stack == nil || ctx.Stack.Options == nil || ctx.Stack.Options.Skip == nil {
return false
}
return ctx.Stack.Options.Skip.History
}

99
agent/context/jsapi.go
View file

@ -44,10 +44,12 @@ func (ctx *Context) NewObject(v8ctx *v8go.Context) (*v8go.Value, error) {
// Set methods
jsObject.Set("Send", ctx.sendMethod(v8ctx.Isolate()))
jsObject.Set("SendStream", ctx.sendStreamMethod(v8ctx.Isolate()))
jsObject.Set("Replace", ctx.replaceMethod(v8ctx.Isolate()))
jsObject.Set("Append", ctx.appendMethod(v8ctx.Isolate()))
jsObject.Set("Merge", ctx.mergeMethod(v8ctx.Isolate()))
jsObject.Set("Set", ctx.setMethod(v8ctx.Isolate()))
jsObject.Set("End", ctx.endMethod(v8ctx.Isolate()))
// Set ID generator methods
jsObject.Set("MessageID", ctx.messageIDMethod(v8ctx.Isolate()))
@ -266,6 +268,103 @@ func (ctx *Context) sendMethod(iso *v8go.Isolate) *v8go.FunctionTemplate {
})
}
// sendStreamMethod implements ctx.SendStream(message)
// Usage: const msgId = ctx.SendStream({ type: "text", props: { content: "Initial content" } })
// Starts a streaming message that can be appended to with ctx.Append()
// Must be finalized with ctx.End(msgId) or ctx.End(msgId, "final content")
// Unlike Send(), this does NOT automatically send message_end event
// Returns: message_id (string)
func (ctx *Context) sendStreamMethod(iso *v8go.Isolate) *v8go.FunctionTemplate {
return v8go.NewFunctionTemplate(iso, func(info *v8go.FunctionCallbackInfo) *v8go.Value {
v8ctx := info.Context()
args := info.Args()
if len(args) < 1 {
return bridge.JsException(v8ctx, "SendStream requires a message argument")
}
// Parse message argument
msg, err := parseMessage(v8ctx, args[0])
if err != nil {
return bridge.JsException(v8ctx, "invalid message: "+err.Error())
}
// Get optional blockId argument (second argument)
if len(args) >= 2 && args[1].IsString() && msg.BlockID == "" {
msg.BlockID = args[1].String()
}
// Call ctx.SendStream
messageID, err := ctx.SendStream(msg)
if err != nil {
return bridge.JsException(v8ctx, "SendStream failed: "+err.Error())
}
// Automatically flush after sending
if err := ctx.Flush(); err != nil {
return bridge.JsException(v8ctx, "Flush failed: "+err.Error())
}
// Return the message ID
returnID, err := v8go.NewValue(iso, messageID)
if err != nil {
return bridge.JsException(v8ctx, "Failed to create return value: "+err.Error())
}
return returnID
})
}
// endMethod implements ctx.End(messageId, finalContent?)
// Usage: ctx.End(msgId) or ctx.End(msgId, "final content to append")
// Finalizes a streaming message started with SendStream()
// Sends message_end event with the complete accumulated content
// Returns: message_id (string)
func (ctx *Context) endMethod(iso *v8go.Isolate) *v8go.FunctionTemplate {
return v8go.NewFunctionTemplate(iso, func(info *v8go.FunctionCallbackInfo) *v8go.Value {
v8ctx := info.Context()
args := info.Args()
if len(args) < 1 {
return bridge.JsException(v8ctx, "End requires a messageId argument")
}
// Get message ID (first argument)
if !args[0].IsString() {
return bridge.JsException(v8ctx, "messageId must be a string")
}
messageID := args[0].String()
// Get optional final content (second argument)
var finalContent string
if len(args) >= 2 && args[1].IsString() {
finalContent = args[1].String()
}
// Call ctx.End
var err error
if finalContent != "" {
err = ctx.End(messageID, finalContent)
} else {
err = ctx.End(messageID)
}
if err != nil {
return bridge.JsException(v8ctx, "End failed: "+err.Error())
}
// Automatically flush after sending
if err := ctx.Flush(); err != nil {
return bridge.JsException(v8ctx, "Flush failed: "+err.Error())
}
// Return the message ID
returnID, err := v8go.NewValue(iso, messageID)
if err != nil {
return bridge.JsException(v8ctx, "Failed to create return value: "+err.Error())
}
return returnID
})
}
// replaceMethod implements ctx.Replace(messageId, message)
// Usage: ctx.Replace(messageId, { type: "text", props: { content: "Updated content" } })
// Replaces the entire message content with the specified message_id

515
agent/context/jsapi_output_test.go
View file

@ -826,3 +826,518 @@ func TestJsValueEndBlock(t *testing.T) {
output := mockWriter.buffer.String()
assert.Contains(t, output, "block_end", "Output should contain block_end event")
}
// TestJsValueSendStream tests the SendStream method on Context
func TestJsValueSendStream(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
// Setup mock writer
mockWriter := newMockResponseWriter()
// Use New() to properly initialize messageMetadata
cxt := New(context.Background(), nil, "test-chat-id")
cxt.AssistantID = "test-assistant-id"
cxt.Accept = AcceptWebCUI
cxt.Locale = "en"
cxt.Writer = mockWriter
// Test SendStream method
res, err := v8.Call(v8.CallOptions{}, `
function test(ctx) {
try {
// Start a streaming message
const msgId = ctx.SendStream({
type: "text",
props: { content: "Initial content" }
});
// Verify msgId is returned
if (typeof msgId !== 'string' || msgId === '') {
throw new Error('SendStream should return a message ID');
}
return { success: true, msgId: msgId };
} catch (error) {
return { success: false, error: error.message };
}
}`, cxt)
if err != nil {
t.Fatalf("Call failed: %v", err)
}
result, ok := res.(map[string]interface{})
if !ok {
t.Fatalf("Expected map result, got %T", res)
}
if !result["success"].(bool) {
t.Logf("Error: %v", result["error"])
}
assert.Equal(t, true, result["success"], "SendStream should work correctly")
// Verify message_start was sent but NOT message_end
output := mockWriter.buffer.String()
assert.Contains(t, output, "message_start", "Output should contain message_start event")
assert.NotContains(t, output, "message_end", "Output should NOT contain message_end event (streaming)")
}
// TestJsValueSendStreamWithBlockID tests SendStream with block_id parameter
func TestJsValueSendStreamWithBlockID(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
mockWriter := newMockResponseWriter()
cxt := New(context.Background(), nil, "test-chat-id")
cxt.AssistantID = "test-assistant-id"
cxt.Accept = AcceptWebCUI
cxt.Locale = "en"
cxt.Writer = mockWriter
res, err := v8.Call(v8.CallOptions{}, `
function test(ctx) {
try {
// Generate block ID
const blockId = ctx.BlockID();
// Start streaming with block_id
const msgId = ctx.SendStream({
type: "text",
props: { content: "Streaming with block" },
block_id: blockId
});
return { success: true, msgId: msgId, blockId: blockId };
} catch (error) {
return { success: false, error: error.message };
}
}`, cxt)
if err != nil {
t.Fatalf("Call failed: %v", err)
}
result, ok := res.(map[string]interface{})
if !ok {
t.Fatalf("Expected map result, got %T", res)
}
assert.Equal(t, true, result["success"], "SendStream with blockId should succeed")
// Verify block_start was also sent
output := mockWriter.buffer.String()
assert.Contains(t, output, "block_start", "Output should contain block_start event")
}
// TestJsValueEnd tests the End method on Context
func TestJsValueEnd(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
mockWriter := newMockResponseWriter()
cxt := New(context.Background(), nil, "test-chat-id")
cxt.AssistantID = "test-assistant-id"
cxt.Accept = AcceptWebCUI
cxt.Locale = "en"
cxt.Writer = mockWriter
res, err := v8.Call(v8.CallOptions{}, `
function test(ctx) {
try {
// Start a streaming message
const msgId = ctx.SendStream({
type: "text",
props: { content: "Hello" }
});
// End the message
ctx.End(msgId);
return { success: true, msgId: msgId };
} catch (error) {
return { success: false, error: error.message };
}
}`, cxt)
if err != nil {
t.Fatalf("Call failed: %v", err)
}
result, ok := res.(map[string]interface{})
if !ok {
t.Fatalf("Expected map result, got %T", res)
}
if !result["success"].(bool) {
t.Logf("Error: %v", result["error"])
}
assert.Equal(t, true, result["success"], "End should work correctly")
// Verify message_end was sent
output := mockWriter.buffer.String()
assert.Contains(t, output, "message_end", "Output should contain message_end event after End()")
}
// TestJsValueEndWithFinalContent tests End with final content parameter
func TestJsValueEndWithFinalContent(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
mockWriter := newMockResponseWriter()
cxt := New(context.Background(), nil, "test-chat-id")
cxt.AssistantID = "test-assistant-id"
cxt.Accept = AcceptWebCUI
cxt.Locale = "en"
cxt.Writer = mockWriter
res, err := v8.Call(v8.CallOptions{}, `
function test(ctx) {
try {
// Start a streaming message
const msgId = ctx.SendStream({
type: "text",
props: { content: "Start" }
});
// End with final content
ctx.End(msgId, " End");
return { success: true, msgId: msgId };
} catch (error) {
return { success: false, error: error.message };
}
}`, cxt)
if err != nil {
t.Fatalf("Call failed: %v", err)
}
result, ok := res.(map[string]interface{})
if !ok {
t.Fatalf("Expected map result, got %T", res)
}
if !result["success"].(bool) {
t.Logf("Error: %v", result["error"])
}
assert.Equal(t, true, result["success"], "End with final content should work correctly")
// Verify message_end was sent
output := mockWriter.buffer.String()
assert.Contains(t, output, "message_end", "Output should contain message_end event")
}
// TestJsValueStreamingWorkflow tests the complete streaming workflow: SendStream -> Append -> End
func TestJsValueStreamingWorkflow(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
mockWriter := newMockResponseWriter()
cxt := New(context.Background(), nil, "test-chat-id")
cxt.AssistantID = "test-assistant-id"
cxt.Accept = AcceptWebCUI
cxt.Locale = "en"
cxt.Writer = mockWriter
res, err := v8.Call(v8.CallOptions{}, `
function test(ctx) {
try {
// Start a streaming message
const msgId = ctx.SendStream({
type: "text",
props: { content: "# Title\n\n" }
});
// Append content in chunks (simulating streaming)
ctx.Append(msgId, "First paragraph. ");
ctx.Append(msgId, "Second sentence. ");
ctx.Append(msgId, "Third sentence.\n\n");
ctx.Append(msgId, "Second paragraph.");
// Finalize the message
ctx.End(msgId);
return { success: true, msgId: msgId };
} catch (error) {
return { success: false, error: error.message };
}
}`, cxt)
if err != nil {
t.Fatalf("Call failed: %v", err)
}
result, ok := res.(map[string]interface{})
if !ok {
t.Fatalf("Expected map result, got %T", res)
}
if !result["success"].(bool) {
t.Logf("Error: %v", result["error"])
}
assert.Equal(t, true, result["success"], "Streaming workflow should work correctly")
// Verify the complete workflow events
output := mockWriter.buffer.String()
assert.Contains(t, output, "message_start", "Output should contain message_start")
assert.Contains(t, output, "message_end", "Output should contain message_end")
assert.Contains(t, output, "# Title", "Output should contain initial content")
assert.Contains(t, output, "First paragraph", "Output should contain appended content")
}
// TestJsValueSendStreamStringShorthand tests SendStream with string shorthand
func TestJsValueSendStreamStringShorthand(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
mockWriter := newMockResponseWriter()
cxt := New(context.Background(), nil, "test-chat-id")
cxt.AssistantID = "test-assistant-id"
cxt.Accept = AcceptWebCUI
cxt.Locale = "en"
cxt.Writer = mockWriter
res, err := v8.Call(v8.CallOptions{}, `
function test(ctx) {
try {
// SendStream with string shorthand
const msgId = ctx.SendStream("Hello streaming");
if (typeof msgId !== 'string' || msgId === '') {
throw new Error('SendStream should return a message ID');
}
ctx.End(msgId);
return { success: true, msgId: msgId };
} catch (error) {
return { success: false, error: error.message };
}
}`, cxt)
if err != nil {
t.Fatalf("Call failed: %v", err)
}
result, ok := res.(map[string]interface{})
if !ok {
t.Fatalf("Expected map result, got %T", res)
}
assert.Equal(t, true, result["success"], "SendStream with string shorthand should succeed")
}
// TestJsValueEndErrorHandling tests error handling in End method
func TestJsValueEndErrorHandling(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
mockWriter := newMockResponseWriter()
cxt := New(context.Background(), nil, "test-chat-id")
cxt.AssistantID = "test-assistant-id"
cxt.Accept = AcceptWebCUI
cxt.Locale = "en"
cxt.Writer = mockWriter
// Test End without arguments
res, err := v8.Call(v8.CallOptions{}, `
function test(ctx) {
try {
ctx.End();
return { success: true };
} catch (error) {
return { success: false, error: error.message };
}
}`, cxt)
if err != nil {
t.Fatalf("Call failed: %v", err)
}
result, ok := res.(map[string]interface{})
if !ok {
t.Fatalf("Expected map result, got %T", res)
}
assert.Equal(t, false, result["success"], "End without arguments should fail")
assert.Contains(t, result["error"], "messageId", "Error should mention missing messageId")
}
// TestJsValueEndWithInvalidMessageID tests End with invalid messageId type
func TestJsValueEndWithInvalidMessageID(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
mockWriter := newMockResponseWriter()
cxt := New(context.Background(), nil, "test-chat-id")
cxt.AssistantID = "test-assistant-id"
cxt.Accept = AcceptWebCUI
cxt.Locale = "en"
cxt.Writer = mockWriter
// Test End with non-string messageId
res, err := v8.Call(v8.CallOptions{}, `
function test(ctx) {
try {
ctx.End(123);
return { success: true };
} catch (error) {
return { success: false, error: error.message };
}
}`, cxt)
if err != nil {
t.Fatalf("Call failed: %v", err)
}
result, ok := res.(map[string]interface{})
if !ok {
t.Fatalf("Expected map result, got %T", res)
}
assert.Equal(t, false, result["success"], "End with non-string messageId should fail")
assert.Contains(t, result["error"], "string", "Error should mention messageId must be string")
}
// TestJsValueSendStreamErrorHandling tests error handling in SendStream method
func TestJsValueSendStreamErrorHandling(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
mockWriter := newMockResponseWriter()
cxt := New(context.Background(), nil, "test-chat-id")
cxt.AssistantID = "test-assistant-id"
cxt.Accept = AcceptWebCUI
cxt.Locale = "en"
cxt.Writer = mockWriter
// Test SendStream without arguments
res, err := v8.Call(v8.CallOptions{}, `
function test(ctx) {
try {
ctx.SendStream();
return { success: true };
} catch (error) {
return { success: false, error: error.message };
}
}`, cxt)
if err != nil {
t.Fatalf("Call failed: %v", err)
}
result, ok := res.(map[string]interface{})
if !ok {
t.Fatalf("Expected map result, got %T", res)
}
assert.Equal(t, false, result["success"], "SendStream without arguments should fail")
assert.Contains(t, result["error"], "SendStream requires a message argument", "Error should mention missing message")
}
// TestJsValueMultipleStreams tests handling multiple concurrent streaming messages
func TestJsValueMultipleStreams(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
mockWriter := newMockResponseWriter()
cxt := New(context.Background(), nil, "test-chat-id")
cxt.AssistantID = "test-assistant-id"
cxt.Accept = AcceptWebCUI
cxt.Locale = "en"
cxt.Writer = mockWriter
res, err := v8.Call(v8.CallOptions{}, `
function test(ctx) {
try {
// Start multiple streaming messages
const msg1 = ctx.SendStream({ type: "text", props: { content: "Stream 1: " } });
const msg2 = ctx.SendStream({ type: "text", props: { content: "Stream 2: " } });
// Interleave appends
ctx.Append(msg1, "A");
ctx.Append(msg2, "X");
ctx.Append(msg1, "B");
ctx.Append(msg2, "Y");
ctx.Append(msg1, "C");
ctx.Append(msg2, "Z");
// End both streams
ctx.End(msg1);
ctx.End(msg2);
return { success: true, msg1: msg1, msg2: msg2 };
} catch (error) {
return { success: false, error: error.message };
}
}`, cxt)
if err != nil {
t.Fatalf("Call failed: %v", err)
}
result, ok := res.(map[string]interface{})
if !ok {
t.Fatalf("Expected map result, got %T", res)
}
if !result["success"].(bool) {
t.Logf("Error: %v", result["error"])
}
assert.Equal(t, true, result["success"], "Multiple streams should work correctly")
assert.NotEqual(t, result["msg1"], result["msg2"], "Message IDs should be different")
}
// TestJsValueSendVsSendStream tests the difference between Send and SendStream
func TestJsValueSendVsSendStream(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
// Test Send - should auto-send message_end
t.Run("Send auto-ends", func(t *testing.T) {
mockWriter := newMockResponseWriter()
cxt := New(context.Background(), nil, "test-chat-id")
cxt.AssistantID = "test-assistant-id"
cxt.Accept = AcceptWebCUI
cxt.Locale = "en"
cxt.Writer = mockWriter
_, err := v8.Call(v8.CallOptions{}, `
function test(ctx) {
ctx.Send("Complete message");
return true;
}`, cxt)
if err != nil {
t.Fatalf("Call failed: %v", err)
}
output := mockWriter.buffer.String()
assert.Contains(t, output, "message_start", "Send should emit message_start")
assert.Contains(t, output, "message_end", "Send should auto-emit message_end")
})
// Test SendStream - should NOT auto-send message_end
t.Run("SendStream requires explicit End", func(t *testing.T) {
mockWriter := newMockResponseWriter()
cxt := New(context.Background(), nil, "test-chat-id")
cxt.AssistantID = "test-assistant-id"
cxt.Accept = AcceptWebCUI
cxt.Locale = "en"
cxt.Writer = mockWriter
_, err := v8.Call(v8.CallOptions{}, `
function test(ctx) {
const msgId = ctx.SendStream("Streaming message");
// Intentionally NOT calling ctx.End(msgId)
return msgId;
}`, cxt)
if err != nil {
t.Fatalf("Call failed: %v", err)
}
output := mockWriter.buffer.String()
assert.Contains(t, output, "message_start", "SendStream should emit message_start")
assert.NotContains(t, output, "message_end", "SendStream should NOT auto-emit message_end")
})
}

210
agent/context/output.go
View file

@ -55,6 +55,13 @@ func (ctx *Context) Send(msg *message.Message) error {
// Increment chunk count for this message
metadata.ChunkCount++
// Update Buffer content for streaming messages (for storage)
if ctx.Buffer != nil && msg.Props != nil {
if content, ok := msg.Props["content"].(string); ok {
ctx.Buffer.AppendMessageContent(msg.MessageID, content)
}
}
}
}
@ -146,6 +153,26 @@ func (ctx *Context) Send(msg *message.Message) error {
return err
}
// === Buffer message for batch saving (non-delta, non-event messages only) ===
// Delta messages are streaming chunks; only final content should be saved
// Event messages are transient lifecycle signals, not stored
// Skip if History is disabled in options
if !msg.Delta && !isEventMessage && ctx.Buffer != nil && !ctx.shouldSkipHistory() {
assistantID := ""
if ctx.Stack != nil {
assistantID = ctx.Stack.AssistantID
}
ctx.Buffer.AddAssistantMessage(
msg.MessageID, // Use the same MessageID as sent to client
msg.Type,
msg.Props,
msg.BlockID,
msg.ThreadID,
assistantID,
nil, // metadata can be added if needed
)
}
// === Auto-send message_end for non-delta messages (complete messages) ===
if !msg.Delta && !isEventMessage && msg.MessageID != "" && ctx.messageMetadata != nil {
metadata := ctx.messageMetadata.getMessage(msg.MessageID)
@ -188,6 +215,189 @@ func (ctx *Context) Send(msg *message.Message) error {
return nil
}
// SendStream sends a streaming message that can be appended to later
// Unlike Send(), this does NOT automatically send message_end event
// Use ctx.Append() to add content, then ctx.End() to finalize
// Returns the message ID for use with Append/End
func (ctx *Context) SendStream(msg *message.Message) (string, error) {
out, err := ctx.getOutput()
if err != nil {
return "", err
}
// Skip lifecycle events for event-type messages
isEventMessage := msg.Type == message.TypeEvent
if isEventMessage {
// Event messages should use Send(), not SendStream()
return "", ctx.Send(msg)
}
// === Auto-generate ChunkID ===
if msg.ChunkID == "" {
if ctx.IDGenerator != nil {
msg.ChunkID = ctx.IDGenerator.GenerateChunkID()
} else {
msg.ChunkID = message.GenerateNanoID()
}
}
// === Auto-set ThreadID for non-root Stack ===
if msg.ThreadID == "" && ctx.Stack != nil && !ctx.Stack.IsRoot() {
msg.ThreadID = ctx.Stack.ID
}
// === Handle BlockID and block_start event ===
if msg.BlockID != "" && ctx.messageMetadata != nil {
if ctx.messageMetadata.getBlock(msg.BlockID) == nil {
blockStartData := message.EventBlockStartData{
BlockID: msg.BlockID,
Type: "mixed",
Timestamp: time.Now().UnixMilli(),
}
blockStartEvent := output.NewEventMessage(message.EventBlockStart, "Block started", blockStartData)
if err := ctx.sendRaw(blockStartEvent); err != nil {
return "", err
}
ctx.messageMetadata.setBlock(msg.BlockID, &BlockMetadata{
BlockID: msg.BlockID,
Type: "mixed",
StartTime: time.Now(),
MessageCount: 0,
})
}
ctx.messageMetadata.updateBlock(msg.BlockID, func(block *BlockMetadata) {
block.MessageCount++
})
}
// === Generate MessageID if not provided ===
if msg.MessageID == "" {
if ctx.IDGenerator != nil {
msg.MessageID = ctx.IDGenerator.GenerateMessageID()
} else {
msg.MessageID = message.GenerateNanoID()
}
}
// === Send message_start event ===
messageStartData := message.EventMessageStartData{
MessageID: msg.MessageID,
Type: msg.Type,
Timestamp: time.Now().UnixMilli(),
ThreadID: msg.ThreadID,
}
messageStartEvent := output.NewEventMessage(message.EventMessageStart, "Message started", messageStartData)
if err := ctx.sendRaw(messageStartEvent); err != nil {
return "", err
}
// === Record message metadata ===
if ctx.messageMetadata != nil {
ctx.messageMetadata.setMessage(msg.MessageID, &MessageMetadata{
MessageID: msg.MessageID,
BlockID: msg.BlockID,
ThreadID: msg.ThreadID,
Type: msg.Type,
StartTime: time.Now(),
ChunkCount: 1,
})
}
// === Actually send the message ===
if err := out.Send(msg); err != nil {
return "", err
}
// === Buffer streaming message (will be completed by End()) ===
if ctx.Buffer != nil && !ctx.shouldSkipHistory() {
assistantID := ""
if ctx.Stack != nil {
assistantID = ctx.Stack.AssistantID
}
ctx.Buffer.AddStreamingMessage(
msg.MessageID,
msg.Type,
msg.Props,
msg.BlockID,
msg.ThreadID,
assistantID,
nil,
)
}
// NOTE: No message_end event here - will be sent by End()
return msg.MessageID, nil
}
// End finalizes a streaming message started with SendStream
// Optionally appends final content before sending message_end event
// This also saves the complete message to the buffer for storage
func (ctx *Context) End(messageID string, finalContent ...string) error {
if messageID == "" {
return nil
}
// Append final content if provided
if len(finalContent) > 0 && finalContent[0] != "" {
// Create a delta message for the final content
deltaMsg := &message.Message{
MessageID: messageID,
Type: message.TypeText,
Delta: true,
DeltaAction: message.DeltaAppend,
Props: map[string]interface{}{
"content": finalContent[0],
},
}
if err := ctx.Send(deltaMsg); err != nil {
return err
}
}
// Get complete content from buffer
var completeContent string
if ctx.Buffer != nil {
completeContent, _ = ctx.Buffer.CompleteStreamingMessage(messageID)
}
// Get metadata for duration calculation
var durationMs int64
var threadID string
var chunkCount int
var msgType string = message.TypeText
if ctx.messageMetadata != nil {
if metadata := ctx.messageMetadata.getMessage(messageID); metadata != nil {
durationMs = time.Since(metadata.StartTime).Milliseconds()
threadID = metadata.ThreadID
chunkCount = metadata.ChunkCount
msgType = metadata.Type
}
}
// Build message_end event data
endData := message.EventMessageEndData{
MessageID: messageID,
Type: msgType,
Timestamp: time.Now().UnixMilli(),
ThreadID: threadID,
DurationMs: durationMs,
ChunkCount: chunkCount,
Status: "completed",
}
// Add complete content to extra
if completeContent != "" {
endData.Extra = map[string]interface{}{
"content": completeContent,
}
}
// Send message_end event
messageEndEvent := output.NewEventMessage(message.EventMessageEnd, "Message completed", endData)
return ctx.sendRaw(messageEndEvent)
}
// EndMessage sends a message_end event for a completed message
// Note: For non-delta messages, message_end is automatically sent by Send()
// This method is primarily for delta streaming scenarios:

3
agent/context/types.go
View file

@ -231,6 +231,9 @@ type Context struct {
Writer Writer `json:"-"` // Writer, it will be used to write response data to the client
IDGenerator *message.IDGenerator `json:"-"` // ID generator for this context (chunk, message, block, thread IDs)
// Chat buffer for batch saving messages and resume steps
Buffer *ChatBuffer `json:"-"` // Chat buffer for batch saving at end of Stream()
// Internal
trace traceTypes.Manager `json:"-"` // Trace manager, lazy initialized on first access
messageMetadata *messageMetadataStore `json:"-"` // Thread-safe message metadata store for delta operations

461
agent/store/CHAT_STORAGE_DESIGN.md
View file

@ -91,6 +91,7 @@ Stores chat metadata and session information.
| `chat_id` | string(64) | No | Unique | Unique chat identifier |
| `title` | string(500) | Yes | - | Chat title |
| `assistant_id` | string(200) | No | Yes | Associated assistant ID |
| `last_connector` | string(200) | Yes | Yes | Last used connector ID |
| `mode` | string(50) | No | - | Chat mode (default: "chat") |
| `status` | enum | No | Yes | Status: `active`, `archived` |
| `public` | boolean | No | - | Whether shared across all teams |
@ -129,6 +130,7 @@ These fields are automatically managed by the framework and used for access cont
| Name | Columns | Type |
| -------------------- | ----------------- | ----- |
| `idx_chat_assistant` | `assistant_id` | index |
| `idx_chat_last_conn` | `last_connector` | index |
| `idx_chat_status` | `status` | index |
| `idx_chat_share` | `share` | index |
| `idx_chat_last_msg` | `last_message_at` | index |
@ -139,33 +141,35 @@ Stores user-visible messages (both user input and assistant responses).
**Table Name:** `agent_message`
| Column | Type | Nullable | Index | Description |
| -------------- | ----------- | -------- | ------ | ----------------------------------------- |
| `id` | ID | No | PK | Auto-increment primary key |
| `message_id` | string(64) | No | Unique | Unique message identifier |
| `chat_id` | string(64) | No | Yes | Parent chat ID |
| `request_id` | string(64) | Yes | Yes | Request ID for grouping |
| `role` | enum | No | Yes | Role: `user`, `assistant` |
| `type` | string(50) | No | - | Message type (text, image, loading, etc.) |
| `props` | json | No | - | Message properties (content, url, etc.) |
| `block_id` | string(64) | Yes | Yes | Block grouping ID |
| `thread_id` | string(64) | Yes | Yes | Thread grouping ID |
| `assistant_id` | string(200) | Yes | Yes | Assistant ID (join to get name/avatar) |
| `sequence` | integer | No | - | Message order within chat (in composite) |
| `metadata` | json | Yes | - | Additional metadata |
| `created_at` | timestamp | No | Yes | Creation timestamp |
| `updated_at` | timestamp | No | - | Last update timestamp |
| Column | Type | Nullable | Index | Description |
| -------------- | ----------- | -------- | ----- | ------------------------------------------ |
| `id` | ID | No | PK | Auto-increment primary key |
| `message_id` | string(64) | No | - | Message identifier (unique within request) |
| `chat_id` | string(64) | No | Yes | Parent chat ID |
| `request_id` | string(64) | Yes | Yes | Request ID for grouping |
| `role` | enum | No | Yes | Role: `user`, `assistant` |
| `type` | string(50) | No | - | Message type (text, image, loading, etc.) |
| `props` | json | No | - | Message properties (content, url, etc.) |
| `block_id` | string(64) | Yes | Yes | Block grouping ID |
| `thread_id` | string(64) | Yes | Yes | Thread grouping ID |
| `assistant_id` | string(200) | Yes | Yes | Assistant ID (join to get name/avatar) |
| `connector` | string(200) | Yes | Yes | Connector ID used for this message |
| `sequence` | integer | No | - | Message order within chat (in composite) |
| `metadata` | json | Yes | - | Additional metadata |
| `created_at` | timestamp | No | Yes | Creation timestamp |
| `updated_at` | timestamp | No | - | Last update timestamp |
**Indexes:**
| Name | Columns | Type |
| ------------------- | --------------------- | ----- |
| `idx_msg_chat_seq` | `chat_id`, `sequence` | index |
| `idx_msg_request` | `request_id` | index |
| `idx_msg_role` | `role` | index |
| `idx_msg_block` | `block_id` | index |
| `idx_msg_thread` | `thread_id` | index |
| `idx_msg_assistant` | `assistant_id` | index |
| Name | Columns | Type |
| ------------------------- | -------------------------- | ------ |
| `idx_msg_chat_seq` | `chat_id`, `sequence` | index |
| `idx_msg_request_message` | `request_id`, `message_id` | unique |
| `idx_msg_request` | `request_id` | index |
| `idx_msg_role` | `role` | index |
| `idx_msg_block` | `block_id` | index |
| `idx_msg_thread` | `thread_id` | index |
| `idx_msg_assistant` | `assistant_id` | index |
**Message Types:**
@ -693,27 +697,101 @@ func createResumeRecord(ctx *Context, stepType, status string, input, output int
```go
// ChatStore defines the chat storage interface
// Provides operations for chat, message, and resume management
type ChatStore interface {
// ==========================================================================
// Chat Management
// ==========================================================================
// CreateChat creates a new chat session
CreateChat(chat *Chat) error
// GetChat retrieves a single chat by ID
GetChat(chatID string) (*Chat, error)
// UpdateChat updates chat fields
UpdateChat(chatID string, updates map[string]interface{}) error
// DeleteChat deletes a chat and its associated messages
DeleteChat(chatID string) error
// ListChats retrieves a paginated list of chats with optional grouping
ListChats(filter ChatFilter) (*ChatList, error)
// ==========================================================================
// Message Management
// ==========================================================================
// SaveMessages batch saves messages for a chat
// This is the primary write method - messages are buffered during execution
// and batch-written at the end of a request
SaveMessages(chatID string, messages []*Message) error
// GetMessages retrieves messages for a chat with filtering
GetMessages(chatID string, filter MessageFilter) ([]*Message, error)
// UpdateMessage updates a single message
UpdateMessage(messageID string, updates map[string]interface{}) error
// DeleteMessages deletes specific messages from a chat
DeleteMessages(chatID string, messageIDs []string) error
// ==========================================================================
// Resume Management (only called on failure/interrupt)
// ==========================================================================
// SaveResume batch saves resume records
// Only called when request is interrupted or failed
SaveResume(records []*Resume) error
// GetResume retrieves all resume records for a chat
GetResume(chatID string) ([]*Resume, error)
// GetLastResume retrieves the last (most recent) resume record for a chat
GetLastResume(chatID string) (*Resume, error)
// GetResumeByStackID retrieves resume records for a specific stack
GetResumeByStackID(stackID string) ([]*Resume, error)
GetStackPath(stackID string) ([]string, error) // Returns [root_stack_id, ..., current_stack_id]
DeleteResume(chatID string) error // Clean up after successful resume
// GetStackPath returns the stack path from root to the given stack
// Returns: [root_stack_id, ..., current_stack_id]
GetStackPath(stackID string) ([]string, error)
// DeleteResume deletes all resume records for a chat
// Called after successful resume to clean up
DeleteResume(chatID string) error
}
// AssistantStore defines the assistant storage interface
// Separated from ChatStore for clearer responsibility
type AssistantStore interface {
// SaveAssistant saves assistant information
SaveAssistant(assistant *AssistantModel) (string, error)
// UpdateAssistant updates assistant fields
UpdateAssistant(assistantID string, updates map[string]interface{}) error
// DeleteAssistant deletes an assistant
DeleteAssistant(assistantID string) error
// GetAssistants retrieves a paginated list of assistants with filtering
GetAssistants(filter AssistantFilter, locale ...string) (*AssistantList, error)
// GetAssistantTags retrieves all unique tags from assistants with filtering
GetAssistantTags(filter AssistantFilter, locale ...string) ([]Tag, error)
// GetAssistant retrieves a single assistant by ID
GetAssistant(assistantID string, fields []string, locale ...string) (*AssistantModel, error)
// DeleteAssistants deletes assistants based on filter conditions
DeleteAssistants(filter AssistantFilter) (int64, error)
}
// Store combines ChatStore and AssistantStore interfaces
// This is the main interface for the storage layer
type Store interface {
ChatStore
AssistantStore
}
// SpaceStore defines the interface for Space snapshot operations
@ -725,7 +803,7 @@ type SpaceStore interface {
// Restore sets multiple key-value pairs from a snapshot
Restore(data map[string]interface{}) error
}
````
```
### Data Structures
@ -735,11 +813,12 @@ type Chat struct {
ChatID string `json:"chat_id"`
Title string `json:"title,omitempty"`
AssistantID string `json:"assistant_id"`
LastConnector string `json:"last_connector,omitempty"` // Last used connector ID
Mode string `json:"mode"`
Status string `json:"status"`
Public bool `json:"public"`
Share string `json:"share"` // "private" or "team"
Sort int `json:"sort"`
Status string `json:"status"` // "active" or "archived"
Public bool `json:"public"` // Whether shared across all teams
Share string `json:"share"` // "private" or "team"
Sort int `json:"sort"` // Sort order for display
LastMessageAt *time.Time `json:"last_message_at,omitempty"`
Metadata map[string]interface{} `json:"metadata,omitempty"`
CreatedAt time.Time `json:"created_at"`
@ -751,19 +830,21 @@ type Message struct {
MessageID string `json:"message_id"`
ChatID string `json:"chat_id"`
RequestID string `json:"request_id,omitempty"`
Role string `json:"role"`
Type string `json:"type"`
Role string `json:"role"` // "user" or "assistant"
Type string `json:"type"` // "text", "image", "loading", "tool_call", "retrieval", etc.
Props map[string]interface{} `json:"props"`
BlockID string `json:"block_id,omitempty"`
ThreadID string `json:"thread_id,omitempty"`
AssistantID string `json:"assistant_id,omitempty"`
Connector string `json:"connector,omitempty"` // Connector ID used for this message
Sequence int `json:"sequence"`
Metadata map[string]interface{} `json:"metadata,omitempty"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
}
// Resume represents an execution state for recovery (only stored on failure/interrupt)
// Resume represents an execution state for recovery
// Only stored when request is interrupted or failed
type Resume struct {
ResumeID string `json:"resume_id"`
ChatID string `json:"chat_id"`
@ -772,7 +853,7 @@ type Resume struct {
StackID string `json:"stack_id"`
StackParentID string `json:"stack_parent_id,omitempty"`
StackDepth int `json:"stack_depth"`
Type string `json:"type"`
Type string `json:"type"` // "input", "hook_create", "llm", "tool", "hook_next", "delegate"
Status string `json:"status"` // "failed" or "interrupted"
Input map[string]interface{} `json:"input,omitempty"`
Output map[string]interface{} `json:"output,omitempty"`
@ -783,6 +864,22 @@ type Resume struct {
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
}
// ResumeStatus constants
const (
ResumeStatusFailed = "failed"
ResumeStatusInterrupted = "interrupted"
)
// ResumeType constants
const (
ResumeTypeInput = "input"
ResumeTypeHookCreate = "hook_create"
ResumeTypeLLM = "llm"
ResumeTypeTool = "tool"
ResumeTypeHookNext = "hook_next"
ResumeTypeDelegate = "delegate"
)
```
### Filter Structures
@ -790,27 +887,34 @@ type Resume struct {
```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"`
// Time range filter
StartTime *time.Time `json:"start_time,omitempty"` // Filter chats after this time
EndTime *time.Time `json:"end_time,omitempty"` // Filter chats before this time
TimeField string `json:"time_field,omitempty"` // Field for time filter: "created_at" or "last_message_at" (default)
StartTime *time.Time `json:"start_time,omitempty"` // Filter chats after this time
EndTime *time.Time `json:"end_time,omitempty"` // Filter chats before this time
TimeField string `json:"time_field,omitempty"` // Field for time filter: "created_at" or "last_message_at" (default)
// Sorting
OrderBy string `json:"order_by,omitempty"` // Field to sort by (default: "last_message_at")
Order string `json:"order,omitempty"` // Sort order: "desc" (default) or "asc"
OrderBy string `json:"order_by,omitempty"` // Field to sort by (default: "last_message_at")
Order string `json:"order,omitempty"` // Sort order: "desc" (default) or "asc"
// Response format
GroupBy string `json:"group_by,omitempty"` // "time" for time-based groups, empty for flat list
GroupBy string `json:"group_by,omitempty"` // "time" for time-based groups, empty for flat list
// Pagination
Page int `json:"page,omitempty"`
PageSize int `json:"pagesize,omitempty"`
Page int `json:"page,omitempty"`
PageSize int `json:"pagesize,omitempty"`
// 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
@ -1068,9 +1172,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",
@ -1079,7 +1183,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
@ -1097,7 +1229,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",
@ -1107,7 +1239,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{
@ -1117,6 +1249,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,
@ -1352,8 +1495,226 @@ Main Agent concurrently calls 3 tasks:
- Within a block, optionally group by `thread_id` to show parallel results
- Use `sequence` for chronological display
## HTTP API
The chat storage provides RESTful HTTP APIs for managing chat sessions and messages.
**Base Path:** `/v1/chat`
### Chat Sessions
| Method | Endpoint | Description |
|--------|----------|-------------|
| `GET` | `/sessions` | List chat sessions with pagination and filtering |
| `GET` | `/sessions/:chat_id` | Get a single chat session |
| `PUT` | `/sessions/:chat_id` | Update chat session (title, status, metadata) |
| `DELETE` | `/sessions/:chat_id` | Delete chat session |
| `GET` | `/sessions/:chat_id/messages` | Get messages for a chat session |
### List Chat Sessions
**Request:**
```
GET /v1/chat/sessions?page=1&pagesize=20&assistant_id=xxx&status=active&keywords=search&group_by=time
```
**Query Parameters:**
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `page` | int | 1 | Page number |
| `pagesize` | int | 20 | Items per page (max 100) |
| `assistant_id` | string | - | Filter by assistant ID |
| `status` | string | - | Filter by status: `active`, `archived` |
| `keywords` | string | - | Search in title |
| `start_time` | RFC3339 | - | Filter chats after this time |
| `end_time` | RFC3339 | - | Filter chats before this time |
| `time_field` | string | `last_message_at` | Field for time filter: `created_at` or `last_message_at` |
| `order_by` | string | `last_message_at` | Sort field |
| `order` | string | `desc` | Sort order: `asc` or `desc` |
| `group_by` | string | - | Set to `time` for time-based grouping |
**Response:**
```json
{
"data": [
{
"chat_id": "chat_123",
"title": "Weather Query",
"assistant_id": "weather_assistant",
"status": "active",
"last_message_at": "2024-01-15T10:30:00Z",
"created_at": "2024-01-15T10:00:00Z"
}
],
"groups": [
{
"key": "today",
"label": "Today",
"chats": [...],
"count": 3
},
{
"key": "yesterday",
"label": "Yesterday",
"chats": [...],
"count": 5
}
],
"page": 1,
"pagesize": 20,
"pagecount": 5,
"total": 100
}
```
### Get Chat Session
**Request:**
```
GET /v1/chat/sessions/chat_123
```
**Response:**
```json
{
"chat_id": "chat_123",
"title": "Weather Query",
"assistant_id": "weather_assistant",
"mode": "chat",
"status": "active",
"public": false,
"share": "private",
"last_message_at": "2024-01-15T10:30:00Z",
"metadata": {},
"created_at": "2024-01-15T10:00:00Z",
"updated_at": "2024-01-15T10:30:00Z"
}
```
### Update Chat Session
**Request:**
```
PUT /v1/chat/sessions/chat_123
Content-Type: application/json
{
"title": "New Title",
"status": "archived",
"metadata": {"custom_field": "value"}
}
```
**Response:**
```json
{
"message": "Chat updated successfully",
"chat_id": "chat_123"
}
```
### Delete Chat Session
**Request:**
```
DELETE /v1/chat/sessions/chat_123
```
**Response:**
```json
{
"message": "Chat deleted successfully",
"chat_id": "chat_123"
}
```
### Get Chat Messages
**Request:**
```
GET /v1/chat/sessions/chat_123/messages?limit=100&offset=0&role=assistant&type=text
```
**Query Parameters:**
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `request_id` | string | - | Filter by request ID |
| `role` | string | - | Filter by role: `user`, `assistant` |
| `block_id` | string | - | Filter by block ID |
| `thread_id` | string | - | Filter by thread ID |
| `type` | string | - | Filter by message type |
| `limit` | int | 100 | Max messages to return (max 1000) |
| `offset` | int | 0 | Offset for pagination |
**Response:**
```json
{
"chat_id": "chat_123",
"messages": [
{
"message_id": "msg_001",
"chat_id": "chat_123",
"request_id": "req_abc",
"role": "user",
"type": "user_input",
"props": {
"content": "What's the weather?",
"role": "user"
},
"sequence": 1,
"created_at": "2024-01-15T10:00:00Z"
},
{
"message_id": "msg_002",
"chat_id": "chat_123",
"request_id": "req_abc",
"role": "assistant",
"type": "text",
"props": {
"content": "The weather in San Francisco is 18°C and sunny."
},
"block_id": "B1",
"assistant_id": "weather_assistant",
"sequence": 2,
"created_at": "2024-01-15T10:00:05Z"
}
],
"count": 2
}
```
### Permission Filtering
All endpoints respect Yao's permission system:
| Constraint | Behavior |
|------------|----------|
| `OwnerOnly` | User can only access their own chats (`__yao_created_by` matches) |
| `TeamOnly` | User can access own chats OR team-shared chats (`share = "team"`) |
| No constraints | Full access (for admin users) |
**Permission Fields Used:**
- `__yao_created_by`: User who created the chat
- `__yao_team_id`: Team ID for team-level access
- `public`: Whether chat is public to all
- `share`: Sharing scope (`private` or `team`)
## Related Documents
- [OpenAPI Request Design](../../openapi/request/REQUEST_DESIGN.md) - Global request tracking, billing, rate limiting
- [Trace Module](../../trace/README.md) - Detailed execution tracing for debugging
- [Agent Context](../context/README.md) - Context and message handling
````

154
agent/store/mongo/mongo.go
View file

@ -10,88 +10,150 @@ func NewMongo() types.Store {
return &Mongo{}
}
// GetChats retrieves a list of chats
func (m *Mongo) GetChats(sid string, filter types.ChatFilter, locale ...string) (*types.ChatGroupResponse, error) {
return &types.ChatGroupResponse{}, nil
}
// =============================================================================
// Chat Management
// =============================================================================
// GetChat retrieves a single chat's information
func (m *Mongo) GetChat(sid string, cid string, locale ...string) (*types.ChatInfo, error) {
return &types.ChatInfo{}, nil
}
// GetChatWithFilter retrieves a single chat's information with filter options
func (m *Mongo) GetChatWithFilter(sid string, cid string, filter types.ChatFilter, locale ...string) (*types.ChatInfo, error) {
return &types.ChatInfo{}, nil
}
// GetHistory retrieves chat history
func (m *Mongo) GetHistory(sid string, cid string, locale ...string) ([]map[string]interface{}, error) {
return []map[string]interface{}{}, nil
}
// GetHistoryWithFilter retrieves chat history with filter options
func (m *Mongo) GetHistoryWithFilter(sid string, cid string, filter types.ChatFilter, locale ...string) ([]map[string]interface{}, error) {
return []map[string]interface{}{}, nil
}
// SaveHistory saves chat history
func (m *Mongo) SaveHistory(sid string, messages []map[string]interface{}, cid string, context map[string]interface{}) error {
// CreateChat creates a new chat session
func (m *Mongo) CreateChat(chat *types.Chat) error {
// TODO: implement
return nil
}
// DeleteChat deletes a single chat
func (m *Mongo) DeleteChat(sid string, cid string) error {
// GetChat retrieves a single chat by ID
func (m *Mongo) GetChat(chatID string) (*types.Chat, error) {
// TODO: implement
return nil, nil
}
// UpdateChat updates chat fields
func (m *Mongo) UpdateChat(chatID string, updates map[string]interface{}) error {
// TODO: implement
return nil
}
// DeleteAllChats deletes all chats
func (m *Mongo) DeleteAllChats(sid string) error {
// DeleteChat deletes a chat and its associated messages
func (m *Mongo) DeleteChat(chatID string) error {
// TODO: implement
return nil
}
// UpdateChatTitle updates chat title
func (m *Mongo) UpdateChatTitle(sid string, cid string, title string) error {
// ListChats retrieves a paginated list of chats with optional grouping
func (m *Mongo) ListChats(filter types.ChatFilter) (*types.ChatList, error) {
// TODO: implement
return nil, nil
}
// =============================================================================
// Message Management
// =============================================================================
// SaveMessages batch saves messages for a chat
func (m *Mongo) SaveMessages(chatID string, messages []*types.Message) error {
// TODO: implement
return nil
}
// GetMessages retrieves messages for a chat with filtering
func (m *Mongo) GetMessages(chatID string, filter types.MessageFilter) ([]*types.Message, error) {
// TODO: implement
return nil, nil
}
// UpdateMessage updates a single message
func (m *Mongo) UpdateMessage(messageID string, updates map[string]interface{}) error {
// TODO: implement
return nil
}
// DeleteMessages deletes specific messages from a chat
func (m *Mongo) DeleteMessages(chatID string, messageIDs []string) error {
// TODO: implement
return nil
}
// =============================================================================
// Resume Management (only called on failure/interrupt)
// =============================================================================
// SaveResume batch saves resume records
func (m *Mongo) SaveResume(records []*types.Resume) error {
// TODO: implement
return nil
}
// GetResume retrieves all resume records for a chat
func (m *Mongo) GetResume(chatID string) ([]*types.Resume, error) {
// TODO: implement
return nil, nil
}
// GetLastResume retrieves the last resume record for a chat
func (m *Mongo) GetLastResume(chatID string) (*types.Resume, error) {
// TODO: implement
return nil, nil
}
// GetResumeByStackID retrieves resume records for a specific stack
func (m *Mongo) GetResumeByStackID(stackID string) ([]*types.Resume, error) {
// TODO: implement
return nil, nil
}
// GetStackPath returns the stack path from root to the given stack
func (m *Mongo) GetStackPath(stackID string) ([]string, error) {
// TODO: implement
return nil, nil
}
// DeleteResume deletes all resume records for a chat
func (m *Mongo) DeleteResume(chatID string) error {
// TODO: implement
return nil
}
// =============================================================================
// Assistant Management
// =============================================================================
// SaveAssistant saves assistant information
func (m *Mongo) SaveAssistant(assistant *types.AssistantModel) (string, error) {
// TODO: implement
return assistant.ID, nil
}
// UpdateAssistant updates specific fields of an assistant
func (m *Mongo) UpdateAssistant(assistantID string, updates map[string]interface{}) error {
// TODO: implement
return nil
}
// DeleteAssistant deletes an assistant
func (m *Mongo) DeleteAssistant(assistantID string) error {
// TODO: implement
return nil
}
// GetAssistants retrieves a list of assistants
func (m *Mongo) GetAssistants(filter types.AssistantFilter, locale ...string) (*types.AssistantList, error) {
// TODO: implement
return &types.AssistantList{}, nil
}
// GetAssistant retrieves a single assistant by ID
// fields: Optional list of fields to retrieve. If empty, a default set of fields will be returned.
func (m *Mongo) GetAssistant(assistantID string, fields []string, locale ...string) (*types.AssistantModel, error) {
return nil, nil
}
// DeleteAssistants deletes assistants based on filter conditions (not implemented)
func (m *Mongo) DeleteAssistants(filter types.AssistantFilter) (int64, error) {
return 0, nil
}
// GetAssistantTags retrieves all unique tags from assistants with filtering
func (m *Mongo) GetAssistantTags(filter types.AssistantFilter, locale ...string) ([]types.Tag, error) {
// TODO: implement
return []types.Tag{}, nil
}
// Close closes the store and releases any resources
func (m *Mongo) Close() error {
return nil
// GetAssistant retrieves a single assistant by ID
func (m *Mongo) GetAssistant(assistantID string, fields []string, locale ...string) (*types.AssistantModel, error) {
// TODO: implement
return nil, nil
}
// DeleteAssistants deletes assistants based on filter conditions
func (m *Mongo) DeleteAssistants(filter types.AssistantFilter) (int64, error) {
// TODO: implement
return 0, nil
}

156
agent/store/redis/redis.go
View file

@ -1,4 +1,4 @@
package store
package redis
import "github.com/yaoapp/yao/agent/store/types"
@ -10,88 +10,150 @@ func NewRedis() types.Store {
return &Redis{}
}
// GetChats retrieves a list of chats
func (r *Redis) GetChats(sid string, filter types.ChatFilter, locale ...string) (*types.ChatGroupResponse, error) {
return &types.ChatGroupResponse{}, nil
}
// =============================================================================
// Chat Management
// =============================================================================
// GetChat retrieves a single chat's information
func (r *Redis) GetChat(sid string, cid string, locale ...string) (*types.ChatInfo, error) {
return &types.ChatInfo{}, nil
}
// GetChatWithFilter retrieves a single chat's information with filter options
func (r *Redis) GetChatWithFilter(sid string, cid string, filter types.ChatFilter, locale ...string) (*types.ChatInfo, error) {
return &types.ChatInfo{}, nil
}
// GetHistory retrieves chat history
func (r *Redis) GetHistory(sid string, cid string, locale ...string) ([]map[string]interface{}, error) {
return []map[string]interface{}{}, nil
}
// GetHistoryWithFilter retrieves chat history with filter options
func (r *Redis) GetHistoryWithFilter(sid string, cid string, filter types.ChatFilter, locale ...string) ([]map[string]interface{}, error) {
return []map[string]interface{}{}, nil
}
// SaveHistory saves chat history
func (r *Redis) SaveHistory(sid string, messages []map[string]interface{}, cid string, context map[string]interface{}) error {
// CreateChat creates a new chat session
func (r *Redis) CreateChat(chat *types.Chat) error {
// TODO: implement
return nil
}
// DeleteChat deletes a single chat
func (r *Redis) DeleteChat(sid string, cid string) error {
// GetChat retrieves a single chat by ID
func (r *Redis) GetChat(chatID string) (*types.Chat, error) {
// TODO: implement
return nil, nil
}
// UpdateChat updates chat fields
func (r *Redis) UpdateChat(chatID string, updates map[string]interface{}) error {
// TODO: implement
return nil
}
// DeleteAllChats deletes all chats
func (r *Redis) DeleteAllChats(sid string) error {
// DeleteChat deletes a chat and its associated messages
func (r *Redis) DeleteChat(chatID string) error {
// TODO: implement
return nil
}
// UpdateChatTitle updates chat title
func (r *Redis) UpdateChatTitle(sid string, cid string, title string) error {
// ListChats retrieves a paginated list of chats with optional grouping
func (r *Redis) ListChats(filter types.ChatFilter) (*types.ChatList, error) {
// TODO: implement
return nil, nil
}
// =============================================================================
// Message Management
// =============================================================================
// SaveMessages batch saves messages for a chat
func (r *Redis) SaveMessages(chatID string, messages []*types.Message) error {
// TODO: implement
return nil
}
// GetMessages retrieves messages for a chat with filtering
func (r *Redis) GetMessages(chatID string, filter types.MessageFilter) ([]*types.Message, error) {
// TODO: implement
return nil, nil
}
// UpdateMessage updates a single message
func (r *Redis) UpdateMessage(messageID string, updates map[string]interface{}) error {
// TODO: implement
return nil
}
// DeleteMessages deletes specific messages from a chat
func (r *Redis) DeleteMessages(chatID string, messageIDs []string) error {
// TODO: implement
return nil
}
// =============================================================================
// Resume Management (only called on failure/interrupt)
// =============================================================================
// SaveResume batch saves resume records
func (r *Redis) SaveResume(records []*types.Resume) error {
// TODO: implement
return nil
}
// GetResume retrieves all resume records for a chat
func (r *Redis) GetResume(chatID string) ([]*types.Resume, error) {
// TODO: implement
return nil, nil
}
// GetLastResume retrieves the last resume record for a chat
func (r *Redis) GetLastResume(chatID string) (*types.Resume, error) {
// TODO: implement
return nil, nil
}
// GetResumeByStackID retrieves resume records for a specific stack
func (r *Redis) GetResumeByStackID(stackID string) ([]*types.Resume, error) {
// TODO: implement
return nil, nil
}
// GetStackPath returns the stack path from root to the given stack
func (r *Redis) GetStackPath(stackID string) ([]string, error) {
// TODO: implement
return nil, nil
}
// DeleteResume deletes all resume records for a chat
func (r *Redis) DeleteResume(chatID string) error {
// TODO: implement
return nil
}
// =============================================================================
// Assistant Management
// =============================================================================
// SaveAssistant saves assistant information
func (r *Redis) SaveAssistant(assistant *types.AssistantModel) (string, error) {
// TODO: implement
return assistant.ID, nil
}
// UpdateAssistant updates specific fields of an assistant
func (r *Redis) UpdateAssistant(assistantID string, updates map[string]interface{}) error {
// TODO: implement
return nil
}
// DeleteAssistant deletes an assistant
func (r *Redis) DeleteAssistant(assistantID string) error {
// TODO: implement
return nil
}
// GetAssistants retrieves a list of assistants
func (r *Redis) GetAssistants(filter types.AssistantFilter, locale ...string) (*types.AssistantList, error) {
// TODO: implement
return &types.AssistantList{}, nil
}
// GetAssistant retrieves a single assistant by ID
// fields: Optional list of fields to retrieve. If empty, a default set of fields will be returned.
func (r *Redis) GetAssistant(assistantID string, fields []string, locale ...string) (*types.AssistantModel, error) {
return nil, nil
}
// DeleteAssistants deletes assistants based on filter conditions (not implemented)
func (r *Redis) DeleteAssistants(filter types.AssistantFilter) (int64, error) {
return 0, nil
}
// GetAssistantTags retrieves all unique tags from assistants with filtering
func (r *Redis) GetAssistantTags(filter types.AssistantFilter, locale ...string) ([]types.Tag, error) {
// TODO: implement
return []types.Tag{}, nil
}
// Close closes the store and releases any resources
func (r *Redis) Close() error {
return nil
// GetAssistant retrieves a single assistant by ID
func (r *Redis) GetAssistant(assistantID string, fields []string, locale ...string) (*types.AssistantModel, error) {
// TODO: implement
return nil, nil
}
// DeleteAssistants deletes assistants based on filter conditions
func (r *Redis) DeleteAssistants(filter types.AssistantFilter) (int64, error) {
// TODO: implement
return 0, nil
}

165
agent/store/types/store.go
View file

@ -1,66 +1,108 @@
package types
// Store defines the conversation storage interface
// Provides basic operations required for conversation management
type Store interface {
// GetChats retrieves a list of chats
// sid: Session ID
// filter: Filter conditions
// Returns: Grouped chat list and potential error
GetChats(sid string, filter ChatFilter, locale ...string) (*ChatGroupResponse, error)
// ChatStore defines the chat storage interface
// Provides operations for chat, message, and resume management
type ChatStore interface {
// ==========================================================================
// Chat Management
// ==========================================================================
// GetChat retrieves a single chat's information
// sid: Session ID
// cid: Chat ID
// CreateChat creates a new chat session
// chat: Chat session to create
// Returns: Potential error
CreateChat(chat *Chat) error
// GetChat retrieves a single chat by ID
// chatID: Chat ID
// Returns: Chat information and potential error
GetChat(sid string, cid string, locale ...string) (*ChatInfo, error)
GetChat(chatID string) (*Chat, error)
// GetChatWithFilter retrieves a single chat's information with filter options
// sid: Session ID
// cid: Chat ID
// filter: Filter conditions
// Returns: Chat information and potential error
GetChatWithFilter(sid string, cid string, filter ChatFilter, locale ...string) (*ChatInfo, error)
// GetHistory retrieves chat history
// sid: Session ID
// cid: Chat ID
// Returns: History record list and potential error
GetHistory(sid string, cid string, locale ...string) ([]map[string]interface{}, error)
// GetHistoryWithFilter retrieves chat history with filter options
// sid: Session ID
// cid: Chat ID
// filter: Filter conditions
// Returns: History record list and potential error
GetHistoryWithFilter(sid string, cid string, filter ChatFilter, locale ...string) ([]map[string]interface{}, error)
// SaveHistory saves chat history
// sid: Session ID
// messages: Message list
// cid: Chat ID
// context: Context information
// UpdateChat updates chat fields
// chatID: Chat ID
// updates: Map of fields to update
// Returns: Potential error
SaveHistory(sid string, messages []map[string]interface{}, cid string, context map[string]interface{}) error
UpdateChat(chatID string, updates map[string]interface{}) error
// DeleteChat deletes a single chat
// sid: Session ID
// cid: Chat ID
// DeleteChat deletes a chat and its associated messages
// chatID: Chat ID
// Returns: Potential error
DeleteChat(sid string, cid string) error
DeleteChat(chatID string) error
// DeleteAllChats deletes all chats
// sid: Session ID
// ListChats retrieves a paginated list of chats with optional grouping
// filter: Filter conditions including time range, sorting, and grouping
// Returns: Paginated chat list (flat or grouped) and potential error
ListChats(filter ChatFilter) (*ChatList, error)
// ==========================================================================
// Message Management
// ==========================================================================
// SaveMessages batch saves messages for a chat
// This is the primary write method - messages are buffered during execution
// and batch-written at the end of a request
// chatID: Parent chat ID
// messages: Messages to save (includes user input and assistant responses)
// Returns: Potential error
DeleteAllChats(sid string) error
SaveMessages(chatID string, messages []*Message) error
// UpdateChatTitle updates chat title
// sid: Session ID
// cid: Chat ID
// title: New title
// GetMessages retrieves messages for a chat with filtering
// chatID: Chat ID
// filter: Filter conditions (role, type, block, thread, etc.)
// Returns: Message list and potential error
GetMessages(chatID string, filter MessageFilter) ([]*Message, error)
// UpdateMessage updates a single message
// messageID: Message ID
// updates: Map of fields to update
// Returns: Potential error
UpdateChatTitle(sid string, cid string, title string) error
UpdateMessage(messageID string, updates map[string]interface{}) error
// DeleteMessages deletes specific messages from a chat
// chatID: Chat ID
// messageIDs: List of message IDs to delete
// Returns: Potential error
DeleteMessages(chatID string, messageIDs []string) error
// ==========================================================================
// Resume Management (only called on failure/interrupt)
// ==========================================================================
// SaveResume batch saves resume records
// Only called when request is interrupted or failed
// records: Resume records to save
// Returns: Potential error
SaveResume(records []*Resume) error
// GetResume retrieves all resume records for a chat
// chatID: Chat ID
// Returns: Resume records and potential error
GetResume(chatID string) ([]*Resume, error)
// GetLastResume retrieves the last (most recent) resume record for a chat
// chatID: Chat ID
// Returns: Last resume record and potential error
GetLastResume(chatID string) (*Resume, error)
// GetResumeByStackID retrieves resume records for a specific stack
// stackID: Stack ID
// Returns: Resume records and potential error
GetResumeByStackID(stackID string) ([]*Resume, error)
// GetStackPath returns the stack path from root to the given stack
// stackID: Current stack ID
// Returns: Stack path [root_stack_id, ..., current_stack_id] and potential error
GetStackPath(stackID string) ([]string, error)
// DeleteResume deletes all resume records for a chat
// Called after successful resume to clean up
// chatID: Chat ID
// Returns: Potential error
DeleteResume(chatID string) error
}
// AssistantStore defines the assistant storage interface
// Separated from ChatStore for clearer responsibility
type AssistantStore interface {
// SaveAssistant saves assistant information
// assistant: Assistant information
// Returns: Assistant ID and potential error
@ -91,7 +133,7 @@ type Store interface {
// GetAssistant retrieves a single assistant by ID
// assistantID: Assistant ID
// fields: List of fields to select, empty/nil means default fields (AssistantDefaultFields)
// fields: List of fields to select, empty/nil means default fields
// locale: Optional locale for i18n translations
// Returns: Assistant information and potential error
GetAssistant(assistantID string, fields []string, locale ...string) (*AssistantModel, error)
@ -100,8 +142,21 @@ type Store interface {
// filter: Filter conditions
// Returns: Number of deleted records and potential error
DeleteAssistants(filter AssistantFilter) (int64, error)
// Close closes the store and releases any resources
// Returns: Potential error
Close() error
}
// Store combines ChatStore and AssistantStore interfaces
// This is the main interface for the storage layer
type Store interface {
ChatStore
AssistantStore
}
// SpaceStore defines the interface for Space snapshot operations
// Note: Space itself uses plan.Space interface, this is for persistence
type SpaceStore interface {
// Snapshot returns all key-value pairs in the space
Snapshot() map[string]interface{}
// Restore sets multiple key-value pairs from a snapshot
Restore(data map[string]interface{}) error
}

163
agent/store/types/types.go
View file

@ -3,6 +3,7 @@ package types
import (
"encoding/json"
"fmt"
"time"
graphragtypes "github.com/yaoapp/gou/graphrag/types"
"github.com/yaoapp/xun/dbal/query"
@ -19,40 +20,154 @@ type Setting struct {
Options map[string]interface{} `json:"optional,omitempty" yaml:"optional,omitempty"` // The options for the store
}
// ChatInfo represents the chat information structure
// Contains basic information and history for a single chat
type ChatInfo struct {
Chat map[string]interface{} `json:"chat"` // Basic chat information
History []map[string]interface{} `json:"history"` // Chat history records
// =============================================================================
// Chat Types
// =============================================================================
// Chat represents a chat session
type Chat struct {
ChatID string `json:"chat_id"`
Title string `json:"title,omitempty"`
AssistantID string `json:"assistant_id"`
LastConnector string `json:"last_connector,omitempty"` // Last used connector ID (updated on each message)
Mode string `json:"mode"`
Status string `json:"status"` // "active" or "archived"
Public bool `json:"public"` // Whether shared across all teams
Share string `json:"share"` // "private" or "team"
Sort int `json:"sort"` // Sort order for display
LastMessageAt *time.Time `json:"last_message_at,omitempty"`
Metadata map[string]interface{} `json:"metadata,omitempty"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
// Permission fields (managed by Yao framework when permission: true)
CreatedBy string `json:"__yao_created_by,omitempty"` // User ID who created the record
UpdatedBy string `json:"__yao_updated_by,omitempty"` // User ID who last updated
TeamID string `json:"__yao_team_id,omitempty"` // Team ID for team-level access
TenantID string `json:"__yao_tenant_id,omitempty"` // Tenant ID for multi-tenancy
}
// ChatFilter represents the chat filter structure
// Used for filtering and pagination when retrieving chat lists
// ChatFilter for listing chats
type ChatFilter struct {
Keywords string `json:"keywords,omitempty"` // Keyword search
Page int `json:"page,omitempty"` // Page number, starting from 1
PageSize int `json:"pagesize,omitempty"` // Number of items per page
Order string `json:"order,omitempty"` // Sort order: desc/asc
Silent *bool `json:"silent,omitempty"` // Include silent messages (default: false)
UserID string `json:"user_id,omitempty"`
TeamID string `json:"team_id,omitempty"`
AssistantID string `json:"assistant_id,omitempty"`
Status string `json:"status,omitempty"`
Keywords string `json:"keywords,omitempty"`
// Time range filter
StartTime *time.Time `json:"start_time,omitempty"` // Filter chats after this time
EndTime *time.Time `json:"end_time,omitempty"` // Filter chats before this time
TimeField string `json:"time_field,omitempty"` // Field for time filter: "created_at" or "last_message_at" (default)
// Sorting
OrderBy string `json:"order_by,omitempty"` // Field to sort by (default: "last_message_at")
Order string `json:"order,omitempty"` // Sort order: "desc" (default) or "asc"
// Response format
GroupBy string `json:"group_by,omitempty"` // "time" for time-based groups, empty for flat list
// Pagination
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
}
// ChatGroup represents the chat group structure
// Groups chats by date
// ChatList paginated response with time-based grouping
type ChatList struct {
Data []*Chat `json:"data"`
Groups []*ChatGroup `json:"groups,omitempty"` // Time-based groups for UI display
Page int `json:"page"`
PageSize int `json:"pagesize"`
PageCount int `json:"pagecount"`
Total int `json:"total"`
}
// ChatGroup represents a time-based group of chats
type ChatGroup struct {
Label string `json:"label"` // Group label (typically a date)
Chats []map[string]interface{} `json:"chats"` // List of chats in this group
Label string `json:"label"` // "Today", "Yesterday", "This Week", "This Month", "Earlier"
Key string `json:"key"` // "today", "yesterday", "this_week", "this_month", "earlier"
Chats []*Chat `json:"chats"` // Chats in this group
Count int `json:"count"` // Number of chats in group
}
// ChatGroupResponse represents the paginated chat group response
// Contains paginated chat group information
type ChatGroupResponse struct {
Groups []ChatGroup `json:"groups"` // List of chat groups
Page int `json:"page"` // Current page number
PageSize int `json:"pagesize"` // Items per page
Total int64 `json:"total"` // Total number of records
LastPage int `json:"last_page"` // Last page number
// =============================================================================
// Message Types
// =============================================================================
// Message represents a chat message
type Message struct {
MessageID string `json:"message_id"`
ChatID string `json:"chat_id"`
RequestID string `json:"request_id,omitempty"`
Role string `json:"role"` // "user" or "assistant"
Type string `json:"type"` // "text", "image", "loading", "tool_call", "retrieval", etc.
Props map[string]interface{} `json:"props"`
BlockID string `json:"block_id,omitempty"`
ThreadID string `json:"thread_id,omitempty"`
AssistantID string `json:"assistant_id,omitempty"`
Connector string `json:"connector,omitempty"` // Connector ID used for this message
Sequence int `json:"sequence"`
Metadata map[string]interface{} `json:"metadata,omitempty"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
}
// MessageFilter for listing messages
type MessageFilter struct {
RequestID string `json:"request_id,omitempty"`
Role string `json:"role,omitempty"`
BlockID string `json:"block_id,omitempty"`
ThreadID string `json:"thread_id,omitempty"`
Type string `json:"type,omitempty"`
Limit int `json:"limit,omitempty"`
Offset int `json:"offset,omitempty"`
}
// =============================================================================
// Resume Types (for recovery from interruption/failure)
// =============================================================================
// Resume represents an execution state for recovery
// Only stored when request is interrupted or failed
type Resume struct {
ResumeID string `json:"resume_id"`
ChatID string `json:"chat_id"`
RequestID string `json:"request_id"`
AssistantID string `json:"assistant_id"`
StackID string `json:"stack_id"`
StackParentID string `json:"stack_parent_id,omitempty"`
StackDepth int `json:"stack_depth"`
Type string `json:"type"` // "input", "hook_create", "llm", "tool", "hook_next", "delegate"
Status string `json:"status"` // "failed" or "interrupted"
Input map[string]interface{} `json:"input,omitempty"`
Output map[string]interface{} `json:"output,omitempty"`
SpaceSnapshot map[string]interface{} `json:"space_snapshot,omitempty"` // Shared space data for recovery
Error string `json:"error,omitempty"`
Sequence int `json:"sequence"`
Metadata map[string]interface{} `json:"metadata,omitempty"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
}
// ResumeStatus constants
const (
ResumeStatusFailed = "failed"
ResumeStatusInterrupted = "interrupted"
)
// ResumeType constants
const (
ResumeTypeInput = "input"
ResumeTypeHookCreate = "hook_create"
ResumeTypeLLM = "llm"
ResumeTypeTool = "tool"
ResumeTypeHookNext = "hook_next"
ResumeTypeDelegate = "delegate"
)
// AssistantFilter represents the assistant filter structure
// Used for filtering and pagination when retrieving assistant lists
type AssistantFilter struct {

68
agent/store/xun/assistant.go
View file

@ -14,7 +14,7 @@ import (
)
// SaveAssistant saves assistant information
func (conv *Xun) SaveAssistant(assistant *types.AssistantModel) (string, error) {
func (store *Xun) SaveAssistant(assistant *types.AssistantModel) (string, error) {
if assistant == nil {
return "", fmt.Errorf("assistant cannot be nil")
}
@ -33,15 +33,15 @@ func (conv *Xun) SaveAssistant(assistant *types.AssistantModel) (string, error)
// Generate assistant_id if not provided
if assistant.ID == "" {
var err error
assistant.ID, err = conv.GenerateAssistantID()
assistant.ID, err = store.GenerateAssistantID()
if err != nil {
return "", err
}
}
// Check if assistant exists
exists, err := conv.query.New().
Table(conv.getAssistantTable()).
exists, err := store.query.New().
Table(store.getAssistantTable()).
Where("assistant_id", assistant.ID).
Exists()
if err != nil {
@ -197,8 +197,8 @@ func (conv *Xun) SaveAssistant(assistant *types.AssistantModel) (string, error)
// Update or insert
if exists {
_, err := conv.query.New().
Table(conv.getAssistantTable()).
_, err := store.query.New().
Table(store.getAssistantTable()).
Where("assistant_id", assistant.ID).
Update(data)
if err != nil {
@ -207,8 +207,8 @@ func (conv *Xun) SaveAssistant(assistant *types.AssistantModel) (string, error)
return assistant.ID, nil
}
err = conv.query.New().
Table(conv.getAssistantTable()).
err = store.query.New().
Table(store.getAssistantTable()).
Insert(data)
if err != nil {
return "", err
@ -217,7 +217,7 @@ func (conv *Xun) SaveAssistant(assistant *types.AssistantModel) (string, error)
}
// UpdateAssistant updates specific fields of an assistant
func (conv *Xun) UpdateAssistant(assistantID string, updates map[string]interface{}) error {
func (store *Xun) UpdateAssistant(assistantID string, updates map[string]interface{}) error {
if assistantID == "" {
return fmt.Errorf("assistant_id is required")
}
@ -226,8 +226,8 @@ func (conv *Xun) UpdateAssistant(assistantID string, updates map[string]interfac
}
// Check if assistant exists
exists, err := conv.query.New().
Table(conv.getAssistantTable()).
exists, err := store.query.New().
Table(store.getAssistantTable()).
Where("assistant_id", assistantID).
Exists()
if err != nil {
@ -291,8 +291,8 @@ func (conv *Xun) UpdateAssistant(assistantID string, updates map[string]interfac
}
// Perform update
_, err = conv.query.New().
Table(conv.getAssistantTable()).
_, err = store.query.New().
Table(store.getAssistantTable()).
Where("assistant_id", assistantID).
Update(data)
@ -300,10 +300,10 @@ func (conv *Xun) UpdateAssistant(assistantID string, updates map[string]interfac
}
// DeleteAssistant deletes an assistant by assistant_id
func (conv *Xun) DeleteAssistant(assistantID string) error {
func (store *Xun) DeleteAssistant(assistantID string) error {
// Check if assistant exists
exists, err := conv.query.New().
Table(conv.getAssistantTable()).
exists, err := store.query.New().
Table(store.getAssistantTable()).
Where("assistant_id", assistantID).
Exists()
if err != nil {
@ -314,17 +314,17 @@ func (conv *Xun) DeleteAssistant(assistantID string) error {
return fmt.Errorf("assistant %s not found", assistantID)
}
_, err = conv.query.New().
Table(conv.getAssistantTable()).
_, err = store.query.New().
Table(store.getAssistantTable()).
Where("assistant_id", assistantID).
Delete()
return err
}
// GetAssistants retrieves assistants with pagination and filtering
func (conv *Xun) GetAssistants(filter types.AssistantFilter, locale ...string) (*types.AssistantList, error) {
qb := conv.query.New().
Table(conv.getAssistantTable())
func (store *Xun) GetAssistants(filter types.AssistantFilter, locale ...string) (*types.AssistantList, error) {
qb := store.query.New().
Table(store.getAssistantTable())
// Apply tag filter if provided
if len(filter.Tags) > 0 {
@ -450,7 +450,7 @@ func (conv *Xun) GetAssistants(filter types.AssistantFilter, locale ...string) (
}
// Parse JSON fields
conv.parseJSONFields(data, jsonFields)
store.parseJSONFields(data, jsonFields)
// Convert map to types.AssistantModel using existing helper function
model, err := types.ToAssistantModel(data)
@ -461,7 +461,7 @@ func (conv *Xun) GetAssistants(filter types.AssistantFilter, locale ...string) (
// Apply i18n translations if locale is provided
if len(locale) > 0 && locale[0] != "" && model != nil {
conv.translate(model, model.ID, locale[0])
store.translate(model, model.ID, locale[0])
}
assistants = append(assistants, model)
@ -479,9 +479,9 @@ func (conv *Xun) GetAssistants(filter types.AssistantFilter, locale ...string) (
}
// GetAssistant retrieves a single assistant by ID
func (conv *Xun) GetAssistant(assistantID string, fields []string, locale ...string) (*types.AssistantModel, error) {
qb := conv.query.New().
Table(conv.getAssistantTable()).
func (store *Xun) GetAssistant(assistantID string, fields []string, locale ...string) (*types.AssistantModel, error) {
qb := store.query.New().
Table(store.getAssistantTable()).
Where("assistant_id", assistantID)
// Apply select fields with security validation
@ -515,7 +515,7 @@ func (conv *Xun) GetAssistant(assistantID string, fields []string, locale ...str
// Parse JSON fields
jsonFields := []string{"tags", "modes", "options", "prompts", "prompt_presets", "connector_options", "workflow", "kb", "db", "mcp", "placeholder", "locales", "uses"}
conv.parseJSONFields(data, jsonFields)
store.parseJSONFields(data, jsonFields)
// Convert map to types.AssistantModel
model := &types.AssistantModel{
@ -661,16 +661,16 @@ func (conv *Xun) GetAssistant(assistantID string, fields []string, locale ...str
// Apply i18n translation if locale is provided
if len(locale) > 0 && locale[0] != "" {
conv.translate(model, assistantID, locale[0])
store.translate(model, assistantID, locale[0])
}
return model, nil
}
// DeleteAssistants deletes assistants based on filter conditions
func (conv *Xun) DeleteAssistants(filter types.AssistantFilter) (int64, error) {
qb := conv.query.New().
Table(conv.getAssistantTable())
func (store *Xun) DeleteAssistants(filter types.AssistantFilter) (int64, error) {
qb := store.query.New().
Table(store.getAssistantTable())
// Apply tag filter if provided
if len(filter.Tags) > 0 {
@ -729,8 +729,8 @@ func (conv *Xun) DeleteAssistants(filter types.AssistantFilter) (int64, error) {
}
// GetAssistantTags retrieves all unique tags from assistants with filtering
func (conv *Xun) GetAssistantTags(filter types.AssistantFilter, locale ...string) ([]types.Tag, error) {
qb := conv.query.New().Table(conv.getAssistantTable())
func (store *Xun) GetAssistantTags(filter types.AssistantFilter, locale ...string) ([]types.Tag, error) {
qb := store.query.New().Table(store.getAssistantTable())
// Apply type filter (default to "assistant")
typeFilter := "assistant"
@ -803,7 +803,7 @@ func (conv *Xun) GetAssistantTags(filter types.AssistantFilter, locale ...string
}
// translate applies i18n translation to assistant model fields
func (conv *Xun) translate(model *types.AssistantModel, assistantID string, locale string) {
func (store *Xun) translate(model *types.AssistantModel, assistantID string, locale string) {
if model == nil {
return
}

81
agent/store/xun/assistant_test.go
View file

@ -1,4 +1,4 @@
package xun
package xun_test
import (
"fmt"
@ -11,6 +11,7 @@ import (
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/i18n"
"github.com/yaoapp/yao/agent/store/types"
"github.com/yaoapp/yao/agent/store/xun"
"github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/test"
)
@ -31,13 +32,12 @@ func TestSaveAssistant(t *testing.T) {
defer test.Clean()
// Create a new xun store
store, err := NewXun(types.Setting{
store, err := xun.NewXun(types.Setting{
Connector: "default",
})
if err != nil {
t.Fatalf("Failed to create store: %v", err)
}
defer store.Close()
t.Run("CreateNewAssistant", func(t *testing.T) {
assistant := &types.AssistantModel{
@ -648,13 +648,12 @@ func TestDeleteAssistant(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
store, err := NewXun(types.Setting{
store, err := xun.NewXun(types.Setting{
Connector: "default",
})
if err != nil {
t.Fatalf("Failed to create store: %v", err)
}
defer store.Close()
t.Run("DeleteExistingAssistant", func(t *testing.T) {
// Create assistant
@ -696,13 +695,12 @@ func TestGetAssistant(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
store, err := NewXun(types.Setting{
store, err := xun.NewXun(types.Setting{
Connector: "default",
})
if err != nil {
t.Fatalf("Failed to create store: %v", err)
}
defer store.Close()
t.Run("GetExistingAssistant", func(t *testing.T) {
// Create assistant
@ -764,13 +762,12 @@ func TestGetAssistants(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
store, err := NewXun(types.Setting{
store, err := xun.NewXun(types.Setting{
Connector: "default",
})
if err != nil {
t.Fatalf("Failed to create store: %v", err)
}
defer store.Close()
// Clean up existing data before creating test assistants
deleted, err := store.DeleteAssistants(types.AssistantFilter{})
@ -1078,13 +1075,12 @@ func TestDeleteAssistants(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
store, err := NewXun(types.Setting{
store, err := xun.NewXun(types.Setting{
Connector: "default",
})
if err != nil {
t.Fatalf("Failed to create store: %v", err)
}
defer store.Close()
t.Run("DeleteByTag", func(t *testing.T) {
// Create assistants with specific tag
@ -1205,13 +1201,12 @@ func TestGetAssistantTags(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
store, err := NewXun(types.Setting{
store, err := xun.NewXun(types.Setting{
Connector: "default",
})
if err != nil {
t.Fatalf("Failed to create store: %v", err)
}
defer store.Close()
t.Run("GetUniqueTags", func(t *testing.T) {
// Create assistants with various tags
@ -1433,57 +1428,17 @@ func TestGetAssistantTags(t *testing.T) {
})
}
// TestGenerateAssistantID tests the ID generation function
func TestGenerateAssistantID(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
store, err := NewXun(types.Setting{
Connector: "default",
})
if err != nil {
t.Fatalf("Failed to create store: %v", err)
}
defer store.Close()
xunStore := store.(*Xun)
t.Run("GenerateUniqueIDs", func(t *testing.T) {
ids := make(map[string]bool)
for i := 0; i < 10; i++ {
id, err := xunStore.GenerateAssistantID()
if err != nil {
t.Fatalf("Failed to generate ID: %v", err)
}
// Verify ID format (6 digits)
if len(id) != 6 {
t.Errorf("Expected 6-digit ID, got %s (length %d)", id, len(id))
}
// Verify ID is unique
if ids[id] {
t.Errorf("Generated duplicate ID: %s", id)
}
ids[id] = true
}
t.Logf("Generated %d unique IDs", len(ids))
})
}
// TestAssistantPermissionFields tests permission management fields
func TestAssistantPermissionFields(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
store, err := NewXun(types.Setting{
store, err := xun.NewXun(types.Setting{
Connector: "default",
})
if err != nil {
t.Fatalf("Failed to create store: %v", err)
}
defer store.Close()
t.Run("SaveWithPermissionFields", func(t *testing.T) {
assistant := &types.AssistantModel{
@ -1608,13 +1563,12 @@ func TestEmptyStringAsNull(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
store, err := NewXun(types.Setting{
store, err := xun.NewXun(types.Setting{
Connector: "default",
})
if err != nil {
t.Fatalf("Failed to create store: %v", err)
}
defer store.Close()
t.Run("EmptyStringsStoredAsNull", func(t *testing.T) {
// Create assistant with empty strings for nullable fields
@ -1711,13 +1665,12 @@ func TestGetAssistantWithLocale(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
store, err := NewXun(types.Setting{
store, err := xun.NewXun(types.Setting{
Connector: "default",
})
if err != nil {
t.Fatalf("Failed to create store: %v", err)
}
defer store.Close()
t.Run("GetAssistantWithLocaleTranslation", func(t *testing.T) {
// Create assistant with i18n locales
@ -1837,13 +1790,12 @@ func TestGetAssistantsWithLocale(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
store, err := NewXun(types.Setting{
store, err := xun.NewXun(types.Setting{
Connector: "default",
})
if err != nil {
t.Fatalf("Failed to create store: %v", err)
}
defer store.Close()
t.Run("GetAssistantsWithLocaleTranslation", func(t *testing.T) {
// Create assistant with i18n locales
@ -1950,13 +1902,12 @@ func TestGetAssistantsWithQueryFilter(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
store, err := NewXun(types.Setting{
store, err := xun.NewXun(types.Setting{
Connector: "default",
})
if err != nil {
t.Fatalf("Failed to create store: %v", err)
}
defer store.Close()
// Create test assistants with different permission settings
assistants := []types.AssistantModel{
@ -2195,13 +2146,12 @@ func TestUpdateAssistant(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
store, err := NewXun(types.Setting{
store, err := xun.NewXun(types.Setting{
Connector: "default",
})
if err != nil {
t.Fatalf("Failed to create store: %v", err)
}
defer store.Close()
t.Run("UpdateSingleField", func(t *testing.T) {
// Create assistant
@ -3028,13 +2978,12 @@ func TestAssistantCompleteWorkflow(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
store, err := NewXun(types.Setting{
store, err := xun.NewXun(types.Setting{
Connector: "default",
})
if err != nil {
t.Fatalf("Failed to create store: %v", err)
}
defer store.Close()
t.Run("CompleteWorkflow", func(t *testing.T) {
// Step 1: Create multiple assistants

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

@ -3,425 +3,471 @@ package xun
import (
"fmt"
"math"
"strings"
"time"
"github.com/yaoapp/yao/agent/i18n"
"github.com/google/uuid"
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/yao/agent/store/types"
)
// UpdateChatTitle update the chat title
func (conv *Xun) UpdateChatTitle(sid string, cid string, title string) error {
userID, err := conv.getUserID(sid)
// =============================================================================
// Chat Management
// =============================================================================
// CreateChat creates a new chat session
func (store *Xun) CreateChat(chat *types.Chat) error {
if chat == nil {
return fmt.Errorf("chat cannot be nil")
}
// Validate required fields
if chat.AssistantID == "" {
return fmt.Errorf("assistant_id is required")
}
// Generate chat_id if not provided
if chat.ChatID == "" {
chat.ChatID = uuid.New().String()
}
// Check if chat already exists
exists, err := store.newQueryChat().
Where("chat_id", chat.ChatID).
Exists()
if err != nil {
return err
}
if exists {
return fmt.Errorf("chat %s already exists", chat.ChatID)
}
_, err = conv.newQueryChat().
Where("sid", userID).
Where("chat_id", cid).
// Set defaults
if chat.Mode == "" {
chat.Mode = "chat"
}
if chat.Status == "" {
chat.Status = "active"
}
if chat.Share == "" {
chat.Share = "private"
}
// Prepare data
data := map[string]interface{}{
"chat_id": chat.ChatID,
"assistant_id": chat.AssistantID,
"mode": chat.Mode,
"status": chat.Status,
"public": chat.Public,
"share": chat.Share,
"sort": chat.Sort,
"created_at": time.Now(),
"updated_at": time.Now(),
}
// Handle nullable fields
if chat.Title != "" {
data["title"] = chat.Title
}
if chat.LastConnector != "" {
data["last_connector"] = chat.LastConnector
}
if chat.LastMessageAt != nil {
data["last_message_at"] = *chat.LastMessageAt
}
if chat.Metadata != nil {
metadataJSON, err := jsoniter.MarshalToString(chat.Metadata)
if err != nil {
return fmt.Errorf("failed to marshal metadata: %w", err)
}
data["metadata"] = metadataJSON
}
// Handle permission fields (Yao framework permission: true)
if chat.CreatedBy != "" {
data["__yao_created_by"] = chat.CreatedBy
}
if chat.UpdatedBy != "" {
data["__yao_updated_by"] = chat.UpdatedBy
}
if chat.TeamID != "" {
data["__yao_team_id"] = chat.TeamID
}
if chat.TenantID != "" {
data["__yao_tenant_id"] = chat.TenantID
}
// Insert
return store.newQueryChat().Insert(data)
}
// GetChat retrieves a single chat by ID
func (store *Xun) GetChat(chatID string) (*types.Chat, error) {
if chatID == "" {
return nil, fmt.Errorf("chat_id is required")
}
row, err := store.newQueryChat().
Where("chat_id", chatID).
WhereNull("deleted_at").
First()
if err != nil {
return nil, err
}
if row == nil {
return nil, fmt.Errorf("chat %s not found", chatID)
}
data := row.ToMap()
if len(data) == 0 || data["chat_id"] == nil {
return nil, fmt.Errorf("chat %s not found", chatID)
}
return store.rowToChat(data)
}
// UpdateChat updates chat fields
func (store *Xun) UpdateChat(chatID string, updates map[string]interface{}) error {
if chatID == "" {
return fmt.Errorf("chat_id is required")
}
if len(updates) == 0 {
return fmt.Errorf("no fields to update")
}
// Check if chat exists
exists, err := store.newQueryChat().
Where("chat_id", chatID).
WhereNull("deleted_at").
Exists()
if err != nil {
return err
}
if !exists {
return fmt.Errorf("chat %s not found", chatID)
}
// Prepare update data
data := make(map[string]interface{})
// Process each update field
for key, value := range updates {
// Skip system fields
if key == "chat_id" || key == "created_at" {
continue
}
// Handle metadata specially
if key == "metadata" {
if value != nil {
metadataJSON, err := jsoniter.MarshalToString(value)
if err != nil {
return fmt.Errorf("failed to marshal metadata: %w", err)
}
data["metadata"] = metadataJSON
} else {
data["metadata"] = 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.newQueryChat().
Where("chat_id", chatID).
Update(data)
return err
}
// DeleteChat deletes a chat and its associated messages (soft delete)
func (store *Xun) DeleteChat(chatID string) error {
if chatID == "" {
return fmt.Errorf("chat_id is required")
}
// Check if chat exists
exists, err := store.newQueryChat().
Where("chat_id", chatID).
WhereNull("deleted_at").
Exists()
if err != nil {
return err
}
if !exists {
return fmt.Errorf("chat %s not found", chatID)
}
// Soft delete the chat
_, err = store.newQueryChat().
Where("chat_id", chatID).
Update(map[string]interface{}{
"title": title,
"deleted_at": time.Now(),
"updated_at": time.Now(),
})
return err
}
// GetChat get the chat info and its history
func (conv *Xun) GetChat(sid string, cid string, locale ...string) (*types.ChatInfo, error) {
// userID, err := conv.getUserID(sid)
// if err != nil {
// return nil, err
// }
// Get chat info
qb := conv.newQueryChat().
Select("chat_id", "title", "assistant_id").
// Where("sid", userID).
Where("chat_id", cid)
row, err := qb.First()
if err != nil {
return nil, err
}
// Return nil if chat_id is nil (means no chat found)
if row.Get("chat_id") == nil {
return nil, nil
}
chat := map[string]interface{}{
"chat_id": row.Get("chat_id"),
"title": row.Get("title"),
"assistant_id": row.Get("assistant_id"),
}
// Get assistant details if assistant_id exists
if assistantID := row.Get("assistant_id"); assistantID != nil && assistantID != "" {
assistant, err := conv.query.New().
Table(conv.getAssistantTable()).
Select("name", "avatar").
Where("assistant_id", assistantID).
First()
if err != nil {
return nil, err
}
name := assistant.Get("name")
if len(locale) > 0 {
lang := strings.ToLower(locale[0])
name = i18n.Translate(assistantID.(string), lang, name).(string)
}
if assistant != nil {
chat["assistant_name"] = name
chat["assistant_avatar"] = assistant.Get("avatar")
}
}
// Get chat history with default filter (silent=false)
history, err := conv.GetHistory(sid, cid, locale...)
if err != nil {
return nil, err
}
return &types.ChatInfo{
Chat: chat,
History: history,
}, nil
}
// GetChatWithFilter get the chat info and its history with filter options
func (conv *Xun) GetChatWithFilter(sid string, cid string, filter types.ChatFilter, locale ...string) (*types.ChatInfo, error) {
// userID, err := conv.getUserID(sid)
// if err != nil {
// return nil, err
// }
// Get chat info
qb := conv.newQueryChat().
Select("chat_id", "title", "assistant_id").
// Where("sid", userID).
Where("chat_id", cid)
row, err := qb.First()
if err != nil {
return nil, err
}
// Return nil if chat_id is nil (means no chat found)
if row.Get("chat_id") == nil {
return nil, nil
}
chat := map[string]interface{}{
"chat_id": row.Get("chat_id"),
"title": row.Get("title"),
"assistant_id": row.Get("assistant_id"),
}
// Get assistant details if assistant_id exists
if assistantID := row.Get("assistant_id"); assistantID != nil && assistantID != "" {
assistant, err := conv.query.New().
Table(conv.getAssistantTable()).
Select("name", "avatar").
Where("assistant_id", assistantID).
First()
if err != nil {
return nil, err
}
if assistant != nil {
chat["assistant_name"] = assistant.Get("name")
chat["assistant_avatar"] = assistant.Get("avatar")
}
}
// Get chat history with filter
history, err := conv.GetHistoryWithFilter(sid, cid, filter, locale...)
if err != nil {
return nil, err
}
return &types.ChatInfo{
Chat: chat,
History: history,
}, nil
}
// DeleteChat deletes a specific chat and its history
func (conv *Xun) DeleteChat(sid string, cid string) error {
userID, err := conv.getUserID(sid)
if err != nil {
return err
}
// Delete history records first
_, err = conv.newQuery().
Where("sid", userID).
Where("cid", cid).
Delete()
if err != nil {
return err
}
// Then delete the chat
_, err = conv.newQueryChat().
Where("sid", userID).
Where("chat_id", cid).
Limit(1).
Delete()
return err
}
// DeleteAllChats deletes all chats and their histories for a user
func (conv *Xun) DeleteAllChats(sid string) error {
userID, err := conv.getUserID(sid)
if err != nil {
return err
}
// Delete history records first
_, err = conv.newQuery().
Where("sid", userID).
Delete()
if err != nil {
return err
}
// Then delete all chats
_, err = conv.newQueryChat().
Where("sid", userID).
Delete()
return err
}
// GetChats get the chat list with grouping by date
func (conv *Xun) GetChats(sid string, filter types.ChatFilter, locale ...string) (*types.ChatGroupResponse, error) {
// Default behavior: exclude silent chats
if filter.Silent == nil {
silentFalse := false
filter.Silent = &silentFalse
}
return conv.getChatsWithFilter(sid, filter, locale...)
}
// getChatsWithFilter get the chats with filter options
func (conv *Xun) getChatsWithFilter(sid string, filter types.ChatFilter, locale ...string) (*types.ChatGroupResponse, error) {
// userID, err := conv.getUserID(sid)
// if err != nil {
// return nil, err
// }
// Set default values
// ListChats retrieves a paginated list of chats with optional grouping
func (store *Xun) ListChats(filter types.ChatFilter) (*types.ChatList, error) {
// Set defaults
if filter.Page <= 0 {
filter.Page = 1
}
if filter.PageSize <= 0 {
filter.PageSize = 20
}
if filter.OrderBy == "" {
filter.OrderBy = "last_message_at"
}
if filter.Order == "" {
filter.Order = "desc"
}
// Get total count
qbCount := conv.newQueryChat()
// Where("sid", userID)
// Apply silent filter if provided
if filter.Silent != nil {
if *filter.Silent {
// Include all chats (both silent and non-silent)
} else {
// Only include non-silent chats
qbCount.Where("silent", false)
}
if filter.TimeField == "" {
filter.TimeField = "last_message_at"
}
// Apply keyword filter if provided
if filter.Keywords != "" {
qbCount.Where("title", "like", fmt.Sprintf("%%%s%%", filter.Keywords))
// Build base query
qb := store.newQueryChat().WhereNull("deleted_at")
// 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)
}
total, err := qbCount.Count()
if err != nil {
return nil, err
// Apply business filters
if filter.AssistantID != "" {
qb.Where("assistant_id", filter.AssistantID)
}
// Calculate last page
lastPage := int(math.Ceil(float64(total) / float64(filter.PageSize)))
if lastPage < 1 {
lastPage = 1
if filter.Status != "" {
qb.Where("status", filter.Status)
}
// Get chats with pagination
qb := conv.newQueryChat().
Select("chat_id", "title", "assistant_id", "silent", "created_at", "updated_at")
// Where("sid", userID)
// Apply silent filter if provided
if filter.Silent != nil {
if *filter.Silent {
// Include all chats (both silent and non-silent)
} else {
// Only include non-silent chats
qb.Where("silent", false)
}
}
// Apply keyword filter if provided
if filter.Keywords != "" {
qb.Where("title", "like", fmt.Sprintf("%%%s%%", filter.Keywords))
}
// Apply pagination
offset := (filter.Page - 1) * filter.PageSize
qb.OrderBy("updated_at", filter.Order).
Offset(offset).
Limit(filter.PageSize)
// Apply time range filter
if filter.StartTime != nil {
qb.Where(filter.TimeField, ">=", *filter.StartTime)
}
if filter.EndTime != nil {
qb.Where(filter.TimeField, "<=", *filter.EndTime)
}
rows, err := qb.Get()
// 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)
}
// Get total count
total, err := qb.Clone().Count()
if err != nil {
return nil, err
}
// Group chats by date
today := time.Now().Truncate(24 * time.Hour)
// Calculate pagination
pageCount := int(math.Ceil(float64(total) / float64(filter.PageSize)))
if pageCount < 1 {
pageCount = 1
}
offset := (filter.Page - 1) * filter.PageSize
// Get paginated results
rows, err := qb.OrderBy(filter.OrderBy, filter.Order).
Offset(offset).
Limit(filter.PageSize).
Get()
if err != nil {
return nil, err
}
// Convert rows to Chat objects
chats := make([]*types.Chat, 0, len(rows))
for _, row := range rows {
data := row.ToMap()
if data == nil || data["chat_id"] == nil {
continue
}
chat, err := store.rowToChat(data)
if err != nil {
continue
}
chats = append(chats, chat)
}
result := &types.ChatList{
Data: chats,
Page: filter.Page,
PageSize: filter.PageSize,
PageCount: pageCount,
Total: int(total),
}
// Apply time-based grouping if requested
if filter.GroupBy == "time" {
result.Groups = store.groupChatsByTime(chats)
}
return result, nil
}
// =============================================================================
// Helper Functions
// =============================================================================
// rowToChat converts a database row to a Chat struct
func (store *Xun) rowToChat(data map[string]interface{}) (*types.Chat, error) {
chat := &types.Chat{
ChatID: getString(data, "chat_id"),
Title: getString(data, "title"),
AssistantID: getString(data, "assistant_id"),
LastConnector: getString(data, "last_connector"),
Mode: getString(data, "mode"),
Status: getString(data, "status"),
Public: getBool(data, "public"),
Share: getString(data, "share"),
Sort: getInt(data, "sort"),
}
// Handle timestamps
if createdAt := getTime(data, "created_at"); createdAt != nil {
chat.CreatedAt = *createdAt
}
if updatedAt := getTime(data, "updated_at"); updatedAt != nil {
chat.UpdatedAt = *updatedAt
}
if lastMsgAt := getTime(data, "last_message_at"); lastMsgAt != nil {
chat.LastMessageAt = lastMsgAt
}
// Handle metadata
if metadata := data["metadata"]; metadata != nil {
if metaStr, ok := metadata.(string); ok && metaStr != "" {
var meta map[string]interface{}
if err := jsoniter.UnmarshalFromString(metaStr, &meta); err == nil {
chat.Metadata = meta
}
} else if metaMap, ok := metadata.(map[string]interface{}); ok {
chat.Metadata = metaMap
}
}
// Handle permission fields
chat.CreatedBy = getString(data, "__yao_created_by")
chat.UpdatedBy = getString(data, "__yao_updated_by")
chat.TeamID = getString(data, "__yao_team_id")
chat.TenantID = getString(data, "__yao_tenant_id")
return chat, nil
}
// groupChatsByTime groups chats by time periods
func (store *Xun) groupChatsByTime(chats []*types.Chat) []*types.ChatGroup {
now := time.Now()
today := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location())
yesterday := today.AddDate(0, 0, -1)
thisWeekStart := today.AddDate(0, 0, -int(today.Weekday()))
lastWeekStart := thisWeekStart.AddDate(0, 0, -7)
lastWeekEnd := thisWeekStart.AddDate(0, 0, -1)
thisMonthStart := time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, now.Location())
groups := map[string][]map[string]interface{}{
"Today": {},
"Yesterday": {},
"This Week": {},
"Last Week": {},
"Even Earlier": {},
groups := map[string]*types.ChatGroup{
"today": {Key: "today", Label: "Today", Chats: []*types.Chat{}},
"yesterday": {Key: "yesterday", Label: "Yesterday", Chats: []*types.Chat{}},
"this_week": {Key: "this_week", Label: "This Week", Chats: []*types.Chat{}},
"this_month": {Key: "this_month", Label: "This Month", Chats: []*types.Chat{}},
"earlier": {Key: "earlier", Label: "Earlier", Chats: []*types.Chat{}},
}
// Collect assistant IDs to fetch their details
assistantIDs := []interface{}{}
for _, row := range rows {
if assistantID := row.Get("assistant_id"); assistantID != nil && assistantID != "" {
assistantIDs = append(assistantIDs, assistantID)
}
}
// Fetch assistant details
assistantMap := map[string]map[string]interface{}{}
if len(assistantIDs) > 0 {
assistants, err := conv.query.New().
Table(conv.getAssistantTable()).
Select("assistant_id", "name", "avatar").
WhereIn("assistant_id", assistantIDs).
Get()
if err != nil {
return nil, err
for _, chat := range chats {
// Use last_message_at if available, otherwise created_at
var chatTime time.Time
if chat.LastMessageAt != nil {
chatTime = *chat.LastMessageAt
} else {
chatTime = chat.CreatedAt
}
for _, assistant := range assistants {
if id := assistant.Get("assistant_id"); id != nil {
name := assistant.Get("name")
if len(locale) > 0 {
lang := strings.ToLower(locale[0])
name = i18n.Translate(id.(string), lang, name).(string)
}
assistantMap[fmt.Sprintf("%v", id)] = map[string]interface{}{
"name": name,
"avatar": assistant.Get("avatar"),
}
}
}
}
for _, row := range rows {
chatID := row.Get("chat_id")
if chatID == nil || chatID == "" {
continue
}
chat := map[string]interface{}{
"chat_id": chatID,
"title": row.Get("title"),
"assistant_id": row.Get("assistant_id"),
"silent": row.Get("silent"),
}
// Add assistant details if available
if assistantID := row.Get("assistant_id"); assistantID != nil && assistantID != "" {
if assistant, ok := assistantMap[fmt.Sprintf("%v", assistantID)]; ok {
name := assistant["name"]
if len(locale) > 0 {
lang := strings.ToLower(locale[0])
name = i18n.Translate(assistantID.(string), lang, name).(string)
}
chat["assistant_name"] = name
chat["assistant_avatar"] = assistant["avatar"]
}
}
var dbDatetime = row.Get("updated_at")
if dbDatetime == nil {
dbDatetime = row.Get("created_at")
}
var createdAt time.Time
switch v := dbDatetime.(type) {
case time.Time:
createdAt = v
case string:
parsed, err := time.Parse("2006-01-02 15:04:05.999999-07:00", v)
if err != nil {
// Try alternative format
parsed, err = time.Parse(time.RFC3339, v)
if err != nil {
continue
}
}
createdAt = parsed
default:
continue
}
createdDate := createdAt.Truncate(24 * time.Hour)
chatDate := time.Date(chatTime.Year(), chatTime.Month(), chatTime.Day(), 0, 0, 0, 0, chatTime.Location())
switch {
case createdDate.Equal(today):
groups["Today"] = append(groups["Today"], chat)
case createdDate.Equal(yesterday):
groups["Yesterday"] = append(groups["Yesterday"], chat)
case createdDate.After(thisWeekStart) && createdDate.Before(today):
groups["This Week"] = append(groups["This Week"], chat)
case createdDate.After(lastWeekStart) && createdDate.Before(lastWeekEnd.AddDate(0, 0, 1)):
groups["Last Week"] = append(groups["Last Week"], chat)
case chatDate.Equal(today) || chatDate.After(today):
groups["today"].Chats = append(groups["today"].Chats, chat)
case chatDate.Equal(yesterday):
groups["yesterday"].Chats = append(groups["yesterday"].Chats, chat)
case chatDate.After(thisWeekStart) || chatDate.Equal(thisWeekStart):
groups["this_week"].Chats = append(groups["this_week"].Chats, chat)
case chatDate.After(thisMonthStart) || chatDate.Equal(thisMonthStart):
groups["this_month"].Chats = append(groups["this_month"].Chats, chat)
default:
groups["Even Earlier"] = append(groups["Even Earlier"], chat)
groups["earlier"].Chats = append(groups["earlier"].Chats, chat)
}
}
// Convert to ordered slice and apply i18n
result := []types.ChatGroup{}
for _, label := range []string{"Today", "Yesterday", "This Week", "Last Week", "Even Earlier"} {
if len(groups[label]) > 0 {
translatedLabel := label
if len(locale) > 0 {
lang := strings.ToLower(locale[0])
translatedLabel = i18n.TranslateGlobal(lang, label).(string)
}
result = append(result, types.ChatGroup{
Label: translatedLabel,
Chats: groups[label],
})
// Update counts and filter empty groups
result := make([]*types.ChatGroup, 0)
for _, key := range []string{"today", "yesterday", "this_week", "this_month", "earlier"} {
group := groups[key]
group.Count = len(group.Chats)
if group.Count > 0 {
result = append(result, group)
}
}
return &types.ChatGroupResponse{
Groups: result,
Page: filter.Page,
PageSize: filter.PageSize,
Total: total,
LastPage: lastPage,
}, nil
return result
}
// getTime helper function to convert database value to time.Time pointer
func getTime(data map[string]interface{}, key string) *time.Time {
if v := data[key]; v != nil {
switch t := v.(type) {
case time.Time:
return &t
case *time.Time:
return t
case string:
// Try parsing various formats
formats := []string{
time.RFC3339,
"2006-01-02 15:04:05",
"2006-01-02 15:04:05.999999-07:00",
"2006-01-02T15:04:05Z",
}
for _, format := range formats {
if parsed, err := time.Parse(format, t); err == nil {
return &parsed
}
}
}
}
return nil
}
// UpdateChatLastMessageAt updates the last_message_at timestamp for a chat
func (store *Xun) UpdateChatLastMessageAt(chatID string, timestamp time.Time) error {
if chatID == "" {
return fmt.Errorf("chat_id is required")
}
_, err := store.newQueryChat().
Where("chat_id", chatID).
Update(map[string]interface{}{
"last_message_at": timestamp,
"updated_at": time.Now(),
})
return err
}

1181
agent/store/xun/chat_test.go Normal file
View file

File diff suppressed because it is too large Load diff

322
agent/store/xun/history.go
View file

@ -1,322 +0,0 @@
package xun
import (
"fmt"
"strings"
"time"
"github.com/google/uuid"
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/yao/agent/i18n"
"github.com/yaoapp/yao/agent/store/types"
)
// GetHistory get the history
func (conv *Xun) GetHistory(sid string, cid string, locale ...string) ([]map[string]interface{}, error) {
// userID, err := conv.getUserID(sid)
// if err != nil {
// return nil, err
// }
qb := conv.newQuery().
Select("role", "name", "content", "context", "assistant_id", "assistant_name", "assistant_avatar", "mentions", "uid", "silent", "created_at", "updated_at").
// Where("sid", userID).
Where("cid", cid).
OrderBy("id", "desc")
// By default, exclude silent messages
qb.Where("silent", false)
if conv.setting.TTL > 0 {
qb.Where("expired_at", ">", time.Now())
}
limit := 20
if conv.setting.MaxSize > 0 {
limit = conv.setting.MaxSize
}
rows, err := qb.Limit(limit).Get()
if err != nil {
return nil, err
}
res := []map[string]interface{}{}
for _, row := range rows {
assistantName := row.Get("assistant_name")
assistantID := row.Get("assistant_id")
if len(locale) > 0 && assistantID != nil {
lang := strings.ToLower(locale[0])
assistantName = i18n.Translate(assistantID.(string), lang, assistantName).(string)
}
message := map[string]interface{}{
"role": row.Get("role"),
"name": row.Get("name"),
"content": row.Get("content"),
"context": row.Get("context"),
"assistant_id": row.Get("assistant_id"),
"assistant_name": assistantName,
"assistant_avatar": row.Get("assistant_avatar"),
"mentions": row.Get("mentions"),
"uid": row.Get("uid"),
"silent": row.Get("silent"),
"created_at": row.Get("created_at"),
"updated_at": row.Get("updated_at"),
}
res = append([]map[string]interface{}{message}, res...)
}
return res, nil
}
// SaveHistory save the history
func (conv *Xun) SaveHistory(sid string, messages []map[string]interface{}, cid string, context map[string]interface{}) error {
if cid == "" {
cid = uuid.New().String() // Generate a new UUID if cid is empty
}
userID, err := conv.getUserID(sid)
if err != nil {
return err
}
// Get assistant_id from context
var assistantID interface{} = nil
if context != nil {
if id, ok := context["assistant_id"].(string); ok && id != "" {
assistantID = id
}
}
// Get silent flag from context
var silent bool = false
var historyVisible bool = true
if context != nil {
if silentVal, ok := context["silent"]; ok {
switch v := silentVal.(type) {
case bool:
silent = v
case string:
silent = v == "true" || v == "1" || v == "yes"
case int:
silent = v != 0
case float64:
silent = v != 0
}
}
// Get history visible from context
if historyVisibleVal, ok := context["history_visible"]; ok {
switch v := historyVisibleVal.(type) {
case bool:
historyVisible = v
case string:
historyVisible = v == "true" || v == "1" || v == "yes"
case int:
historyVisible = v != 0
case float64:
historyVisible = v != 0
}
}
}
// First ensure chat record exists
exists, err := conv.newQueryChat().
Where("chat_id", cid).
Where("sid", userID).
Exists()
if err != nil {
return err
}
if !exists {
// Create new chat record
err = conv.newQueryChat().
Insert(map[string]interface{}{
"chat_id": cid,
"sid": userID,
"assistant_id": assistantID,
"silent": silent || historyVisible == false,
"created_at": time.Now(),
})
if err != nil {
return err
}
} else {
// Update assistant_id and silent if needed
_, err = conv.newQueryChat().
Where("chat_id", cid).
Where("sid", userID).
Update(map[string]interface{}{
"assistant_id": assistantID,
"silent": silent || historyVisible == false,
})
if err != nil {
return err
}
}
// Save message history
var expiredAt interface{} = nil
values := []map[string]interface{}{}
if conv.setting.TTL > 0 {
expiredAt = time.Now().Add(time.Duration(conv.setting.TTL) * time.Second)
}
now := time.Now()
for _, message := range messages {
// Type assertion safety checks
role, ok := message["role"].(string)
if !ok {
return fmt.Errorf("invalid role type in message: %v", message["role"])
}
content, ok := message["content"].(string)
if !ok {
return fmt.Errorf("invalid content type in message: %v", message["content"])
}
var contextRaw interface{} = nil
if context != nil {
contextRaw, err = jsoniter.MarshalToString(context)
if err != nil {
return err
}
}
// Process mentions if present
var mentionsRaw interface{} = nil
if mentions, ok := message["mentions"].([]interface{}); ok && len(mentions) > 0 {
mentionsRaw, err = jsoniter.MarshalToString(mentions)
if err != nil {
return err
}
}
value := map[string]interface{}{
"role": role,
"name": "",
"content": content,
"sid": userID,
"cid": cid,
"uid": userID,
"context": contextRaw,
"mentions": mentionsRaw,
"assistant_id": nil,
"assistant_name": nil,
"assistant_avatar": nil,
"silent": silent,
"created_at": now,
"updated_at": nil,
"expired_at": expiredAt,
}
if name, ok := message["name"].(string); ok {
value["name"] = name
}
// Add assistant fields if present
if assistantID, ok := message["assistant_id"].(string); ok {
value["assistant_id"] = assistantID
}
if assistantName, ok := message["assistant_name"].(string); ok {
value["assistant_name"] = assistantName
}
if assistantAvatar, ok := message["assistant_avatar"].(string); ok {
value["assistant_avatar"] = assistantAvatar
}
values = append(values, value)
}
err = conv.newQuery().Insert(values)
if err != nil {
return err
}
// Update Chat updated_at
_, err = conv.newQueryChat().
Where("chat_id", cid).
Where("sid", userID).
Update(map[string]interface{}{"updated_at": now})
if err != nil {
return err
}
return nil
}
// GetHistoryWithFilter get the history with filter options
func (conv *Xun) GetHistoryWithFilter(sid string, cid string, filter types.ChatFilter, locale ...string) ([]map[string]interface{}, error) {
userID, err := conv.getUserID(sid)
if err != nil {
return nil, err
}
qb := conv.newQuery().
Select("role", "name", "content", "context", "assistant_id", "assistant_name", "assistant_avatar", "mentions", "uid", "silent", "created_at", "updated_at").
Where("sid", userID).
Where("cid", cid).
OrderBy("id", "desc")
// Apply silent filter if provided, otherwise exclude silent messages by default
if filter.Silent != nil {
if *filter.Silent {
// Include all messages (both silent and non-silent)
} else {
// Only include non-silent messages
qb.Where("silent", false)
}
} else {
// Default behavior: exclude silent messages
qb.Where("silent", false)
}
if conv.setting.TTL > 0 {
qb.Where("expired_at", ">", time.Now())
}
limit := 20
if conv.setting.MaxSize > 0 {
limit = conv.setting.MaxSize
}
if filter.PageSize > 0 {
limit = filter.PageSize
}
// Apply pagination if provided
if filter.Page > 0 {
offset := (filter.Page - 1) * limit
qb.Offset(offset)
}
rows, err := qb.Limit(limit).Get()
if err != nil {
return nil, err
}
res := []map[string]interface{}{}
for _, row := range rows {
message := map[string]interface{}{
"role": row.Get("role"),
"name": row.Get("name"),
"content": row.Get("content"),
"context": row.Get("context"),
"assistant_id": row.Get("assistant_id"),
"assistant_name": row.Get("assistant_name"),
"assistant_avatar": row.Get("assistant_avatar"),
"mentions": row.Get("mentions"),
"uid": row.Get("uid"),
"silent": row.Get("silent"),
"created_at": row.Get("created_at"),
"updated_at": row.Get("updated_at"),
}
res = append([]map[string]interface{}{message}, res...)
}
return res, nil
}

371
agent/store/xun/message.go Normal file
View file

@ -0,0 +1,371 @@
package xun
import (
"fmt"
"time"
"github.com/google/uuid"
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/yao/agent/store/types"
)
// =============================================================================
// Message Management
// =============================================================================
// 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 {
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,
"connector": 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.Connector != "" {
row["connector"] = msg.Connector
}
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) {
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 {
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 soft deletes specific messages from a chat
func (store *Xun) DeleteMessages(chatID string, messageIDs []string) error {
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"),
Connector: getString(data, "connector"),
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
}

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

File diff suppressed because it is too large Load diff

379
agent/store/xun/resume.go Normal file
View file

@ -0,0 +1,379 @@
package xun
import (
"fmt"
"time"
"github.com/google/uuid"
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/yao/agent/store/types"
)
// =============================================================================
// Resume Management (only called on failure/interrupt)
// =============================================================================
// 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 {
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) {
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) {
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) {
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) {
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 soft deletes all resume records for a chat
// Called after successful resume to clean up
func (store *Xun) DeleteResume(chatID string) error {
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!")
})
}

304
agent/store/xun/xun.go
View file

@ -7,7 +7,6 @@ import (
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/gou/model"
"github.com/yaoapp/kun/log"
"github.com/yaoapp/xun/capsule"
"github.com/yaoapp/xun/dbal/query"
"github.com/yaoapp/xun/dbal/schema"
@ -18,274 +17,155 @@ import (
// Xun implements the Store interface using a database backend.
// It provides functionality for:
// - Managing chat conversations and their message histories
// - Managing chat sessions and their messages
// - Organizing chats with pagination and date-based grouping
// - Handling chat metadata like titles and creation dates
// - Managing AI assistants with their configurations and metadata
// - Supporting data expiration through TTL settings
// - Managing resume records for recovery from interruptions
type Xun struct {
query query.Query
schema schema.Schema
setting types.Setting
cleanTicker *time.Ticker
cleanStop chan bool
query query.Query
schema schema.Schema
setting types.Setting
}
// Public interface methods:
//
// NewXun creates a new conversation instance with the given settings
// GetChats retrieves a paginated list of chats grouped by date
// GetChat retrieves a specific chat and its message history
// GetChatWithFilter retrieves a specific chat with filter options
// GetHistory retrieves the message history for a specific chat
// GetHistoryWithFilter retrieves the message history with filter options
// SaveHistory saves new messages to a chat's historys
// DeleteChat deletes a specific chat and its history
// DeleteAllChats deletes all chats and their histories for a user
// UpdateChatTitle updates the title of a specific chat
// NewXun creates a new store instance with the given settings
//
// Chat Management:
// CreateChat creates a new chat session
// GetChat retrieves a single chat by ID
// UpdateChat updates chat fields
// DeleteChat deletes a chat and its associated messages
// ListChats retrieves a paginated list of chats with optional grouping
//
// Message Management:
// SaveMessages batch saves messages for a chat
// GetMessages retrieves messages for a chat with filtering
// UpdateMessage updates a single message
// DeleteMessages deletes specific messages from a chat
//
// Resume Management:
// SaveResume batch saves resume records (only on failure/interrupt)
// GetResume retrieves all resume records for a chat
// GetLastResume retrieves the last resume record for a chat
// GetResumeByStackID retrieves resume records for a specific stack
// GetStackPath returns the stack path from root to the given stack
// DeleteResume deletes all resume records for a chat
//
// Assistant Management:
// SaveAssistant creates or updates an assistant
// UpdateAssistant updates assistant fields
// DeleteAssistant deletes an assistant by assistant_id
// GetAssistants retrieves a paginated list of assistants with filtering
// GetAssistant retrieves a single assistant by assistant_id
// DeleteAssistants deletes assistants based on filter conditions
// GetAssistantTags retrieves all unique tags from assistants
// Close closes the store and releases any resources
// NewXun create a new xun store
func NewXun(setting types.Setting) (types.Store, error) {
conv := &Xun{setting: setting}
store := &Xun{setting: setting}
if setting.Connector == "default" || setting.Connector == "" {
conv.query = capsule.Global.Query()
conv.schema = capsule.Global.Schema()
store.query = capsule.Global.Query()
store.schema = capsule.Global.Schema()
} else {
conn, err := connector.Select(setting.Connector)
if err != nil {
return nil, fmt.Errorf("select store connector %s error: %s", setting.Connector, err.Error())
}
conv.query, err = conn.Query()
store.query, err = conn.Query()
if err != nil {
return nil, fmt.Errorf("query store connector %s error: %s", setting.Connector, err.Error())
}
conv.schema, err = conn.Schema()
store.schema, err = conn.Schema()
if err != nil {
return nil, err
}
}
err := conv.initialize()
if err != nil {
return nil, err
}
return conv, nil
return store, nil
}
// Rename the following functions to start with lowercase letters to make them private:
// =============================================================================
// Query Builders
// =============================================================================
func (conv *Xun) newQuery() query.Query {
qb := conv.query.New()
qb.Table(conv.getHistoryTable())
// newQueryChat creates a new query builder for the chat table
func (store *Xun) newQueryChat() query.Query {
qb := store.query.New()
qb.Table(store.getChatTable())
return qb
}
func (conv *Xun) newQueryChat() query.Query {
qb := conv.query.New()
qb.Table(conv.getChatTable())
// newQueryMessage creates a new query builder for the message table
func (store *Xun) newQueryMessage() query.Query {
qb := store.query.New()
qb.Table(store.getMessageTable())
return qb
}
func (conv *Xun) clean() {
nums, err := conv.newQuery().Where("expired_at", "<=", time.Now()).Delete()
if err != nil {
log.Error("Clean the conversation table error: %s", err.Error())
return
}
if nums > 0 {
log.Trace("Clean the conversation table: %d", nums)
}
// newQueryResume creates a new query builder for the resume table
func (store *Xun) newQueryResume() query.Query {
qb := store.query.New()
qb.Table(store.getResumeTable())
return qb
}
// startAutoClean starts the automatic cleanup routine
func (conv *Xun) startAutoClean() {
if conv.cleanTicker != nil {
conv.stopAutoClean() // Stop existing ticker if any
}
conv.cleanTicker = time.NewTicker(1 * time.Hour) // Clean every hour
conv.cleanStop = make(chan bool)
go func() {
for {
select {
case <-conv.cleanTicker.C:
conv.clean()
case <-conv.cleanStop:
return
}
}
}()
log.Trace("Started automatic cleanup")
// newQueryAssistant creates a new query builder for the assistant table
func (store *Xun) newQueryAssistant() query.Query {
qb := store.query.New()
qb.Table(store.getAssistantTable())
return qb
}
// stopAutoClean stops the automatic cleanup routine
func (conv *Xun) stopAutoClean() {
if conv.cleanTicker != nil {
conv.cleanTicker.Stop()
conv.cleanTicker = nil
}
// =============================================================================
// Table Name Getters
// =============================================================================
if conv.cleanStop != nil {
close(conv.cleanStop)
conv.cleanStop = nil
}
log.Trace("Stopped automatic cleanup")
}
// Close stops the automatic cleanup and closes resources
func (conv *Xun) Close() error {
conv.stopAutoClean()
return nil
}
// Rename Init to initialize to avoid conflicts
func (conv *Xun) initialize() error {
// Start automatic cleanup if TTL is enabled
if conv.setting.TTL > 0 {
conv.startAutoClean()
}
return nil
}
func (conv *Xun) initHistoryTable() error {
historyTable := conv.getHistoryTable()
has, err := conv.schema.HasTable(historyTable)
if err != nil {
return err
}
// Create the history table
if !has {
err = conv.schema.CreateTable(historyTable, func(table schema.Blueprint) {
table.ID("id")
table.String("sid", 255).Index()
table.String("cid", 200).Null().Index()
table.String("uid", 255).Null().Index()
table.String("role", 200).Null().Index()
table.String("name", 200).Null().Index()
table.Text("content").Null()
table.JSON("context").Null()
table.String("assistant_id", 200).Null().Index()
table.String("assistant_name", 200).Null()
table.String("assistant_avatar", 200).Null()
table.JSON("mentions").Null()
table.Boolean("silent").SetDefault(false).Index()
table.TimestampTz("created_at").SetDefaultRaw("CURRENT_TIMESTAMP").Index()
table.TimestampTz("updated_at").Null().Index()
table.TimestampTz("expired_at").Null().Index()
})
if err != nil {
return err
}
log.Trace("Create the conversation history table: %s", historyTable)
}
// Validate the table
tab, err := conv.schema.GetTable(historyTable)
if err != nil {
return err
}
fields := []string{"id", "sid", "cid", "uid", "role", "name", "content", "context", "assistant_id", "assistant_name", "assistant_avatar", "mentions", "silent", "created_at", "updated_at", "expired_at"}
for _, field := range fields {
if !tab.HasColumn(field) {
return fmt.Errorf("%s is required", field)
}
}
return nil
}
func (conv *Xun) initChatTable() error {
chatTable := conv.getChatTable()
has, err := conv.schema.HasTable(chatTable)
if err != nil {
return err
}
// Create the chat table
if !has {
err = conv.schema.CreateTable(chatTable, func(table schema.Blueprint) {
table.ID("id")
table.String("chat_id", 200).Unique().Index()
table.String("title", 200).Null()
table.String("assistant_id", 200).Null().Index()
table.String("sid", 255).Index()
table.Boolean("silent").SetDefault(false).Index()
table.TimestampTz("created_at").SetDefaultRaw("CURRENT_TIMESTAMP").Index()
table.TimestampTz("updated_at").Null().Index()
})
if err != nil {
return err
}
log.Trace("Create the chat table: %s", chatTable)
}
// Validate the table
tab, err := conv.schema.GetTable(chatTable)
if err != nil {
return err
}
fields := []string{"id", "chat_id", "title", "assistant_id", "sid", "silent", "created_at", "updated_at"}
for _, field := range fields {
if !tab.HasColumn(field) {
return fmt.Errorf("%s is required", field)
}
}
return nil
}
func (conv *Xun) getUserID(sid string) (string, error) {
// TODO: get the user id from the authentication system
return "guest", nil
}
func (conv *Xun) getHistoryTable() string {
m := model.Select("__yao.agent.history")
if m != nil && m.MetaData.Table.Name != "" {
return m.MetaData.Table.Name
}
return "__yao.agent.history"
}
func (conv *Xun) getChatTable() string {
// getChatTable returns the chat table name
func (store *Xun) getChatTable() string {
m := model.Select("__yao.agent.chat")
if m != nil && m.MetaData.Table.Name != "" {
return m.MetaData.Table.Name
}
return "__yao.agent.chat"
return "agent_chat"
}
func (conv *Xun) getAssistantTable() string {
// getMessageTable returns the message table name
func (store *Xun) getMessageTable() string {
m := model.Select("__yao.agent.message")
if m != nil && m.MetaData.Table.Name != "" {
return m.MetaData.Table.Name
}
return "agent_message"
}
// getResumeTable returns the resume table name
func (store *Xun) getResumeTable() string {
m := model.Select("__yao.agent.resume")
if m != nil && m.MetaData.Table.Name != "" {
return m.MetaData.Table.Name
}
return "agent_resume"
}
// getAssistantTable returns the assistant table name
func (store *Xun) getAssistantTable() string {
m := model.Select("__yao.agent.assistant")
if m != nil && m.MetaData.Table.Name != "" {
return m.MetaData.Table.Name
}
return "__yao.agent.assistant"
return "agent_assistant"
}
// =============================================================================
// Utility Functions
// =============================================================================
// parseJSONFields parses JSON string fields into their corresponding Go types
func (conv *Xun) parseJSONFields(data map[string]interface{}, fields []string) {
func (store *Xun) parseJSONFields(data map[string]interface{}, fields []string) {
for _, field := range fields {
if val := data[field]; val != nil {
if strVal, ok := val.(string); ok && strVal != "" {
@ -299,7 +179,7 @@ func (conv *Xun) parseJSONFields(data map[string]interface{}, fields []string) {
}
// GenerateAssistantID generates a random-looking 6-digit ID
func (conv *Xun) GenerateAssistantID() (string, error) {
func (store *Xun) GenerateAssistantID() (string, error) {
maxAttempts := 10 // Maximum number of attempts to generate a unique ID
for i := 0; i < maxAttempts; i++ {
// Generate a random number using timestamp and some bit operations
@ -308,8 +188,8 @@ func (conv *Xun) GenerateAssistantID() (string, error) {
hash := fmt.Sprintf("%06d", random)
// Check if this ID already exists
exists, err := conv.query.New().
Table(conv.getAssistantTable()).
exists, err := store.query.New().
Table(store.getAssistantTable()).
Where("assistant_id", hash).
Exists()

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

File diff suppressed because it is too large Load diff

290
data/bindata.go
View file

File diff suppressed because it is too large Load diff

28
openapi/chat/chat.go
View file

@ -12,6 +12,10 @@ func Attach(group *gin.RouterGroup, oauth types.OAuth) {
// Protect all endpoints with OAuth
group.Use(oauth.Guard)
// ==========================================================================
// Chat Completions (Streaming API)
// ==========================================================================
// List Chat Completions
group.GET("/completions", placeholder)
@ -24,7 +28,7 @@ func Attach(group *gin.RouterGroup, oauth types.OAuth) {
// Get Chat Completion Details
group.GET("/completions/:completion_id", placeholder)
// Get Chat Messages
// Get Chat Messages (by completion)
group.GET("/completions/:completion_id/messages", placeholder)
// Delete Chat Completion
@ -33,6 +37,28 @@ func Attach(group *gin.RouterGroup, oauth types.OAuth) {
// Append messages to running completion
group.POST("/completions/:context_id/append", GinAppendMessages)
// ==========================================================================
// Chat Sessions (History Management)
// ==========================================================================
// List chat sessions with pagination and filtering
// Query params: page, pagesize, assistant_id, status, keywords,
// start_time, end_time, time_field, order_by, order, group_by
group.GET("/sessions", ListChats)
// Get a single chat session by ID
group.GET("/sessions/:chat_id", GetChat)
// Update chat session (title, status, metadata)
group.PUT("/sessions/:chat_id", UpdateChat)
// Delete chat session
group.DELETE("/sessions/:chat_id", DeleteChat)
// Get messages for a chat session
// Query params: request_id, role, block_id, thread_id, type, limit, offset
group.GET("/sessions/:chat_id/messages", GetMessages)
}
func placeholder(c *gin.Context) {

544
openapi/chat/session.go Normal file
View file

@ -0,0 +1,544 @@
package chat
import (
"strconv"
"strings"
"time"
"github.com/gin-gonic/gin"
"github.com/yaoapp/xun/dbal/query"
"github.com/yaoapp/yao/agent/assistant"
storetypes "github.com/yaoapp/yao/agent/store/types"
"github.com/yaoapp/yao/openapi/oauth/authorized"
oauthtypes "github.com/yaoapp/yao/openapi/oauth/types"
"github.com/yaoapp/yao/openapi/response"
)
// =============================================================================
// Chat Session Handlers
// =============================================================================
// ListChats lists chat sessions with pagination and filtering
// GET /v1/chat/sessions
func ListChats(c *gin.Context) {
// Get chat store
chatStore := assistant.GetChatStore()
if chatStore == nil {
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: "Chat storage not initialized",
}
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
return
}
// Get authorized information
authInfo := authorized.GetInfo(c)
// Build filter from query parameters
filter := buildChatFilter(c, authInfo)
// Call store to list chats
result, err := chatStore.ListChats(filter)
if err != nil {
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: err.Error(),
}
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
return
}
// Return result
response.RespondWithSuccess(c, response.StatusOK, gin.H{
"data": result.Data,
"groups": result.Groups,
"page": result.Page,
"pagesize": result.PageSize,
"pagecount": result.PageCount,
"total": result.Total,
})
}
// GetChat retrieves a single chat session by ID
// GET /v1/chat/sessions/:chat_id
func GetChat(c *gin.Context) {
// Get chat store
chatStore := assistant.GetChatStore()
if chatStore == nil {
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: "Chat storage not initialized",
}
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
return
}
// Get chat ID from URL parameter
chatID := c.Param("chat_id")
if chatID == "" {
errorResp := &response.ErrorResponse{
Code: response.ErrInvalidRequest.Code,
ErrorDescription: "Chat ID is required",
}
response.RespondWithError(c, response.StatusBadRequest, errorResp)
return
}
// Get authorized information
authInfo := authorized.GetInfo(c)
// Check permission
hasPermission, err := checkChatPermission(chatStore, authInfo, chatID, true)
if err != nil {
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: err.Error(),
}
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
return
}
if !hasPermission {
errorResp := &response.ErrorResponse{
Code: response.ErrAccessDenied.Code,
ErrorDescription: "Forbidden: No permission to access this chat",
}
response.RespondWithError(c, response.StatusForbidden, errorResp)
return
}
// Get chat
chat, err := chatStore.GetChat(chatID)
if err != nil {
// Check if it's a "not found" error
if strings.Contains(err.Error(), "not found") {
errorResp := &response.ErrorResponse{
Code: response.ErrInvalidRequest.Code,
ErrorDescription: "Chat not found",
}
response.RespondWithError(c, response.StatusNotFound, errorResp)
return
}
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: err.Error(),
}
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
return
}
response.RespondWithSuccess(c, response.StatusOK, chat)
}
// UpdateChat updates a chat session
// PUT /v1/chat/sessions/:chat_id
func UpdateChat(c *gin.Context) {
// Get chat store
chatStore := assistant.GetChatStore()
if chatStore == nil {
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: "Chat storage not initialized",
}
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
return
}
// Get chat ID from URL parameter
chatID := c.Param("chat_id")
if chatID == "" {
errorResp := &response.ErrorResponse{
Code: response.ErrInvalidRequest.Code,
ErrorDescription: "Chat ID is required",
}
response.RespondWithError(c, response.StatusBadRequest, errorResp)
return
}
// Parse request body
var req UpdateChatRequest
if err := c.ShouldBindJSON(&req); err != nil {
errorResp := &response.ErrorResponse{
Code: response.ErrInvalidRequest.Code,
ErrorDescription: "Invalid request format: " + err.Error(),
}
response.RespondWithError(c, response.StatusBadRequest, errorResp)
return
}
// Get authorized information
authInfo := authorized.GetInfo(c)
// Check permission (write access)
hasPermission, err := checkChatPermission(chatStore, authInfo, chatID, false)
if err != nil {
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: err.Error(),
}
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
return
}
if !hasPermission {
errorResp := &response.ErrorResponse{
Code: response.ErrAccessDenied.Code,
ErrorDescription: "Forbidden: No permission to update this chat",
}
response.RespondWithError(c, response.StatusForbidden, errorResp)
return
}
// Build updates map
updates := make(map[string]interface{})
if req.Title != nil {
updates["title"] = *req.Title
}
if req.Status != nil {
updates["status"] = *req.Status
}
if req.Metadata != nil {
updates["metadata"] = req.Metadata
}
// Add update scope
if authInfo != nil {
updates["__yao_updated_by"] = authInfo.UserID
}
if len(updates) == 0 {
errorResp := &response.ErrorResponse{
Code: response.ErrInvalidRequest.Code,
ErrorDescription: "No fields to update",
}
response.RespondWithError(c, response.StatusBadRequest, errorResp)
return
}
// Update chat
if err := chatStore.UpdateChat(chatID, updates); err != nil {
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: err.Error(),
}
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
return
}
response.RespondWithSuccess(c, response.StatusOK, gin.H{
"message": "Chat updated successfully",
"chat_id": chatID,
})
}
// DeleteChat deletes a chat session
// DELETE /v1/chat/sessions/:chat_id
func DeleteChat(c *gin.Context) {
// Get chat store
chatStore := assistant.GetChatStore()
if chatStore == nil {
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: "Chat storage not initialized",
}
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
return
}
// Get chat ID from URL parameter
chatID := c.Param("chat_id")
if chatID == "" {
errorResp := &response.ErrorResponse{
Code: response.ErrInvalidRequest.Code,
ErrorDescription: "Chat ID is required",
}
response.RespondWithError(c, response.StatusBadRequest, errorResp)
return
}
// Get authorized information
authInfo := authorized.GetInfo(c)
// Check permission (write access)
hasPermission, err := checkChatPermission(chatStore, authInfo, chatID, false)
if err != nil {
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: err.Error(),
}
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
return
}
if !hasPermission {
errorResp := &response.ErrorResponse{
Code: response.ErrAccessDenied.Code,
ErrorDescription: "Forbidden: No permission to delete this chat",
}
response.RespondWithError(c, response.StatusForbidden, errorResp)
return
}
// Delete chat
if err := chatStore.DeleteChat(chatID); err != nil {
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: err.Error(),
}
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
return
}
response.RespondWithSuccess(c, response.StatusOK, gin.H{
"message": "Chat deleted successfully",
"chat_id": chatID,
})
}
// =============================================================================
// Message Handlers
// =============================================================================
// GetMessages retrieves messages for a chat session
// GET /v1/chat/sessions/:chat_id/messages
func GetMessages(c *gin.Context) {
// Get chat store
chatStore := assistant.GetChatStore()
if chatStore == nil {
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: "Chat storage not initialized",
}
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
return
}
// Get chat ID from URL parameter
chatID := c.Param("chat_id")
if chatID == "" {
errorResp := &response.ErrorResponse{
Code: response.ErrInvalidRequest.Code,
ErrorDescription: "Chat ID is required",
}
response.RespondWithError(c, response.StatusBadRequest, errorResp)
return
}
// Get authorized information
authInfo := authorized.GetInfo(c)
// Check permission (read access)
hasPermission, err := checkChatPermission(chatStore, authInfo, chatID, true)
if err != nil {
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: err.Error(),
}
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
return
}
if !hasPermission {
errorResp := &response.ErrorResponse{
Code: response.ErrAccessDenied.Code,
ErrorDescription: "Forbidden: No permission to access this chat",
}
response.RespondWithError(c, response.StatusForbidden, errorResp)
return
}
// Build message filter
filter := buildMessageFilter(c)
// Get messages
messages, err := chatStore.GetMessages(chatID, filter)
if err != nil {
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: err.Error(),
}
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
return
}
response.RespondWithSuccess(c, response.StatusOK, gin.H{
"chat_id": chatID,
"messages": messages,
"count": len(messages),
})
}
// =============================================================================
// Helper Functions
// =============================================================================
// buildChatFilter builds ChatFilter from query parameters
func buildChatFilter(c *gin.Context, authInfo *oauthtypes.AuthorizedInfo) storetypes.ChatFilter {
filter := storetypes.ChatFilter{}
// Pagination
if pageStr := c.Query("page"); pageStr != "" {
if p, err := strconv.Atoi(pageStr); err == nil && p > 0 {
filter.Page = p
}
}
if filter.Page == 0 {
filter.Page = 1
}
if pagesizeStr := c.Query("pagesize"); pagesizeStr != "" {
if ps, err := strconv.Atoi(pagesizeStr); err == nil && ps > 0 && ps <= 100 {
filter.PageSize = ps
}
}
if filter.PageSize == 0 {
filter.PageSize = 20
}
// Business filters
filter.AssistantID = strings.TrimSpace(c.Query("assistant_id"))
filter.Status = strings.TrimSpace(c.Query("status"))
filter.Keywords = strings.TrimSpace(c.Query("keywords"))
// Time range filter
if startTimeStr := c.Query("start_time"); startTimeStr != "" {
if t, err := time.Parse(time.RFC3339, startTimeStr); err == nil {
filter.StartTime = &t
}
}
if endTimeStr := c.Query("end_time"); endTimeStr != "" {
if t, err := time.Parse(time.RFC3339, endTimeStr); err == nil {
filter.EndTime = &t
}
}
filter.TimeField = strings.TrimSpace(c.Query("time_field"))
if filter.TimeField == "" {
filter.TimeField = "last_message_at"
}
// Sorting
filter.OrderBy = strings.TrimSpace(c.Query("order_by"))
if filter.OrderBy == "" {
filter.OrderBy = "last_message_at"
}
filter.Order = strings.TrimSpace(c.Query("order"))
if filter.Order == "" {
filter.Order = "desc"
}
// Grouping
filter.GroupBy = strings.TrimSpace(c.Query("group_by"))
// Permission filters based on auth constraints
if authInfo != nil {
// Direct permission filters (AND logic)
if authInfo.Constraints.OwnerOnly {
filter.UserID = authInfo.UserID
}
if authInfo.Constraints.TeamOnly {
filter.TeamID = authInfo.TeamID
}
// For complex permission logic (OR conditions), use QueryFilter
// Example: user can see their own chats OR team shared chats
if authInfo.Constraints.TeamOnly && !authInfo.Constraints.OwnerOnly {
// Team member can see: own chats OR team shared chats
filter.QueryFilter = func(qb query.Query) {
qb.Where(func(sub query.Query) {
sub.Where("__yao_created_by", authInfo.UserID).
OrWhere(func(inner query.Query) {
inner.Where("__yao_team_id", authInfo.TeamID).
Where("share", "team")
})
})
}
// Clear direct filters since we're using QueryFilter
filter.UserID = ""
filter.TeamID = ""
}
}
return filter
}
// buildMessageFilter builds MessageFilter from query parameters
func buildMessageFilter(c *gin.Context) storetypes.MessageFilter {
filter := storetypes.MessageFilter{}
// Filter parameters
filter.RequestID = strings.TrimSpace(c.Query("request_id"))
filter.Role = strings.TrimSpace(c.Query("role"))
filter.BlockID = strings.TrimSpace(c.Query("block_id"))
filter.ThreadID = strings.TrimSpace(c.Query("thread_id"))
filter.Type = strings.TrimSpace(c.Query("type"))
// Pagination
if limitStr := c.Query("limit"); limitStr != "" {
if l, err := strconv.Atoi(limitStr); err == nil && l > 0 && l <= 1000 {
filter.Limit = l
}
}
if filter.Limit == 0 {
filter.Limit = 100
}
if offsetStr := c.Query("offset"); offsetStr != "" {
if o, err := strconv.Atoi(offsetStr); err == nil && o >= 0 {
filter.Offset = o
}
}
return filter
}
// checkChatPermission checks if the user has permission to access the chat
// readable: true for read access, false for write access
func checkChatPermission(chatStore storetypes.ChatStore, authInfo *oauthtypes.AuthorizedInfo, chatID string, readable bool) (bool, error) {
// No auth info means no constraints (for internal calls)
if authInfo == nil {
return true, nil
}
// No constraints means full access
if !authInfo.Constraints.TeamOnly && !authInfo.Constraints.OwnerOnly {
return true, nil
}
// Get chat to check permissions
chat, err := chatStore.GetChat(chatID)
if err != nil {
return false, err
}
// For read access, check if chat is public or shared with team
if readable {
if chat.Public {
return true, nil
}
if chat.Share == "team" && authInfo.Constraints.TeamOnly && chat.TeamID == authInfo.TeamID {
return true, nil
}
}
// Combined Team and Owner permission validation
if authInfo.Constraints.TeamOnly && authInfo.Constraints.OwnerOnly {
if chat.CreatedBy == authInfo.UserID && chat.TeamID == authInfo.TeamID {
return true, nil
}
return false, nil
}
// Owner only permission validation
if authInfo.Constraints.OwnerOnly && chat.CreatedBy == authInfo.UserID {
return true, nil
}
// Team only permission validation
if authInfo.Constraints.TeamOnly && chat.TeamID == authInfo.TeamID {
return true, nil
}
return false, nil
}

15
openapi/chat/types.go
View file

@ -2,9 +2,24 @@ package chat
import "github.com/yaoapp/yao/agent/context"
// =============================================================================
// Completion Types
// =============================================================================
// AppendMessagesRequest represents the request body for appending messages to running completion
type AppendMessagesRequest struct {
Type context.InterruptType `json:"type" binding:"required"` // Interrupt type: "graceful" or "force"
Messages []context.Message `json:"messages" binding:"required"`
Metadata map[string]interface{} `json:"metadata,omitempty"`
}
// =============================================================================
// Chat Session Types
// =============================================================================
// UpdateChatRequest represents the request for updating a chat session
type UpdateChatRequest struct {
Title *string `json:"title,omitempty"` // Chat title
Status *string `json:"status,omitempty"` // Status: "active" or "archived"
Metadata map[string]interface{} `json:"metadata,omitempty"` // Additional metadata
}

766
openapi/tests/chat/session_test.go Normal file
View file

@ -0,0 +1,766 @@
package openapi_test
import (
"bytes"
"encoding/json"
"fmt"
"net/http"
"testing"
"time"
"github.com/google/uuid"
"github.com/stretchr/testify/assert"
"github.com/yaoapp/yao/agent/assistant"
storetypes "github.com/yaoapp/yao/agent/store/types"
"github.com/yaoapp/yao/openapi"
"github.com/yaoapp/yao/openapi/tests/testutils"
)
// =============================================================================
// Test Setup Helpers
// =============================================================================
// createTestChat creates a test chat session in the database
func createTestChat(t *testing.T, title string, assistantID string) string {
chatStore := assistant.GetChatStore()
if chatStore == nil {
t.Skip("Chat store not initialized")
}
chatID := uuid.New().String()
chat := &storetypes.Chat{
ChatID: chatID,
AssistantID: assistantID,
Title: title,
Status: "active",
CreatedAt: time.Now(),
UpdatedAt: time.Now(),
}
err := chatStore.CreateChat(chat)
if err != nil {
t.Fatalf("Failed to create test chat: %v", err)
}
t.Logf("Created test chat: %s (title: %s)", chatID, title)
return chatID
}
// createTestMessage creates a test message in the database
func createTestMessage(t *testing.T, chatID, role, msgType, content string) string {
chatStore := assistant.GetChatStore()
if chatStore == nil {
t.Skip("Chat store not initialized")
}
msgID := uuid.New().String()
msg := &storetypes.Message{
MessageID: msgID,
ChatID: chatID,
Role: role,
Type: msgType,
Props: map[string]interface{}{
"content": content,
},
Sequence: 1,
CreatedAt: time.Now(),
}
err := chatStore.SaveMessages(chatID, []*storetypes.Message{msg})
if err != nil {
t.Fatalf("Failed to create test message: %v", err)
}
t.Logf("Created test message: %s (role: %s)", msgID, role)
return msgID
}
// cleanupTestChat deletes a test chat session
func cleanupTestChat(t *testing.T, chatID string) {
chatStore := assistant.GetChatStore()
if chatStore == nil {
return
}
err := chatStore.DeleteChat(chatID)
if err != nil {
t.Logf("Warning: Failed to cleanup test chat %s: %v", chatID, err)
} else {
t.Logf("Cleaned up test chat: %s", chatID)
}
}
// =============================================================================
// List Chat Sessions Tests
// =============================================================================
// TestListChatSessions tests the chat sessions listing endpoint
func TestListChatSessions(t *testing.T) {
serverURL := testutils.Prepare(t)
defer testutils.Clean()
// Get base URL from server config
baseURL := ""
if openapi.Server != nil && openapi.Server.Config != nil {
baseURL = openapi.Server.Config.BaseURL
}
// Register test client and get token
client := testutils.RegisterTestClient(t, "Chat Session Test Client", []string{"https://localhost/callback"})
defer testutils.CleanupTestClient(t, client.ClientID)
tokenInfo := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile")
// Create test chats
chatID1 := createTestChat(t, "Test Chat 1", "test-assistant")
defer cleanupTestChat(t, chatID1)
chatID2 := createTestChat(t, "Test Chat 2", "test-assistant")
defer cleanupTestChat(t, chatID2)
t.Run("ListChatsSuccess", func(t *testing.T) {
req, err := http.NewRequest("GET", serverURL+baseURL+"/chat/sessions", nil)
assert.NoError(t, err)
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
resp, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
assert.NotNil(t, resp)
defer resp.Body.Close()
assert.Equal(t, http.StatusOK, resp.StatusCode, "Should successfully retrieve chat sessions")
var response map[string]interface{}
err = json.NewDecoder(resp.Body).Decode(&response)
assert.NoError(t, err)
// Check response structure
assert.Contains(t, response, "data")
assert.Contains(t, response, "page")
assert.Contains(t, response, "pagesize")
assert.Contains(t, response, "total")
data, hasData := response["data"].([]interface{})
if hasData {
t.Logf("Successfully retrieved %d chat sessions", len(data))
}
})
t.Run("ListChatsWithPagination", func(t *testing.T) {
req, err := http.NewRequest("GET", serverURL+baseURL+"/chat/sessions?page=1&pagesize=10", nil)
assert.NoError(t, err)
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
resp, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
assert.NotNil(t, resp)
defer resp.Body.Close()
assert.Equal(t, http.StatusOK, resp.StatusCode)
var response map[string]interface{}
err = json.NewDecoder(resp.Body).Decode(&response)
assert.NoError(t, err)
// Verify pagination values
page, hasPage := response["page"].(float64)
pagesize, hasPagesize := response["pagesize"].(float64)
if hasPage && hasPagesize {
assert.Equal(t, float64(1), page, "Page should be 1")
assert.Equal(t, float64(10), pagesize, "Pagesize should be 10")
t.Logf("Pagination working correctly: page=%d, pagesize=%d", int(page), int(pagesize))
}
})
t.Run("ListChatsWithKeywords", func(t *testing.T) {
req, err := http.NewRequest("GET", serverURL+baseURL+"/chat/sessions?keywords=Test", nil)
assert.NoError(t, err)
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
resp, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
assert.NotNil(t, resp)
defer resp.Body.Close()
assert.Equal(t, http.StatusOK, resp.StatusCode)
var response map[string]interface{}
err = json.NewDecoder(resp.Body).Decode(&response)
assert.NoError(t, err)
data, hasData := response["data"].([]interface{})
if hasData {
t.Logf("Successfully retrieved %d chat sessions with keywords filter", len(data))
}
})
t.Run("ListChatsWithStatusFilter", func(t *testing.T) {
req, err := http.NewRequest("GET", serverURL+baseURL+"/chat/sessions?status=active", nil)
assert.NoError(t, err)
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
resp, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
assert.NotNil(t, resp)
defer resp.Body.Close()
assert.Equal(t, http.StatusOK, resp.StatusCode)
var response map[string]interface{}
err = json.NewDecoder(resp.Body).Decode(&response)
assert.NoError(t, err)
data, hasData := response["data"].([]interface{})
if hasData {
t.Logf("Successfully retrieved %d active chat sessions", len(data))
}
})
t.Run("ListChatsWithAssistantFilter", func(t *testing.T) {
req, err := http.NewRequest("GET", serverURL+baseURL+"/chat/sessions?assistant_id=test-assistant", nil)
assert.NoError(t, err)
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
resp, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
assert.NotNil(t, resp)
defer resp.Body.Close()
assert.Equal(t, http.StatusOK, resp.StatusCode)
var response map[string]interface{}
err = json.NewDecoder(resp.Body).Decode(&response)
assert.NoError(t, err)
data, hasData := response["data"].([]interface{})
if hasData {
t.Logf("Successfully retrieved %d chat sessions with assistant filter", len(data))
}
})
t.Run("ListChatsWithTimeRange", func(t *testing.T) {
startTime := time.Now().Add(-24 * time.Hour).Format(time.RFC3339)
endTime := time.Now().Add(time.Hour).Format(time.RFC3339)
req, err := http.NewRequest("GET", fmt.Sprintf("%s%s/chat/sessions?start_time=%s&end_time=%s", serverURL, baseURL, startTime, endTime), nil)
assert.NoError(t, err)
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
resp, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
assert.NotNil(t, resp)
defer resp.Body.Close()
assert.Equal(t, http.StatusOK, resp.StatusCode)
var response map[string]interface{}
err = json.NewDecoder(resp.Body).Decode(&response)
assert.NoError(t, err)
data, hasData := response["data"].([]interface{})
if hasData {
t.Logf("Successfully retrieved %d chat sessions within time range", len(data))
}
})
t.Run("ListChatsWithSorting", func(t *testing.T) {
req, err := http.NewRequest("GET", serverURL+baseURL+"/chat/sessions?order_by=created_at&order=desc", nil)
assert.NoError(t, err)
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
resp, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
assert.NotNil(t, resp)
defer resp.Body.Close()
assert.Equal(t, http.StatusOK, resp.StatusCode)
var response map[string]interface{}
err = json.NewDecoder(resp.Body).Decode(&response)
assert.NoError(t, err)
data, hasData := response["data"].([]interface{})
if hasData {
t.Logf("Successfully retrieved %d chat sessions with sorting", len(data))
}
})
t.Run("ListChatsWithGroupBy", func(t *testing.T) {
req, err := http.NewRequest("GET", serverURL+baseURL+"/chat/sessions?group_by=time", nil)
assert.NoError(t, err)
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
resp, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
assert.NotNil(t, resp)
defer resp.Body.Close()
assert.Equal(t, http.StatusOK, resp.StatusCode)
var response map[string]interface{}
err = json.NewDecoder(resp.Body).Decode(&response)
assert.NoError(t, err)
// Check for groups in response
_, hasGroups := response["groups"]
assert.True(t, hasGroups, "Response should contain groups when group_by=time")
t.Logf("Successfully retrieved chat sessions with time grouping")
})
t.Run("ListChatsUnauthorized", func(t *testing.T) {
req, err := http.NewRequest("GET", serverURL+baseURL+"/chat/sessions", nil)
assert.NoError(t, err)
// No Authorization header
resp, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
assert.NotNil(t, resp)
defer resp.Body.Close()
// Should fail without authorization
assert.NotEqual(t, http.StatusOK, resp.StatusCode, "Should fail without authorization")
})
}
// =============================================================================
// Get Chat Session Tests
// =============================================================================
// TestGetChatSession tests the get single chat session endpoint
func TestGetChatSession(t *testing.T) {
serverURL := testutils.Prepare(t)
defer testutils.Clean()
// Get base URL from server config
baseURL := ""
if openapi.Server != nil && openapi.Server.Config != nil {
baseURL = openapi.Server.Config.BaseURL
}
// Register test client and get token
client := testutils.RegisterTestClient(t, "Chat Get Test Client", []string{"https://localhost/callback"})
defer testutils.CleanupTestClient(t, client.ClientID)
tokenInfo := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile")
// Create test chat
chatID := createTestChat(t, "Test Chat for Get", "test-assistant")
defer cleanupTestChat(t, chatID)
t.Run("GetChatSuccess", func(t *testing.T) {
req, err := http.NewRequest("GET", serverURL+baseURL+"/chat/sessions/"+chatID, nil)
assert.NoError(t, err)
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
resp, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
assert.NotNil(t, resp)
defer resp.Body.Close()
assert.Equal(t, http.StatusOK, resp.StatusCode, "Should successfully retrieve chat session")
var response map[string]interface{}
err = json.NewDecoder(resp.Body).Decode(&response)
assert.NoError(t, err)
// Check response contains chat data
data, hasData := response["data"].(map[string]interface{})
if hasData {
assert.Equal(t, chatID, data["chat_id"], "Chat ID should match")
t.Logf("Successfully retrieved chat: %s", chatID)
}
})
t.Run("GetChatNotFound", func(t *testing.T) {
req, err := http.NewRequest("GET", serverURL+baseURL+"/chat/sessions/non-existent-chat-id", nil)
assert.NoError(t, err)
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
resp, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
assert.NotNil(t, resp)
defer resp.Body.Close()
assert.Equal(t, http.StatusNotFound, resp.StatusCode, "Should return 404 for non-existent chat")
})
t.Run("GetChatUnauthorized", func(t *testing.T) {
req, err := http.NewRequest("GET", serverURL+baseURL+"/chat/sessions/"+chatID, nil)
assert.NoError(t, err)
// No Authorization header
resp, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
assert.NotNil(t, resp)
defer resp.Body.Close()
assert.NotEqual(t, http.StatusOK, resp.StatusCode, "Should fail without authorization")
})
}
// =============================================================================
// Update Chat Session Tests
// =============================================================================
// TestUpdateChatSession tests the update chat session endpoint
func TestUpdateChatSession(t *testing.T) {
serverURL := testutils.Prepare(t)
defer testutils.Clean()
// Get base URL from server config
baseURL := ""
if openapi.Server != nil && openapi.Server.Config != nil {
baseURL = openapi.Server.Config.BaseURL
}
// Register test client and get token
client := testutils.RegisterTestClient(t, "Chat Update Test Client", []string{"https://localhost/callback"})
defer testutils.CleanupTestClient(t, client.ClientID)
tokenInfo := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile")
// Create test chat
chatID := createTestChat(t, "Test Chat for Update", "test-assistant")
defer cleanupTestChat(t, chatID)
t.Run("UpdateChatTitleSuccess", func(t *testing.T) {
body := map[string]interface{}{
"title": "Updated Chat Title",
}
bodyBytes, _ := json.Marshal(body)
req, err := http.NewRequest("PUT", serverURL+baseURL+"/chat/sessions/"+chatID, bytes.NewReader(bodyBytes))
assert.NoError(t, err)
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
req.Header.Set("Content-Type", "application/json")
resp, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
assert.NotNil(t, resp)
defer resp.Body.Close()
assert.Equal(t, http.StatusOK, resp.StatusCode, "Should successfully update chat title")
var response map[string]interface{}
err = json.NewDecoder(resp.Body).Decode(&response)
assert.NoError(t, err)
t.Logf("Successfully updated chat title: %s", chatID)
})
t.Run("UpdateChatStatusSuccess", func(t *testing.T) {
body := map[string]interface{}{
"status": "archived",
}
bodyBytes, _ := json.Marshal(body)
req, err := http.NewRequest("PUT", serverURL+baseURL+"/chat/sessions/"+chatID, bytes.NewReader(bodyBytes))
assert.NoError(t, err)
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
req.Header.Set("Content-Type", "application/json")
resp, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
assert.NotNil(t, resp)
defer resp.Body.Close()
assert.Equal(t, http.StatusOK, resp.StatusCode, "Should successfully update chat status")
t.Logf("Successfully updated chat status: %s", chatID)
})
t.Run("UpdateChatMetadataSuccess", func(t *testing.T) {
body := map[string]interface{}{
"metadata": map[string]interface{}{
"custom_key": "custom_value",
},
}
bodyBytes, _ := json.Marshal(body)
req, err := http.NewRequest("PUT", serverURL+baseURL+"/chat/sessions/"+chatID, bytes.NewReader(bodyBytes))
assert.NoError(t, err)
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
req.Header.Set("Content-Type", "application/json")
resp, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
assert.NotNil(t, resp)
defer resp.Body.Close()
assert.Equal(t, http.StatusOK, resp.StatusCode, "Should successfully update chat metadata")
t.Logf("Successfully updated chat metadata: %s", chatID)
})
t.Run("UpdateChatNoFields", func(t *testing.T) {
body := map[string]interface{}{}
bodyBytes, _ := json.Marshal(body)
req, err := http.NewRequest("PUT", serverURL+baseURL+"/chat/sessions/"+chatID, bytes.NewReader(bodyBytes))
assert.NoError(t, err)
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
req.Header.Set("Content-Type", "application/json")
resp, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
assert.NotNil(t, resp)
defer resp.Body.Close()
// Note: Server may still return 200 if it adds __yao_updated_by automatically
// This is acceptable behavior - the update still happens with the updater field
assert.Contains(t, []int{http.StatusOK, http.StatusBadRequest}, resp.StatusCode, "Should either succeed with auto-fields or fail with no fields")
})
t.Run("UpdateChatNotFound", func(t *testing.T) {
body := map[string]interface{}{
"title": "Updated Title",
}
bodyBytes, _ := json.Marshal(body)
req, err := http.NewRequest("PUT", serverURL+baseURL+"/chat/sessions/non-existent-chat-id", bytes.NewReader(bodyBytes))
assert.NoError(t, err)
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
req.Header.Set("Content-Type", "application/json")
resp, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
assert.NotNil(t, resp)
defer resp.Body.Close()
// Should fail for non-existent chat
assert.NotEqual(t, http.StatusOK, resp.StatusCode, "Should fail for non-existent chat")
})
t.Run("UpdateChatUnauthorized", func(t *testing.T) {
body := map[string]interface{}{
"title": "Updated Title",
}
bodyBytes, _ := json.Marshal(body)
req, err := http.NewRequest("PUT", serverURL+baseURL+"/chat/sessions/"+chatID, bytes.NewReader(bodyBytes))
assert.NoError(t, err)
req.Header.Set("Content-Type", "application/json")
// No Authorization header
resp, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
assert.NotNil(t, resp)
defer resp.Body.Close()
assert.NotEqual(t, http.StatusOK, resp.StatusCode, "Should fail without authorization")
})
}
// =============================================================================
// Delete Chat Session Tests
// =============================================================================
// TestDeleteChatSession tests the delete chat session endpoint
func TestDeleteChatSession(t *testing.T) {
serverURL := testutils.Prepare(t)
defer testutils.Clean()
// Get base URL from server config
baseURL := ""
if openapi.Server != nil && openapi.Server.Config != nil {
baseURL = openapi.Server.Config.BaseURL
}
// Register test client and get token
client := testutils.RegisterTestClient(t, "Chat Delete Test Client", []string{"https://localhost/callback"})
defer testutils.CleanupTestClient(t, client.ClientID)
tokenInfo := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile")
t.Run("DeleteChatSuccess", func(t *testing.T) {
// Create a chat to delete
chatID := createTestChat(t, "Test Chat for Delete", "test-assistant")
req, err := http.NewRequest("DELETE", serverURL+baseURL+"/chat/sessions/"+chatID, nil)
assert.NoError(t, err)
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
resp, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
assert.NotNil(t, resp)
defer resp.Body.Close()
assert.Equal(t, http.StatusOK, resp.StatusCode, "Should successfully delete chat session")
var response map[string]interface{}
err = json.NewDecoder(resp.Body).Decode(&response)
assert.NoError(t, err)
t.Logf("Successfully deleted chat: %s", chatID)
})
t.Run("DeleteChatNotFound", func(t *testing.T) {
req, err := http.NewRequest("DELETE", serverURL+baseURL+"/chat/sessions/non-existent-chat-id", nil)
assert.NoError(t, err)
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
resp, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
assert.NotNil(t, resp)
defer resp.Body.Close()
// Should fail for non-existent chat
assert.NotEqual(t, http.StatusOK, resp.StatusCode, "Should fail for non-existent chat")
})
t.Run("DeleteChatUnauthorized", func(t *testing.T) {
// Create a chat to attempt to delete
chatID := createTestChat(t, "Test Chat for Unauthorized Delete", "test-assistant")
defer cleanupTestChat(t, chatID)
req, err := http.NewRequest("DELETE", serverURL+baseURL+"/chat/sessions/"+chatID, nil)
assert.NoError(t, err)
// No Authorization header
resp, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
assert.NotNil(t, resp)
defer resp.Body.Close()
assert.NotEqual(t, http.StatusOK, resp.StatusCode, "Should fail without authorization")
})
}
// =============================================================================
// Get Messages Tests
// =============================================================================
// TestGetMessages tests the get messages endpoint
func TestGetMessages(t *testing.T) {
serverURL := testutils.Prepare(t)
defer testutils.Clean()
// Get base URL from server config
baseURL := ""
if openapi.Server != nil && openapi.Server.Config != nil {
baseURL = openapi.Server.Config.BaseURL
}
// Register test client and get token
client := testutils.RegisterTestClient(t, "Chat Messages Test Client", []string{"https://localhost/callback"})
defer testutils.CleanupTestClient(t, client.ClientID)
tokenInfo := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile")
// Create test chat with messages
chatID := createTestChat(t, "Test Chat for Messages", "test-assistant")
defer cleanupTestChat(t, chatID)
// Create test messages
createTestMessage(t, chatID, "user", "text", "Hello, how are you?")
createTestMessage(t, chatID, "assistant", "text", "I'm doing well, thank you!")
t.Run("GetMessagesSuccess", func(t *testing.T) {
req, err := http.NewRequest("GET", serverURL+baseURL+"/chat/sessions/"+chatID+"/messages", nil)
assert.NoError(t, err)
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
resp, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
assert.NotNil(t, resp)
defer resp.Body.Close()
assert.Equal(t, http.StatusOK, resp.StatusCode, "Should successfully retrieve messages")
var response map[string]interface{}
err = json.NewDecoder(resp.Body).Decode(&response)
assert.NoError(t, err)
// Check response structure
data, hasData := response["data"].(map[string]interface{})
if hasData {
messages, hasMessages := data["messages"].([]interface{})
if hasMessages {
t.Logf("Successfully retrieved %d messages", len(messages))
}
}
})
t.Run("GetMessagesWithRoleFilter", func(t *testing.T) {
req, err := http.NewRequest("GET", serverURL+baseURL+"/chat/sessions/"+chatID+"/messages?role=user", nil)
assert.NoError(t, err)
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
resp, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
assert.NotNil(t, resp)
defer resp.Body.Close()
assert.Equal(t, http.StatusOK, resp.StatusCode)
var response map[string]interface{}
err = json.NewDecoder(resp.Body).Decode(&response)
assert.NoError(t, err)
t.Logf("Successfully retrieved messages with role filter")
})
t.Run("GetMessagesWithTypeFilter", func(t *testing.T) {
req, err := http.NewRequest("GET", serverURL+baseURL+"/chat/sessions/"+chatID+"/messages?type=text", nil)
assert.NoError(t, err)
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
resp, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
assert.NotNil(t, resp)
defer resp.Body.Close()
assert.Equal(t, http.StatusOK, resp.StatusCode)
var response map[string]interface{}
err = json.NewDecoder(resp.Body).Decode(&response)
assert.NoError(t, err)
t.Logf("Successfully retrieved messages with type filter")
})
t.Run("GetMessagesWithPagination", func(t *testing.T) {
req, err := http.NewRequest("GET", serverURL+baseURL+"/chat/sessions/"+chatID+"/messages?limit=10&offset=0", nil)
assert.NoError(t, err)
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
resp, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
assert.NotNil(t, resp)
defer resp.Body.Close()
assert.Equal(t, http.StatusOK, resp.StatusCode)
var response map[string]interface{}
err = json.NewDecoder(resp.Body).Decode(&response)
assert.NoError(t, err)
t.Logf("Successfully retrieved messages with pagination")
})
t.Run("GetMessagesNotFound", func(t *testing.T) {
req, err := http.NewRequest("GET", serverURL+baseURL+"/chat/sessions/non-existent-chat-id/messages", nil)
assert.NoError(t, err)
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
resp, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
assert.NotNil(t, resp)
defer resp.Body.Close()
// For non-existent chat, the API may return:
// - 200 with empty messages (if permission check passes first)
// - 403 Forbidden (if permission check fails on non-existent chat)
// - 404 Not Found (if explicitly checking chat existence)
// All are acceptable behaviors depending on implementation
t.Logf("Response status for non-existent chat messages: %d", resp.StatusCode)
})
t.Run("GetMessagesUnauthorized", func(t *testing.T) {
req, err := http.NewRequest("GET", serverURL+baseURL+"/chat/sessions/"+chatID+"/messages", nil)
assert.NoError(t, err)
// No Authorization header
resp, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
assert.NotNil(t, resp)
defer resp.Body.Close()
assert.NotEqual(t, http.StatusOK, resp.StatusCode, "Should fail without authorization")
})
}

10
yao/models/agent/chat.mod.yao
View file

@ -40,6 +40,15 @@
"nullable": false,
"index": true
},
{
"name": "last_connector",
"type": "string",
"label": "Last Connector",
"comment": "Last used connector ID (updated on each message)",
"length": 200,
"nullable": true,
"index": true
},
{
"name": "mode",
"type": "string",
@ -122,4 +131,3 @@
},
"option": { "timestamps": true, "soft_deletes": true, "permission": true }
}

23
yao/models/agent/message.mod.yao
View file

@ -18,10 +18,9 @@
"name": "message_id",
"type": "string",
"label": "Message ID",
"comment": "Unique message identifier",
"comment": "Message identifier (unique within request)",
"length": 64,
"nullable": false,
"unique": true
"nullable": false
},
{
"name": "chat_id",
@ -92,6 +91,15 @@
"nullable": true,
"index": true
},
{
"name": "connector",
"type": "string",
"label": "Connector",
"comment": "Connector ID used for this message",
"length": 200,
"nullable": true,
"index": true
},
{
"name": "sequence",
"type": "integer",
@ -127,8 +135,13 @@
"columns": ["chat_id", "sequence"],
"type": "index",
"comment": "Index for message ordering within chat"
},
{
"name": "idx_msg_request_message",
"columns": ["request_id", "message_id"],
"type": "unique",
"comment": "Unique constraint for message_id within request"
}
],
"option": { "timestamps": true, "soft_deletes": false }
"option": { "timestamps": true, "soft_deletes": true }
}

2
yao/models/agent/resume.mod.yao
View file

@ -166,5 +166,5 @@
"comment": "Index for resume ordering within request"
}
],
"option": { "timestamps": true, "soft_deletes": false }
"option": { "timestamps": true, "soft_deletes": true }
}