ai-agent/internal/agent/agent.go

package agent

import (
    "sync"
    "time"

    "ai-agent/internal/config"
    "ai-agent/internal/ice"
    "ai-agent/internal/llm"
    "ai-agent/internal/mcp"
    "ai-agent/internal/memory"
    "ai-agent/internal/permission"
)

type Agent struct {
    mu          sync.RWMutex
    llmClient    llm.Client
    registry     *mcp.Registry
    messages     []llm.Message
    skillContent string
    loadedCtx    string
    numCtx       int
    memoryStore  *memory.Store
    iceEngine    *ice.Engine
    router       *config.Router
    modePrefix   string
    toolsEnabled      bool
    workDir           string
    ignoreContent     string
    permChecker       *permission.Checker
    approvalCallback  func(permission.ApprovalRequest)
    toolsConfig       config.ToolsConfig
}

func New(llmClient llm.Client, registry *mcp.Registry, numCtx int) *Agent {
    return &Agent{
        llmClient:    llmClient,
        registry:     registry,
        numCtx:       numCtx,
        toolsEnabled: true,
    }
}

func (a *Agent) SetRouter(router *config.Router) {
    a.router = router
}

func (a *Agent) SetModeContext(prefix string, allowTools bool) {
    a.modePrefix = prefix
    a.toolsEnabled = allowTools
}

func (a *Agent) AppendLoadedContext(content string) {
    if a.loadedCtx == "" {
        a.loadedCtx = content
    } else {
        a.loadedCtx += content
    }
}

func (a *Agent) Router() *config.Router {
    return a.router
}

func (a *Agent) NumCtx() int {
    return a.numCtx
}

func (a *Agent) SetMemoryStore(store *memory.Store) {
    a.memoryStore = store
}

func (a *Agent) AddUserMessage(content string) {
    a.mu.Lock()
    defer a.mu.Unlock()
    a.messages = append(a.messages, llm.Message{
        Role:    "user",
        Content: content,
    })
}

func (a *Agent) Messages() []llm.Message {
    a.mu.RLock()
    defer a.mu.RUnlock()
    return a.messages
}

func (a *Agent) ClearHistory() {
    a.mu.Lock()
    defer a.mu.Unlock()
    a.messages = nil
}

func (a *Agent) AppendMessage(msg llm.Message) {
    a.mu.Lock()
    defer a.mu.Unlock()
    a.messages = append(a.messages, msg)
}

func (a *Agent) ReplaceMessages(msgs []llm.Message) {
    a.mu.Lock()
    defer a.mu.Unlock()
    a.messages = msgs
}

func (a *Agent) SetSkillContent(content string) {
    a.skillContent = content
}

func (a *Agent) SetLoadedContext(content string) {
    a.loadedCtx = content
}

func (a *Agent) Model() string {
    return a.llmClient.Model()
}

func (a *Agent) LLMClient() llm.Client {
    return a.llmClient
}

func (a *Agent) ToolCount() int {
    count := a.registry.ToolCount()
    if a.memoryStore != nil {
        count += 2
    }
    return count
}

func (a *Agent) ServerCount() int {
    return a.registry.ServerCount()
}

func (a *Agent) ServerNames() []string {
    return a.registry.ServerNames()
}

func (a *Agent) SetWorkDir(dir string) {
    a.workDir = dir
}

func (a *Agent) SetIgnoreContent(content string) {
    a.ignoreContent = content
}

func (a *Agent) SetPermissionChecker(checker *permission.Checker) {
    a.permChecker = checker
}

func (a *Agent) SetApprovalCallback(cb func(permission.ApprovalRequest)) {
    a.approvalCallback = cb
}

func (a *Agent) SetICEEngine(engine *ice.Engine) {
    a.iceEngine = engine
}

func (a *Agent) ICEEngine() *ice.Engine {
    return a.iceEngine
}

func (a *Agent) SetToolsConfig(cfg config.ToolsConfig) {
    a.toolsConfig = cfg
}

func (a *Agent) MaxIterations() int {
    if a.toolsConfig.MaxIterations > 0 {
        return a.toolsConfig.MaxIterations
    }
    return 10
}

func (a *Agent) ToolTimeout() time.Duration {
    if a.toolsConfig.Timeout != "" {
        if d, err := time.ParseDuration(a.toolsConfig.Timeout); err == nil {
            return d
        }
    }
    return 30 * time.Second
}

func (a *Agent) MaxGrepResults() int {
    if a.toolsConfig.MaxGrepResults > 0 {
        return a.toolsConfig.MaxGrepResults
    }
    return 500
}

func (a *Agent) Close() {
    if a.iceEngine != nil {
        _ = a.iceEngine.Flush()
    }
    a.registry.Close()
}