diff --git a/agent/assistant/agent.go b/agent/assistant/agent.go index 8ad10082..0e32c212 100644 --- a/agent/assistant/agent.go +++ b/agent/assistant/agent.go @@ -3,11 +3,31 @@ 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 { +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) { +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 } diff --git a/agent/assistant/hooks/hooks.go b/agent/assistant/hooks/hooks.go new file mode 100644 index 00000000..6de678e5 --- /dev/null +++ b/agent/assistant/hooks/hooks.go @@ -0,0 +1 @@ +package hooks diff --git a/agent/context/context.go b/agent/context/context.go index ed418b3b..26c7a16e 100644 --- a/agent/context/context.go +++ b/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() - ctx.Space = nil + // Clear space + 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 } diff --git a/agent/context/interfaces.go b/agent/context/interfaces.go index 20683fbf..ddf3197d 100644 --- a/agent/context/interfaces.go +++ b/agent/context/interfaces.go @@ -1,14 +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 + Stream(ctx *Context, messages []Message, handler StreamFunc) error // Run run the agent - Run(ctx Context, messages []Message) (*Response, error) + Run(ctx *Context, messages []Message) (*Response, error) } diff --git a/agent/context/openapi.go b/agent/context/openapi.go index d32370de..e6ffd830 100644 --- a/agent/context/openapi.go +++ b/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 - chatID = "" + // Fallback: Generate a new chat ID if extraction fails + 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, diff --git a/agent/context/openapi_test.go b/agent/context/openapi_test.go index 5f479aad..65cc1bda 100644 --- a/agent/context/openapi_test.go +++ b/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)) + } +} diff --git a/agent/context/stack.go b/agent/context/stack.go new file mode 100644 index 00000000..6ee09e05 --- /dev/null +++ b/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 +} diff --git a/agent/context/stack_test.go b/agent/context/stack_test.go new file mode 100644 index 00000000..80f43653 --- /dev/null +++ b/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") + } + } +} diff --git a/agent/context/types.go b/agent/context/types.go index d1737f5d..74068f14 100644 --- a/agent/context/types.go +++ b/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 - Cache store.Store `json:"-"` // Cache store, it will be used to store the message cache, default is "__yao.agent.cache" + 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" + 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,32 @@ 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{} diff --git a/agent/assistant/js/js.go b/agent/jsapi/jsapi.go similarity index 73% rename from agent/assistant/js/js.go rename to agent/jsapi/jsapi.go index 63ff0ccc..4c244c45 100644 --- a/agent/assistant/js/js.go +++ b/agent/jsapi/jsapi.go @@ -1,3 +1,3 @@ -package js +package jsapi // JSAPI Register the JavaScript API diff --git a/agent/llm/llm.go b/agent/llm/llm.go new file mode 100644 index 00000000..006c3ded --- /dev/null +++ b/agent/llm/llm.go @@ -0,0 +1 @@ +package llm diff --git a/agent/mcp/fetch/fetch.go b/agent/mcp/fetch/fetch.go new file mode 100644 index 00000000..96ccdabc --- /dev/null +++ b/agent/mcp/fetch/fetch.go @@ -0,0 +1 @@ +package fetch diff --git a/agent/mcp/mcp.go b/agent/mcp/mcp.go new file mode 100644 index 00000000..87468ebd --- /dev/null +++ b/agent/mcp/mcp.go @@ -0,0 +1 @@ +package mcp diff --git a/agent/mcp/search/search.go b/agent/mcp/search/search.go new file mode 100644 index 00000000..5ed85155 --- /dev/null +++ b/agent/mcp/search/search.go @@ -0,0 +1 @@ +package search diff --git a/agent/output/output.go b/agent/output/output.go new file mode 100644 index 00000000..ad89311b --- /dev/null +++ b/agent/output/output.go @@ -0,0 +1 @@ +package output diff --git a/agent/plan/plan.go b/agent/plan/plan.go new file mode 100644 index 00000000..ec2d2c6b --- /dev/null +++ b/agent/plan/plan.go @@ -0,0 +1 @@ +package plan diff --git a/agent/trace/trace.go b/agent/trace/trace.go new file mode 100644 index 00000000..ea39374f --- /dev/null +++ b/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 +} diff --git a/openapi/chat/completions.go b/openapi/chat/completions.go index a1491bc2..c6bd23e7 100644 --- a/openapi/chat/completions.go +++ b/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)