diff --git a/agent/assistant/agent.go b/agent/assistant/agent.go index 4d655f63..b3411118 100644 --- a/agent/assistant/agent.go +++ b/agent/assistant/agent.go @@ -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 diff --git a/agent/assistant/chat.go b/agent/assistant/chat.go index 2091507b..e8166ecd 100644 --- a/agent/assistant/chat.go +++ b/agent/assistant/chat.go @@ -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,261 @@ 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() + } + + ctx.Buffer = agentcontext.NewChatBuffer(ctx.ChatID, requestID, ast.ID) + log.Trace("[CHAT] Buffer initialized: chatID=%s, requestID=%s, assistantID=%s", ctx.ChatID, requestID, ast.ID) +} + +// 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 + if len(messages) > 0 { + now := time.Now() + if updateErr := chatStore.UpdateChat(ctx.ChatID, map[string]interface{}{ + "last_message_at": now, + }); updateErr != nil { + log.Trace("[CHAT] Failed to update last_message_at: %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, + 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 + } + + 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 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 diff --git a/agent/assistant/chat_test.go b/agent/assistant/chat_test.go index 7fb02da3..d15855d7 100644 --- a/agent/assistant/chat_test.go +++ b/agent/assistant/chat_test.go @@ -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,566 @@ 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 + ctx.Buffer.AddUserInput("Test question", "") + ctx.Buffer.AddAssistantMessage("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) + }) +} + +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") + }) +} diff --git a/agent/assistant/handlers/stream.go b/agent/assistant/handlers/stream.go index 54111605..36bf2c21 100644 --- a/agent/assistant/handlers/stream.go +++ b/agent/assistant/handlers/stream.go @@ -323,6 +323,57 @@ 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( + msgType, + props, + blockID, + threadID, + assistantID, + nil, + ) + } + // Build EventMessageEndData with complete content endData := message.EventMessageEndData{ MessageID: s.currentGroupID, // Use the message ID diff --git a/agent/context/buffer.go b/agent/context/buffer.go new file mode 100644 index 00000000..d095ad6e --- /dev/null +++ b/agent/context/buffer.go @@ -0,0 +1,359 @@ +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 + + // 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"` + Sequence int `json:"sequence"` + Metadata map[string]interface{} `json:"metadata,omitempty"` + CreatedAt time.Time `json:"created_at"` +} + +// 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 string) *ChatBuffer { + return &ChatBuffer{ + chatID: chatID, + requestID: requestID, + assistantID: assistantID, + 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(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{ + Role: "assistant", + Type: msgType, + Props: props, + BlockID: blockID, + ThreadID: threadID, + AssistantID: assistantID, + Metadata: metadata, + }) +} + +// 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 +} + +// ============================================================================= +// 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 +} diff --git a/agent/context/buffer_test.go b/agent/context/buffer_test.go new file mode 100644 index 00000000..af226bbb --- /dev/null +++ b/agent/context/buffer_test.go @@ -0,0 +1,1074 @@ +package context_test + +import ( + "fmt" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/yaoapp/yao/agent/context" +) + +// ============================================================================= +// ChatBuffer Creation Tests +// ============================================================================= + +func TestBufferNewChatBuffer(t *testing.T) { + t.Run("CreateWithAllFields", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-123", "req-456", "assistant-789") + + assert.NotNil(t, buffer) + assert.Equal(t, "chat-123", buffer.ChatID()) + assert.Equal(t, "req-456", buffer.RequestID()) + assert.Equal(t, "assistant-789", buffer.AssistantID()) + assert.Empty(t, buffer.GetMessages()) + assert.Empty(t, buffer.GetAllSteps()) + assert.Equal(t, 0, buffer.GetMessageCount()) + }) + + t.Run("CreateWithEmptyFields", func(t *testing.T) { + buffer := context.NewChatBuffer("", "", "") + + assert.NotNil(t, buffer) + assert.Empty(t, buffer.ChatID()) + assert.Empty(t, buffer.RequestID()) + assert.Empty(t, buffer.AssistantID()) + }) +} + +// ============================================================================= +// Message Buffer Tests +// ============================================================================= + +func TestBufferAddMessage(t *testing.T) { + t.Run("AddSingleMessage", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1") + + msg := &context.BufferedMessage{ + Role: "assistant", + Type: "text", + Props: map[string]interface{}{"content": "Hello"}, + } + buffer.AddMessage(msg) + + messages := buffer.GetMessages() + require.Len(t, messages, 1) + assert.Equal(t, "assistant", messages[0].Role) + assert.Equal(t, "text", messages[0].Type) + assert.Equal(t, 1, messages[0].Sequence) + assert.NotEmpty(t, messages[0].MessageID) // Auto-generated + assert.Equal(t, "chat-1", messages[0].ChatID) + assert.Equal(t, "req-1", messages[0].RequestID) + assert.False(t, messages[0].CreatedAt.IsZero()) + }) + + t.Run("AddMultipleMessages", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2") + + for i := 0; i < 5; i++ { + buffer.AddMessage(&context.BufferedMessage{ + Role: "assistant", + Type: "text", + Props: map[string]interface{}{"content": fmt.Sprintf("Message %d", i+1)}, + }) + } + + messages := buffer.GetMessages() + require.Len(t, messages, 5) + + // Verify sequence numbers + for i, msg := range messages { + assert.Equal(t, i+1, msg.Sequence) + } + }) + + t.Run("AddNilMessage", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3") + buffer.AddMessage(nil) + + assert.Equal(t, 0, buffer.GetMessageCount()) + }) + + t.Run("AddMessageWithExistingID", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-4", "req-4", "assistant-4") + + msg := &context.BufferedMessage{ + MessageID: "custom-id-123", + Role: "assistant", + Type: "text", + } + buffer.AddMessage(msg) + + messages := buffer.GetMessages() + require.Len(t, messages, 1) + assert.Equal(t, "custom-id-123", messages[0].MessageID) // Preserved + }) + + t.Run("AddMessageWithExistingTimestamp", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-5", "req-5", "assistant-5") + + customTime := time.Date(2024, 1, 1, 12, 0, 0, 0, time.UTC) + msg := &context.BufferedMessage{ + Role: "assistant", + Type: "text", + CreatedAt: customTime, + } + buffer.AddMessage(msg) + + messages := buffer.GetMessages() + require.Len(t, messages, 1) + assert.Equal(t, customTime, messages[0].CreatedAt) // Preserved + }) +} + +func TestBufferAddUserInput(t *testing.T) { + t.Run("AddStringContent", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1") + buffer.AddUserInput("What is the weather?", "") + + messages := buffer.GetMessages() + require.Len(t, messages, 1) + assert.Equal(t, "user", messages[0].Role) + assert.Equal(t, "user_input", messages[0].Type) + assert.Equal(t, "What is the weather?", messages[0].Props["content"]) + assert.Equal(t, "user", messages[0].Props["role"]) + }) + + t.Run("AddUserInputWithName", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2") + buffer.AddUserInput("Hello", "John") + + messages := buffer.GetMessages() + require.Len(t, messages, 1) + assert.Equal(t, "John", messages[0].Props["name"]) + }) + + t.Run("AddComplexContent", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3") + complexContent := []map[string]interface{}{ + {"type": "text", "text": "Look at this image"}, + {"type": "image_url", "image_url": map[string]string{"url": "https://example.com/image.jpg"}}, + } + buffer.AddUserInput(complexContent, "") + + messages := buffer.GetMessages() + require.Len(t, messages, 1) + content, ok := messages[0].Props["content"].([]map[string]interface{}) + require.True(t, ok) + assert.Len(t, content, 2) + }) +} + +func TestBufferAddAssistantMessage(t *testing.T) { + t.Run("AddTextMessage", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1") + buffer.AddAssistantMessage( + "text", + map[string]interface{}{"content": "Hello, how can I help?"}, + "block-1", + "thread-1", + "assistant-1", + map[string]interface{}{"model": "gpt-4"}, + ) + + messages := buffer.GetMessages() + require.Len(t, messages, 1) + assert.Equal(t, "assistant", messages[0].Role) + assert.Equal(t, "text", messages[0].Type) + assert.Equal(t, "block-1", messages[0].BlockID) + assert.Equal(t, "thread-1", messages[0].ThreadID) + assert.Equal(t, "assistant-1", messages[0].AssistantID) + assert.Equal(t, "gpt-4", messages[0].Metadata["model"]) + }) + + t.Run("SkipEventMessage", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2") + buffer.AddAssistantMessage( + "event", + map[string]interface{}{"event": "message_start"}, + "", "", "", nil, + ) + + // Event messages should be skipped + assert.Equal(t, 0, buffer.GetMessageCount()) + }) + + t.Run("AddRetrievalMessage", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3") + buffer.AddAssistantMessage( + "retrieval", + map[string]interface{}{ + "sources": []map[string]interface{}{ + {"title": "Doc 1", "score": 0.95}, + {"title": "Doc 2", "score": 0.87}, + }, + }, + "block-1", "", "assistant-3", nil, + ) + + messages := buffer.GetMessages() + require.Len(t, messages, 1) + assert.Equal(t, "retrieval", messages[0].Type) + }) + + t.Run("AddToolCallMessage", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-4", "req-4", "assistant-4") + buffer.AddAssistantMessage( + "tool_call", + map[string]interface{}{ + "name": "get_weather", + "arguments": `{"location": "San Francisco"}`, + }, + "block-1", "", "assistant-4", nil, + ) + + messages := buffer.GetMessages() + require.Len(t, messages, 1) + assert.Equal(t, "tool_call", messages[0].Type) + assert.Equal(t, "get_weather", messages[0].Props["name"]) + }) + + t.Run("AddCustomTypeMessage", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-5", "req-5", "assistant-5") + buffer.AddAssistantMessage( + "custom_chart", + map[string]interface{}{ + "chart_type": "bar", + "data": []int{1, 2, 3, 4, 5}, + }, + "block-1", "", "assistant-5", nil, + ) + + messages := buffer.GetMessages() + require.Len(t, messages, 1) + assert.Equal(t, "custom_chart", messages[0].Type) + }) +} + +func TestBufferGetMessages(t *testing.T) { + t.Run("GetMessagesReturnsSliceCopy", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1") + buffer.AddUserInput("Hello", "") + + messages1 := buffer.GetMessages() + messages2 := buffer.GetMessages() + + // Slices should be different (copy of slice) + // But pointers point to same underlying objects (shallow copy) + assert.Len(t, messages1, 1) + assert.Len(t, messages2, 1) + }) + + t.Run("GetEmptyMessages", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2") + messages := buffer.GetMessages() + + assert.NotNil(t, messages) + assert.Empty(t, messages) + }) +} + +func TestBufferGetMessageCount(t *testing.T) { + buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1") + assert.Equal(t, 0, buffer.GetMessageCount()) + + buffer.AddUserInput("Message 1", "") + assert.Equal(t, 1, buffer.GetMessageCount()) + + buffer.AddAssistantMessage("text", map[string]interface{}{"content": "Reply"}, "", "", "", nil) + assert.Equal(t, 2, buffer.GetMessageCount()) +} + +// ============================================================================= +// Step Buffer Tests (for Resume) +// ============================================================================= + +func TestBufferBeginStep(t *testing.T) { + t.Run("BeginStepWithStack", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1") + + stack := &context.Stack{ + ID: "stack-123", + ParentID: "stack-parent-456", + Depth: 2, + } + + step := buffer.BeginStep(context.StepTypeLLM, map[string]interface{}{"prompt": "Hello"}, stack) + + require.NotNil(t, step) + assert.NotEmpty(t, step.ResumeID) + assert.Equal(t, "chat-1", step.ChatID) + assert.Equal(t, "req-1", step.RequestID) + assert.Equal(t, "assistant-1", step.AssistantID) + assert.Equal(t, "stack-123", step.StackID) + assert.Equal(t, "stack-parent-456", step.StackParentID) + assert.Equal(t, 2, step.StackDepth) + assert.Equal(t, context.StepTypeLLM, step.Type) + assert.Equal(t, context.StepStatusRunning, step.Status) + assert.Equal(t, 1, step.Sequence) + assert.Equal(t, "Hello", step.Input["prompt"]) + assert.False(t, step.CreatedAt.IsZero()) + }) + + t.Run("BeginStepWithNilStack", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2") + + step := buffer.BeginStep(context.StepTypeInput, nil, nil) + + require.NotNil(t, step) + assert.Empty(t, step.StackID) + assert.Empty(t, step.StackParentID) + assert.Equal(t, 0, step.StackDepth) + }) + + t.Run("BeginMultipleSteps", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3") + + step1 := buffer.BeginStep(context.StepTypeInput, nil, nil) + step2 := buffer.BeginStep(context.StepTypeHookCreate, nil, nil) + step3 := buffer.BeginStep(context.StepTypeLLM, nil, nil) + + assert.Equal(t, 1, step1.Sequence) + assert.Equal(t, 2, step2.Sequence) + assert.Equal(t, 3, step3.Sequence) + + steps := buffer.GetAllSteps() + require.Len(t, steps, 3) + }) + + t.Run("BeginStepWithSpaceSnapshot", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-4", "req-4", "assistant-4") + + // Set space snapshot before beginning step + buffer.SetSpaceSnapshot(map[string]interface{}{ + "key1": "value1", + "key2": 42, + }) + + step := buffer.BeginStep(context.StepTypeLLM, nil, nil) + + require.NotNil(t, step.SpaceSnapshot) + assert.Equal(t, "value1", step.SpaceSnapshot["key1"]) + assert.Equal(t, 42, step.SpaceSnapshot["key2"]) + }) +} + +func TestBufferCompleteStep(t *testing.T) { + t.Run("CompleteCurrentStep", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1") + + buffer.BeginStep(context.StepTypeLLM, map[string]interface{}{"prompt": "Hello"}, nil) + buffer.CompleteStep(map[string]interface{}{"response": "Hi there!"}) + + steps := buffer.GetAllSteps() + require.Len(t, steps, 1) + assert.Equal(t, context.StepStatusCompleted, steps[0].Status) + assert.Equal(t, "Hi there!", steps[0].Output["response"]) + assert.Nil(t, buffer.GetCurrentStep()) // Current step cleared + }) + + t.Run("CompleteWithNoCurrentStep", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2") + + // Should not panic + buffer.CompleteStep(map[string]interface{}{"response": "test"}) + assert.Nil(t, buffer.GetCurrentStep()) + }) + + t.Run("CompleteMultipleStepsSequentially", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3") + + buffer.BeginStep(context.StepTypeInput, nil, nil) + buffer.CompleteStep(map[string]interface{}{"done": true}) + + buffer.BeginStep(context.StepTypeHookCreate, nil, nil) + buffer.CompleteStep(map[string]interface{}{"hook_result": "ok"}) + + buffer.BeginStep(context.StepTypeLLM, nil, nil) + buffer.CompleteStep(map[string]interface{}{"llm_response": "hello"}) + + steps := buffer.GetAllSteps() + require.Len(t, steps, 3) + for _, step := range steps { + assert.Equal(t, context.StepStatusCompleted, step.Status) + } + }) +} + +func TestBufferFailCurrentStep(t *testing.T) { + t.Run("FailWithError", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1") + + buffer.BeginStep(context.StepTypeLLM, nil, nil) + buffer.FailCurrentStep(context.ResumeStatusFailed, fmt.Errorf("API error: rate limit exceeded")) + + steps := buffer.GetAllSteps() + require.Len(t, steps, 1) + assert.Equal(t, context.ResumeStatusFailed, steps[0].Status) + assert.Equal(t, "API error: rate limit exceeded", steps[0].Error) + }) + + t.Run("FailWithInterrupted", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2") + + buffer.BeginStep(context.StepTypeLLM, nil, nil) + buffer.FailCurrentStep(context.ResumeStatusInterrupted, nil) + + steps := buffer.GetAllSteps() + require.Len(t, steps, 1) + assert.Equal(t, context.ResumeStatusInterrupted, steps[0].Status) + assert.Empty(t, steps[0].Error) + }) + + t.Run("FailAlreadyCompletedStep", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3") + + buffer.BeginStep(context.StepTypeLLM, nil, nil) + buffer.CompleteStep(map[string]interface{}{"done": true}) + + // Try to fail completed step (should be no-op since currentStep is nil) + buffer.FailCurrentStep(context.ResumeStatusFailed, fmt.Errorf("late error")) + + steps := buffer.GetAllSteps() + require.Len(t, steps, 1) + assert.Equal(t, context.StepStatusCompleted, steps[0].Status) // Still completed + }) + + t.Run("FailWithNoCurrentStep", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-4", "req-4", "assistant-4") + + // Should not panic + buffer.FailCurrentStep(context.ResumeStatusFailed, fmt.Errorf("error")) + }) +} + +func TestBufferGetCurrentStep(t *testing.T) { + t.Run("NoCurrentStep", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1") + assert.Nil(t, buffer.GetCurrentStep()) + }) + + t.Run("HasCurrentStep", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2") + buffer.BeginStep(context.StepTypeLLM, nil, nil) + + current := buffer.GetCurrentStep() + require.NotNil(t, current) + assert.Equal(t, context.StepTypeLLM, current.Type) + }) + + t.Run("CurrentStepClearedAfterComplete", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3") + buffer.BeginStep(context.StepTypeLLM, nil, nil) + buffer.CompleteStep(nil) + + assert.Nil(t, buffer.GetCurrentStep()) + }) +} + +func TestBufferGetStepsForResume(t *testing.T) { + t.Run("CompletedSuccessfully", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1") + + buffer.BeginStep(context.StepTypeInput, nil, nil) + buffer.CompleteStep(nil) + buffer.BeginStep(context.StepTypeLLM, nil, nil) + buffer.CompleteStep(nil) + + // Completed successfully - no steps need to be saved + steps := buffer.GetStepsForResume(context.StepStatusCompleted) + assert.Nil(t, steps) + }) + + t.Run("FailedRequest", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2") + + buffer.BeginStep(context.StepTypeInput, nil, nil) + buffer.CompleteStep(nil) + buffer.BeginStep(context.StepTypeLLM, nil, nil) + // Step still running when failure occurs + + steps := buffer.GetStepsForResume(context.ResumeStatusFailed) + require.NotNil(t, steps) + assert.Len(t, steps, 2) + + // Current step should be marked as failed + assert.Equal(t, context.ResumeStatusFailed, steps[1].Status) + }) + + t.Run("InterruptedRequest", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3") + + buffer.BeginStep(context.StepTypeInput, nil, nil) + buffer.CompleteStep(nil) + buffer.BeginStep(context.StepTypeHookCreate, nil, nil) + buffer.CompleteStep(nil) + buffer.BeginStep(context.StepTypeLLM, nil, nil) + // Interrupted during LLM + + steps := buffer.GetStepsForResume(context.ResumeStatusInterrupted) + require.NotNil(t, steps) + assert.Len(t, steps, 3) + assert.Equal(t, context.ResumeStatusInterrupted, steps[2].Status) + }) +} + +func TestBufferGetAllSteps(t *testing.T) { + t.Run("GetStepsReturnsSliceCopy", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1") + buffer.BeginStep(context.StepTypeLLM, nil, nil) + + steps1 := buffer.GetAllSteps() + steps2 := buffer.GetAllSteps() + + // Slices should be different (copy of slice) + assert.Len(t, steps1, 1) + assert.Len(t, steps2, 1) + }) + + t.Run("GetEmptySteps", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2") + steps := buffer.GetAllSteps() + + assert.NotNil(t, steps) + assert.Empty(t, steps) + }) +} + +// ============================================================================= +// Space Snapshot Tests +// ============================================================================= + +func TestBufferSpaceSnapshot(t *testing.T) { + t.Run("SetAndGetSnapshot", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1") + + snapshot := map[string]interface{}{ + "user_id": "user-123", + "session": map[string]interface{}{"token": "abc"}, + "counter": 42, + "is_active": true, + } + buffer.SetSpaceSnapshot(snapshot) + + retrieved := buffer.GetSpaceSnapshot() + assert.Equal(t, "user-123", retrieved["user_id"]) + assert.Equal(t, 42, retrieved["counter"]) + assert.Equal(t, true, retrieved["is_active"]) + }) + + t.Run("SnapshotIsCopy", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2") + + original := map[string]interface{}{"key": "original"} + buffer.SetSpaceSnapshot(original) + + // Modify original + original["key"] = "modified" + + // Buffer should have original value + retrieved := buffer.GetSpaceSnapshot() + assert.Equal(t, "original", retrieved["key"]) + }) + + t.Run("GetSnapshotReturnsCopy", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3") + buffer.SetSpaceSnapshot(map[string]interface{}{"key": "value"}) + + retrieved1 := buffer.GetSpaceSnapshot() + retrieved1["key"] = "modified" + + retrieved2 := buffer.GetSpaceSnapshot() + assert.Equal(t, "value", retrieved2["key"]) // Original unchanged + }) + + t.Run("GetNilSnapshot", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-4", "req-4", "assistant-4") + snapshot := buffer.GetSpaceSnapshot() + assert.Nil(t, snapshot) + }) + + t.Run("SetNilSnapshot", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-5", "req-5", "assistant-5") + buffer.SetSpaceSnapshot(map[string]interface{}{"key": "value"}) + buffer.SetSpaceSnapshot(nil) + + snapshot := buffer.GetSpaceSnapshot() + assert.Nil(t, snapshot) + }) +} + +// ============================================================================= +// Identity Methods Tests +// ============================================================================= + +func TestBufferIdentityMethods(t *testing.T) { + t.Run("SetAssistantID", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-original") + + assert.Equal(t, "assistant-original", buffer.AssistantID()) + + buffer.SetAssistantID("assistant-new") + assert.Equal(t, "assistant-new", buffer.AssistantID()) + }) + + t.Run("ChatID", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-test", "req-test", "assistant-test") + assert.Equal(t, "chat-test", buffer.ChatID()) + }) + + t.Run("RequestID", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-test", "req-test", "assistant-test") + assert.Equal(t, "req-test", buffer.RequestID()) + }) +} + +// ============================================================================= +// Concurrency Tests +// ============================================================================= + +func TestBufferConcurrentMessageOperations(t *testing.T) { + buffer := context.NewChatBuffer("chat-concurrent", "req-concurrent", "assistant-concurrent") + + var wg sync.WaitGroup + numGoroutines := 100 + + // Concurrent writes + for i := 0; i < numGoroutines; i++ { + wg.Add(1) + go func(idx int) { + defer wg.Done() + buffer.AddMessage(&context.BufferedMessage{ + Role: "assistant", + Type: "text", + Props: map[string]interface{}{"content": fmt.Sprintf("Message %d", idx)}, + }) + }(i) + } + + wg.Wait() + + // Verify all messages were added + messages := buffer.GetMessages() + assert.Len(t, messages, numGoroutines) + + // Verify sequences are unique + sequences := make(map[int]bool) + for _, msg := range messages { + assert.False(t, sequences[msg.Sequence], "Duplicate sequence found: %d", msg.Sequence) + sequences[msg.Sequence] = true + } +} + +func TestBufferConcurrentStepOperations(t *testing.T) { + buffer := context.NewChatBuffer("chat-concurrent", "req-concurrent", "assistant-concurrent") + + var wg sync.WaitGroup + numGoroutines := 50 + + // Concurrent step operations + for i := 0; i < numGoroutines; i++ { + wg.Add(1) + go func(idx int) { + defer wg.Done() + buffer.BeginStep(context.StepTypeLLM, map[string]interface{}{"idx": idx}, nil) + time.Sleep(time.Millisecond) // Simulate some work + buffer.CompleteStep(map[string]interface{}{"result": idx}) + }(i) + } + + wg.Wait() + + // Verify all steps were recorded + steps := buffer.GetAllSteps() + assert.Len(t, steps, numGoroutines) +} + +func TestBufferConcurrentReadWrite(t *testing.T) { + buffer := context.NewChatBuffer("chat-rw", "req-rw", "assistant-rw") + + var wg sync.WaitGroup + done := make(chan bool) + + // Writer goroutine + wg.Add(1) + go func() { + defer wg.Done() + for i := 0; i < 100; i++ { + buffer.AddMessage(&context.BufferedMessage{ + Role: "assistant", + Type: "text", + Props: map[string]interface{}{"content": fmt.Sprintf("Message %d", i)}, + }) + time.Sleep(time.Microsecond) + } + }() + + // Reader goroutine + wg.Add(1) + go func() { + defer wg.Done() + for { + select { + case <-done: + return + default: + _ = buffer.GetMessages() + _ = buffer.GetMessageCount() + time.Sleep(time.Microsecond) + } + } + }() + + // Let it run for a bit + time.Sleep(50 * time.Millisecond) + close(done) + wg.Wait() + + // Should complete without race conditions + assert.Equal(t, 100, buffer.GetMessageCount()) +} + +// ============================================================================= +// Step Type Constants Tests +// ============================================================================= + +func TestBufferStepTypeConstants(t *testing.T) { + // Verify all step types are defined + assert.Equal(t, "input", context.StepTypeInput) + assert.Equal(t, "hook_create", context.StepTypeHookCreate) + assert.Equal(t, "llm", context.StepTypeLLM) + assert.Equal(t, "tool", context.StepTypeTool) + assert.Equal(t, "hook_next", context.StepTypeHookNext) + assert.Equal(t, "delegate", context.StepTypeDelegate) +} + +func TestBufferResumeStatusConstants(t *testing.T) { + assert.Equal(t, "failed", context.ResumeStatusFailed) + assert.Equal(t, "interrupted", context.ResumeStatusInterrupted) +} + +func TestBufferStepStatusConstants(t *testing.T) { + assert.Equal(t, "running", context.StepStatusRunning) + assert.Equal(t, "completed", context.StepStatusCompleted) +} + +// ============================================================================= +// Edge Cases and Error Handling Tests +// ============================================================================= + +func TestBufferEdgeCases(t *testing.T) { + t.Run("LargeNumberOfMessages", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-large", "req-large", "assistant-large") + + // Add 10000 messages + for i := 0; i < 10000; i++ { + buffer.AddMessage(&context.BufferedMessage{ + Role: "assistant", + Type: "text", + Props: map[string]interface{}{"content": fmt.Sprintf("Message %d", i)}, + }) + } + + assert.Equal(t, 10000, buffer.GetMessageCount()) + messages := buffer.GetMessages() + assert.Len(t, messages, 10000) + }) + + t.Run("MessageWithEmptyProps", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-empty", "req-empty", "assistant-empty") + + buffer.AddMessage(&context.BufferedMessage{ + Role: "assistant", + Type: "text", + Props: nil, + }) + + messages := buffer.GetMessages() + require.Len(t, messages, 1) + assert.Nil(t, messages[0].Props) + }) + + t.Run("StepWithEmptyInput", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-step", "req-step", "assistant-step") + + step := buffer.BeginStep(context.StepTypeLLM, nil, nil) + assert.Nil(t, step.Input) + + buffer.CompleteStep(nil) + steps := buffer.GetAllSteps() + assert.Nil(t, steps[0].Output) + }) + + t.Run("AllMessageTypes", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-types", "req-types", "assistant-types") + + messageTypes := []string{ + "text", "image", "loading", "tool_call", "tool_result", + "retrieval", "thinking", "action", "chart", "table", + "custom_type_1", "custom_type_2", + } + + for _, msgType := range messageTypes { + buffer.AddAssistantMessage(msgType, map[string]interface{}{"type": msgType}, "", "", "", nil) + } + + assert.Equal(t, len(messageTypes), buffer.GetMessageCount()) + }) + + t.Run("AllStepTypes", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-step-types", "req-step-types", "assistant-step-types") + + stepTypes := []string{ + context.StepTypeInput, context.StepTypeHookCreate, context.StepTypeLLM, + context.StepTypeTool, context.StepTypeHookNext, context.StepTypeDelegate, + } + + for _, stepType := range stepTypes { + buffer.BeginStep(stepType, nil, nil) + buffer.CompleteStep(nil) + } + + steps := buffer.GetAllSteps() + assert.Len(t, steps, len(stepTypes)) + }) +} + +// ============================================================================= +// Integration-like Tests (Simulating Real Workflow) +// ============================================================================= + +func TestBufferCompleteWorkflow(t *testing.T) { + t.Run("SuccessfulChatFlow", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-workflow", "req-workflow", "assistant-main") + + // 1. User input + buffer.AddUserInput("What's the weather in San Francisco?", "John") + buffer.BeginStep(context.StepTypeInput, map[string]interface{}{"content": "What's the weather in San Francisco?"}, nil) + buffer.CompleteStep(nil) + + // 2. Create hook + buffer.BeginStep(context.StepTypeHookCreate, nil, nil) + buffer.AddAssistantMessage("thinking", map[string]interface{}{"content": "Processing your request..."}, "block-1", "", "assistant-main", nil) + buffer.CompleteStep(nil) + + // 3. LLM call with tool + buffer.BeginStep(context.StepTypeLLM, map[string]interface{}{"model": "gpt-4"}, nil) + buffer.AddAssistantMessage("tool_call", map[string]interface{}{ + "name": "get_weather", + "arguments": `{"location":"San Francisco"}`, + }, "block-2", "", "assistant-main", nil) + buffer.CompleteStep(map[string]interface{}{"tool_calls": 1}) + + // 4. Tool execution + buffer.BeginStep(context.StepTypeTool, map[string]interface{}{"tool": "get_weather"}, nil) + buffer.AddAssistantMessage("tool_result", map[string]interface{}{ + "result": "72°F, Sunny", + }, "block-2", "", "assistant-main", nil) + buffer.CompleteStep(map[string]interface{}{"result": "72°F, Sunny"}) + + // 5. Final LLM response + buffer.BeginStep(context.StepTypeLLM, nil, nil) + buffer.AddAssistantMessage("text", map[string]interface{}{ + "content": "The weather in San Francisco is currently 72°F and sunny.", + }, "block-3", "", "assistant-main", nil) + buffer.CompleteStep(nil) + + // Verify: 1 user_input + 4 assistant messages (thinking, tool_call, tool_result, text) + assert.Equal(t, 5, buffer.GetMessageCount()) + assert.Len(t, buffer.GetAllSteps(), 5) // 5 steps (no hook_next in this flow) + + // All steps should be completed + steps := buffer.GetStepsForResume(context.StepStatusCompleted) + assert.Nil(t, steps) + }) + + t.Run("InterruptedChatFlow", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-interrupted", "req-interrupted", "assistant-main") + + // Set space snapshot + buffer.SetSpaceSnapshot(map[string]interface{}{ + "user_context": "previous conversation", + "preferences": map[string]interface{}{"language": "en"}, + }) + + // 1. User input + buffer.AddUserInput("Generate a long story", "") + buffer.BeginStep(context.StepTypeInput, nil, nil) + buffer.CompleteStep(nil) + + // 2. LLM starts generating + buffer.BeginStep(context.StepTypeLLM, map[string]interface{}{"model": "gpt-4"}, nil) + buffer.AddAssistantMessage("text", map[string]interface{}{"content": "Once upon a time..."}, "block-1", "", "assistant-main", nil) + // User interrupts here! + + // Get steps for resume + steps := buffer.GetStepsForResume(context.ResumeStatusInterrupted) + require.NotNil(t, steps) + assert.Len(t, steps, 2) + + // Last step should be interrupted with space snapshot + lastStep := steps[len(steps)-1] + assert.Equal(t, context.ResumeStatusInterrupted, lastStep.Status) + assert.NotNil(t, lastStep.SpaceSnapshot) + assert.Equal(t, "previous conversation", lastStep.SpaceSnapshot["user_context"]) + }) + + t.Run("A2ACallWithDelegation", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-a2a", "req-a2a", "assistant-main") + + mainStack := &context.Stack{ID: "stack-main", Depth: 0} + childStack := &context.Stack{ID: "stack-child", ParentID: "stack-main", Depth: 1} + + // Main assistant starts + buffer.BeginStep(context.StepTypeInput, nil, mainStack) + buffer.CompleteStep(nil) + + // Delegate to child assistant + buffer.SetAssistantID("assistant-child") + buffer.BeginStep(context.StepTypeDelegate, map[string]interface{}{"delegate_to": "assistant-child"}, childStack) + + // Child assistant messages + buffer.AddAssistantMessage("text", map[string]interface{}{"content": "Child assistant responding"}, "block-child", "", "assistant-child", nil) + buffer.CompleteStep(map[string]interface{}{"delegate_result": "success"}) + + // Return to main assistant + buffer.SetAssistantID("assistant-main") + buffer.BeginStep(context.StepTypeLLM, nil, mainStack) + buffer.AddAssistantMessage("text", map[string]interface{}{"content": "Main assistant continuing"}, "block-main", "", "assistant-main", nil) + buffer.CompleteStep(nil) + + // Verify + messages := buffer.GetMessages() + assert.Len(t, messages, 2) + assert.Equal(t, "assistant-child", messages[0].AssistantID) + assert.Equal(t, "assistant-main", messages[1].AssistantID) + + steps := buffer.GetAllSteps() + assert.Len(t, steps, 3) + assert.Equal(t, "stack-child", steps[1].StackID) + assert.Equal(t, "stack-main", steps[1].StackParentID) + }) + + t.Run("ConcurrentAgentCalls", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-concurrent-a2a", "req-concurrent-a2a", "assistant-main") + + // Main assistant spawns multiple concurrent calls + buffer.BeginStep(context.StepTypeInput, nil, nil) + buffer.CompleteStep(nil) + + // Simulate concurrent responses with thread IDs + var wg sync.WaitGroup + for i := 0; i < 3; i++ { + wg.Add(1) + go func(idx int) { + defer wg.Done() + threadID := fmt.Sprintf("thread-%d", idx) + buffer.AddAssistantMessage( + "text", + map[string]interface{}{"content": fmt.Sprintf("Response from thread %d", idx)}, + "block-concurrent", + threadID, + fmt.Sprintf("assistant-%d", idx), + nil, + ) + }(i) + } + wg.Wait() + + messages := buffer.GetMessages() + assert.Len(t, messages, 3) + + // Verify all have same block ID but different thread IDs + threadIDs := make(map[string]bool) + for _, msg := range messages { + assert.Equal(t, "block-concurrent", msg.BlockID) + assert.False(t, threadIDs[msg.ThreadID], "Duplicate thread ID") + threadIDs[msg.ThreadID] = true + } + }) +} + +// ============================================================================= +// Message Sequence Tests +// ============================================================================= + +func TestBufferMessageSequence(t *testing.T) { + t.Run("SequenceAutoIncrement", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-seq", "req-seq", "assistant-seq") + + for i := 0; i < 10; i++ { + buffer.AddMessage(&context.BufferedMessage{ + Role: "assistant", + Type: "text", + }) + } + + messages := buffer.GetMessages() + for i, msg := range messages { + assert.Equal(t, i+1, msg.Sequence) + } + }) + + t.Run("MixedMessageTypes", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-mixed", "req-mixed", "assistant-mixed") + + buffer.AddUserInput("Hello", "") + buffer.AddAssistantMessage("text", nil, "", "", "", nil) + buffer.AddUserInput("Follow up", "") + buffer.AddAssistantMessage("tool_call", nil, "", "", "", nil) + + messages := buffer.GetMessages() + assert.Len(t, messages, 4) + for i, msg := range messages { + assert.Equal(t, i+1, msg.Sequence) + } + }) +} + +// ============================================================================= +// Step Sequence Tests +// ============================================================================= + +func TestBufferStepSequence(t *testing.T) { + t.Run("SequenceAutoIncrement", func(t *testing.T) { + buffer := context.NewChatBuffer("chat-step-seq", "req-step-seq", "assistant-step-seq") + + for i := 0; i < 5; i++ { + buffer.BeginStep(context.StepTypeLLM, nil, nil) + buffer.CompleteStep(nil) + } + + steps := buffer.GetAllSteps() + for i, step := range steps { + assert.Equal(t, i+1, step.Sequence) + } + }) +} + +// ============================================================================= +// Buffer Reset/Clear Tests (if needed in future) +// ============================================================================= + +func TestBufferMultipleRequests(t *testing.T) { + t.Run("NewBufferPerRequest", func(t *testing.T) { + // Simulate multiple requests with separate buffers + buffer1 := context.NewChatBuffer("chat-1", "req-1", "assistant-1") + buffer1.AddUserInput("Request 1", "") + + buffer2 := context.NewChatBuffer("chat-1", "req-2", "assistant-1") + buffer2.AddUserInput("Request 2", "") + + // Buffers should be independent + assert.Equal(t, 1, buffer1.GetMessageCount()) + assert.Equal(t, 1, buffer2.GetMessageCount()) + + msg1 := buffer1.GetMessages()[0] + msg2 := buffer2.GetMessages()[0] + + assert.Equal(t, "req-1", msg1.RequestID) + assert.Equal(t, "req-2", msg2.RequestID) + }) +} diff --git a/agent/context/context.go b/agent/context/context.go index 0ed8fd32..65948595 100644 --- a/agent/context/context.go +++ b/agent/context/context.go @@ -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 string) *ChatBuffer { + ctx.Buffer = NewChatBuffer(ctx.ChatID, ctx.RequestID(), assistantID) + 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(msgType string, props map[string]interface{}, blockID, threadID string, metadata map[string]interface{}) { + if ctx.Buffer == nil { + return + } + + ctx.Buffer.AddAssistantMessage(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 +} diff --git a/agent/context/output.go b/agent/context/output.go index debd39f5..6bf96c82 100644 --- a/agent/context/output.go +++ b/agent/context/output.go @@ -146,6 +146,25 @@ 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.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) diff --git a/agent/context/types.go b/agent/context/types.go index 8815c312..51b7070b 100644 --- a/agent/context/types.go +++ b/agent/context/types.go @@ -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 diff --git a/agent/store/types/types.go b/agent/store/types/types.go index e4ee154f..a29b9b2a 100644 --- a/agent/store/types/types.go +++ b/agent/store/types/types.go @@ -38,6 +38,12 @@ type Chat struct { 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 for listing chats diff --git a/agent/store/xun/chat.go b/agent/store/xun/chat.go index e850fc11..07362879 100644 --- a/agent/store/xun/chat.go +++ b/agent/store/xun/chat.go @@ -80,6 +80,20 @@ func (store *Xun) CreateChat(chat *types.Chat) error { 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) } @@ -349,6 +363,12 @@ func (store *Xun) rowToChat(data map[string]interface{}) (*types.Chat, error) { } } + // 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 }