kit/internal/agent/agent.go

package agent

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"time"

	tea "charm.land/bubbletea/v2"
	"charm.land/fantasy"

	"github.com/mark3labs/mcphost/internal/config"
	"github.com/mark3labs/mcphost/internal/models"
	"github.com/mark3labs/mcphost/internal/tools"
)

// AgentConfig holds configuration options for creating a new Agent.
type AgentConfig struct {
	ModelConfig      *models.ProviderConfig
	MCPConfig        *config.Config
	SystemPrompt     string
	MaxSteps         int
	StreamingEnabled bool
	DebugLogger      tools.DebugLogger
}

// ToolCallHandler is a function type for handling tool calls as they happen.
type ToolCallHandler func(toolName, toolArgs string)

// ToolExecutionHandler is a function type for handling tool execution start/end events.
type ToolExecutionHandler func(toolName string, isStarting bool)

// ToolResultHandler is a function type for handling tool results.
type ToolResultHandler func(toolName, toolArgs, result string, isError bool)

// ResponseHandler is a function type for handling LLM responses.
type ResponseHandler func(content string)

// StreamingResponseHandler is a function type for handling streaming LLM responses.
type StreamingResponseHandler func(content string)

// ToolCallContentHandler is a function type for handling content that accompanies tool calls.
type ToolCallContentHandler func(content string)

// ToolApprovalHandler is a function type for handling user approval of tool calls.
type ToolApprovalHandler func(toolName, toolArgs string) (bool, error)

// Agent represents an AI agent with MCP tool integration using the fantasy library.
// It manages the interaction between an LLM and various tools through the MCP protocol.
type Agent struct {
	toolManager      *tools.MCPToolManager
	fantasyAgent     fantasy.Agent
	model            fantasy.LanguageModel
	maxSteps         int
	systemPrompt     string
	loadingMessage   string
	providerType     string
	streamingEnabled bool
}

// GenerateWithLoopResult contains the result and conversation history from an agent interaction.
type GenerateWithLoopResult struct {
	// FinalResponse is the last message generated by the model
	FinalResponse *fantasy.Response
	// ConversationMessages contains all messages in the conversation including tool calls and results
	ConversationMessages []fantasy.Message
	// TotalUsage contains aggregate token usage across all steps
	TotalUsage fantasy.Usage
}

// NewAgent creates a new Agent with MCP tool integration and streaming support.
func NewAgent(ctx context.Context, agentConfig *AgentConfig) (*Agent, error) {
	// Create the LLM provider via fantasy
	providerResult, err := models.CreateProvider(ctx, agentConfig.ModelConfig)
	if err != nil {
		return nil, fmt.Errorf("failed to create model provider: %v", err)
	}

	// Create and load MCP tools
	toolManager := tools.NewMCPToolManager()
	toolManager.SetModel(providerResult.Model)

	if agentConfig.DebugLogger != nil {
		toolManager.SetDebugLogger(agentConfig.DebugLogger)
	}

	if err := toolManager.LoadTools(ctx, agentConfig.MCPConfig); err != nil {
		return nil, fmt.Errorf("failed to load MCP tools: %v", err)
	}

	// Build fantasy agent options
	var agentOpts []fantasy.AgentOption

	if agentConfig.SystemPrompt != "" {
		agentOpts = append(agentOpts, fantasy.WithSystemPrompt(agentConfig.SystemPrompt))
	}

	// Register all MCP tools with the fantasy agent
	mcpTools := toolManager.GetTools()
	if len(mcpTools) > 0 {
		agentOpts = append(agentOpts, fantasy.WithTools(mcpTools...))
	}

	// Set max steps as stop condition
	if agentConfig.MaxSteps > 0 {
		agentOpts = append(agentOpts, fantasy.WithStopConditions(
			fantasy.StepCountIs(agentConfig.MaxSteps),
		))
	}

	// Create the fantasy agent
	fantasyAgent := fantasy.NewAgent(providerResult.Model, agentOpts...)

	// Determine provider type from model string
	providerType := "default"
	if agentConfig.ModelConfig != nil && agentConfig.ModelConfig.ModelString != "" {
		parts := strings.SplitN(agentConfig.ModelConfig.ModelString, ":", 2)
		if len(parts) >= 1 {
			providerType = parts[0]
		}
	}

	return &Agent{
		toolManager:      toolManager,
		fantasyAgent:     fantasyAgent,
		model:            providerResult.Model,
		maxSteps:         agentConfig.MaxSteps,
		systemPrompt:     agentConfig.SystemPrompt,
		loadingMessage:   providerResult.Message,
		providerType:     providerType,
		streamingEnabled: agentConfig.StreamingEnabled,
	}, nil
}

// GenerateWithLoop processes messages with a custom loop that displays tool calls in real-time.
func (a *Agent) GenerateWithLoop(ctx context.Context, messages []fantasy.Message,
	onToolCall ToolCallHandler, onToolExecution ToolExecutionHandler, onToolResult ToolResultHandler,
	onResponse ResponseHandler, onToolCallContent ToolCallContentHandler, onToolApproval ToolApprovalHandler,
) (*GenerateWithLoopResult, error) {
	return a.GenerateWithLoopAndStreaming(ctx, messages, onToolCall, onToolExecution, onToolResult,
		onResponse, onToolCallContent, nil, onToolApproval)
}

// GenerateWithLoopAndStreaming processes messages using the fantasy agent with streaming and callbacks.
// Fantasy handles the tool call loop internally. We map fantasy's rich callback system
// to mcphost's existing callback interface for UI integration.
func (a *Agent) GenerateWithLoopAndStreaming(ctx context.Context, messages []fantasy.Message,
	onToolCall ToolCallHandler, onToolExecution ToolExecutionHandler, onToolResult ToolResultHandler,
	onResponse ResponseHandler, onToolCallContent ToolCallContentHandler,
	onStreamingResponse StreamingResponseHandler, onToolApproval ToolApprovalHandler,
) (*GenerateWithLoopResult, error) {

	// Fantasy requires the current user input as Prompt, with prior messages as history.
	// Extract the last user message text as the prompt, and pass everything before it as Messages.
	prompt, history := splitPromptAndHistory(messages)

	// Track current tool call info for callbacks
	var currentToolName string
	var currentToolArgs string

	if a.streamingEnabled {
		// Use fantasy's streaming agent
		result, err := a.fantasyAgent.Stream(ctx, fantasy.AgentStreamCall{
			Prompt:   prompt,
			Messages: history,

			// Text streaming callback
			OnTextDelta: func(id, text string) error {
				if onStreamingResponse != nil {
					onStreamingResponse(text)
				}
				return nil
			},

			// Tool call complete - the tool has been parsed and is about to execute
			OnToolCall: func(tc fantasy.ToolCallContent) error {
				currentToolName = tc.ToolName
				currentToolArgs = tc.Input

				// Check approval if handler is set
				if onToolApproval != nil {
					approved, err := onToolApproval(tc.ToolName, tc.Input)
					if err != nil {
						return err
					}
					if !approved {
						return fmt.Errorf("tool call %s rejected by user", tc.ToolName)
					}
				}

				// Notify about the tool call
				if onToolCall != nil {
					onToolCall(tc.ToolName, tc.Input)
				}

				// Notify tool execution starting
				if onToolExecution != nil {
					onToolExecution(tc.ToolName, true)
				}

				return nil
			},

			// Tool result - tool execution completed
			OnToolResult: func(tr fantasy.ToolResultContent) error {
				// Notify tool execution finished
				if onToolExecution != nil {
					onToolExecution(tr.ToolName, false)
				}

				if onToolResult != nil {
					// Extract result text and error status
					resultText, isError := extractToolResultText(tr)
					onToolResult(tr.ToolName, currentToolArgs, resultText, isError)
				}

				return nil
			},

			// Step callbacks for content that accompanies tool calls
			OnStepFinish: func(step fantasy.StepResult) error {
				// Check if step has text content alongside tool calls
				text := step.Content.Text()
				toolCalls := step.Content.ToolCalls()
				if text != "" && len(toolCalls) > 0 && onToolCallContent != nil {
					onToolCallContent(text)
				}
				return nil
			},
		})
		if err != nil {
			return nil, err
		}

		return convertAgentResult(result, messages), nil
	}

	// Non-streaming path
	result, err := a.fantasyAgent.Generate(ctx, fantasy.AgentCall{
		Prompt:   prompt,
		Messages: history,
	})
	if err != nil {
		return nil, err
	}

	// For non-streaming, fire the response callback with the final text
	if onResponse != nil && result.Response.Content.Text() != "" {
		onResponse(result.Response.Content.Text())
	}

	_ = currentToolName // satisfy compiler for non-streaming path

	return convertAgentResult(result, messages), nil
}

// splitPromptAndHistory extracts the last user message as the prompt string,
// and returns everything before it as conversation history. Fantasy's agent
// requires the current turn's input as Prompt (string), with prior messages
// passed separately as Messages (history).
func splitPromptAndHistory(messages []fantasy.Message) (string, []fantasy.Message) {
	if len(messages) == 0 {
		return "", nil
	}

	// Walk backwards to find the last user message
	for i := len(messages) - 1; i >= 0; i-- {
		if messages[i].Role == fantasy.MessageRoleUser {
			// Extract text from the user message parts
			var prompt string
			for _, part := range messages[i].Content {
				if tp, ok := part.(fantasy.TextPart); ok {
					prompt = tp.Text
					break
				}
			}
			// History is everything except this last user message
			history := make([]fantasy.Message, 0, len(messages)-1)
			history = append(history, messages[:i]...)
			history = append(history, messages[i+1:]...)
			return prompt, history
		}
	}

	// No user message found — use the last message's text as prompt
	last := messages[len(messages)-1]
	for _, part := range last.Content {
		if tp, ok := part.(fantasy.TextPart); ok {
			return tp.Text, messages[:len(messages)-1]
		}
	}

	return "", messages
}

// convertAgentResult converts a fantasy AgentResult to our GenerateWithLoopResult.
func convertAgentResult(result *fantasy.AgentResult, originalMessages []fantasy.Message) *GenerateWithLoopResult {
	// Collect all conversation messages: original + all step messages
	var allMessages []fantasy.Message
	allMessages = append(allMessages, originalMessages...)

	for _, step := range result.Steps {
		allMessages = append(allMessages, step.Messages...)
	}

	return &GenerateWithLoopResult{
		FinalResponse:        &result.Response,
		ConversationMessages: allMessages,
		TotalUsage:           result.TotalUsage,
	}
}

// extractToolResultText extracts the text and error status from a fantasy ToolResultContent.
func extractToolResultText(tr fantasy.ToolResultContent) (string, bool) {
	if tr.Result == nil {
		return "", false
	}

	// Marshal the result to JSON for display
	resultBytes, err := json.Marshal(tr.Result)
	if err != nil {
		return fmt.Sprintf("%v", tr.Result), false
	}

	resultText := string(resultBytes)

	// Check if this is an error result by examining the type
	if errResult, ok := tr.Result.(fantasy.ToolResultOutputContentError); ok {
		return errResult.Error.Error(), true
	}

	return resultText, false
}

// GetTools returns the list of available tools loaded in the agent.
func (a *Agent) GetTools() []fantasy.AgentTool {
	return a.toolManager.GetTools()
}

// GetLoadingMessage returns the loading message from provider creation.
func (a *Agent) GetLoadingMessage() string {
	return a.loadingMessage
}

// GetLoadedServerNames returns the names of successfully loaded MCP servers.
func (a *Agent) GetLoadedServerNames() []string {
	return a.toolManager.GetLoadedServerNames()
}

// GetModel returns the underlying fantasy LanguageModel.
func (a *Agent) GetModel() fantasy.LanguageModel {
	return a.model
}

// Close closes the agent and cleans up resources.
func (a *Agent) Close() error {
	return a.toolManager.Close()
}

// escListenerModel is a simple Bubble Tea model for ESC key detection
type escListenerModel struct {
	escPressed chan bool
}

func (m escListenerModel) Init() tea.Cmd {
	return nil
}

func (m escListenerModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
	if msg, ok := msg.(tea.KeyPressMsg); ok {
		if msg.String() == "esc" {
			select {
			case m.escPressed <- true:
			default:
			}
			return m, tea.Quit
		}
	}
	return m, nil
}

func (m escListenerModel) View() tea.View {
	return tea.NewView("")
}

// listenForESC listens for ESC key press using Bubble Tea and returns true if detected
func (a *Agent) listenForESC(stopChan chan bool, readyChan chan bool) bool {
	escPressed := make(chan bool, 1)

	model := escListenerModel{
		escPressed: escPressed,
	}

	p := tea.NewProgram(model, tea.WithoutRenderer())

	go func() {
		if _, err := p.Run(); err != nil {
			select {
			case escPressed <- false:
			default:
			}
		}
	}()

	go func() {
		time.Sleep(10 * time.Millisecond)
		select {
		case readyChan <- true:
		default:
		}
	}()

	select {
	case <-stopChan:
		p.Kill()
		time.Sleep(50 * time.Millisecond)
		return false
	case pressed := <-escPressed:
		p.Kill()
		time.Sleep(50 * time.Millisecond)
		return pressed
	case <-time.After(30 * time.Second):
		p.Kill()
		time.Sleep(50 * time.Millisecond)
		return false
	}
}