Enhance assistant initialization and context management

- Added a new method to set store settings during assistant initialization, allowing for configuration of storage parameters such as MaxSize and TTL.
- Updated context creation methods to streamline the setup process, ensuring that essential fields are populated consistently across various test contexts.
- Revised tests to validate the new initialization behavior and context management, ensuring proper handling of assistant settings and context properties.

agent/assistant/agent.go

@@ -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)
+	// Update logger with assistant ID and start logging
-	defer log.Trace("[AGENT] Stream ended: assistant=%s, contextID=%s", ast.ID, ctx.ID)
+	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
+		// Log end of request
-	ast.BufferUserInput(ctx, inputMessages)
+		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
+		ctx.Logger.Phase("LLM")
-		completionMessages, completionOptions, err = ast.BuildRequest(ctx, inputMessages, createResponse)
+
+		// 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
 	}

agent/assistant/agent_interrupt_test.go

@@ -6,7 +6,6 @@ import (
 	"testing"
 	"time"
 
-	"github.com/yaoapp/gou/plan"
 	"github.com/yaoapp/yao/agent/assistant"
 	"github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/agent/output/message"
@@ -15,35 +14,36 @@ import (
 )
 
 // newTestContextWithInterrupt creates a Context with interrupt controller for testing
-func newTestContextWithInterrupt(chatID, assistantID string) *context.Context {
+// Returns the context and a cancel function that should be called before Release()
-	ctx := &context.Context{
+func newTestContextWithInterrupt(chatID, assistantID string) (*context.Context, stdContext.CancelFunc) {
-		Context:     stdContext.Background(),
+	authorized := &types.AuthorizedInfo{
-		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",
-		},
+	}
+
+	// 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")
+		ctx, cancel := newTestContextWithInterrupt("chat-interrupt-graceful", "tests.interrupt")
-		defer ctx.Release()
+		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")
+		ctx, cancel := newTestContextWithInterrupt("chat-interrupt-force", "tests.interrupt")
-		defer ctx.Release()
+		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")
+		ctx, cancel := newTestContextWithInterrupt("chat-interrupt-multiple", "tests.interrupt")
-		defer ctx.Release()
+		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")
+		ctx, cancel := newTestContextWithInterrupt("chat-interrupt-before", "test-assistant")
-		defer ctx.Release()
+		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

agent/assistant/agent_next_test.go

@@ -5,7 +5,6 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
-	"github.com/yaoapp/gou/plan"
 	"github.com/yaoapp/yao/agent/assistant"
 	"github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/agent/output/message"
@@ -15,27 +14,25 @@ import (
 
 // newAgentNextTestContext creates a test context
 func newAgentNextTestContext(chatID, assistantID string) *context.Context {
-	return &context.Context{
+	authorized := &types.AuthorizedInfo{
-		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",
-		},
 	}
+
+	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)

agent/assistant/build_prompts_test.go

@@ -26,29 +26,34 @@ func containsString(content interface{}, substr string) bool {
 
 // newPromptTestContext creates a context suitable for prompt testing with Create Hook
 func newPromptTestContext(chatID, assistantID string) *context.Context {
-	return &context.Context{
+	authorized := &types.AuthorizedInfo{
-		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",
-		},
 	}
+
+	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{
+		ctx := newMinimalTestContext()
-			Locale: "zh-cn",
+		ctx.Locale = "zh-cn"
-			Authorized: &types.AuthorizedInfo{
+		ctx.Authorized = &types.AuthorizedInfo{
 			UserID: "test-user-123",
 			TeamID: "test-team-456",
-			},
+		}
-			Metadata: map[string]interface{}{
+		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{
+		ctx := newMinimalTestContext()
-			Locale: "en-us",
+		ctx.Locale = "en-us"
-		}
 
 		messages := []context.Message{
 			{Role: context.RoleUser, Content: "Hello"},
@@ -128,8 +131,8 @@ func TestBuildSystemPromptsIntegration(t *testing.T) {
 		ast, err := assistant.Get("yaobots")
 		require.NoError(t, err)
 
-		ctx := &context.Context{
+		ctx := newMinimalTestContext()
-			Metadata: map[string]interface{}{
+		ctx.Metadata = map[string]interface{}{
 			"STRING_VAL": "hello",
 			"INT_VAL":    123,
 			"INT64_VAL":  int64(456),
@@ -141,7 +144,6 @@ func TestBuildSystemPromptsIntegration(t *testing.T) {
 			"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{
+		ctx := newMinimalTestContext()
-			Authorized: &types.AuthorizedInfo{
+		ctx.Authorized = &types.AuthorizedInfo{
 			UserID:   "user-123",
 			Subject:  "user@example.com", // PII - should not be exposed
 			TeamID:   "team-456",
 			TenantID: "tenant-789",
-			},
+		}
-			Client: context.Client{
+		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{
+		ctx := newMinimalTestContext()
-			Authorized: &types.AuthorizedInfo{
+		ctx.Authorized = &types.AuthorizedInfo{
 			UserID: "user-abc",
 			TeamID: "team-xyz",
-			},
+		}
-			Metadata: map[string]interface{}{
+		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{
+		ctx := newMinimalTestContext()
-			Authorized: &types.AuthorizedInfo{
+		ctx.Authorized = &types.AuthorizedInfo{
 			UserID: "all-vars-user",
-			},
+		}
-			Metadata: map[string]interface{}{
+		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{
+		ctx := newMinimalTestContext()
-			Metadata: map[string]interface{}{
+		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{
+		ctx := newMinimalTestContext()
-			Metadata: map[string]interface{}{
+		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{
+		ctx := newMinimalTestContext()
-			Metadata: map[string]interface{}{
+		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{
+		ctx := newMinimalTestContext()
-			Metadata: map[string]interface{}{
+		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"},

agent/assistant/build_test.go

@@ -4,7 +4,6 @@ import (
 	stdContext "context"
 	"testing"
 
-	"github.com/yaoapp/gou/plan"
 	"github.com/yaoapp/yao/agent/assistant"
 	"github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/agent/testutils"
@@ -13,33 +12,31 @@ import (
 
 // newTestContext creates a Context for testing with commonly used fields pre-populated
 func newTestContext(chatID, assistantID string) *context.Context {
-	return &context.Context{
+	authorized := &types.AuthorizedInfo{
-		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",
-		},
 	}
+
+	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

agent/assistant/chat.go

@@ -7,39 +7,16 @@ import (
 	"time"
 
 	"github.com/google/uuid"
-	"github.com/yaoapp/kun/log"
 	agentcontext "github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/agent/i18n"
 	storetypes "github.com/yaoapp/yao/agent/store/types"
 	"github.com/yaoapp/yao/kb"
 	kbapi "github.com/yaoapp/yao/kb/api"
-	"github.com/yaoapp/yao/trace/types"
 )
 
 // kbCollectionCreating tracks collections currently being created to avoid duplicate creation
 var kbCollectionCreating sync.Map
 
-// WithHistory merges the input messages with chat history and traces it
-// This method can be overridden or extended to implement actual history loading
-func (ast *Assistant) WithHistory(ctx *agentcontext.Context, input []agentcontext.Message, agentNode types.Node, options ...*agentcontext.Options) ([]agentcontext.Message, error) {
-
-	// TODO: Implement actual history loading logic here
-	// For now, just simulate a check and return the input messages as is
-
-	// Simulate error check (this is where actual history loading would happen)
-	// if some_condition {
-	//     ast.traceAgentFail(agentNode, err)
-	//     return nil, err
-	// }
-
-	fullMessages := input
-
-	// Log the chat history
-	ast.traceAgentHistory(ctx, agentNode, fullMessages)
-
-	return fullMessages, nil
-}
-
 // InitializeConversation prepares KB collection for the conversation (synchronous)
 func (ast *Assistant) InitializeConversation(ctx *agentcontext.Context, options ...*agentcontext.Options) error {
 
@@ -57,7 +34,7 @@ func (ast *Assistant) InitializeConversation(ctx *agentcontext.Context, options
 
 	// Check if authorized info is available
 	if ctx.Authorized == nil {
-		fmt.Printf(">>> Warning: no authorized info, skipping KB collection preparation\n")
+		ctx.Logger.Warn("no authorized info, skipping KB collection preparation")
 		return nil
 	}
 
@@ -65,7 +42,7 @@ func (ast *Assistant) InitializeConversation(ctx *agentcontext.Context, options
 	err := ast.prepareKBCollection(ctx, opts)
 	if err != nil {
 		// Log but don't fail the chat
-		fmt.Printf(">>> Warning: failed to prepare KB collection: %v\n", err)
+		ctx.Logger.Warn("failed to prepare KB collection: %v", err)
 	}
 
 	return nil
@@ -98,7 +75,7 @@ func (ast *Assistant) prepareKBCollection(ctx *agentcontext.Context, opts *agent
 	chatKB := kbSetting.Chat
 
 	// Debug: log locale information
-	fmt.Printf(">>> prepareKBCollection: locale=%s\n", ctx.Locale)
+	ctx.Logger.Debug("prepareKBCollection: locale=%s", ctx.Locale)
 
 	// Get KB collection ID for this chat session
 	// Same team + user always produces the same ID (idempotent)
@@ -106,7 +83,7 @@ func (ast *Assistant) prepareKBCollection(ctx *agentcontext.Context, opts *agent
 
 	// Check if this collection is currently being created by another goroutine
 	if _, isCreating := kbCollectionCreating.LoadOrStore(collectionID, true); isCreating {
-		fmt.Printf(">>> KB collection %s is already being created, skipping\n", collectionID)
+		ctx.Logger.Debug("KB collection %s is already being created, skipping", collectionID)
 		return nil
 	}
 	// Ensure cleanup even if panic occurs
@@ -116,10 +93,10 @@ func (ast *Assistant) prepareKBCollection(ctx *agentcontext.Context, opts *agent
 	existsResult, err := kb.API.CollectionExists(ctx.Context, collectionID)
 	if err != nil {
 		// If check fails, log and continue to create (let create handle conflicts)
-		fmt.Printf(">>> Warning: failed to check collection existence: %v, will attempt to create\n", err)
+		ctx.Logger.Warn("failed to check collection existence: %v, will attempt to create", err)
 	} else if existsResult != nil && existsResult.Exists {
 		// Collection exists, no need to create
-		fmt.Printf(">>> KB collection already exists: %s\n", collectionID)
+		ctx.Logger.Debug("KB collection already exists: %s", collectionID)
 		return nil
 	}
 
@@ -139,7 +116,7 @@ func (ast *Assistant) prepareKBCollection(ctx *agentcontext.Context, opts *agent
 		return fmt.Errorf("failed to create KB collection: %w", err)
 	}
 
-	fmt.Printf(">>> Created KB collection: %s for team=%s, user=%s\n",
+	ctx.Logger.Info("Created KB collection: %s for team=%s, user=%s",
 		collectionID, ctx.Authorized.TeamID, ctx.Authorized.UserID)
 
 	_ = opts
@@ -209,8 +186,6 @@ func mergeChatMetadata(defaultMetadata map[string]interface{}, ctx *agentcontext
 		metadata["description"] = i18n.T(locale, "kb.chat.description")
 	}
 
-	fmt.Printf(">>> mergeChatMetadata: locale=%s, name=%v, description=%v\n", locale, metadata["name"], metadata["description"]) // Debug log
-
 	return metadata
 }
 
@@ -233,7 +208,7 @@ func (ast *Assistant) InitBuffer(ctx *agentcontext.Context) {
 
 	// Skip if History is disabled in options
 	if ctx.Stack.Options != nil && ctx.Stack.Options.Skip != nil && ctx.Stack.Options.Skip.History {
-		log.Trace("[CHAT] Buffer skipped: Skip.History is true")
+		ctx.Logger.Debug("Buffer skipped: Skip.History is true")
 		return
 	}
 
@@ -252,7 +227,7 @@ func (ast *Assistant) InitBuffer(ctx *agentcontext.Context) {
 	}
 
 	ctx.Buffer = agentcontext.NewChatBuffer(ctx.ChatID, requestID, ast.ID, connector, mode)
-	log.Trace("[CHAT] Buffer initialized: chatID=%s, requestID=%s, assistantID=%s, connector=%s, mode=%s", ctx.ChatID, requestID, ast.ID, connector, mode)
+	ctx.Logger.Debug("Buffer initialized: chatID=%s, requestID=%s, assistantID=%s", ctx.ChatID, requestID, ast.ID)
 }
 
 // BufferUserInput adds user input messages to the buffer
@@ -324,7 +299,7 @@ func (ast *Assistant) FlushBuffer(ctx *agentcontext.Context, finalStatus string,
 	// Get chat store
 	chatStore := GetChatStore()
 	if chatStore == nil {
-		log.Error("[CHAT] Chat store not available, cannot flush buffer")
+		ctx.Logger.Error("Chat store not available, cannot flush buffer")
 		return
 	}
 
@@ -337,9 +312,9 @@ func (ast *Assistant) FlushBuffer(ctx *agentcontext.Context, finalStatus string,
 	messages := ast.convertBufferedMessages(ctx.Buffer.GetMessages())
 	if len(messages) > 0 {
 		if saveErr := chatStore.SaveMessages(ctx.ChatID, messages); saveErr != nil {
-			log.Error("[CHAT] Failed to save messages: %v", saveErr)
+			ctx.Logger.Error("Failed to save messages: %v", saveErr)
 		} else {
-			log.Trace("[CHAT] Saved %d messages for chat=%s", len(messages), ctx.ChatID)
+			ctx.Logger.Debug("Saved %d messages for chat=%s", len(messages), ctx.ChatID)
 		}
 	}
 
@@ -358,7 +333,7 @@ func (ast *Assistant) FlushBuffer(ctx *agentcontext.Context, finalStatus string,
 			updates["last_mode"] = mode
 		}
 		if updateErr := chatStore.UpdateChat(ctx.ChatID, updates); updateErr != nil {
-			log.Trace("[CHAT] Failed to update chat: %v", updateErr)
+			ctx.Logger.Debug("Failed to update chat: %v", updateErr)
 		}
 	}
 
@@ -367,9 +342,9 @@ func (ast *Assistant) FlushBuffer(ctx *agentcontext.Context, finalStatus string,
 		steps := ast.convertBufferedSteps(ctx.Buffer.GetStepsForResume(finalStatus))
 		if len(steps) > 0 {
 			if saveErr := chatStore.SaveResume(steps); saveErr != nil {
-				log.Error("[CHAT] Failed to save resume steps: %v", saveErr)
+				ctx.Logger.Error("Failed to save resume steps: %v", saveErr)
 			} else {
-				log.Trace("[CHAT] Saved %d resume steps for chat=%s (status=%s)", len(steps), ctx.ChatID, finalStatus)
+				ctx.Logger.Debug("Saved %d resume steps for chat=%s (status=%s)", len(steps), ctx.ChatID, finalStatus)
 			}
 		}
 	}

agent/assistant/chat_test.go

@@ -103,14 +103,10 @@ func TestPrepareKBCollection(t *testing.T) {
 		teamID := fmt.Sprintf("test_team_%s", timestamp)
 		userID := fmt.Sprintf("test_user_%s", timestamp)
 
-		ctx := &agentcontext.Context{
+		ctx := agentcontext.New(context.Background(), &oauthtypes.AuthorizedInfo{
-			Context: context.Background(),
-			ChatID:  "test_chat_prepare_001",
-			Authorized: &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{
+		ctx := agentcontext.New(context.Background(), &oauthtypes.AuthorizedInfo{
-			Context: context.Background(),
-			ChatID:  "test_chat_idempotent",
-			Authorized: &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{
+		ctx := agentcontext.New(context.Background(), nil, "test_chat_no_auth") // Missing authorized info
-			Context:    context.Background(),
-			ChatID:     "test_chat_no_auth",
-			Authorized: nil, // 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{
+		ctx := agentcontext.New(context.Background(), &oauthtypes.AuthorizedInfo{
-			Context: context.Background(),
-			ChatID:  "test_chat_concurrent",
-			Authorized: &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{
+		ctx := agentcontext.New(context.Background(), &oauthtypes.AuthorizedInfo{
-			Context: context.Background(),
-			ChatID:  "test_init_chat_001",
-			Authorized: &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{
+		ctx := agentcontext.New(context.Background(), &oauthtypes.AuthorizedInfo{
-			Context: context.Background(),
-			ChatID:  "test_skip_history",
-			Authorized: &oauthtypes.AuthorizedInfo{
 			TeamID: "skip_team",
 			UserID: "skip_user",
-			},
+		}, "test_skip_history")
-		}
 
 		opts := &agentcontext.Options{
 			Skip: &agentcontext.Skip{

agent/assistant/history.go

@@ -0,0 +1,285 @@
+package assistant
+
+import (
+	"fmt"
+	"reflect"
+
+	agentcontext "github.com/yaoapp/yao/agent/context"
+	storetypes "github.com/yaoapp/yao/agent/store/types"
+	"github.com/yaoapp/yao/trace/types"
+)
+
+// =============================================================================
+// Chat History Management
+// =============================================================================
+
+// HistoryResult represents the result of history processing
+type HistoryResult struct {
+	InputMessages []agentcontext.Message // Clean input messages (without overlap)
+	FullMessages  []agentcontext.Message // Full messages (history + clean input)
+}
+
+// WithHistory merges the input messages with chat history and traces it
+// Returns HistoryResult containing:
+// - InputMessages: cleaned input (overlap removed)
+// - FullMessages: history + clean input merged
+func (ast *Assistant) WithHistory(ctx *agentcontext.Context, input []agentcontext.Message, agentNode types.Node, options ...*agentcontext.Options) (*HistoryResult, error) {
+
+	// Get options
+	var opts *agentcontext.Options
+	if len(options) > 0 && options[0] != nil {
+		opts = options[0]
+	}
+
+	// SKIP: History (for internal calls like title/prompt etc.)
+	if opts != nil && opts.Skip != nil && opts.Skip.History {
+		result := &HistoryResult{
+			InputMessages: input,
+			FullMessages:  input,
+		}
+		ast.traceAgentHistory(ctx, agentNode, result.FullMessages)
+		return result, nil
+	}
+
+	// Get MaxSize from store setting
+	maxSize := 20 // default
+	if storeSetting := GetStoreSetting(); storeSetting != nil && storeSetting.MaxSize > 0 {
+		maxSize = storeSetting.MaxSize
+	}
+
+	// Load history from store
+	historyMessages, err := ast.loadHistory(ctx)
+	if err != nil {
+		// Log warning but continue without history
+		ctx.Logger.Warn("Failed to load history for chat=%s: %v", ctx.ChatID, err)
+		result := &HistoryResult{
+			InputMessages: input,
+			FullMessages:  input,
+		}
+		ast.traceAgentHistory(ctx, agentNode, result.FullMessages)
+		return result, nil
+	}
+
+	// If no history, return input as is
+	if len(historyMessages) == 0 {
+		ctx.Logger.HistoryLoad(0, maxSize)
+		result := &HistoryResult{
+			InputMessages: input,
+			FullMessages:  input,
+		}
+		ast.traceAgentHistory(ctx, agentNode, result.FullMessages)
+		return result, nil
+	}
+
+	// Log history loaded
+	ctx.Logger.HistoryLoad(len(historyMessages), maxSize)
+
+	// Find overlap between history and input
+	// Some external clients may include history in their requests
+	overlapIndex := ast.findOverlapIndex(historyMessages, input)
+
+	// Remove overlap from input
+	cleanInput := input
+	if overlapIndex > 0 {
+		cleanInput = input[overlapIndex:]
+		ctx.Logger.HistoryOverlap(overlapIndex)
+	}
+
+	// Merge history with clean input
+	fullMessages := make([]agentcontext.Message, 0, len(historyMessages)+len(cleanInput))
+	fullMessages = append(fullMessages, historyMessages...)
+	fullMessages = append(fullMessages, cleanInput...)
+
+	result := &HistoryResult{
+		InputMessages: cleanInput,
+		FullMessages:  fullMessages,
+	}
+
+	// Log the chat history
+	ast.traceAgentHistory(ctx, agentNode, result.FullMessages)
+
+	return result, nil
+}
+
+// loadHistory loads chat history from the store
+// Returns the most recent MaxSize messages, ordered by time (oldest first)
+func (ast *Assistant) loadHistory(ctx *agentcontext.Context) ([]agentcontext.Message, error) {
+	// Check if chat ID is available
+	if ctx.ChatID == "" {
+		return nil, nil
+	}
+
+	// Get chat store
+	chatStore := GetChatStore()
+	if chatStore == nil {
+		return nil, nil
+	}
+
+	// Get store setting for MaxSize
+	setting := GetStoreSetting()
+	maxSize := 20 // default
+	if setting != nil && setting.MaxSize > 0 {
+		maxSize = setting.MaxSize
+	}
+
+	// Load messages from store with limit
+	filter := storetypes.MessageFilter{
+		Limit: maxSize,
+	}
+
+	storeMessages, err := chatStore.GetMessages(ctx.ChatID, filter)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get messages: %w", err)
+	}
+
+	if len(storeMessages) == 0 {
+		return nil, nil
+	}
+
+	// Convert store messages to context messages
+	messages := make([]agentcontext.Message, 0, len(storeMessages))
+	for _, msg := range storeMessages {
+		// Only include user and assistant messages for LLM context
+		// Skip internal types like loading, event, etc.
+		if msg.Role != "user" && msg.Role != "assistant" {
+			continue
+		}
+
+		// Convert store message to context message
+		ctxMsg := ast.convertStoreMessageToContext(msg)
+		if ctxMsg != nil {
+			messages = append(messages, *ctxMsg)
+		}
+	}
+
+	return messages, nil
+}
+
+// convertStoreMessageToContext converts a store message to a context message
+func (ast *Assistant) convertStoreMessageToContext(msg *storetypes.Message) *agentcontext.Message {
+	if msg == nil {
+		return nil
+	}
+
+	// Extract content from Props
+	content := ast.extractContentFromProps(msg.Props, msg.Type)
+	if content == nil {
+		return nil
+	}
+
+	// Build context message
+	ctxMsg := &agentcontext.Message{
+		Role:    agentcontext.MessageRole(msg.Role),
+		Content: content,
+	}
+
+	// Handle name field
+	if msg.Props != nil {
+		if name, ok := msg.Props["name"].(string); ok && name != "" {
+			ctxMsg.Name = &name
+		}
+	}
+
+	return ctxMsg
+}
+
+// extractContentFromProps extracts the content from message Props based on message type
+func (ast *Assistant) extractContentFromProps(props map[string]interface{}, msgType string) interface{} {
+	if props == nil {
+		return nil
+	}
+
+	// For user input, content is stored directly in props["content"]
+	if msgType == "user_input" {
+		return props["content"]
+	}
+
+	// For text type messages
+	if msgType == "text" {
+		if text, ok := props["text"].(string); ok {
+			return text
+		}
+		// Also try content field
+		if content, ok := props["content"].(string); ok {
+			return content
+		}
+	}
+
+	// For other types, try to extract content or text
+	if content, ok := props["content"]; ok {
+		return content
+	}
+	if text, ok := props["text"]; ok {
+		return text
+	}
+
+	return nil
+}
+
+// findOverlapIndex finds the index in input where history messages end
+// Returns the number of input messages that overlap with history
+func (ast *Assistant) findOverlapIndex(history, input []agentcontext.Message) int {
+	if len(history) == 0 || len(input) == 0 {
+		return 0
+	}
+
+	// We need to find the longest suffix of history that matches a prefix of input
+	// Start from the end of history and try to match with the beginning of input
+
+	maxOverlap := len(history)
+	if maxOverlap > len(input) {
+		maxOverlap = len(input)
+	}
+
+	// Try different overlap lengths, starting from the largest possible
+	for overlapLen := maxOverlap; overlapLen > 0; overlapLen-- {
+		// Check if the last 'overlapLen' messages of history match the first 'overlapLen' of input
+		historyStart := len(history) - overlapLen
+		matched := true
+
+		for i := 0; i < overlapLen; i++ {
+			if !ast.messagesMatch(history[historyStart+i], input[i]) {
+				matched = false
+				break
+			}
+		}
+
+		if matched {
+			return overlapLen
+		}
+	}
+
+	return 0
+}
+
+// messagesMatch checks if two messages are equivalent
+func (ast *Assistant) messagesMatch(a, b agentcontext.Message) bool {
+	// Must have same role
+	if a.Role != b.Role {
+		return false
+	}
+
+	// Compare content
+	return ast.contentMatches(a.Content, b.Content)
+}
+
+// contentMatches compares two content values for equality
+func (ast *Assistant) contentMatches(a, b interface{}) bool {
+	// Handle nil cases
+	if a == nil && b == nil {
+		return true
+	}
+	if a == nil || b == nil {
+		return false
+	}
+
+	// If both are strings, compare directly
+	aStr, aIsStr := a.(string)
+	bStr, bIsStr := b.(string)
+	if aIsStr && bIsStr {
+		return aStr == bStr
+	}
+
+	// For complex content (arrays, etc.), use deep equal
+	return reflect.DeepEqual(a, b)
+}

agent/assistant/history_test.go

@@ -0,0 +1,791 @@
+package assistant_test
+
+import (
+	"context"
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/yaoapp/yao/agent/assistant"
+	agentcontext "github.com/yaoapp/yao/agent/context"
+	"github.com/yaoapp/yao/agent/output/message"
+	storetypes "github.com/yaoapp/yao/agent/store/types"
+	"github.com/yaoapp/yao/agent/testutils"
+	oauthtypes "github.com/yaoapp/yao/openapi/oauth/types"
+)
+
+// =============================================================================
+// Helper Functions
+// =============================================================================
+
+// newHistoryTestContext creates a test context for history tests
+func newHistoryTestContext(chatID string) *agentcontext.Context {
+	authorized := &oauthtypes.AuthorizedInfo{
+		Subject:  "test-user",
+		UserID:   "history-test-user",
+		TeamID:   "history-test-team",
+		TenantID: "history-test-tenant",
+	}
+
+	ctx := agentcontext.New(context.Background(), authorized, chatID)
+	ctx.AssistantID = "tests.history"
+	ctx.Locale = "en-us"
+	ctx.Client = agentcontext.Client{
+		Type: "web",
+		IP:   "127.0.0.1",
+	}
+	ctx.Referer = agentcontext.RefererAPI
+	ctx.Accept = agentcontext.AcceptWebCUI
+	ctx.IDGenerator = message.NewIDGenerator()
+	ctx.Metadata = make(map[string]interface{})
+	return ctx
+}
+
+// =============================================================================
+// WithHistory Tests
+// =============================================================================
+
+func TestWithHistory(t *testing.T) {
+	testutils.Prepare(t)
+	defer testutils.Clean(t)
+
+	// Get assistant
+	ast, err := assistant.Get("tests.history")
+	require.NoError(t, err)
+	require.NotNil(t, ast)
+
+	// Get chat store for setup/cleanup
+	chatStore := assistant.GetChatStore()
+	if chatStore == nil {
+		t.Skip("Chat store not configured, skipping history tests")
+	}
+
+	t.Run("NoHistory", func(t *testing.T) {
+		chatID := fmt.Sprintf("test_history_none_%s", uuid.New().String()[:8])
+		ctx := newHistoryTestContext(chatID)
+
+		// Create chat without any messages
+		err := chatStore.CreateChat(&storetypes.Chat{
+			ChatID:      chatID,
+			AssistantID: ast.ID,
+			Status:      "active",
+			Share:       "private",
+			CreatedAt:   time.Now(),
+			UpdatedAt:   time.Now(),
+		})
+		require.NoError(t, err)
+		defer chatStore.DeleteChat(chatID)
+
+		input := []agentcontext.Message{
+			{Role: agentcontext.RoleUser, Content: "Hello, this is my first message"},
+		}
+
+		result, err := ast.WithHistory(ctx, input, nil)
+		require.NoError(t, err)
+		require.NotNil(t, result)
+
+		// With no history, InputMessages and FullMessages should be the same as input
+		assert.Equal(t, input, result.InputMessages)
+		assert.Equal(t, input, result.FullMessages)
+		t.Log("✓ No history: input returned as is")
+	})
+
+	t.Run("WithExistingHistory", func(t *testing.T) {
+		chatID := fmt.Sprintf("test_history_exist_%s", uuid.New().String()[:8])
+		ctx := newHistoryTestContext(chatID)
+		reqID := uuid.New().String()[:8]
+
+		// Create chat
+		err := chatStore.CreateChat(&storetypes.Chat{
+			ChatID:      chatID,
+			AssistantID: ast.ID,
+			Status:      "active",
+			Share:       "private",
+			CreatedAt:   time.Now(),
+			UpdatedAt:   time.Now(),
+		})
+		require.NoError(t, err)
+		defer func() {
+			chatStore.DeleteMessages(chatID, nil)
+			chatStore.DeleteChat(chatID)
+		}()
+
+		// Add history messages
+		historyMessages := []*storetypes.Message{
+			{
+				MessageID:   fmt.Sprintf("hist_msg_1_%s", reqID),
+				ChatID:      chatID,
+				RequestID:   fmt.Sprintf("req_%s", reqID),
+				Role:        "user",
+				Type:        "user_input",
+				Props:       map[string]interface{}{"content": "Previous question"},
+				Sequence:    1,
+				AssistantID: ast.ID,
+				CreatedAt:   time.Now().Add(-2 * time.Minute),
+			},
+			{
+				MessageID:   fmt.Sprintf("hist_msg_2_%s", reqID),
+				ChatID:      chatID,
+				RequestID:   fmt.Sprintf("req_%s", reqID),
+				Role:        "assistant",
+				Type:        "text",
+				Props:       map[string]interface{}{"text": "Previous answer"},
+				Sequence:    2,
+				AssistantID: ast.ID,
+				CreatedAt:   time.Now().Add(-1 * time.Minute),
+			},
+		}
+
+		err = chatStore.SaveMessages(chatID, historyMessages)
+		require.NoError(t, err)
+
+		// New input message
+		input := []agentcontext.Message{
+			{Role: agentcontext.RoleUser, Content: "New question"},
+		}
+
+		result, err := ast.WithHistory(ctx, input, nil)
+		require.NoError(t, err)
+		require.NotNil(t, result)
+
+		// InputMessages should be unchanged (no overlap)
+		assert.Equal(t, input, result.InputMessages)
+
+		// FullMessages should have history + input
+		assert.Len(t, result.FullMessages, 3) // 2 history + 1 new
+
+		// Verify order: history first, then input
+		assert.Equal(t, agentcontext.RoleUser, result.FullMessages[0].Role)
+		assert.Equal(t, "Previous question", result.FullMessages[0].Content)
+		assert.Equal(t, agentcontext.RoleAssistant, result.FullMessages[1].Role)
+		assert.Equal(t, "Previous answer", result.FullMessages[1].Content)
+		assert.Equal(t, agentcontext.RoleUser, result.FullMessages[2].Role)
+		assert.Equal(t, "New question", result.FullMessages[2].Content)
+
+		t.Log("✓ History merged correctly with new input")
+	})
+
+	t.Run("SkipHistoryOption", func(t *testing.T) {
+		chatID := fmt.Sprintf("test_history_skip_%s", uuid.New().String()[:8])
+		ctx := newHistoryTestContext(chatID)
+		reqID := uuid.New().String()[:8]
+
+		// Create chat with history
+		err := chatStore.CreateChat(&storetypes.Chat{
+			ChatID:      chatID,
+			AssistantID: ast.ID,
+			Status:      "active",
+			Share:       "private",
+			CreatedAt:   time.Now(),
+			UpdatedAt:   time.Now(),
+		})
+		require.NoError(t, err)
+		defer func() {
+			chatStore.DeleteMessages(chatID, nil)
+			chatStore.DeleteChat(chatID)
+		}()
+
+		// Add history message
+		err = chatStore.SaveMessages(chatID, []*storetypes.Message{
+			{
+				MessageID:   fmt.Sprintf("skip_hist_1_%s", reqID),
+				ChatID:      chatID,
+				RequestID:   fmt.Sprintf("req_skip_%s", reqID),
+				Role:        "user",
+				Type:        "user_input",
+				Props:       map[string]interface{}{"content": "Should be skipped"},
+				Sequence:    1,
+				AssistantID: ast.ID,
+				CreatedAt:   time.Now(),
+			},
+		})
+		require.NoError(t, err)
+
+		input := []agentcontext.Message{
+			{Role: agentcontext.RoleUser, Content: "Only this should appear"},
+		}
+
+		// Use Skip.History option
+		opts := &agentcontext.Options{
+			Skip: &agentcontext.Skip{
+				History: true,
+			},
+		}
+
+		result, err := ast.WithHistory(ctx, input, nil, opts)
+		require.NoError(t, err)
+		require.NotNil(t, result)
+
+		// Both should be same as input (history skipped)
+		assert.Equal(t, input, result.InputMessages)
+		assert.Equal(t, input, result.FullMessages)
+		assert.Len(t, result.FullMessages, 1)
+
+		t.Log("✓ History skipped when Skip.History=true")
+	})
+
+	t.Run("OverlapDetection", func(t *testing.T) {
+		chatID := fmt.Sprintf("test_history_overlap_%s", uuid.New().String()[:8])
+		ctx := newHistoryTestContext(chatID)
+		reqID := uuid.New().String()[:8]
+
+		// Create chat
+		err := chatStore.CreateChat(&storetypes.Chat{
+			ChatID:      chatID,
+			AssistantID: ast.ID,
+			Status:      "active",
+			Share:       "private",
+			CreatedAt:   time.Now(),
+			UpdatedAt:   time.Now(),
+		})
+		require.NoError(t, err)
+		defer func() {
+			chatStore.DeleteMessages(chatID, nil)
+			chatStore.DeleteChat(chatID)
+		}()
+
+		// Add history messages
+		err = chatStore.SaveMessages(chatID, []*storetypes.Message{
+			{
+				MessageID:   fmt.Sprintf("overlap_1_%s", reqID),
+				ChatID:      chatID,
+				RequestID:   fmt.Sprintf("req_overlap_%s", reqID),
+				Role:        "user",
+				Type:        "user_input",
+				Props:       map[string]interface{}{"content": "Message one"},
+				Sequence:    1,
+				AssistantID: ast.ID,
+				CreatedAt:   time.Now().Add(-3 * time.Minute),
+			},
+			{
+				MessageID:   fmt.Sprintf("overlap_2_%s", reqID),
+				ChatID:      chatID,
+				RequestID:   fmt.Sprintf("req_overlap_%s", reqID),
+				Role:        "assistant",
+				Type:        "text",
+				Props:       map[string]interface{}{"text": "Response one"},
+				Sequence:    2,
+				AssistantID: ast.ID,
+				CreatedAt:   time.Now().Add(-2 * time.Minute),
+			},
+			{
+				MessageID:   fmt.Sprintf("overlap_3_%s", reqID),
+				ChatID:      chatID,
+				RequestID:   fmt.Sprintf("req_overlap_2_%s", reqID),
+				Role:        "user",
+				Type:        "user_input",
+				Props:       map[string]interface{}{"content": "Message two"},
+				Sequence:    3,
+				AssistantID: ast.ID,
+				CreatedAt:   time.Now().Add(-1 * time.Minute),
+			},
+		})
+		require.NoError(t, err)
+
+		// Input that overlaps with history (includes last messages)
+		// Some clients send full history + new message
+		input := []agentcontext.Message{
+			{Role: agentcontext.RoleAssistant, Content: "Response one"}, // Overlap
+			{Role: agentcontext.RoleUser, Content: "Message two"},       // Overlap
+			{Role: agentcontext.RoleUser, Content: "New message"},       // New
+		}
+
+		result, err := ast.WithHistory(ctx, input, nil)
+		require.NoError(t, err)
+		require.NotNil(t, result)
+
+		// InputMessages should have overlap removed
+		assert.Len(t, result.InputMessages, 1, "Should remove 2 overlapping messages")
+		assert.Equal(t, "New message", result.InputMessages[0].Content)
+
+		// FullMessages should be history + clean input
+		assert.Len(t, result.FullMessages, 4) // 3 history + 1 new
+
+		t.Log("✓ Overlap detected and removed from input")
+	})
+
+	t.Run("EmptyChatID", func(t *testing.T) {
+		ctx := newHistoryTestContext("")
+
+		input := []agentcontext.Message{
+			{Role: agentcontext.RoleUser, Content: "No chat ID"},
+		}
+
+		result, err := ast.WithHistory(ctx, input, nil)
+		require.NoError(t, err)
+		require.NotNil(t, result)
+
+		// With empty chat ID, should return input as is
+		assert.Equal(t, input, result.InputMessages)
+		assert.Equal(t, input, result.FullMessages)
+
+		t.Log("✓ Empty chat ID handled gracefully")
+	})
+
+	t.Run("MultipleUserMessages", func(t *testing.T) {
+		chatID := fmt.Sprintf("test_history_multi_%s", uuid.New().String()[:8])
+		ctx := newHistoryTestContext(chatID)
+		reqID := uuid.New().String()[:8]
+
+		// Create chat with history
+		err := chatStore.CreateChat(&storetypes.Chat{
+			ChatID:      chatID,
+			AssistantID: ast.ID,
+			Status:      "active",
+			Share:       "private",
+			CreatedAt:   time.Now(),
+			UpdatedAt:   time.Now(),
+		})
+		require.NoError(t, err)
+		defer func() {
+			chatStore.DeleteMessages(chatID, nil)
+			chatStore.DeleteChat(chatID)
+		}()
+
+		// Add history
+		err = chatStore.SaveMessages(chatID, []*storetypes.Message{
+			{
+				MessageID:   fmt.Sprintf("multi_1_%s", reqID),
+				ChatID:      chatID,
+				RequestID:   fmt.Sprintf("req_multi_%s", reqID),
+				Role:        "user",
+				Type:        "user_input",
+				Props:       map[string]interface{}{"content": "First"},
+				Sequence:    1,
+				AssistantID: ast.ID,
+				CreatedAt:   time.Now().Add(-1 * time.Minute),
+			},
+		})
+		require.NoError(t, err)
+
+		// Multiple input messages
+		input := []agentcontext.Message{
+			{Role: agentcontext.RoleUser, Content: "Second"},
+			{Role: agentcontext.RoleUser, Content: "Third"},
+		}
+
+		result, err := ast.WithHistory(ctx, input, nil)
+		require.NoError(t, err)
+		require.NotNil(t, result)
+
+		assert.Len(t, result.InputMessages, 2)
+		assert.Len(t, result.FullMessages, 3) // 1 history + 2 new
+
+		t.Log("✓ Multiple input messages handled correctly")
+	})
+}
+
+// =============================================================================
+// History Load Tests
+// =============================================================================
+
+func TestHistoryLoading(t *testing.T) {
+	testutils.Prepare(t)
+	defer testutils.Clean(t)
+
+	ast, err := assistant.Get("tests.history")
+	require.NoError(t, err)
+
+	chatStore := assistant.GetChatStore()
+	if chatStore == nil {
+		t.Skip("Chat store not configured")
+	}
+
+	t.Run("FilterNonConversationTypes", func(t *testing.T) {
+		chatID := fmt.Sprintf("test_filter_%s", uuid.New().String()[:8])
+		ctx := newHistoryTestContext(chatID)
+		reqID := uuid.New().String()[:8]
+
+		// Create chat
+		err := chatStore.CreateChat(&storetypes.Chat{
+			ChatID:      chatID,
+			AssistantID: ast.ID,
+			Status:      "active",
+			Share:       "private",
+			CreatedAt:   time.Now(),
+			UpdatedAt:   time.Now(),
+		})
+		require.NoError(t, err)
+		defer func() {
+			chatStore.DeleteMessages(chatID, nil)
+			chatStore.DeleteChat(chatID)
+		}()
+
+		// Add various message types
+		err = chatStore.SaveMessages(chatID, []*storetypes.Message{
+			{
+				MessageID:   fmt.Sprintf("filter_1_%s", reqID),
+				ChatID:      chatID,
+				RequestID:   fmt.Sprintf("req_filter_%s", reqID),
+				Role:        "user",
+				Type:        "user_input",
+				Props:       map[string]interface{}{"content": "User message"},
+				Sequence:    1,
+				AssistantID: ast.ID,
+				CreatedAt:   time.Now().Add(-4 * time.Minute),
+			},
+			{
+				MessageID:   fmt.Sprintf("filter_2_%s", reqID),
+				ChatID:      chatID,
+				RequestID:   fmt.Sprintf("req_filter_%s", reqID),
+				Role:        "assistant",
+				Type:        "loading",
+				Props:       map[string]interface{}{"text": "Loading..."},
+				Sequence:    2,
+				AssistantID: ast.ID,
+				CreatedAt:   time.Now().Add(-3 * time.Minute),
+			},
+			{
+				MessageID:   fmt.Sprintf("filter_3_%s", reqID),
+				ChatID:      chatID,
+				RequestID:   fmt.Sprintf("req_filter_%s", reqID),
+				Role:        "assistant",
+				Type:        "text",
+				Props:       map[string]interface{}{"text": "Assistant response"},
+				Sequence:    3,
+				AssistantID: ast.ID,
+				CreatedAt:   time.Now().Add(-2 * time.Minute),
+			},
+			{
+				MessageID:   fmt.Sprintf("filter_4_%s", reqID),
+				ChatID:      chatID,
+				RequestID:   fmt.Sprintf("req_filter_%s", reqID),
+				Role:        "system",
+				Type:        "event",
+				Props:       map[string]interface{}{"event": "stream_end"},
+				Sequence:    4,
+				AssistantID: ast.ID,
+				CreatedAt:   time.Now().Add(-1 * time.Minute),
+			},
+		})
+		require.NoError(t, err)
+
+		input := []agentcontext.Message{
+			{Role: agentcontext.RoleUser, Content: "New input"},
+		}
+
+		result, err := ast.WithHistory(ctx, input, nil)
+		require.NoError(t, err)
+		require.NotNil(t, result)
+
+		// User and assistant roles are included, system is filtered out
+		// Note: loading type messages with role=assistant are included (role-based filtering)
+		// Only system role messages are filtered out
+		assert.GreaterOrEqual(t, len(result.FullMessages), 3) // At least 1 user + 1 assistant from history + 1 new
+
+		// Verify no system role messages
+		for _, msg := range result.FullMessages {
+			assert.NotEqual(t, "system", string(msg.Role))
+		}
+
+		t.Log("✓ System role messages filtered out")
+	})
+
+	t.Run("ContentExtraction", func(t *testing.T) {
+		chatID := fmt.Sprintf("test_extract_%s", uuid.New().String()[:8])
+		ctx := newHistoryTestContext(chatID)
+		reqID := uuid.New().String()[:8]
+
+		// Create chat
+		err := chatStore.CreateChat(&storetypes.Chat{
+			ChatID:      chatID,
+			AssistantID: ast.ID,
+			Status:      "active",
+			Share:       "private",
+			CreatedAt:   time.Now(),
+			UpdatedAt:   time.Now(),
+		})
+		require.NoError(t, err)
+		defer func() {
+			chatStore.DeleteMessages(chatID, nil)
+			chatStore.DeleteChat(chatID)
+		}()
+
+		// Add messages with different content formats
+		err = chatStore.SaveMessages(chatID, []*storetypes.Message{
+			{
+				MessageID:   fmt.Sprintf("extract_1_%s", reqID),
+				ChatID:      chatID,
+				RequestID:   fmt.Sprintf("req_extract_%s", reqID),
+				Role:        "user",
+				Type:        "user_input",
+				Props:       map[string]interface{}{"content": "User content from props.content"},
+				Sequence:    1,
+				AssistantID: ast.ID,
+				CreatedAt:   time.Now().Add(-2 * time.Minute),
+			},
+			{
+				MessageID:   fmt.Sprintf("extract_2_%s", reqID),
+				ChatID:      chatID,
+				RequestID:   fmt.Sprintf("req_extract_%s", reqID),
+				Role:        "assistant",
+				Type:        "text",
+				Props:       map[string]interface{}{"text": "Assistant content from props.text"},
+				Sequence:    2,
+				AssistantID: ast.ID,
+				CreatedAt:   time.Now().Add(-1 * time.Minute),
+			},
+		})
+		require.NoError(t, err)
+
+		input := []agentcontext.Message{
+			{Role: agentcontext.RoleUser, Content: "New message"},
+		}
+
+		result, err := ast.WithHistory(ctx, input, nil)
+		require.NoError(t, err)
+		require.NotNil(t, result)
+
+		// Verify content was extracted correctly
+		assert.Len(t, result.FullMessages, 3)
+		assert.Equal(t, "User content from props.content", result.FullMessages[0].Content)
+		assert.Equal(t, "Assistant content from props.text", result.FullMessages[1].Content)
+
+		t.Log("✓ Content extracted correctly from different formats")
+	})
+}
+
+// =============================================================================
+// Edge Cases Tests
+// =============================================================================
+
+func TestHistoryEdgeCases(t *testing.T) {
+	testutils.Prepare(t)
+	defer testutils.Clean(t)
+
+	ast, err := assistant.Get("tests.history")
+	require.NoError(t, err)
+
+	chatStore := assistant.GetChatStore()
+	if chatStore == nil {
+		t.Skip("Chat store not configured")
+	}
+
+	t.Run("EmptyInput", func(t *testing.T) {
+		chatID := fmt.Sprintf("test_empty_input_%s", uuid.New().String()[:8])
+		ctx := newHistoryTestContext(chatID)
+		reqID := uuid.New().String()[:8]
+
+		// Create chat with history
+		err := chatStore.CreateChat(&storetypes.Chat{
+			ChatID:      chatID,
+			AssistantID: ast.ID,
+			Status:      "active",
+			Share:       "private",
+			CreatedAt:   time.Now(),
+			UpdatedAt:   time.Now(),
+		})
+		require.NoError(t, err)
+		defer func() {
+			chatStore.DeleteMessages(chatID, nil)
+			chatStore.DeleteChat(chatID)
+		}()
+
+		err = chatStore.SaveMessages(chatID, []*storetypes.Message{
+			{
+				MessageID:   fmt.Sprintf("empty_input_1_%s", reqID),
+				ChatID:      chatID,
+				RequestID:   fmt.Sprintf("req_empty_%s", reqID),
+				Role:        "user",
+				Type:        "user_input",
+				Props:       map[string]interface{}{"content": "Previous"},
+				Sequence:    1,
+				AssistantID: ast.ID,
+				CreatedAt:   time.Now(),
+			},
+		})
+		require.NoError(t, err)
+
+		// Empty input
+		input := []agentcontext.Message{}
+
+		result, err := ast.WithHistory(ctx, input, nil)
+		require.NoError(t, err)
+		require.NotNil(t, result)
+
+		// Should return history only
+		assert.Empty(t, result.InputMessages)
+		assert.Len(t, result.FullMessages, 1)
+
+		t.Log("✓ Empty input handled correctly")
+	})
+
+	t.Run("FullOverlap", func(t *testing.T) {
+		chatID := fmt.Sprintf("test_full_overlap_%s", uuid.New().String()[:8])
+		ctx := newHistoryTestContext(chatID)
+		reqID := uuid.New().String()[:8]
+
+		// Create chat
+		err := chatStore.CreateChat(&storetypes.Chat{
+			ChatID:      chatID,
+			AssistantID: ast.ID,
+			Status:      "active",
+			Share:       "private",
+			CreatedAt:   time.Now(),
+			UpdatedAt:   time.Now(),
+		})
+		require.NoError(t, err)
+		defer func() {
+			chatStore.DeleteMessages(chatID, nil)
+			chatStore.DeleteChat(chatID)
+		}()
+
+		// Add history
+		err = chatStore.SaveMessages(chatID, []*storetypes.Message{
+			{
+				MessageID:   fmt.Sprintf("full_overlap_1_%s", reqID),
+				ChatID:      chatID,
+				RequestID:   fmt.Sprintf("req_full_%s", reqID),
+				Role:        "user",
+				Type:        "user_input",
+				Props:       map[string]interface{}{"content": "Exact same message"},
+				Sequence:    1,
+				AssistantID: ast.ID,
+				CreatedAt:   time.Now(),
+			},
+		})
+		require.NoError(t, err)
+
+		// Input is exactly the same as history
+		input := []agentcontext.Message{
+			{Role: agentcontext.RoleUser, Content: "Exact same message"},
+		}
+
+		result, err := ast.WithHistory(ctx, input, nil)
+		require.NoError(t, err)
+		require.NotNil(t, result)
+
+		// Full overlap: clean input should be empty
+		assert.Empty(t, result.InputMessages)
+		// FullMessages should be just history (no duplicates)
+		assert.Len(t, result.FullMessages, 1)
+
+		t.Log("✓ Full overlap handled correctly")
+	})
+
+	t.Run("NonExistentChat", func(t *testing.T) {
+		chatID := "non_existent_chat_12345"
+		ctx := newHistoryTestContext(chatID)
+
+		input := []agentcontext.Message{
+			{Role: agentcontext.RoleUser, Content: "Message to non-existent chat"},
+		}
+
+		result, err := ast.WithHistory(ctx, input, nil)
+		require.NoError(t, err)
+		require.NotNil(t, result)
+
+		// Should return input as is (no history found)
+		assert.Equal(t, input, result.InputMessages)
+		assert.Equal(t, input, result.FullMessages)
+
+		t.Log("✓ Non-existent chat handled gracefully")
+	})
+
+	t.Run("MessageWithName", func(t *testing.T) {
+		chatID := fmt.Sprintf("test_name_%s", uuid.New().String()[:8])
+		ctx := newHistoryTestContext(chatID)
+		reqID := uuid.New().String()[:8]
+
+		// Create chat
+		err := chatStore.CreateChat(&storetypes.Chat{
+			ChatID:      chatID,
+			AssistantID: ast.ID,
+			Status:      "active",
+			Share:       "private",
+			CreatedAt:   time.Now(),
+			UpdatedAt:   time.Now(),
+		})
+		require.NoError(t, err)
+		defer func() {
+			chatStore.DeleteMessages(chatID, nil)
+			chatStore.DeleteChat(chatID)
+		}()
+
+		// Add message with name
+		err = chatStore.SaveMessages(chatID, []*storetypes.Message{
+			{
+				MessageID:   fmt.Sprintf("name_msg_1_%s", reqID),
+				ChatID:      chatID,
+				RequestID:   fmt.Sprintf("req_name_%s", reqID),
+				Role:        "user",
+				Type:        "user_input",
+				Props:       map[string]interface{}{"content": "Message with name", "name": "John"},
+				Sequence:    1,
+				AssistantID: ast.ID,
+				CreatedAt:   time.Now(),
+			},
+		})
+		require.NoError(t, err)
+
+		input := []agentcontext.Message{
+			{Role: agentcontext.RoleUser, Content: "New message"},
+		}
+
+		result, err := ast.WithHistory(ctx, input, nil)
+		require.NoError(t, err)
+		require.NotNil(t, result)
+
+		// First message should have name
+		assert.Len(t, result.FullMessages, 2)
+		assert.NotNil(t, result.FullMessages[0].Name)
+		assert.Equal(t, "John", *result.FullMessages[0].Name)
+
+		t.Log("✓ Message name field preserved")
+	})
+
+	t.Run("EmptyContent", func(t *testing.T) {
+		chatID := fmt.Sprintf("test_empty_content_%s", uuid.New().String()[:8])
+		ctx := newHistoryTestContext(chatID)
+		reqID := uuid.New().String()[:8]
+
+		// Create chat
+		err := chatStore.CreateChat(&storetypes.Chat{
+			ChatID:      chatID,
+			AssistantID: ast.ID,
+			Status:      "active",
+			Share:       "private",
+			CreatedAt:   time.Now(),
+			UpdatedAt:   time.Now(),
+		})
+		require.NoError(t, err)
+		defer func() {
+			chatStore.DeleteMessages(chatID, nil)
+			chatStore.DeleteChat(chatID)
+		}()
+
+		// Add message with empty content in props
+		err = chatStore.SaveMessages(chatID, []*storetypes.Message{
+			{
+				MessageID:   fmt.Sprintf("empty_content_1_%s", reqID),
+				ChatID:      chatID,
+				RequestID:   fmt.Sprintf("req_empty_content_%s", reqID),
+				Role:        "user",
+				Type:        "user_input",
+				Props:       map[string]interface{}{}, // empty props (no content)
+				Sequence:    1,
+				AssistantID: ast.ID,
+				CreatedAt:   time.Now(),
+			},
+		})
+		require.NoError(t, err)
+
+		input := []agentcontext.Message{
+			{Role: agentcontext.RoleUser, Content: "New message"},
+		}
+
+		result, err := ast.WithHistory(ctx, input, nil)
+		require.NoError(t, err)
+		require.NotNil(t, result)
+
+		// Message with empty props should be skipped (no content extractable)
+		// Only new input should be present
+		assert.Len(t, result.FullMessages, 1)
+		assert.Equal(t, "New message", result.FullMessages[0].Content)
+
+		t.Log("✓ Empty content handled gracefully (message skipped)")
+	})
+}

agent/assistant/hook/create_bench_test.go

@@ -4,7 +4,6 @@ import (
 	stdContext "context"
 	"testing"
 
-	"github.com/yaoapp/gou/plan"
 	"github.com/yaoapp/yao/agent/assistant"
 	"github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/agent/testutils"
@@ -296,23 +295,7 @@ func getBusinessScenarios() []struct {
 
 // newBenchContext creates a minimal context for benchmarking
 func newBenchContext(chatID, assistantID string) *context.Context {
-	return &context.Context{
+	authorized := &types.AuthorizedInfo{
-		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",
@@ -324,6 +307,20 @@ func newBenchContext(chatID, assistantID string) *context.Context {
 				"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
 }

agent/assistant/hook/create_mem_test.go

@@ -6,7 +6,6 @@ import (
 	"testing"
 	"time"
 
-	"github.com/yaoapp/gou/plan"
 	"github.com/yaoapp/yao/agent/assistant"
 	"github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/agent/testutils"
@@ -610,23 +609,7 @@ func TestIsolateDisposal(t *testing.T) {
 
 // newMemTestContext creates a context for memory leak testing
 func newMemTestContext(chatID, assistantID string) *context.Context {
-	return &context.Context{
+	authorized := &types.AuthorizedInfo{
-		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",
@@ -638,6 +621,20 @@ func newMemTestContext(chatID, assistantID string) *context.Context {
 				"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
 }

agent/assistant/hook/create_test.go

@@ -4,7 +4,6 @@ import (
 	stdContext "context"
 	"testing"
 
-	"github.com/yaoapp/gou/plan"
 	"github.com/yaoapp/yao/agent/assistant"
 	"github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/agent/testutils"
@@ -14,23 +13,7 @@ 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{
+	authorized := &types.AuthorizedInfo{
-		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",
@@ -50,8 +33,22 @@ func newTestContext(chatID, assistantID string) *context.Context {
 				"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

agent/assistant/hook/goroutine_leak_test.go

@@ -10,7 +10,6 @@ import (
 	"testing"
 	"time"
 
-	"github.com/yaoapp/gou/plan"
 	"github.com/yaoapp/yao/agent/assistant"
 	"github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/agent/testutils"
@@ -294,23 +293,7 @@ func truncate(s string, max int) string {
 }
 
 func newLeakTestContext(chatID, assistantID string) *context.Context {
-	return &context.Context{
+	authorized := &types.AuthorizedInfo{
-		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",
@@ -322,6 +305,20 @@ func newLeakTestContext(chatID, assistantID string) *context.Context {
 				"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
 }

agent/assistant/hook/next_test.go

@@ -4,7 +4,6 @@ import (
 	stdContext "context"
 	"testing"
 
-	"github.com/yaoapp/gou/plan"
 	"github.com/yaoapp/yao/agent/assistant"
 	"github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/agent/output/message"
@@ -15,23 +14,7 @@ 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{
+	authorized := &types.AuthorizedInfo{
-		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",
@@ -51,8 +34,22 @@ func newTestContextForNext(chatID, assistantID string) *context.Context {
 				"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

agent/assistant/hook/realworld_next_test.go

@@ -5,7 +5,6 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
-	"github.com/yaoapp/gou/plan"
 	"github.com/yaoapp/yao/agent/assistant"
 	"github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/agent/testutils"
@@ -14,23 +13,7 @@ import (
 
 // newRealWorldNextContext creates a Context for real world Next Hook testing
 func newRealWorldNextContext(chatID, assistantID string) *context.Context {
-	return &context.Context{
+	authorized := &types.AuthorizedInfo{
-		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",
@@ -38,8 +21,22 @@ func newRealWorldNextContext(chatID, assistantID string) *context.Context {
 		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)

agent/assistant/hook/realworld_stress_test.go

@@ -9,7 +9,6 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/assert"
-	"github.com/yaoapp/gou/plan"
 	"github.com/yaoapp/yao/agent/assistant"
 	"github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/agent/testutils"
@@ -721,23 +720,7 @@ func TestRealWorldStressResourceHeavy(t *testing.T) {
 
 // newRealWorldContext creates a Context for real-world testing
 func newRealWorldContext(chatID, assistantID string) *context.Context {
-	return &context.Context{
+	authorized := &types.AuthorizedInfo{
-		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",
@@ -757,8 +740,22 @@ func newRealWorldContext(chatID, assistantID string) *context.Context {
 				"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

agent/assistant/llm.go

@@ -1,9 +1,6 @@
 package assistant
 
 import (
-	"fmt"
-
-	"github.com/yaoapp/kun/log"
 	"github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/agent/llm"
 	"github.com/yaoapp/yao/agent/output/message"
@@ -21,12 +18,6 @@ func (ast *Assistant) executeLLMStream(
 	opts *context.Options,
 ) (*context.CompletionResponse, error) {
 
-	// === Debug LLM Stream Start ===
-	fmt.Println(">>> executeLLMStream: STARTING")
-	fmt.Printf(">>> Messages count: %d\n", len(completionMessages))
-	fmt.Printf(">>> Tools count: %d\n", len(completionOptions.Tools))
-	// === End Debug ===
-
 	// Get connector object (capabilities were already set above, before stream_start)
 	conn, capabilities, err := ast.GetConnector(ctx, opts)
 	if err != nil {
@@ -45,6 +36,9 @@ func (ast *Assistant) executeLLMStream(
 	// Trace Add LLM request
 	ast.traceLLMRequest(ctx, conn.ID(), completionMessages, completionOptions)
 
+	// Log LLM call start
+	ctx.Logger.LLMStart(conn.ID(), "", len(completionMessages))
+
 	// Create LLM instance with connector and options
 	llmInstance, err := llm.New(conn, completionOptions)
 	if err != nil {
@@ -54,25 +48,9 @@ func (ast *Assistant) executeLLMStream(
 	}
 
 	// Call the LLM Completion Stream (streamHandler was set earlier)
-	log.Trace("[AGENT] Calling LLM Stream: assistant=%s", ast.ID)
-
-	// === Debug LLM Stream Call ===
-	fmt.Println(">>> executeLLMStream: CALLING llmInstance.Stream()")
-	// === End Debug ===
-
 	completionResponse, err := llmInstance.Stream(ctx, completionMessages, completionOptions, streamHandler)
 
-	// === Debug LLM Stream Return ===
-	fmt.Println(">>> executeLLMStream: llmInstance.Stream() RETURNED")
-	fmt.Printf(">>> err: %v\n", err)
-	if completionResponse != nil {
-		fmt.Printf(">>> ToolCalls: %d\n", len(completionResponse.ToolCalls))
-	}
-	// === End Debug ===
-
-	log.Trace("[AGENT] LLM Stream returned: assistant=%s, err=%v", ast.ID, err)
 	if err != nil {
-		log.Trace("[AGENT] Calling sendStreamEndOnError")
 		// Mark LLM Request as failed in trace
 		ast.traceLLMFail(ctx, err)
 		return nil, err
@@ -111,6 +89,9 @@ func (ast *Assistant) executeLLMForToolRetry(
 	// Trace Add LLM retry request
 	ast.traceLLMRetryRequest(ctx, conn.ID(), completionMessages, completionOptions)
 
+	// Log LLM call start (retry)
+	ctx.Logger.LLMStart(conn.ID(), "", len(completionMessages))
+
 	// Create LLM instance with connector and options
 	llmInstance, err := llm.New(conn, completionOptions)
 	if err != nil {
@@ -120,9 +101,7 @@ func (ast *Assistant) executeLLMForToolRetry(
 	}
 
 	// Call the LLM Completion Stream (still streaming for tool retry)
-	log.Trace("[AGENT] Calling LLM Stream for tool retry: assistant=%s", ast.ID)
 	completionResponse, err := llmInstance.Stream(ctx, completionMessages, completionOptions, streamHandler)
-	log.Trace("[AGENT] LLM tool retry stream returned: assistant=%s, err=%v", ast.ID, err)
 	if err != nil {
 		// Mark LLM Retry Request as failed in trace
 		ast.traceLLMFail(ctx, err)

agent/assistant/load.go

@@ -21,6 +21,7 @@ import (
 // loaded the loaded assistant
 var loaded = NewCache(200) // 200 is the default capacity
 var storage store.Store = nil
+var storeSetting *store.Setting = nil // store setting from agent.yml
 var search interface{} = nil
 var modelCapabilities map[string]gouOpenAI.Capabilities = map[string]gouOpenAI.Capabilities{}
 var defaultConnector string = ""           // default connector
@@ -153,6 +154,16 @@ func SetGlobalPrompts(prompts []store.Prompt) {
 	globalPrompts = prompts
 }
 
+// SetStoreSetting set the store setting from agent.yml
+func SetStoreSetting(setting *store.Setting) {
+	storeSetting = setting
+}
+
+// GetStoreSetting returns the store setting
+func GetStoreSetting() *store.Setting {
+	return storeSetting
+}
+
 // GetGlobalPrompts returns the global prompts with variables parsed
 // ctx: context variables for parsing $CTX.* variables
 func GetGlobalPrompts(ctx map[string]string) []store.Prompt {

agent/assistant/load_store_test.go

@@ -176,22 +176,7 @@ 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{
+	authorized := &types.AuthorizedInfo{
-		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",
@@ -199,8 +184,22 @@ func newStoreTestContext(chatID, assistantID string) *context.Context {
 		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

agent/assistant/mcp.go

@@ -9,7 +9,6 @@ import (
 	gouJson "github.com/yaoapp/gou/json"
 	"github.com/yaoapp/gou/mcp"
 	mcpTypes "github.com/yaoapp/gou/mcp/types"
-	"github.com/yaoapp/kun/log"
 	agentContext "github.com/yaoapp/yao/agent/context"
 	storeTypes "github.com/yaoapp/yao/agent/store/types"
 	"github.com/yaoapp/yao/trace/types"
@@ -109,21 +108,21 @@ func (ast *Assistant) buildMCPTools(ctx *agentContext.Context, createResponse *a
 	// Process each MCP server in order
 	for _, serverConfig := range servers {
 		if len(allTools) >= MaxMCPTools {
-			log.Warn("[Assistant MCP] Reached maximum tool limit (%d), skipping remaining servers", MaxMCPTools)
+			ctx.Logger.Warn("Reached maximum tool limit (%d), skipping remaining servers", MaxMCPTools)
 			break
 		}
 
 		// Get MCP client
 		client, err := mcp.Select(serverConfig.ServerID)
 		if err != nil {
-			log.Warn("[Assistant MCP] Failed to select MCP client '%s': %v", serverConfig.ServerID, err)
+			ctx.Logger.Warn("Failed to select MCP client '%s': %v", serverConfig.ServerID, err)
 			continue
 		}
 
 		// Get tools list (filter by serverConfig.Tools if specified)
 		toolsResponse, err := client.ListTools(mcpCtx, "")
 		if err != nil {
-			log.Warn("[Assistant MCP] Failed to list tools for '%s': %v", serverConfig.ServerID, err)
+			ctx.Logger.Warn("Failed to list tools for '%s': %v", serverConfig.ServerID, err)
 			continue
 		}
 
@@ -204,7 +203,7 @@ func (ast *Assistant) buildMCPTools(ctx *agentContext.Context, createResponse *a
 			}
 		}
 
-		log.Trace("[Assistant MCP] Loaded %d tools from server '%s'", len(toolsResponse.Tools), serverConfig.ServerID)
+		ctx.Logger.Debug("Loaded %d tools from server '%s'", len(toolsResponse.Tools), serverConfig.ServerID)
 	}
 
 	samplesPrompt := ""
@@ -212,7 +211,7 @@ func (ast *Assistant) buildMCPTools(ctx *agentContext.Context, createResponse *a
 		samplesPrompt = samplesBuilder.String()
 	}
 
-	log.Trace("[Assistant MCP] Total MCP tools loaded: %d", len(allTools))
+	ctx.Logger.Debug("Total MCP tools loaded: %d", len(allTools))
 	return allTools, samplesPrompt, nil
 }
 
@@ -226,32 +225,20 @@ func (ast *Assistant) executeToolCalls(ctx *agentContext.Context, toolCalls []ag
 		return nil, false
 	}
 
-	// === Debug ===
+	ctx.Logger.Debug("Executing %d tool calls (attempt %d)", len(toolCalls), attempt)
-	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)
 
 	// Single tool call
 	if len(toolCalls) == 1 {
-		fmt.Println(">>> executeToolCalls: Calling executeSingleToolCall")
+		return ast.executeSingleToolCall(ctx, toolCalls[0])
-		results, hasErrors := ast.executeSingleToolCall(ctx, toolCalls[0])
-		fmt.Printf(">>> executeToolCalls: executeSingleToolCall RETURNED (hasErrors: %v)\n", hasErrors)
-		return results, hasErrors
 	}
 
 	// Multiple tool calls - try parallel first
-	fmt.Println(">>> executeToolCalls: Calling executeMultipleToolCallsParallel")
+	return ast.executeMultipleToolCallsParallel(ctx, toolCalls)
-	results, hasErrors := ast.executeMultipleToolCallsParallel(ctx, toolCalls)
-	fmt.Printf(">>> executeToolCalls: executeMultipleToolCallsParallel RETURNED (hasErrors: %v)\n", hasErrors)
-	return results, hasErrors
 }
 
 // executeSingleToolCall executes a single tool call with trace logging
 func (ast *Assistant) executeSingleToolCall(ctx *agentContext.Context, toolCall agentContext.ToolCall) ([]ToolCallResult, bool) {
-	// === Debug ===
+	ctx.Logger.ToolStart(toolCall.Function.Name)
-	fmt.Printf(">>> executeSingleToolCall: START (tool: %s)\n", toolCall.Function.Name)
-	// === End Debug ===
 
 	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)
+	ctx.Logger.Debug("Calling tool: %s (server: %s)", toolName, serverID)
-	fmt.Printf(">>> executeSingleToolCall: CALLING client.CallTool (tool: %s, server: %s)\n", 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)
+			ctx.Logger.Debug("trace.Parallel() failed: %v", err)
-		} else {
-			fmt.Printf(">>> executeServerToolsParallelWithTrace: Created %d trace nodes\n", len(toolNodes))
 		}
-	} 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)
 			}

agent/context/authorized_test.go

@@ -1,10 +1,11 @@
-package context
+package context_test
 
 import (
-	"context"
+	stdContext "context"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/config"
 	"github.com/yaoapp/yao/openapi/oauth/types"
 	"github.com/yaoapp/yao/test"
@@ -23,7 +24,8 @@ func TestContextNew_PreservesAuthorizedInfo(t *testing.T) {
 	}
 
 	// Create context using New()
-	ctx := New(context.Background(), authInfo, "test-chat-123")
+	ctx := context.New(stdContext.Background(), authInfo, "test-chat-123")
+	defer ctx.Release()
 
 	// Verify authorized info is preserved
 	assert.NotNil(t, ctx)
@@ -45,9 +47,15 @@ func TestContextTrace_SavesAuthorizedInfo(t *testing.T) {
 		TenantID: "tenant-001",
 	}
 
-	// Create context
+	// Create context using New
-	ctx := New(context.Background(), authInfo, "test-chat-456")
+	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)

agent/context/chat_test.go

@@ -1,9 +1,10 @@
-package context
+package context_test
 
 import (
 	"testing"
 
 	"github.com/yaoapp/gou/store"
+	"github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/config"
 	"github.com/yaoapp/yao/test"
 )
@@ -23,15 +24,15 @@ func TestGetChatIDByMessages_NewConversation(t *testing.T) {
 
 	cache := getTestCache(t)
 
-	messages := []Message{
+	messages := []context.Message{
 		{
-			Role:    RoleUser,
+			Role:    context.RoleUser,
 			Content: "Hello, how are you?",
 		},
 	}
 
 	// First request - should generate new chat ID
-	chatID1, err := GetChatIDByMessages(cache, messages)
+	chatID1, err := context.GetChatIDByMessages(cache, messages)
 	if err != nil {
 		t.Fatalf("Failed to get chat ID: %v", err)
 	}
@@ -42,7 +43,7 @@ func TestGetChatIDByMessages_NewConversation(t *testing.T) {
 
 	// Second request with same single user message - should generate DIFFERENT chat ID
 	// (single user message always generates new chat ID to avoid false matches)
-	chatID2, err := GetChatIDByMessages(cache, messages)
+	chatID2, err := context.GetChatIDByMessages(cache, messages)
 	if err != nil {
 		t.Fatalf("Failed to get chat ID: %v", err)
 	}
@@ -65,10 +66,10 @@ func TestGetChatIDByMessages_ContinuousConversation(t *testing.T) {
 
 	// Scenario: User conversation with incrementally added messages
 	// Request 1: [user1]
-	messages1 := []Message{
+	messages1 := []context.Message{
-		{Role: RoleUser, Content: "First 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{
+	messages2 := []context.Message{
-		{Role: RoleUser, Content: "First message"},
+		{Role: context.RoleUser, Content: "First message"},
-		{Role: RoleUser, Content: "Second 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{
+	messages3 := []context.Message{
-		{Role: RoleUser, Content: "First message"},
+		{Role: context.RoleUser, Content: "First message"},
-		{Role: RoleUser, Content: "Second message"},
+		{Role: context.RoleUser, Content: "Second message"},
-		{Role: RoleUser, Content: "Third 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{
+	messages4 := []context.Message{
-		{Role: RoleUser, Content: "First message"},
+		{Role: context.RoleUser, Content: "First message"},
-		{Role: RoleUser, Content: "Second message"},
+		{Role: context.RoleUser, Content: "Second message"},
-		{Role: RoleUser, Content: "Third message"},
+		{Role: context.RoleUser, Content: "Third message"},
-		{Role: RoleUser, Content: "Fourth 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,
+			Role: context.RoleUser,
-			Content: []ContentPart{
+			Content: []context.ContentPart{
 				{
-					Type: ContentText,
+					Type: context.ContentText,
 					Text: "What's in this image?",
 				},
 				{
-					Type: ContentImageURL,
+					Type: context.ContentImageURL,
-					ImageURL: &ImageURL{
+					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{
+	messages2 := append(messages1, context.Message{
-		Role:    RoleUser,
+		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,
+					Type: context.ToolTypeFunction,
-					Function: Function{
+					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")

agent/context/context.go

@@ -27,14 +27,17 @@ func New(parent context.Context, authorized *types.AuthorizedInfo, chatID string
 		parent = context.Background()
 	}
 
+	contextID := generateContextID()
+
 	ctx := &Context{
 		Context:         parent,
-		ID:              generateContextID(), // Generate unique ID for the context
+		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
+		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())
+			trace.MarkCancelled(ctx.Stack.TraceID, ctx.Context.Err().Error())
-
+			trace.Release(ctx.Stack.TraceID)
-			// 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")
+			ctx.trace.MarkComplete()
+			trace.Release(ctx.Stack.TraceID)
 		}
-
-			// 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")
-			}
-		} 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 = 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)
+	// Close logger (MUST be last)
-	ctx = nil
+	if ctx.Logger != nil {
+		ctx.Logger.Close()
+		ctx.Logger = nil
+	}
 }
 
 // Send sends data to the context's writer

agent/context/context_test.go

@@ -1,7 +1,8 @@
-package context
+package context_test
 
 import (
 	"bytes"
+	stdContext "context"
 	"encoding/json"
 	"net/http"
 	"net/http/httptest"
@@ -10,6 +11,7 @@ import (
 	"github.com/gin-gonic/gin"
 	"github.com/stretchr/testify/assert"
 	"github.com/yaoapp/gou/store"
+	"github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/config"
 	"github.com/yaoapp/yao/test"
 )
@@ -37,7 +39,7 @@ func TestGetCompletionRequest(t *testing.T) {
 		expectedLocale      string
 		expectedTheme       string
 		expectedReferer     string
-		expectedAccept      Accept
+		expectedAccept      context.Accept
 		expectedAssistantID string
 		expectError         bool
 	}{
@@ -64,8 +66,8 @@ func TestGetCompletionRequest(t *testing.T) {
 			expectedStream:      boolPtr(true),
 			expectedLocale:      "zh-cn",
 			expectedTheme:       "dark",
-			expectedReferer:     RefererProcess,
+			expectedReferer:     context.RefererProcess,
-			expectedAccept:      AcceptWebCUI,
+			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,
+			expectedReferer:     context.RefererAPI,
-			expectedAccept:      AcceptStandard,
+			expectedAccept:      context.AcceptStandard,
 			expectedAssistantID: "test456",
 			expectError:         false,
 		},
@@ -114,8 +116,8 @@ func TestGetCompletionRequest(t *testing.T) {
 			expectedMsgCount:    1,
 			expectedLocale:      "",
 			expectedTheme:       "",
-			expectedReferer:     RefererMCP,
+			expectedReferer:     context.RefererMCP,
-			expectedAccept:      AcceptDesktopCUI,
+			expectedAccept:      context.AcceptDesktopCUI,
 			expectedAssistantID: "header789",
 			expectError:         false,
 		},
@@ -131,8 +133,8 @@ func TestGetCompletionRequest(t *testing.T) {
 			expectedMsgCount:    1,
 			expectedLocale:      "",
 			expectedTheme:       "",
-			expectedReferer:     RefererAPI,
+			expectedReferer:     context.RefererAPI,
-			expectedAccept:      AcceptStandard,
+			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) {
+func TestContextNew_WithAuthorized(t *testing.T) {
-	tests := []struct {
+	test.Prepare(t, config.Conf)
-		name      string
+	defer test.Clean()
-		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"},
-	}
 
-	for _, tt := range tests {
+	// Create context using New()
-		t.Run(tt.name, func(t *testing.T) {
+	ctx := context.New(stdContext.Background(), nil, "test-chat-id")
-			result := getClientType(tt.userAgent)
+	defer ctx.Release()
-			assert.Equal(t, tt.expected, result)
+
-		})
+	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) {
+// Helper functions for context_test package
-	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
 func floatPtr(f float64) *float64 {
 	return &f
 }

agent/context/interrupt_test.go

@@ -1,4 +1,4 @@
-package context
+package context_test
 
 import (
 	stdContext "context"
@@ -6,48 +6,41 @@ import (
 	"testing"
 	"time"
 
-	"github.com/yaoapp/gou/plan"
+	"github.com/yaoapp/yao/agent/context"
-	"github.com/yaoapp/yao/agent/output/message"
 	"github.com/yaoapp/yao/openapi/oauth/types"
 )
 
 // newTestContextWithInterrupt creates a Context with interrupt controller for testing
-func newTestContextWithInterrupt(chatID, assistantID string) *Context {
+func newTestContextWithInterrupt(chatID, assistantID string) *context.Context {
-	ctx := &Context{
+	ctx := context.New(stdContext.Background(), &types.AuthorizedInfo{
-		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",
-		},
+	}, 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{
+		signal := &context.InterruptSignal{
-			Type: InterruptGraceful,
+			Type: context.InterruptGraceful,
-			Messages: []Message{
+			Messages: []context.Message{
-				{Role: RoleUser, Content: "This is a graceful interrupt"},
+				{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{
+		signal := &context.InterruptSignal{
-			Type: InterruptForce,
+			Type: context.InterruptForce,
-			Messages: []Message{
+			Messages: []context.Message{
-				{Role: RoleUser, Content: "This is a force interrupt"},
+				{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{
+			signal := &context.InterruptSignal{
-				Type: InterruptGraceful,
+				Type: context.InterruptGraceful,
-				Messages: []Message{
+				Messages: []context.Message{
-					{Role: RoleUser, Content: fmt.Sprintf("Message %d", i+1)},
+					{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{
+		signal := &context.InterruptSignal{
-			Type: InterruptGraceful,
+			Type: context.InterruptGraceful,
-			Messages: []Message{
+			Messages: []context.Message{
-				{Role: RoleUser, Content: "Test handler invocation"},
+				{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{
+		signal := &context.InterruptSignal{
-			Type: InterruptForce,
+			Type: context.InterruptForce,
-			Messages: []Message{
+			Messages: []context.Message{
-				{Role: RoleUser, Content: "Test error handling"},
+				{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{
+		signal := &context.InterruptSignal{
-			Type:      InterruptGraceful,
+			Type:      context.InterruptGraceful,
-			Messages:  []Message{{Role: RoleUser, Content: "test"}},
+			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{
+		signal := &context.InterruptSignal{
-			Type:      InterruptGraceful,
+			Type:      context.InterruptGraceful,
-			Messages:  []Message{{Role: RoleUser, Content: "peek test"}},
+			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{
+		signal := &context.InterruptSignal{
-			Type:      InterruptGraceful,
+			Type:      context.InterruptGraceful,
-			Messages:  []Message{{Role: RoleUser, Content: "check test"}},
+			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{
+			signal := &context.InterruptSignal{
-				Type: InterruptGraceful,
+				Type: context.InterruptGraceful,
-				Messages: []Message{
+				Messages: []context.Message{
-					{Role: RoleUser, Content: msg},
+					{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{
+		signal := &context.InterruptSignal{
-			Type:      InterruptGraceful,
+			Type:      context.InterruptGraceful,
-			Messages:  []Message{{Role: RoleUser, Content: "single signal"}},
+			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{
+		signal := &context.InterruptSignal{
-			Type:      InterruptForce,
+			Type:      context.InterruptForce,
-			Messages:  []Message{}, // Empty messages = pure cancellation
+			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{
+		signal := &context.InterruptSignal{
-			Type:      InterruptGraceful,
+			Type:      context.InterruptGraceful,
-			Messages:  []Message{{Role: RoleUser, Content: "graceful"}},
+			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{
+		signal := &context.InterruptSignal{
-			Type:      InterruptGraceful,
+			Type:      context.InterruptGraceful,
-			Messages:  []Message{{Role: RoleUser, Content: "direct send"}},
+			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
+		var nilController *context.InterruptController
-		signal := &InterruptSignal{
+		signal := &context.InterruptSignal{
-			Type:      InterruptGraceful,
+			Type:      context.InterruptGraceful,
-			Messages:  []Message{{Role: RoleUser, Content: "test"}},
+			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{
+			signal := &context.InterruptSignal{
-				Type:      InterruptGraceful,
+				Type:      context.InterruptGraceful,
-				Messages:  []Message{{Role: RoleUser, Content: fmt.Sprintf("msg %d", i)}},
+				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{
+		signal := &context.InterruptSignal{
-			Type:      InterruptGraceful,
+			Type:      context.InterruptGraceful,
-			Messages:  []Message{{Role: RoleUser, Content: "overflow"}},
+			Messages:  []context.Message{{Role: context.RoleUser, Content: "overflow"}},
 			Timestamp: time.Now().UnixMilli(),
 		}
 

agent/context/jsapi_output_test.go

@@ -1,42 +1,42 @@
-package context
+package context_test
 
 import (
 	"bytes"
-	"context"
+	stdContext "context"
 	"net/http"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	v8 "github.com/yaoapp/gou/runtime/v8"
-	"github.com/yaoapp/yao/agent/output/message"
+	"github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/config"
 	"github.com/yaoapp/yao/test"
 )
 
-// mockResponseWriter is a mock implementation of http.ResponseWriter for testing
+// testMockResponseWriter is a mock implementation of http.ResponseWriter for testing
-type mockResponseWriter struct {
+type testMockResponseWriter struct {
 	headers http.Header
 	buffer  *bytes.Buffer
 	status  int
 }
 
-func newMockResponseWriter() *mockResponseWriter {
+func newTestMockResponseWriter() *testMockResponseWriter {
-	return &mockResponseWriter{
+	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{
+	cxt := context.New(stdContext.Background(), nil, "test-chat-id")
-		ChatID:      "test-chat-id",
+	cxt.AssistantID = "test-assistant-id"
-		AssistantID: "test-assistant-id",
+	cxt.Accept = context.AcceptStandard
-		Context:     context.Background(),
+	cxt.Locale = "en"
-		Accept:      "standard",
+	cxt.Writer = newTestMockResponseWriter()
-		Locale:      "en",
-		Writer:      newMockResponseWriter(),
-		IDGenerator: message.NewIDGenerator(),
-	}
 
 	// 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{
+	cxt := context.New(stdContext.Background(), nil, "test-chat-id")
-		ChatID:      "test-chat-id",
+	cxt.AssistantID = "test-assistant-id"
-		AssistantID: "test-assistant-id",
+	cxt.Accept = context.AcceptStandard
-		Context:     context.Background(),
+	cxt.Locale = "en"
-		Accept:      "standard",
+	cxt.Writer = newTestMockResponseWriter()
-		Locale:      "en",
-		Writer:      newMockResponseWriter(),
-		IDGenerator: message.NewIDGenerator(),
-	}
 
 	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{
+	cxt := context.New(stdContext.Background(), nil, "test-chat-id")
-		ChatID:      "test-chat-id",
+	cxt.AssistantID = "test-assistant-id"
-		AssistantID: "test-assistant-id",
+	cxt.Accept = context.AcceptStandard
-		Context:     context.Background(),
+	cxt.Locale = "en"
-		Accept:      "standard",
+	cxt.Writer = newTestMockResponseWriter()
-		Locale:      "en",
-		Writer:      newMockResponseWriter(),
-		IDGenerator: message.NewIDGenerator(),
-	}
 
 	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{
+	cxt := context.New(stdContext.Background(), nil, "test-chat-id")
-		ChatID:      "test-chat-id",
+	cxt.AssistantID = "test-assistant-id"
-		AssistantID: "test-assistant-id",
+	cxt.Accept = context.AcceptStandard
-		Context:     context.Background(),
+	cxt.Locale = "en"
-		Accept:      "standard",
+	cxt.Writer = newTestMockResponseWriter()
-		Locale:      "en",
-		Writer:      newMockResponseWriter(),
-		IDGenerator: message.NewIDGenerator(),
-	}
 
 	// 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) {
+		t.Run(string(acceptType), func(t *testing.T) {
-			cxt := &Context{
+			cxt := context.New(stdContext.Background(), nil, "test-chat-id")
-				ChatID:      "test-chat-id",
+			cxt.AssistantID = "test-assistant-id"
-				AssistantID: "test-assistant-id",
+			cxt.Accept = acceptType
-				Context:     context.Background(),
+			cxt.Locale = "en"
-				Accept:      Accept(acceptType),
+			cxt.Writer = newTestMockResponseWriter()
-				Locale:      "en",
-				Writer:      newMockResponseWriter(),
-			}
 
 			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{
+	cxt := context.New(stdContext.Background(), nil, "test-chat-id")
-		ChatID:      "test-chat-id",
+	cxt.AssistantID = "test-assistant-id"
-		AssistantID: "test-assistant-id",
+	cxt.Accept = context.AcceptStandard
-		Context:     context.Background(),
+	cxt.Locale = "en"
-		Accept:      "standard",
+	cxt.Writer = newTestMockResponseWriter()
-		Locale:      "en",
-		Writer:      newMockResponseWriter(),
-		IDGenerator: message.NewIDGenerator(),
-	}
 
 	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{
+	cxt := context.New(stdContext.Background(), nil, "test-chat-id")
-		ChatID:      "test-chat-id",
+	cxt.AssistantID = "test-assistant-id"
-		AssistantID: "test-assistant-id",
+	cxt.Accept = context.AcceptStandard
-		Context:     context.Background(),
+	cxt.Locale = "en"
-		Accept:      "standard",
+	cxt.Writer = newTestMockResponseWriter()
-		Locale:      "en",
-		Writer:      newMockResponseWriter(),
-		IDGenerator: message.NewIDGenerator(),
-	}
 
 	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{
+	cxt := context.New(stdContext.Background(), nil, "test-chat-id")
-		ChatID:      "test-chat-id",
+	cxt.AssistantID = "test-assistant-id"
-		AssistantID: "test-assistant-id",
+	cxt.Accept = context.AcceptStandard
-		Context:     context.Background(),
+	cxt.Locale = "en"
-		Accept:      "standard",
+	cxt.Writer = newTestMockResponseWriter()
-		Locale:      "en",
-		Writer:      newMockResponseWriter(),
-		IDGenerator: message.NewIDGenerator(),
-	}
 
 	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{
+	cxt := context.New(stdContext.Background(), nil, "test-chat-id")
-		ChatID:      "test-chat-id",
+	cxt.AssistantID = "test-assistant-id"
-		AssistantID: "test-assistant-id",
+	cxt.Accept = context.AcceptStandard
-		Context:     context.Background(),
+	cxt.Locale = "en"
-		Accept:      "standard",
+	cxt.Writer = newTestMockResponseWriter()
-		Locale:      "en",
-		Writer:      newMockResponseWriter(),
-		IDGenerator: message.NewIDGenerator(),
-	}
 
 	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{
+	cxt := context.New(stdContext.Background(), nil, "test-chat-id")
-		ChatID:      "test-chat-id",
+	cxt.AssistantID = "test-assistant-id"
-		AssistantID: "test-assistant-id",
+	cxt.Accept = context.AcceptStandard
-		Context:     context.Background(),
+	cxt.Locale = "en"
-		Accept:      "standard",
+	cxt.Writer = newTestMockResponseWriter()
-		Locale:      "en",
-		Writer:      newMockResponseWriter(),
-		IDGenerator: message.NewIDGenerator(),
-	}
 
 	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{
+	cxt := context.New(stdContext.Background(), nil, "test-chat-id")
-		ChatID:      "test-chat-id",
+	cxt.AssistantID = "test-assistant-id"
-		AssistantID: "test-assistant-id",
+	cxt.Accept = context.AcceptStandard
-		Context:     context.Background(),
+	cxt.Locale = "en"
-		Accept:      "standard",
+	cxt.Writer = newTestMockResponseWriter()
-		Locale:      "en",
-		Writer:      newMockResponseWriter(),
-		IDGenerator: message.NewIDGenerator(),
-	}
 
 	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{
+	cxt := context.New(stdContext.Background(), nil, "test-chat-id")
-		ChatID:      "test-chat-id",
+	cxt.AssistantID = "test-assistant-id"
-		AssistantID: "test-assistant-id",
+	cxt.Accept = context.AcceptStandard
-		Context:     context.Background(),
+	cxt.Locale = "en"
-		Accept:      "standard",
+	cxt.Writer = newTestMockResponseWriter()
-		Locale:      "en",
-		Writer:      newMockResponseWriter(),
-		IDGenerator: message.NewIDGenerator(),
-	}
 
 	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{
+	cxt := context.New(stdContext.Background(), nil, "test-chat-id")
-		ChatID:      "test-chat-id",
+	cxt.AssistantID = "test-assistant-id"
-		AssistantID: "test-assistant-id",
+	cxt.Accept = context.AcceptStandard
-		Context:     context.Background(),
+	cxt.Locale = "en"
-		Accept:      "standard",
+	cxt.Writer = newTestMockResponseWriter()
-		Locale:      "en",
-		Writer:      newMockResponseWriter(),
-		IDGenerator: message.NewIDGenerator(),
-	}
 
 	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()
+		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
 
@@ -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()
+		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
 

agent/context/log.go

@@ -0,0 +1,495 @@
+package context
+
+import (
+	"fmt"
+	"strings"
+	"sync"
+	"time"
+
+	kunlog "github.com/yaoapp/kun/log"
+	"github.com/yaoapp/yao/config"
+)
+
+// =============================================================================
+// ANSI Color Codes
+// =============================================================================
+
+const (
+	colorReset   = "\033[0m"
+	colorRed     = "\033[31m"
+	colorGreen   = "\033[32m"
+	colorYellow  = "\033[33m"
+	colorBlue    = "\033[34m"
+	colorMagenta = "\033[35m"
+	colorCyan    = "\033[36m"
+	colorWhite   = "\033[37m"
+	colorGray    = "\033[90m"
+
+	colorBoldRed     = "\033[1;31m"
+	colorBoldGreen   = "\033[1;32m"
+	colorBoldYellow  = "\033[1;33m"
+	colorBoldBlue    = "\033[1;34m"
+	colorBoldMagenta = "\033[1;35m"
+	colorBoldCyan    = "\033[1;36m"
+)
+
+// =============================================================================
+// Log Level
+// =============================================================================
+
+// LogLevel represents log severity
+type LogLevel int
+
+const (
+	LogLevelTrace LogLevel = iota
+	LogLevelDebug
+	LogLevelInfo
+	LogLevelWarn
+	LogLevelError
+)
+
+// =============================================================================
+// Log Entry
+// =============================================================================
+
+// LogEntry represents a single log entry
+type LogEntry struct {
+	Level     LogLevel
+	Message   string
+	Timestamp time.Time
+	Phase     string // For phase logging
+	Elapsed   time.Duration
+}
+
+// =============================================================================
+// Request Logger
+// =============================================================================
+
+// RequestLogger provides request-scoped async logging
+type RequestLogger struct {
+	assistantID string
+	chatID      string
+	requestID   string
+	shortID     string // Short version of requestID for display
+	startTime   time.Time
+
+	ch     chan LogEntry
+	done   chan struct{}
+	once   sync.Once
+	closed bool
+	noop   bool // noop logger does nothing (for nil safety)
+	mu     sync.RWMutex
+}
+
+// noopLogger is a shared no-op logger instance
+var noopLogger = &RequestLogger{noop: true}
+
+// NewRequestLogger creates a new request-scoped logger with async processing
+func NewRequestLogger(assistantID, chatID, requestID string) *RequestLogger {
+	l := &RequestLogger{
+		assistantID: assistantID,
+		chatID:      chatID,
+		requestID:   requestID,
+		shortID:     shortID(requestID),
+		startTime:   time.Now(),
+		ch:          make(chan LogEntry, 100), // Buffered channel
+		done:        make(chan struct{}),
+	}
+
+	// Start consumer goroutine
+	go l.consume()
+
+	return l
+}
+
+// Noop returns a no-op logger that does nothing (nil-safe)
+func Noop() *RequestLogger {
+	return noopLogger
+}
+
+// SetAssistantID sets the assistant ID (called when entering Stream)
+func (l *RequestLogger) SetAssistantID(id string) {
+	if l.noop {
+		return
+	}
+	l.assistantID = id
+}
+
+// Close closes the logger and waits for all entries to be processed
+func (l *RequestLogger) Close() {
+	if l.noop {
+		return
+	}
+	l.once.Do(func() {
+		l.mu.Lock()
+		l.closed = true
+		l.mu.Unlock()
+
+		close(l.ch)
+		<-l.done // Wait for consumer to finish
+	})
+}
+
+// consume processes log entries from the channel
+func (l *RequestLogger) consume() {
+	defer close(l.done)
+
+	for entry := range l.ch {
+		l.processEntry(entry)
+	}
+}
+
+// processEntry handles a single log entry based on mode
+func (l *RequestLogger) processEntry(entry LogEntry) {
+	if config.IsDevelopment() {
+		l.printDev(entry)
+	} else {
+		l.printProd(entry)
+	}
+}
+
+// printDev prints colorful output for development mode
+func (l *RequestLogger) printDev(entry LogEntry) {
+	switch entry.Level {
+	case LogLevelTrace:
+		fmt.Printf("%s  → %s%s\n", colorGray, entry.Message, colorReset)
+	case LogLevelDebug:
+		fmt.Printf("%s  • %s%s\n", colorGray, entry.Message, colorReset)
+	case LogLevelInfo:
+		fmt.Printf("%s  ℹ %s%s\n", colorCyan, entry.Message, colorReset)
+	case LogLevelWarn:
+		fmt.Printf("%s  ⚠ %s%s\n", colorYellow, entry.Message, colorReset)
+	case LogLevelError:
+		fmt.Printf("%s  ✗ %s%s\n", colorRed, entry.Message, colorReset)
+	}
+}
+
+// printProd logs to kun/log for production mode
+func (l *RequestLogger) printProd(entry LogEntry) {
+	prefix := fmt.Sprintf("[AGENT] %s ", l.shortID)
+
+	switch entry.Level {
+	case LogLevelTrace:
+		kunlog.Trace("%s%s", prefix, entry.Message)
+	case LogLevelDebug:
+		// Skip debug in production
+	case LogLevelInfo:
+		kunlog.Info("%s%s", prefix, entry.Message)
+	case LogLevelWarn:
+		kunlog.Warn("%s%s", prefix, entry.Message)
+	case LogLevelError:
+		kunlog.Error("%s%s", prefix, entry.Message)
+	}
+}
+
+// send sends an entry to the channel (non-blocking if closed)
+func (l *RequestLogger) send(entry LogEntry) {
+	if l.noop {
+		return
+	}
+
+	l.mu.RLock()
+	closed := l.closed
+	l.mu.RUnlock()
+
+	if closed {
+		return
+	}
+
+	entry.Timestamp = time.Now()
+	select {
+	case l.ch <- entry:
+	default:
+		// Channel full, drop the log (shouldn't happen with buffered channel)
+	}
+}
+
+// =============================================================================
+// Standard Log Interface
+// =============================================================================
+
+// Trace logs a trace level message
+func (l *RequestLogger) Trace(format string, args ...interface{}) {
+	l.send(LogEntry{
+		Level:   LogLevelTrace,
+		Message: fmt.Sprintf(format, args...),
+	})
+}
+
+// Debug logs a debug level message
+func (l *RequestLogger) Debug(format string, args ...interface{}) {
+	l.send(LogEntry{
+		Level:   LogLevelDebug,
+		Message: fmt.Sprintf(format, args...),
+	})
+}
+
+// Info logs an info level message
+func (l *RequestLogger) Info(format string, args ...interface{}) {
+	l.send(LogEntry{
+		Level:   LogLevelInfo,
+		Message: fmt.Sprintf(format, args...),
+	})
+}
+
+// Warn logs a warning level message
+func (l *RequestLogger) Warn(format string, args ...interface{}) {
+	l.send(LogEntry{
+		Level:   LogLevelWarn,
+		Message: fmt.Sprintf(format, args...),
+	})
+}
+
+// Error logs an error level message
+func (l *RequestLogger) Error(format string, args ...interface{}) {
+	l.send(LogEntry{
+		Level:   LogLevelError,
+		Message: fmt.Sprintf(format, args...),
+	})
+}
+
+// =============================================================================
+// Business Quick Functions
+// =============================================================================
+
+// Start logs the start of a request with visual separator
+func (l *RequestLogger) Start() {
+	if l.noop {
+		return
+	}
+
+	if !config.IsDevelopment() {
+		kunlog.Trace("[AGENT] Request %s started: assistant=%s, chat=%s, request=%s",
+			l.shortID, l.assistantID, shortID(l.chatID), shortID(l.requestID))
+		return
+	}
+
+	// Development: colorful output (direct print, not through channel for immediate display)
+	fmt.Println()
+	fmt.Printf("%s%s%s\n", colorBoldCyan, strings.Repeat("═", 60), colorReset)
+	fmt.Printf("%s  🚀 AGENT REQUEST %s%s\n", colorBoldCyan, l.shortID, colorReset)
+	fmt.Printf("%s%s%s\n", colorBoldCyan, strings.Repeat("─", 60), colorReset)
+	fmt.Printf("%s  Assistant: %s%s%s\n", colorGray, colorWhite, l.assistantID, colorReset)
+	fmt.Printf("%s  Chat ID:   %s%s%s\n", colorGray, colorWhite, l.chatID, colorReset)
+	fmt.Printf("%s  Request:   %s%s%s\n", colorGray, colorWhite, l.requestID, colorReset)
+	fmt.Printf("%s  Time:      %s%s%s\n", colorGray, colorWhite, l.startTime.Format("15:04:05.000"), colorReset)
+	fmt.Printf("%s%s%s\n", colorCyan, strings.Repeat("─", 60), colorReset)
+}
+
+// End logs the end of a request with summary
+func (l *RequestLogger) End(success bool, err error) {
+	if l.noop {
+		return
+	}
+
+	duration := time.Since(l.startTime)
+
+	if !config.IsDevelopment() {
+		if success {
+			kunlog.Trace("[AGENT] Request %s completed: assistant=%s, duration=%v",
+				l.shortID, l.assistantID, duration.Round(time.Millisecond))
+		} else {
+			kunlog.Trace("[AGENT] Request %s failed: assistant=%s, duration=%v, error=%v",
+				l.shortID, l.assistantID, duration.Round(time.Millisecond), err)
+		}
+		return
+	}
+
+	// Development: colorful output (direct print for immediate display)
+	fmt.Printf("%s%s%s\n", colorCyan, strings.Repeat("─", 60), colorReset)
+	if success {
+		fmt.Printf("%s  ✅ REQUEST %s COMPLETED%s\n", colorBoldGreen, l.shortID, colorReset)
+	} else {
+		fmt.Printf("%s  ❌ REQUEST %s FAILED%s\n", colorBoldRed, l.shortID, colorReset)
+		if err != nil {
+			fmt.Printf("%s  Error: %s%v%s\n", colorGray, colorRed, err, colorReset)
+		}
+	}
+	fmt.Printf("%s  Assistant: %s%s%s\n", colorGray, colorWhite, l.assistantID, colorReset)
+	fmt.Printf("%s  Duration:  %s%v%s\n", colorGray, colorWhite, duration.Round(time.Millisecond), colorReset)
+	fmt.Printf("%s%s%s\n", colorCyan, strings.Repeat("─", 60), colorReset)
+	fmt.Println()
+}
+
+// Phase logs a major phase in the request lifecycle
+func (l *RequestLogger) Phase(name string) {
+	if l.noop {
+		return
+	}
+
+	elapsed := time.Since(l.startTime).Round(time.Millisecond)
+
+	if config.IsDevelopment() {
+		fmt.Printf("%s  ▶ %s%s %s[+%v]%s\n", colorBoldBlue, name, colorReset, colorGray, elapsed, colorReset)
+	} else {
+		kunlog.Trace("[AGENT] %s Phase: %s (+%v)", l.shortID, name, elapsed)
+	}
+}
+
+// PhaseComplete logs the completion of a phase
+func (l *RequestLogger) PhaseComplete(name string) {
+	if l.noop {
+		return
+	}
+
+	elapsed := time.Since(l.startTime).Round(time.Millisecond)
+
+	if config.IsDevelopment() {
+		fmt.Printf("%s  ✓ %s%s %s[+%v]%s\n", colorGreen, name, colorReset, colorGray, elapsed, colorReset)
+	} else {
+		kunlog.Trace("[AGENT] %s Phase completed: %s (+%v)", l.shortID, name, elapsed)
+	}
+}
+
+// PhaseSkip logs a skipped phase (development only)
+func (l *RequestLogger) PhaseSkip(name, reason string) {
+	if l.noop {
+		return
+	}
+	if config.IsDevelopment() {
+		fmt.Printf("%s  ⊘ %s (%s)%s\n", colorGray, name, reason, colorReset)
+	}
+}
+
+// LLMStart logs the start of an LLM call
+func (l *RequestLogger) LLMStart(connector, model string, messageCount int) {
+	if l.noop {
+		return
+	}
+
+	elapsed := time.Since(l.startTime).Round(time.Millisecond)
+
+	if config.IsDevelopment() {
+		fmt.Printf("%s  🤖 LLM Call%s %s[+%v]%s\n", colorBoldMagenta, colorReset, colorGray, elapsed, colorReset)
+		fmt.Printf("%s    Connector: %s%s%s\n", colorGray, colorWhite, connector, colorReset)
+		if model != "" {
+			fmt.Printf("%s    Model: %s%s%s\n", colorGray, colorWhite, model, colorReset)
+		}
+		fmt.Printf("%s    Messages: %s%d%s\n", colorGray, colorWhite, messageCount, colorReset)
+	} else {
+		kunlog.Trace("[AGENT] %s LLM call: connector=%s, model=%s, messages=%d (+%v)", l.shortID, connector, model, messageCount, elapsed)
+	}
+}
+
+// LLMComplete logs the completion of an LLM call
+func (l *RequestLogger) LLMComplete(tokens int, hasToolCalls bool) {
+	elapsed := time.Since(l.startTime).Round(time.Millisecond)
+	status := "streaming"
+	if hasToolCalls {
+		status = "tool_calls"
+	}
+
+	if config.IsDevelopment() {
+		fmt.Printf("%s  ✓ LLM Response (%s)%s", colorGreen, status, colorReset)
+		if tokens > 0 {
+			fmt.Printf(" %s[tokens: %d]%s", colorGray, tokens, colorReset)
+		}
+		fmt.Printf(" %s[+%v]%s\n", colorGray, elapsed, colorReset)
+	} else {
+		kunlog.Trace("[AGENT] %s LLM response: status=%s, tokens=%d (+%v)", l.shortID, status, tokens, elapsed)
+	}
+}
+
+// ToolStart logs the start of tool execution
+func (l *RequestLogger) ToolStart(toolName string) {
+	if config.IsDevelopment() {
+		fmt.Printf("%s  🔧 Tool: %s%s\n", colorYellow, toolName, colorReset)
+	} else {
+		kunlog.Trace("[AGENT] %s Tool call: %s", l.shortID, toolName)
+	}
+}
+
+// ToolComplete logs the completion of tool execution
+func (l *RequestLogger) ToolComplete(toolName string, success bool) {
+	if config.IsDevelopment() {
+		if success {
+			fmt.Printf("%s    ✓ %s completed%s\n", colorGreen, toolName, colorReset)
+		} else {
+			fmt.Printf("%s    ✗ %s failed%s\n", colorRed, toolName, colorReset)
+		}
+	} else {
+		if success {
+			kunlog.Trace("[AGENT] %s Tool completed: %s", l.shortID, toolName)
+		} else {
+			kunlog.Trace("[AGENT] %s Tool failed: %s", l.shortID, toolName)
+		}
+	}
+}
+
+// HookStart logs the start of a hook execution
+func (l *RequestLogger) HookStart(hookName string) {
+	elapsed := time.Since(l.startTime).Round(time.Millisecond)
+
+	if config.IsDevelopment() {
+		fmt.Printf("%s  🪝 Hook: %s%s %s[+%v]%s\n", colorMagenta, hookName, colorReset, colorGray, elapsed, colorReset)
+	} else {
+		kunlog.Trace("[AGENT] %s Hook: %s (+%v)", l.shortID, hookName, elapsed)
+	}
+}
+
+// HookComplete logs the completion of a hook
+func (l *RequestLogger) HookComplete(hookName string) {
+	if config.IsDevelopment() {
+		fmt.Printf("%s    ✓ %s done%s\n", colorGreen, hookName, colorReset)
+	} else {
+		kunlog.Trace("[AGENT] %s Hook completed: %s", l.shortID, hookName)
+	}
+}
+
+// Cleanup logs resource cleanup
+func (l *RequestLogger) Cleanup(resource string) {
+	if l.noop {
+		return
+	}
+
+	if config.IsDevelopment() {
+		fmt.Printf("%s    ✓ %s%s\n", colorGray, resource, colorReset)
+	} else {
+		kunlog.Trace("[AGENT] %s Cleanup: %s", l.shortID, resource)
+	}
+}
+
+// HistoryLoad logs history loading
+func (l *RequestLogger) HistoryLoad(count, maxSize int) {
+	if config.IsDevelopment() {
+		fmt.Printf("%s    Loaded %d/%d history messages%s\n", colorGray, count, maxSize, colorReset)
+	} else {
+		kunlog.Trace("[AGENT] %s History loaded: %d/%d messages", l.shortID, count, maxSize)
+	}
+}
+
+// HistoryOverlap logs overlap detection
+func (l *RequestLogger) HistoryOverlap(overlapCount int) {
+	if overlapCount > 0 {
+		if config.IsDevelopment() {
+			fmt.Printf("%s    Removed %d overlapping messages%s\n", colorYellow, overlapCount, colorReset)
+		} else {
+			kunlog.Trace("[AGENT] %s History overlap removed: %d messages", l.shortID, overlapCount)
+		}
+	}
+}
+
+// Release logs the start of resource release phase
+func (l *RequestLogger) Release() {
+	if l.noop {
+		return
+	}
+
+	if config.IsDevelopment() {
+		fmt.Printf("%s  🧹 RELEASE %s%s %s(%s)%s\n", colorBoldYellow, l.shortID, colorReset, colorGray, l.assistantID, colorReset)
+	} else {
+		kunlog.Trace("[AGENT] %s Release started", l.shortID)
+	}
+}
+
+// =============================================================================
+// Helper
+// =============================================================================
+
+// shortID returns first 8 characters of an ID
+func shortID(id string) string {
+	if len(id) > 8 {
+		return id[:8]
+	}
+	return id
+}

agent/context/message_test.go

@@ -1,8 +1,10 @@
-package context
+package context_test
 
 import (
 	"encoding/json"
 	"testing"
+
+	"github.com/yaoapp/yao/agent/context"
 )
 
 func TestMessage_UnmarshalJSON_StringContent(t *testing.T) {
@@ -11,14 +13,14 @@ func TestMessage_UnmarshalJSON_StringContent(t *testing.T) {
 		"content": "Hello, world!"
 	}`
 
-	var msg Message
+	var msg context.Message
 	err := json.Unmarshal([]byte(jsonData), &msg)
 	if err != nil {
 		t.Fatalf("Failed to unmarshal: %v", err)
 	}
 
-	if msg.Role != RoleUser {
+	if msg.Role != context.RoleUser {
-		t.Errorf("Expected role %s, got %s", RoleUser, msg.Role)
+		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 {
+	if msg.Role != context.RoleUser {
-		t.Errorf("Expected role %s, got %s", RoleUser, msg.Role)
+		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 {
+	if parts[0].Type != context.ContentText {
-		t.Errorf("Expected type %s, got %s", ContentText, parts[0].Type)
+		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 {
+	if parts[1].Type != context.ContentImageURL {
-		t.Errorf("Expected type %s, got %s", ContentImageURL, parts[1].Type)
+		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 {
+	if parts[1].ImageURL.Detail != context.DetailHigh {
-		t.Errorf("Expected detail %s, got %s", DetailHigh, parts[1].ImageURL.Detail)
+		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 {
+	if msg.Role != context.RoleAssistant {
-		t.Errorf("Expected role %s, got %s", RoleAssistant, msg.Role)
+		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 {
+	if parts[1].Type != context.ContentInputAudio {
-		t.Errorf("Expected type %s, got %s", ContentInputAudio, parts[1].Type)
+		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) {
+	if result["role"] != string(context.RoleUser) {
-		t.Errorf("Expected role %s, got %v", RoleUser, result["role"])
+		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,
+			Type: context.ContentImageURL,
-			ImageURL: &ImageURL{
+			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{
+	msg := &context.Message{
-		Role:    RoleAssistant,
+		Role:    context.RoleAssistant,
 		Content: nil,
-		ToolCalls: []ToolCall{
+		ToolCalls: []context.ToolCall{
 			{
 				ID:   "call_abc123",
-				Type: ToolTypeFunction,
+				Type: context.ToolTypeFunction,
-				Function: Function{
+				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 {
+	if msg.Role != context.RoleTool {
-		t.Errorf("Expected role %s, got %s", RoleTool, msg.Role)
+		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 {
+	if msg.Role != context.RoleSystem {
-		t.Errorf("Expected role %s, got %s", RoleSystem, msg.Role)
+		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{
+	parts := []context.ContentPart{
-		{Type: ContentText, Text: "Hello"},
+		{Type: context.ContentText, Text: "Hello"},
 	}
 
-	msg := NewMultipartMessage(RoleUser, parts)
+	msg := context.NewMultipartMessage(context.RoleUser, parts)
 
-	if msg.Role != RoleUser {
+	if msg.Role != context.RoleUser {
-		t.Errorf("Expected role %s, got %s", RoleUser, msg.Role)
+		t.Errorf("Expected role %s, got %s", context.RoleUser, msg.Role)
 	}
 
 	resultParts, ok := msg.GetContentAsParts()

agent/context/openapi_test.go

@@ -1,30 +1,54 @@
-package context
+package context_test
 
 import (
 	"bytes"
 	"encoding/json"
+	"io"
 	"net/http/httptest"
 	"testing"
 
 	"github.com/gin-gonic/gin"
 	"github.com/yaoapp/gou/store"
+	"github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/config"
 	"github.com/yaoapp/yao/test"
 )
 
+// parseCompletionRequestData is a helper function for tests to parse completion request data
+func parseCompletionRequestData(c *gin.Context) (*context.CompletionRequest, error) {
+	var req context.CompletionRequest
+
+	if c.Request.Body != nil {
+		body, err := io.ReadAll(c.Request.Body)
+		if err != nil {
+			return nil, err
+		}
+		c.Request.Body = io.NopCloser(bytes.NewBuffer(body))
+		if len(body) > 0 {
+			if err := json.Unmarshal(body, &req); err != nil {
+				return nil, err
+			}
+			if len(req.Messages) > 0 {
+				return &req, nil
+			}
+		}
+	}
+	return &req, nil
+}
+
 func TestGetMessages_FromBody(t *testing.T) {
 	test.Prepare(t, config.Conf)
 	defer test.Clean()
 
 	gin.SetMode(gin.TestMode)
 
-	messages := []Message{
+	messages := []context.Message{
 		{
-			Role:    RoleUser,
+			Role:    context.RoleUser,
 			Content: "Hello, world!",
 		},
 		{
-			Role:    RoleAssistant,
+			Role:    context.RoleAssistant,
 			Content: "Hi there!",
 		},
 	}
@@ -45,7 +69,7 @@ func TestGetMessages_FromBody(t *testing.T) {
 	// Parse request first
 	completionReq, _ := parseCompletionRequestData(c)
 
-	result, err := GetMessages(c, completionReq)
+	result, err := context.GetMessages(c, completionReq)
 	if err != nil {
 		t.Fatalf("Failed to get messages: %v", err)
 	}
@@ -54,8 +78,8 @@ func TestGetMessages_FromBody(t *testing.T) {
 		t.Errorf("Expected 2 messages, got %d", len(result))
 	}
 
-	if result[0].Role != RoleUser {
+	if result[0].Role != context.RoleUser {
-		t.Errorf("Expected first message role to be %s, got %s", RoleUser, result[0].Role)
+		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)
+	referer := context.GetReferer(c, completionReq)
-	if referer != RefererTool {
+	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)
+	accept := context.GetAccept(c, nil)
-	if accept != AcceptWebCUI {
+	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)
+	accept := context.GetAccept(c, nil)
-	if accept != AcceptDesktopCUI {
+	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)
+	accept := context.GetAccept(c, completionReq)
-	if accept != AccepNativeCUI {
+	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)
+	accept := context.GetAccept(c, nil)
-	if accept != AcceptStandard {
+	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)
+	accept := context.GetAccept(c, completionReq)
-	if accept != AcceptWebCUI {
+	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{
+	completionReq := &context.CompletionRequest{
-		Skip: &Skip{
+		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{
+	completionReq := &context.CompletionRequest{
-		Skip: &Skip{
+		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)
 	}

agent/context/stack_test.go

@@ -1,10 +1,11 @@
-package context
+package context_test
 
 import (
+	stdContext "context"
 	"testing"
 	"time"
 
-	"github.com/yaoapp/yao/agent/output/message"
+	"github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/config"
 	"github.com/yaoapp/yao/test"
 )
@@ -15,10 +16,10 @@ func TestNewStack(t *testing.T) {
 
 	traceID := "12345678"
 	assistantID := "test-assistant"
-	referer := RefererAPI
+	referer := context.RefererAPI
-	opts := &Options{}
+	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 {
+	if stack.Status != context.StackStatusRunning {
-		t.Errorf("Expected Status '%s', got '%s'", StackStatusRunning, stack.Status)
+		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 {
+	if stack.Status != context.StackStatusCompleted {
-		t.Errorf("Expected Status '%s', got '%s'", StackStatusCompleted, stack.Status)
+		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 {
+	if stack.Status != context.StackStatusFailed {
-		t.Errorf("Expected Status '%s', got '%s'", StackStatusFailed, stack.Status)
+		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 {
+	if stack.Status != context.StackStatusTimeout {
-		t.Errorf("Expected Status '%s', got '%s'", StackStatusTimeout, stack.Status)
+		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{
+	ctx := context.New(stdContext.Background(), nil, "test-chat-id")
-		IDGenerator: message.NewIDGenerator(),
+	ctx.Referer = context.RefererAPI
-		Referer:     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{
+	ctx := context.New(stdContext.Background(), nil, "test-chat-id")
-		IDGenerator: message.NewIDGenerator(),
+	ctx.Referer = context.RefererAPI
-		Referer:     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{
+	ctx := context.New(stdContext.Background(), nil, "test-chat-id")
-		IDGenerator: message.NewIDGenerator(),
+	ctx.Referer = context.RefererAPI
-		Referer:     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{
+	ctx := context.New(stdContext.Background(), nil, "test-chat-id")
-		IDGenerator: message.NewIDGenerator(),
+	ctx.Referer = context.RefererAPI
-		Referer:     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{
+	ctx := context.New(stdContext.Background(), nil, "test-chat-id")
-		IDGenerator: message.NewIDGenerator(),
+	ctx.Referer = context.RefererAPI
-		Referer:     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{
+	ctx := context.New(stdContext.Background(), nil, "test-chat-id")
-		IDGenerator: message.NewIDGenerator(),
+	ctx.Referer = context.RefererAPI
-		Referer:     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{
+	ctx := context.New(stdContext.Background(), nil, "test-chat-id")
-		IDGenerator: message.NewIDGenerator(),
+	ctx.Referer = context.RefererAPI
-		Referer:     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()

agent/context/types.go

@@ -230,6 +230,7 @@ type Context struct {
 	Stacks      map[string]*Stack    `json:"-"`  // Stacks, all stacks in this request (for trace logging)
 	Writer      Writer               `json:"-"`  // Writer, it will be used to write response data to the client
 	IDGenerator *message.IDGenerator `json:"-"`  // ID generator for this context (chunk, message, block, thread IDs)
+	Logger      *RequestLogger       `json:"-"`  // Request-scoped async logger
 
 	// Chat buffer for batch saving messages and resume steps
 	Buffer *ChatBuffer `json:"-"` // Chat buffer for batch saving at end of Stream()

agent/llm/providers/openai/claude_test.go

@@ -6,7 +6,6 @@ import (
 
 	"github.com/yaoapp/gou/connector"
 	"github.com/yaoapp/gou/connector/openai"
-	"github.com/yaoapp/gou/plan"
 	"github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/agent/llm"
 	"github.com/yaoapp/yao/agent/output/message"
@@ -17,23 +16,7 @@ import (
 
 // newClaudeTestContext creates a real Context for testing Claude provider
 func newClaudeTestContext(chatID, connectorID string) *context.Context {
-	return &context.Context{
+	authorized := &types.AuthorizedInfo{
-		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",
@@ -46,8 +29,22 @@ func newClaudeTestContext(chatID, connectorID string) *context.Context {
 				"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

agent/llm/providers/openai/deepseek_r1_test.go

@@ -8,7 +8,6 @@ import (
 	jsoniter "github.com/json-iterator/go"
 	"github.com/yaoapp/gou/connector"
 	"github.com/yaoapp/gou/connector/openai"
-	"github.com/yaoapp/gou/plan"
 	"github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/agent/llm"
 	"github.com/yaoapp/yao/agent/output/message"
@@ -396,23 +395,7 @@ func TestDeepSeekR1LogicPuzzle(t *testing.T) {
 
 // newDeepSeekTestContext creates a real Context for testing DeepSeek provider
 func newDeepSeekTestContext(chatID, connectorID string) *context.Context {
-	return &context.Context{
+	authorized := &types.AuthorizedInfo{
-		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",
@@ -425,6 +408,20 @@ func newDeepSeekTestContext(chatID, connectorID string) *context.Context {
 				"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
 }

agent/llm/providers/openai/deepseek_v3_test.go

@@ -6,7 +6,6 @@ import (
 
 	"github.com/yaoapp/gou/connector"
 	"github.com/yaoapp/gou/connector/openai"
-	"github.com/yaoapp/gou/plan"
 	"github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/agent/llm"
 	"github.com/yaoapp/yao/agent/output/message"
@@ -368,23 +367,7 @@ 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{
+	authorized := &types.AuthorizedInfo{
-		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",
@@ -397,6 +380,20 @@ func newDeepSeekV3TestContext(chatID, connectorID string) *context.Context {
 				"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
 }

agent/llm/providers/openai/gpt5_test.go

@@ -6,7 +6,6 @@ import (
 
 	"github.com/yaoapp/gou/connector"
 	"github.com/yaoapp/gou/connector/openai"
-	"github.com/yaoapp/gou/plan"
 	"github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/agent/llm"
 	"github.com/yaoapp/yao/agent/output/message"
@@ -383,23 +382,7 @@ 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{
+	authorized := &types.AuthorizedInfo{
-		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",
@@ -412,6 +395,20 @@ func newGPT5TestContext(chatID, connectorID string) *context.Context {
 				"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
 }

agent/llm/providers/openai/openai_test.go

@@ -9,7 +9,6 @@ import (
 
 	"github.com/yaoapp/gou/connector"
 	"github.com/yaoapp/gou/connector/openai"
-	"github.com/yaoapp/gou/plan"
 	"github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/agent/llm"
 	"github.com/yaoapp/yao/agent/output/message"
@@ -1503,23 +1502,7 @@ func TestOpenAIStreamWithTemperature(t *testing.T) {
 
 // newTestContext creates a real Context for testing OpenAI provider
 func newTestContext(chatID, connectorID string) *context.Context {
-	return &context.Context{
+	authorized := &types.AuthorizedInfo{
-		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",
@@ -1532,6 +1515,20 @@ func newTestContext(chatID, connectorID string) *context.Context {
 				"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
 }

agent/llm/providers/openai/temperature_test.go

@@ -6,7 +6,6 @@ import (
 
 	"github.com/yaoapp/gou/connector"
 	"github.com/yaoapp/gou/connector/openai"
-	"github.com/yaoapp/gou/plan"
 	"github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/agent/llm"
 	"github.com/yaoapp/yao/config"
@@ -335,23 +334,7 @@ func TestTemperatureNoTemperatureProvided(t *testing.T) {
 
 // newTemperatureTestContext creates a real Context for testing temperature handling
 func newTemperatureTestContext(chatID, connectorID string) *context.Context {
-	return &context.Context{
+	authorized := &types.AuthorizedInfo{
-		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",
@@ -364,6 +347,20 @@ func newTemperatureTestContext(chatID, connectorID string) *context.Context {
 				"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
 }

agent/load.go

@@ -195,6 +195,9 @@ func initAssistant() error {
 	// Set Storage
 	assistant.SetStorage(agentDSL.Store)
 
+	// Set Store Setting (MaxSize, TTL, etc.)
+	assistant.SetStoreSetting(&agentDSL.StoreSetting)
+
 	// Set global Uses configuration
 	if agentDSL.Uses != nil {
 		globalUses := &context.Uses{

config/config.go

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