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

Enhance assistant initialization and context management

Browse source 
- 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.
This commit is contained in:
Max 2025-12-11 15:17:17 +08:00
parent 77fbfc1651
commit 4a1c0ec100
39 changed files with 2844 additions and 1532 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

59
agent/assistant/agent.go
View file

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

104
agent/assistant/agent_interrupt_test.go
View file

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

39
agent/assistant/agent_next_test.go
View file

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

189
agent/assistant/build_prompts_test.go
View file

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

51
agent/assistant/build_test.go
View file

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

55
agent/assistant/chat.go
View file

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

66
agent/assistant/chat_test.go
View file

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

285
agent/assistant/history.go Normal file
View file

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

791
agent/assistant/history_test.go Normal file
View file

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

53
agent/assistant/hook/create_bench_test.go
View file

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

53
agent/assistant/hook/create_mem_test.go
View file

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

69
agent/assistant/hook/create_test.go
View file

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

53
agent/assistant/hook/goroutine_leak_test.go
View file

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

69
agent/assistant/hook/next_test.go
View file

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

49
agent/assistant/hook/realworld_next_test.go
View file

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

69
agent/assistant/hook/realworld_stress_test.go
View file

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

33
agent/assistant/llm.go
View file

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

11
agent/assistant/load.go
View file

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

47
agent/assistant/load_store_test.go
View file

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

97
agent/assistant/mcp.go
View file

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

23
agent/context/authorized_test.go
View file

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

156
agent/context/chat_test.go
View file

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

71
agent/context/context.go
View file

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

137
agent/context/context_test.go
View file

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

213
agent/context/interrupt_test.go
View file

@ -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(),
}

285
agent/context/jsapi_output_test.go
View file

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

495
agent/context/log.go Normal file
View file

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

96
agent/context/message_test.go
View file

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

210
agent/context/openapi_test.go
View file

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

109
agent/context/stack_test.go
View file

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

1
agent/context/types.go
View file

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

55
agent/llm/providers/openai/claude_test.go
View file

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

55
agent/llm/providers/openai/deepseek_r1_test.go
View file

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

55
agent/llm/providers/openai/deepseek_v3_test.go
View file

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

55
agent/llm/providers/openai/gpt5_test.go
View file

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

55
agent/llm/providers/openai/openai_test.go
View file

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

55
agent/llm/providers/openai/temperature_test.go
View file

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

3
agent/load.go
View file

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

5
config/config.go
View file

@ -226,3 +226,8 @@ func CloseLog() {
}
}
}
// IsDevelopment returns true if the current mode is development
func IsDevelopment() bool {
return Conf.Mode == "development"
}