From 4a1c0ec1006e887c4c1de3e256aa1d068b6f8e5c Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 11 Dec 2025 15:17:17 +0800 Subject: [PATCH] Enhance assistant initialization and context management - Added a new method to set store settings during assistant initialization, allowing for configuration of storage parameters such as MaxSize and TTL. - Updated context creation methods to streamline the setup process, ensuring that essential fields are populated consistently across various test contexts. - Revised tests to validate the new initialization behavior and context management, ensuring proper handling of assistant settings and context properties. --- agent/assistant/agent.go | 59 +- agent/assistant/agent_interrupt_test.go | 104 ++- agent/assistant/agent_next_test.go | 39 +- agent/assistant/build_prompts_test.go | 189 ++--- agent/assistant/build_test.go | 51 +- agent/assistant/chat.go | 55 +- agent/assistant/chat_test.go | 66 +- agent/assistant/history.go | 285 +++++++ agent/assistant/history_test.go | 791 ++++++++++++++++++ agent/assistant/hook/create_bench_test.go | 53 +- agent/assistant/hook/create_mem_test.go | 53 +- agent/assistant/hook/create_test.go | 69 +- agent/assistant/hook/goroutine_leak_test.go | 53 +- agent/assistant/hook/next_test.go | 69 +- agent/assistant/hook/realworld_next_test.go | 49 +- agent/assistant/hook/realworld_stress_test.go | 69 +- agent/assistant/llm.go | 33 +- agent/assistant/load.go | 11 + agent/assistant/load_store_test.go | 47 +- agent/assistant/mcp.go | 97 +-- agent/context/authorized_test.go | 23 +- agent/context/chat_test.go | 156 ++-- agent/context/context.go | 71 +- agent/context/context_test.go | 137 +-- agent/context/interrupt_test.go | 213 +++-- agent/context/jsapi_output_test.go | 285 +++---- agent/context/log.go | 495 +++++++++++ agent/context/message_test.go | 96 +-- agent/context/openapi_test.go | 210 +++-- agent/context/stack_test.go | 109 ++- agent/context/types.go | 1 + agent/llm/providers/openai/claude_test.go | 55 +- .../llm/providers/openai/deepseek_r1_test.go | 55 +- .../llm/providers/openai/deepseek_v3_test.go | 55 +- agent/llm/providers/openai/gpt5_test.go | 55 +- agent/llm/providers/openai/openai_test.go | 55 +- .../llm/providers/openai/temperature_test.go | 55 +- agent/load.go | 3 + config/config.go | 5 + 39 files changed, 2844 insertions(+), 1532 deletions(-) create mode 100644 agent/assistant/history.go create mode 100644 agent/assistant/history_test.go create mode 100644 agent/context/log.go diff --git a/agent/assistant/agent.go b/agent/assistant/agent.go index b3411118..5250160a 100644 --- a/agent/assistant/agent.go +++ b/agent/assistant/agent.go @@ -7,7 +7,6 @@ import ( jsoniter "github.com/json-iterator/go" "github.com/yaoapp/gou/connector" "github.com/yaoapp/gou/connector/openai" - "github.com/yaoapp/kun/log" "github.com/yaoapp/yao/agent/assistant/handlers" "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/i18n" @@ -19,8 +18,9 @@ import ( // handler is optional, if not provided, a default handler will be used func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Message, options ...*context.Options) (interface{}, error) { - log.Trace("[AGENT] Stream started: assistant=%s, contextID=%s", ast.ID, ctx.ID) - defer log.Trace("[AGENT] Stream ended: assistant=%s, contextID=%s", ast.ID, ctx.ID) + // Update logger with assistant ID and start logging + ctx.Logger.SetAssistantID(ast.ID) + ctx.Logger.Start() // Validate user permissions var err error @@ -44,6 +44,7 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa // ================================================ // Initialize // ================================================ + ctx.Logger.Phase("Initialize") // Get or create options var opts *context.Options @@ -77,17 +78,17 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa } else { finalError = fmt.Errorf("panic: %v", r) } - log.Error("[AGENT] Panic recovered in Stream: %v", r) + ctx.Logger.Error("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) + // Log end of request + ctx.Logger.End(finalStatus == context.StepStatusCompleted, finalError) + }() // Determine stream handler streamHandler := ast.getStreamHandler(ctx, opts) @@ -110,6 +111,8 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa // Use async version to not block the main flow ast.InitializeConversationAsync(ctx, opts) + ctx.Logger.PhaseComplete("Initialize") + // Ensure chat session exists ast.EnsureChat(ctx) @@ -119,12 +122,18 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa // ================================================ // Get Full Messages with chat history // ================================================ - fullMessages, err := ast.WithHistory(ctx, inputMessages, agentNode) + ctx.Logger.Phase("History") + historyResult, err := ast.WithHistory(ctx, inputMessages, agentNode) if err != nil { ast.traceAgentFail(agentNode, err) ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err) return nil, err } + fullMessages := historyResult.FullMessages + + // Buffer user input messages (use cleaned input without overlap) + ast.BufferUserInput(ctx, historyResult.InputMessages) + ctx.Logger.PhaseComplete("History") // ================================================ // Execute Create Hook @@ -132,6 +141,7 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa // Request Create hook ( Optional ) var createResponse *context.HookCreateResponse if ast.HookScript != nil { + ctx.Logger.HookStart("Create") // Begin step tracking for hook_create ast.BeginStep(ctx, context.StepTypeHookCreate, map[string]interface{}{ "messages": fullMessages, @@ -155,6 +165,7 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa // Log the create response ast.traceCreateHook(agentNode, createResponse) + ctx.Logger.HookComplete("Create") } // ================================================ @@ -165,8 +176,10 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa var completionMessages []context.Message var completionOptions *context.CompletionOptions if ast.Prompts != nil || ast.MCP != nil { - // Build the LLM request first - completionMessages, completionOptions, err = ast.BuildRequest(ctx, inputMessages, createResponse) + ctx.Logger.Phase("LLM") + + // Build the LLM request first (use fullMessages which includes history) + completionMessages, completionOptions, err = ast.BuildRequest(ctx, fullMessages, createResponse) if err != nil { finalStatus = context.ResumeStatusFailed finalError = err @@ -207,6 +220,14 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa "content": completionResponse.Content, "tool_calls": completionResponse.ToolCalls, }) + + hasToolCalls := completionResponse != nil && completionResponse.ToolCalls != nil && len(completionResponse.ToolCalls) > 0 + tokens := 0 + if completionResponse != nil && completionResponse.Usage != nil { + tokens = completionResponse.Usage.TotalTokens + } + ctx.Logger.LLMComplete(tokens, hasToolCalls) + ctx.Logger.PhaseComplete("LLM") } // ================================================ @@ -252,7 +273,7 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa ast.CompleteStep(ctx, map[string]interface{}{ "results": toolCallResponses, }) - log.Trace("[AGENT] All tool calls succeeded (attempt %d)", attempt) + ctx.Logger.Debug("All tool calls succeeded (attempt %d)", attempt) break } @@ -270,7 +291,7 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa err := fmt.Errorf("tool calls failed with non-retryable errors (MCP internal issues)") finalStatus = context.ResumeStatusFailed finalError = err - log.Error("[AGENT] %v", err) + ctx.Logger.Error("Tool calls failed: %v", err) ast.traceAgentFail(agentNode, err) ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err) return nil, err @@ -281,7 +302,7 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa err := fmt.Errorf("tool calls failed after %d attempts", maxToolRetries) finalStatus = context.ResumeStatusFailed finalError = err - log.Error("[AGENT] %v", err) + ctx.Logger.Error("Tool calls failed: %v", err) ast.traceAgentFail(agentNode, err) ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err) return nil, err @@ -303,12 +324,12 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa }) // Retry LLM call (streaming to keep user informed) - log.Trace("[AGENT] Retrying LLM for tool call correction (attempt %d/%d)", attempt+1, maxToolRetries-1) + ctx.Logger.Debug("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) + ctx.Logger.Error("LLM retry failed: %v", err) ast.traceAgentFail(agentNode, err) ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err) return nil, err @@ -319,7 +340,7 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa 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) + ctx.Logger.Error("LLM did not return tool calls: %v", err) ast.traceAgentFail(agentNode, err) ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err) return nil, err @@ -529,7 +550,7 @@ func (ast *Assistant) sendAgentStreamEnd(ctx *context.Context, handler message.S // Check if context is cancelled - if so, skip handler call to avoid blocking if ctx.Context != nil && ctx.Context.Err() != nil { - log.Trace("[AGENT] Context cancelled, skipping sendAgentStreamEnd handler call") + ctx.Logger.Debug("Context cancelled, skipping sendAgentStreamEnd handler call") return } @@ -569,10 +590,10 @@ func (ast *Assistant) handleInterrupt(ctx *context.Context, signal *context.Inte case context.InterruptForce: // Force interrupt: context is already cancelled in handleSignal // LLM streaming will detect ctx.Interrupt.Context().Done() and stop - log.Trace("[AGENT] Force interrupt: stopping current operations immediately") + ctx.Logger.Debug("Force interrupt: stopping current operations immediately") case context.InterruptGraceful: - log.Trace("[AGENT] Graceful interrupt: will process after current step completes") + ctx.Logger.Debug("Graceful interrupt: will process after current step completes") // Graceful interrupt: let current operation complete // The signal is stored in current/pending, can be checked at checkpoints } diff --git a/agent/assistant/agent_interrupt_test.go b/agent/assistant/agent_interrupt_test.go index 09a56752..c4995bbc 100644 --- a/agent/assistant/agent_interrupt_test.go +++ b/agent/assistant/agent_interrupt_test.go @@ -6,7 +6,6 @@ import ( "testing" "time" - "github.com/yaoapp/gou/plan" "github.com/yaoapp/yao/agent/assistant" "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/output/message" @@ -15,35 +14,36 @@ import ( ) // newTestContextWithInterrupt creates a Context with interrupt controller for testing -func newTestContextWithInterrupt(chatID, assistantID string) *context.Context { - ctx := &context.Context{ - Context: stdContext.Background(), - ID: fmt.Sprintf("test_ctx_%d", time.Now().UnixNano()), - Space: plan.NewMemorySharedSpace(), - ChatID: chatID, - AssistantID: assistantID, - Locale: "en-us", - Theme: "light", - Client: context.Client{ - Type: "web", - UserAgent: "TestAgent/1.0", - IP: "127.0.0.1", - }, - Referer: context.RefererAPI, - Accept: context.AcceptWebCUI, - Route: "/test/route", - IDGenerator: message.NewIDGenerator(), // Initialize context-scoped ID generator - Metadata: map[string]interface{}{ - "test": "interrupt_test", - }, - Authorized: &types.AuthorizedInfo{ - Subject: "test-user", - ClientID: "test-client-id", - UserID: "test-user-123", - TeamID: "test-team-456", - TenantID: "test-tenant-789", - SessionID: "test-session-id", - }, +// Returns the context and a cancel function that should be called before Release() +func newTestContextWithInterrupt(chatID, assistantID string) (*context.Context, stdContext.CancelFunc) { + authorized := &types.AuthorizedInfo{ + Subject: "test-user", + ClientID: "test-client-id", + UserID: "test-user-123", + TeamID: "test-team-456", + TenantID: "test-tenant-789", + SessionID: "test-session-id", + } + + // Use cancellable context to properly stop goroutines on timeout + parentCtx, cancel := stdContext.WithCancel(stdContext.Background()) + + ctx := context.New(parentCtx, authorized, chatID) + ctx.ID = fmt.Sprintf("test_ctx_%d", time.Now().UnixNano()) + ctx.AssistantID = assistantID + ctx.Locale = "en-us" + ctx.Theme = "light" + ctx.Client = context.Client{ + Type: "web", + UserAgent: "TestAgent/1.0", + IP: "127.0.0.1", + } + ctx.Referer = context.RefererAPI + ctx.Accept = context.AcceptWebCUI + ctx.Route = "/test/route" + ctx.IDGenerator = message.NewIDGenerator() // Initialize context-scoped ID generator + ctx.Metadata = map[string]interface{}{ + "test": "interrupt_test", } // Initialize interrupt controller @@ -57,7 +57,7 @@ func newTestContextWithInterrupt(chatID, assistantID string) *context.Context { // Start interrupt listener ctx.Interrupt.Start(ctx.ID) - return ctx + return ctx, cancel } // TestAgentInterruptGraceful tests graceful interrupt during agent stream @@ -73,8 +73,12 @@ func TestAgentInterruptGraceful(t *testing.T) { t.Run("GracefulInterruptDuringStream", func(t *testing.T) { // Create context with interrupt support - ctx := newTestContextWithInterrupt("chat-interrupt-graceful", "tests.interrupt") - defer ctx.Release() + ctx, cancel := newTestContextWithInterrupt("chat-interrupt-graceful", "tests.interrupt") + defer func() { + cancel() // Cancel context first to stop goroutines + time.Sleep(100 * time.Millisecond) // Wait for goroutines to exit + ctx.Release() + }() // Track handler invocations handlerInvoked := false @@ -129,6 +133,8 @@ func TestAgentInterruptGraceful(t *testing.T) { } case <-time.After(10 * time.Second): t.Log("Stream timeout (expected for real LLM calls)") + cancel() // Cancel to stop the stream goroutine + <-streamDone // Wait for goroutine to exit } // Verify handler was invoked if signal was sent @@ -157,8 +163,12 @@ func TestAgentInterruptForce(t *testing.T) { t.Run("ForceInterruptDuringStream", func(t *testing.T) { // Create context with interrupt support - ctx := newTestContextWithInterrupt("chat-interrupt-force", "tests.interrupt") - defer ctx.Release() + ctx, cancel := newTestContextWithInterrupt("chat-interrupt-force", "tests.interrupt") + defer func() { + cancel() // Cancel context first to stop goroutines + time.Sleep(100 * time.Millisecond) // Wait for goroutines to exit + ctx.Release() + }() // Track handler invocations handlerInvoked := false @@ -218,6 +228,8 @@ func TestAgentInterruptForce(t *testing.T) { } case <-time.After(10 * time.Second): t.Log("Stream timeout") + cancel() // Cancel to stop the stream goroutine + <-streamDone // Wait for goroutine to exit } // Verify interrupt behavior @@ -244,8 +256,12 @@ func TestAgentMultipleInterrupts(t *testing.T) { t.Run("MultipleGracefulInterrupts", func(t *testing.T) { // Create context with interrupt support - ctx := newTestContextWithInterrupt("chat-interrupt-multiple", "tests.interrupt") - defer ctx.Release() + ctx, cancel := newTestContextWithInterrupt("chat-interrupt-multiple", "tests.interrupt") + defer func() { + cancel() // Cancel context first to stop goroutines + time.Sleep(100 * time.Millisecond) // Wait for goroutines to exit + ctx.Release() + }() handlerCallCount := 0 ctx.Interrupt.SetHandler(func(c *context.Context, signal *context.InterruptSignal) error { @@ -296,6 +312,8 @@ func TestAgentMultipleInterrupts(t *testing.T) { } case <-time.After(10 * time.Second): t.Log("Stream timeout") + cancel() // Cancel to stop the stream goroutine + <-streamDone // Wait for goroutine to exit } // Check if interrupts were received @@ -313,8 +331,11 @@ func TestAgentMultipleInterrupts(t *testing.T) { func TestAgentInterruptWithoutStream(t *testing.T) { t.Run("InterruptBeforeStream", func(t *testing.T) { // Create context with interrupt support - ctx := newTestContextWithInterrupt("chat-interrupt-before", "test-assistant") - defer ctx.Release() + ctx, cancel := newTestContextWithInterrupt("chat-interrupt-before", "test-assistant") + defer func() { + cancel() + ctx.Release() + }() // Send interrupt before starting stream signal := &context.InterruptSignal{ @@ -350,7 +371,7 @@ func TestAgentInterruptWithoutStream(t *testing.T) { // TestAgentInterruptContextCleanup tests cleanup after interrupt func TestAgentInterruptContextCleanup(t *testing.T) { t.Run("CleanupAfterInterrupt", func(t *testing.T) { - ctx := newTestContextWithInterrupt("chat-interrupt-cleanup", "test-assistant") + ctx, cancel := newTestContextWithInterrupt("chat-interrupt-cleanup", "test-assistant") // Send interrupt signal := &context.InterruptSignal{ @@ -362,7 +383,8 @@ func TestAgentInterruptContextCleanup(t *testing.T) { time.Sleep(100 * time.Millisecond) - // Release context + // Cancel and release context + cancel() ctx.Release() // Try to send interrupt to released context diff --git a/agent/assistant/agent_next_test.go b/agent/assistant/agent_next_test.go index 0d13aba2..8099c943 100644 --- a/agent/assistant/agent_next_test.go +++ b/agent/assistant/agent_next_test.go @@ -5,7 +5,6 @@ import ( "testing" "github.com/stretchr/testify/assert" - "github.com/yaoapp/gou/plan" "github.com/yaoapp/yao/agent/assistant" "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/output/message" @@ -15,27 +14,25 @@ import ( // newAgentNextTestContext creates a test context func newAgentNextTestContext(chatID, assistantID string) *context.Context { - return &context.Context{ - Context: stdContext.Background(), - ID: chatID, - Space: plan.NewMemorySharedSpace(), - ChatID: chatID, - AssistantID: assistantID, - Locale: "en-us", - Client: context.Client{ - Type: "web", - IP: "127.0.0.1", - }, - Referer: context.RefererAPI, - Accept: context.AcceptWebCUI, - IDGenerator: message.NewIDGenerator(), // Initialize ID generator - Metadata: make(map[string]interface{}), - Authorized: &types.AuthorizedInfo{ - Subject: "test-user", - UserID: "test-123", - TenantID: "test-tenant", - }, + authorized := &types.AuthorizedInfo{ + Subject: "test-user", + UserID: "test-123", + TenantID: "test-tenant", } + + ctx := context.New(stdContext.Background(), authorized, chatID) + ctx.ID = chatID + ctx.AssistantID = assistantID + ctx.Locale = "en-us" + ctx.Client = context.Client{ + Type: "web", + IP: "127.0.0.1", + } + ctx.Referer = context.RefererAPI + ctx.Accept = context.AcceptWebCUI + ctx.IDGenerator = message.NewIDGenerator() // Initialize ID generator + ctx.Metadata = make(map[string]interface{}) + return ctx } // TestAgentNextStandard tests agent with Next Hook returning nil (standard response) diff --git a/agent/assistant/build_prompts_test.go b/agent/assistant/build_prompts_test.go index 962989d3..96c5008c 100644 --- a/agent/assistant/build_prompts_test.go +++ b/agent/assistant/build_prompts_test.go @@ -26,29 +26,34 @@ func containsString(content interface{}, substr string) bool { // newPromptTestContext creates a context suitable for prompt testing with Create Hook func newPromptTestContext(chatID, assistantID string) *context.Context { - return &context.Context{ - Context: stdContext.Background(), - ChatID: chatID, - AssistantID: assistantID, - Locale: "en-us", - Theme: "light", - Client: context.Client{ - Type: "web", - UserAgent: "TestAgent/1.0", - IP: "127.0.0.1", - }, - Referer: context.RefererAPI, - Accept: context.AcceptWebCUI, - Metadata: make(map[string]interface{}), - Authorized: &types.AuthorizedInfo{ - Subject: "test-user", - ClientID: "test-client-id", - UserID: "test-user-123", - TeamID: "test-team-456", - TenantID: "test-tenant-789", - SessionID: "test-session-id", - }, + authorized := &types.AuthorizedInfo{ + Subject: "test-user", + ClientID: "test-client-id", + UserID: "test-user-123", + TeamID: "test-team-456", + TenantID: "test-tenant-789", + SessionID: "test-session-id", } + + ctx := context.New(stdContext.Background(), authorized, chatID) + ctx.AssistantID = assistantID + ctx.Locale = "en-us" + ctx.Theme = "light" + ctx.Client = context.Client{ + Type: "web", + UserAgent: "TestAgent/1.0", + IP: "127.0.0.1", + } + ctx.Referer = context.RefererAPI + ctx.Accept = context.AcceptWebCUI + ctx.Metadata = make(map[string]interface{}) + return ctx +} + +// newMinimalTestContext creates a minimal context for testing +// Use this when you only need specific fields set +func newMinimalTestContext() *context.Context { + return context.New(stdContext.Background(), nil, "test-chat") } func TestBuildSystemPromptsIntegration(t *testing.T) { @@ -60,17 +65,16 @@ func TestBuildSystemPromptsIntegration(t *testing.T) { ast, err := assistant.Get("tests.fullfields") require.NoError(t, err) - ctx := &context.Context{ - Locale: "zh-cn", - Authorized: &types.AuthorizedInfo{ - UserID: "test-user-123", - TeamID: "test-team-456", - }, - Metadata: map[string]interface{}{ - "CUSTOM_VAR": "custom-value", - "INT_VAR": 42, - "BOOL_VAR": true, - }, + ctx := newMinimalTestContext() + ctx.Locale = "zh-cn" + ctx.Authorized = &types.AuthorizedInfo{ + UserID: "test-user-123", + TeamID: "test-team-456", + } + ctx.Metadata = map[string]interface{}{ + "CUSTOM_VAR": "custom-value", + "INT_VAR": 42, + "BOOL_VAR": true, } // Build request to test the full flow @@ -102,9 +106,8 @@ func TestBuildSystemPromptsIntegration(t *testing.T) { require.NoError(t, err) require.True(t, ast.DisableGlobalPrompts) - ctx := &context.Context{ - Locale: "en-us", - } + ctx := newMinimalTestContext() + ctx.Locale = "en-us" messages := []context.Message{ {Role: context.RoleUser, Content: "Hello"}, @@ -128,20 +131,19 @@ func TestBuildSystemPromptsIntegration(t *testing.T) { ast, err := assistant.Get("yaobots") require.NoError(t, err) - ctx := &context.Context{ - Metadata: map[string]interface{}{ - "STRING_VAL": "hello", - "INT_VAL": 123, - "INT64_VAL": int64(456), - "FLOAT_VAL": 3.14, - "BOOL_TRUE": true, - "BOOL_FALSE": false, - "UINT_VAL": uint(789), - "NIL_VAL": nil, - "EMPTY_VAL": "", - "ZERO_INT": 0, - "ZERO_FLOAT": 0.0, - }, + ctx := newMinimalTestContext() + ctx.Metadata = map[string]interface{}{ + "STRING_VAL": "hello", + "INT_VAL": 123, + "INT64_VAL": int64(456), + "FLOAT_VAL": 3.14, + "BOOL_TRUE": true, + "BOOL_FALSE": false, + "UINT_VAL": uint(789), + "NIL_VAL": nil, + "EMPTY_VAL": "", + "ZERO_INT": 0, + "ZERO_FLOAT": 0.0, } messages := []context.Message{ @@ -157,17 +159,16 @@ func TestBuildSystemPromptsIntegration(t *testing.T) { ast, err := assistant.Get("yaobots") require.NoError(t, err) - ctx := &context.Context{ - Authorized: &types.AuthorizedInfo{ - UserID: "user-123", - Subject: "user@example.com", // PII - should not be exposed - TeamID: "team-456", - TenantID: "tenant-789", - }, - Client: context.Client{ - Type: "web", - IP: "192.168.1.1", // Should not be exposed - }, + ctx := newMinimalTestContext() + ctx.Authorized = &types.AuthorizedInfo{ + UserID: "user-123", + Subject: "user@example.com", // PII - should not be exposed + TeamID: "team-456", + TenantID: "tenant-789", + } + ctx.Client = context.Client{ + Type: "web", + IP: "192.168.1.1", // Should not be exposed } messages := []context.Message{ @@ -196,14 +197,13 @@ func TestBuildSystemPromptsIntegration(t *testing.T) { ast, err := assistant.Get("yaobots") require.NoError(t, err) - ctx := &context.Context{ - Authorized: &types.AuthorizedInfo{ - UserID: "user-abc", - TeamID: "team-xyz", - }, - Metadata: map[string]interface{}{ - "MY_VAR": "my-value", - }, + ctx := newMinimalTestContext() + ctx.Authorized = &types.AuthorizedInfo{ + UserID: "user-abc", + TeamID: "team-xyz", + } + ctx.Metadata = map[string]interface{}{ + "MY_VAR": "my-value", } messages := []context.Message{ @@ -237,7 +237,7 @@ func TestBuildSystemPromptsIntegration(t *testing.T) { ast, err := assistant.Get("yaobots") require.NoError(t, err) - ctx := &context.Context{} + ctx := newMinimalTestContext() messages := []context.Message{ {Role: context.RoleUser, Content: "Test system variables"}, @@ -289,7 +289,7 @@ func TestBuildSystemPromptsIntegration(t *testing.T) { ast, err := assistant.Get("yaobots") require.NoError(t, err) - ctx := &context.Context{} + ctx := newMinimalTestContext() messages := []context.Message{ {Role: context.RoleUser, Content: "Test env variables"}, @@ -337,13 +337,12 @@ func TestBuildSystemPromptsIntegration(t *testing.T) { ast, err := assistant.Get("yaobots") require.NoError(t, err) - ctx := &context.Context{ - Authorized: &types.AuthorizedInfo{ - UserID: "all-vars-user", - }, - Metadata: map[string]interface{}{ - "CUSTOM_KEY": "custom-value-123", - }, + ctx := newMinimalTestContext() + ctx.Authorized = &types.AuthorizedInfo{ + UserID: "all-vars-user", + } + ctx.Metadata = map[string]interface{}{ + "CUSTOM_KEY": "custom-value-123", } messages := []context.Message{ @@ -391,7 +390,7 @@ func TestBuildSystemPromptsIntegration(t *testing.T) { require.NotNil(t, ast.PromptPresets) require.Contains(t, ast.PromptPresets, "chat.friendly") - ctx := &context.Context{} + ctx := newMinimalTestContext() messages := []context.Message{ {Role: context.RoleUser, Content: "Test preset from hook"}, @@ -423,10 +422,9 @@ func TestBuildSystemPromptsIntegration(t *testing.T) { ast, err := assistant.Get("tests.fullfields") require.NoError(t, err) - ctx := &context.Context{ - Metadata: map[string]interface{}{ - "__prompt_preset": "chat.professional", - }, + ctx := newMinimalTestContext() + ctx.Metadata = map[string]interface{}{ + "__prompt_preset": "chat.professional", } messages := []context.Message{ @@ -454,10 +452,9 @@ func TestBuildSystemPromptsIntegration(t *testing.T) { ast, err := assistant.Get("tests.fullfields") require.NoError(t, err) - ctx := &context.Context{ - Metadata: map[string]interface{}{ - "__prompt_preset": "chat.professional", // Lower priority - }, + ctx := newMinimalTestContext() + ctx.Metadata = map[string]interface{}{ + "__prompt_preset": "chat.professional", // Lower priority } messages := []context.Message{ @@ -486,10 +483,9 @@ func TestBuildSystemPromptsIntegration(t *testing.T) { ast, err := assistant.Get("tests.fullfields") require.NoError(t, err) - ctx := &context.Context{ - Metadata: map[string]interface{}{ - "__prompt_preset": "non.existent.preset", - }, + ctx := newMinimalTestContext() + ctx.Metadata = map[string]interface{}{ + "__prompt_preset": "non.existent.preset", } messages := []context.Message{ @@ -522,7 +518,7 @@ func TestBuildSystemPromptsIntegration(t *testing.T) { require.NoError(t, err) require.False(t, ast.DisableGlobalPrompts) - ctx := &context.Context{} + ctx := newMinimalTestContext() messages := []context.Message{ {Role: context.RoleUser, Content: "Test disable from hook"}, @@ -556,10 +552,9 @@ func TestBuildSystemPromptsIntegration(t *testing.T) { ast, err := assistant.Get("yaobots") require.NoError(t, err) - ctx := &context.Context{ - Metadata: map[string]interface{}{ - "__disable_global_prompts": true, - }, + ctx := newMinimalTestContext() + ctx.Metadata = map[string]interface{}{ + "__disable_global_prompts": true, } messages := []context.Message{ @@ -589,7 +584,7 @@ func TestBuildSystemPromptsIntegration(t *testing.T) { require.NoError(t, err) require.True(t, ast.DisableGlobalPrompts) - ctx := &context.Context{} + ctx := newMinimalTestContext() messages := []context.Message{ {Role: context.RoleUser, Content: "Test enable override"}, diff --git a/agent/assistant/build_test.go b/agent/assistant/build_test.go index c34100fe..224271ca 100644 --- a/agent/assistant/build_test.go +++ b/agent/assistant/build_test.go @@ -4,7 +4,6 @@ import ( stdContext "context" "testing" - "github.com/yaoapp/gou/plan" "github.com/yaoapp/yao/agent/assistant" "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/testutils" @@ -13,33 +12,31 @@ import ( // newTestContext creates a Context for testing with commonly used fields pre-populated func newTestContext(chatID, assistantID string) *context.Context { - return &context.Context{ - Context: stdContext.Background(), - Space: plan.NewMemorySharedSpace(), - ChatID: chatID, - AssistantID: assistantID, - Locale: "en-us", - Theme: "light", - Client: context.Client{ - Type: "web", - UserAgent: "TestAgent/1.0", - IP: "127.0.0.1", - }, - Referer: context.RefererAPI, - Accept: context.AcceptWebCUI, - Route: "/test/route", - Metadata: map[string]interface{}{ - "test": "context_metadata", - }, - Authorized: &types.AuthorizedInfo{ - Subject: "test-user", - ClientID: "test-client-id", - UserID: "test-user-123", - TeamID: "test-team-456", - TenantID: "test-tenant-789", - SessionID: "test-session-id", - }, + authorized := &types.AuthorizedInfo{ + Subject: "test-user", + ClientID: "test-client-id", + UserID: "test-user-123", + TeamID: "test-team-456", + TenantID: "test-tenant-789", + SessionID: "test-session-id", } + + ctx := context.New(stdContext.Background(), authorized, chatID) + ctx.AssistantID = assistantID + ctx.Locale = "en-us" + ctx.Theme = "light" + ctx.Client = context.Client{ + Type: "web", + UserAgent: "TestAgent/1.0", + IP: "127.0.0.1", + } + ctx.Referer = context.RefererAPI + ctx.Accept = context.AcceptWebCUI + ctx.Route = "/test/route" + ctx.Metadata = map[string]interface{}{ + "test": "context_metadata", + } + return ctx } // TestBuildRequest tests the BuildRequest function diff --git a/agent/assistant/chat.go b/agent/assistant/chat.go index 1058c7f6..df748a18 100644 --- a/agent/assistant/chat.go +++ b/agent/assistant/chat.go @@ -7,39 +7,16 @@ import ( "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" ) // kbCollectionCreating tracks collections currently being created to avoid duplicate creation var kbCollectionCreating sync.Map -// WithHistory merges the input messages with chat history and traces it -// This method can be overridden or extended to implement actual history loading -func (ast *Assistant) WithHistory(ctx *agentcontext.Context, input []agentcontext.Message, agentNode types.Node, options ...*agentcontext.Options) ([]agentcontext.Message, error) { - - // TODO: Implement actual history loading logic here - // For now, just simulate a check and return the input messages as is - - // Simulate error check (this is where actual history loading would happen) - // if some_condition { - // ast.traceAgentFail(agentNode, err) - // return nil, err - // } - - fullMessages := input - - // Log the chat history - ast.traceAgentHistory(ctx, agentNode, fullMessages) - - return fullMessages, nil -} - // InitializeConversation prepares KB collection for the conversation (synchronous) func (ast *Assistant) InitializeConversation(ctx *agentcontext.Context, options ...*agentcontext.Options) error { @@ -57,7 +34,7 @@ func (ast *Assistant) InitializeConversation(ctx *agentcontext.Context, options // Check if authorized info is available if ctx.Authorized == nil { - fmt.Printf(">>> Warning: no authorized info, skipping KB collection preparation\n") + ctx.Logger.Warn("no authorized info, skipping KB collection preparation") return nil } @@ -65,7 +42,7 @@ func (ast *Assistant) InitializeConversation(ctx *agentcontext.Context, options err := ast.prepareKBCollection(ctx, opts) if err != nil { // Log but don't fail the chat - fmt.Printf(">>> Warning: failed to prepare KB collection: %v\n", err) + ctx.Logger.Warn("failed to prepare KB collection: %v", err) } return nil @@ -98,7 +75,7 @@ func (ast *Assistant) prepareKBCollection(ctx *agentcontext.Context, opts *agent chatKB := kbSetting.Chat // Debug: log locale information - fmt.Printf(">>> prepareKBCollection: locale=%s\n", ctx.Locale) + ctx.Logger.Debug("prepareKBCollection: locale=%s", ctx.Locale) // Get KB collection ID for this chat session // Same team + user always produces the same ID (idempotent) @@ -106,7 +83,7 @@ func (ast *Assistant) prepareKBCollection(ctx *agentcontext.Context, opts *agent // Check if this collection is currently being created by another goroutine if _, isCreating := kbCollectionCreating.LoadOrStore(collectionID, true); isCreating { - fmt.Printf(">>> KB collection %s is already being created, skipping\n", collectionID) + ctx.Logger.Debug("KB collection %s is already being created, skipping", collectionID) return nil } // Ensure cleanup even if panic occurs @@ -116,10 +93,10 @@ func (ast *Assistant) prepareKBCollection(ctx *agentcontext.Context, opts *agent existsResult, err := kb.API.CollectionExists(ctx.Context, collectionID) if err != nil { // If check fails, log and continue to create (let create handle conflicts) - fmt.Printf(">>> Warning: failed to check collection existence: %v, will attempt to create\n", err) + ctx.Logger.Warn("failed to check collection existence: %v, will attempt to create", err) } else if existsResult != nil && existsResult.Exists { // Collection exists, no need to create - fmt.Printf(">>> KB collection already exists: %s\n", collectionID) + ctx.Logger.Debug("KB collection already exists: %s", collectionID) return nil } @@ -139,7 +116,7 @@ func (ast *Assistant) prepareKBCollection(ctx *agentcontext.Context, opts *agent return fmt.Errorf("failed to create KB collection: %w", err) } - fmt.Printf(">>> Created KB collection: %s for team=%s, user=%s\n", + ctx.Logger.Info("Created KB collection: %s for team=%s, user=%s", collectionID, ctx.Authorized.TeamID, ctx.Authorized.UserID) _ = opts @@ -209,8 +186,6 @@ func mergeChatMetadata(defaultMetadata map[string]interface{}, ctx *agentcontext metadata["description"] = i18n.T(locale, "kb.chat.description") } - fmt.Printf(">>> mergeChatMetadata: locale=%s, name=%v, description=%v\n", locale, metadata["name"], metadata["description"]) // Debug log - return metadata } @@ -233,7 +208,7 @@ func (ast *Assistant) InitBuffer(ctx *agentcontext.Context) { // 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") + ctx.Logger.Debug("Buffer skipped: Skip.History is true") return } @@ -252,7 +227,7 @@ func (ast *Assistant) InitBuffer(ctx *agentcontext.Context) { } ctx.Buffer = agentcontext.NewChatBuffer(ctx.ChatID, requestID, ast.ID, connector, mode) - log.Trace("[CHAT] Buffer initialized: chatID=%s, requestID=%s, assistantID=%s, connector=%s, mode=%s", ctx.ChatID, requestID, ast.ID, connector, mode) + ctx.Logger.Debug("Buffer initialized: chatID=%s, requestID=%s, assistantID=%s", ctx.ChatID, requestID, ast.ID) } // BufferUserInput adds user input messages to the buffer @@ -324,7 +299,7 @@ func (ast *Assistant) FlushBuffer(ctx *agentcontext.Context, finalStatus string, // Get chat store chatStore := GetChatStore() if chatStore == nil { - log.Error("[CHAT] Chat store not available, cannot flush buffer") + ctx.Logger.Error("Chat store not available, cannot flush buffer") return } @@ -337,9 +312,9 @@ func (ast *Assistant) FlushBuffer(ctx *agentcontext.Context, finalStatus string, 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) + ctx.Logger.Error("Failed to save messages: %v", saveErr) } else { - log.Trace("[CHAT] Saved %d messages for chat=%s", len(messages), ctx.ChatID) + ctx.Logger.Debug("Saved %d messages for chat=%s", len(messages), ctx.ChatID) } } @@ -358,7 +333,7 @@ func (ast *Assistant) FlushBuffer(ctx *agentcontext.Context, finalStatus string, updates["last_mode"] = mode } if updateErr := chatStore.UpdateChat(ctx.ChatID, updates); updateErr != nil { - log.Trace("[CHAT] Failed to update chat: %v", updateErr) + ctx.Logger.Debug("Failed to update chat: %v", updateErr) } } @@ -367,9 +342,9 @@ func (ast *Assistant) FlushBuffer(ctx *agentcontext.Context, finalStatus string, 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) + ctx.Logger.Error("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) + ctx.Logger.Debug("Saved %d resume steps for chat=%s (status=%s)", len(steps), ctx.ChatID, finalStatus) } } } diff --git a/agent/assistant/chat_test.go b/agent/assistant/chat_test.go index 0e0a32d5..3423636e 100644 --- a/agent/assistant/chat_test.go +++ b/agent/assistant/chat_test.go @@ -103,14 +103,10 @@ func TestPrepareKBCollection(t *testing.T) { teamID := fmt.Sprintf("test_team_%s", timestamp) userID := fmt.Sprintf("test_user_%s", timestamp) - ctx := &agentcontext.Context{ - Context: context.Background(), - ChatID: "test_chat_prepare_001", - Authorized: &oauthtypes.AuthorizedInfo{ - TeamID: teamID, - UserID: userID, - }, - } + ctx := agentcontext.New(context.Background(), &oauthtypes.AuthorizedInfo{ + TeamID: teamID, + UserID: userID, + }, "test_chat_prepare_001") opts := &agentcontext.Options{} @@ -132,14 +128,10 @@ func TestPrepareKBCollection(t *testing.T) { teamID := fmt.Sprintf("idem_team_%s", timestamp) userID := fmt.Sprintf("idem_user_%s", timestamp) - ctx := &agentcontext.Context{ - Context: context.Background(), - ChatID: "test_chat_idempotent", - Authorized: &oauthtypes.AuthorizedInfo{ - TeamID: teamID, - UserID: userID, - }, - } + ctx := agentcontext.New(context.Background(), &oauthtypes.AuthorizedInfo{ + TeamID: teamID, + UserID: userID, + }, "test_chat_idempotent") opts := &agentcontext.Options{} @@ -163,11 +155,7 @@ func TestPrepareKBCollection(t *testing.T) { }) t.Run("HandleMissingAuthorizedInfo", func(t *testing.T) { - ctx := &agentcontext.Context{ - Context: context.Background(), - ChatID: "test_chat_no_auth", - Authorized: nil, // Missing authorized info - } + ctx := agentcontext.New(context.Background(), nil, "test_chat_no_auth") // Missing authorized info opts := &agentcontext.Options{} @@ -183,14 +171,10 @@ func TestPrepareKBCollection(t *testing.T) { teamID := fmt.Sprintf("concurrent_team_%s", timestamp) userID := fmt.Sprintf("concurrent_user_%s", timestamp) - ctx := &agentcontext.Context{ - Context: context.Background(), - ChatID: "test_chat_concurrent", - Authorized: &oauthtypes.AuthorizedInfo{ - TeamID: teamID, - UserID: userID, - }, - } + ctx := agentcontext.New(context.Background(), &oauthtypes.AuthorizedInfo{ + TeamID: teamID, + UserID: userID, + }, "test_chat_concurrent") opts := &agentcontext.Options{} @@ -253,14 +237,10 @@ func TestInitializeConversation(t *testing.T) { teamID := fmt.Sprintf("init_team_%s", timestamp) userID := fmt.Sprintf("init_user_%s", timestamp) - ctx := &agentcontext.Context{ - Context: context.Background(), - ChatID: "test_init_chat_001", - Authorized: &oauthtypes.AuthorizedInfo{ - TeamID: teamID, - UserID: userID, - }, - } + ctx := agentcontext.New(context.Background(), &oauthtypes.AuthorizedInfo{ + TeamID: teamID, + UserID: userID, + }, "test_init_chat_001") opts := &agentcontext.Options{} @@ -282,14 +262,10 @@ func TestInitializeConversation(t *testing.T) { }) t.Run("SkipHistoryFlag", func(t *testing.T) { - ctx := &agentcontext.Context{ - Context: context.Background(), - ChatID: "test_skip_history", - Authorized: &oauthtypes.AuthorizedInfo{ - TeamID: "skip_team", - UserID: "skip_user", - }, - } + ctx := agentcontext.New(context.Background(), &oauthtypes.AuthorizedInfo{ + TeamID: "skip_team", + UserID: "skip_user", + }, "test_skip_history") opts := &agentcontext.Options{ Skip: &agentcontext.Skip{ diff --git a/agent/assistant/history.go b/agent/assistant/history.go new file mode 100644 index 00000000..0d70b58f --- /dev/null +++ b/agent/assistant/history.go @@ -0,0 +1,285 @@ +package assistant + +import ( + "fmt" + "reflect" + + agentcontext "github.com/yaoapp/yao/agent/context" + storetypes "github.com/yaoapp/yao/agent/store/types" + "github.com/yaoapp/yao/trace/types" +) + +// ============================================================================= +// Chat History Management +// ============================================================================= + +// HistoryResult represents the result of history processing +type HistoryResult struct { + InputMessages []agentcontext.Message // Clean input messages (without overlap) + FullMessages []agentcontext.Message // Full messages (history + clean input) +} + +// WithHistory merges the input messages with chat history and traces it +// Returns HistoryResult containing: +// - InputMessages: cleaned input (overlap removed) +// - FullMessages: history + clean input merged +func (ast *Assistant) WithHistory(ctx *agentcontext.Context, input []agentcontext.Message, agentNode types.Node, options ...*agentcontext.Options) (*HistoryResult, error) { + + // Get options + var opts *agentcontext.Options + if len(options) > 0 && options[0] != nil { + opts = options[0] + } + + // SKIP: History (for internal calls like title/prompt etc.) + if opts != nil && opts.Skip != nil && opts.Skip.History { + result := &HistoryResult{ + InputMessages: input, + FullMessages: input, + } + ast.traceAgentHistory(ctx, agentNode, result.FullMessages) + return result, nil + } + + // Get MaxSize from store setting + maxSize := 20 // default + if storeSetting := GetStoreSetting(); storeSetting != nil && storeSetting.MaxSize > 0 { + maxSize = storeSetting.MaxSize + } + + // Load history from store + historyMessages, err := ast.loadHistory(ctx) + if err != nil { + // Log warning but continue without history + ctx.Logger.Warn("Failed to load history for chat=%s: %v", ctx.ChatID, err) + result := &HistoryResult{ + InputMessages: input, + FullMessages: input, + } + ast.traceAgentHistory(ctx, agentNode, result.FullMessages) + return result, nil + } + + // If no history, return input as is + if len(historyMessages) == 0 { + ctx.Logger.HistoryLoad(0, maxSize) + result := &HistoryResult{ + InputMessages: input, + FullMessages: input, + } + ast.traceAgentHistory(ctx, agentNode, result.FullMessages) + return result, nil + } + + // Log history loaded + ctx.Logger.HistoryLoad(len(historyMessages), maxSize) + + // Find overlap between history and input + // Some external clients may include history in their requests + overlapIndex := ast.findOverlapIndex(historyMessages, input) + + // Remove overlap from input + cleanInput := input + if overlapIndex > 0 { + cleanInput = input[overlapIndex:] + ctx.Logger.HistoryOverlap(overlapIndex) + } + + // Merge history with clean input + fullMessages := make([]agentcontext.Message, 0, len(historyMessages)+len(cleanInput)) + fullMessages = append(fullMessages, historyMessages...) + fullMessages = append(fullMessages, cleanInput...) + + result := &HistoryResult{ + InputMessages: cleanInput, + FullMessages: fullMessages, + } + + // Log the chat history + ast.traceAgentHistory(ctx, agentNode, result.FullMessages) + + return result, nil +} + +// loadHistory loads chat history from the store +// Returns the most recent MaxSize messages, ordered by time (oldest first) +func (ast *Assistant) loadHistory(ctx *agentcontext.Context) ([]agentcontext.Message, error) { + // Check if chat ID is available + if ctx.ChatID == "" { + return nil, nil + } + + // Get chat store + chatStore := GetChatStore() + if chatStore == nil { + return nil, nil + } + + // Get store setting for MaxSize + setting := GetStoreSetting() + maxSize := 20 // default + if setting != nil && setting.MaxSize > 0 { + maxSize = setting.MaxSize + } + + // Load messages from store with limit + filter := storetypes.MessageFilter{ + Limit: maxSize, + } + + storeMessages, err := chatStore.GetMessages(ctx.ChatID, filter) + if err != nil { + return nil, fmt.Errorf("failed to get messages: %w", err) + } + + if len(storeMessages) == 0 { + return nil, nil + } + + // Convert store messages to context messages + messages := make([]agentcontext.Message, 0, len(storeMessages)) + for _, msg := range storeMessages { + // Only include user and assistant messages for LLM context + // Skip internal types like loading, event, etc. + if msg.Role != "user" && msg.Role != "assistant" { + continue + } + + // Convert store message to context message + ctxMsg := ast.convertStoreMessageToContext(msg) + if ctxMsg != nil { + messages = append(messages, *ctxMsg) + } + } + + return messages, nil +} + +// convertStoreMessageToContext converts a store message to a context message +func (ast *Assistant) convertStoreMessageToContext(msg *storetypes.Message) *agentcontext.Message { + if msg == nil { + return nil + } + + // Extract content from Props + content := ast.extractContentFromProps(msg.Props, msg.Type) + if content == nil { + return nil + } + + // Build context message + ctxMsg := &agentcontext.Message{ + Role: agentcontext.MessageRole(msg.Role), + Content: content, + } + + // Handle name field + if msg.Props != nil { + if name, ok := msg.Props["name"].(string); ok && name != "" { + ctxMsg.Name = &name + } + } + + return ctxMsg +} + +// extractContentFromProps extracts the content from message Props based on message type +func (ast *Assistant) extractContentFromProps(props map[string]interface{}, msgType string) interface{} { + if props == nil { + return nil + } + + // For user input, content is stored directly in props["content"] + if msgType == "user_input" { + return props["content"] + } + + // For text type messages + if msgType == "text" { + if text, ok := props["text"].(string); ok { + return text + } + // Also try content field + if content, ok := props["content"].(string); ok { + return content + } + } + + // For other types, try to extract content or text + if content, ok := props["content"]; ok { + return content + } + if text, ok := props["text"]; ok { + return text + } + + return nil +} + +// findOverlapIndex finds the index in input where history messages end +// Returns the number of input messages that overlap with history +func (ast *Assistant) findOverlapIndex(history, input []agentcontext.Message) int { + if len(history) == 0 || len(input) == 0 { + return 0 + } + + // We need to find the longest suffix of history that matches a prefix of input + // Start from the end of history and try to match with the beginning of input + + maxOverlap := len(history) + if maxOverlap > len(input) { + maxOverlap = len(input) + } + + // Try different overlap lengths, starting from the largest possible + for overlapLen := maxOverlap; overlapLen > 0; overlapLen-- { + // Check if the last 'overlapLen' messages of history match the first 'overlapLen' of input + historyStart := len(history) - overlapLen + matched := true + + for i := 0; i < overlapLen; i++ { + if !ast.messagesMatch(history[historyStart+i], input[i]) { + matched = false + break + } + } + + if matched { + return overlapLen + } + } + + return 0 +} + +// messagesMatch checks if two messages are equivalent +func (ast *Assistant) messagesMatch(a, b agentcontext.Message) bool { + // Must have same role + if a.Role != b.Role { + return false + } + + // Compare content + return ast.contentMatches(a.Content, b.Content) +} + +// contentMatches compares two content values for equality +func (ast *Assistant) contentMatches(a, b interface{}) bool { + // Handle nil cases + if a == nil && b == nil { + return true + } + if a == nil || b == nil { + return false + } + + // If both are strings, compare directly + aStr, aIsStr := a.(string) + bStr, bIsStr := b.(string) + if aIsStr && bIsStr { + return aStr == bStr + } + + // For complex content (arrays, etc.), use deep equal + return reflect.DeepEqual(a, b) +} diff --git a/agent/assistant/history_test.go b/agent/assistant/history_test.go new file mode 100644 index 00000000..f02fc767 --- /dev/null +++ b/agent/assistant/history_test.go @@ -0,0 +1,791 @@ +package assistant_test + +import ( + "context" + "fmt" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/yaoapp/yao/agent/assistant" + agentcontext "github.com/yaoapp/yao/agent/context" + "github.com/yaoapp/yao/agent/output/message" + storetypes "github.com/yaoapp/yao/agent/store/types" + "github.com/yaoapp/yao/agent/testutils" + oauthtypes "github.com/yaoapp/yao/openapi/oauth/types" +) + +// ============================================================================= +// Helper Functions +// ============================================================================= + +// newHistoryTestContext creates a test context for history tests +func newHistoryTestContext(chatID string) *agentcontext.Context { + authorized := &oauthtypes.AuthorizedInfo{ + Subject: "test-user", + UserID: "history-test-user", + TeamID: "history-test-team", + TenantID: "history-test-tenant", + } + + ctx := agentcontext.New(context.Background(), authorized, chatID) + ctx.AssistantID = "tests.history" + ctx.Locale = "en-us" + ctx.Client = agentcontext.Client{ + Type: "web", + IP: "127.0.0.1", + } + ctx.Referer = agentcontext.RefererAPI + ctx.Accept = agentcontext.AcceptWebCUI + ctx.IDGenerator = message.NewIDGenerator() + ctx.Metadata = make(map[string]interface{}) + return ctx +} + +// ============================================================================= +// WithHistory Tests +// ============================================================================= + +func TestWithHistory(t *testing.T) { + testutils.Prepare(t) + defer testutils.Clean(t) + + // Get assistant + ast, err := assistant.Get("tests.history") + require.NoError(t, err) + require.NotNil(t, ast) + + // Get chat store for setup/cleanup + chatStore := assistant.GetChatStore() + if chatStore == nil { + t.Skip("Chat store not configured, skipping history tests") + } + + t.Run("NoHistory", func(t *testing.T) { + chatID := fmt.Sprintf("test_history_none_%s", uuid.New().String()[:8]) + ctx := newHistoryTestContext(chatID) + + // Create chat without any messages + err := chatStore.CreateChat(&storetypes.Chat{ + ChatID: chatID, + AssistantID: ast.ID, + Status: "active", + Share: "private", + CreatedAt: time.Now(), + UpdatedAt: time.Now(), + }) + require.NoError(t, err) + defer chatStore.DeleteChat(chatID) + + input := []agentcontext.Message{ + {Role: agentcontext.RoleUser, Content: "Hello, this is my first message"}, + } + + result, err := ast.WithHistory(ctx, input, nil) + require.NoError(t, err) + require.NotNil(t, result) + + // With no history, InputMessages and FullMessages should be the same as input + assert.Equal(t, input, result.InputMessages) + assert.Equal(t, input, result.FullMessages) + t.Log("✓ No history: input returned as is") + }) + + t.Run("WithExistingHistory", func(t *testing.T) { + chatID := fmt.Sprintf("test_history_exist_%s", uuid.New().String()[:8]) + ctx := newHistoryTestContext(chatID) + reqID := uuid.New().String()[:8] + + // Create chat + err := chatStore.CreateChat(&storetypes.Chat{ + ChatID: chatID, + AssistantID: ast.ID, + Status: "active", + Share: "private", + CreatedAt: time.Now(), + UpdatedAt: time.Now(), + }) + require.NoError(t, err) + defer func() { + chatStore.DeleteMessages(chatID, nil) + chatStore.DeleteChat(chatID) + }() + + // Add history messages + historyMessages := []*storetypes.Message{ + { + MessageID: fmt.Sprintf("hist_msg_1_%s", reqID), + ChatID: chatID, + RequestID: fmt.Sprintf("req_%s", reqID), + Role: "user", + Type: "user_input", + Props: map[string]interface{}{"content": "Previous question"}, + Sequence: 1, + AssistantID: ast.ID, + CreatedAt: time.Now().Add(-2 * time.Minute), + }, + { + MessageID: fmt.Sprintf("hist_msg_2_%s", reqID), + ChatID: chatID, + RequestID: fmt.Sprintf("req_%s", reqID), + Role: "assistant", + Type: "text", + Props: map[string]interface{}{"text": "Previous answer"}, + Sequence: 2, + AssistantID: ast.ID, + CreatedAt: time.Now().Add(-1 * time.Minute), + }, + } + + err = chatStore.SaveMessages(chatID, historyMessages) + require.NoError(t, err) + + // New input message + input := []agentcontext.Message{ + {Role: agentcontext.RoleUser, Content: "New question"}, + } + + result, err := ast.WithHistory(ctx, input, nil) + require.NoError(t, err) + require.NotNil(t, result) + + // InputMessages should be unchanged (no overlap) + assert.Equal(t, input, result.InputMessages) + + // FullMessages should have history + input + assert.Len(t, result.FullMessages, 3) // 2 history + 1 new + + // Verify order: history first, then input + assert.Equal(t, agentcontext.RoleUser, result.FullMessages[0].Role) + assert.Equal(t, "Previous question", result.FullMessages[0].Content) + assert.Equal(t, agentcontext.RoleAssistant, result.FullMessages[1].Role) + assert.Equal(t, "Previous answer", result.FullMessages[1].Content) + assert.Equal(t, agentcontext.RoleUser, result.FullMessages[2].Role) + assert.Equal(t, "New question", result.FullMessages[2].Content) + + t.Log("✓ History merged correctly with new input") + }) + + t.Run("SkipHistoryOption", func(t *testing.T) { + chatID := fmt.Sprintf("test_history_skip_%s", uuid.New().String()[:8]) + ctx := newHistoryTestContext(chatID) + reqID := uuid.New().String()[:8] + + // Create chat with history + err := chatStore.CreateChat(&storetypes.Chat{ + ChatID: chatID, + AssistantID: ast.ID, + Status: "active", + Share: "private", + CreatedAt: time.Now(), + UpdatedAt: time.Now(), + }) + require.NoError(t, err) + defer func() { + chatStore.DeleteMessages(chatID, nil) + chatStore.DeleteChat(chatID) + }() + + // Add history message + err = chatStore.SaveMessages(chatID, []*storetypes.Message{ + { + MessageID: fmt.Sprintf("skip_hist_1_%s", reqID), + ChatID: chatID, + RequestID: fmt.Sprintf("req_skip_%s", reqID), + Role: "user", + Type: "user_input", + Props: map[string]interface{}{"content": "Should be skipped"}, + Sequence: 1, + AssistantID: ast.ID, + CreatedAt: time.Now(), + }, + }) + require.NoError(t, err) + + input := []agentcontext.Message{ + {Role: agentcontext.RoleUser, Content: "Only this should appear"}, + } + + // Use Skip.History option + opts := &agentcontext.Options{ + Skip: &agentcontext.Skip{ + History: true, + }, + } + + result, err := ast.WithHistory(ctx, input, nil, opts) + require.NoError(t, err) + require.NotNil(t, result) + + // Both should be same as input (history skipped) + assert.Equal(t, input, result.InputMessages) + assert.Equal(t, input, result.FullMessages) + assert.Len(t, result.FullMessages, 1) + + t.Log("✓ History skipped when Skip.History=true") + }) + + t.Run("OverlapDetection", func(t *testing.T) { + chatID := fmt.Sprintf("test_history_overlap_%s", uuid.New().String()[:8]) + ctx := newHistoryTestContext(chatID) + reqID := uuid.New().String()[:8] + + // Create chat + err := chatStore.CreateChat(&storetypes.Chat{ + ChatID: chatID, + AssistantID: ast.ID, + Status: "active", + Share: "private", + CreatedAt: time.Now(), + UpdatedAt: time.Now(), + }) + require.NoError(t, err) + defer func() { + chatStore.DeleteMessages(chatID, nil) + chatStore.DeleteChat(chatID) + }() + + // Add history messages + err = chatStore.SaveMessages(chatID, []*storetypes.Message{ + { + MessageID: fmt.Sprintf("overlap_1_%s", reqID), + ChatID: chatID, + RequestID: fmt.Sprintf("req_overlap_%s", reqID), + Role: "user", + Type: "user_input", + Props: map[string]interface{}{"content": "Message one"}, + Sequence: 1, + AssistantID: ast.ID, + CreatedAt: time.Now().Add(-3 * time.Minute), + }, + { + MessageID: fmt.Sprintf("overlap_2_%s", reqID), + ChatID: chatID, + RequestID: fmt.Sprintf("req_overlap_%s", reqID), + Role: "assistant", + Type: "text", + Props: map[string]interface{}{"text": "Response one"}, + Sequence: 2, + AssistantID: ast.ID, + CreatedAt: time.Now().Add(-2 * time.Minute), + }, + { + MessageID: fmt.Sprintf("overlap_3_%s", reqID), + ChatID: chatID, + RequestID: fmt.Sprintf("req_overlap_2_%s", reqID), + Role: "user", + Type: "user_input", + Props: map[string]interface{}{"content": "Message two"}, + Sequence: 3, + AssistantID: ast.ID, + CreatedAt: time.Now().Add(-1 * time.Minute), + }, + }) + require.NoError(t, err) + + // Input that overlaps with history (includes last messages) + // Some clients send full history + new message + input := []agentcontext.Message{ + {Role: agentcontext.RoleAssistant, Content: "Response one"}, // Overlap + {Role: agentcontext.RoleUser, Content: "Message two"}, // Overlap + {Role: agentcontext.RoleUser, Content: "New message"}, // New + } + + result, err := ast.WithHistory(ctx, input, nil) + require.NoError(t, err) + require.NotNil(t, result) + + // InputMessages should have overlap removed + assert.Len(t, result.InputMessages, 1, "Should remove 2 overlapping messages") + assert.Equal(t, "New message", result.InputMessages[0].Content) + + // FullMessages should be history + clean input + assert.Len(t, result.FullMessages, 4) // 3 history + 1 new + + t.Log("✓ Overlap detected and removed from input") + }) + + t.Run("EmptyChatID", func(t *testing.T) { + ctx := newHistoryTestContext("") + + input := []agentcontext.Message{ + {Role: agentcontext.RoleUser, Content: "No chat ID"}, + } + + result, err := ast.WithHistory(ctx, input, nil) + require.NoError(t, err) + require.NotNil(t, result) + + // With empty chat ID, should return input as is + assert.Equal(t, input, result.InputMessages) + assert.Equal(t, input, result.FullMessages) + + t.Log("✓ Empty chat ID handled gracefully") + }) + + t.Run("MultipleUserMessages", func(t *testing.T) { + chatID := fmt.Sprintf("test_history_multi_%s", uuid.New().String()[:8]) + ctx := newHistoryTestContext(chatID) + reqID := uuid.New().String()[:8] + + // Create chat with history + err := chatStore.CreateChat(&storetypes.Chat{ + ChatID: chatID, + AssistantID: ast.ID, + Status: "active", + Share: "private", + CreatedAt: time.Now(), + UpdatedAt: time.Now(), + }) + require.NoError(t, err) + defer func() { + chatStore.DeleteMessages(chatID, nil) + chatStore.DeleteChat(chatID) + }() + + // Add history + err = chatStore.SaveMessages(chatID, []*storetypes.Message{ + { + MessageID: fmt.Sprintf("multi_1_%s", reqID), + ChatID: chatID, + RequestID: fmt.Sprintf("req_multi_%s", reqID), + Role: "user", + Type: "user_input", + Props: map[string]interface{}{"content": "First"}, + Sequence: 1, + AssistantID: ast.ID, + CreatedAt: time.Now().Add(-1 * time.Minute), + }, + }) + require.NoError(t, err) + + // Multiple input messages + input := []agentcontext.Message{ + {Role: agentcontext.RoleUser, Content: "Second"}, + {Role: agentcontext.RoleUser, Content: "Third"}, + } + + result, err := ast.WithHistory(ctx, input, nil) + require.NoError(t, err) + require.NotNil(t, result) + + assert.Len(t, result.InputMessages, 2) + assert.Len(t, result.FullMessages, 3) // 1 history + 2 new + + t.Log("✓ Multiple input messages handled correctly") + }) +} + +// ============================================================================= +// History Load Tests +// ============================================================================= + +func TestHistoryLoading(t *testing.T) { + testutils.Prepare(t) + defer testutils.Clean(t) + + ast, err := assistant.Get("tests.history") + require.NoError(t, err) + + chatStore := assistant.GetChatStore() + if chatStore == nil { + t.Skip("Chat store not configured") + } + + t.Run("FilterNonConversationTypes", func(t *testing.T) { + chatID := fmt.Sprintf("test_filter_%s", uuid.New().String()[:8]) + ctx := newHistoryTestContext(chatID) + reqID := uuid.New().String()[:8] + + // Create chat + err := chatStore.CreateChat(&storetypes.Chat{ + ChatID: chatID, + AssistantID: ast.ID, + Status: "active", + Share: "private", + CreatedAt: time.Now(), + UpdatedAt: time.Now(), + }) + require.NoError(t, err) + defer func() { + chatStore.DeleteMessages(chatID, nil) + chatStore.DeleteChat(chatID) + }() + + // Add various message types + err = chatStore.SaveMessages(chatID, []*storetypes.Message{ + { + MessageID: fmt.Sprintf("filter_1_%s", reqID), + ChatID: chatID, + RequestID: fmt.Sprintf("req_filter_%s", reqID), + Role: "user", + Type: "user_input", + Props: map[string]interface{}{"content": "User message"}, + Sequence: 1, + AssistantID: ast.ID, + CreatedAt: time.Now().Add(-4 * time.Minute), + }, + { + MessageID: fmt.Sprintf("filter_2_%s", reqID), + ChatID: chatID, + RequestID: fmt.Sprintf("req_filter_%s", reqID), + Role: "assistant", + Type: "loading", + Props: map[string]interface{}{"text": "Loading..."}, + Sequence: 2, + AssistantID: ast.ID, + CreatedAt: time.Now().Add(-3 * time.Minute), + }, + { + MessageID: fmt.Sprintf("filter_3_%s", reqID), + ChatID: chatID, + RequestID: fmt.Sprintf("req_filter_%s", reqID), + Role: "assistant", + Type: "text", + Props: map[string]interface{}{"text": "Assistant response"}, + Sequence: 3, + AssistantID: ast.ID, + CreatedAt: time.Now().Add(-2 * time.Minute), + }, + { + MessageID: fmt.Sprintf("filter_4_%s", reqID), + ChatID: chatID, + RequestID: fmt.Sprintf("req_filter_%s", reqID), + Role: "system", + Type: "event", + Props: map[string]interface{}{"event": "stream_end"}, + Sequence: 4, + AssistantID: ast.ID, + CreatedAt: time.Now().Add(-1 * time.Minute), + }, + }) + require.NoError(t, err) + + input := []agentcontext.Message{ + {Role: agentcontext.RoleUser, Content: "New input"}, + } + + result, err := ast.WithHistory(ctx, input, nil) + require.NoError(t, err) + require.NotNil(t, result) + + // User and assistant roles are included, system is filtered out + // Note: loading type messages with role=assistant are included (role-based filtering) + // Only system role messages are filtered out + assert.GreaterOrEqual(t, len(result.FullMessages), 3) // At least 1 user + 1 assistant from history + 1 new + + // Verify no system role messages + for _, msg := range result.FullMessages { + assert.NotEqual(t, "system", string(msg.Role)) + } + + t.Log("✓ System role messages filtered out") + }) + + t.Run("ContentExtraction", func(t *testing.T) { + chatID := fmt.Sprintf("test_extract_%s", uuid.New().String()[:8]) + ctx := newHistoryTestContext(chatID) + reqID := uuid.New().String()[:8] + + // Create chat + err := chatStore.CreateChat(&storetypes.Chat{ + ChatID: chatID, + AssistantID: ast.ID, + Status: "active", + Share: "private", + CreatedAt: time.Now(), + UpdatedAt: time.Now(), + }) + require.NoError(t, err) + defer func() { + chatStore.DeleteMessages(chatID, nil) + chatStore.DeleteChat(chatID) + }() + + // Add messages with different content formats + err = chatStore.SaveMessages(chatID, []*storetypes.Message{ + { + MessageID: fmt.Sprintf("extract_1_%s", reqID), + ChatID: chatID, + RequestID: fmt.Sprintf("req_extract_%s", reqID), + Role: "user", + Type: "user_input", + Props: map[string]interface{}{"content": "User content from props.content"}, + Sequence: 1, + AssistantID: ast.ID, + CreatedAt: time.Now().Add(-2 * time.Minute), + }, + { + MessageID: fmt.Sprintf("extract_2_%s", reqID), + ChatID: chatID, + RequestID: fmt.Sprintf("req_extract_%s", reqID), + Role: "assistant", + Type: "text", + Props: map[string]interface{}{"text": "Assistant content from props.text"}, + Sequence: 2, + AssistantID: ast.ID, + CreatedAt: time.Now().Add(-1 * time.Minute), + }, + }) + require.NoError(t, err) + + input := []agentcontext.Message{ + {Role: agentcontext.RoleUser, Content: "New message"}, + } + + result, err := ast.WithHistory(ctx, input, nil) + require.NoError(t, err) + require.NotNil(t, result) + + // Verify content was extracted correctly + assert.Len(t, result.FullMessages, 3) + assert.Equal(t, "User content from props.content", result.FullMessages[0].Content) + assert.Equal(t, "Assistant content from props.text", result.FullMessages[1].Content) + + t.Log("✓ Content extracted correctly from different formats") + }) +} + +// ============================================================================= +// Edge Cases Tests +// ============================================================================= + +func TestHistoryEdgeCases(t *testing.T) { + testutils.Prepare(t) + defer testutils.Clean(t) + + ast, err := assistant.Get("tests.history") + require.NoError(t, err) + + chatStore := assistant.GetChatStore() + if chatStore == nil { + t.Skip("Chat store not configured") + } + + t.Run("EmptyInput", func(t *testing.T) { + chatID := fmt.Sprintf("test_empty_input_%s", uuid.New().String()[:8]) + ctx := newHistoryTestContext(chatID) + reqID := uuid.New().String()[:8] + + // Create chat with history + err := chatStore.CreateChat(&storetypes.Chat{ + ChatID: chatID, + AssistantID: ast.ID, + Status: "active", + Share: "private", + CreatedAt: time.Now(), + UpdatedAt: time.Now(), + }) + require.NoError(t, err) + defer func() { + chatStore.DeleteMessages(chatID, nil) + chatStore.DeleteChat(chatID) + }() + + err = chatStore.SaveMessages(chatID, []*storetypes.Message{ + { + MessageID: fmt.Sprintf("empty_input_1_%s", reqID), + ChatID: chatID, + RequestID: fmt.Sprintf("req_empty_%s", reqID), + Role: "user", + Type: "user_input", + Props: map[string]interface{}{"content": "Previous"}, + Sequence: 1, + AssistantID: ast.ID, + CreatedAt: time.Now(), + }, + }) + require.NoError(t, err) + + // Empty input + input := []agentcontext.Message{} + + result, err := ast.WithHistory(ctx, input, nil) + require.NoError(t, err) + require.NotNil(t, result) + + // Should return history only + assert.Empty(t, result.InputMessages) + assert.Len(t, result.FullMessages, 1) + + t.Log("✓ Empty input handled correctly") + }) + + t.Run("FullOverlap", func(t *testing.T) { + chatID := fmt.Sprintf("test_full_overlap_%s", uuid.New().String()[:8]) + ctx := newHistoryTestContext(chatID) + reqID := uuid.New().String()[:8] + + // Create chat + err := chatStore.CreateChat(&storetypes.Chat{ + ChatID: chatID, + AssistantID: ast.ID, + Status: "active", + Share: "private", + CreatedAt: time.Now(), + UpdatedAt: time.Now(), + }) + require.NoError(t, err) + defer func() { + chatStore.DeleteMessages(chatID, nil) + chatStore.DeleteChat(chatID) + }() + + // Add history + err = chatStore.SaveMessages(chatID, []*storetypes.Message{ + { + MessageID: fmt.Sprintf("full_overlap_1_%s", reqID), + ChatID: chatID, + RequestID: fmt.Sprintf("req_full_%s", reqID), + Role: "user", + Type: "user_input", + Props: map[string]interface{}{"content": "Exact same message"}, + Sequence: 1, + AssistantID: ast.ID, + CreatedAt: time.Now(), + }, + }) + require.NoError(t, err) + + // Input is exactly the same as history + input := []agentcontext.Message{ + {Role: agentcontext.RoleUser, Content: "Exact same message"}, + } + + result, err := ast.WithHistory(ctx, input, nil) + require.NoError(t, err) + require.NotNil(t, result) + + // Full overlap: clean input should be empty + assert.Empty(t, result.InputMessages) + // FullMessages should be just history (no duplicates) + assert.Len(t, result.FullMessages, 1) + + t.Log("✓ Full overlap handled correctly") + }) + + t.Run("NonExistentChat", func(t *testing.T) { + chatID := "non_existent_chat_12345" + ctx := newHistoryTestContext(chatID) + + input := []agentcontext.Message{ + {Role: agentcontext.RoleUser, Content: "Message to non-existent chat"}, + } + + result, err := ast.WithHistory(ctx, input, nil) + require.NoError(t, err) + require.NotNil(t, result) + + // Should return input as is (no history found) + assert.Equal(t, input, result.InputMessages) + assert.Equal(t, input, result.FullMessages) + + t.Log("✓ Non-existent chat handled gracefully") + }) + + t.Run("MessageWithName", func(t *testing.T) { + chatID := fmt.Sprintf("test_name_%s", uuid.New().String()[:8]) + ctx := newHistoryTestContext(chatID) + reqID := uuid.New().String()[:8] + + // Create chat + err := chatStore.CreateChat(&storetypes.Chat{ + ChatID: chatID, + AssistantID: ast.ID, + Status: "active", + Share: "private", + CreatedAt: time.Now(), + UpdatedAt: time.Now(), + }) + require.NoError(t, err) + defer func() { + chatStore.DeleteMessages(chatID, nil) + chatStore.DeleteChat(chatID) + }() + + // Add message with name + err = chatStore.SaveMessages(chatID, []*storetypes.Message{ + { + MessageID: fmt.Sprintf("name_msg_1_%s", reqID), + ChatID: chatID, + RequestID: fmt.Sprintf("req_name_%s", reqID), + Role: "user", + Type: "user_input", + Props: map[string]interface{}{"content": "Message with name", "name": "John"}, + Sequence: 1, + AssistantID: ast.ID, + CreatedAt: time.Now(), + }, + }) + require.NoError(t, err) + + input := []agentcontext.Message{ + {Role: agentcontext.RoleUser, Content: "New message"}, + } + + result, err := ast.WithHistory(ctx, input, nil) + require.NoError(t, err) + require.NotNil(t, result) + + // First message should have name + assert.Len(t, result.FullMessages, 2) + assert.NotNil(t, result.FullMessages[0].Name) + assert.Equal(t, "John", *result.FullMessages[0].Name) + + t.Log("✓ Message name field preserved") + }) + + t.Run("EmptyContent", func(t *testing.T) { + chatID := fmt.Sprintf("test_empty_content_%s", uuid.New().String()[:8]) + ctx := newHistoryTestContext(chatID) + reqID := uuid.New().String()[:8] + + // Create chat + err := chatStore.CreateChat(&storetypes.Chat{ + ChatID: chatID, + AssistantID: ast.ID, + Status: "active", + Share: "private", + CreatedAt: time.Now(), + UpdatedAt: time.Now(), + }) + require.NoError(t, err) + defer func() { + chatStore.DeleteMessages(chatID, nil) + chatStore.DeleteChat(chatID) + }() + + // Add message with empty content in props + err = chatStore.SaveMessages(chatID, []*storetypes.Message{ + { + MessageID: fmt.Sprintf("empty_content_1_%s", reqID), + ChatID: chatID, + RequestID: fmt.Sprintf("req_empty_content_%s", reqID), + Role: "user", + Type: "user_input", + Props: map[string]interface{}{}, // empty props (no content) + Sequence: 1, + AssistantID: ast.ID, + CreatedAt: time.Now(), + }, + }) + require.NoError(t, err) + + input := []agentcontext.Message{ + {Role: agentcontext.RoleUser, Content: "New message"}, + } + + result, err := ast.WithHistory(ctx, input, nil) + require.NoError(t, err) + require.NotNil(t, result) + + // Message with empty props should be skipped (no content extractable) + // Only new input should be present + assert.Len(t, result.FullMessages, 1) + assert.Equal(t, "New message", result.FullMessages[0].Content) + + t.Log("✓ Empty content handled gracefully (message skipped)") + }) +} diff --git a/agent/assistant/hook/create_bench_test.go b/agent/assistant/hook/create_bench_test.go index 618a8efa..a4062723 100644 --- a/agent/assistant/hook/create_bench_test.go +++ b/agent/assistant/hook/create_bench_test.go @@ -4,7 +4,6 @@ import ( stdContext "context" "testing" - "github.com/yaoapp/gou/plan" "github.com/yaoapp/yao/agent/assistant" "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/testutils" @@ -296,34 +295,32 @@ func getBusinessScenarios() []struct { // newBenchContext creates a minimal context for benchmarking func newBenchContext(chatID, assistantID string) *context.Context { - return &context.Context{ - Context: stdContext.Background(), - Space: plan.NewMemorySharedSpace(), - ChatID: chatID, - AssistantID: assistantID, - Locale: "en-us", - Theme: "light", - Client: context.Client{ - Type: "web", - UserAgent: "BenchAgent/1.0", - IP: "127.0.0.1", - }, - Referer: context.RefererAPI, - Accept: context.AcceptWebCUI, - Route: "", - Metadata: make(map[string]interface{}), - Authorized: &types.AuthorizedInfo{ - Subject: "bench-user", - ClientID: "bench-client", - UserID: "bench-user-123", - TeamID: "bench-team-456", - TenantID: "bench-tenant-789", - Constraints: types.DataConstraints{ - TeamOnly: true, - Extra: map[string]interface{}{ - "department": "engineering", - }, + authorized := &types.AuthorizedInfo{ + Subject: "bench-user", + ClientID: "bench-client", + UserID: "bench-user-123", + TeamID: "bench-team-456", + TenantID: "bench-tenant-789", + Constraints: types.DataConstraints{ + TeamOnly: true, + Extra: map[string]interface{}{ + "department": "engineering", }, }, } + + ctx := context.New(stdContext.Background(), authorized, chatID) + ctx.AssistantID = assistantID + ctx.Locale = "en-us" + ctx.Theme = "light" + ctx.Client = context.Client{ + Type: "web", + UserAgent: "BenchAgent/1.0", + IP: "127.0.0.1", + } + ctx.Referer = context.RefererAPI + ctx.Accept = context.AcceptWebCUI + ctx.Route = "" + ctx.Metadata = make(map[string]interface{}) + return ctx } diff --git a/agent/assistant/hook/create_mem_test.go b/agent/assistant/hook/create_mem_test.go index 3c6f1bd8..abeb4946 100644 --- a/agent/assistant/hook/create_mem_test.go +++ b/agent/assistant/hook/create_mem_test.go @@ -6,7 +6,6 @@ import ( "testing" "time" - "github.com/yaoapp/gou/plan" "github.com/yaoapp/yao/agent/assistant" "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/testutils" @@ -610,34 +609,32 @@ func TestIsolateDisposal(t *testing.T) { // newMemTestContext creates a context for memory leak testing func newMemTestContext(chatID, assistantID string) *context.Context { - return &context.Context{ - Context: stdContext.Background(), - Space: plan.NewMemorySharedSpace(), - ChatID: chatID, - AssistantID: assistantID, - Locale: "en-us", - Theme: "light", - Client: context.Client{ - Type: "web", - UserAgent: "MemTestAgent/1.0", - IP: "127.0.0.1", - }, - Referer: context.RefererAPI, - Accept: context.AcceptWebCUI, - Route: "", - Metadata: make(map[string]interface{}), - Authorized: &types.AuthorizedInfo{ - Subject: "mem-test-user", - ClientID: "mem-test-client", - UserID: "mem-user-123", - TeamID: "mem-team-456", - TenantID: "mem-tenant-789", - Constraints: types.DataConstraints{ - TeamOnly: true, - Extra: map[string]interface{}{ - "department": "engineering", - }, + authorized := &types.AuthorizedInfo{ + Subject: "mem-test-user", + ClientID: "mem-test-client", + UserID: "mem-user-123", + TeamID: "mem-team-456", + TenantID: "mem-tenant-789", + Constraints: types.DataConstraints{ + TeamOnly: true, + Extra: map[string]interface{}{ + "department": "engineering", }, }, } + + ctx := context.New(stdContext.Background(), authorized, chatID) + ctx.AssistantID = assistantID + ctx.Locale = "en-us" + ctx.Theme = "light" + ctx.Client = context.Client{ + Type: "web", + UserAgent: "MemTestAgent/1.0", + IP: "127.0.0.1", + } + ctx.Referer = context.RefererAPI + ctx.Accept = context.AcceptWebCUI + ctx.Route = "" + ctx.Metadata = make(map[string]interface{}) + return ctx } diff --git a/agent/assistant/hook/create_test.go b/agent/assistant/hook/create_test.go index e1f446c9..e459df1f 100644 --- a/agent/assistant/hook/create_test.go +++ b/agent/assistant/hook/create_test.go @@ -4,7 +4,6 @@ import ( stdContext "context" "testing" - "github.com/yaoapp/gou/plan" "github.com/yaoapp/yao/agent/assistant" "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/testutils" @@ -14,44 +13,42 @@ import ( // newTestContext creates a Context for testing with commonly used fields pre-populated. // You can override any fields after creation as needed for specific test scenarios. func newTestContext(chatID, assistantID string) *context.Context { - return &context.Context{ - Context: stdContext.Background(), - Space: plan.NewMemorySharedSpace(), - ChatID: chatID, - AssistantID: assistantID, - Locale: "en-us", - Theme: "light", - Client: context.Client{ - Type: "web", - UserAgent: "TestAgent/1.0", - IP: "127.0.0.1", - }, - Referer: context.RefererAPI, - Accept: context.AcceptWebCUI, - Route: "", - Metadata: make(map[string]interface{}), - Authorized: &types.AuthorizedInfo{ - Subject: "test-user", - ClientID: "test-client-id", - Scope: "openid profile email", - SessionID: "test-session-id", - UserID: "test-user-123", - TeamID: "test-team-456", - TenantID: "test-tenant-789", - RememberMe: true, - Constraints: types.DataConstraints{ - OwnerOnly: false, - CreatorOnly: false, - EditorOnly: false, - TeamOnly: true, - Extra: map[string]interface{}{ - "department": "engineering", - "region": "us-west", - "project": "yao", - }, + authorized := &types.AuthorizedInfo{ + Subject: "test-user", + ClientID: "test-client-id", + Scope: "openid profile email", + SessionID: "test-session-id", + UserID: "test-user-123", + TeamID: "test-team-456", + TenantID: "test-tenant-789", + RememberMe: true, + Constraints: types.DataConstraints{ + OwnerOnly: false, + CreatorOnly: false, + EditorOnly: false, + TeamOnly: true, + Extra: map[string]interface{}{ + "department": "engineering", + "region": "us-west", + "project": "yao", }, }, } + + ctx := context.New(stdContext.Background(), authorized, chatID) + ctx.AssistantID = assistantID + ctx.Locale = "en-us" + ctx.Theme = "light" + ctx.Client = context.Client{ + Type: "web", + UserAgent: "TestAgent/1.0", + IP: "127.0.0.1", + } + ctx.Referer = context.RefererAPI + ctx.Accept = context.AcceptWebCUI + ctx.Route = "" + ctx.Metadata = make(map[string]interface{}) + return ctx } // TestCreate test the create hook diff --git a/agent/assistant/hook/goroutine_leak_test.go b/agent/assistant/hook/goroutine_leak_test.go index 9673001e..fc26f6b4 100644 --- a/agent/assistant/hook/goroutine_leak_test.go +++ b/agent/assistant/hook/goroutine_leak_test.go @@ -10,7 +10,6 @@ import ( "testing" "time" - "github.com/yaoapp/gou/plan" "github.com/yaoapp/yao/agent/assistant" "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/testutils" @@ -294,34 +293,32 @@ func truncate(s string, max int) string { } func newLeakTestContext(chatID, assistantID string) *context.Context { - return &context.Context{ - Context: stdContext.Background(), - Space: plan.NewMemorySharedSpace(), - ChatID: chatID, - AssistantID: assistantID, - Locale: "en-us", - Theme: "light", - Client: context.Client{ - Type: "web", - UserAgent: "LeakTestAgent/1.0", - IP: "127.0.0.1", - }, - Referer: context.RefererAPI, - Accept: context.AcceptWebCUI, - Route: "", - Metadata: make(map[string]interface{}), - Authorized: &types.AuthorizedInfo{ - Subject: "leak-test-user", - ClientID: "leak-test-client", - UserID: "leak-user-123", - TeamID: "leak-team-456", - TenantID: "leak-tenant-789", - Constraints: types.DataConstraints{ - TeamOnly: true, - Extra: map[string]interface{}{ - "department": "testing", - }, + authorized := &types.AuthorizedInfo{ + Subject: "leak-test-user", + ClientID: "leak-test-client", + UserID: "leak-user-123", + TeamID: "leak-team-456", + TenantID: "leak-tenant-789", + Constraints: types.DataConstraints{ + TeamOnly: true, + Extra: map[string]interface{}{ + "department": "testing", }, }, } + + ctx := context.New(stdContext.Background(), authorized, chatID) + ctx.AssistantID = assistantID + ctx.Locale = "en-us" + ctx.Theme = "light" + ctx.Client = context.Client{ + Type: "web", + UserAgent: "LeakTestAgent/1.0", + IP: "127.0.0.1", + } + ctx.Referer = context.RefererAPI + ctx.Accept = context.AcceptWebCUI + ctx.Route = "" + ctx.Metadata = make(map[string]interface{}) + return ctx } diff --git a/agent/assistant/hook/next_test.go b/agent/assistant/hook/next_test.go index becaedb2..15f4d16b 100644 --- a/agent/assistant/hook/next_test.go +++ b/agent/assistant/hook/next_test.go @@ -4,7 +4,6 @@ import ( stdContext "context" "testing" - "github.com/yaoapp/gou/plan" "github.com/yaoapp/yao/agent/assistant" "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/output/message" @@ -15,44 +14,42 @@ import ( // newTestContextForNext creates a Context for testing Next Hook with commonly used fields pre-populated. // You can override any fields after creation as needed for specific test scenarios. func newTestContextForNext(chatID, assistantID string) *context.Context { - return &context.Context{ - Context: stdContext.Background(), - Space: plan.NewMemorySharedSpace(), - ChatID: chatID, - AssistantID: assistantID, - Locale: "en-us", - Theme: "light", - Client: context.Client{ - Type: "web", - UserAgent: "TestAgent/1.0", - IP: "127.0.0.1", - }, - Referer: context.RefererAPI, - Accept: context.AcceptWebCUI, - Route: "", - Metadata: make(map[string]interface{}), - Authorized: &types.AuthorizedInfo{ - Subject: "test-user", - ClientID: "test-client-id", - Scope: "openid profile email", - SessionID: "test-session-id", - UserID: "test-user-123", - TeamID: "test-team-456", - TenantID: "test-tenant-789", - RememberMe: true, - Constraints: types.DataConstraints{ - OwnerOnly: false, - CreatorOnly: false, - EditorOnly: false, - TeamOnly: true, - Extra: map[string]interface{}{ - "department": "engineering", - "region": "us-west", - "project": "yao", - }, + authorized := &types.AuthorizedInfo{ + Subject: "test-user", + ClientID: "test-client-id", + Scope: "openid profile email", + SessionID: "test-session-id", + UserID: "test-user-123", + TeamID: "test-team-456", + TenantID: "test-tenant-789", + RememberMe: true, + Constraints: types.DataConstraints{ + OwnerOnly: false, + CreatorOnly: false, + EditorOnly: false, + TeamOnly: true, + Extra: map[string]interface{}{ + "department": "engineering", + "region": "us-west", + "project": "yao", }, }, } + + ctx := context.New(stdContext.Background(), authorized, chatID) + ctx.AssistantID = assistantID + ctx.Locale = "en-us" + ctx.Theme = "light" + ctx.Client = context.Client{ + Type: "web", + UserAgent: "TestAgent/1.0", + IP: "127.0.0.1", + } + ctx.Referer = context.RefererAPI + ctx.Accept = context.AcceptWebCUI + ctx.Route = "" + ctx.Metadata = make(map[string]interface{}) + return ctx } // TestNext tests the Next hook diff --git a/agent/assistant/hook/realworld_next_test.go b/agent/assistant/hook/realworld_next_test.go index 1142b749..474d3fc4 100644 --- a/agent/assistant/hook/realworld_next_test.go +++ b/agent/assistant/hook/realworld_next_test.go @@ -5,7 +5,6 @@ import ( "testing" "github.com/stretchr/testify/assert" - "github.com/yaoapp/gou/plan" "github.com/yaoapp/yao/agent/assistant" "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/testutils" @@ -14,32 +13,30 @@ import ( // newRealWorldNextContext creates a Context for real world Next Hook testing func newRealWorldNextContext(chatID, assistantID string) *context.Context { - return &context.Context{ - Context: stdContext.Background(), - Space: plan.NewMemorySharedSpace(), - ChatID: chatID, - AssistantID: assistantID, - Locale: "en-us", - Theme: "light", - Client: context.Client{ - Type: "web", - UserAgent: "RealWorldTest/1.0", - IP: "127.0.0.1", - }, - Referer: context.RefererAPI, - Accept: context.AcceptWebCUI, - Route: "", - Metadata: make(map[string]interface{}), - Authorized: &types.AuthorizedInfo{ - Subject: "realworld-test-user", - ClientID: "realworld-test-client", - Scope: "openid profile", - SessionID: "realworld-test-session", - UserID: "realworld-user-123", - TeamID: "realworld-team-456", - TenantID: "realworld-tenant-789", - }, + authorized := &types.AuthorizedInfo{ + Subject: "realworld-test-user", + ClientID: "realworld-test-client", + Scope: "openid profile", + SessionID: "realworld-test-session", + UserID: "realworld-user-123", + TeamID: "realworld-team-456", + TenantID: "realworld-tenant-789", } + + ctx := context.New(stdContext.Background(), authorized, chatID) + ctx.AssistantID = assistantID + ctx.Locale = "en-us" + ctx.Theme = "light" + ctx.Client = context.Client{ + Type: "web", + UserAgent: "RealWorldTest/1.0", + IP: "127.0.0.1", + } + ctx.Referer = context.RefererAPI + ctx.Accept = context.AcceptWebCUI + ctx.Route = "" + ctx.Metadata = make(map[string]interface{}) + return ctx } // TestRealWorldNextStandard tests standard response (nil return) diff --git a/agent/assistant/hook/realworld_stress_test.go b/agent/assistant/hook/realworld_stress_test.go index fa170a2a..b86c30d8 100644 --- a/agent/assistant/hook/realworld_stress_test.go +++ b/agent/assistant/hook/realworld_stress_test.go @@ -9,7 +9,6 @@ import ( "time" "github.com/stretchr/testify/assert" - "github.com/yaoapp/gou/plan" "github.com/yaoapp/yao/agent/assistant" "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/testutils" @@ -721,44 +720,42 @@ func TestRealWorldStressResourceHeavy(t *testing.T) { // newRealWorldContext creates a Context for real-world testing func newRealWorldContext(chatID, assistantID string) *context.Context { - return &context.Context{ - Context: stdContext.Background(), - Space: plan.NewMemorySharedSpace(), - ChatID: chatID, - AssistantID: assistantID, - Locale: "en-us", - Theme: "light", - Client: context.Client{ - Type: "web", - UserAgent: "RealWorldTest/1.0", - IP: "127.0.0.1", - }, - Referer: context.RefererAPI, - Accept: context.AcceptWebCUI, - Route: "", - Metadata: make(map[string]interface{}), - Authorized: &types.AuthorizedInfo{ - Subject: "realworld-test-user", - ClientID: "realworld-test-client", - Scope: "openid profile email", - SessionID: "realworld-test-session", - UserID: "realworld-user-123", - TeamID: "realworld-team-456", - TenantID: "realworld-tenant-789", - RememberMe: true, - Constraints: types.DataConstraints{ - OwnerOnly: false, - CreatorOnly: false, - EditorOnly: false, - TeamOnly: true, - Extra: map[string]interface{}{ - "department": "engineering", - "region": "us-west", - "project": "yao-realworld-test", - }, + authorized := &types.AuthorizedInfo{ + Subject: "realworld-test-user", + ClientID: "realworld-test-client", + Scope: "openid profile email", + SessionID: "realworld-test-session", + UserID: "realworld-user-123", + TeamID: "realworld-team-456", + TenantID: "realworld-tenant-789", + RememberMe: true, + Constraints: types.DataConstraints{ + OwnerOnly: false, + CreatorOnly: false, + EditorOnly: false, + TeamOnly: true, + Extra: map[string]interface{}{ + "department": "engineering", + "region": "us-west", + "project": "yao-realworld-test", }, }, } + + ctx := context.New(stdContext.Background(), authorized, chatID) + ctx.AssistantID = assistantID + ctx.Locale = "en-us" + ctx.Theme = "light" + ctx.Client = context.Client{ + Type: "web", + UserAgent: "RealWorldTest/1.0", + IP: "127.0.0.1", + } + ctx.Referer = context.RefererAPI + ctx.Accept = context.AcceptWebCUI + ctx.Route = "" + ctx.Metadata = make(map[string]interface{}) + return ctx } // getMemStats returns current memory allocation in bytes diff --git a/agent/assistant/llm.go b/agent/assistant/llm.go index eab35e2e..415f3ee7 100644 --- a/agent/assistant/llm.go +++ b/agent/assistant/llm.go @@ -1,9 +1,6 @@ package assistant import ( - "fmt" - - "github.com/yaoapp/kun/log" "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/llm" "github.com/yaoapp/yao/agent/output/message" @@ -21,12 +18,6 @@ func (ast *Assistant) executeLLMStream( opts *context.Options, ) (*context.CompletionResponse, error) { - // === Debug LLM Stream Start === - fmt.Println(">>> executeLLMStream: STARTING") - fmt.Printf(">>> Messages count: %d\n", len(completionMessages)) - fmt.Printf(">>> Tools count: %d\n", len(completionOptions.Tools)) - // === End Debug === - // Get connector object (capabilities were already set above, before stream_start) conn, capabilities, err := ast.GetConnector(ctx, opts) if err != nil { @@ -45,6 +36,9 @@ func (ast *Assistant) executeLLMStream( // Trace Add LLM request ast.traceLLMRequest(ctx, conn.ID(), completionMessages, completionOptions) + // Log LLM call start + ctx.Logger.LLMStart(conn.ID(), "", len(completionMessages)) + // Create LLM instance with connector and options llmInstance, err := llm.New(conn, completionOptions) if err != nil { @@ -54,25 +48,9 @@ func (ast *Assistant) executeLLMStream( } // Call the LLM Completion Stream (streamHandler was set earlier) - log.Trace("[AGENT] Calling LLM Stream: assistant=%s", ast.ID) - - // === Debug LLM Stream Call === - fmt.Println(">>> executeLLMStream: CALLING llmInstance.Stream()") - // === End Debug === - completionResponse, err := llmInstance.Stream(ctx, completionMessages, completionOptions, streamHandler) - // === Debug LLM Stream Return === - fmt.Println(">>> executeLLMStream: llmInstance.Stream() RETURNED") - fmt.Printf(">>> err: %v\n", err) - if completionResponse != nil { - fmt.Printf(">>> ToolCalls: %d\n", len(completionResponse.ToolCalls)) - } - // === End Debug === - - log.Trace("[AGENT] LLM Stream returned: assistant=%s, err=%v", ast.ID, err) if err != nil { - log.Trace("[AGENT] Calling sendStreamEndOnError") // Mark LLM Request as failed in trace ast.traceLLMFail(ctx, err) return nil, err @@ -111,6 +89,9 @@ func (ast *Assistant) executeLLMForToolRetry( // Trace Add LLM retry request ast.traceLLMRetryRequest(ctx, conn.ID(), completionMessages, completionOptions) + // Log LLM call start (retry) + ctx.Logger.LLMStart(conn.ID(), "", len(completionMessages)) + // Create LLM instance with connector and options llmInstance, err := llm.New(conn, completionOptions) if err != nil { @@ -120,9 +101,7 @@ func (ast *Assistant) executeLLMForToolRetry( } // Call the LLM Completion Stream (still streaming for tool retry) - log.Trace("[AGENT] Calling LLM Stream for tool retry: assistant=%s", ast.ID) completionResponse, err := llmInstance.Stream(ctx, completionMessages, completionOptions, streamHandler) - log.Trace("[AGENT] LLM tool retry stream returned: assistant=%s, err=%v", ast.ID, err) if err != nil { // Mark LLM Retry Request as failed in trace ast.traceLLMFail(ctx, err) diff --git a/agent/assistant/load.go b/agent/assistant/load.go index 2ac9974d..6d62130f 100644 --- a/agent/assistant/load.go +++ b/agent/assistant/load.go @@ -21,6 +21,7 @@ import ( // loaded the loaded assistant var loaded = NewCache(200) // 200 is the default capacity var storage store.Store = nil +var storeSetting *store.Setting = nil // store setting from agent.yml var search interface{} = nil var modelCapabilities map[string]gouOpenAI.Capabilities = map[string]gouOpenAI.Capabilities{} var defaultConnector string = "" // default connector @@ -153,6 +154,16 @@ func SetGlobalPrompts(prompts []store.Prompt) { globalPrompts = prompts } +// SetStoreSetting set the store setting from agent.yml +func SetStoreSetting(setting *store.Setting) { + storeSetting = setting +} + +// GetStoreSetting returns the store setting +func GetStoreSetting() *store.Setting { + return storeSetting +} + // GetGlobalPrompts returns the global prompts with variables parsed // ctx: context variables for parsing $CTX.* variables func GetGlobalPrompts(ctx map[string]string) []store.Prompt { diff --git a/agent/assistant/load_store_test.go b/agent/assistant/load_store_test.go index e8dfb173..75d376ac 100644 --- a/agent/assistant/load_store_test.go +++ b/agent/assistant/load_store_test.go @@ -176,31 +176,30 @@ func TestLoadStoreWithoutSource(t *testing.T) { // newStoreTestContext creates a Context for testing with commonly used fields pre-populated. func newStoreTestContext(chatID, assistantID string) *context.Context { - return &context.Context{ - Context: stdContext.Background(), - ChatID: chatID, - AssistantID: assistantID, - Locale: "en-us", - Theme: "light", - Client: context.Client{ - Type: "web", - UserAgent: "TestAgent/1.0", - IP: "127.0.0.1", - }, - Referer: context.RefererAPI, - Accept: context.AcceptWebCUI, - Route: "", - Metadata: make(map[string]interface{}), - Authorized: &types.AuthorizedInfo{ - Subject: "test-user", - ClientID: "test-client-id", - Scope: "openid profile email", - SessionID: "test-session-id", - UserID: "test-user-123", - TeamID: "test-team-456", - TenantID: "test-tenant-789", - }, + authorized := &types.AuthorizedInfo{ + Subject: "test-user", + ClientID: "test-client-id", + Scope: "openid profile email", + SessionID: "test-session-id", + UserID: "test-user-123", + TeamID: "test-team-456", + TenantID: "test-tenant-789", } + + ctx := context.New(stdContext.Background(), authorized, chatID) + ctx.AssistantID = assistantID + ctx.Locale = "en-us" + ctx.Theme = "light" + ctx.Client = context.Client{ + Type: "web", + UserAgent: "TestAgent/1.0", + IP: "127.0.0.1", + } + ctx.Referer = context.RefererAPI + ctx.Accept = context.AcceptWebCUI + ctx.Route = "" + ctx.Metadata = make(map[string]interface{}) + return ctx } // TestLoadStoreWithSourceExecuteHook tests that Source-based script is properly compiled and can execute diff --git a/agent/assistant/mcp.go b/agent/assistant/mcp.go index 53c5df85..1e544e00 100644 --- a/agent/assistant/mcp.go +++ b/agent/assistant/mcp.go @@ -9,7 +9,6 @@ import ( gouJson "github.com/yaoapp/gou/json" "github.com/yaoapp/gou/mcp" mcpTypes "github.com/yaoapp/gou/mcp/types" - "github.com/yaoapp/kun/log" agentContext "github.com/yaoapp/yao/agent/context" storeTypes "github.com/yaoapp/yao/agent/store/types" "github.com/yaoapp/yao/trace/types" @@ -109,21 +108,21 @@ func (ast *Assistant) buildMCPTools(ctx *agentContext.Context, createResponse *a // Process each MCP server in order for _, serverConfig := range servers { if len(allTools) >= MaxMCPTools { - log.Warn("[Assistant MCP] Reached maximum tool limit (%d), skipping remaining servers", MaxMCPTools) + ctx.Logger.Warn("Reached maximum tool limit (%d), skipping remaining servers", MaxMCPTools) break } // Get MCP client client, err := mcp.Select(serverConfig.ServerID) if err != nil { - log.Warn("[Assistant MCP] Failed to select MCP client '%s': %v", serverConfig.ServerID, err) + ctx.Logger.Warn("Failed to select MCP client '%s': %v", serverConfig.ServerID, err) continue } // Get tools list (filter by serverConfig.Tools if specified) toolsResponse, err := client.ListTools(mcpCtx, "") if err != nil { - log.Warn("[Assistant MCP] Failed to list tools for '%s': %v", serverConfig.ServerID, err) + ctx.Logger.Warn("Failed to list tools for '%s': %v", serverConfig.ServerID, err) continue } @@ -204,7 +203,7 @@ func (ast *Assistant) buildMCPTools(ctx *agentContext.Context, createResponse *a } } - log.Trace("[Assistant MCP] Loaded %d tools from server '%s'", len(toolsResponse.Tools), serverConfig.ServerID) + ctx.Logger.Debug("Loaded %d tools from server '%s'", len(toolsResponse.Tools), serverConfig.ServerID) } samplesPrompt := "" @@ -212,7 +211,7 @@ func (ast *Assistant) buildMCPTools(ctx *agentContext.Context, createResponse *a samplesPrompt = samplesBuilder.String() } - log.Trace("[Assistant MCP] Total MCP tools loaded: %d", len(allTools)) + ctx.Logger.Debug("Total MCP tools loaded: %d", len(allTools)) return allTools, samplesPrompt, nil } @@ -226,32 +225,20 @@ func (ast *Assistant) executeToolCalls(ctx *agentContext.Context, toolCalls []ag return nil, false } - // === Debug === - fmt.Printf(">>> executeToolCalls: START (attempt %d, toolCalls count: %d)\n", attempt, len(toolCalls)) - // === End Debug === - - log.Trace("[Assistant MCP] Executing %d tool calls (attempt %d)", len(toolCalls), attempt) + ctx.Logger.Debug("Executing %d tool calls (attempt %d)", len(toolCalls), attempt) // Single tool call if len(toolCalls) == 1 { - fmt.Println(">>> executeToolCalls: Calling executeSingleToolCall") - results, hasErrors := ast.executeSingleToolCall(ctx, toolCalls[0]) - fmt.Printf(">>> executeToolCalls: executeSingleToolCall RETURNED (hasErrors: %v)\n", hasErrors) - return results, hasErrors + return ast.executeSingleToolCall(ctx, toolCalls[0]) } // Multiple tool calls - try parallel first - fmt.Println(">>> executeToolCalls: Calling executeMultipleToolCallsParallel") - results, hasErrors := ast.executeMultipleToolCallsParallel(ctx, toolCalls) - fmt.Printf(">>> executeToolCalls: executeMultipleToolCallsParallel RETURNED (hasErrors: %v)\n", hasErrors) - return results, hasErrors + return ast.executeMultipleToolCallsParallel(ctx, toolCalls) } // executeSingleToolCall executes a single tool call with trace logging func (ast *Assistant) executeSingleToolCall(ctx *agentContext.Context, toolCall agentContext.ToolCall) ([]ToolCallResult, bool) { - // === Debug === - fmt.Printf(">>> executeSingleToolCall: START (tool: %s)\n", toolCall.Function.Name) - // === End Debug === + ctx.Logger.ToolStart(toolCall.Function.Name) trace, _ := ctx.Trace() @@ -267,12 +254,12 @@ func (ast *Assistant) executeSingleToolCall(ctx *agentContext.Context, toolCall } // Parse tool name - fmt.Println(">>> executeSingleToolCall: Parsing tool name") serverID, toolName, ok := ParseMCPToolName(toolCall.Function.Name) if !ok { result.Error = fmt.Errorf("invalid MCP tool name format: %s", toolCall.Function.Name) result.Content = result.Error.Error() - log.Error("[Assistant MCP] %v", result.Error) + ctx.Logger.Error("Invalid MCP tool name format: %s", toolCall.Function.Name) + ctx.Logger.ToolComplete(toolCall.Function.Name, false) return []ToolCallResult{result}, true } @@ -282,7 +269,8 @@ func (ast *Assistant) executeSingleToolCall(ctx *agentContext.Context, toolCall result.Error = fmt.Errorf("failed to select MCP client '%s': %w", serverID, err) result.Content = result.Error.Error() result.IsRetryableError = false // MCP client selection error is not retryable - log.Error("[Assistant MCP] %v", result.Error) + ctx.Logger.Error("Failed to select MCP client '%s': %v", serverID, err) + ctx.Logger.ToolComplete(toolCall.Function.Name, false) return []ToolCallResult{result}, true } @@ -330,7 +318,7 @@ func (ast *Assistant) executeSingleToolCall(ctx *agentContext.Context, toolCall result.Error = fmt.Errorf("failed to parse arguments: %w", err) result.Content = result.Error.Error() result.IsRetryableError = true // Argument parsing error is retryable by LLM - log.Error("[Assistant MCP] %v", result.Error) + ctx.Logger.Error("Failed to parse arguments: %v", err) if toolNode != nil { toolNode.Fail(result.Error) } @@ -344,7 +332,7 @@ func (ast *Assistant) executeSingleToolCall(ctx *agentContext.Context, toolCall result.Error = fmt.Errorf("arguments must be an object, got %T", parsed) result.Content = result.Error.Error() result.IsRetryableError = true // Type error is retryable by LLM - log.Error("[Assistant MCP] %v", result.Error) + ctx.Logger.Error("Arguments must be an object, got %T", parsed) if toolNode != nil { toolNode.Fail(result.Error) } @@ -353,42 +341,34 @@ func (ast *Assistant) executeSingleToolCall(ctx *agentContext.Context, toolCall // Validate arguments against tool schema if available if toolSchema != nil { - fmt.Println(">>> executeSingleToolCall: Validating arguments against schema") if err := gouJson.Validate(args, toolSchema); err != nil { - fmt.Printf(">>> executeSingleToolCall: Validation FAILED: %v\n", err) result.Error = fmt.Errorf("argument validation failed: %w", err) result.Content = result.Error.Error() result.IsRetryableError = true // Validation error is retryable by LLM - log.Error("[Assistant MCP] %v", result.Error) + ctx.Logger.Error("Argument validation failed: %v", err) if toolNode != nil { - fmt.Println(">>> executeSingleToolCall: Failing toolNode due to validation error") toolNode.Fail(result.Error) - fmt.Println(">>> executeSingleToolCall: toolNode.Fail() finished") } - fmt.Println(">>> executeSingleToolCall: RETURNING with validation error") return []ToolCallResult{result}, true } - fmt.Println(">>> executeSingleToolCall: Validation PASSED") } } // Call the tool with agent context as extra argument - log.Trace("[Assistant MCP] Calling tool: %s (server: %s)", toolName, serverID) - fmt.Printf(">>> executeSingleToolCall: CALLING client.CallTool (tool: %s, server: %s)\n", toolName, serverID) + ctx.Logger.Debug("Calling tool: %s (server: %s)", toolName, serverID) // Pass agent context as extra argument (only used for Process transport) callResult, err := client.CallTool(mcpCtx, toolName, args, ctx) - fmt.Printf(">>> executeSingleToolCall: client.CallTool RETURNED (err: %v)\n", err) if err != nil { result.Error = fmt.Errorf("tool call failed: %w", err) result.Content = result.Error.Error() // Check if error is retryable (parameter/validation errors) result.IsRetryableError = isRetryableToolError(err) - log.Error("[Assistant MCP] Tool call failed: %v (retryable: %v)", err, result.IsRetryableError) + ctx.Logger.Error("Tool call failed: %v (retryable: %v)", err, result.IsRetryableError) + ctx.Logger.ToolComplete(toolCall.Function.Name, false) if toolNode != nil { toolNode.Fail(result.Error) } - fmt.Println(">>> executeSingleToolCall: RETURNING with error") return []ToolCallResult{result}, true } @@ -404,7 +384,8 @@ func (ast *Assistant) executeSingleToolCall(ctx *agentContext.Context, toolCall result.Error = fmt.Errorf("failed to serialize result: %w", err) result.Content = result.Error.Error() result.IsRetryableError = false - log.Error("[Assistant MCP] %v", result.Error) + ctx.Logger.Error("Failed to serialize result: %v", err) + ctx.Logger.ToolComplete(toolCall.Function.Name, false) if toolNode != nil { toolNode.Fail(result.Error) } @@ -412,17 +393,14 @@ func (ast *Assistant) executeSingleToolCall(ctx *agentContext.Context, toolCall } result.Content = string(contentBytes) - log.Trace("[Assistant MCP] Tool call succeeded: %s", toolName) + ctx.Logger.ToolComplete(toolName, true) if toolNode != nil { - fmt.Println(">>> executeSingleToolCall: Completing toolNode") toolNode.Complete(map[string]any{ "result": callResult, }) - fmt.Println(">>> executeSingleToolCall: toolNode.Complete() finished") } - fmt.Println(">>> executeSingleToolCall: RETURNING success") return []ToolCallResult{result}, false } @@ -441,7 +419,7 @@ func (ast *Assistant) executeMultipleToolCallsParallel(ctx *agentContext.Context for _, tc := range toolCalls { serverID, _, ok := ParseMCPToolName(tc.Function.Name) if !ok { - log.Warn("[Assistant MCP] Invalid tool name format: %s", tc.Function.Name) + ctx.Logger.Warn("Invalid tool name format: %s", tc.Function.Name) continue } serverGroups[serverID] = append(serverGroups[serverID], tc) @@ -454,7 +432,7 @@ func (ast *Assistant) executeMultipleToolCallsParallel(ctx *agentContext.Context for serverID, calls := range serverGroups { client, err := mcp.Select(serverID) if err != nil { - log.Error("[Assistant MCP] Failed to select MCP client '%s': %v", serverID, err) + ctx.Logger.Error("Failed to select MCP client '%s': %v", serverID, err) // Add error results for all calls to this server for _, tc := range calls { results = append(results, ToolCallResult{ @@ -475,7 +453,7 @@ func (ast *Assistant) executeMultipleToolCallsParallel(ctx *agentContext.Context // If parallel execution failed with retryable error, try sequential if serverHasErrors && ast.shouldRetrySequential(serverResults) { - log.Warn("[Assistant MCP] Parallel execution had parameter errors for server '%s', retrying sequentially", serverID) + ctx.Logger.Warn("Parallel execution had parameter errors for server '%s', retrying sequentially", serverID) serverResults, serverHasErrors = ast.executeServerToolsSequentialWithTrace( mcpCtx, ctx, trace, client, serverID, calls, ) @@ -575,7 +553,7 @@ func (ast *Assistant) executeServerToolsParallelWithTrace(mcpCtx context.Context var args map[string]interface{} if tc.Function.Arguments != "" { if err := jsoniter.UnmarshalFromString(tc.Function.Arguments, &args); err != nil { - log.Error("[Assistant MCP] Failed to parse arguments for %s: %v", toolName, err) + ctx.Logger.Error("Failed to parse arguments for %s: %v", toolName, err) continue } } @@ -606,25 +584,20 @@ func (ast *Assistant) executeServerToolsParallelWithTrace(mcpCtx context.Context // Create parallel trace nodes var toolNodes []types.Node if trace != nil && len(parallelInputs) > 0 { - fmt.Printf(">>> executeServerToolsParallelWithTrace: Creating %d parallel trace nodes\n", len(parallelInputs)) var err error toolNodes, err = trace.Parallel(parallelInputs) if err != nil { - fmt.Printf(">>> executeServerToolsParallelWithTrace: trace.Parallel() FAILED: %v\n", err) - } else { - fmt.Printf(">>> executeServerToolsParallelWithTrace: Created %d trace nodes\n", len(toolNodes)) + ctx.Logger.Debug("trace.Parallel() failed: %v", err) } - } else { - fmt.Printf(">>> executeServerToolsParallelWithTrace: NOT creating trace nodes (trace: %v, inputs: %d)\n", trace != nil, len(parallelInputs)) } // Call tools in parallel with agent context as extra argument - log.Trace("[Assistant MCP] Calling %d tools in parallel on server '%s'", len(mcpCalls), serverID) + ctx.Logger.Debug("Calling %d tools in parallel on server '%s'", len(mcpCalls), serverID) // Pass agent context as extra argument (only used for Process transport) mcpResponse, err := client.CallToolsParallel(mcpCtx, mcpCalls, ctx) if err != nil { - log.Error("[Assistant MCP] Parallel call failed: %v", err) + ctx.Logger.Error("Parallel call failed: %v", err) // Mark all trace nodes as failed for _, node := range toolNodes { if node != nil { @@ -669,20 +642,16 @@ func (ast *Assistant) executeServerToolsParallelWithTrace(mcpCtx context.Context result.Error = fmt.Errorf("tool call error: %s", result.Content) result.IsRetryableError = isRetryableToolError(result.Error) hasErrors = true - log.Error("[Assistant MCP] Tool call failed: %s - %s (retryable: %v)", toolName, result.Content, result.IsRetryableError) + ctx.Logger.Error("Tool call failed: %s - %s (retryable: %v)", toolName, result.Content, result.IsRetryableError) if toolNode != nil { toolNode.Fail(result.Error) } } else { // Success if toolNode != nil { - fmt.Printf(">>> executeServerToolsParallelWithTrace: Completing toolNode %d\n", i) toolNode.Complete(map[string]any{ "result": mcpResult.Content, }) - fmt.Printf(">>> executeServerToolsParallelWithTrace: toolNode %d completed\n", i) - } else { - fmt.Printf(">>> executeServerToolsParallelWithTrace: toolNode %d is nil!\n", i) } } } @@ -698,7 +667,7 @@ func (ast *Assistant) executeServerToolsSequentialWithTrace(mcpCtx context.Conte results := make([]ToolCallResult, 0, len(toolCalls)) hasErrors := false - log.Trace("[Assistant MCP] Calling %d tools sequentially on server '%s'", len(toolCalls), serverID) + ctx.Logger.Debug("Calling %d tools sequentially on server '%s'", len(toolCalls), serverID) for _, tc := range toolCalls { _, toolName, ok := ParseMCPToolName(tc.Function.Name) @@ -805,7 +774,7 @@ func (ast *Assistant) executeServerToolsSequentialWithTrace(mcpCtx context.Conte } // Call single tool with agent context as extra argument - log.Trace("[Assistant MCP] Calling tool: %s", toolName) + ctx.Logger.Debug("Calling tool: %s", toolName) mcpResult, err := client.CallTool(mcpCtx, toolName, args, ctx) result := ToolCallResult{ @@ -818,7 +787,7 @@ func (ast *Assistant) executeServerToolsSequentialWithTrace(mcpCtx context.Conte result.Content = fmt.Sprintf("Tool call failed: %v", err) result.IsRetryableError = isRetryableToolError(err) hasErrors = true - log.Error("[Assistant MCP] Tool call failed: %s - %v (retryable: %v)", toolName, err, result.IsRetryableError) + ctx.Logger.Error("Tool call failed: %s - %v (retryable: %v)", toolName, err, result.IsRetryableError) if toolNode != nil { toolNode.Fail(err) } diff --git a/agent/context/authorized_test.go b/agent/context/authorized_test.go index 47eeeea2..ded93ff1 100644 --- a/agent/context/authorized_test.go +++ b/agent/context/authorized_test.go @@ -1,10 +1,11 @@ -package context +package context_test import ( - "context" + stdContext "context" "testing" "github.com/stretchr/testify/assert" + "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/config" "github.com/yaoapp/yao/openapi/oauth/types" "github.com/yaoapp/yao/test" @@ -23,7 +24,8 @@ func TestContextNew_PreservesAuthorizedInfo(t *testing.T) { } // Create context using New() - ctx := New(context.Background(), authInfo, "test-chat-123") + ctx := context.New(stdContext.Background(), authInfo, "test-chat-123") + defer ctx.Release() // Verify authorized info is preserved assert.NotNil(t, ctx) @@ -45,9 +47,15 @@ func TestContextTrace_SavesAuthorizedInfo(t *testing.T) { TenantID: "tenant-001", } - // Create context - ctx := New(context.Background(), authInfo, "test-chat-456") + // Create context using New + ctx := context.New(stdContext.Background(), authInfo, "test-chat-456") ctx.AssistantID = "test-assistant" + ctx.Referer = context.RefererAPI + + // Initialize stack (required for trace) + stack, _, done := context.EnterStack(ctx, "test-assistant", &context.Options{}) + ctx.Stack = stack + defer done() // Initialize trace manager, err := ctx.Trace() @@ -67,7 +75,7 @@ func TestContextTrace_SavesAuthorizedInfo(t *testing.T) { // Clean up if ctx.Stack != nil && ctx.Stack.TraceID != "" { trace.Release(ctx.Stack.TraceID) - trace.Remove(context.Background(), trace.Local, ctx.Stack.TraceID) + trace.Remove(stdContext.Background(), trace.Local, ctx.Stack.TraceID) } } @@ -76,7 +84,8 @@ func TestContextNew_NilAuthorized(t *testing.T) { defer test.Clean() // Create context with nil authorized info (should not panic) - ctx := New(context.Background(), nil, "test-chat-789") + ctx := context.New(stdContext.Background(), nil, "test-chat-789") + defer ctx.Release() assert.NotNil(t, ctx) assert.Nil(t, ctx.Authorized) diff --git a/agent/context/chat_test.go b/agent/context/chat_test.go index f9adbf9f..f3125bb5 100644 --- a/agent/context/chat_test.go +++ b/agent/context/chat_test.go @@ -1,9 +1,10 @@ -package context +package context_test import ( "testing" "github.com/yaoapp/gou/store" + "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/config" "github.com/yaoapp/yao/test" ) @@ -23,15 +24,15 @@ func TestGetChatIDByMessages_NewConversation(t *testing.T) { cache := getTestCache(t) - messages := []Message{ + messages := []context.Message{ { - Role: RoleUser, + Role: context.RoleUser, Content: "Hello, how are you?", }, } // First request - should generate new chat ID - chatID1, err := GetChatIDByMessages(cache, messages) + chatID1, err := context.GetChatIDByMessages(cache, messages) if err != nil { t.Fatalf("Failed to get chat ID: %v", err) } @@ -42,7 +43,7 @@ func TestGetChatIDByMessages_NewConversation(t *testing.T) { // Second request with same single user message - should generate DIFFERENT chat ID // (single user message always generates new chat ID to avoid false matches) - chatID2, err := GetChatIDByMessages(cache, messages) + chatID2, err := context.GetChatIDByMessages(cache, messages) if err != nil { t.Fatalf("Failed to get chat ID: %v", err) } @@ -65,10 +66,10 @@ func TestGetChatIDByMessages_ContinuousConversation(t *testing.T) { // Scenario: User conversation with incrementally added messages // Request 1: [user1] - messages1 := []Message{ - {Role: RoleUser, Content: "First message"}, + messages1 := []context.Message{ + {Role: context.RoleUser, Content: "First message"}, } - chatID1, err := GetChatIDByMessages(cache, messages1) + chatID1, err := context.GetChatIDByMessages(cache, messages1) if err != nil { t.Fatalf("Failed to get chat ID: %v", err) } @@ -76,11 +77,11 @@ func TestGetChatIDByMessages_ContinuousConversation(t *testing.T) { // Request 2: [user1, user2] // For 2 messages, matches last 1 message // Should match chatID1 because last message is cached - messages2 := []Message{ - {Role: RoleUser, Content: "First message"}, - {Role: RoleUser, Content: "Second message"}, + messages2 := []context.Message{ + {Role: context.RoleUser, Content: "First message"}, + {Role: context.RoleUser, Content: "Second message"}, } - chatID2, err := GetChatIDByMessages(cache, messages2) + chatID2, err := context.GetChatIDByMessages(cache, messages2) if err != nil { t.Fatalf("Failed to get chat ID: %v", err) } @@ -92,12 +93,12 @@ func TestGetChatIDByMessages_ContinuousConversation(t *testing.T) { // Request 3: [user1, user2, user3] // For 3+ messages, matches last 2 messages // Should match chatID2 because last 2 messages are cached - messages3 := []Message{ - {Role: RoleUser, Content: "First message"}, - {Role: RoleUser, Content: "Second message"}, - {Role: RoleUser, Content: "Third message"}, + messages3 := []context.Message{ + {Role: context.RoleUser, Content: "First message"}, + {Role: context.RoleUser, Content: "Second message"}, + {Role: context.RoleUser, Content: "Third message"}, } - chatID3, err := GetChatIDByMessages(cache, messages3) + chatID3, err := context.GetChatIDByMessages(cache, messages3) if err != nil { t.Fatalf("Failed to get chat ID: %v", err) } @@ -108,13 +109,13 @@ func TestGetChatIDByMessages_ContinuousConversation(t *testing.T) { // Request 4: [user1, user2, user3, user4] // Should match chatID3 because last 2 messages are cached - messages4 := []Message{ - {Role: RoleUser, Content: "First message"}, - {Role: RoleUser, Content: "Second message"}, - {Role: RoleUser, Content: "Third message"}, - {Role: RoleUser, Content: "Fourth message"}, + messages4 := []context.Message{ + {Role: context.RoleUser, Content: "First message"}, + {Role: context.RoleUser, Content: "Second message"}, + {Role: context.RoleUser, Content: "Third message"}, + {Role: context.RoleUser, Content: "Fourth message"}, } - chatID4, err := GetChatIDByMessages(cache, messages4) + chatID4, err := context.GetChatIDByMessages(cache, messages4) if err != nil { t.Fatalf("Failed to get chat ID: %v", err) } @@ -136,32 +137,32 @@ func TestGetChatIDByMessages_DifferentConversations(t *testing.T) { cache := getTestCache(t) // First conversation - messages1 := []Message{ + messages1 := []context.Message{ { - Role: RoleUser, + Role: context.RoleUser, Content: "Hello", }, } - chatID1, err := GetChatIDByMessages(cache, messages1) + chatID1, err := context.GetChatIDByMessages(cache, messages1) if err != nil { t.Fatalf("Failed to get chat ID: %v", err) } - err = CacheChatID(cache, messages1, chatID1) + err = context.CacheChatID(cache, messages1, chatID1) if err != nil { t.Fatalf("Failed to cache chat ID: %v", err) } // Different conversation - messages2 := []Message{ + messages2 := []context.Message{ { - Role: RoleUser, + Role: context.RoleUser, Content: "Goodbye", }, } - chatID2, err := GetChatIDByMessages(cache, messages2) + chatID2, err := context.GetChatIDByMessages(cache, messages2) if err != nil { t.Fatalf("Failed to get chat ID: %v", err) } @@ -178,37 +179,37 @@ func TestGetChatIDByMessages_MultiModalContent(t *testing.T) { cache := getTestCache(t) // First request with multimodal content - messages1 := []Message{ + messages1 := []context.Message{ { - Role: RoleUser, - Content: []ContentPart{ + Role: context.RoleUser, + Content: []context.ContentPart{ { - Type: ContentText, + Type: context.ContentText, Text: "What's in this image?", }, { - Type: ContentImageURL, - ImageURL: &ImageURL{ + Type: context.ContentImageURL, + ImageURL: &context.ImageURL{ URL: "https://example.com/image.jpg", - Detail: DetailHigh, + Detail: context.DetailHigh, }, }, }, }, } - chatID1, err := GetChatIDByMessages(cache, messages1) + chatID1, err := context.GetChatIDByMessages(cache, messages1) if err != nil { t.Fatalf("Failed to get chat ID: %v", err) } // Second request - add another message to continue conversation - messages2 := append(messages1, Message{ - Role: RoleUser, + messages2 := append(messages1, context.Message{ + Role: context.RoleUser, Content: "Tell me more details", }) - chatID2, err := GetChatIDByMessages(cache, messages2) + chatID2, err := context.GetChatIDByMessages(cache, messages2) if err != nil { t.Fatalf("Failed to get chat ID: %v", err) } @@ -226,32 +227,32 @@ func TestGetChatIDByMessages_WithToolCalls(t *testing.T) { cache := getTestCache(t) // First request with user message - messages1 := []Message{ + messages1 := []context.Message{ { - Role: RoleUser, + Role: context.RoleUser, Content: "What's the weather in Tokyo?", }, } - chatID1, err := GetChatIDByMessages(cache, messages1) + chatID1, err := context.GetChatIDByMessages(cache, messages1) if err != nil { t.Fatalf("Failed to get chat ID: %v", err) } // Second request - add assistant response and another user message - messages2 := []Message{ + messages2 := []context.Message{ { - Role: RoleUser, + Role: context.RoleUser, Content: "What's the weather in Tokyo?", }, { - Role: RoleAssistant, + Role: context.RoleAssistant, Content: nil, - ToolCalls: []ToolCall{ + ToolCalls: []context.ToolCall{ { ID: "call_123", - Type: ToolTypeFunction, - Function: Function{ + Type: context.ToolTypeFunction, + Function: context.Function{ Name: "get_weather", Arguments: `{"location":"Tokyo"}`, }, @@ -259,12 +260,12 @@ func TestGetChatIDByMessages_WithToolCalls(t *testing.T) { }, }, { - Role: RoleUser, + Role: context.RoleUser, Content: "How about tomorrow?", }, } - chatID2, err := GetChatIDByMessages(cache, messages2) + chatID2, err := context.GetChatIDByMessages(cache, messages2) if err != nil { t.Fatalf("Failed to get chat ID: %v", err) } @@ -281,7 +282,7 @@ func TestCacheChatID_EmptyMessages(t *testing.T) { cache := getTestCache(t) - err := CacheChatID(cache, []Message{}, "chat_123") + err := context.CacheChatID(cache, []context.Message{}, "chat_123") if err == nil { t.Error("Expected error for empty messages") } @@ -293,14 +294,14 @@ func TestCacheChatID_EmptyChatID(t *testing.T) { cache := getTestCache(t) - messages := []Message{ + messages := []context.Message{ { - Role: RoleUser, + Role: context.RoleUser, Content: "Hello", }, } - err := CacheChatID(cache, messages, "") + err := context.CacheChatID(cache, messages, "") if err == nil { t.Error("Expected error for empty chat ID") } @@ -312,53 +313,14 @@ func TestGetChatIDByMessages_EmptyMessages(t *testing.T) { cache := getTestCache(t) - _, err := GetChatIDByMessages(cache, []Message{}) + _, err := context.GetChatIDByMessages(cache, []context.Message{}) if err == nil { t.Error("Expected error for empty messages") } } -func TestHashMessage_Consistency(t *testing.T) { - msg := Message{ - Role: RoleUser, - Content: "Test message", - } - - hash1, err := hashMessage(msg) - if err != nil { - t.Fatalf("Failed to hash message: %v", err) - } - - hash2, err := hashMessage(msg) - if err != nil { - t.Fatalf("Failed to hash message: %v", err) - } - - if hash1 != hash2 { - t.Errorf("Expected consistent hashes, got %s and %s", hash1, hash2) - } -} - -func TestGetKey(t *testing.T) { - hash := "abc123" - key := getKey(hash) - - expectedPrefix := chatCachePrefix - if len(key) <= len(expectedPrefix) { - t.Errorf("Expected key to have prefix, got %s", key) - } - - if key[:len(expectedPrefix)] != expectedPrefix { - t.Errorf("Expected key to start with %s, got %s", expectedPrefix, key) - } - - if key != chatCachePrefix+hash { - t.Errorf("Expected key %s, got %s", chatCachePrefix+hash, key) - } -} - func TestGenChatID(t *testing.T) { - id1 := GenChatID() + id1 := context.GenChatID() if id1 == "" { t.Error("Expected non-empty chat ID") diff --git a/agent/context/context.go b/agent/context/context.go index a3d383b2..fa2b91b4 100644 --- a/agent/context/context.go +++ b/agent/context/context.go @@ -27,14 +27,17 @@ func New(parent context.Context, authorized *types.AuthorizedInfo, chatID string parent = context.Background() } + contextID := generateContextID() + ctx := &Context{ Context: parent, - ID: generateContextID(), // Generate unique ID for the context - Authorized: authorized, // Set authorized info + ID: contextID, // Generate unique ID for the context + Authorized: authorized, // Set authorized info Space: plan.NewMemorySharedSpace(), ChatID: chatID, - IDGenerator: message.NewIDGenerator(), // Initialize ID generator for this context - messageMetadata: newMessageMetadataStore(), // Initialize message metadata store + IDGenerator: message.NewIDGenerator(), // Initialize ID generator for this context + messageMetadata: newMessageMetadataStore(), // Initialize message metadata store + Logger: NewRequestLogger("", chatID, contextID), // Initialize logger (assistantID set later) } return ctx @@ -82,7 +85,9 @@ func WithTimeout(parent *Context, timeout time.Duration) (*Context, context.Canc // Release the context and clean up all resources including stacks and trace func (ctx *Context) Release() { - log.Trace("[RELEASE] Context cleanup started: contextID=%s, assistantID=%s", ctx.ID, ctx.AssistantID) + if ctx.Logger != nil { + ctx.Logger.Release() + } // Unregister from global registry if ctx.ID != "" { @@ -91,61 +96,44 @@ func (ctx *Context) Release() { // Stop interrupt controller if ctx.Interrupt != nil { - log.Trace("[RELEASE] Stopping interrupt controller") + if ctx.Logger != nil { + ctx.Logger.Cleanup("Interrupt controller") + } ctx.Interrupt.Stop() ctx.Interrupt = nil } // Complete and release trace if exists if ctx.trace != nil && ctx.Stack != nil && ctx.Stack.TraceID != "" { - log.Trace("[RELEASE] Releasing trace: traceID=%s", ctx.Stack.TraceID) + if ctx.Logger != nil { + ctx.Logger.Cleanup("Trace: " + ctx.Stack.TraceID) + } // Check if context is cancelled - if so, mark as cancelled instead of complete if ctx.Context != nil && ctx.Context.Err() != nil { - log.Trace("[RELEASE] Context cancelled, marking trace as cancelled: err=%v", ctx.Context.Err()) - - // Mark trace as cancelled (saves to disk and broadcasts to subscribers) - log.Trace("[RELEASE] Calling trace.MarkCancelled: traceID=%s", ctx.Stack.TraceID) - if err := trace.MarkCancelled(ctx.Stack.TraceID, ctx.Context.Err().Error()); err != nil { - log.Trace("[RELEASE] Failed to mark trace as cancelled: %v", err) - } else { - log.Trace("[RELEASE] Successfully marked trace as cancelled") - } - - // Release trace from registry - // Subscribers will be notified via channel close and will cleanup automatically - log.Trace("[RELEASE] Calling trace.Release: traceID=%s", ctx.Stack.TraceID) - if err := trace.Release(ctx.Stack.TraceID); err != nil { - log.Trace("[RELEASE] Failed to release trace from registry: %v", err) - } else { - log.Trace("[RELEASE] Successfully released trace from registry") - } + trace.MarkCancelled(ctx.Stack.TraceID, ctx.Context.Err().Error()) + trace.Release(ctx.Stack.TraceID) } else { - // Normal case: mark complete then release - if err := ctx.trace.MarkComplete(); err != nil { - log.Trace("[RELEASE] Failed to mark trace complete: %v", err) - } - - if err := trace.Release(ctx.Stack.TraceID); err != nil { - log.Trace("[RELEASE] Failed to release trace: %v", err) - } + ctx.trace.MarkComplete() + trace.Release(ctx.Stack.TraceID) } - ctx.trace = nil - } else { - log.Trace("[RELEASE] No trace to release (trace=%v, stack=%v)", ctx.trace != nil, ctx.Stack != nil) } // Clear space if ctx.Space != nil { - log.Trace("[RELEASE] Clearing space") + if ctx.Logger != nil { + ctx.Logger.Cleanup("Space") + } ctx.Space.Clear() ctx.Space = nil } // Clear stacks if ctx.Stacks != nil { - log.Trace("[RELEASE] Clearing %d stacks", len(ctx.Stacks)) + if ctx.Logger != nil { + ctx.Logger.Cleanup(fmt.Sprintf("Stacks (%d)", len(ctx.Stacks))) + } for k := range ctx.Stacks { delete(ctx.Stacks, k) } @@ -158,8 +146,11 @@ func (ctx *Context) Release() { // Clear writer reference ctx.Writer = nil - log.Trace("[RELEASE] Context cleanup completed: contextID=%s", ctx.ID) - ctx = nil + // Close logger (MUST be last) + if ctx.Logger != nil { + ctx.Logger.Close() + ctx.Logger = nil + } } // Send sends data to the context's writer diff --git a/agent/context/context_test.go b/agent/context/context_test.go index f4c48aa3..623cfb1c 100644 --- a/agent/context/context_test.go +++ b/agent/context/context_test.go @@ -1,7 +1,8 @@ -package context +package context_test import ( "bytes" + stdContext "context" "encoding/json" "net/http" "net/http/httptest" @@ -10,6 +11,7 @@ import ( "github.com/gin-gonic/gin" "github.com/stretchr/testify/assert" "github.com/yaoapp/gou/store" + "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/config" "github.com/yaoapp/yao/test" ) @@ -37,7 +39,7 @@ func TestGetCompletionRequest(t *testing.T) { expectedLocale string expectedTheme string expectedReferer string - expectedAccept Accept + expectedAccept context.Accept expectedAssistantID string expectError bool }{ @@ -64,8 +66,8 @@ func TestGetCompletionRequest(t *testing.T) { expectedStream: boolPtr(true), expectedLocale: "zh-cn", expectedTheme: "dark", - expectedReferer: RefererProcess, - expectedAccept: AcceptWebCUI, + expectedReferer: context.RefererProcess, + expectedAccept: context.AcceptWebCUI, expectedAssistantID: "assistant123", expectError: false, }, @@ -89,8 +91,8 @@ func TestGetCompletionRequest(t *testing.T) { expectedMsgCount: 1, expectedLocale: "fr-fr", expectedTheme: "auto", - expectedReferer: RefererAPI, - expectedAccept: AcceptStandard, + expectedReferer: context.RefererAPI, + expectedAccept: context.AcceptStandard, expectedAssistantID: "test456", expectError: false, }, @@ -114,8 +116,8 @@ func TestGetCompletionRequest(t *testing.T) { expectedMsgCount: 1, expectedLocale: "", expectedTheme: "", - expectedReferer: RefererMCP, - expectedAccept: AcceptDesktopCUI, + expectedReferer: context.RefererMCP, + expectedAccept: context.AcceptDesktopCUI, expectedAssistantID: "header789", expectError: false, }, @@ -131,8 +133,8 @@ func TestGetCompletionRequest(t *testing.T) { expectedMsgCount: 1, expectedLocale: "", expectedTheme: "", - expectedReferer: RefererAPI, - expectedAccept: AcceptStandard, + expectedReferer: context.RefererAPI, + expectedAccept: context.AcceptStandard, expectedAssistantID: "minimal", expectError: false, }, @@ -179,7 +181,7 @@ func TestGetCompletionRequest(t *testing.T) { c.Request = req // Call GetCompletionRequest - completionReq, ctx, opts, err := GetCompletionRequest(c, cache) + completionReq, ctx, opts, err := context.GetCompletionRequest(c, cache) if tt.expectError { assert.Error(t, err) @@ -215,110 +217,21 @@ func TestGetCompletionRequest(t *testing.T) { } } -func TestParseClientType(t *testing.T) { - tests := []struct { - name string - userAgent string - expected string - }{ - {"Empty user agent", "", "web"}, - {"Standard web browser", "Mozilla/5.0", "web"}, - {"Android", "Mozilla/5.0 (Linux; Android 10)", "android"}, - {"iPhone", "Mozilla/5.0 (iPhone; CPU iPhone OS 14_0)", "ios"}, - {"iPad", "Mozilla/5.0 (iPad; CPU OS 14_0)", "ios"}, - {"Windows", "Mozilla/5.0 (Windows NT 10.0)", "windows"}, - {"macOS", "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7)", "macos"}, - {"Linux", "Mozilla/5.0 (X11; Linux x86_64)", "linux"}, - {"Yao Agent", "Yao-Agent/1.0", "agent"}, - {"JSSDK", "Yao-JSSDK/2.0", "jssdk"}, - } +func TestContextNew_WithAuthorized(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - result := getClientType(tt.userAgent) - assert.Equal(t, tt.expected, result) - }) - } + // Create context using New() + ctx := context.New(stdContext.Background(), nil, "test-chat-id") + defer ctx.Release() + + assert.NotNil(t, ctx) + assert.Equal(t, "test-chat-id", ctx.ChatID) + assert.NotNil(t, ctx.Space) + assert.NotNil(t, ctx.IDGenerator) } -func TestParseAccept(t *testing.T) { - tests := []struct { - name string - clientType string - expected Accept - }{ - {"Web client", "web", AcceptWebCUI}, - {"Android client", "android", AccepNativeCUI}, - {"iOS client", "ios", AccepNativeCUI}, - {"Windows client", "windows", AcceptDesktopCUI}, - {"macOS client", "macos", AcceptDesktopCUI}, - {"Linux client", "linux", AcceptDesktopCUI}, - {"Agent client", "agent", AcceptStandard}, - {"JSSDK client", "jssdk", AcceptStandard}, - {"Unknown client", "unknown", AcceptStandard}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - result := parseAccept(tt.clientType) - assert.Equal(t, tt.expected, result) - }) - } -} - -func TestValidateAccept(t *testing.T) { - tests := []struct { - name string - accept string - expected Accept - }{ - {"Valid standard", "standard", AcceptStandard}, - {"Valid cui-web", "cui-web", AcceptWebCUI}, - {"Valid cui-native", "cui-native", AccepNativeCUI}, - {"Valid cui-desktop", "cui-desktop", AcceptDesktopCUI}, - {"Invalid value", "invalid", AcceptStandard}, - {"Empty string", "", AcceptStandard}, - {"Random string", "random-accept", AcceptStandard}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - result := validateAccept(tt.accept) - assert.Equal(t, tt.expected, result) - }) - } -} - -func TestValidateReferer(t *testing.T) { - tests := []struct { - name string - referer string - expected string - }{ - {"Valid api", "api", RefererAPI}, - {"Valid process", "process", RefererProcess}, - {"Valid mcp", "mcp", RefererMCP}, - {"Valid jssdk", "jssdk", RefererJSSDK}, - {"Valid agent", "agent", RefererAgent}, - {"Valid tool", "tool", RefererTool}, - {"Valid hook", "hook", RefererHook}, - {"Valid schedule", "schedule", RefererSchedule}, - {"Valid script", "script", RefererScript}, - {"Valid internal", "internal", RefererInternal}, - {"Invalid value", "invalid", RefererAPI}, - {"Empty string", "", RefererAPI}, - {"Random string", "random-referer", RefererAPI}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - result := validateReferer(tt.referer) - assert.Equal(t, tt.expected, result) - }) - } -} - -// Helper functions +// Helper functions for context_test package func floatPtr(f float64) *float64 { return &f } diff --git a/agent/context/interrupt_test.go b/agent/context/interrupt_test.go index 89bebddb..70e3190c 100644 --- a/agent/context/interrupt_test.go +++ b/agent/context/interrupt_test.go @@ -1,4 +1,4 @@ -package context +package context_test import ( stdContext "context" @@ -6,48 +6,41 @@ import ( "testing" "time" - "github.com/yaoapp/gou/plan" - "github.com/yaoapp/yao/agent/output/message" + "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/openapi/oauth/types" ) // newTestContextWithInterrupt creates a Context with interrupt controller for testing -func newTestContextWithInterrupt(chatID, assistantID string) *Context { - ctx := &Context{ - Context: stdContext.Background(), - ID: fmt.Sprintf("test_ctx_%d", time.Now().UnixNano()), - Space: plan.NewMemorySharedSpace(), - ChatID: chatID, - AssistantID: assistantID, - Locale: "en-us", - Theme: "light", - Client: Client{ - Type: "web", - UserAgent: "TestAgent/1.0", - IP: "127.0.0.1", - }, - Referer: RefererAPI, - Accept: AcceptWebCUI, - Route: "/test/route", - IDGenerator: message.NewIDGenerator(), // Initialize context-scoped ID generator - Metadata: map[string]interface{}{ - "test": "context_metadata", - }, - Authorized: &types.AuthorizedInfo{ - Subject: "test-user", - ClientID: "test-client-id", - UserID: "test-user-123", - TeamID: "test-team-456", - TenantID: "test-tenant-789", - SessionID: "test-session-id", - }, +func newTestContextWithInterrupt(chatID, assistantID string) *context.Context { + ctx := context.New(stdContext.Background(), &types.AuthorizedInfo{ + Subject: "test-user", + ClientID: "test-client-id", + UserID: "test-user-123", + TeamID: "test-team-456", + TenantID: "test-tenant-789", + SessionID: "test-session-id", + }, chatID) + + ctx.AssistantID = assistantID + ctx.Locale = "en-us" + ctx.Theme = "light" + ctx.Client = context.Client{ + Type: "web", + UserAgent: "TestAgent/1.0", + IP: "127.0.0.1", + } + ctx.Referer = context.RefererAPI + ctx.Accept = context.AcceptWebCUI + ctx.Route = "/test/route" + ctx.Metadata = map[string]interface{}{ + "test": "context_metadata", } // Initialize interrupt controller - ctx.Interrupt = NewInterruptController() + ctx.Interrupt = context.NewInterruptController() // Register context globally - if err := Register(ctx); err != nil { + if err := context.Register(ctx); err != nil { panic(fmt.Sprintf("Failed to register context: %v", err)) } @@ -65,16 +58,16 @@ func TestInterruptBasic(t *testing.T) { t.Run("SendGracefulInterrupt", func(t *testing.T) { // Create a graceful interrupt signal - signal := &InterruptSignal{ - Type: InterruptGraceful, - Messages: []Message{ - {Role: RoleUser, Content: "This is a graceful interrupt"}, + signal := &context.InterruptSignal{ + Type: context.InterruptGraceful, + Messages: []context.Message{ + {Role: context.RoleUser, Content: "This is a graceful interrupt"}, }, Timestamp: time.Now().UnixMilli(), } // Send interrupt signal - err := SendInterrupt(ctx.ID, signal) + err := context.SendInterrupt(ctx.ID, signal) if err != nil { t.Fatalf("Failed to send interrupt signal: %v", err) } @@ -88,7 +81,7 @@ func TestInterruptBasic(t *testing.T) { t.Fatal("Expected interrupt signal to be received, got nil") } - if receivedSignal.Type != InterruptGraceful { + if receivedSignal.Type != context.InterruptGraceful { t.Errorf("Expected interrupt type 'graceful', got: %s", receivedSignal.Type) } @@ -108,16 +101,16 @@ func TestInterruptBasic(t *testing.T) { ctx.Interrupt.Clear() // Create a force interrupt signal - signal := &InterruptSignal{ - Type: InterruptForce, - Messages: []Message{ - {Role: RoleUser, Content: "This is a force interrupt"}, + signal := &context.InterruptSignal{ + Type: context.InterruptForce, + Messages: []context.Message{ + {Role: context.RoleUser, Content: "This is a force interrupt"}, }, Timestamp: time.Now().UnixMilli(), } // Send interrupt signal - err := SendInterrupt(ctx.ID, signal) + err := context.SendInterrupt(ctx.ID, signal) if err != nil { t.Fatalf("Failed to send interrupt signal: %v", err) } @@ -131,7 +124,7 @@ func TestInterruptBasic(t *testing.T) { t.Fatal("Expected interrupt signal to be received, got nil") } - if receivedSignal.Type != InterruptForce { + if receivedSignal.Type != context.InterruptForce { t.Errorf("Expected interrupt type 'force', got: %s", receivedSignal.Type) } @@ -144,15 +137,15 @@ func TestInterruptBasic(t *testing.T) { // Send multiple interrupt signals for i := 0; i < 3; i++ { - signal := &InterruptSignal{ - Type: InterruptGraceful, - Messages: []Message{ - {Role: RoleUser, Content: fmt.Sprintf("Message %d", i+1)}, + signal := &context.InterruptSignal{ + Type: context.InterruptGraceful, + Messages: []context.Message{ + {Role: context.RoleUser, Content: fmt.Sprintf("Message %d", i+1)}, }, Timestamp: time.Now().UnixMilli(), } - err := SendInterrupt(ctx.ID, signal) + err := context.SendInterrupt(ctx.ID, signal) if err != nil { t.Fatalf("Failed to send interrupt signal %d: %v", i+1, err) } @@ -198,10 +191,10 @@ func TestInterruptHandler(t *testing.T) { t.Run("HandlerInvocation", func(t *testing.T) { // Track if handler was called handlerCalled := false - var receivedSignal *InterruptSignal + var receivedSignal *context.InterruptSignal // Set up handler - ctx.Interrupt.SetHandler(func(c *Context, signal *InterruptSignal) error { + ctx.Interrupt.SetHandler(func(c *context.Context, signal *context.InterruptSignal) error { handlerCalled = true receivedSignal = signal t.Logf("Handler called with signal type: %s, messages: %d", signal.Type, len(signal.Messages)) @@ -209,15 +202,15 @@ func TestInterruptHandler(t *testing.T) { }) // Send interrupt signal - signal := &InterruptSignal{ - Type: InterruptGraceful, - Messages: []Message{ - {Role: RoleUser, Content: "Test handler invocation"}, + signal := &context.InterruptSignal{ + Type: context.InterruptGraceful, + Messages: []context.Message{ + {Role: context.RoleUser, Content: "Test handler invocation"}, }, Timestamp: time.Now().UnixMilli(), } - err := SendInterrupt(ctx.ID, signal) + err := context.SendInterrupt(ctx.ID, signal) if err != nil { t.Fatalf("Failed to send interrupt signal: %v", err) } @@ -234,7 +227,7 @@ func TestInterruptHandler(t *testing.T) { t.Fatal("Expected signal in handler, got nil") } - if receivedSignal.Type != InterruptGraceful { + if receivedSignal.Type != context.InterruptGraceful { t.Errorf("Expected graceful interrupt in handler, got: %s", receivedSignal.Type) } @@ -252,21 +245,21 @@ func TestInterruptHandler(t *testing.T) { // Set up handler that returns error handlerCalled := false - ctx2.Interrupt.SetHandler(func(c *Context, signal *InterruptSignal) error { + ctx2.Interrupt.SetHandler(func(c *context.Context, signal *context.InterruptSignal) error { handlerCalled = true return fmt.Errorf("test error from handler") }) // Send interrupt signal - signal := &InterruptSignal{ - Type: InterruptForce, - Messages: []Message{ - {Role: RoleUser, Content: "Test error handling"}, + signal := &context.InterruptSignal{ + Type: context.InterruptForce, + Messages: []context.Message{ + {Role: context.RoleUser, Content: "Test error handling"}, }, Timestamp: time.Now().UnixMilli(), } - err := SendInterrupt(ctx2.ID, signal) + err := context.SendInterrupt(ctx2.ID, signal) if err != nil { t.Fatalf("Failed to send interrupt signal: %v", err) } @@ -289,7 +282,7 @@ func TestInterruptContextLifecycle(t *testing.T) { ctx := newTestContextWithInterrupt("chat-test-lifecycle", "test-assistant") // Verify context can be retrieved - retrievedCtx, err := Get(ctx.ID) + retrievedCtx, err := context.Get(ctx.ID) if err != nil { t.Fatalf("Failed to retrieve context: %v", err) } @@ -301,7 +294,7 @@ func TestInterruptContextLifecycle(t *testing.T) { ctx.Release() // After release, context should be removed - _, err = Get(ctx.ID) + _, err = context.Get(ctx.ID) if err == nil { t.Error("Expected error when retrieving released context") } @@ -310,13 +303,13 @@ func TestInterruptContextLifecycle(t *testing.T) { }) t.Run("SendToNonExistentContext", func(t *testing.T) { - signal := &InterruptSignal{ - Type: InterruptGraceful, - Messages: []Message{{Role: RoleUser, Content: "test"}}, + signal := &context.InterruptSignal{ + Type: context.InterruptGraceful, + Messages: []context.Message{{Role: context.RoleUser, Content: "test"}}, Timestamp: time.Now().UnixMilli(), } - err := SendInterrupt("non-existent-id", signal) + err := context.SendInterrupt("non-existent-id", signal) if err == nil { t.Error("Expected error when sending to non-existent context") } @@ -332,12 +325,12 @@ func TestInterruptCheckMethods(t *testing.T) { t.Run("PeekDoesNotRemove", func(t *testing.T) { // Send signal - signal := &InterruptSignal{ - Type: InterruptGraceful, - Messages: []Message{{Role: RoleUser, Content: "peek test"}}, + signal := &context.InterruptSignal{ + Type: context.InterruptGraceful, + Messages: []context.Message{{Role: context.RoleUser, Content: "peek test"}}, Timestamp: time.Now().UnixMilli(), } - SendInterrupt(ctx.ID, signal) + context.SendInterrupt(ctx.ID, signal) time.Sleep(100 * time.Millisecond) // Peek should return signal but not remove it @@ -362,12 +355,12 @@ func TestInterruptCheckMethods(t *testing.T) { ctx.Interrupt.Clear() // Send signal - signal := &InterruptSignal{ - Type: InterruptGraceful, - Messages: []Message{{Role: RoleUser, Content: "check test"}}, + signal := &context.InterruptSignal{ + Type: context.InterruptGraceful, + Messages: []context.Message{{Role: context.RoleUser, Content: "check test"}}, Timestamp: time.Now().UnixMilli(), } - SendInterrupt(ctx.ID, signal) + context.SendInterrupt(ctx.ID, signal) time.Sleep(100 * time.Millisecond) // Check should return and remove signal @@ -398,17 +391,17 @@ func TestInterruptCheckMethods(t *testing.T) { } for i, msg := range messages { - signal := &InterruptSignal{ - Type: InterruptGraceful, - Messages: []Message{ - {Role: RoleUser, Content: msg}, + signal := &context.InterruptSignal{ + Type: context.InterruptGraceful, + Messages: []context.Message{ + {Role: context.RoleUser, Content: msg}, }, Timestamp: time.Now().UnixMilli(), Metadata: map[string]interface{}{ "sequence": i + 1, }, } - err := SendInterrupt(ctx.ID, signal) + err := context.SendInterrupt(ctx.ID, signal) if err != nil { t.Fatalf("Failed to send signal %d: %v", i+1, err) } @@ -461,12 +454,12 @@ func TestInterruptCheckMethods(t *testing.T) { ctx.Interrupt.Clear() // Send single signal - signal := &InterruptSignal{ - Type: InterruptGraceful, - Messages: []Message{{Role: RoleUser, Content: "single signal"}}, + signal := &context.InterruptSignal{ + Type: context.InterruptGraceful, + Messages: []context.Message{{Role: context.RoleUser, Content: "single signal"}}, Timestamp: time.Now().UnixMilli(), } - SendInterrupt(ctx.ID, signal) + context.SendInterrupt(ctx.ID, signal) time.Sleep(100 * time.Millisecond) // CheckWithMerge with single signal should return it without merge metadata @@ -523,12 +516,12 @@ func TestInterruptContext(t *testing.T) { // Send force interrupt with empty messages (pure cancellation) // This is the pattern for stopping streaming without appending messages - signal := &InterruptSignal{ - Type: InterruptForce, - Messages: []Message{}, // Empty messages = pure cancellation + signal := &context.InterruptSignal{ + Type: context.InterruptForce, + Messages: []context.Message{}, // Empty messages = pure cancellation Timestamp: time.Now().UnixMilli(), } - err := SendInterrupt(ctx.ID, signal) + err := context.SendInterrupt(ctx.ID, signal) if err != nil { t.Fatalf("Failed to send interrupt: %v", err) } @@ -557,12 +550,12 @@ func TestInterruptContext(t *testing.T) { interruptCtx := ctx2.Interrupt.Context() // Send graceful interrupt - signal := &InterruptSignal{ - Type: InterruptGraceful, - Messages: []Message{{Role: RoleUser, Content: "graceful"}}, + signal := &context.InterruptSignal{ + Type: context.InterruptGraceful, + Messages: []context.Message{{Role: context.RoleUser, Content: "graceful"}}, Timestamp: time.Now().UnixMilli(), } - SendInterrupt(ctx2.ID, signal) + context.SendInterrupt(ctx2.ID, signal) time.Sleep(100 * time.Millisecond) // Context should NOT be cancelled for graceful interrupt @@ -588,9 +581,9 @@ func TestInterruptSendSignalDirectly(t *testing.T) { defer ctx.Release() t.Run("SendSignalSuccess", func(t *testing.T) { - signal := &InterruptSignal{ - Type: InterruptGraceful, - Messages: []Message{{Role: RoleUser, Content: "direct send"}}, + signal := &context.InterruptSignal{ + Type: context.InterruptGraceful, + Messages: []context.Message{{Role: context.RoleUser, Content: "direct send"}}, Timestamp: time.Now().UnixMilli(), } @@ -615,10 +608,10 @@ func TestInterruptSendSignalDirectly(t *testing.T) { }) t.Run("SendSignalToNilController", func(t *testing.T) { - var nilController *InterruptController - signal := &InterruptSignal{ - Type: InterruptGraceful, - Messages: []Message{{Role: RoleUser, Content: "test"}}, + var nilController *context.InterruptController + signal := &context.InterruptSignal{ + Type: context.InterruptGraceful, + Messages: []context.Message{{Role: context.RoleUser, Content: "test"}}, Timestamp: time.Now().UnixMilli(), } @@ -632,23 +625,23 @@ func TestInterruptSendSignalDirectly(t *testing.T) { t.Run("SendSignalTimeout", func(t *testing.T) { // Create controller but don't start listener - testCtrl := NewInterruptController() + testCtrl := context.NewInterruptController() // Don't call Start(), so channel won't be read // Fill the buffer (capacity is 10) for i := 0; i < 10; i++ { - signal := &InterruptSignal{ - Type: InterruptGraceful, - Messages: []Message{{Role: RoleUser, Content: fmt.Sprintf("msg %d", i)}}, + signal := &context.InterruptSignal{ + Type: context.InterruptGraceful, + Messages: []context.Message{{Role: context.RoleUser, Content: fmt.Sprintf("msg %d", i)}}, Timestamp: time.Now().UnixMilli(), } testCtrl.SendSignal(signal) } // This should timeout since buffer is full and no listener - signal := &InterruptSignal{ - Type: InterruptGraceful, - Messages: []Message{{Role: RoleUser, Content: "overflow"}}, + signal := &context.InterruptSignal{ + Type: context.InterruptGraceful, + Messages: []context.Message{{Role: context.RoleUser, Content: "overflow"}}, Timestamp: time.Now().UnixMilli(), } diff --git a/agent/context/jsapi_output_test.go b/agent/context/jsapi_output_test.go index ed109ab6..128b18a4 100644 --- a/agent/context/jsapi_output_test.go +++ b/agent/context/jsapi_output_test.go @@ -1,42 +1,42 @@ -package context +package context_test import ( "bytes" - "context" + stdContext "context" "net/http" "testing" "github.com/stretchr/testify/assert" v8 "github.com/yaoapp/gou/runtime/v8" - "github.com/yaoapp/yao/agent/output/message" + "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/config" "github.com/yaoapp/yao/test" ) -// mockResponseWriter is a mock implementation of http.ResponseWriter for testing -type mockResponseWriter struct { +// testMockResponseWriter is a mock implementation of http.ResponseWriter for testing +type testMockResponseWriter struct { headers http.Header buffer *bytes.Buffer status int } -func newMockResponseWriter() *mockResponseWriter { - return &mockResponseWriter{ +func newTestMockResponseWriter() *testMockResponseWriter { + return &testMockResponseWriter{ headers: make(http.Header), buffer: &bytes.Buffer{}, status: http.StatusOK, } } -func (m *mockResponseWriter) Header() http.Header { +func (m *testMockResponseWriter) Header() http.Header { return m.headers } -func (m *mockResponseWriter) Write(b []byte) (int, error) { +func (m *testMockResponseWriter) Write(b []byte) (int, error) { return m.buffer.Write(b) } -func (m *mockResponseWriter) WriteHeader(statusCode int) { +func (m *testMockResponseWriter) WriteHeader(statusCode int) { m.status = statusCode } @@ -46,15 +46,11 @@ func TestJsValueSend(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - cxt := &Context{ - ChatID: "test-chat-id", - AssistantID: "test-assistant-id", - Context: context.Background(), - Accept: "standard", - Locale: "en", - Writer: newMockResponseWriter(), - IDGenerator: message.NewIDGenerator(), - } + cxt := context.New(stdContext.Background(), nil, "test-chat-id") + cxt.AssistantID = "test-assistant-id" + cxt.Accept = context.AcceptStandard + cxt.Locale = "en" + cxt.Writer = newTestMockResponseWriter() // Test sending string shorthand res, err := v8.Call(v8.CallOptions{}, ` @@ -116,15 +112,11 @@ func TestJsValueSendDeltaUpdates(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - cxt := &Context{ - ChatID: "test-chat-id", - AssistantID: "test-assistant-id", - Context: context.Background(), - Accept: "standard", - Locale: "en", - Writer: newMockResponseWriter(), - IDGenerator: message.NewIDGenerator(), - } + cxt := context.New(stdContext.Background(), nil, "test-chat-id") + cxt.AssistantID = "test-assistant-id" + cxt.Accept = context.AcceptStandard + cxt.Locale = "en" + cxt.Writer = newTestMockResponseWriter() res, err := v8.Call(v8.CallOptions{}, ` function test(ctx) { @@ -171,15 +163,11 @@ func TestJsValueSendMultipleTypes(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - cxt := &Context{ - ChatID: "test-chat-id", - AssistantID: "test-assistant-id", - Context: context.Background(), - Accept: "standard", - Locale: "en", - Writer: newMockResponseWriter(), - IDGenerator: message.NewIDGenerator(), - } + cxt := context.New(stdContext.Background(), nil, "test-chat-id") + cxt.AssistantID = "test-assistant-id" + cxt.Accept = context.AcceptStandard + cxt.Locale = "en" + cxt.Writer = newTestMockResponseWriter() res, err := v8.Call(v8.CallOptions{}, ` function test(ctx) { @@ -254,15 +242,11 @@ func TestJsValueSendErrorHandling(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - cxt := &Context{ - ChatID: "test-chat-id", - AssistantID: "test-assistant-id", - Context: context.Background(), - Accept: "standard", - Locale: "en", - Writer: newMockResponseWriter(), - IDGenerator: message.NewIDGenerator(), - } + cxt := context.New(stdContext.Background(), nil, "test-chat-id") + cxt.AssistantID = "test-assistant-id" + cxt.Accept = context.AcceptStandard + cxt.Locale = "en" + cxt.Writer = newTestMockResponseWriter() // Test invalid argument - no arguments res, err := v8.Call(v8.CallOptions{}, ` @@ -293,18 +277,15 @@ func TestJsValueSendWithCUIAccept(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - acceptTypes := []string{"cui-web", "cui-native", "cui-desktop"} + acceptTypes := []context.Accept{context.AcceptWebCUI, context.AccepNativeCUI, context.AcceptDesktopCUI} for _, acceptType := range acceptTypes { - t.Run(acceptType, func(t *testing.T) { - cxt := &Context{ - ChatID: "test-chat-id", - AssistantID: "test-assistant-id", - Context: context.Background(), - Accept: Accept(acceptType), - Locale: "en", - Writer: newMockResponseWriter(), - } + t.Run(string(acceptType), func(t *testing.T) { + cxt := context.New(stdContext.Background(), nil, "test-chat-id") + cxt.AssistantID = "test-assistant-id" + cxt.Accept = acceptType + cxt.Locale = "en" + cxt.Writer = newTestMockResponseWriter() res, err := v8.Call(v8.CallOptions{}, ` function test(ctx) { @@ -326,7 +307,7 @@ func TestJsValueSendWithCUIAccept(t *testing.T) { if !ok { t.Fatalf("Expected map result, got %T", res) } - assert.Equal(t, true, result["success"], "Send with "+acceptType+" should succeed") + assert.Equal(t, true, result["success"], "Send with "+string(acceptType)+" should succeed") }) } } @@ -338,15 +319,11 @@ func TestJsValueSendChainedCalls(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - cxt := &Context{ - ChatID: "test-chat-id", - AssistantID: "test-assistant-id", - Context: context.Background(), - Accept: "standard", - Locale: "en", - Writer: newMockResponseWriter(), - IDGenerator: message.NewIDGenerator(), - } + cxt := context.New(stdContext.Background(), nil, "test-chat-id") + cxt.AssistantID = "test-assistant-id" + cxt.Accept = context.AcceptStandard + cxt.Locale = "en" + cxt.Writer = newTestMockResponseWriter() res, err := v8.Call(v8.CallOptions{}, ` function test(ctx) { @@ -379,15 +356,11 @@ func TestJsValueIDGenerators(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - cxt := &Context{ - ChatID: "test-chat-id", - AssistantID: "test-assistant-id", - Context: context.Background(), - Accept: "standard", - Locale: "en", - Writer: newMockResponseWriter(), - IDGenerator: message.NewIDGenerator(), - } + cxt := context.New(stdContext.Background(), nil, "test-chat-id") + cxt.AssistantID = "test-assistant-id" + cxt.Accept = context.AcceptStandard + cxt.Locale = "en" + cxt.Writer = newTestMockResponseWriter() res, err := v8.Call(v8.CallOptions{}, ` function test(ctx) { @@ -456,15 +429,11 @@ func TestJsValueSendWithBlockID(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - cxt := &Context{ - ChatID: "test-chat-id", - AssistantID: "test-assistant-id", - Context: context.Background(), - Accept: "standard", - Locale: "en", - Writer: newMockResponseWriter(), - IDGenerator: message.NewIDGenerator(), - } + cxt := context.New(stdContext.Background(), nil, "test-chat-id") + cxt.AssistantID = "test-assistant-id" + cxt.Accept = context.AcceptStandard + cxt.Locale = "en" + cxt.Writer = newTestMockResponseWriter() res, err := v8.Call(v8.CallOptions{}, ` function test(ctx) { @@ -513,15 +482,11 @@ func TestJsValueReplace(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - cxt := &Context{ - ChatID: "test-chat-id", - AssistantID: "test-assistant-id", - Context: context.Background(), - Accept: "standard", - Locale: "en", - Writer: newMockResponseWriter(), - IDGenerator: message.NewIDGenerator(), - } + cxt := context.New(stdContext.Background(), nil, "test-chat-id") + cxt.AssistantID = "test-assistant-id" + cxt.Accept = context.AcceptStandard + cxt.Locale = "en" + cxt.Writer = newTestMockResponseWriter() res, err := v8.Call(v8.CallOptions{}, ` function test(ctx) { @@ -560,15 +525,11 @@ func TestJsValueAppend(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - cxt := &Context{ - ChatID: "test-chat-id", - AssistantID: "test-assistant-id", - Context: context.Background(), - Accept: "standard", - Locale: "en", - Writer: newMockResponseWriter(), - IDGenerator: message.NewIDGenerator(), - } + cxt := context.New(stdContext.Background(), nil, "test-chat-id") + cxt.AssistantID = "test-assistant-id" + cxt.Accept = context.AcceptStandard + cxt.Locale = "en" + cxt.Writer = newTestMockResponseWriter() res, err := v8.Call(v8.CallOptions{}, ` function test(ctx) { @@ -610,15 +571,11 @@ func TestJsValueMerge(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - cxt := &Context{ - ChatID: "test-chat-id", - AssistantID: "test-assistant-id", - Context: context.Background(), - Accept: "standard", - Locale: "en", - Writer: newMockResponseWriter(), - IDGenerator: message.NewIDGenerator(), - } + cxt := context.New(stdContext.Background(), nil, "test-chat-id") + cxt.AssistantID = "test-assistant-id" + cxt.Accept = context.AcceptStandard + cxt.Locale = "en" + cxt.Writer = newTestMockResponseWriter() res, err := v8.Call(v8.CallOptions{}, ` function test(ctx) { @@ -668,15 +625,11 @@ func TestJsValueSet(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - cxt := &Context{ - ChatID: "test-chat-id", - AssistantID: "test-assistant-id", - Context: context.Background(), - Accept: "standard", - Locale: "en", - Writer: newMockResponseWriter(), - IDGenerator: message.NewIDGenerator(), - } + cxt := context.New(stdContext.Background(), nil, "test-chat-id") + cxt.AssistantID = "test-assistant-id" + cxt.Accept = context.AcceptStandard + cxt.Locale = "en" + cxt.Writer = newTestMockResponseWriter() res, err := v8.Call(v8.CallOptions{}, ` function test(ctx) { @@ -726,15 +679,11 @@ func TestJsValueBlockIDInheritance(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - cxt := &Context{ - ChatID: "test-chat-id", - AssistantID: "test-assistant-id", - Context: context.Background(), - Accept: "standard", - Locale: "en", - Writer: newMockResponseWriter(), - IDGenerator: message.NewIDGenerator(), - } + cxt := context.New(stdContext.Background(), nil, "test-chat-id") + cxt.AssistantID = "test-assistant-id" + cxt.Accept = context.AcceptStandard + cxt.Locale = "en" + cxt.Writer = newTestMockResponseWriter() res, err := v8.Call(v8.CallOptions{}, ` function test(ctx) { @@ -781,12 +730,12 @@ func TestJsValueEndBlock(t *testing.T) { defer test.Clean() // Setup mock writer - mockWriter := newMockResponseWriter() + mockWriter := newTestMockResponseWriter() // Use New() to properly initialize messageMetadata - cxt := New(context.Background(), nil, "test-chat-id") + cxt := context.New(stdContext.Background(), nil, "test-chat-id") cxt.AssistantID = "test-assistant-id" - cxt.Accept = AcceptWebCUI + cxt.Accept = context.AcceptWebCUI cxt.Locale = "en" cxt.Writer = mockWriter @@ -834,12 +783,12 @@ func TestJsValueSendStream(t *testing.T) { defer test.Clean() // Setup mock writer - mockWriter := newMockResponseWriter() + mockWriter := newTestMockResponseWriter() // Use New() to properly initialize messageMetadata - cxt := New(context.Background(), nil, "test-chat-id") + cxt := context.New(stdContext.Background(), nil, "test-chat-id") cxt.AssistantID = "test-assistant-id" - cxt.Accept = AcceptWebCUI + cxt.Accept = context.AcceptWebCUI cxt.Locale = "en" cxt.Writer = mockWriter @@ -888,11 +837,11 @@ func TestJsValueSendStreamWithBlockID(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - mockWriter := newMockResponseWriter() + mockWriter := newTestMockResponseWriter() - cxt := New(context.Background(), nil, "test-chat-id") + cxt := context.New(stdContext.Background(), nil, "test-chat-id") cxt.AssistantID = "test-assistant-id" - cxt.Accept = AcceptWebCUI + cxt.Accept = context.AcceptWebCUI cxt.Locale = "en" cxt.Writer = mockWriter @@ -935,11 +884,11 @@ func TestJsValueEnd(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - mockWriter := newMockResponseWriter() + mockWriter := newTestMockResponseWriter() - cxt := New(context.Background(), nil, "test-chat-id") + cxt := context.New(stdContext.Background(), nil, "test-chat-id") cxt.AssistantID = "test-assistant-id" - cxt.Accept = AcceptWebCUI + cxt.Accept = context.AcceptWebCUI cxt.Locale = "en" cxt.Writer = mockWriter @@ -984,11 +933,11 @@ func TestJsValueEndWithFinalContent(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - mockWriter := newMockResponseWriter() + mockWriter := newTestMockResponseWriter() - cxt := New(context.Background(), nil, "test-chat-id") + cxt := context.New(stdContext.Background(), nil, "test-chat-id") cxt.AssistantID = "test-assistant-id" - cxt.Accept = AcceptWebCUI + cxt.Accept = context.AcceptWebCUI cxt.Locale = "en" cxt.Writer = mockWriter @@ -1033,11 +982,11 @@ func TestJsValueStreamingWorkflow(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - mockWriter := newMockResponseWriter() + mockWriter := newTestMockResponseWriter() - cxt := New(context.Background(), nil, "test-chat-id") + cxt := context.New(stdContext.Background(), nil, "test-chat-id") cxt.AssistantID = "test-assistant-id" - cxt.Accept = AcceptWebCUI + cxt.Accept = context.AcceptWebCUI cxt.Locale = "en" cxt.Writer = mockWriter @@ -1091,11 +1040,11 @@ func TestJsValueSendStreamStringShorthand(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - mockWriter := newMockResponseWriter() + mockWriter := newTestMockResponseWriter() - cxt := New(context.Background(), nil, "test-chat-id") + cxt := context.New(stdContext.Background(), nil, "test-chat-id") cxt.AssistantID = "test-assistant-id" - cxt.Accept = AcceptWebCUI + cxt.Accept = context.AcceptWebCUI cxt.Locale = "en" cxt.Writer = mockWriter @@ -1133,11 +1082,11 @@ func TestJsValueEndErrorHandling(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - mockWriter := newMockResponseWriter() + mockWriter := newTestMockResponseWriter() - cxt := New(context.Background(), nil, "test-chat-id") + cxt := context.New(stdContext.Background(), nil, "test-chat-id") cxt.AssistantID = "test-assistant-id" - cxt.Accept = AcceptWebCUI + cxt.Accept = context.AcceptWebCUI cxt.Locale = "en" cxt.Writer = mockWriter @@ -1169,11 +1118,11 @@ func TestJsValueEndWithInvalidMessageID(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - mockWriter := newMockResponseWriter() + mockWriter := newTestMockResponseWriter() - cxt := New(context.Background(), nil, "test-chat-id") + cxt := context.New(stdContext.Background(), nil, "test-chat-id") cxt.AssistantID = "test-assistant-id" - cxt.Accept = AcceptWebCUI + cxt.Accept = context.AcceptWebCUI cxt.Locale = "en" cxt.Writer = mockWriter @@ -1205,11 +1154,11 @@ func TestJsValueSendStreamErrorHandling(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - mockWriter := newMockResponseWriter() + mockWriter := newTestMockResponseWriter() - cxt := New(context.Background(), nil, "test-chat-id") + cxt := context.New(stdContext.Background(), nil, "test-chat-id") cxt.AssistantID = "test-assistant-id" - cxt.Accept = AcceptWebCUI + cxt.Accept = context.AcceptWebCUI cxt.Locale = "en" cxt.Writer = mockWriter @@ -1241,11 +1190,11 @@ func TestJsValueMultipleStreams(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - mockWriter := newMockResponseWriter() + mockWriter := newTestMockResponseWriter() - cxt := New(context.Background(), nil, "test-chat-id") + cxt := context.New(stdContext.Background(), nil, "test-chat-id") cxt.AssistantID = "test-assistant-id" - cxt.Accept = AcceptWebCUI + cxt.Accept = context.AcceptWebCUI cxt.Locale = "en" cxt.Writer = mockWriter @@ -1296,10 +1245,10 @@ func TestJsValueSendVsSendStream(t *testing.T) { // 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") + mockWriter := newTestMockResponseWriter() + cxt := context.New(stdContext.Background(), nil, "test-chat-id") cxt.AssistantID = "test-assistant-id" - cxt.Accept = AcceptWebCUI + cxt.Accept = context.AcceptWebCUI cxt.Locale = "en" cxt.Writer = mockWriter @@ -1319,10 +1268,10 @@ func TestJsValueSendVsSendStream(t *testing.T) { // 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") + mockWriter := newTestMockResponseWriter() + cxt := context.New(stdContext.Background(), nil, "test-chat-id") cxt.AssistantID = "test-assistant-id" - cxt.Accept = AcceptWebCUI + cxt.Accept = context.AcceptWebCUI cxt.Locale = "en" cxt.Writer = mockWriter diff --git a/agent/context/log.go b/agent/context/log.go new file mode 100644 index 00000000..d6cb7128 --- /dev/null +++ b/agent/context/log.go @@ -0,0 +1,495 @@ +package context + +import ( + "fmt" + "strings" + "sync" + "time" + + kunlog "github.com/yaoapp/kun/log" + "github.com/yaoapp/yao/config" +) + +// ============================================================================= +// ANSI Color Codes +// ============================================================================= + +const ( + colorReset = "\033[0m" + colorRed = "\033[31m" + colorGreen = "\033[32m" + colorYellow = "\033[33m" + colorBlue = "\033[34m" + colorMagenta = "\033[35m" + colorCyan = "\033[36m" + colorWhite = "\033[37m" + colorGray = "\033[90m" + + colorBoldRed = "\033[1;31m" + colorBoldGreen = "\033[1;32m" + colorBoldYellow = "\033[1;33m" + colorBoldBlue = "\033[1;34m" + colorBoldMagenta = "\033[1;35m" + colorBoldCyan = "\033[1;36m" +) + +// ============================================================================= +// Log Level +// ============================================================================= + +// LogLevel represents log severity +type LogLevel int + +const ( + LogLevelTrace LogLevel = iota + LogLevelDebug + LogLevelInfo + LogLevelWarn + LogLevelError +) + +// ============================================================================= +// Log Entry +// ============================================================================= + +// LogEntry represents a single log entry +type LogEntry struct { + Level LogLevel + Message string + Timestamp time.Time + Phase string // For phase logging + Elapsed time.Duration +} + +// ============================================================================= +// Request Logger +// ============================================================================= + +// RequestLogger provides request-scoped async logging +type RequestLogger struct { + assistantID string + chatID string + requestID string + shortID string // Short version of requestID for display + startTime time.Time + + ch chan LogEntry + done chan struct{} + once sync.Once + closed bool + noop bool // noop logger does nothing (for nil safety) + mu sync.RWMutex +} + +// noopLogger is a shared no-op logger instance +var noopLogger = &RequestLogger{noop: true} + +// NewRequestLogger creates a new request-scoped logger with async processing +func NewRequestLogger(assistantID, chatID, requestID string) *RequestLogger { + l := &RequestLogger{ + assistantID: assistantID, + chatID: chatID, + requestID: requestID, + shortID: shortID(requestID), + startTime: time.Now(), + ch: make(chan LogEntry, 100), // Buffered channel + done: make(chan struct{}), + } + + // Start consumer goroutine + go l.consume() + + return l +} + +// Noop returns a no-op logger that does nothing (nil-safe) +func Noop() *RequestLogger { + return noopLogger +} + +// SetAssistantID sets the assistant ID (called when entering Stream) +func (l *RequestLogger) SetAssistantID(id string) { + if l.noop { + return + } + l.assistantID = id +} + +// Close closes the logger and waits for all entries to be processed +func (l *RequestLogger) Close() { + if l.noop { + return + } + l.once.Do(func() { + l.mu.Lock() + l.closed = true + l.mu.Unlock() + + close(l.ch) + <-l.done // Wait for consumer to finish + }) +} + +// consume processes log entries from the channel +func (l *RequestLogger) consume() { + defer close(l.done) + + for entry := range l.ch { + l.processEntry(entry) + } +} + +// processEntry handles a single log entry based on mode +func (l *RequestLogger) processEntry(entry LogEntry) { + if config.IsDevelopment() { + l.printDev(entry) + } else { + l.printProd(entry) + } +} + +// printDev prints colorful output for development mode +func (l *RequestLogger) printDev(entry LogEntry) { + switch entry.Level { + case LogLevelTrace: + fmt.Printf("%s → %s%s\n", colorGray, entry.Message, colorReset) + case LogLevelDebug: + fmt.Printf("%s • %s%s\n", colorGray, entry.Message, colorReset) + case LogLevelInfo: + fmt.Printf("%s ℹ %s%s\n", colorCyan, entry.Message, colorReset) + case LogLevelWarn: + fmt.Printf("%s ⚠ %s%s\n", colorYellow, entry.Message, colorReset) + case LogLevelError: + fmt.Printf("%s ✗ %s%s\n", colorRed, entry.Message, colorReset) + } +} + +// printProd logs to kun/log for production mode +func (l *RequestLogger) printProd(entry LogEntry) { + prefix := fmt.Sprintf("[AGENT] %s ", l.shortID) + + switch entry.Level { + case LogLevelTrace: + kunlog.Trace("%s%s", prefix, entry.Message) + case LogLevelDebug: + // Skip debug in production + case LogLevelInfo: + kunlog.Info("%s%s", prefix, entry.Message) + case LogLevelWarn: + kunlog.Warn("%s%s", prefix, entry.Message) + case LogLevelError: + kunlog.Error("%s%s", prefix, entry.Message) + } +} + +// send sends an entry to the channel (non-blocking if closed) +func (l *RequestLogger) send(entry LogEntry) { + if l.noop { + return + } + + l.mu.RLock() + closed := l.closed + l.mu.RUnlock() + + if closed { + return + } + + entry.Timestamp = time.Now() + select { + case l.ch <- entry: + default: + // Channel full, drop the log (shouldn't happen with buffered channel) + } +} + +// ============================================================================= +// Standard Log Interface +// ============================================================================= + +// Trace logs a trace level message +func (l *RequestLogger) Trace(format string, args ...interface{}) { + l.send(LogEntry{ + Level: LogLevelTrace, + Message: fmt.Sprintf(format, args...), + }) +} + +// Debug logs a debug level message +func (l *RequestLogger) Debug(format string, args ...interface{}) { + l.send(LogEntry{ + Level: LogLevelDebug, + Message: fmt.Sprintf(format, args...), + }) +} + +// Info logs an info level message +func (l *RequestLogger) Info(format string, args ...interface{}) { + l.send(LogEntry{ + Level: LogLevelInfo, + Message: fmt.Sprintf(format, args...), + }) +} + +// Warn logs a warning level message +func (l *RequestLogger) Warn(format string, args ...interface{}) { + l.send(LogEntry{ + Level: LogLevelWarn, + Message: fmt.Sprintf(format, args...), + }) +} + +// Error logs an error level message +func (l *RequestLogger) Error(format string, args ...interface{}) { + l.send(LogEntry{ + Level: LogLevelError, + Message: fmt.Sprintf(format, args...), + }) +} + +// ============================================================================= +// Business Quick Functions +// ============================================================================= + +// Start logs the start of a request with visual separator +func (l *RequestLogger) Start() { + if l.noop { + return + } + + if !config.IsDevelopment() { + kunlog.Trace("[AGENT] Request %s started: assistant=%s, chat=%s, request=%s", + l.shortID, l.assistantID, shortID(l.chatID), shortID(l.requestID)) + return + } + + // Development: colorful output (direct print, not through channel for immediate display) + fmt.Println() + fmt.Printf("%s%s%s\n", colorBoldCyan, strings.Repeat("═", 60), colorReset) + fmt.Printf("%s 🚀 AGENT REQUEST %s%s\n", colorBoldCyan, l.shortID, colorReset) + fmt.Printf("%s%s%s\n", colorBoldCyan, strings.Repeat("─", 60), colorReset) + fmt.Printf("%s Assistant: %s%s%s\n", colorGray, colorWhite, l.assistantID, colorReset) + fmt.Printf("%s Chat ID: %s%s%s\n", colorGray, colorWhite, l.chatID, colorReset) + fmt.Printf("%s Request: %s%s%s\n", colorGray, colorWhite, l.requestID, colorReset) + fmt.Printf("%s Time: %s%s%s\n", colorGray, colorWhite, l.startTime.Format("15:04:05.000"), colorReset) + fmt.Printf("%s%s%s\n", colorCyan, strings.Repeat("─", 60), colorReset) +} + +// End logs the end of a request with summary +func (l *RequestLogger) End(success bool, err error) { + if l.noop { + return + } + + duration := time.Since(l.startTime) + + if !config.IsDevelopment() { + if success { + kunlog.Trace("[AGENT] Request %s completed: assistant=%s, duration=%v", + l.shortID, l.assistantID, duration.Round(time.Millisecond)) + } else { + kunlog.Trace("[AGENT] Request %s failed: assistant=%s, duration=%v, error=%v", + l.shortID, l.assistantID, duration.Round(time.Millisecond), err) + } + return + } + + // Development: colorful output (direct print for immediate display) + fmt.Printf("%s%s%s\n", colorCyan, strings.Repeat("─", 60), colorReset) + if success { + fmt.Printf("%s ✅ REQUEST %s COMPLETED%s\n", colorBoldGreen, l.shortID, colorReset) + } else { + fmt.Printf("%s ❌ REQUEST %s FAILED%s\n", colorBoldRed, l.shortID, colorReset) + if err != nil { + fmt.Printf("%s Error: %s%v%s\n", colorGray, colorRed, err, colorReset) + } + } + fmt.Printf("%s Assistant: %s%s%s\n", colorGray, colorWhite, l.assistantID, colorReset) + fmt.Printf("%s Duration: %s%v%s\n", colorGray, colorWhite, duration.Round(time.Millisecond), colorReset) + fmt.Printf("%s%s%s\n", colorCyan, strings.Repeat("─", 60), colorReset) + fmt.Println() +} + +// Phase logs a major phase in the request lifecycle +func (l *RequestLogger) Phase(name string) { + if l.noop { + return + } + + elapsed := time.Since(l.startTime).Round(time.Millisecond) + + if config.IsDevelopment() { + fmt.Printf("%s ▶ %s%s %s[+%v]%s\n", colorBoldBlue, name, colorReset, colorGray, elapsed, colorReset) + } else { + kunlog.Trace("[AGENT] %s Phase: %s (+%v)", l.shortID, name, elapsed) + } +} + +// PhaseComplete logs the completion of a phase +func (l *RequestLogger) PhaseComplete(name string) { + if l.noop { + return + } + + elapsed := time.Since(l.startTime).Round(time.Millisecond) + + if config.IsDevelopment() { + fmt.Printf("%s ✓ %s%s %s[+%v]%s\n", colorGreen, name, colorReset, colorGray, elapsed, colorReset) + } else { + kunlog.Trace("[AGENT] %s Phase completed: %s (+%v)", l.shortID, name, elapsed) + } +} + +// PhaseSkip logs a skipped phase (development only) +func (l *RequestLogger) PhaseSkip(name, reason string) { + if l.noop { + return + } + if config.IsDevelopment() { + fmt.Printf("%s ⊘ %s (%s)%s\n", colorGray, name, reason, colorReset) + } +} + +// LLMStart logs the start of an LLM call +func (l *RequestLogger) LLMStart(connector, model string, messageCount int) { + if l.noop { + return + } + + elapsed := time.Since(l.startTime).Round(time.Millisecond) + + if config.IsDevelopment() { + fmt.Printf("%s 🤖 LLM Call%s %s[+%v]%s\n", colorBoldMagenta, colorReset, colorGray, elapsed, colorReset) + fmt.Printf("%s Connector: %s%s%s\n", colorGray, colorWhite, connector, colorReset) + if model != "" { + fmt.Printf("%s Model: %s%s%s\n", colorGray, colorWhite, model, colorReset) + } + fmt.Printf("%s Messages: %s%d%s\n", colorGray, colorWhite, messageCount, colorReset) + } else { + kunlog.Trace("[AGENT] %s LLM call: connector=%s, model=%s, messages=%d (+%v)", l.shortID, connector, model, messageCount, elapsed) + } +} + +// LLMComplete logs the completion of an LLM call +func (l *RequestLogger) LLMComplete(tokens int, hasToolCalls bool) { + elapsed := time.Since(l.startTime).Round(time.Millisecond) + status := "streaming" + if hasToolCalls { + status = "tool_calls" + } + + if config.IsDevelopment() { + fmt.Printf("%s ✓ LLM Response (%s)%s", colorGreen, status, colorReset) + if tokens > 0 { + fmt.Printf(" %s[tokens: %d]%s", colorGray, tokens, colorReset) + } + fmt.Printf(" %s[+%v]%s\n", colorGray, elapsed, colorReset) + } else { + kunlog.Trace("[AGENT] %s LLM response: status=%s, tokens=%d (+%v)", l.shortID, status, tokens, elapsed) + } +} + +// ToolStart logs the start of tool execution +func (l *RequestLogger) ToolStart(toolName string) { + if config.IsDevelopment() { + fmt.Printf("%s 🔧 Tool: %s%s\n", colorYellow, toolName, colorReset) + } else { + kunlog.Trace("[AGENT] %s Tool call: %s", l.shortID, toolName) + } +} + +// ToolComplete logs the completion of tool execution +func (l *RequestLogger) ToolComplete(toolName string, success bool) { + if config.IsDevelopment() { + if success { + fmt.Printf("%s ✓ %s completed%s\n", colorGreen, toolName, colorReset) + } else { + fmt.Printf("%s ✗ %s failed%s\n", colorRed, toolName, colorReset) + } + } else { + if success { + kunlog.Trace("[AGENT] %s Tool completed: %s", l.shortID, toolName) + } else { + kunlog.Trace("[AGENT] %s Tool failed: %s", l.shortID, toolName) + } + } +} + +// HookStart logs the start of a hook execution +func (l *RequestLogger) HookStart(hookName string) { + elapsed := time.Since(l.startTime).Round(time.Millisecond) + + if config.IsDevelopment() { + fmt.Printf("%s 🪝 Hook: %s%s %s[+%v]%s\n", colorMagenta, hookName, colorReset, colorGray, elapsed, colorReset) + } else { + kunlog.Trace("[AGENT] %s Hook: %s (+%v)", l.shortID, hookName, elapsed) + } +} + +// HookComplete logs the completion of a hook +func (l *RequestLogger) HookComplete(hookName string) { + if config.IsDevelopment() { + fmt.Printf("%s ✓ %s done%s\n", colorGreen, hookName, colorReset) + } else { + kunlog.Trace("[AGENT] %s Hook completed: %s", l.shortID, hookName) + } +} + +// Cleanup logs resource cleanup +func (l *RequestLogger) Cleanup(resource string) { + if l.noop { + return + } + + if config.IsDevelopment() { + fmt.Printf("%s ✓ %s%s\n", colorGray, resource, colorReset) + } else { + kunlog.Trace("[AGENT] %s Cleanup: %s", l.shortID, resource) + } +} + +// HistoryLoad logs history loading +func (l *RequestLogger) HistoryLoad(count, maxSize int) { + if config.IsDevelopment() { + fmt.Printf("%s Loaded %d/%d history messages%s\n", colorGray, count, maxSize, colorReset) + } else { + kunlog.Trace("[AGENT] %s History loaded: %d/%d messages", l.shortID, count, maxSize) + } +} + +// HistoryOverlap logs overlap detection +func (l *RequestLogger) HistoryOverlap(overlapCount int) { + if overlapCount > 0 { + if config.IsDevelopment() { + fmt.Printf("%s Removed %d overlapping messages%s\n", colorYellow, overlapCount, colorReset) + } else { + kunlog.Trace("[AGENT] %s History overlap removed: %d messages", l.shortID, overlapCount) + } + } +} + +// Release logs the start of resource release phase +func (l *RequestLogger) Release() { + if l.noop { + return + } + + if config.IsDevelopment() { + fmt.Printf("%s 🧹 RELEASE %s%s %s(%s)%s\n", colorBoldYellow, l.shortID, colorReset, colorGray, l.assistantID, colorReset) + } else { + kunlog.Trace("[AGENT] %s Release started", l.shortID) + } +} + +// ============================================================================= +// Helper +// ============================================================================= + +// shortID returns first 8 characters of an ID +func shortID(id string) string { + if len(id) > 8 { + return id[:8] + } + return id +} diff --git a/agent/context/message_test.go b/agent/context/message_test.go index 7c20b712..f250b381 100644 --- a/agent/context/message_test.go +++ b/agent/context/message_test.go @@ -1,8 +1,10 @@ -package context +package context_test import ( "encoding/json" "testing" + + "github.com/yaoapp/yao/agent/context" ) func TestMessage_UnmarshalJSON_StringContent(t *testing.T) { @@ -11,14 +13,14 @@ func TestMessage_UnmarshalJSON_StringContent(t *testing.T) { "content": "Hello, world!" }` - var msg Message + var msg context.Message err := json.Unmarshal([]byte(jsonData), &msg) if err != nil { t.Fatalf("Failed to unmarshal: %v", err) } - if msg.Role != RoleUser { - t.Errorf("Expected role %s, got %s", RoleUser, msg.Role) + if msg.Role != context.RoleUser { + t.Errorf("Expected role %s, got %s", context.RoleUser, msg.Role) } content, ok := msg.GetContentAsString() @@ -49,14 +51,14 @@ func TestMessage_UnmarshalJSON_ArrayContent(t *testing.T) { ] }` - var msg Message + var msg context.Message err := json.Unmarshal([]byte(jsonData), &msg) if err != nil { t.Fatalf("Failed to unmarshal: %v", err) } - if msg.Role != RoleUser { - t.Errorf("Expected role %s, got %s", RoleUser, msg.Role) + if msg.Role != context.RoleUser { + t.Errorf("Expected role %s, got %s", context.RoleUser, msg.Role) } parts, ok := msg.GetContentAsParts() @@ -69,16 +71,16 @@ func TestMessage_UnmarshalJSON_ArrayContent(t *testing.T) { } // Check first part (text) - if parts[0].Type != ContentText { - t.Errorf("Expected type %s, got %s", ContentText, parts[0].Type) + if parts[0].Type != context.ContentText { + t.Errorf("Expected type %s, got %s", context.ContentText, parts[0].Type) } if parts[0].Text != "What's in this image?" { t.Errorf("Expected text 'What's in this image?', got '%s'", parts[0].Text) } // Check second part (image) - if parts[1].Type != ContentImageURL { - t.Errorf("Expected type %s, got %s", ContentImageURL, parts[1].Type) + if parts[1].Type != context.ContentImageURL { + t.Errorf("Expected type %s, got %s", context.ContentImageURL, parts[1].Type) } if parts[1].ImageURL == nil { t.Fatal("Expected ImageURL to be non-nil") @@ -86,8 +88,8 @@ func TestMessage_UnmarshalJSON_ArrayContent(t *testing.T) { if parts[1].ImageURL.URL != "https://example.com/image.jpg" { t.Errorf("Expected URL 'https://example.com/image.jpg', got '%s'", parts[1].ImageURL.URL) } - if parts[1].ImageURL.Detail != DetailHigh { - t.Errorf("Expected detail %s, got %s", DetailHigh, parts[1].ImageURL.Detail) + if parts[1].ImageURL.Detail != context.DetailHigh { + t.Errorf("Expected detail %s, got %s", context.DetailHigh, parts[1].ImageURL.Detail) } } @@ -107,14 +109,14 @@ func TestMessage_UnmarshalJSON_NullContent(t *testing.T) { ] }` - var msg Message + var msg context.Message err := json.Unmarshal([]byte(jsonData), &msg) if err != nil { t.Fatalf("Failed to unmarshal: %v", err) } - if msg.Role != RoleAssistant { - t.Errorf("Expected role %s, got %s", RoleAssistant, msg.Role) + if msg.Role != context.RoleAssistant { + t.Errorf("Expected role %s, got %s", context.RoleAssistant, msg.Role) } if msg.Content != nil { @@ -142,7 +144,7 @@ func TestMessage_UnmarshalJSON_WithRefusal(t *testing.T) { "refusal": "I cannot help with that request." }` - var msg Message + var msg context.Message err := json.Unmarshal([]byte(jsonData), &msg) if err != nil { t.Fatalf("Failed to unmarshal: %v", err) @@ -179,7 +181,7 @@ func TestMessage_UnmarshalJSON_AudioContent(t *testing.T) { ] }` - var msg Message + var msg context.Message err := json.Unmarshal([]byte(jsonData), &msg) if err != nil { t.Fatalf("Failed to unmarshal: %v", err) @@ -195,8 +197,8 @@ func TestMessage_UnmarshalJSON_AudioContent(t *testing.T) { } // Check audio part - if parts[1].Type != ContentInputAudio { - t.Errorf("Expected type %s, got %s", ContentInputAudio, parts[1].Type) + if parts[1].Type != context.ContentInputAudio { + t.Errorf("Expected type %s, got %s", context.ContentInputAudio, parts[1].Type) } if parts[1].InputAudio == nil { t.Fatal("Expected InputAudio to be non-nil") @@ -210,7 +212,7 @@ func TestMessage_UnmarshalJSON_AudioContent(t *testing.T) { } func TestMessage_MarshalJSON_StringContent(t *testing.T) { - msg := NewTextMessage(RoleUser, "Hello, AI!") + msg := context.NewTextMessage(context.RoleUser, "Hello, AI!") data, err := json.Marshal(msg) if err != nil { @@ -223,8 +225,8 @@ func TestMessage_MarshalJSON_StringContent(t *testing.T) { t.Fatalf("Failed to unmarshal result: %v", err) } - if result["role"] != string(RoleUser) { - t.Errorf("Expected role %s, got %v", RoleUser, result["role"]) + if result["role"] != string(context.RoleUser) { + t.Errorf("Expected role %s, got %v", context.RoleUser, result["role"]) } if result["content"] != "Hello, AI!" { @@ -233,21 +235,21 @@ func TestMessage_MarshalJSON_StringContent(t *testing.T) { } func TestMessage_MarshalJSON_ArrayContent(t *testing.T) { - parts := []ContentPart{ + parts := []context.ContentPart{ { - Type: ContentText, + Type: context.ContentText, Text: "Describe this image", }, { - Type: ContentImageURL, - ImageURL: &ImageURL{ + Type: context.ContentImageURL, + ImageURL: &context.ImageURL{ URL: "https://example.com/test.jpg", - Detail: DetailLow, + Detail: context.DetailLow, }, }, } - msg := NewMultipartMessage(RoleUser, parts) + msg := context.NewMultipartMessage(context.RoleUser, parts) data, err := json.Marshal(msg) if err != nil { @@ -255,7 +257,7 @@ func TestMessage_MarshalJSON_ArrayContent(t *testing.T) { } // Unmarshal back to verify - var result Message + var result context.Message err = json.Unmarshal(data, &result) if err != nil { t.Fatalf("Failed to unmarshal result: %v", err) @@ -272,14 +274,14 @@ func TestMessage_MarshalJSON_ArrayContent(t *testing.T) { } func TestMessage_MarshalJSON_WithToolCalls(t *testing.T) { - msg := &Message{ - Role: RoleAssistant, + msg := &context.Message{ + Role: context.RoleAssistant, Content: nil, - ToolCalls: []ToolCall{ + ToolCalls: []context.ToolCall{ { ID: "call_abc123", - Type: ToolTypeFunction, - Function: Function{ + Type: context.ToolTypeFunction, + Function: context.Function{ Name: "get_weather", Arguments: `{"location":"San Francisco"}`, }, @@ -293,7 +295,7 @@ func TestMessage_MarshalJSON_WithToolCalls(t *testing.T) { } // Unmarshal back to verify - var result Message + var result context.Message err = json.Unmarshal(data, &result) if err != nil { t.Fatalf("Failed to unmarshal result: %v", err) @@ -320,14 +322,14 @@ func TestMessage_ToolMessage(t *testing.T) { "content": "The weather in San Francisco is sunny, 72°F" }` - var msg Message + var msg context.Message err := json.Unmarshal([]byte(jsonData), &msg) if err != nil { t.Fatalf("Failed to unmarshal: %v", err) } - if msg.Role != RoleTool { - t.Errorf("Expected role %s, got %s", RoleTool, msg.Role) + if msg.Role != context.RoleTool { + t.Errorf("Expected role %s, got %s", context.RoleTool, msg.Role) } if msg.ToolCallID == nil { @@ -349,10 +351,10 @@ func TestMessage_ToolMessage(t *testing.T) { } func TestNewTextMessage(t *testing.T) { - msg := NewTextMessage(RoleSystem, "You are a helpful assistant.") + msg := context.NewTextMessage(context.RoleSystem, "You are a helpful assistant.") - if msg.Role != RoleSystem { - t.Errorf("Expected role %s, got %s", RoleSystem, msg.Role) + if msg.Role != context.RoleSystem { + t.Errorf("Expected role %s, got %s", context.RoleSystem, msg.Role) } content, ok := msg.GetContentAsString() @@ -366,14 +368,14 @@ func TestNewTextMessage(t *testing.T) { } func TestNewMultipartMessage(t *testing.T) { - parts := []ContentPart{ - {Type: ContentText, Text: "Hello"}, + parts := []context.ContentPart{ + {Type: context.ContentText, Text: "Hello"}, } - msg := NewMultipartMessage(RoleUser, parts) + msg := context.NewMultipartMessage(context.RoleUser, parts) - if msg.Role != RoleUser { - t.Errorf("Expected role %s, got %s", RoleUser, msg.Role) + if msg.Role != context.RoleUser { + t.Errorf("Expected role %s, got %s", context.RoleUser, msg.Role) } resultParts, ok := msg.GetContentAsParts() diff --git a/agent/context/openapi_test.go b/agent/context/openapi_test.go index 639b55a4..f65bfe67 100644 --- a/agent/context/openapi_test.go +++ b/agent/context/openapi_test.go @@ -1,30 +1,54 @@ -package context +package context_test import ( "bytes" "encoding/json" + "io" "net/http/httptest" "testing" "github.com/gin-gonic/gin" "github.com/yaoapp/gou/store" + "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/config" "github.com/yaoapp/yao/test" ) +// parseCompletionRequestData is a helper function for tests to parse completion request data +func parseCompletionRequestData(c *gin.Context) (*context.CompletionRequest, error) { + var req context.CompletionRequest + + if c.Request.Body != nil { + body, err := io.ReadAll(c.Request.Body) + if err != nil { + return nil, err + } + c.Request.Body = io.NopCloser(bytes.NewBuffer(body)) + if len(body) > 0 { + if err := json.Unmarshal(body, &req); err != nil { + return nil, err + } + if len(req.Messages) > 0 { + return &req, nil + } + } + } + return &req, nil +} + func TestGetMessages_FromBody(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() gin.SetMode(gin.TestMode) - messages := []Message{ + messages := []context.Message{ { - Role: RoleUser, + Role: context.RoleUser, Content: "Hello, world!", }, { - Role: RoleAssistant, + Role: context.RoleAssistant, Content: "Hi there!", }, } @@ -45,7 +69,7 @@ func TestGetMessages_FromBody(t *testing.T) { // Parse request first completionReq, _ := parseCompletionRequestData(c) - result, err := GetMessages(c, completionReq) + result, err := context.GetMessages(c, completionReq) if err != nil { t.Fatalf("Failed to get messages: %v", err) } @@ -54,8 +78,8 @@ func TestGetMessages_FromBody(t *testing.T) { t.Errorf("Expected 2 messages, got %d", len(result)) } - if result[0].Role != RoleUser { - t.Errorf("Expected first message role to be %s, got %s", RoleUser, result[0].Role) + if result[0].Role != context.RoleUser { + t.Errorf("Expected first message role to be %s, got %s", context.RoleUser, result[0].Role) } } @@ -65,9 +89,9 @@ func TestGetMessages_FromQuery(t *testing.T) { gin.SetMode(gin.TestMode) - messages := []Message{ + messages := []context.Message{ { - Role: RoleUser, + Role: context.RoleUser, Content: "Test message", }, } @@ -83,7 +107,7 @@ func TestGetMessages_FromQuery(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - result, err := GetMessages(c, nil) + result, err := context.GetMessages(c, nil) if err != nil { t.Fatalf("Failed to get messages: %v", err) } @@ -100,7 +124,7 @@ func TestGetMessages_EmptyMessages(t *testing.T) { gin.SetMode(gin.TestMode) requestBody := map[string]interface{}{ - "messages": []Message{}, + "messages": []context.Message{}, "model": "gpt-4", } @@ -114,7 +138,7 @@ func TestGetMessages_EmptyMessages(t *testing.T) { completionReq, _ := parseCompletionRequestData(c) - _, err := GetMessages(c, completionReq) + _, err := context.GetMessages(c, completionReq) if err == nil { t.Error("Expected error for empty messages") } @@ -138,7 +162,7 @@ func TestGetChatID_FromQuery(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - chatID, err := GetChatID(c, cache, nil) + chatID, err := context.GetChatID(c, cache, nil) if err != nil { t.Fatalf("Failed to get chat ID: %v", err) } @@ -167,7 +191,7 @@ func TestGetChatID_FromHeader(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - chatID, err := GetChatID(c, cache, nil) + chatID, err := context.GetChatID(c, cache, nil) if err != nil { t.Fatalf("Failed to get chat ID: %v", err) } @@ -210,7 +234,7 @@ func TestGetChatID_FromMetadata(t *testing.T) { completionReq, _ := parseCompletionRequestData(c) - chatID, err := GetChatID(c, cache, completionReq) + chatID, err := context.GetChatID(c, cache, completionReq) if err != nil { t.Fatalf("Failed to get chat ID: %v", err) } @@ -233,9 +257,9 @@ func TestGetChatID_FromMessages(t *testing.T) { cache.Clear() // First request with one user message - messages1 := []Message{ + messages1 := []context.Message{ { - Role: RoleUser, + Role: context.RoleUser, Content: "First message", }, } @@ -255,7 +279,7 @@ func TestGetChatID_FromMessages(t *testing.T) { completionReq1, _ := parseCompletionRequestData(c) - chatID1, err := GetChatID(c, cache, completionReq1) + chatID1, err := context.GetChatID(c, cache, completionReq1) if err != nil { t.Fatalf("Failed to get chat ID: %v", err) } @@ -265,13 +289,13 @@ func TestGetChatID_FromMessages(t *testing.T) { } // Second request with two user messages (continuation) - messages2 := []Message{ + messages2 := []context.Message{ { - Role: RoleUser, + Role: context.RoleUser, Content: "First message", }, { - Role: RoleUser, + Role: context.RoleUser, Content: "Second message", }, } @@ -291,7 +315,7 @@ func TestGetChatID_FromMessages(t *testing.T) { completionReq2, _ := parseCompletionRequestData(c2) - chatID2, err := GetChatID(c2, cache, completionReq2) + chatID2, err := context.GetChatID(c2, cache, completionReq2) if err != nil { t.Fatalf("Failed to get chat ID second time: %v", err) } @@ -317,9 +341,9 @@ func TestGetChatID_Priority(t *testing.T) { headerChatID := "header-chat-id" metadataChatID := "metadata-chat-id" - messages := []Message{ + messages := []context.Message{ { - Role: RoleUser, + Role: context.RoleUser, Content: "This should not be used", }, } @@ -344,7 +368,7 @@ func TestGetChatID_Priority(t *testing.T) { completionReq, _ := parseCompletionRequestData(c) - chatID, err := GetChatID(c, cache, completionReq) + chatID, err := context.GetChatID(c, cache, completionReq) if err != nil { t.Fatalf("Failed to get chat ID: %v", err) } @@ -362,7 +386,7 @@ func TestGetLocale_FromQuery(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - locale := GetLocale(c, nil) + locale := context.GetLocale(c, nil) if locale != "zh-cn" { t.Errorf("Expected locale 'zh-cn', got '%s'", locale) } @@ -377,7 +401,7 @@ func TestGetLocale_FromHeader(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - locale := GetLocale(c, nil) + locale := context.GetLocale(c, nil) if locale != "en-us" { t.Errorf("Expected locale 'en-us', got '%s'", locale) } @@ -391,13 +415,13 @@ func TestGetLocale_FromMetadata(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - completionReq := &CompletionRequest{ + completionReq := &context.CompletionRequest{ Metadata: map[string]interface{}{ "locale": "ja-JP", }, } - locale := GetLocale(c, completionReq) + locale := context.GetLocale(c, completionReq) if locale != "ja-jp" { t.Errorf("Expected locale 'ja-jp' from metadata, got '%s'", locale) } @@ -412,13 +436,13 @@ func TestGetLocale_Priority(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - completionReq := &CompletionRequest{ + completionReq := &context.CompletionRequest{ Metadata: map[string]interface{}{ "locale": "de-DE", }, } - locale := GetLocale(c, completionReq) + locale := context.GetLocale(c, completionReq) if locale != "fr-fr" { t.Errorf("Expected query parameter to take priority, got '%s'", locale) } @@ -432,7 +456,7 @@ func TestGetTheme_FromQuery(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - theme := GetTheme(c, nil) + theme := context.GetTheme(c, nil) if theme != "dark" { t.Errorf("Expected theme 'dark', got '%s'", theme) } @@ -447,7 +471,7 @@ func TestGetTheme_FromHeader(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - theme := GetTheme(c, nil) + theme := context.GetTheme(c, nil) if theme != "light" { t.Errorf("Expected theme 'light', got '%s'", theme) } @@ -461,13 +485,13 @@ func TestGetTheme_FromMetadata(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - completionReq := &CompletionRequest{ + completionReq := &context.CompletionRequest{ Metadata: map[string]interface{}{ "theme": "auto", }, } - theme := GetTheme(c, completionReq) + theme := context.GetTheme(c, completionReq) if theme != "auto" { t.Errorf("Expected theme 'auto' from metadata, got '%s'", theme) } @@ -481,14 +505,14 @@ func TestGetReferer_FromMetadata(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - completionReq := &CompletionRequest{ + completionReq := &context.CompletionRequest{ Metadata: map[string]interface{}{ "referer": "tool", }, } - referer := GetReferer(c, completionReq) - if referer != RefererTool { + referer := context.GetReferer(c, completionReq) + if referer != context.RefererTool { t.Errorf("Expected referer 'tool' from metadata, got '%s'", referer) } } @@ -501,8 +525,8 @@ func TestGetAccept_FromQuery(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - accept := GetAccept(c, nil) - if accept != AcceptWebCUI { + accept := context.GetAccept(c, nil) + if accept != context.AcceptWebCUI { t.Errorf("Expected accept 'cui-web' from query, got '%s'", accept) } } @@ -516,8 +540,8 @@ func TestGetAccept_FromHeader(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - accept := GetAccept(c, nil) - if accept != AcceptDesktopCUI { + accept := context.GetAccept(c, nil) + if accept != context.AcceptDesktopCUI { t.Errorf("Expected accept 'cui-desktop' from header, got '%s'", accept) } } @@ -530,14 +554,14 @@ func TestGetAccept_FromMetadata(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - completionReq := &CompletionRequest{ + completionReq := &context.CompletionRequest{ Metadata: map[string]interface{}{ "accept": "cui-native", }, } - accept := GetAccept(c, completionReq) - if accept != AccepNativeCUI { + accept := context.GetAccept(c, completionReq) + if accept != context.AccepNativeCUI { t.Errorf("Expected accept 'cui-native' from metadata, got '%s'", accept) } } @@ -550,8 +574,8 @@ func TestGetAccept_Default(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - accept := GetAccept(c, nil) - if accept != AcceptStandard { + accept := context.GetAccept(c, nil) + if accept != context.AcceptStandard { t.Errorf("Expected default accept 'standard', got '%s'", accept) } } @@ -565,14 +589,14 @@ func TestGetAccept_Priority(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - completionReq := &CompletionRequest{ + completionReq := &context.CompletionRequest{ Metadata: map[string]interface{}{ "accept": "cui-native", }, } - accept := GetAccept(c, completionReq) - if accept != AcceptWebCUI { + accept := context.GetAccept(c, completionReq) + if accept != context.AcceptWebCUI { t.Errorf("Expected query parameter to take priority, got '%s'", accept) } } @@ -585,11 +609,11 @@ func TestGetAssistantID_FromModel(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - completionReq := &CompletionRequest{ + completionReq := &context.CompletionRequest{ Model: "gpt-4-turbo-yao_myassistant", } - assistantID, err := GetAssistantID(c, completionReq) + assistantID, err := context.GetAssistantID(c, completionReq) if err != nil { t.Fatalf("Failed to get assistant ID: %v", err) } @@ -608,11 +632,11 @@ func TestGetAssistantID_Priority(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - completionReq := &CompletionRequest{ + completionReq := &context.CompletionRequest{ Model: "gpt-4-yao_from_model", } - assistantID, err := GetAssistantID(c, completionReq) + assistantID, err := context.GetAssistantID(c, completionReq) if err != nil { t.Fatalf("Failed to get assistant ID: %v", err) } @@ -630,7 +654,7 @@ func TestGetRoute_FromQuery(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - route := GetRoute(c, nil) + route := context.GetRoute(c, nil) if route != "/dashboard/home" { t.Errorf("Expected route '/dashboard/home', got '%s'", route) } @@ -645,7 +669,7 @@ func TestGetRoute_FromHeader(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - route := GetRoute(c, nil) + route := context.GetRoute(c, nil) if route != "/settings/profile" { t.Errorf("Expected route '/settings/profile', got '%s'", route) } @@ -659,11 +683,11 @@ func TestGetRoute_FromPayload(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - completionReq := &CompletionRequest{ + completionReq := &context.CompletionRequest{ Route: "/admin/users", } - route := GetRoute(c, completionReq) + route := context.GetRoute(c, completionReq) if route != "/admin/users" { t.Errorf("Expected route '/admin/users' from payload, got '%s'", route) } @@ -678,11 +702,11 @@ func TestGetRoute_Priority(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - completionReq := &CompletionRequest{ + completionReq := &context.CompletionRequest{ Route: "/from/payload", } - route := GetRoute(c, completionReq) + route := context.GetRoute(c, completionReq) if route != "/from/query" { t.Errorf("Expected query parameter to take priority, got '%s'", route) } @@ -702,7 +726,7 @@ func TestGetMetadata_FromQuery(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - result := GetMetadata(c, nil) + result := context.GetMetadata(c, nil) if result == nil { t.Fatal("Expected data to be returned") } @@ -727,7 +751,7 @@ func TestGetMetadata_FromHeader_Base64(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - result := GetMetadata(c, nil) + result := context.GetMetadata(c, nil) if result == nil { t.Fatal("Expected data to be returned") } @@ -754,11 +778,11 @@ func TestGetMetadata_FromPayload(t *testing.T) { "limit": float64(10), } - completionReq := &CompletionRequest{ + completionReq := &context.CompletionRequest{ Metadata: data, } - result := GetMetadata(c, completionReq) + result := context.GetMetadata(c, completionReq) if result == nil { t.Fatal("Expected data to be returned") } @@ -792,11 +816,11 @@ func TestGetMetadata_Priority(t *testing.T) { "source": "payload", } - completionReq := &CompletionRequest{ + completionReq := &context.CompletionRequest{ Metadata: payloadData, } - result := GetMetadata(c, completionReq) + result := context.GetMetadata(c, completionReq) if result == nil { t.Fatal("Expected data to be returned") } @@ -814,7 +838,7 @@ func TestGetMetadata_EmptyData(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - result := GetMetadata(c, nil) + result := context.GetMetadata(c, nil) if result != nil { t.Errorf("Expected nil data, got '%v'", result) } @@ -831,9 +855,9 @@ func TestGetCompletionRequest_WriterInitialized(t *testing.T) { t.Fatalf("Failed to get cache: %v", err) } - messages := []Message{ + messages := []context.Message{ { - Role: RoleUser, + Role: context.RoleUser, Content: "Test message", }, } @@ -851,7 +875,7 @@ func TestGetCompletionRequest_WriterInitialized(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - completionReq, ctx, opts, err := GetCompletionRequest(c, cache) + completionReq, ctx, opts, err := context.GetCompletionRequest(c, cache) if err != nil { t.Fatalf("Failed to get completion request: %v", err) } @@ -899,9 +923,9 @@ func TestGetCompletionRequest_ChatIDFallback(t *testing.T) { } // Request without explicit chat_id should generate one - messages := []Message{ + messages := []context.Message{ { - Role: RoleUser, + Role: context.RoleUser, Content: "Test message", }, } @@ -919,7 +943,7 @@ func TestGetCompletionRequest_ChatIDFallback(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - _, ctx, opts, err := GetCompletionRequest(c, cache) + _, ctx, opts, err := context.GetCompletionRequest(c, cache) if err != nil { t.Fatalf("Failed to get completion request: %v", err) } @@ -949,14 +973,14 @@ func TestGetSkip_FromBody(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - completionReq := &CompletionRequest{ - Skip: &Skip{ + completionReq := &context.CompletionRequest{ + Skip: &context.Skip{ History: true, Trace: false, }, } - skip := GetSkip(c, completionReq) + skip := context.GetSkip(c, completionReq) if skip == nil { t.Fatal("Expected skip to be returned") } @@ -978,7 +1002,7 @@ func TestGetSkip_FromQueryParams(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - skip := GetSkip(c, nil) + skip := context.GetSkip(c, nil) if skip == nil { t.Fatal("Expected skip to be returned") } @@ -1000,7 +1024,7 @@ func TestGetSkip_FromQueryParams_ShortForm(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - skip := GetSkip(c, nil) + skip := context.GetSkip(c, nil) if skip == nil { t.Fatal("Expected skip to be returned") } @@ -1023,14 +1047,14 @@ func TestGetSkip_Priority(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - completionReq := &CompletionRequest{ - Skip: &Skip{ + completionReq := &context.CompletionRequest{ + Skip: &context.Skip{ History: true, Trace: true, }, } - skip := GetSkip(c, completionReq) + skip := context.GetSkip(c, completionReq) if skip == nil { t.Fatal("Expected skip to be returned") } @@ -1053,7 +1077,7 @@ func TestGetSkip_Nil(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - skip := GetSkip(c, nil) + skip := context.GetSkip(c, nil) if skip != nil { t.Errorf("Expected skip to be nil, got %v", skip) } @@ -1067,7 +1091,7 @@ func TestGetSkip_OnlyHistorySet(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - skip := GetSkip(c, nil) + skip := context.GetSkip(c, nil) if skip == nil { t.Fatal("Expected skip to be returned") } @@ -1088,9 +1112,9 @@ func TestGetSkip_FromBodyViaParseRequest(t *testing.T) { gin.SetMode(gin.TestMode) // Test parsing Skip from full request body - messages := []Message{ + messages := []context.Message{ { - Role: RoleUser, + Role: context.RoleUser, Content: "Generate a title for this chat", }, } @@ -1132,7 +1156,7 @@ func TestGetSkip_FromBodyViaParseRequest(t *testing.T) { } // Now test GetSkip function with the parsed request - skip := GetSkip(c, completionReq) + skip := context.GetSkip(c, completionReq) if skip == nil { t.Fatal("Expected GetSkip to return skip configuration") } @@ -1157,7 +1181,7 @@ func TestGetMode_FromQuery(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - mode := GetMode(c, nil) + mode := context.GetMode(c, nil) if mode != "task" { t.Errorf("Expected mode 'task' from query, got '%s'", mode) } @@ -1175,7 +1199,7 @@ func TestGetMode_FromHeader(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - mode := GetMode(c, nil) + mode := context.GetMode(c, nil) if mode != "chat" { t.Errorf("Expected mode 'chat' from header, got '%s'", mode) } @@ -1192,13 +1216,13 @@ func TestGetMode_FromMetadata(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - completionReq := &CompletionRequest{ + completionReq := &context.CompletionRequest{ Metadata: map[string]interface{}{ "mode": "task", }, } - mode := GetMode(c, completionReq) + mode := context.GetMode(c, completionReq) if mode != "task" { t.Errorf("Expected mode 'task' from metadata, got '%s'", mode) } @@ -1217,13 +1241,13 @@ func TestGetMode_Priority(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - completionReq := &CompletionRequest{ + completionReq := &context.CompletionRequest{ Metadata: map[string]interface{}{ "mode": "metadata_mode", }, } - mode := GetMode(c, completionReq) + mode := context.GetMode(c, completionReq) if mode != "query_mode" { t.Errorf("Expected mode 'query_mode' (query has priority), got '%s'", mode) } @@ -1240,7 +1264,7 @@ func TestGetMode_Empty(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - mode := GetMode(c, nil) + mode := context.GetMode(c, nil) if mode != "" { t.Errorf("Expected empty mode, got '%s'", mode) } diff --git a/agent/context/stack_test.go b/agent/context/stack_test.go index b678332c..1226a667 100644 --- a/agent/context/stack_test.go +++ b/agent/context/stack_test.go @@ -1,10 +1,11 @@ -package context +package context_test import ( + stdContext "context" "testing" "time" - "github.com/yaoapp/yao/agent/output/message" + "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/config" "github.com/yaoapp/yao/test" ) @@ -15,10 +16,10 @@ func TestNewStack(t *testing.T) { traceID := "12345678" assistantID := "test-assistant" - referer := RefererAPI - opts := &Options{} + referer := context.RefererAPI + opts := &context.Options{} - stack := NewStack(traceID, assistantID, referer, opts) + stack := context.NewStack(traceID, assistantID, referer, opts) if stack == nil { t.Fatal("Expected stack to be created, got nil") @@ -48,8 +49,8 @@ func TestNewStack(t *testing.T) { t.Error("Expected stack to be root") } - if stack.Status != StackStatusRunning { - t.Errorf("Expected Status '%s', got '%s'", StackStatusRunning, stack.Status) + if stack.Status != context.StackStatusRunning { + t.Errorf("Expected Status '%s', got '%s'", context.StackStatusRunning, stack.Status) } } @@ -58,7 +59,7 @@ func TestNewStack_GenerateTraceID(t *testing.T) { defer test.Clean() // Empty traceID should generate a UUID - stack := NewStack("", "test-assistant", RefererAPI, &Options{}) + stack := context.NewStack("", "test-assistant", context.RefererAPI, &context.Options{}) if stack.TraceID == "" { t.Error("Expected TraceID to be generated, got empty string") @@ -75,10 +76,10 @@ func TestNewChildStack(t *testing.T) { defer test.Clean() // Create parent stack - parentStack := NewStack("12345678", "parent-assistant", RefererAPI, &Options{}) + parentStack := context.NewStack("12345678", "parent-assistant", context.RefererAPI, &context.Options{}) // Create child stack - childStack := parentStack.NewChildStack("child-assistant", RefererAgent, &Options{}) + childStack := parentStack.NewChildStack("child-assistant", context.RefererAgent, &context.Options{}) if childStack == nil { t.Fatal("Expected child stack to be created, got nil") @@ -122,15 +123,15 @@ func TestStackComplete(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - stack := NewStack("12345678", "test-assistant", RefererAPI, &Options{}) + stack := context.NewStack("12345678", "test-assistant", context.RefererAPI, &context.Options{}) // Wait a bit to have measurable duration time.Sleep(10 * time.Millisecond) stack.Complete() - if stack.Status != StackStatusCompleted { - t.Errorf("Expected Status '%s', got '%s'", StackStatusCompleted, stack.Status) + if stack.Status != context.StackStatusCompleted { + t.Errorf("Expected Status '%s', got '%s'", context.StackStatusCompleted, stack.Status) } if stack.CompletedAt == nil { @@ -158,14 +159,14 @@ func TestStackFail(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - stack := NewStack("12345678", "test-assistant", RefererAPI, &Options{}) + stack := context.NewStack("12345678", "test-assistant", context.RefererAPI, &context.Options{}) testError := "test error message" stack.Fail(nil) stack.Error = testError - if stack.Status != StackStatusFailed { - t.Errorf("Expected Status '%s', got '%s'", StackStatusFailed, stack.Status) + if stack.Status != context.StackStatusFailed { + t.Errorf("Expected Status '%s', got '%s'", context.StackStatusFailed, stack.Status) } if stack.Error != testError { @@ -181,12 +182,12 @@ func TestStackTimeout(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - stack := NewStack("12345678", "test-assistant", RefererAPI, &Options{}) + stack := context.NewStack("12345678", "test-assistant", context.RefererAPI, &context.Options{}) stack.Timeout() - if stack.Status != StackStatusTimeout { - t.Errorf("Expected Status '%s', got '%s'", StackStatusTimeout, stack.Status) + if stack.Status != context.StackStatusTimeout { + t.Errorf("Expected Status '%s', got '%s'", context.StackStatusTimeout, stack.Status) } if !stack.IsCompleted() { @@ -198,12 +199,10 @@ func TestEnterStack_RootCreation(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - ctx := &Context{ - IDGenerator: message.NewIDGenerator(), - Referer: RefererAPI, - } + ctx := context.New(stdContext.Background(), nil, "test-chat-id") + ctx.Referer = context.RefererAPI - stack, traceID, done := EnterStack(ctx, "test-assistant", &Options{}) + stack, traceID, done := context.EnterStack(ctx, "test-assistant", &context.Options{}) defer done() if stack == nil { @@ -244,13 +243,11 @@ func TestEnterStack_ChildCreation(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - ctx := &Context{ - IDGenerator: message.NewIDGenerator(), - Referer: RefererAPI, - } + ctx := context.New(stdContext.Background(), nil, "test-chat-id") + ctx.Referer = context.RefererAPI // Create parent - parentStack, parentTraceID, parentDone := EnterStack(ctx, "parent-assistant", &Options{}) + parentStack, parentTraceID, parentDone := context.EnterStack(ctx, "parent-assistant", &context.Options{}) defer parentDone() if parentStack == nil { @@ -258,7 +255,7 @@ func TestEnterStack_ChildCreation(t *testing.T) { } // Create child - childStack, childTraceID, childDone := EnterStack(ctx, "child-assistant", &Options{}) + childStack, childTraceID, childDone := context.EnterStack(ctx, "child-assistant", &context.Options{}) defer childDone() if childStack == nil { @@ -290,16 +287,14 @@ func TestEnterStack_DoneCallback(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - ctx := &Context{ - IDGenerator: message.NewIDGenerator(), - Referer: RefererAPI, - } + ctx := context.New(stdContext.Background(), nil, "test-chat-id") + ctx.Referer = context.RefererAPI // Create parent - parentStack, _, parentDone := EnterStack(ctx, "parent-assistant", &Options{}) + parentStack, _, parentDone := context.EnterStack(ctx, "parent-assistant", &context.Options{}) // Create child - childStack, _, childDone := EnterStack(ctx, "child-assistant", &Options{}) + childStack, _, childDone := context.EnterStack(ctx, "child-assistant", &context.Options{}) // Child should be current if ctx.Stack != childStack { @@ -332,19 +327,17 @@ func TestContextGetAllStacks(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - ctx := &Context{ - IDGenerator: message.NewIDGenerator(), - Referer: RefererAPI, - } + ctx := context.New(stdContext.Background(), nil, "test-chat-id") + ctx.Referer = context.RefererAPI // Create multiple stacks - _, _, done1 := EnterStack(ctx, "assistant1", &Options{}) + _, _, done1 := context.EnterStack(ctx, "assistant1", &context.Options{}) defer done1() - _, _, done2 := EnterStack(ctx, "assistant2", &Options{}) + _, _, done2 := context.EnterStack(ctx, "assistant2", &context.Options{}) defer done2() - _, _, done3 := EnterStack(ctx, "assistant3", &Options{}) + _, _, done3 := context.EnterStack(ctx, "assistant3", &context.Options{}) defer done3() // Get all stacks @@ -359,12 +352,10 @@ func TestContextGetStackByID(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - ctx := &Context{ - IDGenerator: message.NewIDGenerator(), - Referer: RefererAPI, - } + ctx := context.New(stdContext.Background(), nil, "test-chat-id") + ctx.Referer = context.RefererAPI - stack, _, done := EnterStack(ctx, "test-assistant", &Options{}) + stack, _, done := context.EnterStack(ctx, "test-assistant", &context.Options{}) defer done() // Get stack by ID @@ -389,16 +380,14 @@ func TestContextGetStacksByTraceID(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - ctx := &Context{ - IDGenerator: message.NewIDGenerator(), - Referer: RefererAPI, - } + ctx := context.New(stdContext.Background(), nil, "test-chat-id") + ctx.Referer = context.RefererAPI // Create parent and child (same trace ID) - _, traceID, done1 := EnterStack(ctx, "parent-assistant", &Options{}) + _, traceID, done1 := context.EnterStack(ctx, "parent-assistant", &context.Options{}) defer done1() - _, _, done2 := EnterStack(ctx, "child-assistant", &Options{}) + _, _, done2 := context.EnterStack(ctx, "child-assistant", &context.Options{}) defer done2() // Get stacks by trace ID @@ -420,17 +409,15 @@ func TestContextGetRootStack(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - ctx := &Context{ - IDGenerator: message.NewIDGenerator(), - Referer: RefererAPI, - } + ctx := context.New(stdContext.Background(), nil, "test-chat-id") + ctx.Referer = context.RefererAPI // Create parent - parentStack, _, done1 := EnterStack(ctx, "parent-assistant", &Options{}) + parentStack, _, done1 := context.EnterStack(ctx, "parent-assistant", &context.Options{}) defer done1() // Create child - _, _, done2 := EnterStack(ctx, "child-assistant", &Options{}) + _, _, done2 := context.EnterStack(ctx, "child-assistant", &context.Options{}) defer done2() // Get root stack @@ -453,7 +440,7 @@ func TestStackClone(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() - original := NewStack("12345678", "test-assistant", RefererAPI, &Options{}) + original := context.NewStack("12345678", "test-assistant", context.RefererAPI, &context.Options{}) original.Complete() clone := original.Clone() diff --git a/agent/context/types.go b/agent/context/types.go index 51b7070b..e136b365 100644 --- a/agent/context/types.go +++ b/agent/context/types.go @@ -230,6 +230,7 @@ type Context struct { Stacks map[string]*Stack `json:"-"` // Stacks, all stacks in this request (for trace logging) 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) + Logger *RequestLogger `json:"-"` // Request-scoped async logger // Chat buffer for batch saving messages and resume steps Buffer *ChatBuffer `json:"-"` // Chat buffer for batch saving at end of Stream() diff --git a/agent/llm/providers/openai/claude_test.go b/agent/llm/providers/openai/claude_test.go index 15fc39c1..cc8884ab 100644 --- a/agent/llm/providers/openai/claude_test.go +++ b/agent/llm/providers/openai/claude_test.go @@ -6,7 +6,6 @@ import ( "github.com/yaoapp/gou/connector" "github.com/yaoapp/gou/connector/openai" - "github.com/yaoapp/gou/plan" "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/llm" "github.com/yaoapp/yao/agent/output/message" @@ -17,37 +16,35 @@ import ( // newClaudeTestContext creates a real Context for testing Claude provider func newClaudeTestContext(chatID, connectorID string) *context.Context { - return &context.Context{ - Context: gocontext.Background(), - Space: plan.NewMemorySharedSpace(), - ChatID: chatID, - AssistantID: "test-assistant", - Locale: "en-us", - Theme: "light", - Client: context.Client{ - Type: "web", - UserAgent: "ClaudeProviderTest/1.0", - IP: "127.0.0.1", - }, - Referer: context.RefererAPI, - Accept: context.AcceptStandard, - Route: "/api/test", - Metadata: make(map[string]interface{}), - Authorized: &types.AuthorizedInfo{ - Subject: "test-user", - ClientID: "test-client", - UserID: "test-user-123", - TeamID: "test-team-456", - TenantID: "test-tenant-789", - SessionID: "test-session-id", - Constraints: types.DataConstraints{ - TeamOnly: true, - Extra: map[string]interface{}{ - "test": "claude-provider", - }, + authorized := &types.AuthorizedInfo{ + Subject: "test-user", + ClientID: "test-client", + UserID: "test-user-123", + TeamID: "test-team-456", + TenantID: "test-tenant-789", + SessionID: "test-session-id", + Constraints: types.DataConstraints{ + TeamOnly: true, + Extra: map[string]interface{}{ + "test": "claude-provider", }, }, } + + ctx := context.New(gocontext.Background(), authorized, chatID) + ctx.AssistantID = "test-assistant" + ctx.Locale = "en-us" + ctx.Theme = "light" + ctx.Client = context.Client{ + Type: "web", + UserAgent: "ClaudeProviderTest/1.0", + IP: "127.0.0.1", + } + ctx.Referer = context.RefererAPI + ctx.Accept = context.AcceptStandard + ctx.Route = "/api/test" + ctx.Metadata = make(map[string]interface{}) + return ctx } // TestClaudeSonnet4StreamBasic tests basic streaming completion with Claude Sonnet 4 diff --git a/agent/llm/providers/openai/deepseek_r1_test.go b/agent/llm/providers/openai/deepseek_r1_test.go index bf790e3a..099e59cf 100644 --- a/agent/llm/providers/openai/deepseek_r1_test.go +++ b/agent/llm/providers/openai/deepseek_r1_test.go @@ -8,7 +8,6 @@ import ( jsoniter "github.com/json-iterator/go" "github.com/yaoapp/gou/connector" "github.com/yaoapp/gou/connector/openai" - "github.com/yaoapp/gou/plan" "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/llm" "github.com/yaoapp/yao/agent/output/message" @@ -396,35 +395,33 @@ func TestDeepSeekR1LogicPuzzle(t *testing.T) { // newDeepSeekTestContext creates a real Context for testing DeepSeek provider func newDeepSeekTestContext(chatID, connectorID string) *context.Context { - return &context.Context{ - Context: gocontext.Background(), - Space: plan.NewMemorySharedSpace(), - ChatID: chatID, - AssistantID: "test-assistant", - Locale: "en-us", - Theme: "light", - Client: context.Client{ - Type: "web", - UserAgent: "DeepSeekProviderTest/1.0", - IP: "127.0.0.1", - }, - Referer: context.RefererAPI, - Accept: context.AcceptStandard, - Route: "/api/test", - Metadata: make(map[string]interface{}), - Authorized: &types.AuthorizedInfo{ - Subject: "test-user", - ClientID: "test-client", - UserID: "test-user-123", - TeamID: "test-team-456", - TenantID: "test-tenant-789", - SessionID: "test-session-id", - Constraints: types.DataConstraints{ - TeamOnly: true, - Extra: map[string]interface{}{ - "test": "deepseek-provider", - }, + authorized := &types.AuthorizedInfo{ + Subject: "test-user", + ClientID: "test-client", + UserID: "test-user-123", + TeamID: "test-team-456", + TenantID: "test-tenant-789", + SessionID: "test-session-id", + Constraints: types.DataConstraints{ + TeamOnly: true, + Extra: map[string]interface{}{ + "test": "deepseek-provider", }, }, } + + ctx := context.New(gocontext.Background(), authorized, chatID) + ctx.AssistantID = "test-assistant" + ctx.Locale = "en-us" + ctx.Theme = "light" + ctx.Client = context.Client{ + Type: "web", + UserAgent: "DeepSeekProviderTest/1.0", + IP: "127.0.0.1", + } + ctx.Referer = context.RefererAPI + ctx.Accept = context.AcceptStandard + ctx.Route = "/api/test" + ctx.Metadata = make(map[string]interface{}) + return ctx } diff --git a/agent/llm/providers/openai/deepseek_v3_test.go b/agent/llm/providers/openai/deepseek_v3_test.go index 3f944735..b16f58dd 100644 --- a/agent/llm/providers/openai/deepseek_v3_test.go +++ b/agent/llm/providers/openai/deepseek_v3_test.go @@ -6,7 +6,6 @@ import ( "github.com/yaoapp/gou/connector" "github.com/yaoapp/gou/connector/openai" - "github.com/yaoapp/gou/plan" "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/llm" "github.com/yaoapp/yao/agent/output/message" @@ -368,35 +367,33 @@ func TestDeepSeekV3NoReasoningEffort(t *testing.T) { // newDeepSeekV3TestContext creates a real Context for testing DeepSeek V3 provider func newDeepSeekV3TestContext(chatID, connectorID string) *context.Context { - return &context.Context{ - Context: gocontext.Background(), - Space: plan.NewMemorySharedSpace(), - ChatID: chatID, - AssistantID: "test-assistant", - Locale: "en-us", - Theme: "light", - Client: context.Client{ - Type: "web", - UserAgent: "DeepSeekV3ProviderTest/1.0", - IP: "127.0.0.1", - }, - Referer: context.RefererAPI, - Accept: context.AcceptStandard, - Route: "/api/test", - Metadata: make(map[string]interface{}), - Authorized: &types.AuthorizedInfo{ - Subject: "test-user", - ClientID: "test-client", - UserID: "test-user-123", - TeamID: "test-team-456", - TenantID: "test-tenant-789", - SessionID: "test-session-id", - Constraints: types.DataConstraints{ - TeamOnly: true, - Extra: map[string]interface{}{ - "test": "deepseek-v3-provider", - }, + authorized := &types.AuthorizedInfo{ + Subject: "test-user", + ClientID: "test-client", + UserID: "test-user-123", + TeamID: "test-team-456", + TenantID: "test-tenant-789", + SessionID: "test-session-id", + Constraints: types.DataConstraints{ + TeamOnly: true, + Extra: map[string]interface{}{ + "test": "deepseek-v3-provider", }, }, } + + ctx := context.New(gocontext.Background(), authorized, chatID) + ctx.AssistantID = "test-assistant" + ctx.Locale = "en-us" + ctx.Theme = "light" + ctx.Client = context.Client{ + Type: "web", + UserAgent: "DeepSeekV3ProviderTest/1.0", + IP: "127.0.0.1", + } + ctx.Referer = context.RefererAPI + ctx.Accept = context.AcceptStandard + ctx.Route = "/api/test" + ctx.Metadata = make(map[string]interface{}) + return ctx } diff --git a/agent/llm/providers/openai/gpt5_test.go b/agent/llm/providers/openai/gpt5_test.go index a04ea943..f0d4f9cb 100644 --- a/agent/llm/providers/openai/gpt5_test.go +++ b/agent/llm/providers/openai/gpt5_test.go @@ -6,7 +6,6 @@ import ( "github.com/yaoapp/gou/connector" "github.com/yaoapp/gou/connector/openai" - "github.com/yaoapp/gou/plan" "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/llm" "github.com/yaoapp/yao/agent/output/message" @@ -383,35 +382,33 @@ func TestGPT5ReasoningEffortWithGPT4o(t *testing.T) { // newGPT5TestContext creates a real Context for testing GPT-5 provider func newGPT5TestContext(chatID, connectorID string) *context.Context { - return &context.Context{ - Context: gocontext.Background(), - Space: plan.NewMemorySharedSpace(), - ChatID: chatID, - AssistantID: "test-assistant", - Locale: "en-us", - Theme: "light", - Client: context.Client{ - Type: "web", - UserAgent: "GPT5ProviderTest/1.0", - IP: "127.0.0.1", - }, - Referer: context.RefererAPI, - Accept: context.AcceptStandard, - Route: "/api/test", - Metadata: make(map[string]interface{}), - Authorized: &types.AuthorizedInfo{ - Subject: "test-user", - ClientID: "test-client", - UserID: "test-user-123", - TeamID: "test-team-456", - TenantID: "test-tenant-789", - SessionID: "test-session-id", - Constraints: types.DataConstraints{ - TeamOnly: true, - Extra: map[string]interface{}{ - "test": "gpt5-provider", - }, + authorized := &types.AuthorizedInfo{ + Subject: "test-user", + ClientID: "test-client", + UserID: "test-user-123", + TeamID: "test-team-456", + TenantID: "test-tenant-789", + SessionID: "test-session-id", + Constraints: types.DataConstraints{ + TeamOnly: true, + Extra: map[string]interface{}{ + "test": "gpt5-provider", }, }, } + + ctx := context.New(gocontext.Background(), authorized, chatID) + ctx.AssistantID = "test-assistant" + ctx.Locale = "en-us" + ctx.Theme = "light" + ctx.Client = context.Client{ + Type: "web", + UserAgent: "GPT5ProviderTest/1.0", + IP: "127.0.0.1", + } + ctx.Referer = context.RefererAPI + ctx.Accept = context.AcceptStandard + ctx.Route = "/api/test" + ctx.Metadata = make(map[string]interface{}) + return ctx } diff --git a/agent/llm/providers/openai/openai_test.go b/agent/llm/providers/openai/openai_test.go index 111e7e3e..c17d2dd8 100644 --- a/agent/llm/providers/openai/openai_test.go +++ b/agent/llm/providers/openai/openai_test.go @@ -9,7 +9,6 @@ import ( "github.com/yaoapp/gou/connector" "github.com/yaoapp/gou/connector/openai" - "github.com/yaoapp/gou/plan" "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/llm" "github.com/yaoapp/yao/agent/output/message" @@ -1503,35 +1502,33 @@ func TestOpenAIStreamWithTemperature(t *testing.T) { // newTestContext creates a real Context for testing OpenAI provider func newTestContext(chatID, connectorID string) *context.Context { - return &context.Context{ - Context: gocontext.Background(), - Space: plan.NewMemorySharedSpace(), - ChatID: chatID, - AssistantID: "test-assistant", - Locale: "en-us", - Theme: "light", - Client: context.Client{ - Type: "web", - UserAgent: "OpenAIProviderTest/1.0", - IP: "127.0.0.1", - }, - Referer: context.RefererAPI, - Accept: context.AcceptStandard, - Route: "/api/test", - Metadata: make(map[string]interface{}), - Authorized: &types.AuthorizedInfo{ - Subject: "test-user", - ClientID: "test-client", - UserID: "test-user-123", - TeamID: "test-team-456", - TenantID: "test-tenant-789", - SessionID: "test-session-id", - Constraints: types.DataConstraints{ - TeamOnly: true, - Extra: map[string]interface{}{ - "test": "openai-provider", - }, + authorized := &types.AuthorizedInfo{ + Subject: "test-user", + ClientID: "test-client", + UserID: "test-user-123", + TeamID: "test-team-456", + TenantID: "test-tenant-789", + SessionID: "test-session-id", + Constraints: types.DataConstraints{ + TeamOnly: true, + Extra: map[string]interface{}{ + "test": "openai-provider", }, }, } + + ctx := context.New(gocontext.Background(), authorized, chatID) + ctx.AssistantID = "test-assistant" + ctx.Locale = "en-us" + ctx.Theme = "light" + ctx.Client = context.Client{ + Type: "web", + UserAgent: "OpenAIProviderTest/1.0", + IP: "127.0.0.1", + } + ctx.Referer = context.RefererAPI + ctx.Accept = context.AcceptStandard + ctx.Route = "/api/test" + ctx.Metadata = make(map[string]interface{}) + return ctx } diff --git a/agent/llm/providers/openai/temperature_test.go b/agent/llm/providers/openai/temperature_test.go index b0a8bd98..bd0d9e39 100644 --- a/agent/llm/providers/openai/temperature_test.go +++ b/agent/llm/providers/openai/temperature_test.go @@ -6,7 +6,6 @@ import ( "github.com/yaoapp/gou/connector" "github.com/yaoapp/gou/connector/openai" - "github.com/yaoapp/gou/plan" "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/llm" "github.com/yaoapp/yao/config" @@ -335,35 +334,33 @@ func TestTemperatureNoTemperatureProvided(t *testing.T) { // newTemperatureTestContext creates a real Context for testing temperature handling func newTemperatureTestContext(chatID, connectorID string) *context.Context { - return &context.Context{ - Context: gocontext.Background(), - Space: plan.NewMemorySharedSpace(), - ChatID: chatID, - AssistantID: "test-assistant", - Locale: "en-us", - Theme: "light", - Client: context.Client{ - Type: "web", - UserAgent: "TemperatureTest/1.0", - IP: "127.0.0.1", - }, - Referer: context.RefererAPI, - Accept: context.AcceptStandard, - Route: "/api/test", - Metadata: make(map[string]interface{}), - Authorized: &types.AuthorizedInfo{ - Subject: "test-user", - ClientID: "test-client", - UserID: "test-user-123", - TeamID: "test-team-456", - TenantID: "test-tenant-789", - SessionID: "test-session-id", - Constraints: types.DataConstraints{ - TeamOnly: true, - Extra: map[string]interface{}{ - "test": "temperature", - }, + authorized := &types.AuthorizedInfo{ + Subject: "test-user", + ClientID: "test-client", + UserID: "test-user-123", + TeamID: "test-team-456", + TenantID: "test-tenant-789", + SessionID: "test-session-id", + Constraints: types.DataConstraints{ + TeamOnly: true, + Extra: map[string]interface{}{ + "test": "temperature", }, }, } + + ctx := context.New(gocontext.Background(), authorized, chatID) + ctx.AssistantID = "test-assistant" + ctx.Locale = "en-us" + ctx.Theme = "light" + ctx.Client = context.Client{ + Type: "web", + UserAgent: "TemperatureTest/1.0", + IP: "127.0.0.1", + } + ctx.Referer = context.RefererAPI + ctx.Accept = context.AcceptStandard + ctx.Route = "/api/test" + ctx.Metadata = make(map[string]interface{}) + return ctx } diff --git a/agent/load.go b/agent/load.go index 63738a52..dafac72e 100644 --- a/agent/load.go +++ b/agent/load.go @@ -195,6 +195,9 @@ func initAssistant() error { // Set Storage assistant.SetStorage(agentDSL.Store) + // Set Store Setting (MaxSize, TTL, etc.) + assistant.SetStoreSetting(&agentDSL.StoreSetting) + // Set global Uses configuration if agentDSL.Uses != nil { globalUses := &context.Uses{ diff --git a/config/config.go b/config/config.go index 26effa96..d7239494 100644 --- a/config/config.go +++ b/config/config.go @@ -226,3 +226,8 @@ func CloseLog() { } } } + +// IsDevelopment returns true if the current mode is development +func IsDevelopment() bool { + return Conf.Mode == "development" +}