yao/agent/assistant/agent.go
Max 10e004bfad feat(logging): enhance logging across various components for improved traceability
- Added detailed trace logging in the Assistant's Stream function to capture computer and workspace information.
- Replaced fmt.Printf statements with kunlog.Trace in the Telegram polling adapter for consistent logging.
- Updated sandbox lifecycle management functions to utilize kunlog for logging workspace and node resolution details.
- Enhanced error handling in the Claude parser with trace logging for JSON unmarshal and scanner errors.
- Improved file handling logging in the workspace API to provide clearer insights into file read operations.
2026-03-21 20:44:59 +08:00

813 lines
28 KiB
Go

package assistant
import (
"fmt"
"log"
"time"
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/gou/connector"
goullm "github.com/yaoapp/gou/llm"
"github.com/yaoapp/yao/agent/assistant/handlers"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/i18n"
"github.com/yaoapp/yao/agent/llm"
"github.com/yaoapp/yao/agent/output/message"
agentsandbox "github.com/yaoapp/yao/agent/sandbox"
sandboxTypes "github.com/yaoapp/yao/agent/sandbox/v2/types"
infraV2 "github.com/yaoapp/yao/sandbox/v2"
)
// Stream stream the agent
// 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) (*context.Response, error) {
// Update logger with assistant ID and start logging
ctx.Logger.SetAssistantID(ast.ID)
ctx.Logger.Start()
// Validate user permissions
var err error
err = ast.checkPermissions(ctx)
if err != nil {
return nil, err
}
// Start stream time
streamStartTime := time.Now()
// Set up interrupt handler if interrupt controller is available
// InterruptController handles user interrupt signals (stop button) for appending messages
// HTTP context cancellation is handled naturally by LLM/Agent layers
if ctx.Interrupt != nil {
ctx.Interrupt.SetHandler(func(c *context.Context, signal *context.InterruptSignal) error {
return ast.handleInterrupt(c, signal)
})
}
// ================================================
// Initialize
// ================================================
ctx.Logger.Phase("Initialize")
// Get or create options
var opts *context.Options
if len(options) > 0 && options[0] != nil {
opts = options[0]
} else {
opts = &context.Options{}
}
// Merge caller-provided metadata into ctx so sub-agent hooks can read it via ctx.metadata
ctx.MergeMetadata(opts.Metadata)
// Initialize stack and auto-handle completion/failure/restore
_, _, done := context.EnterStack(ctx, ast.ID, opts)
defer done()
// Auto-skip history for forked Agent-to-Agent calls (ctx.agent.Call/All/Any/Race)
// This ensures forked A2A messages don't pollute chat history.
// Delegate calls (RefererAgent) still save history as they are part of the main conversation flow.
// Note: Output is NOT skipped - sub-agents output normally with ThreadID for UI separation.
if ctx.IsForkedA2ACall() {
if opts == nil {
opts = &context.Options{}
}
opts.ForceA2A()
}
// ================================================
// Initialize Chat Buffer (for root stack only)
// Buffer is flushed in defer block at the end
// ================================================
ast.InitBuffer(ctx)
// Track final status for buffer flush
var finalStatus = context.StepStatusCompleted
var finalError error
// Defer buffer flush - always executes on exit (success, error, interrupt, panic)
defer func() {
// Handle panic recovery for status tracking
if r := recover(); r != nil {
finalStatus = context.ResumeStatusFailed
if e, ok := r.(error); ok {
finalError = e
} else {
finalError = fmt.Errorf("panic: %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)
// Log end of request
ctx.Logger.End(finalStatus == context.StepStatusCompleted, finalError)
ctx.Logger.RestoreAssistantID()
}()
// Determine stream handler
streamHandler := ast.getStreamHandler(ctx, opts)
// Get connector and capabilities early (before sending stream_start)
// so that output adapters can use them when converting stream_start event
err = ast.initializeCapabilities(ctx, opts)
if err != nil {
finalStatus = context.ResumeStatusFailed
finalError = err
ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
return nil, err
}
// Send ChunkStreamStart only for root stack (agent-level stream start)
// Now ctx.Capabilities is set, so output adapters can use it
ast.sendAgentStreamStart(ctx, streamHandler, streamStartTime)
// Initialize chat, prepare kb collection (optional) etc.
// Use async version to not block the main flow
ast.InitializeConversationAsync(ctx, opts)
ctx.Logger.PhaseComplete("Initialize")
// Ensure chat session exists
ast.EnsureChat(ctx)
// Initialize agent trace node
agentNode := ast.initAgentTraceNode(ctx, inputMessages)
// ================================================
// Get Full Messages with chat history
// ================================================
ctx.Logger.Phase("History")
historyResult, err := ast.WithHistory(ctx, inputMessages, agentNode, opts)
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)
// Skip if History is disabled in options (for internal calls like needsearch)
// Note: For A2A calls, ForceA2A() sets skip.history = true, so this will be skipped
if opts == nil || opts.Skip == nil || !opts.Skip.History {
ast.BufferUserInput(ctx, historyResult.InputMessages)
}
ctx.Logger.PhaseComplete("History")
// ================================================
// Initialize Sandbox (if configured)
// ================================================
// Sandbox must be created BEFORE hooks so that hooks can access ctx.sandbox
var sandboxExecutor agentsandbox.Executor
var sandboxCleanup func()
var sandboxLoadingMsgID string
// V2 sandbox state
var v2Runner sandboxTypes.Runner
var v2Computer infraV2.Computer
var v2LoadingMsgID string
if ast.HasSandboxV2() {
ctx.Logger.Phase("Sandbox V2")
var err error
var v2Cleanup func()
v2Runner, v2Computer, v2Cleanup, v2LoadingMsgID, err = ast.initSandboxV2(ctx, opts)
if err != nil {
ast.traceAgentFail(agentNode, err)
ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
return nil, err
}
sandboxCleanup = v2Cleanup
ctx.Logger.PhaseComplete("Sandbox V2")
if v2Computer != nil {
ci := v2Computer.ComputerInfo()
ctx.Logger.Trace("Node: %s (%s)", ci.NodeID, ci.Kind)
if ci.BoxID != "" {
ctx.Logger.Trace("Computer: %s", ci.BoxID)
}
ctx.Logger.Trace("Workspace: %s", ast.SandboxV2.WorkspaceID)
}
} else if ast.HasSandbox() {
ctx.Logger.Phase("Sandbox")
var err error
sandboxExecutor, sandboxCleanup, sandboxLoadingMsgID, err = ast.initSandbox(ctx, opts)
if err != nil {
ast.traceAgentFail(agentNode, err)
ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
return nil, err
}
// Set sandbox executor in context so hooks can access ctx.sandbox
// The executor implements both agentsandbox.Executor and context.SandboxExecutor
ctx.SetSandboxExecutor(sandboxExecutor)
ctx.Logger.PhaseComplete("Sandbox")
}
// Ensure sandbox cleanup on exit
defer func() {
if sandboxCleanup != nil {
sandboxCleanup()
}
}()
// ================================================
// Standalone Workspace Loading (no sandbox required)
// ================================================
// When no sandbox is configured but the user selected a workspace,
// load the workspace FS into context so hooks can access ctx.workspace.
if !ctx.HasWorkspace() {
ast.initStandaloneWorkspace(ctx)
}
// ================================================
// Execute Create Hook
// ================================================
// 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,
})
var err error
createResponse, opts, err = ast.HookScript.Create(ctx, fullMessages, opts)
if err != nil {
finalStatus = context.ResumeStatusFailed
finalError = err
ast.traceAgentFail(agentNode, err)
// Send error stream_end for root stack
ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
return nil, err
}
// Complete step
ast.CompleteStep(ctx, map[string]interface{}{
"response": createResponse,
})
// Log the create response
ast.traceCreateHook(agentNode, createResponse)
ctx.Logger.HookComplete("Create")
// Check if Create hook wants to delegate to another agent
// This allows early routing to sub-agents without LLM call
if createResponse != nil && createResponse.Delegate != nil {
ctx.Logger.Debug("Create hook delegating to agent: %s", createResponse.Delegate.AgentID)
// Delegate to target agent (reuse existing delegation logic from next.go)
// Note: User input is already buffered by root agent, delegated agent will skip buffering
delegateResponse, err := ast.handleDelegation(ctx, createResponse.Delegate, streamHandler)
if err != nil {
finalStatus = context.ResumeStatusFailed
finalError = err
ast.traceAgentFail(agentNode, err)
ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
return nil, err
}
// For root stack, send stream_end and close output
// (delegated agent handles its own stream events, but root needs to close)
if ctx.Stack != nil && ctx.Stack.IsRoot() {
ast.sendAgentStreamEnd(ctx, streamHandler, streamStartTime, "completed", nil, nil)
if err := ctx.CloseOutput(); err != nil {
if trace, _ := ctx.Trace(); trace != nil {
trace.Error(i18n.Tr(ast.ID, ctx.Locale, "assistant.agent.stream.close_error"), map[string]any{"error": err.Error()})
}
}
}
// Return delegated response directly (skip LLM call and Next hook)
return delegateResponse, nil
}
}
// ================================================
// Execute LLM Call Stream
// ================================================
// LLM Call Stream ( Optional )
var completionResponse *context.CompletionResponse
var completionMessages []context.Message
var completionOptions *context.CompletionOptions
if ast.Prompts != nil || ast.MCP != nil {
ctx.Logger.Phase("LLM")
// Build the LLM request first (use fullMessages which includes history)
// Note: completionMessages here are still in original format (with __yao.attachment:// URLs)
// Content conversion (BuildContent) happens inside executeLLMStream, right before LLM call
// This ensures autoSearch and delegate receive original messages, not converted ones
completionMessages, completionOptions, err = ast.BuildRequest(ctx, fullMessages, createResponse)
if err != nil {
finalStatus = context.ResumeStatusFailed
finalError = err
ast.traceAgentFail(agentNode, err)
// Send error stream_end for root stack
ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
return nil, err
}
// ================================================
// Execute Auto Search (if enabled)
// ================================================
if intent := ast.shouldAutoSearch(ctx, completionMessages, createResponse, opts); intent != nil {
refCtx := ast.executeAutoSearch(ctx, completionMessages, createResponse, intent, opts)
if refCtx != nil && len(refCtx.References) > 0 {
completionMessages = ast.injectSearchContext(completionMessages, refCtx)
}
}
// Begin step tracking for LLM call
ast.BeginStep(ctx, context.StepTypeLLM, map[string]interface{}{
"messages": completionMessages,
})
// Execute the LLM streaming call
// Choose between sandbox execution or direct LLM execution
if ast.HasSandboxV2() && v2Runner != nil && v2Computer != nil && v2Runner.Name() != "yao" {
// V2 Sandbox execution path (non-yao runners replace LLM.Stream)
completionResponse, err = ast.executeSandboxV2Stream(ctx, completionMessages, agentNode, streamHandler, v2Runner, v2Computer, v2LoadingMsgID)
} else if ast.HasSandboxV2() && v2Runner != nil && v2Runner.Name() == "yao" {
// V2 yao runner: Prepare is done, close loading, fall through to LLM
if v2LoadingMsgID != "" {
closeLoadingV2(ctx, v2LoadingMsgID, "")
}
completionResponse, err = ast.executeLLMStream(ctx, completionMessages, completionOptions, agentNode, streamHandler, opts)
} else if ast.HasSandbox() {
// V1 Sandbox execution path (Claude CLI, Cursor CLI, etc.)
completionResponse, err = ast.executeSandboxStream(ctx, completionMessages, agentNode, streamHandler, sandboxExecutor, sandboxLoadingMsgID)
} else {
// Direct LLM execution path
completionResponse, err = ast.executeLLMStream(ctx, completionMessages, completionOptions, agentNode, streamHandler, opts)
}
if err != nil {
finalStatus = context.ResumeStatusFailed
finalError = err
ast.traceAgentFail(agentNode, err)
// Send error stream_end for root stack
ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
return nil, err
}
// Complete LLM step
ast.CompleteStep(ctx, map[string]interface{}{
"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")
}
// ================================================
// Execute tool calls with retry
// ================================================
// Note: Skip MCP tool calls execution for sandbox mode - Claude CLI handles them internally
var toolCallResponses []context.ToolCallResponse = nil
if completionResponse != nil && completionResponse.ToolCalls != nil && !ast.HasSandbox() {
maxToolRetries := 3
currentMessages := completionMessages
currentResponse := completionResponse
for attempt := 0; attempt < maxToolRetries; attempt++ {
// Begin step tracking for tool calls
ast.BeginStep(ctx, context.StepTypeTool, map[string]interface{}{
"tool_calls": currentResponse.ToolCalls,
"attempt": attempt,
})
// Execute all tool calls
toolResults, hasErrors := ast.executeToolCalls(ctx, currentResponse.ToolCalls, attempt)
// Build a map of tool call ID to arguments for quick lookup
toolCallArgsMap := make(map[string]interface{})
for _, tc := range currentResponse.ToolCalls {
toolCallArgsMap[tc.ID] = tc.Function.Arguments
}
// Convert toolResults to toolCallResponses
toolCallResponses = make([]context.ToolCallResponse, len(toolResults))
for i, result := range toolResults {
parsedContent, _ := result.ParsedContent()
toolCallResponses[i] = context.ToolCallResponse{
ToolCallID: result.ToolCallID,
Server: result.Server(),
Tool: result.Tool(),
Arguments: toolCallArgsMap[result.ToolCallID],
Result: parsedContent,
Error: "",
}
if result.Error != nil {
toolCallResponses[i].Error = result.Error.Error()
}
}
// If all successful, complete step and break out
if !hasErrors {
ast.CompleteStep(ctx, map[string]interface{}{
"results": toolCallResponses,
})
ctx.Logger.Debug("All tool calls succeeded (attempt %d)", attempt)
break
}
// Check if any errors are retryable (parameter/validation issues)
hasRetryableErrors := false
for _, result := range toolResults {
if result.Error != nil && result.IsRetryableError {
hasRetryableErrors = true
break
}
}
// If no retryable errors, don't retry (MCP internal issues)
if !hasRetryableErrors {
err := fmt.Errorf("tool calls failed with non-retryable errors (MCP internal issues)")
finalStatus = context.ResumeStatusFailed
finalError = err
ctx.Logger.Error("Tool calls failed: %v", err)
ast.traceAgentFail(agentNode, err)
ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
return nil, err
}
// If it's the last attempt, return error
if attempt == maxToolRetries-1 {
err := fmt.Errorf("tool calls failed after %d attempts", maxToolRetries)
finalStatus = context.ResumeStatusFailed
finalError = err
ctx.Logger.Error("Tool calls failed: %v", err)
ast.traceAgentFail(agentNode, err)
ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
return nil, err
}
// Complete current step (with partial results)
ast.CompleteStep(ctx, map[string]interface{}{
"results": toolCallResponses,
"has_errors": true,
})
// Build retry messages with tool call results (including errors)
retryMessages := ast.buildToolRetryMessages(currentMessages, currentResponse, toolResults)
// Begin LLM retry step
ast.BeginStep(ctx, context.StepTypeLLM, map[string]interface{}{
"messages": retryMessages,
"retry_attempt": attempt + 1,
})
// Retry LLM call (streaming to keep user informed)
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
ctx.Logger.Error("LLM retry failed: %v", err)
ast.traceAgentFail(agentNode, err)
ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
return nil, err
}
// If LLM didn't return tool calls, it might have given up
if currentResponse.ToolCalls == nil {
err := fmt.Errorf("LLM did not return tool calls in retry attempt %d", attempt+1)
finalStatus = context.ResumeStatusFailed
finalError = 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
}
// Complete LLM retry step
ast.CompleteStep(ctx, map[string]interface{}{
"content": currentResponse.Content,
"tool_calls": currentResponse.ToolCalls,
})
// Update messages for next iteration
currentMessages = retryMessages
}
// Update completionResponse with the final successful response
completionResponse = currentResponse
}
// ================================================
// Execute Next Hook and Process Response
// ================================================
var finalResponse *context.Response
var nextResponse *context.NextHookResponse = nil
if ast.HookScript != nil {
ctx.Logger.HookStart("Next")
// Begin step tracking for hook_next
ast.BeginStep(ctx, context.StepTypeHookNext, map[string]interface{}{
"messages": fullMessages,
"completion": completionResponse,
"tools": toolCallResponses,
})
var err error
nextResponse, opts, err = ast.HookScript.Next(ctx, &context.NextHookPayload{
Messages: fullMessages,
Completion: completionResponse,
Tools: toolCallResponses,
}, opts)
if err != nil {
finalStatus = context.ResumeStatusFailed
finalError = err
ast.traceAgentFail(agentNode, err)
ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
return nil, err
}
// Complete hook_next step
ast.CompleteStep(ctx, map[string]interface{}{
"response": nextResponse,
})
ctx.Logger.HookComplete("Next")
// Process Next hook response
finalResponse, err = ast.processNextResponse(&NextProcessContext{
Context: ctx,
NextResponse: nextResponse,
CompletionResponse: completionResponse,
FullMessages: fullMessages,
ToolCallResponses: toolCallResponses,
StreamHandler: streamHandler,
CreateResponse: createResponse,
})
if err != nil {
finalStatus = context.ResumeStatusFailed
finalError = err
ast.traceAgentFail(agentNode, err)
ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
return nil, err
}
} else {
// No Next hook: use standard response
finalResponse = ast.buildStandardResponse(&NextProcessContext{
Context: ctx,
NextResponse: nil,
CompletionResponse: completionResponse,
FullMessages: fullMessages,
ToolCallResponses: toolCallResponses,
StreamHandler: streamHandler,
CreateResponse: createResponse,
})
}
// Create completion node to report final output
ast.traceAgentCompletion(ctx, createResponse, nextResponse, completionResponse, finalResponse)
// Only close output and send stream_end if this is the root call (entry point)
// Nested calls (from MCP, hooks, etc.) should not close the output or send stream_end
// Note: Flush is already handled by the stream handler (handleStreamEnd)
if ctx.Stack != nil && ctx.Stack.IsRoot() {
// Log closing output for root call
if trace, _ := ctx.Trace(); trace != nil {
trace.Debug("Agent: Closing output (root call)", map[string]any{
"stack_id": ctx.Stack.ID,
"depth": ctx.Stack.Depth,
"assistant_id": ctx.Stack.AssistantID,
})
}
// Send ChunkStreamEnd (agent-level stream completion)
ast.sendAgentStreamEnd(ctx, streamHandler, streamStartTime, "completed", nil, completionResponse)
// Close the output writer to send [DONE] marker
if err := ctx.CloseOutput(); err != nil {
if trace, _ := ctx.Trace(); trace != nil {
trace.Error(i18n.Tr(ast.ID, ctx.Locale, "assistant.agent.stream.close_error"), map[string]any{"error": err.Error()}) // "Failed to close output"
}
}
} else {
// Log skipping close for nested call
if trace, _ := ctx.Trace(); trace != nil && ctx.Stack != nil {
trace.Debug("Agent: Skipping output close (nested call)", map[string]any{
"stack_id": ctx.Stack.ID,
"depth": ctx.Stack.Depth,
"parent_id": ctx.Stack.ParentID,
"assistant_id": ctx.Stack.AssistantID,
})
}
}
// Return finalResponse which could be:
// 1. Result from delegated agent call (already a Response)
// 2. Custom data from Next hook (wrapped in standard Response)
// 3. Standard response
return finalResponse, nil
}
// GetConnector get the connector object, capabilities, and error with priority:
// opts.Connector > ast.Connector > defaultConnector (fallback)
// Note: opts.Connector may be set by Create hook's applyOptionsAdjustments
// Returns: (connector, capabilities, error)
func (ast *Assistant) GetConnector(ctx *context.Context, opts ...*context.Options) (connector.Connector, *goullm.Capabilities, error) {
connectorID := ast.Connector
if len(opts) > 0 && opts[0] != nil && opts[0].Connector != "" {
connectorID = opts[0].Connector
}
if connectorID == "" {
connectorID = defaultConnector
}
if connectorID == "" {
return nil, nil, fmt.Errorf("connector not specified")
}
conn, err := connector.Select(connectorID)
if err != nil && connectorID != defaultConnector && defaultConnector != "" {
log.Printf("[Assistant] connector %q not found, falling back to default %q", connectorID, defaultConnector)
conn, err = connector.Select(defaultConnector)
}
if err != nil {
return nil, nil, err
}
capabilities := llm.GetCapabilitiesFromConn(conn)
return conn, capabilities, nil
}
// Info get the assistant information
func (ast *Assistant) Info(locale ...string) *message.AssistantInfo {
lc := "en"
if len(locale) > 0 {
lc = locale[0]
}
return &message.AssistantInfo{
ID: ast.ID,
Type: ast.Type,
Name: i18n.Tr(ast.ID, lc, ast.Name),
Avatar: ast.Avatar,
Description: i18n.Tr(ast.ID, lc, ast.Description),
}
}
// getStreamHandler returns the stream handler from options or a default one
func (ast *Assistant) getStreamHandler(ctx *context.Context, opts ...*context.Options) message.StreamFunc {
// Check if handler is provided in options
if len(opts) > 0 && opts[0] != nil && opts[0].Writer != nil {
return handlers.DefaultStreamHandler(ctx)
}
return handlers.DefaultStreamHandler(ctx)
}
// sendAgentStreamStart sends ChunkStreamStart for root stack only (agent-level stream start)
// This ensures only one stream_start per agent execution, even with multiple LLM calls
func (ast *Assistant) sendAgentStreamStart(ctx *context.Context, handler message.StreamFunc, startTime time.Time) {
if ctx.Stack == nil || !ctx.Stack.IsRoot() || handler == nil {
return
}
// Build the start data
startData := message.EventStreamStartData{
ContextID: ctx.ID,
ChatID: ctx.ChatID,
TraceID: ctx.TraceID(),
RequestID: ctx.RequestID(),
Timestamp: startTime.UnixMilli(),
Assistant: ast.Info(ctx.Locale),
Metadata: ctx.Metadata,
}
if startJSON, err := jsoniter.Marshal(startData); err == nil {
handler(message.ChunkStreamStart, startJSON)
}
}
// sendAgentStreamEnd sends ChunkStreamEnd for root stack only (agent-level stream completion)
func (ast *Assistant) sendAgentStreamEnd(ctx *context.Context, handler message.StreamFunc, startTime time.Time, status string, err error, response *context.CompletionResponse) {
if ctx.Stack == nil || !ctx.Stack.IsRoot() || handler == nil {
return
}
endData := &message.EventStreamEndData{
RequestID: ctx.RequestID(),
ContextID: ctx.ID,
Timestamp: time.Now().UnixMilli(),
DurationMs: time.Since(startTime).Milliseconds(),
Status: status,
TraceID: ctx.TraceID(),
Metadata: ctx.Metadata,
}
if err != nil {
endData.Error = err.Error()
}
if response != nil && response.Usage != nil {
endData.Usage = response.Usage
}
if endJSON, marshalErr := jsoniter.Marshal(endData); marshalErr == nil {
handler(message.ChunkStreamEnd, endJSON)
}
}
// sendStreamEndOnError sends ChunkStreamEnd with error status for root stack only
func (ast *Assistant) sendStreamEndOnError(ctx *context.Context, handler message.StreamFunc, startTime time.Time, err error) {
ast.sendAgentStreamEnd(ctx, handler, startTime, "error", err, nil)
}
// handleInterrupt handles the interrupt signal
// This is called by the interrupt listener when a signal is received
func (ast *Assistant) handleInterrupt(ctx *context.Context, signal *context.InterruptSignal) error {
switch signal.Type {
case context.InterruptForce:
ctx.Logger.Debug("Force interrupt received")
if ctx.Buffer != nil {
ctx.Buffer.FailCurrentStep(context.ResumeStatusInterrupted,
fmt.Errorf("interrupted by user"))
}
case context.InterruptGraceful:
ctx.Logger.Debug("Graceful interrupt received: messages=%d", len(signal.Messages))
}
return nil
}
// initializeCapabilities gets connector and capabilities, then sets them in context
// This should be called early (before sending stream_start) so that output adapters
// can use capabilities when converting stream_start event
func (ast *Assistant) initializeCapabilities(ctx *context.Context, opts *context.Options) error {
if ast.Prompts == nil && ast.MCP == nil {
return nil
}
_, capabilities, err := ast.GetConnector(ctx, opts)
if err != nil {
return err
}
// Set capabilities in context for output adapters to use
if capabilities != nil {
ctx.Capabilities = capabilities
}
return nil
}
// buildToolRetryMessages builds messages for LLM retry with tool call results
// Format follows OpenAI's tool call response pattern:
// 1. Assistant message with tool calls
// 2. Tool messages with results (one per tool call)
// 3. System message explaining the retry
func (ast *Assistant) buildToolRetryMessages(
previousMessages []context.Message,
completionResponse *context.CompletionResponse,
toolResults []ToolCallResult,
) []context.Message {
retryMessages := make([]context.Message, 0, len(previousMessages)+len(toolResults)+2)
// Add all previous messages
retryMessages = append(retryMessages, previousMessages...)
// Add assistant message with tool calls
assistantMsg := context.Message{
Role: context.RoleAssistant,
Content: completionResponse.Content,
ToolCalls: completionResponse.ToolCalls,
}
retryMessages = append(retryMessages, assistantMsg)
// Add tool result messages (one per tool call)
for _, result := range toolResults {
toolMsg := context.Message{
Role: context.RoleTool,
Content: result.Content,
ToolCallID: &result.ToolCallID,
}
// Add tool name if available
if result.Name != "" {
name := result.Name
toolMsg.Name = &name
}
retryMessages = append(retryMessages, toolMsg)
}
// Add system message explaining the retry (optional, helps LLM understand context)
systemMsg := context.Message{
Role: context.RoleSystem,
Content: i18n.Tr(ast.ID, "en", "assistant.agent.tool_retry_prompt"),
}
retryMessages = append(retryMessages, systemMsg)
return retryMessages
}