Merge pull request #1294 from trheyi/main

Refactor assistant

agent/assistant/agent.go

@@ -0,0 +1,33 @@
+package assistant
+
+import "github.com/yaoapp/yao/agent/context"
+
+// Stream stream the agent
+func (ast *Assistant) Stream(ctx *context.Context, messages []context.Message, handler context.StreamFunc) error {
+
+	// Initialize stack and auto-handle completion/failure/restore
+	_, traceID, done := context.EnterStack(ctx, ast.ID, ctx.Referer)
+	defer done()
+
+	_ = traceID // traceID is available for trace logging
+
+	// Request Create hook ( Optional )
+
+	// LLM Call Stream ( Optional )
+
+	// Request Done hook ( Optional )
+
+	return nil
+}
+
+// Run run the agent
+func (ast *Assistant) Run(ctx *context.Context, messages []context.Message) (*context.Response, error) {
+
+	// Initialize stack and auto-handle completion/failure/restore
+	_, traceID, done := context.EnterStack(ctx, ast.ID, ctx.Referer)
+	defer done()
+
+	_ = traceID // traceID is available for trace logging
+
+	return &context.Response{}, nil
+}

agent/assistant/hooks/hooks.go

@@ -0,0 +1 @@
+package hooks

agent/assistant/types.go

@@ -20,10 +20,6 @@ const (
 // API the assistant API interface
 type API interface {
 	Chat(ctx context.Context, messages []message.Message, option map[string]interface{}, cb func(data []byte) int) error
-	// Upload(ctx context.Context, file *multipart.FileHeader, reader io.Reader, option map[string]interface{}) (*File, error)
-	// Download(ctx context.Context, fileID string) (*FileResponse, error)
-	// ReadBase64(ctx context.Context, fileID string) (string, error)
-
 	GetPlaceholder(locale string) *store.Placeholder
 	Execute(c *gin.Context, ctx chatctx.Context, input interface{}, options map[string]interface{}, callback ...interface{}) (interface{}, error)
 	Call(c *gin.Context, payload APIPayload) (interface{}, error)

agent/context/context.go

@@ -62,10 +62,28 @@ func WithTimeout(parent Context, timeout time.Duration) (Context, context.Cancel
 	return parent, cancel
 }
 
-// Release the context
+// Release the context and clean up all resources including stacks
 func (ctx *Context) Release() {
-	ctx.Space.Clear()
+	// Clear space
-	ctx.Space = nil
+	if ctx.Space != nil {
+		ctx.Space.Clear()
+		ctx.Space = nil
+	}
+
+	// Clear stacks
+	if ctx.Stacks != nil {
+		for k := range ctx.Stacks {
+			delete(ctx.Stacks, k)
+		}
+		ctx.Stacks = nil
+	}
+
+	// Clear current stack reference
+	ctx.Stack = nil
+
+	// Clear writer reference
+	ctx.Writer = nil
+
 	ctx = nil
 }
 

agent/context/interfaces.go

@@ -0,0 +1,20 @@
+package context
+
+import "net/http"
+
+// StreamFunc the streaming function
+type StreamFunc func(data []byte) int
+
+// Writer is an alias for http.ResponseWriter interface used by an agent to construct a response.
+// A Writer may not be used after the agent execution has completed.
+type Writer = http.ResponseWriter
+
+// Agent the agent interface
+type Agent interface {
+
+	// Stream stream the agent
+	Stream(ctx *Context, messages []Message, handler StreamFunc) error
+
+	// Run run the agent
+	Run(ctx *Context, messages []Message) (*Response, error)
+}

agent/context/openapi.go

@@ -35,8 +35,8 @@ func GetCompletionRequest(c *gin.Context, cache store.Store) (*CompletionRequest
 	// Extract chat ID (may generate from messages if not provided)
 	chatID, err := GetChatID(c, cache, completionReq)
 	if err != nil {
-		// If chat ID generation fails, it's not critical - allow empty chatID
+		// Fallback: Generate a new chat ID if extraction fails
-		chatID = ""
+		chatID = GenChatID()
 	}
 
 	// Parse client information from User-Agent header
@@ -49,6 +49,7 @@ func GetCompletionRequest(c *gin.Context, cache store.Store) (*CompletionRequest
 		Context:     c.Request.Context(),
 		Space:       plan.NewMemorySharedSpace(),
 		Cache:       cache,
+		Writer:      c.Writer,
 		Authorized:  authInfo,
 		ChatID:      chatID,
 		AssistantID: assistantID,

agent/context/openapi_test.go

@@ -755,3 +755,114 @@ func TestGetData_EmptyData(t *testing.T) {
 		t.Errorf("Expected nil data, got '%v'", result)
 	}
 }
+
+func TestGetCompletionRequest_WriterInitialized(t *testing.T) {
+	test.Prepare(t, config.Conf)
+	defer test.Clean()
+
+	gin.SetMode(gin.TestMode)
+
+	cache, err := store.Get("__yao.agent.cache")
+	if err != nil {
+		t.Fatalf("Failed to get cache: %v", err)
+	}
+
+	messages := []Message{
+		{
+			Role:    RoleUser,
+			Content: "Test message",
+		},
+	}
+
+	requestBody := map[string]interface{}{
+		"model":    "gpt-4-yao_test",
+		"messages": messages,
+	}
+
+	bodyBytes, _ := json.Marshal(requestBody)
+
+	req := httptest.NewRequest("POST", "/chat/completions", bytes.NewBuffer(bodyBytes))
+	req.Header.Set("Content-Type", "application/json")
+	w := httptest.NewRecorder()
+	c, _ := gin.CreateTestContext(w)
+	c.Request = req
+
+	completionReq, ctx, err := GetCompletionRequest(c, cache)
+	if err != nil {
+		t.Fatalf("Failed to get completion request: %v", err)
+	}
+	defer ctx.Release()
+
+	// Check that Writer is initialized
+	if ctx.Writer == nil {
+		t.Error("Expected ctx.Writer to be initialized, got nil")
+	}
+
+	// Check that Writer is the same as gin context writer
+	if ctx.Writer != c.Writer {
+		t.Error("Expected ctx.Writer to be the same as gin context writer")
+	}
+
+	// Check other fields
+	if completionReq.Model != "gpt-4-yao_test" {
+		t.Errorf("Expected model 'gpt-4-yao_test', got '%s'", completionReq.Model)
+	}
+
+	if ctx.AssistantID != "test" {
+		t.Errorf("Expected assistant ID 'test', got '%s'", ctx.AssistantID)
+	}
+
+	// Check that ChatID was generated (fallback)
+	if ctx.ChatID == "" {
+		t.Error("Expected ChatID to be generated, got empty string")
+	}
+}
+
+func TestGetCompletionRequest_ChatIDFallback(t *testing.T) {
+	test.Prepare(t, config.Conf)
+	defer test.Clean()
+
+	gin.SetMode(gin.TestMode)
+
+	cache, err := store.Get("__yao.agent.cache")
+	if err != nil {
+		t.Fatalf("Failed to get cache: %v", err)
+	}
+
+	// Request without explicit chat_id should generate one
+	messages := []Message{
+		{
+			Role:    RoleUser,
+			Content: "Test message",
+		},
+	}
+
+	requestBody := map[string]interface{}{
+		"model":    "gpt-4-yao_assistant1",
+		"messages": messages,
+	}
+
+	bodyBytes, _ := json.Marshal(requestBody)
+
+	req := httptest.NewRequest("POST", "/chat/completions", bytes.NewBuffer(bodyBytes))
+	req.Header.Set("Content-Type", "application/json")
+	w := httptest.NewRecorder()
+	c, _ := gin.CreateTestContext(w)
+	c.Request = req
+
+	_, ctx, err := GetCompletionRequest(c, cache)
+	if err != nil {
+		t.Fatalf("Failed to get completion request: %v", err)
+	}
+	defer ctx.Release()
+
+	// ChatID should be generated (not empty)
+	if ctx.ChatID == "" {
+		t.Error("Expected ChatID to be generated via fallback, got empty string")
+	}
+
+	// ChatID should be a valid NanoID format (16 characters)
+	if len(ctx.ChatID) < 8 {
+		t.Errorf("Expected ChatID to be at least 8 characters, got %d", len(ctx.ChatID))
+	}
+}

agent/context/stack.go

@@ -0,0 +1,271 @@
+package context
+
+import (
+	"fmt"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/yaoapp/yao/agent/trace"
+)
+
+// NewStack creates a new root stack with the given trace ID and assistant ID
+func NewStack(traceID, assistantID, referer string) *Stack {
+	if traceID == "" {
+		traceID = uuid.New().String()
+	}
+
+	stackID := uuid.New().String()
+	now := time.Now().UnixMilli()
+
+	return &Stack{
+		ID:          stackID,
+		TraceID:     traceID,
+		AssistantID: assistantID,
+		Referer:     referer,
+		Depth:       0,
+		ParentID:    "",
+		Path:        []string{stackID},
+		CreatedAt:   now,
+		Status:      StackStatusRunning,
+	}
+}
+
+// NewChildStack creates a child stack from the current stack
+func (s *Stack) NewChildStack(assistantID, referer string) *Stack {
+	stackID := uuid.New().String()
+	now := time.Now().UnixMilli()
+
+	// Build path by appending current stack's path with new ID
+	path := make([]string, len(s.Path)+1)
+	copy(path, s.Path)
+	path[len(s.Path)] = stackID
+
+	return &Stack{
+		ID:          stackID,
+		TraceID:     s.TraceID, // Inherit trace ID
+		AssistantID: assistantID,
+		Referer:     referer,
+		Depth:       s.Depth + 1,
+		ParentID:    s.ID,
+		Path:        path,
+		CreatedAt:   now,
+		Status:      StackStatusRunning,
+	}
+}
+
+// Complete marks the stack as completed and calculates duration
+func (s *Stack) Complete() {
+	now := time.Now().UnixMilli()
+	s.CompletedAt = &now
+	s.Status = StackStatusCompleted
+	duration := now - s.CreatedAt
+	s.DurationMs = &duration
+}
+
+// Fail marks the stack as failed with an error message
+func (s *Stack) Fail(err error) {
+	now := time.Now().UnixMilli()
+	s.CompletedAt = &now
+	s.Status = StackStatusFailed
+	if err != nil {
+		s.Error = err.Error()
+	}
+	duration := now - s.CreatedAt
+	s.DurationMs = &duration
+}
+
+// Timeout marks the stack as timeout
+func (s *Stack) Timeout() {
+	now := time.Now().UnixMilli()
+	s.CompletedAt = &now
+	s.Status = StackStatusTimeout
+	duration := now - s.CreatedAt
+	s.DurationMs = &duration
+}
+
+// IsRoot returns true if this is a root stack (no parent)
+func (s *Stack) IsRoot() bool {
+	return s.ParentID == ""
+}
+
+// IsCompleted returns true if the stack has completed (success, failed, or timeout)
+func (s *Stack) IsCompleted() bool {
+	return s.Status == StackStatusCompleted ||
+		s.Status == StackStatusFailed ||
+		s.Status == StackStatusTimeout
+}
+
+// IsRunning returns true if the stack is currently running
+func (s *Stack) IsRunning() bool {
+	return s.Status == StackStatusRunning
+}
+
+// GetPathString returns the path as a string (e.g., "root_id -> parent_id -> current_id")
+func (s *Stack) GetPathString() string {
+	if len(s.Path) == 0 {
+		return s.ID
+	}
+
+	result := s.Path[0]
+	for i := 1; i < len(s.Path); i++ {
+		result += " -> " + s.Path[i]
+	}
+	return result
+}
+
+// String returns a string representation of the stack for debugging
+func (s *Stack) String() string {
+	status := s.Status
+	if s.IsCompleted() && s.DurationMs != nil {
+		status = fmt.Sprintf("%s (%dms)", s.Status, *s.DurationMs)
+	}
+
+	return fmt.Sprintf("Stack[ID=%s, TraceID=%s, Assistant=%s, Depth=%d, Status=%s]",
+		s.ID[:8], s.TraceID[:8], s.AssistantID, s.Depth, status)
+}
+
+// Clone creates a deep copy of the stack
+func (s *Stack) Clone() *Stack {
+	clone := &Stack{
+		ID:          s.ID,
+		TraceID:     s.TraceID,
+		AssistantID: s.AssistantID,
+		Referer:     s.Referer,
+		Depth:       s.Depth,
+		ParentID:    s.ParentID,
+		Path:        make([]string, len(s.Path)),
+		CreatedAt:   s.CreatedAt,
+		Status:      s.Status,
+		Error:       s.Error,
+	}
+
+	copy(clone.Path, s.Path)
+
+	if s.CompletedAt != nil {
+		completedAt := *s.CompletedAt
+		clone.CompletedAt = &completedAt
+	}
+
+	if s.DurationMs != nil {
+		durationMs := *s.DurationMs
+		clone.DurationMs = &durationMs
+	}
+
+	return clone
+}
+
+// EnterStack initializes or creates a child stack and returns it along with trace ID and completion function
+// This is a helper function to manage stack context for nested calls
+// The stack will be automatically saved to ctx.Stacks for trace logging
+//
+// Returns:
+//   - *Stack: current stack
+//   - string: trace ID (generated for root, inherited for children)
+//   - func(): completion function to be deferred
+//
+// Usage:
+//
+//	stack, traceID, done := context.EnterStack(ctx, assistantID, referer)
+//	defer done()
+//	// ... your code here ...
+func EnterStack(ctx *Context, assistantID, referer string) (*Stack, string, func()) {
+	var stack *Stack
+	var parentStack *Stack
+	var traceID string
+
+	// Initialize Stacks map if not exists
+	if ctx.Stacks == nil {
+		ctx.Stacks = make(map[string]*Stack)
+	}
+
+	if ctx.Stack == nil {
+		// Create root stack for this assistant call (entry point)
+		// Generate a new trace ID for root
+		traceID = trace.GenTraceID()
+		stack = NewStack(traceID, assistantID, referer)
+		ctx.Stack = stack
+	} else {
+		// Create child stack for nested agent call
+		// Inherit trace ID from parent
+		parentStack = ctx.Stack
+		traceID = parentStack.TraceID
+		stack = ctx.Stack.NewChildStack(assistantID, referer)
+		ctx.Stack = stack
+	}
+
+	// Mark stack as running (in case it was pending)
+	if stack.Status == StackStatusPending {
+		stack.Status = StackStatusRunning
+	}
+
+	// Save stack to collection for trace logging
+	ctx.Stacks[stack.ID] = stack
+
+	// Return completion function
+	done := func() {
+		// Mark as completed if no panic occurred
+		if !stack.IsCompleted() {
+			stack.Complete()
+		}
+
+		// Restore parent stack
+		if parentStack != nil {
+			ctx.Stack = parentStack
+		}
+	}
+
+	return stack, traceID, done
+}
+
+// GetAllStacks returns all stacks collected during the request
+// This is useful for trace logging after the request completes
+func (ctx *Context) GetAllStacks() []*Stack {
+	if ctx.Stacks == nil {
+		return nil
+	}
+
+	stacks := make([]*Stack, 0, len(ctx.Stacks))
+	for _, s := range ctx.Stacks {
+		stacks = append(stacks, s)
+	}
+	return stacks
+}
+
+// GetStackByID returns a specific stack by its ID
+// This is useful for querying stack information during request processing
+func (ctx *Context) GetStackByID(id string) *Stack {
+	if ctx.Stacks == nil {
+		return nil
+	}
+	return ctx.Stacks[id]
+}
+
+// GetStacksByTraceID returns all stacks with the given trace ID
+// This is useful for getting the complete call tree for a trace
+func (ctx *Context) GetStacksByTraceID(traceID string) []*Stack {
+	if ctx.Stacks == nil {
+		return nil
+	}
+
+	stacks := make([]*Stack, 0)
+	for _, s := range ctx.Stacks {
+		if s.TraceID == traceID {
+			stacks = append(stacks, s)
+		}
+	}
+	return stacks
+}
+
+// GetRootStack returns the root stack (depth = 0) of current trace
+func (ctx *Context) GetRootStack() *Stack {
+	if ctx.Stacks == nil {
+		return nil
+	}
+
+	for _, s := range ctx.Stacks {
+		if s.IsRoot() {
+			return s
+		}
+	}
+	return nil
+}

agent/context/stack_test.go

@@ -0,0 +1,471 @@
+package context
+
+import (
+	"testing"
+	"time"
+
+	"github.com/yaoapp/yao/config"
+	"github.com/yaoapp/yao/test"
+)
+
+func TestNewStack(t *testing.T) {
+	test.Prepare(t, config.Conf)
+	defer test.Clean()
+
+	traceID := "12345678"
+	assistantID := "test-assistant"
+	referer := RefererAPI
+
+	stack := NewStack(traceID, assistantID, referer)
+
+	if stack == nil {
+		t.Fatal("Expected stack to be created, got nil")
+	}
+
+	if stack.TraceID != traceID {
+		t.Errorf("Expected TraceID '%s', got '%s'", traceID, stack.TraceID)
+	}
+
+	if stack.AssistantID != assistantID {
+		t.Errorf("Expected AssistantID '%s', got '%s'", assistantID, stack.AssistantID)
+	}
+
+	if stack.Referer != referer {
+		t.Errorf("Expected Referer '%s', got '%s'", referer, stack.Referer)
+	}
+
+	if stack.Depth != 0 {
+		t.Errorf("Expected Depth 0, got %d", stack.Depth)
+	}
+
+	if stack.ParentID != "" {
+		t.Errorf("Expected empty ParentID, got '%s'", stack.ParentID)
+	}
+
+	if !stack.IsRoot() {
+		t.Error("Expected stack to be root")
+	}
+
+	if stack.Status != StackStatusRunning {
+		t.Errorf("Expected Status '%s', got '%s'", StackStatusRunning, stack.Status)
+	}
+}
+
+func TestNewStack_GenerateTraceID(t *testing.T) {
+	test.Prepare(t, config.Conf)
+	defer test.Clean()
+
+	// Empty traceID should generate a UUID
+	stack := NewStack("", "test-assistant", RefererAPI)
+
+	if stack.TraceID == "" {
+		t.Error("Expected TraceID to be generated, got empty string")
+	}
+
+	// Should be a valid UUID (36 characters with dashes)
+	if len(stack.TraceID) < 8 {
+		t.Errorf("Expected TraceID to be at least 8 characters, got %d", len(stack.TraceID))
+	}
+}
+
+func TestNewChildStack(t *testing.T) {
+	test.Prepare(t, config.Conf)
+	defer test.Clean()
+
+	// Create parent stack
+	parentStack := NewStack("12345678", "parent-assistant", RefererAPI)
+
+	// Create child stack
+	childStack := parentStack.NewChildStack("child-assistant", RefererAgent)
+
+	if childStack == nil {
+		t.Fatal("Expected child stack to be created, got nil")
+	}
+
+	// Child should inherit TraceID
+	if childStack.TraceID != parentStack.TraceID {
+		t.Errorf("Expected child TraceID '%s', got '%s'", parentStack.TraceID, childStack.TraceID)
+	}
+
+	// Child should have parent ID
+	if childStack.ParentID != parentStack.ID {
+		t.Errorf("Expected ParentID '%s', got '%s'", parentStack.ID, childStack.ParentID)
+	}
+
+	// Child should have incremented depth
+	if childStack.Depth != parentStack.Depth+1 {
+		t.Errorf("Expected Depth %d, got %d", parentStack.Depth+1, childStack.Depth)
+	}
+
+	// Child should not be root
+	if childStack.IsRoot() {
+		t.Error("Expected child stack not to be root")
+	}
+
+	// Path should include both parent and child
+	if len(childStack.Path) != 2 {
+		t.Errorf("Expected Path length 2, got %d", len(childStack.Path))
+	}
+
+	if childStack.Path[0] != parentStack.ID {
+		t.Errorf("Expected first path element '%s', got '%s'", parentStack.ID, childStack.Path[0])
+	}
+
+	if childStack.Path[1] != childStack.ID {
+		t.Errorf("Expected second path element '%s', got '%s'", childStack.ID, childStack.Path[1])
+	}
+}
+
+func TestStackComplete(t *testing.T) {
+	test.Prepare(t, config.Conf)
+	defer test.Clean()
+
+	stack := NewStack("12345678", "test-assistant", RefererAPI)
+
+	// Wait a bit to have measurable duration
+	time.Sleep(10 * time.Millisecond)
+
+	stack.Complete()
+
+	if stack.Status != StackStatusCompleted {
+		t.Errorf("Expected Status '%s', got '%s'", StackStatusCompleted, stack.Status)
+	}
+
+	if stack.CompletedAt == nil {
+		t.Error("Expected CompletedAt to be set, got nil")
+	}
+
+	if stack.DurationMs == nil {
+		t.Error("Expected DurationMs to be set, got nil")
+	}
+
+	if *stack.DurationMs < 10 {
+		t.Errorf("Expected DurationMs to be at least 10ms, got %d", *stack.DurationMs)
+	}
+
+	if !stack.IsCompleted() {
+		t.Error("Expected stack to be completed")
+	}
+
+	if stack.IsRunning() {
+		t.Error("Expected stack not to be running")
+	}
+}
+
+func TestStackFail(t *testing.T) {
+	test.Prepare(t, config.Conf)
+	defer test.Clean()
+
+	stack := NewStack("12345678", "test-assistant", RefererAPI)
+
+	testError := "test error message"
+	stack.Fail(nil)
+	stack.Error = testError
+
+	if stack.Status != StackStatusFailed {
+		t.Errorf("Expected Status '%s', got '%s'", StackStatusFailed, stack.Status)
+	}
+
+	if stack.Error != testError {
+		t.Errorf("Expected Error '%s', got '%s'", testError, stack.Error)
+	}
+
+	if !stack.IsCompleted() {
+		t.Error("Expected failed stack to be completed")
+	}
+}
+
+func TestStackTimeout(t *testing.T) {
+	test.Prepare(t, config.Conf)
+	defer test.Clean()
+
+	stack := NewStack("12345678", "test-assistant", RefererAPI)
+
+	stack.Timeout()
+
+	if stack.Status != StackStatusTimeout {
+		t.Errorf("Expected Status '%s', got '%s'", StackStatusTimeout, stack.Status)
+	}
+
+	if !stack.IsCompleted() {
+		t.Error("Expected timeout stack to be completed")
+	}
+}
+
+func TestEnterStack_RootCreation(t *testing.T) {
+	test.Prepare(t, config.Conf)
+	defer test.Clean()
+
+	ctx := &Context{}
+
+	stack, traceID, done := EnterStack(ctx, "test-assistant", RefererAPI)
+	defer done()
+
+	if stack == nil {
+		t.Fatal("Expected stack to be created, got nil")
+	}
+
+	if traceID == "" {
+		t.Error("Expected traceID to be generated, got empty string")
+	}
+
+	// TraceID should be 8 digits (from trace.GenTraceID)
+	if len(traceID) != 8 {
+		t.Errorf("Expected traceID length 8, got %d", len(traceID))
+	}
+
+	if stack.TraceID != traceID {
+		t.Errorf("Expected stack TraceID '%s', got '%s'", traceID, stack.TraceID)
+	}
+
+	if ctx.Stack != stack {
+		t.Error("Expected ctx.Stack to be set to created stack")
+	}
+
+	if ctx.Stacks == nil {
+		t.Fatal("Expected ctx.Stacks to be initialized, got nil")
+	}
+
+	if ctx.Stacks[stack.ID] != stack {
+		t.Error("Expected stack to be saved in ctx.Stacks")
+	}
+
+	if !stack.IsRoot() {
+		t.Error("Expected stack to be root")
+	}
+}
+
+func TestEnterStack_ChildCreation(t *testing.T) {
+	test.Prepare(t, config.Conf)
+	defer test.Clean()
+
+	ctx := &Context{}
+
+	// Create parent
+	parentStack, parentTraceID, parentDone := EnterStack(ctx, "parent-assistant", RefererAPI)
+	defer parentDone()
+
+	if parentStack == nil {
+		t.Fatal("Expected parent stack to be created, got nil")
+	}
+
+	// Create child
+	childStack, childTraceID, childDone := EnterStack(ctx, "child-assistant", RefererAgent)
+	defer childDone()
+
+	if childStack == nil {
+		t.Fatal("Expected child stack to be created, got nil")
+	}
+
+	// Child should inherit trace ID
+	if childTraceID != parentTraceID {
+		t.Errorf("Expected child traceID '%s', got '%s'", parentTraceID, childTraceID)
+	}
+
+	// Child should have parent ID
+	if childStack.ParentID != parentStack.ID {
+		t.Errorf("Expected child ParentID '%s', got '%s'", parentStack.ID, childStack.ParentID)
+	}
+
+	// Both should be saved in ctx.Stacks
+	if len(ctx.Stacks) != 2 {
+		t.Errorf("Expected 2 stacks in ctx.Stacks, got %d", len(ctx.Stacks))
+	}
+
+	// Current stack should be child
+	if ctx.Stack != childStack {
+		t.Error("Expected ctx.Stack to be child stack")
+	}
+}
+
+func TestEnterStack_DoneCallback(t *testing.T) {
+	test.Prepare(t, config.Conf)
+	defer test.Clean()
+
+	ctx := &Context{}
+
+	// Create parent
+	parentStack, _, parentDone := EnterStack(ctx, "parent-assistant", RefererAPI)
+
+	// Create child
+	childStack, _, childDone := EnterStack(ctx, "child-assistant", RefererAgent)
+
+	// Child should be current
+	if ctx.Stack != childStack {
+		t.Error("Expected ctx.Stack to be child stack before done")
+	}
+
+	// Call child done
+	childDone()
+
+	// Parent should be restored
+	if ctx.Stack != parentStack {
+		t.Error("Expected ctx.Stack to be restored to parent stack after child done")
+	}
+
+	// Child should be completed
+	if !childStack.IsCompleted() {
+		t.Error("Expected child stack to be completed after done")
+	}
+
+	// Call parent done
+	parentDone()
+
+	// Parent should be completed
+	if !parentStack.IsCompleted() {
+		t.Error("Expected parent stack to be completed after done")
+	}
+}
+
+func TestContextGetAllStacks(t *testing.T) {
+	test.Prepare(t, config.Conf)
+	defer test.Clean()
+
+	ctx := &Context{}
+
+	// Create multiple stacks
+	_, _, done1 := EnterStack(ctx, "assistant1", RefererAPI)
+	defer done1()
+
+	_, _, done2 := EnterStack(ctx, "assistant2", RefererAgent)
+	defer done2()
+
+	_, _, done3 := EnterStack(ctx, "assistant3", RefererAgent)
+	defer done3()
+
+	// Get all stacks
+	allStacks := ctx.GetAllStacks()
+
+	if len(allStacks) != 3 {
+		t.Errorf("Expected 3 stacks, got %d", len(allStacks))
+	}
+}
+
+func TestContextGetStackByID(t *testing.T) {
+	test.Prepare(t, config.Conf)
+	defer test.Clean()
+
+	ctx := &Context{}
+
+	stack, _, done := EnterStack(ctx, "test-assistant", RefererAPI)
+	defer done()
+
+	// Get stack by ID
+	found := ctx.GetStackByID(stack.ID)
+
+	if found == nil {
+		t.Fatal("Expected to find stack, got nil")
+	}
+
+	if found.ID != stack.ID {
+		t.Errorf("Expected stack ID '%s', got '%s'", stack.ID, found.ID)
+	}
+
+	// Try to get non-existent stack
+	notFound := ctx.GetStackByID("non-existent-id")
+	if notFound != nil {
+		t.Error("Expected nil for non-existent stack ID")
+	}
+}
+
+func TestContextGetStacksByTraceID(t *testing.T) {
+	test.Prepare(t, config.Conf)
+	defer test.Clean()
+
+	ctx := &Context{}
+
+	// Create parent and child (same trace ID)
+	_, traceID, done1 := EnterStack(ctx, "parent-assistant", RefererAPI)
+	defer done1()
+
+	_, _, done2 := EnterStack(ctx, "child-assistant", RefererAgent)
+	defer done2()
+
+	// Get stacks by trace ID
+	stacks := ctx.GetStacksByTraceID(traceID)
+
+	if len(stacks) != 2 {
+		t.Errorf("Expected 2 stacks with trace ID '%s', got %d", traceID, len(stacks))
+	}
+
+	// All should have same trace ID
+	for _, s := range stacks {
+		if s.TraceID != traceID {
+			t.Errorf("Expected TraceID '%s', got '%s'", traceID, s.TraceID)
+		}
+	}
+}
+
+func TestContextGetRootStack(t *testing.T) {
+	test.Prepare(t, config.Conf)
+	defer test.Clean()
+
+	ctx := &Context{}
+
+	// Create parent
+	parentStack, _, done1 := EnterStack(ctx, "parent-assistant", RefererAPI)
+	defer done1()
+
+	// Create child
+	_, _, done2 := EnterStack(ctx, "child-assistant", RefererAgent)
+	defer done2()
+
+	// Get root stack
+	rootStack := ctx.GetRootStack()
+
+	if rootStack == nil {
+		t.Fatal("Expected to find root stack, got nil")
+	}
+
+	if rootStack.ID != parentStack.ID {
+		t.Errorf("Expected root stack ID '%s', got '%s'", parentStack.ID, rootStack.ID)
+	}
+
+	if !rootStack.IsRoot() {
+		t.Error("Expected returned stack to be root")
+	}
+}
+
+func TestStackClone(t *testing.T) {
+	test.Prepare(t, config.Conf)
+	defer test.Clean()
+
+	original := NewStack("12345678", "test-assistant", RefererAPI)
+	original.Complete()
+
+	clone := original.Clone()
+
+	if clone == nil {
+		t.Fatal("Expected clone to be created, got nil")
+	}
+
+	// Check all fields are copied
+	if clone.ID != original.ID {
+		t.Error("ID not cloned correctly")
+	}
+
+	if clone.TraceID != original.TraceID {
+		t.Error("TraceID not cloned correctly")
+	}
+
+	if clone.AssistantID != original.AssistantID {
+		t.Error("AssistantID not cloned correctly")
+	}
+
+	if clone.Status != original.Status {
+		t.Error("Status not cloned correctly")
+	}
+
+	// Check deep copy of Path
+	if len(clone.Path) != len(original.Path) {
+		t.Error("Path length not cloned correctly")
+	}
+
+	// Modify clone's path shouldn't affect original
+	if len(clone.Path) > 0 {
+		clone.Path[0] = "modified"
+		if original.Path[0] == "modified" {
+			t.Error("Path is not deeply copied")
+		}
+	}
+}

agent/context/types.go

@@ -89,13 +89,42 @@ var ValidReferers = map[string]bool{
 	RefererInternal: true,
 }
 
+const (
+	// StackStatusPending stack is created but not started yet
+	StackStatusPending = "pending"
+
+	// StackStatusRunning stack is currently executing
+	StackStatusRunning = "running"
+
+	// StackStatusCompleted stack completed successfully
+	StackStatusCompleted = "completed"
+
+	// StackStatusFailed stack failed with error
+	StackStatusFailed = "failed"
+
+	// StackStatusTimeout stack execution timeout
+	StackStatusTimeout = "timeout"
+)
+
+// ValidStackStatus is the map of valid stack status types
+var ValidStackStatus = map[string]bool{
+	StackStatusPending:   true,
+	StackStatusRunning:   true,
+	StackStatusCompleted: true,
+	StackStatusFailed:    true,
+	StackStatusTimeout:   true,
+}
+
 // Context the context
 type Context struct {
 
 	// Context
 	context.Context
-	Space plan.Space  `json:"-"` // Shared data space, it will be used to share data between the request and the call
+	Space  plan.Space        `json:"-"` // Shared data space, it will be used to share data between the request and the call
-	Cache store.Store `json:"-"` // Cache store, it will be used to store the message cache, default is "__yao.agent.cache"
+	Cache  store.Store       `json:"-"` // Cache store, it will be used to store the message cache, default is "__yao.agent.cache"
+	Stack  *Stack            `json:"-"` // Stack, current active stack of the request
+	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
 
 	// Authorized information
 	Authorized  *types.AuthorizedInfo `json:"authorized,omitempty"`   // Authorized information
@@ -124,6 +153,36 @@ type Context struct {
 	Silent bool `json:"silent,omitempty"` // Silent mode (Deprecated, use Referer instead)
 }
 
+// Stack represents the call stack node for tracing agent-to-agent calls
+// Uses a flat structure to avoid circular references and memory overhead
+type Stack struct {
+	// Identity
+	ID      string `json:"id"`       // Unique stack node ID, used to identify this specific call
+	TraceID string `json:"trace_id"` // Shared trace ID for entire call tree, inherited from root
+
+	// Call context
+	AssistantID string `json:"assistant_id"`      // Assistant handling this call
+	Referer     string `json:"referer,omitempty"` // Call source: api, agent, tool, process, etc.
+	Depth       int    `json:"depth"`             // Call depth in the tree (0=root)
+
+	// Relationships
+	ParentID string   `json:"parent_id,omitempty"` // Parent stack ID (empty for root call)
+	Path     []string `json:"path"`                // Full path from root: [root_id, parent_id, ..., this_id]
+
+	// Tracking
+	CreatedAt   int64  `json:"created_at"`             // Unix timestamp in milliseconds
+	CompletedAt *int64 `json:"completed_at,omitempty"` // Unix timestamp when completed (nil if ongoing)
+	Status      string `json:"status"`                 // Status: pending, running, completed, failed, timeout
+	Error       string `json:"error,omitempty"`        // Error message if failed
+
+	// Metrics
+	DurationMs *int64 `json:"duration_ms,omitempty"` // Duration in milliseconds (calculated when completed)
+}
+
+// Response the response
+// 100% compatible with the OpenAI API
+type Response struct{}
+
 // Message Structure ( OpenAI Chat Completion Input Message Structure, https://platform.openai.com/docs/api-reference/chat/create#chat/create-messages )
 // ===============================
 

agent/jsapi/jsapi.go

@@ -0,0 +1,3 @@
+package jsapi
+
+// JSAPI Register the JavaScript API

agent/llm/llm.go

@@ -0,0 +1 @@
+package llm

agent/mcp/fetch/fetch.go

@@ -0,0 +1 @@
+package fetch

agent/mcp/mcp.go

@@ -0,0 +1 @@
+package mcp

agent/mcp/search/search.go

@@ -0,0 +1 @@
+package search

agent/output/output.go

@@ -0,0 +1 @@
+package output

agent/plan/plan.go

@@ -0,0 +1 @@
+package plan

agent/trace/trace.go

@@ -0,0 +1,27 @@
+package trace
+
+import (
+	"fmt"
+	"time"
+
+	gonanoid "github.com/matoous/go-nanoid/v2"
+)
+
+// GenTraceID generate a new trace ID using NanoID algorithm
+// safe: optional parameter, reserved for future safe mode implementation (collision detection)
+func GenTraceID(safe ...bool) string {
+	// TODO: Implement safe mode with collision detection when needed
+	// For now, NanoID provides sufficient uniqueness without collision checking
+
+	// URL-safe alphabet (no ambiguous characters like 0/O, 1/l/I)
+	const alphabet = "1234567890"
+	const length = 8 // 8 characters provides good balance of uniqueness and readability
+
+	id, err := gonanoid.Generate(alphabet, length)
+	if err != nil {
+		// Fallback to timestamp-based ID if NanoID generation fails
+		return fmt.Sprintf("%d", time.Now().UnixNano())
+	}
+
+	return id
+}

openapi/chat/completions.go

@@ -36,6 +36,8 @@ func GinCreateCompletions(c *gin.Context) {
 		return
 	}
 
+	defer ctx.Release() // Release the context after the request is complete
+
 	fmt.Println("-----------------------------------------------")
 	fmt.Println("Chat ID: ", ctx.ChatID)
 	fmt.Println("Assistant ID: ", ctx.AssistantID)