diff --git a/pkg/agent/context.go b/pkg/agent/context.go index cfa40e69a..13f52b44f 100644 --- a/pkg/agent/context.go +++ b/pkg/agent/context.go @@ -89,32 +89,14 @@ func (cb *ContextBuilder) getIdentity() string { // Build tools section dynamically toolsSection := cb.buildToolsSection() - // Build prompt with optional orchestration banner - var prompt string - if cb.orchestrationEnabled { - prompt = ` /_/_/_/_/_/_/_/_/_/_/_/_/_/_/ - - O R C H E S T R A M O D E - -/_/_/_/_/_/_/_/_/_/_/_/_/_/_/ - - - -` + // Orchestration overrides (banner, identity, executing rule) + prompt, identity, executingRule := cb.extIdentityOverrides() + if identity == "" { + identity = "a helpful AI assistant" } - - // Conditional identity and plan executing rule for orchestration mode - identity := "a helpful AI assistant" - executingRule := `Work through the current Phase's steps. + if executingRule == "" { + executingRule = `Work through the current Phase's steps. Mark each "- [x]" via edit_file. The system will auto-advance phases.` - if cb.orchestrationEnabled { - identity = "a conductor AI agent that orchestrates subagents" - executingRule = `Delegate the current Phase's steps to subagents using spawn. - For each step: spawn a subagent with the appropriate preset (scout for investigation, - coder for implementation, analyst for review). Spawn multiple independent steps in parallel. - When a subagent completes, mark "- [x]" via edit_file and record findings in - ## Orchestration > Findings in MEMORY.md. - Only do a step inline if it's a single quick tool call (e.g., reading one file).` } return fmt.Sprintf(prompt+`# picoclaw 🦞 (%s) @@ -175,29 +157,6 @@ Your workspace is at: %s toolsSection, executingRule, toolDiscovery) } -func (cb *ContextBuilder) buildToolsSection() string { - if cb.tools == nil { - return "" - } - - summaries := cb.tools.GetSummaries() - if len(summaries) == 0 { - return "" - } - - var sb strings.Builder - sb.WriteString("## Available Tools\n\n") - sb.WriteString( - "**CRITICAL**: You MUST use tools to perform actions. Do NOT pretend to execute commands or schedule tasks.\n\n", - ) - sb.WriteString("You have access to the following tools:\n\n") - for _, s := range summaries { - sb.WriteString(s) - sb.WriteString("\n") - } - return sb.String() -} - func (cb *ContextBuilder) getDiscoveryRule() string { if !cb.toolDiscoveryBM25 && !cb.toolDiscoveryRegex { return "" @@ -223,12 +182,8 @@ func (cb *ContextBuilder) BuildSystemPrompt() string { // Core identity section parts = append(parts, cb.getIdentity()) - // Orchestration guidance — injected only when spawn tool is registered - if cb.tools != nil { - if _, hasSpawn := cb.tools.Get("spawn"); hasSpawn { - parts = append(parts, orchestrationGuidance) - } - } + // Fork-specific prompt sections (orchestration guidance, peer note) + parts = append(parts, cb.extPromptSections()...) // Bootstrap files bootstrapContent := cb.LoadBootstrapFiles() @@ -253,11 +208,6 @@ The following skills extend your capabilities. To use a skill, read its SKILL.md } } - // Peer session coordination - if cb.peerNote != "" { - parts = append(parts, "## Active Sessions\n\n"+cb.peerNote) - } - // Memory context memoryContext := cb.memory.GetMemoryContext() if memoryContext != "" { diff --git a/pkg/agent/context_ext.go b/pkg/agent/context_ext.go index 61b245488..5355612e1 100644 --- a/pkg/agent/context_ext.go +++ b/pkg/agent/context_ext.go @@ -1,6 +1,10 @@ package agent -import "github.com/sipeed/picoclaw/pkg/tools" +import ( + "strings" + + "github.com/sipeed/picoclaw/pkg/tools" +) // contextBuilderExt holds fork-specific fields for ContextBuilder. // Embedded in ContextBuilder so existing field access (cb.workDir, cb.tools, etc.) continues to work. @@ -33,6 +37,77 @@ func (cb *ContextBuilder) SetOrchestrationEnabled(enabled bool) { cb.orchestrationEnabled = enabled } +// buildToolsSection generates the "Available Tools" section for the system prompt. +func (cb *ContextBuilder) buildToolsSection() string { + if cb.tools == nil { + return "" + } + + summaries := cb.tools.GetSummaries() + if len(summaries) == 0 { + return "" + } + + var sb strings.Builder + sb.WriteString("## Available Tools\n\n") + sb.WriteString( + "**CRITICAL**: You MUST use tools to perform actions. Do NOT pretend to execute commands or schedule tasks.\n\n", + ) + sb.WriteString("You have access to the following tools:\n\n") + for _, s := range summaries { + sb.WriteString(s) + sb.WriteString("\n") + } + return sb.String() +} + +// extIdentityOverrides returns the orchestration-specific overrides for +// getIdentity: banner prefix, identity string, and plan executing rule. +// When orchestration is disabled, all return values are empty strings. +func (cb *ContextBuilder) extIdentityOverrides() (banner, identity, executingRule string) { + if !cb.orchestrationEnabled { + return "", "", "" + } + + banner = ` /_/_/_/_/_/_/_/_/_/_/_/_/_/_/ + + O R C H E S T R A M O D E + +/_/_/_/_/_/_/_/_/_/_/_/_/_/_/ + + + +` + identity = "a conductor AI agent that orchestrates subagents" + executingRule = `Delegate the current Phase's steps to subagents using spawn. + For each step: spawn a subagent with the appropriate preset (scout for investigation, + coder for implementation, analyst for review). Spawn multiple independent steps in parallel. + When a subagent completes, mark "- [x]" via edit_file and record findings in + ## Orchestration > Findings in MEMORY.md. + Only do a step inline if it's a single quick tool call (e.g., reading one file).` + return banner, identity, executingRule +} + +// extPromptSections returns fork-specific prompt sections to append to +// BuildSystemPrompt: orchestration guidance and peer session note. +func (cb *ContextBuilder) extPromptSections() []string { + var sections []string + + // Orchestration guidance — injected only when spawn tool is registered + if cb.tools != nil { + if _, hasSpawn := cb.tools.Get("spawn"); hasSpawn { + sections = append(sections, orchestrationGuidance) + } + } + + // Peer session coordination + if cb.peerNote != "" { + sections = append(sections, "## Active Sessions\n\n"+cb.peerNote) + } + + return sections +} + // Memory returns the underlying MemoryStore for direct plan queries. func (cb *ContextBuilder) Memory() *MemoryStore { return cb.memory diff --git a/pkg/agent/instance.go b/pkg/agent/instance.go index 0ac3d8894..3bd0fa7e0 100644 --- a/pkg/agent/instance.go +++ b/pkg/agent/instance.go @@ -9,7 +9,6 @@ import ( "strings" "github.com/sipeed/picoclaw/pkg/config" - "github.com/sipeed/picoclaw/pkg/git" "github.com/sipeed/picoclaw/pkg/providers" "github.com/sipeed/picoclaw/pkg/routing" "github.com/sipeed/picoclaw/pkg/session" @@ -38,6 +37,8 @@ type AgentInstance struct { Sessions *session.LegacyAdapter ContextBuilder *ContextBuilder Tools *tools.ToolRegistry + Subagents *config.SubagentsConfig + SkillsFilter []string Candidates []providers.FallbackCandidate PlanModel string PlanFallbacks []string @@ -141,23 +142,10 @@ func NewAgentInstance( agentID := routing.DefaultAgentID agentName := "" - var subagents *config.SubagentsConfig - var skillsFilter []string if agentCfg != nil { agentID = routing.NormalizeAgentID(agentCfg.ID) agentName = agentCfg.Name - subagents = agentCfg.Subagents - skillsFilter = agentCfg.Skills - } - - // Apply defaults.Orchestration: if the flag is set, ensure orchestration is enabled. - if defaults.Orchestration { - if subagents == nil { - subagents = &config.SubagentsConfig{Enabled: true} - } else { - subagents.Enabled = true - } } maxIter := defaults.MaxToolIterations @@ -243,19 +231,6 @@ func NewAgentInstance( candidates := providers.ResolveCandidatesWithLookup(modelCfg, defaults.Provider, resolveFromModelList) - // Resolve plan model (for interviewing/review phases) - planModel := resolvePlanModel(agentCfg, defaults) - planFallbacks := resolvePlanFallbacks(agentCfg, defaults) - - var planCandidates []providers.FallbackCandidate - if planModel != "" { - planModelCfg := providers.ModelConfig{ - Primary: planModel, - Fallbacks: planFallbacks, - } - planCandidates = providers.ResolveCandidates(planModelCfg, defaults.Provider) - } - // Model routing setup: pre-resolve light model candidates at creation time // to avoid repeated model_list lookups on every incoming message. var router *routing.Router @@ -275,17 +250,7 @@ func NewAgentInstance( } } - // Startup cleanup: prune orphaned worktrees - worktreesDir := filepath.Join(workspace, ".worktrees") - if repoRoot := git.FindRepoRoot(workspace); repoRoot != "" { - git.PruneOrphaned(repoRoot, worktreesDir) - } - - return &AgentInstance{ - instanceExt: instanceExt{ - Subagents: subagents, - SkillsFilter: skillsFilter, - }, + agent := &AgentInstance{ ID: agentID, Name: agentName, Model: model, @@ -304,12 +269,14 @@ func NewAgentInstance( ContextBuilder: contextBuilder, Tools: toolsRegistry, Candidates: candidates, - PlanModel: planModel, - PlanFallbacks: planFallbacks, - PlanCandidates: planCandidates, Router: router, LightCandidates: lightCandidates, } + + // Initialize fork-specific fields (subagents, plan model, worktree pruning). + agent.initInstanceExt(agentCfg, defaults, cfg) + + return agent } // resolveAgentWorkspace determines the workspace directory for an agent. diff --git a/pkg/agent/instance_ext.go b/pkg/agent/instance_ext.go index 10aa835e9..830a3bb2c 100644 --- a/pkg/agent/instance_ext.go +++ b/pkg/agent/instance_ext.go @@ -2,10 +2,12 @@ package agent import ( "fmt" + "path/filepath" "sync" "github.com/sipeed/picoclaw/pkg/config" "github.com/sipeed/picoclaw/pkg/git" + "github.com/sipeed/picoclaw/pkg/providers" "github.com/sipeed/picoclaw/pkg/tools" ) @@ -16,9 +18,6 @@ type instanceExt struct { // Used by runAgentLoop to wait for spawned subagents before worktree cleanup. SubagentMgr *tools.SubagentManager - Subagents *config.SubagentsConfig - SkillsFilter []string - // Interview staleness tracking: consecutive turns where MEMORY.md was not updated. interviewStaleCount int interviewMemoryLen int @@ -28,6 +27,47 @@ type instanceExt struct { worktreeMu sync.RWMutex } +// initInstanceExt initializes fork-specific fields: subagents config, +// skills filter, plan model resolution, and worktree pruning. +func (ai *AgentInstance) initInstanceExt( + agentCfg *config.AgentConfig, + defaults *config.AgentDefaults, + cfg *config.Config, +) { + // Extract subagents and skills filter from agent config + if agentCfg != nil { + ai.Subagents = agentCfg.Subagents + ai.SkillsFilter = agentCfg.Skills + } + + // Apply defaults.Orchestration: if the flag is set, ensure orchestration is enabled. + if defaults.Orchestration { + if ai.Subagents == nil { + ai.Subagents = &config.SubagentsConfig{Enabled: true} + } else { + ai.Subagents.Enabled = true + } + } + + // Resolve plan model (for interviewing/review phases) + ai.PlanModel = resolvePlanModel(agentCfg, defaults) + ai.PlanFallbacks = resolvePlanFallbacks(agentCfg, defaults) + + if ai.PlanModel != "" { + planModelCfg := providers.ModelConfig{ + Primary: ai.PlanModel, + Fallbacks: ai.PlanFallbacks, + } + ai.PlanCandidates = providers.ResolveCandidates(planModelCfg, defaults.Provider) + } + + // Startup cleanup: prune orphaned worktrees + worktreesDir := filepath.Join(ai.Workspace, ".worktrees") + if repoRoot := git.FindRepoRoot(ai.Workspace); repoRoot != "" { + git.PruneOrphaned(repoRoot, worktreesDir) + } +} + // ActivateWorktree creates a worktree for a session. // projectDir is the git repository to create the worktree in. // If empty, falls back to ai.Workspace. diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index b3e92144e..a16fc5cec 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -8,7 +8,6 @@ package agent import ( "context" - "encoding/json" "errors" "fmt" "path/filepath" @@ -23,16 +22,12 @@ import ( "github.com/sipeed/picoclaw/pkg/commands" "github.com/sipeed/picoclaw/pkg/config" "github.com/sipeed/picoclaw/pkg/constants" - "github.com/sipeed/picoclaw/pkg/git" "github.com/sipeed/picoclaw/pkg/logger" "github.com/sipeed/picoclaw/pkg/media" - "github.com/sipeed/picoclaw/pkg/orch" "github.com/sipeed/picoclaw/pkg/providers" "github.com/sipeed/picoclaw/pkg/routing" - "github.com/sipeed/picoclaw/pkg/session" "github.com/sipeed/picoclaw/pkg/skills" "github.com/sipeed/picoclaw/pkg/state" - "github.com/sipeed/picoclaw/pkg/stats" "github.com/sipeed/picoclaw/pkg/tools" "github.com/sipeed/picoclaw/pkg/utils" "github.com/sipeed/picoclaw/pkg/voice" @@ -139,43 +134,7 @@ func NewAgentLoop( providerCache := make(map[string]providers.LLMProvider) - // Create stats tracker if enabled - - var statsTracker *stats.Tracker - - if len(enableStats) > 0 && enableStats[0] && defaultAgent != nil { - statsTracker = stats.NewTracker(defaultAgent.Workspace) - } - - // Determine if orchestration broadcaster is needed (any agent has subagents enabled). - - // Note: instance.go maps defaults.Orchestration → Subagents.Enabled, so --orchestration - - // is automatically reflected here. - - var orchBroadcaster *orch.Broadcaster - - var orchReporter orch.AgentReporter = orch.Noop - - for _, id := range registry.ListAgentIDs() { - if a, ok := registry.GetAgent(id); ok && a.Subagents != nil && a.Subagents.Enabled { - orchBroadcaster = orch.NewBroadcaster() - - orchReporter = orchBroadcaster - - break - } - } - al := &AgentLoop{ - loopExt: loopExt{ - stats: statsTracker, - sessions: NewSessionTracker(), - orchBroadcaster: orchBroadcaster, - orchReporter: orchReporter, - done: make(chan struct{}), - }, - bus: msgBus, cfg: cfg, @@ -193,12 +152,12 @@ func NewAgentLoop( cmdRegistry: commands.NewRegistry(commands.BuiltinDefinitions()), } - // Register shared tools to all agents (needs al for reporter injection). + // Initialize fork-specific fields (stats, sessions, orchestration, gcLoop). + al.initLoopExt(cfg, registry, len(enableStats) > 0 && enableStats[0]) + // Register shared tools to all agents (needs al for orchestration reporter). registerSharedTools(cfg, msgBus, registry, provider, al) - go al.gcLoop() - return al } @@ -208,7 +167,6 @@ func registerSharedTools( msgBus *bus.MessageBus, registry *AgentRegistry, provider providers.LLMProvider, - al *AgentLoop, ) { for _, agentID := range registry.ListAgentIDs() { @@ -339,86 +297,11 @@ func registerSharedTools( } } - // Spawn tool — only registered when orchestration is explicitly enabled. - - if agent.Subagents != nil && agent.Subagents.Enabled { - webSearchOpts := tools.WebSearchToolOptions{ - BraveAPIKeys: config.MergeAPIKeys(cfg.Tools.Web.Brave.APIKey, cfg.Tools.Web.Brave.APIKeys), - BraveMaxResults: cfg.Tools.Web.Brave.MaxResults, - BraveEnabled: cfg.Tools.Web.Brave.Enabled, - TavilyAPIKeys: config.MergeAPIKeys(cfg.Tools.Web.Tavily.APIKey, cfg.Tools.Web.Tavily.APIKeys), - TavilyBaseURL: cfg.Tools.Web.Tavily.BaseURL, - TavilyMaxResults: cfg.Tools.Web.Tavily.MaxResults, - TavilyEnabled: cfg.Tools.Web.Tavily.Enabled, - DuckDuckGoMaxResults: cfg.Tools.Web.DuckDuckGo.MaxResults, - DuckDuckGoEnabled: cfg.Tools.Web.DuckDuckGo.Enabled, - PerplexityAPIKeys: config.MergeAPIKeys( - cfg.Tools.Web.Perplexity.APIKey, - cfg.Tools.Web.Perplexity.APIKeys, - ), - PerplexityMaxResults: cfg.Tools.Web.Perplexity.MaxResults, - PerplexityEnabled: cfg.Tools.Web.Perplexity.Enabled, - } - - subagentManager := tools.NewSubagentManager( - - provider, - - agent.Model, - - agent.Workspace, - - msgBus, - - al.reporter(), - - webSearchOpts, - ) - - subagentManager.SetLLMOptions(agent.MaxTokens, agent.Temperature) - - // Wire session recorder for DAG persistence. - - recorder := newSessionRecorder(agent.Sessions) - - conductorKey := routing.BuildAgentMainSessionKey(agent.ID) - - subagentManager.SetSessionRecorder(recorder, conductorKey) - - agent.SubagentMgr = subagentManager - - spawnTool := tools.NewSpawnTool(subagentManager) - - currentAgentID := agentID - - spawnTool.SetAllowlistChecker(func(targetAgentID string) bool { - return registry.CanSpawnSubagent(currentAgentID, targetAgentID) - }) - - agent.Tools.Register(spawnTool) - - // Register blocking subagent tool alongside spawn - - subagentTool := tools.NewSubagentTool(subagentManager) - - agent.Tools.Register(subagentTool) - - // Register conductor-side escalation tools (answer questions, review plans) - - agent.Tools.Register(tools.NewAnswerSubagentTool(subagentManager)) - - agent.Tools.Register(tools.NewReviewSubagentPlanTool(subagentManager)) - } + // Orchestration tools (spawn, subagent, answer, review_plan) + registerOrchestrationTools(cfg, agent, agentID, registry, provider, msgBus, al) // Update context builder with the complete tools registry - agent.ContextBuilder.SetToolsRegistry(agent.Tools) - - // Set orchestration mode if enabled - - if agent.Subagents != nil && agent.Subagents.Enabled { - agent.ContextBuilder.SetOrchestrationEnabled(true) - } } } @@ -636,15 +519,7 @@ func (al *AgentLoop) Stop() { // and dirty session data). Should be called during graceful shutdown. func (al *AgentLoop) Close() { - select { - case <-al.done: - - // already closed - - default: - - close(al.done) - } + al.closeExt() mcpManager := al.mcp.takeManager() if mcpManager != nil { @@ -656,18 +531,7 @@ func (al *AgentLoop) Close() { } } - if al.stats != nil { - al.stats.Close() - } - - registry := al.GetRegistry() - for _, agentID := range registry.ListAgentIDs() { - if agent, ok := registry.GetAgent(agentID); ok { - agent.Sessions.Close() - } - } - - registry.Close() + al.GetRegistry().Close() } func (al *AgentLoop) RegisterTool(tool tools.Tool) { @@ -1114,64 +978,8 @@ func (al *AgentLoop) processMessage(ctx context.Context, msg bus.InboundMessage) } // Handle reply-based intervention for active tasks - - if taskID, ok := msg.Metadata["task_id"]; ok && taskID != "" { - if val, found := al.activeTasks.Load(taskID); found { - task := val.(*activeTask) - - content := strings.TrimSpace(msg.Content) - - lower := strings.ToLower(content) - - // Check for stop keywords - - stopKeywords := []string{ - "stop", "cancel", "abort", - - "停止", "中止", "やめて", //nolint:gosmopolitan // intentional CJK stop words - - } - - isStop := false - - for _, kw := range stopKeywords { - if lower == kw { - isStop = true - - break - } - } - - if isStop { - task.cancel() - - logger.InfoCF("agent", "Task canceled by user intervention", - - map[string]any{"task_id": taskID}) - - return "Task canceled.", nil - } - - // Inject message into interrupt channel for the tool loop - - select { - case task.interrupt <- content: - - logger.InfoCF("agent", "User intervention queued", - - map[string]any{"task_id": taskID, "content": utils.Truncate(content, 80)}) - - default: - - logger.WarnCF("agent", "Interrupt channel full, message dropped", - - map[string]any{"task_id": taskID}) - } - - return "Intervention sent to running task.", nil - } - - // Task not found — fall through to normal processing + if response, handled := al.handleTaskIntervention(msg); handled { + return response, nil } // Route system messages to processSystemMessage @@ -1185,23 +993,8 @@ func (al *AgentLoop) processMessage(ctx context.Context, msg bus.InboundMessage) al.OnUserMessage() } - // Expand /skill command: inject SKILL.md content into message, then continue to LLM - - var expansionCompact string - - if expanded, compact, ok := al.expandSkillCommand(msg); ok { - msg.Content = expanded - - expansionCompact = compact - } - - // Expand /plan : write interview seed, rewrite for LLM interview - - if expanded, compact, ok := al.expandPlanCommand(msg); ok { - msg.Content = expanded - - expansionCompact = compact - } + // Expand fork-specific /skill and /plan commands + expansionCompact := al.expandForkCommands(&msg) // Check for commands @@ -1301,773 +1094,6 @@ func (al *AgentLoop) withTelegramThread(channel, chatID string, threadID int) st return fmt.Sprintf("%s/%d", baseChatID, threadID) } -func (al *AgentLoop) runAgentLoop(ctx context.Context, agent *AgentInstance, opts processOptions) (string, error) { - // -1. Acquire per-session lock to prevent concurrent access on the same session - - if !al.acquireSessionLock(ctx, opts.SessionKey) { - return "", fmt.Errorf("context canceled while waiting for session lock") - } - - defer al.releaseSessionLock(opts.SessionKey) - - // Report session lifecycle to canvas. - - al.reporter().ReportSpawn(opts.SessionKey, opts.Channel, opts.UserMessage) - - defer al.reporter().ReportGC(opts.SessionKey, "completed") - - // -0. Create cancelable child context and register active task - - taskCtx, taskCancel := context.WithCancel(ctx) - - defer taskCancel() - - task := &activeTask{ - Description: utils.Truncate(opts.UserMessage, 80), - - MaxIter: agent.MaxIterations, - - StartedAt: time.Now(), - - cancel: taskCancel, - - interrupt: make(chan string, 1), - } - - // Guarantee heartbeat worktree cleanup on ALL exit paths (error, panic, normal). - - // Wait for spawned subagents first so they aren't killed mid-flight. - - // After auto-commit, attempt to merge the worktree branch into main. - - defer func() { - if opts.Background { - if agent.SubagentMgr != nil { - agent.SubagentMgr.WaitAll(35 * time.Minute) // slightly above spawnTimeout - } - - wt := agent.GetWorktree(opts.SessionKey) - - if wt != nil { - // 1. Auto-commit uncommitted changes in worktree - - if git.HasUncommittedChanges(wt.Path) { - _ = git.AutoCommit(wt.Path, "heartbeat: auto-save") - } - - // 2. Check if there are unique commits worth merging - - repoRoot := git.FindRepoRoot(agent.Workspace) - - ahead := git.CommitsAhead(repoRoot, wt.BaseBranch, wt.Branch) - - if ahead > 0 && repoRoot != "" { - // 3. Try fast-forward merge into base branch - - mr := git.MergeWorktreeBranch(repoRoot, wt) - - // 4. Notify based on merge result - - if !constants.IsInternalChannel(opts.Channel) { - cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 5*time.Second) - - if mr.Merged { - _ = al.bus.PublishOutbound(cleanupCtx, bus.OutboundMessage{ - Channel: opts.Channel, - - ChatID: opts.ChatID, - - Content: fmt.Sprintf("Heartbeat: merged %d commit(s) to %s.", - - ahead, wt.BaseBranch), - }) - } else if mr.Conflict { - _ = al.bus.PublishOutbound(cleanupCtx, bus.OutboundMessage{ - Channel: opts.Channel, - - ChatID: opts.ChatID, - - Content: fmt.Sprintf("Heartbeat: merge conflict on branch `%s` — manual merge needed.", - - mr.Branch), - }) - } - - cleanupCancel() - } - } - - // 5. Dispose worktree (branch auto-deleted if merged, kept if conflict) - - agent.DeactivateWorktree(opts.SessionKey, "", false) - } - } - }() - - // For background tasks (cron/heartbeat), generate a TaskID and send notification - - isBackgroundTask := opts.Background && al.state != nil - - if isBackgroundTask && opts.TaskID == "" { - opts.TaskID = fmt.Sprintf("task-%s-%d", opts.SessionKey, time.Now().UnixMilli()) - - // Determine notification channel: use opts.Channel if already a real channel, - - // otherwise resolve from last active channel - - notifyChannel := opts.Channel - - notifyChatID := opts.ChatID - - if constants.IsInternalChannel(notifyChannel) || notifyChannel == "" { - if lastChannel := al.state.GetLastChannel(); lastChannel != "" { - if idx := strings.Index(lastChannel, ":"); idx > 0 { - notifyChannel = lastChannel[:idx] - - notifyChatID = lastChannel[idx+1:] - } - } - } - - if notifyChannel != "" && notifyChatID != "" && !constants.IsInternalChannel(notifyChannel) { - // Override opts channel/chatID for status updates - - opts.Channel = notifyChannel - - opts.ChatID = notifyChatID - - // Send initial task notification - - _ = al.bus.PublishOutbound(ctx, bus.OutboundMessage{ - Channel: notifyChannel, - - ChatID: notifyChatID, - - Content: fmt.Sprintf("\U0001F916 Background task started\n%s", task.Description), - - IsTaskStatus: true, - - TaskID: opts.TaskID, - }) - } - } - - // Shared variable for capturing LLM's final response. The defer below reads it - - // to include the response in the task completion message. - - var finalContent string - - // Use TaskID as key if available (for background tasks), else sessionKey - - taskKey := opts.SessionKey - - if opts.TaskID != "" { - taskKey = opts.TaskID - } - - al.activeTasks.Store(taskKey, task) - - defer func() { - al.activeTasks.Delete(taskKey) - - // Publish final task status on completion for background tasks. - - // Include finalContent so the LLM response appears in the same bubble - - // as the completion status, avoiding duplicate messages. - - if opts.TaskID != "" { - elapsed := time.Since(task.StartedAt) - - completionMsg := fmt.Sprintf("\u2705 Task completed (%.1fs)", elapsed.Seconds()) - - // Determine the best content to show in the completion bubble. - - // Priority: message tool content > finalContent > task.Result - - task.mu.Lock() - - msgContent := task.messageContent - - task.mu.Unlock() - - var resultContent string - - switch { - case msgContent != "": - - // The message tool already sent this to the user via the - - // task bubble; re-include it so the completion doesn't erase it. - - resultContent = msgContent - - case finalContent != "" && finalContent != defaultResponse && finalContent != "HEARTBEAT_OK": - - resultContent = finalContent - - default: - - summary := task.Result - - if summary == "" { - summary = task.Description - } - - resultContent = summary - } - - if resultContent != "" { - combined := completionMsg + "\n\n" + resultContent - - if len([]rune(combined)) <= 4096 { - completionMsg = combined - } else { - // Too long for one bubble: send header as task status, - - // body as regular message (auto-split by SplitMessage). - - doneCtx, doneCancel := context.WithTimeout(context.Background(), 5*time.Second) - - _ = al.bus.PublishOutbound(doneCtx, bus.OutboundMessage{ - Channel: opts.Channel, - - ChatID: opts.ChatID, - - Content: completionMsg, - - IsTaskStatus: true, - - TaskID: opts.TaskID, - - Final: true, - }) - - _ = al.bus.PublishOutbound(doneCtx, bus.OutboundMessage{ - Channel: opts.Channel, - - ChatID: opts.ChatID, - - Content: resultContent, - }) - - doneCancel() - - return - } - } - - doneCtx, doneCancel := context.WithTimeout(context.Background(), 5*time.Second) - - _ = al.bus.PublishOutbound(doneCtx, bus.OutboundMessage{ - Channel: opts.Channel, - - ChatID: opts.ChatID, - - Content: completionMsg, - - IsTaskStatus: true, - - TaskID: opts.TaskID, - - Final: true, - }) - - doneCancel() - } - }() - - // Replace ctx with the cancelable child context - - ctx = taskCtx - - // 0. Record last channel for heartbeat notifications (skip internal channels) - - if opts.Channel != "" && opts.ChatID != "" { - // Don't record internal channels (cli, system, subagent) - - if !constants.IsInternalChannel(opts.Channel) { - channelKey := fmt.Sprintf("%s:%s", opts.Channel, opts.ChatID) - if err := al.RecordLastChannel(channelKey); err != nil { - logger.WarnCF("agent", "Failed to record last channel", map[string]any{"error": err.Error()}) - } - - if err := al.RecordLastHeartbeatTarget(channelKey); err != nil { - logger.WarnCF("agent", "Failed to record last heartbeat target", map[string]any{"error": err.Error()}) - } - } - } - - // 1. Update tool contexts - - al.updateToolContexts(agent, opts.Channel, opts.ChatID) - - // 1-bis. For background tasks that don't send a final response (e.g. heartbeat), - - // redirect the message tool to publish as IsTaskStatus so its output lands in - - // the same bubble as the task status instead of creating a separate message. - - if opts.Background && !opts.SendResponse && opts.TaskID != "" { - if tool, ok := agent.Tools.Get("message"); ok { - if mt, ok := tool.(*tools.MessageTool); ok { - taskID := opts.TaskID - - mt.SetSendCallback(func(channel, chatID, content string) error { - // Capture the message tool's content so the completion - - // defer can include it instead of losing it to an overwrite. - - if task != nil { - task.mu.Lock() - - task.messageContent = content - - task.mu.Unlock() - } - - pubCtx, pubCancel := context.WithTimeout(context.Background(), 5*time.Second) - - defer pubCancel() - - return al.bus.PublishOutbound(pubCtx, bus.OutboundMessage{ - Channel: channel, - - ChatID: chatID, - - Content: content, - - IsTaskStatus: true, - - TaskID: taskID, - }) - }) - } - } - } - - // 1a. Set session-specific working directory for bootstrap file lookup. - - // Prefer the tool-detected project directory (touch_dir) from the session tracker, - - // resolved as an absolute path under workspace. Fall back to worktree or workspace. - - if active := al.sessions.ListActive(); len(active) > 0 && active[0].SessionKey == opts.SessionKey && - - active[0].TouchDir != "" { - agent.ContextBuilder.SetWorkDir(filepath.Join(agent.Workspace, active[0].TouchDir)) - } else { - agent.ContextBuilder.SetWorkDir(agent.EffectiveWorkspace(opts.SessionKey)) - } - - // 1b. Inject peer session awareness into system prompt - - projectPath := agent.ContextBuilder.GetPlanWorkDir() - - if projectPath == "" { - projectPath = agent.Workspace - } - - peers := al.sessions.GetPeerPurposes(opts.SessionKey, projectPath) - - if len(peers) > 0 { - var peerNote strings.Builder - - peerNote.WriteString("Other sessions working on this project:\n") - - for _, p := range peers { - peerNote.WriteString(fmt.Sprintf("- %s: %s (branch: %s)\n", p.SessionKey, p.Purpose, p.Branch)) - } - - peerNote.WriteString("\nAvoid conflicting changes with these sessions.") - - agent.ContextBuilder.SetPeerNote(peerNote.String()) - } else { - agent.ContextBuilder.SetPeerNote("") - } - - // 2. Build messages (skip history for heartbeat) - - var history []providers.Message - var summary string - if !opts.NoHistory { - history = agent.Sessions.GetHistory(opts.SessionKey) - summary = agent.Sessions.GetSummary(opts.SessionKey) - - // Sanitize history to remove orphaned tool calls (from crashes/session collisions) - - var removedCount int - - history, removedCount = session.SanitizeHistory(history) - - if removedCount > 0 { - logger.WarnCF("agent", "Sanitized session history: removed orphaned messages", - - map[string]any{ - "session_key": opts.SessionKey, - - "removed_count": removedCount, - }) - - // Persist the sanitized history - - agent.Sessions.SetHistory(opts.SessionKey, history) - - _ = agent.Sessions.Save(opts.SessionKey) - } - } - messages := agent.ContextBuilder.BuildMessages( - history, - summary, - opts.UserMessage, - - opts.Media, - - opts.Channel, - opts.ChatID, - ) - - // Resolve media:// refs: images→base64 data URLs, non-images→local paths in content - cfg := al.GetConfig() - maxMediaSize := cfg.Agents.Defaults.GetMaxMediaSize() - messages = resolveMediaRefs(messages, al.mediaStore, maxMediaSize) - - // 2b. Interview staleness nudge: if MEMORY.md hasn't been updated for - - // several consecutive turns, inject a reminder so the AI writes its findings. - - const interviewStaleThreshold = 2 - - if agent.ContextBuilder.GetPlanStatus() == "interviewing" && agent.interviewStaleCount >= interviewStaleThreshold { - messages = append(messages, providers.Message{ - Role: "user", - - Content: "[System] You have been interviewing for several turns without updating memory/MEMORY.md. Please use edit_file now to save your findings to the ## Context section, or organize the plan into ## Phase sections with `- [ ]` checkbox steps if you have enough information.", - }) - } - - // 2c. Background plan preamble: append to system prompt (high attention) - - // so the LLM knows from the start that it must mark steps [x]. - - // Skip if a chat session is actively working on the plan directory. - - if opts.Background && agent.ContextBuilder.HasActivePlan() && agent.ContextBuilder.GetPlanStatus() == "executing" { - planDir := agent.ContextBuilder.GetPlanWorkDir() - - skipPreamble := planDir != "" && al.sessions.IsActiveInDir(planDir, "heartbeat") - - if !skipPreamble && len(messages) > 0 && messages[0].Role == "system" { - var sb strings.Builder - - sb.WriteString(messages[0].Content) - - sb.WriteString("\n\n## Background Execution\n") - - sb.WriteString("You are running as a background heartbeat with no conversation history. ") - - sb.WriteString("MEMORY.md is the only shared state between heartbeats. ") - - sb.WriteString( - "After completing each plan step, immediately use edit_file to mark it [x] in memory/MEMORY.md.", - ) - - messages[0].Content = sb.String() - } - } - - // 2d. Snapshot plan status and MEMORY.md size before LLM iteration. - - preStatus := agent.ContextBuilder.GetPlanStatus() - - var preMemoryLen int - - if preStatus == "interviewing" { - preMemoryLen = len(agent.ContextBuilder.ReadMemory()) - } - - // 3. Save user message to session (use compact form if available) - - historyMsg := opts.UserMessage - - if opts.HistoryMessage != "" { - historyMsg = opts.HistoryMessage - } - - agent.Sessions.AddMessage(opts.SessionKey, "user", historyMsg) - - // 4. Record user prompt for stats - - if al.stats != nil { - al.stats.RecordPrompt() - } - - // Capture the finalized system prompt for Mini App inspection - - if len(messages) > 0 { - al.lastSystemPrompt.Store(messages[0].Content) - - al.promptDirty.Store(false) - } - - // 5. Run LLM iteration loop (with automatic phase transitions) - - var iteration int - - const maxPhaseTransitions = 10 - - for phaseLoop := 0; ; phaseLoop++ { - // On phase transition: rebuild system prompt with new phase context + nudge - - if phaseLoop > 0 { - messages = agent.ContextBuilder.BuildMessages( - - agent.Sessions.GetHistory(opts.SessionKey), - - agent.Sessions.GetSummary(opts.SessionKey), - - "", opts.Media, opts.Channel, opts.ChatID, - ) - - messages = append(messages, providers.Message{ - Role: "user", - - Content: fmt.Sprintf( - - "[System] Phase %d is now active. Continue working on the next steps.", - - agent.ContextBuilder.GetCurrentPhase(), - ), - }) - - if len(messages) > 0 { - al.lastSystemPrompt.Store(messages[0].Content) - } - } - - curPlanStatus := preStatus - - if phaseLoop > 0 { - curPlanStatus = agent.ContextBuilder.GetPlanStatus() - } - - var err error - - finalContent, iteration, err = al.runLLMIteration(ctx, agent, messages, opts, task, curPlanStatus) - if err != nil { - return "", err - } - - // 5a. Auto-advance plan phases after LLM iteration - - postStatus := agent.ContextBuilder.GetPlanStatus() - - if !agent.ContextBuilder.HasActivePlan() || - - !(postStatus == "executing" || postStatus == "review" || postStatus == "completed") { - break - } - - // Intercept: if AI changed status to executing or review without user approval - - // (from interviewing or review), validate and hold at "review". - - if preStatus == "interviewing" || (preStatus == "review" && postStatus == "executing") { - if err := agent.ContextBuilder.ValidatePlanStructure(); err != nil { - _ = agent.ContextBuilder.SetPlanStatus("interviewing") - - logger.WarnCF("agent", "Reverted plan to interviewing: "+err.Error(), - - map[string]any{"agent_id": agent.ID}) - - rejectionMsg := "[System] Plan rejected: " + err.Error() + ". Fix and try again." - - agent.Sessions.AddMessage(opts.SessionKey, "user", rejectionMsg) - } else { - _ = agent.ContextBuilder.SetPlanStatus("review") - - al.reporter().ReportStateChange(opts.SessionKey, orch.AgentStatePlanReview, "") - - if !constants.IsInternalChannel(opts.Channel) { - planDisplay := agent.ContextBuilder.FormatPlanDisplay() - - _ = al.bus.PublishOutbound(ctx, bus.OutboundMessage{ - Channel: opts.Channel, - - ChatID: opts.ChatID, - - Content: planDisplay + "\n\nUse /plan start to approve, or continue chatting to refine.", - - SkipPlaceholder: true, - }) - } - } - - break - } - - if postStatus == "executing" && agent.ContextBuilder.GetTotalPhases() == 0 { - _ = agent.ContextBuilder.SetPlanStatus("interviewing") - - logger.WarnCF("agent", "Reverted plan to interviewing: no phases defined", - - map[string]any{"agent_id": agent.ID}) - - break - } - - if agent.ContextBuilder.IsPlanComplete() { - total := agent.ContextBuilder.GetTotalPhases() - - _ = agent.ContextBuilder.SetCurrentPhase(total) - - if preStatus != "completed" { - _ = agent.ContextBuilder.SetPlanStatus("completed") - - al.reporter().ReportStateChange(opts.SessionKey, orch.AgentStatePlanCompleted, "") - - // Deactivate worktree on plan completion - - commitMsg := "plan: " + agent.ContextBuilder.Memory().GetPlanTaskName() - - wtResult, _ := agent.DeactivateWorktree(opts.SessionKey, commitMsg, false) - - if !constants.IsInternalChannel(opts.Channel) { - msg := "\u2705 Plan completed!" - - if wtResult != nil && wtResult.CommitsAhead > 0 { - msg += fmt.Sprintf("\nBranch `%s` retained (%d commits). To merge: `git merge %s`", - - wtResult.Branch, wtResult.CommitsAhead, wtResult.Branch) - } - - _ = al.bus.PublishOutbound(ctx, bus.OutboundMessage{ - Channel: opts.Channel, - - ChatID: opts.ChatID, - - Content: msg, - - SkipPlaceholder: true, - }) - } - } - - break - } - - if agent.ContextBuilder.IsCurrentPhaseComplete() { - if phaseLoop >= maxPhaseTransitions { - logger.WarnCF("agent", "Max phase transitions reached, stopping", - - map[string]any{"agent_id": agent.ID, "transitions": phaseLoop}) - - break - } - - prev := agent.ContextBuilder.GetCurrentPhase() - - _ = agent.ContextBuilder.AdvancePhase() - - next := agent.ContextBuilder.GetCurrentPhase() - - if !constants.IsInternalChannel(opts.Channel) { - _ = al.bus.PublishOutbound(ctx, bus.OutboundMessage{ - Channel: opts.Channel, - - ChatID: opts.ChatID, - - Content: fmt.Sprintf("Phase %d complete. Moving to Phase %d.", prev, next), - - SkipPlaceholder: true, - }) - } - - al.notifyStateChange() - - continue - } - - break - } - - al.notifyStateChange() - - // 5b. Interview staleness detection: compare MEMORY.md size after iteration. - - if agent.ContextBuilder.GetPlanStatus() == "interviewing" { - postMemoryLen := len(agent.ContextBuilder.ReadMemory()) - - if postMemoryLen == preMemoryLen { - agent.interviewStaleCount++ - } else { - agent.interviewStaleCount = 0 - } - - agent.interviewMemoryLen = postMemoryLen - } else { - // Reset counter when not interviewing. - - agent.interviewStaleCount = 0 - } - - // 5c. Handle empty response - - if finalContent == "" { - finalContent = opts.DefaultResponse - } - - // 5d. Store result summary for task completion notification - - if task != nil { - task.Result = utils.Truncate(finalContent, 280) - } - - // 6. Save final assistant message to session (deferred write-behind) - - agent.Sessions.AddMessage(opts.SessionKey, "assistant", finalContent) - - agent.Sessions.MarkDirty(opts.SessionKey) - - // 7. Optional: summarization - - if opts.EnableSummary { - al.maybeSummarize(agent, opts.SessionKey, opts.Channel, opts.ChatID) - } - - // 8. Optional: send response via bus - - if opts.SendResponse { - _ = al.bus.PublishOutbound(ctx, bus.OutboundMessage{ - Channel: opts.Channel, - - ChatID: opts.ChatID, - - Content: finalContent, - - SkipPlaceholder: opts.SystemMessage, // suppress Telegram "Thinking..." for system messages - - }) - } - - // 9. Log response - - responsePreview := utils.Truncate(finalContent, 120) - logger.InfoCF("agent", fmt.Sprintf("Response: %s", responsePreview), - map[string]any{ - "agent_id": agent.ID, - - "session_key": opts.SessionKey, - - "iterations": iteration, - - "final_length": len(finalContent), - }) - - return finalContent, nil -} - func (al *AgentLoop) targetReasoningChannelID(channelName string) (chatID string) { if al.channelManager == nil { return "" @@ -2078,189 +1104,6 @@ func (al *AgentLoop) targetReasoningChannelID(channelName string) (chatID string return "" } -// runLLMIteration executes the LLM call loop with tool handling using hooks. -func (al *AgentLoop) runLLMIteration( - ctx context.Context, - agent *AgentInstance, - messages []providers.Message, - opts processOptions, - task *activeTask, - planSnapshot string, -) (string, int, error) { - hooks := al.buildHooks(agent, opts, task, planSnapshot) - - iteration := 0 - var finalContent string - - for iteration < agent.MaxIterations { - iteration++ - - if msg := hooks.OnIterationStart(iteration); msg != "" { - messages = append(messages, providers.Message{Role: "user", Content: msg}) - } - - logger.DebugCF("agent", "LLM iteration", - map[string]any{ - "agent_id": agent.ID, - "iteration": iteration, - "max": agent.MaxIterations, - }) - - // Build tool definitions - providerToolDefs := hooks.FilterTools(agent.Tools.ToProviderDefs()) - - // Resolve model and candidates for this call - candidates := agent.Candidates - activeModel := agent.Model - if m, c := hooks.SelectModel(); m != "" { - activeModel = m - candidates = c - } - - // Log LLM request details - logger.DebugCF("agent", "LLM request", - map[string]any{ - "agent_id": agent.ID, - "iteration": iteration, - "model": activeModel, - "messages_count": len(messages), - "tools_count": len(providerToolDefs), - "max_tokens": agent.MaxTokens, - "temperature": agent.Temperature, - "system_prompt_len": len(messages[0].Content), - }) - logger.DebugCF("agent", "Full LLM request", - map[string]any{ - "iteration": iteration, - "messages_json": formatMessagesForLog(messages), - "tools_json": formatToolsForLog(providerToolDefs), - }) - - // Streaming setup - onChunk, streamCleanup := hooks.SetupStreaming() - - hooks.OnPreLLMCall() - - // Call LLM with retry - response, err := al.callLLMWithRetry(ctx, agent, &messages, opts, - providerToolDefs, candidates, activeModel, onChunk, iteration) - - // Streaming cleanup - if streamCleanup != nil { - onChunk = nil - streamCleanup() - } - - if err != nil { - logger.ErrorCF("agent", "LLM call failed", - map[string]any{ - "agent_id": agent.ID, - "iteration": iteration, - "model": activeModel, - "error": err.Error(), - }) - return "", iteration, fmt.Errorf("LLM call failed after retries: %w", err) - } - - // Record token usage - if response.Usage != nil && al.stats != nil { - al.stats.RecordUsage( - response.Usage.PromptTokens, - response.Usage.CompletionTokens, - response.Usage.TotalTokens, - ) - } - - go al.handleReasoning(ctx, response.Reasoning, opts.Channel, al.targetReasoningChannelID(opts.Channel)) - - logger.DebugCF("agent", "LLM response", - map[string]any{ - "agent_id": agent.ID, - "iteration": iteration, - "content_chars": len(response.Content), - "tool_calls": len(response.ToolCalls), - "reasoning": response.Reasoning, - "target_channel": al.targetReasoningChannelID(opts.Channel), - "channel": opts.Channel, - }) - - // Clean up response content - response = al.cleanLLMResponse(ctx, response, &messages, agent, iteration, - providerToolDefs, candidates, activeModel, onChunk) - - // No tool calls — check for plan nudge or return - if len(response.ToolCalls) == 0 { - if nudge, cont := hooks.OnNoToolCalls(response.Content, iteration); cont { - messages = append(messages, - providers.Message{Role: "assistant", Content: response.Content}, - providers.Message{Role: "user", Content: nudge}, - ) - continue - } - - finalContent = response.Content - if finalContent == "" && response.ReasoningContent != "" { - finalContent = response.ReasoningContent - } - logger.InfoCF("agent", "LLM response without tool calls (direct answer)", - map[string]any{ - "agent_id": agent.ID, - "iteration": iteration, - "content_chars": len(finalContent), - }) - break - } - - // Normalize and filter tool calls - normalizedToolCalls := make([]providers.ToolCall, 0, len(response.ToolCalls)) - for _, tc := range response.ToolCalls { - normalizedToolCalls = append(normalizedToolCalls, providers.NormalizeToolCall(tc)) - } - - filtered, rejMsg := hooks.FilterToolCalls(normalizedToolCalls) - if len(filtered) < len(normalizedToolCalls) && rejMsg != "" { - messages = append(messages, providers.Message{Role: "user", Content: rejMsg}) - } - normalizedToolCalls = filtered - if len(normalizedToolCalls) == 0 { - continue - } - - // Log tool calls - toolNames := make([]string, 0, len(normalizedToolCalls)) - for _, tc := range normalizedToolCalls { - toolNames = append(toolNames, tc.Name) - } - logger.InfoCF("agent", "LLM requested tool calls", - map[string]any{ - "agent_id": agent.ID, - "tools": toolNames, - "count": len(normalizedToolCalls), - "iteration": iteration, - }) - - hooks.OnToolsProcessed(ctx, iteration, normalizedToolCalls) - - // Build and save assistant message - assistantMsg := buildAssistantMessage(response, normalizedToolCalls) - messages = append(messages, assistantMsg) - agent.Sessions.AddFullMessage(opts.SessionKey, assistantMsg) - - // Execute tool calls and collect results - lastBlocker := al.executeToolCalls(ctx, agent, normalizedToolCalls, &messages, opts, hooks, iteration) - - hooks.InjectReminders(iteration, &messages, lastBlocker) - hooks.RefreshSystemPrompt(messages) - } - - // Force a final text response if max iterations exhausted - if finalContent == "" && iteration >= agent.MaxIterations { - finalContent = al.forceTextResponse(ctx, agent, messages) - } - - return finalContent, iteration, nil -} - // callLLMWithRetry calls the LLM with streaming support, fallback chain, // and retry logic for timeout and context window errors. func (al *AgentLoop) callLLMWithRetry( @@ -2488,90 +1331,6 @@ func buildAssistantMessage(response *providers.LLMResponse, toolCalls []provider return msg } -// executeToolCalls runs each tool call sequentially, publishes results, -// and returns the last blocker (error content) for reminder injection. -func (al *AgentLoop) executeToolCalls( - ctx context.Context, - agent *AgentInstance, - toolCalls []providers.ToolCall, - messages *[]providers.Message, - opts processOptions, - hooks iterationHooks, - iteration int, -) string { - var lastBlocker string - for _, tc := range toolCalls { - argsJSON, _ := json.Marshal(tc.Arguments) - argsPreview := utils.Truncate(string(argsJSON), 200) - logger.InfoCF("agent", fmt.Sprintf("Tool call: %s(%s)", tc.Name, argsPreview), - map[string]any{ - "agent_id": agent.ID, - "tool": tc.Name, - "iteration": iteration, - }) - - // Heartbeat lazy worktree: create worktree on first write-tool call - // Always use ai.Workspace (not GetPlanWorkDir) to avoid creating worktrees - // against stale project paths from previous plans. - if opts.Background && isWriteTool(tc.Name) && !agent.IsInWorktree(opts.SessionKey) { - taskName := "heartbeat-" + time.Now().Format("20060102") - if wt, wtErr := agent.ActivateWorktree(opts.SessionKey, taskName, agent.Workspace); wtErr == nil { - logger.InfoCF("agent", "Heartbeat worktree created", map[string]any{"branch": wt.Branch}) - } - } - - asyncCallback := hooks.OnPreToolExec(ctx, tc) - - toolStart := time.Now() - toolCtx := ctx - if wt := agent.GetWorktree(opts.SessionKey); wt != nil { - toolCtx = tools.WithWorkspaceOverride(toolCtx, wt.Path) - toolCtx = tools.WithWorktreeInfo(toolCtx, wt) - } - - toolResult := agent.Tools.ExecuteWithContext( - toolCtx, tc.Name, tc.Arguments, - opts.Channel, opts.ChatID, asyncCallback, - ) - toolDuration := time.Since(toolStart) - - hooks.OnToolExecDone(tc, toolResult, toolDuration) - - // Publish results to user - if !toolResult.Silent && toolResult.ForUser != "" && opts.SendResponse { - _ = al.bus.PublishOutbound(ctx, bus.OutboundMessage{ - Channel: opts.Channel, - ChatID: opts.ChatID, - Content: toolResult.ForUser, - }) - logger.DebugCF("agent", "Sent tool result to user", - map[string]any{"tool": tc.Name, "content_len": len(toolResult.ForUser)}) - } - - if len(toolResult.Media) > 0 && opts.SendResponse { - al.publishToolMedia(ctx, toolResult, opts) - } - - // Build tool result message - contentForLLM := toolResult.ForLLM - if contentForLLM == "" && toolResult.Err != nil { - contentForLLM = toolResult.Err.Error() - } - if toolResult.IsError || toolResult.Err != nil { - lastBlocker = contentForLLM - } - - toolResultMsg := providers.Message{ - Role: "tool", - Content: contentForLLM, - ToolCallID: tc.ID, - } - *messages = append(*messages, toolResultMsg) - agent.Sessions.AddFullMessage(opts.SessionKey, toolResultMsg) - } - return lastBlocker -} - // publishToolMedia publishes media refs from a tool result as outbound media. func (al *AgentLoop) publishToolMedia(ctx context.Context, result *tools.ToolResult, opts processOptions) { parts := make([]bus.MediaPart, 0, len(result.Media)) @@ -2593,55 +1352,6 @@ func (al *AgentLoop) publishToolMedia(ctx context.Context, result *tools.ToolRes }) } -// forceTextResponse makes a final LLM call without tools when max iterations -// are exhausted, forcing a text response. -func (al *AgentLoop) forceTextResponse(ctx context.Context, agent *AgentInstance, messages []providers.Message) string { - logger.WarnCF("agent", "Max iterations reached, forcing final response without tools", - map[string]any{"agent_id": agent.ID}) - - forceResp, forceErr := agent.Provider.Chat(ctx, messages, nil, agent.Model, map[string]any{ - "max_tokens": agent.MaxTokens, - "temperature": agent.Temperature, - "prompt_cache_key": agent.ID, - }) - if forceErr != nil || forceResp.Content == "" { - return "" - } - content := utils.StripThinkBlocks(forceResp.Content) - if forceResp.Usage != nil && al.stats != nil { - al.stats.RecordUsage( - forceResp.Usage.PromptTokens, - forceResp.Usage.CompletionTokens, - forceResp.Usage.TotalTokens, - ) - } - return content -} - -// updateToolContexts updates the context for tools that need channel/chatID info. - -func (al *AgentLoop) updateToolContexts(agent *AgentInstance, channel, chatID string) { - // Use ContextualTool interface instead of type assertions - - if tool, ok := agent.Tools.Get("message"); ok { - if mt, ok := tool.(tools.ContextualTool); ok { - mt.SetContext(channel, chatID) - } - } - - if tool, ok := agent.Tools.Get("spawn"); ok { - if st, ok := tool.(tools.ContextualTool); ok { - st.SetContext(channel, chatID) - } - } - - if tool, ok := agent.Tools.Get("subagent"); ok { - if st, ok := tool.(tools.ContextualTool); ok { - st.SetContext(channel, chatID) - } - } -} - // Helper to extract provider from registry for cleanup func extractProvider(registry *AgentRegistry) (providers.LLMProvider, bool) { if registry == nil { diff --git a/pkg/agent/loop_ext.go b/pkg/agent/loop_ext.go index 67d99a89b..4ef4aaecf 100644 --- a/pkg/agent/loop_ext.go +++ b/pkg/agent/loop_ext.go @@ -1,11 +1,18 @@ package agent import ( + "strings" "sync" + "github.com/sipeed/picoclaw/pkg/bus" "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/logger" "github.com/sipeed/picoclaw/pkg/orch" + "github.com/sipeed/picoclaw/pkg/providers" + "github.com/sipeed/picoclaw/pkg/routing" "github.com/sipeed/picoclaw/pkg/stats" + "github.com/sipeed/picoclaw/pkg/tools" + "github.com/sipeed/picoclaw/pkg/utils" ) // loopExt holds fork-specific fields for AgentLoop. @@ -35,6 +42,61 @@ type loopExt struct { onHeartbeatThreadUpdate func(int) } +// initLoopExt initializes all fork-specific fields: stats tracker, +// session tracker, orchestration broadcaster, and background goroutines. +// Called from NewAgentLoop after the struct is constructed. +func (al *AgentLoop) initLoopExt(cfg *config.Config, registry *AgentRegistry, enableStats bool) { + defaultAgent := registry.GetDefaultAgent() + + // Stats tracker + if enableStats && defaultAgent != nil { + al.stats = stats.NewTracker(defaultAgent.Workspace) + } + + // Session tracker + al.sessions = NewSessionTracker() + + // Orchestration broadcaster — needed if any agent has subagents enabled. + // Note: instance.go maps defaults.Orchestration → Subagents.Enabled, + // so --orchestration is automatically reflected here. + al.orchReporter = orch.Noop + for _, id := range registry.ListAgentIDs() { + if a, ok := registry.GetAgent(id); ok && a.Subagents != nil && a.Subagents.Enabled { + al.orchBroadcaster = orch.NewBroadcaster() + al.orchReporter = al.orchBroadcaster + break + } + } + + // Shutdown signal channel + al.done = make(chan struct{}) + + // Background GC goroutine + go al.gcLoop() +} + +// closeExt releases fork-specific resources: done channel, stats tracker, +// and session stores for all agents. +func (al *AgentLoop) closeExt() { + select { + case <-al.done: + // already closed + default: + close(al.done) + } + + if al.stats != nil { + al.stats.Close() + } + + registry := al.GetRegistry() + for _, agentID := range registry.ListAgentIDs() { + if agent, ok := registry.GetAgent(agentID); ok { + agent.Sessions.Close() + } + } +} + // SetConfigSaver registers a callback to persist config changes. func (al *AgentLoop) SetConfigSaver(fn func(*config.Config) error) { al.saveConfig = fn @@ -44,3 +106,139 @@ func (al *AgentLoop) SetConfigSaver(fn func(*config.Config) error) { func (al *AgentLoop) SetHeartbeatThreadUpdater(fn func(int)) { al.onHeartbeatThreadUpdate = fn } + +// registerOrchestrationTools registers spawn, subagent, answer, and review_plan +// tools for agents with orchestration enabled. +func registerOrchestrationTools( + cfg *config.Config, + agent *AgentInstance, + agentID string, + registry *AgentRegistry, + provider providers.LLMProvider, + msgBus *bus.MessageBus, + al *AgentLoop, +) { + if agent.Subagents == nil || !agent.Subagents.Enabled { + return + } + + webSearchOpts := tools.WebSearchToolOptions{ + BraveAPIKeys: config.MergeAPIKeys(cfg.Tools.Web.Brave.APIKey, cfg.Tools.Web.Brave.APIKeys), + BraveMaxResults: cfg.Tools.Web.Brave.MaxResults, + BraveEnabled: cfg.Tools.Web.Brave.Enabled, + TavilyAPIKeys: config.MergeAPIKeys(cfg.Tools.Web.Tavily.APIKey, cfg.Tools.Web.Tavily.APIKeys), + TavilyBaseURL: cfg.Tools.Web.Tavily.BaseURL, + TavilyMaxResults: cfg.Tools.Web.Tavily.MaxResults, + TavilyEnabled: cfg.Tools.Web.Tavily.Enabled, + DuckDuckGoMaxResults: cfg.Tools.Web.DuckDuckGo.MaxResults, + DuckDuckGoEnabled: cfg.Tools.Web.DuckDuckGo.Enabled, + PerplexityAPIKeys: config.MergeAPIKeys( + cfg.Tools.Web.Perplexity.APIKey, + cfg.Tools.Web.Perplexity.APIKeys, + ), + PerplexityMaxResults: cfg.Tools.Web.Perplexity.MaxResults, + PerplexityEnabled: cfg.Tools.Web.Perplexity.Enabled, + } + + subagentManager := tools.NewSubagentManager( + provider, + agent.Model, + agent.Workspace, + msgBus, + al.reporter(), + webSearchOpts, + ) + + subagentManager.SetLLMOptions(agent.MaxTokens, agent.Temperature) + + // Wire session recorder for DAG persistence. + recorder := newSessionRecorder(agent.Sessions) + conductorKey := routing.BuildAgentMainSessionKey(agent.ID) + subagentManager.SetSessionRecorder(recorder, conductorKey) + + agent.SubagentMgr = subagentManager + + spawnTool := tools.NewSpawnTool(subagentManager) + currentAgentID := agentID + spawnTool.SetAllowlistChecker(func(targetAgentID string) bool { + return registry.CanSpawnSubagent(currentAgentID, targetAgentID) + }) + agent.Tools.Register(spawnTool) + + // Register blocking subagent tool alongside spawn + agent.Tools.Register(tools.NewSubagentTool(subagentManager)) + + // Register conductor-side escalation tools (answer questions, review plans) + agent.Tools.Register(tools.NewAnswerSubagentTool(subagentManager)) + agent.Tools.Register(tools.NewReviewSubagentPlanTool(subagentManager)) + + // Set orchestration mode on context builder + agent.ContextBuilder.SetOrchestrationEnabled(true) +} + +// handleTaskIntervention checks if a message is a reply to an active task and +// either cancels the task or injects a user intervention. Returns (response, handled). +func (al *AgentLoop) handleTaskIntervention(msg bus.InboundMessage) (string, bool) { + taskID, ok := msg.Metadata["task_id"] + if !ok || taskID == "" { + return "", false + } + + val, found := al.activeTasks.Load(taskID) + if !found { + // Task not found — fall through to normal processing + return "", false + } + + task := val.(*activeTask) + + content := strings.TrimSpace(msg.Content) + lower := strings.ToLower(content) + + // Check for stop keywords + stopKeywords := []string{ + "stop", "cancel", "abort", + "停止", "中止", "やめて", //nolint:gosmopolitan // intentional CJK stop words + } + + for _, kw := range stopKeywords { + if lower == kw { + task.cancel() + + logger.InfoCF("agent", "Task canceled by user intervention", + map[string]any{"task_id": taskID}) + + return "Task canceled.", true + } + } + + // Inject message into interrupt channel for the tool loop + select { + case task.interrupt <- content: + logger.InfoCF("agent", "User intervention queued", + map[string]any{"task_id": taskID, "content": utils.Truncate(content, 80)}) + default: + logger.WarnCF("agent", "Interrupt channel full, message dropped", + map[string]any{"task_id": taskID}) + } + + return "Intervention sent to running task.", true +} + +// expandForkCommands expands fork-specific /skill and /plan commands in the message. +// Returns the modified message and the compact form for history. +func (al *AgentLoop) expandForkCommands(msg *bus.InboundMessage) string { + var expansionCompact string + + if expanded, compact, ok := al.expandSkillCommand(*msg); ok { + msg.Content = expanded + expansionCompact = compact + } + + if expanded, compact, ok := al.expandPlanCommand(*msg); ok { + msg.Content = expanded + expansionCompact = compact + } + + return expansionCompact +} diff --git a/pkg/agent/loop_run.go b/pkg/agent/loop_run.go new file mode 100644 index 000000000..8e0a82620 --- /dev/null +++ b/pkg/agent/loop_run.go @@ -0,0 +1,1073 @@ +package agent + +import ( + "context" + "encoding/json" + "fmt" + "path/filepath" + "strings" + "time" + + "github.com/sipeed/picoclaw/pkg/bus" + "github.com/sipeed/picoclaw/pkg/constants" + "github.com/sipeed/picoclaw/pkg/git" + "github.com/sipeed/picoclaw/pkg/logger" + "github.com/sipeed/picoclaw/pkg/orch" + "github.com/sipeed/picoclaw/pkg/providers" + "github.com/sipeed/picoclaw/pkg/session" + "github.com/sipeed/picoclaw/pkg/tools" + "github.com/sipeed/picoclaw/pkg/utils" +) + +// runAgentLoop is the main message processing loop for a single agent session. +func (al *AgentLoop) runAgentLoop(ctx context.Context, agent *AgentInstance, opts processOptions) (string, error) { + // -1. Acquire per-session lock to prevent concurrent access on the same session + + if !al.acquireSessionLock(ctx, opts.SessionKey) { + return "", fmt.Errorf("context canceled while waiting for session lock") + } + + defer al.releaseSessionLock(opts.SessionKey) + + // Report session lifecycle to canvas. + + al.reporter().ReportSpawn(opts.SessionKey, opts.Channel, opts.UserMessage) + + defer al.reporter().ReportGC(opts.SessionKey, "completed") + + // -0. Create cancelable child context and register active task + + taskCtx, taskCancel := context.WithCancel(ctx) + + defer taskCancel() + + task := &activeTask{ + Description: utils.Truncate(opts.UserMessage, 80), + + MaxIter: agent.MaxIterations, + + StartedAt: time.Now(), + + cancel: taskCancel, + + interrupt: make(chan string, 1), + } + + // Guarantee heartbeat worktree cleanup on ALL exit paths (error, panic, normal). + + // Wait for spawned subagents first so they aren't killed mid-flight. + + // After auto-commit, attempt to merge the worktree branch into main. + + defer al.cleanupHeartbeatWorktree(agent, opts) + + // For background tasks (cron/heartbeat), generate a TaskID and send notification + + isBackgroundTask := opts.Background && al.state != nil + + if isBackgroundTask && opts.TaskID == "" { + opts.TaskID = fmt.Sprintf("task-%s-%d", opts.SessionKey, time.Now().UnixMilli()) + + // Determine notification channel: use opts.Channel if already a real channel, + + // otherwise resolve from last active channel + + notifyChannel := opts.Channel + + notifyChatID := opts.ChatID + + if constants.IsInternalChannel(notifyChannel) || notifyChannel == "" { + if lastChannel := al.state.GetLastChannel(); lastChannel != "" { + if idx := strings.Index(lastChannel, ":"); idx > 0 { + notifyChannel = lastChannel[:idx] + + notifyChatID = lastChannel[idx+1:] + } + } + } + + if notifyChannel != "" && notifyChatID != "" && !constants.IsInternalChannel(notifyChannel) { + // Override opts channel/chatID for status updates + + opts.Channel = notifyChannel + + opts.ChatID = notifyChatID + + // Send initial task notification + + _ = al.bus.PublishOutbound(ctx, bus.OutboundMessage{ + Channel: notifyChannel, + + ChatID: notifyChatID, + + Content: fmt.Sprintf("\U0001F916 Background task started\n%s", task.Description), + + IsTaskStatus: true, + + TaskID: opts.TaskID, + }) + } + } + + // Shared variable for capturing LLM's final response. The defer below reads it + + // to include the response in the task completion message. + + var finalContent string + + // Use TaskID as key if available (for background tasks), else sessionKey + + taskKey := opts.SessionKey + + if opts.TaskID != "" { + taskKey = opts.TaskID + } + + al.activeTasks.Store(taskKey, task) + + defer al.publishTaskCompletion(task, &finalContent, opts, taskKey) + + // Replace ctx with the cancelable child context + + ctx = taskCtx + + // 0. Record last channel for heartbeat notifications (skip internal channels) + + if opts.Channel != "" && opts.ChatID != "" { + // Don't record internal channels (cli, system, subagent) + + if !constants.IsInternalChannel(opts.Channel) { + channelKey := fmt.Sprintf("%s:%s", opts.Channel, opts.ChatID) + if err := al.RecordLastChannel(channelKey); err != nil { + logger.WarnCF("agent", "Failed to record last channel", map[string]any{"error": err.Error()}) + } + + if err := al.RecordLastHeartbeatTarget(channelKey); err != nil { + logger.WarnCF("agent", "Failed to record last heartbeat target", map[string]any{"error": err.Error()}) + } + } + } + + // 1. Update tool contexts + + al.updateToolContexts(agent, opts.Channel, opts.ChatID) + + // 1-bis. For background tasks that don't send a final response (e.g. heartbeat), + + // redirect the message tool to publish as IsTaskStatus so its output lands in + + // the same bubble as the task status instead of creating a separate message. + + if opts.Background && !opts.SendResponse && opts.TaskID != "" { + al.redirectMessageToolForTask(agent, task, opts) + } + + // 1a. Set session-specific working directory for bootstrap file lookup. + + // Prefer the tool-detected project directory (touch_dir) from the session tracker, + + // resolved as an absolute path under workspace. Fall back to worktree or workspace. + + if active := al.sessions.ListActive(); len(active) > 0 && active[0].SessionKey == opts.SessionKey && + + active[0].TouchDir != "" { + agent.ContextBuilder.SetWorkDir(filepath.Join(agent.Workspace, active[0].TouchDir)) + } else { + agent.ContextBuilder.SetWorkDir(agent.EffectiveWorkspace(opts.SessionKey)) + } + + // 1b. Inject peer session awareness into system prompt + + projectPath := agent.ContextBuilder.GetPlanWorkDir() + + if projectPath == "" { + projectPath = agent.Workspace + } + + peers := al.sessions.GetPeerPurposes(opts.SessionKey, projectPath) + + if len(peers) > 0 { + var peerNote strings.Builder + + peerNote.WriteString("Other sessions working on this project:\n") + + for _, p := range peers { + peerNote.WriteString(fmt.Sprintf("- %s: %s (branch: %s)\n", p.SessionKey, p.Purpose, p.Branch)) + } + + peerNote.WriteString("\nAvoid conflicting changes with these sessions.") + + agent.ContextBuilder.SetPeerNote(peerNote.String()) + } else { + agent.ContextBuilder.SetPeerNote("") + } + + // 2. Build messages (skip history for heartbeat) + + var history []providers.Message + var summary string + if !opts.NoHistory { + history = agent.Sessions.GetHistory(opts.SessionKey) + summary = agent.Sessions.GetSummary(opts.SessionKey) + + // Sanitize history to remove orphaned tool calls (from crashes/session collisions) + + var removedCount int + + history, removedCount = session.SanitizeHistory(history) + + if removedCount > 0 { + logger.WarnCF("agent", "Sanitized session history: removed orphaned messages", + + map[string]any{ + "session_key": opts.SessionKey, + + "removed_count": removedCount, + }) + + // Persist the sanitized history + + agent.Sessions.SetHistory(opts.SessionKey, history) + + _ = agent.Sessions.Save(opts.SessionKey) + } + } + messages := agent.ContextBuilder.BuildMessages( + history, + summary, + opts.UserMessage, + + opts.Media, + + opts.Channel, + opts.ChatID, + ) + + // Resolve media:// refs: images→base64 data URLs, non-images→local paths in content + cfg := al.GetConfig() + maxMediaSize := cfg.Agents.Defaults.GetMaxMediaSize() + messages = resolveMediaRefs(messages, al.mediaStore, maxMediaSize) + + // 2b. Interview staleness nudge: if MEMORY.md hasn't been updated for + + // several consecutive turns, inject a reminder so the AI writes its findings. + + const interviewStaleThreshold = 2 + + if agent.ContextBuilder.GetPlanStatus() == "interviewing" && agent.interviewStaleCount >= interviewStaleThreshold { + messages = append(messages, providers.Message{ + Role: "user", + + Content: "[System] You have been interviewing for several turns without updating memory/MEMORY.md. Please use edit_file now to save your findings to the ## Context section, or organize the plan into ## Phase sections with `- [ ]` checkbox steps if you have enough information.", + }) + } + + // 2c. Background plan preamble: append to system prompt (high attention) + + // so the LLM knows from the start that it must mark steps [x]. + + // Skip if a chat session is actively working on the plan directory. + + if opts.Background && agent.ContextBuilder.HasActivePlan() && agent.ContextBuilder.GetPlanStatus() == "executing" { + planDir := agent.ContextBuilder.GetPlanWorkDir() + + skipPreamble := planDir != "" && al.sessions.IsActiveInDir(planDir, "heartbeat") + + if !skipPreamble && len(messages) > 0 && messages[0].Role == "system" { + var sb strings.Builder + + sb.WriteString(messages[0].Content) + + sb.WriteString("\n\n## Background Execution\n") + + sb.WriteString("You are running as a background heartbeat with no conversation history. ") + + sb.WriteString("MEMORY.md is the only shared state between heartbeats. ") + + sb.WriteString( + "After completing each plan step, immediately use edit_file to mark it [x] in memory/MEMORY.md.", + ) + + messages[0].Content = sb.String() + } + } + + // 2d. Snapshot plan status and MEMORY.md size before LLM iteration. + + preStatus := agent.ContextBuilder.GetPlanStatus() + + var preMemoryLen int + + if preStatus == "interviewing" { + preMemoryLen = len(agent.ContextBuilder.ReadMemory()) + } + + // 3. Save user message to session (use compact form if available) + + historyMsg := opts.UserMessage + + if opts.HistoryMessage != "" { + historyMsg = opts.HistoryMessage + } + + agent.Sessions.AddMessage(opts.SessionKey, "user", historyMsg) + + // 4. Record user prompt for stats + + if al.stats != nil { + al.stats.RecordPrompt() + } + + // Capture the finalized system prompt for Mini App inspection + + if len(messages) > 0 { + al.lastSystemPrompt.Store(messages[0].Content) + + al.promptDirty.Store(false) + } + + // 5. Run LLM iteration loop (with automatic phase transitions) + + var iteration int + + const maxPhaseTransitions = 10 + + for phaseLoop := 0; ; phaseLoop++ { + // On phase transition: rebuild system prompt with new phase context + nudge + + if phaseLoop > 0 { + messages = agent.ContextBuilder.BuildMessages( + + agent.Sessions.GetHistory(opts.SessionKey), + + agent.Sessions.GetSummary(opts.SessionKey), + + "", opts.Media, opts.Channel, opts.ChatID, + ) + + messages = append(messages, providers.Message{ + Role: "user", + + Content: fmt.Sprintf( + + "[System] Phase %d is now active. Continue working on the next steps.", + + agent.ContextBuilder.GetCurrentPhase(), + ), + }) + + if len(messages) > 0 { + al.lastSystemPrompt.Store(messages[0].Content) + } + } + + curPlanStatus := preStatus + + if phaseLoop > 0 { + curPlanStatus = agent.ContextBuilder.GetPlanStatus() + } + + var err error + + finalContent, iteration, err = al.runLLMIteration(ctx, agent, messages, opts, task, curPlanStatus) + if err != nil { + return "", err + } + + // 5a. Auto-advance plan phases after LLM iteration + + postStatus := agent.ContextBuilder.GetPlanStatus() + + if !agent.ContextBuilder.HasActivePlan() || + + !(postStatus == "executing" || postStatus == "review" || postStatus == "completed") { + break + } + + // Intercept: if AI changed status to executing or review without user approval + + // (from interviewing or review), validate and hold at "review". + + if preStatus == "interviewing" || (preStatus == "review" && postStatus == "executing") { + if err := agent.ContextBuilder.ValidatePlanStructure(); err != nil { + _ = agent.ContextBuilder.SetPlanStatus("interviewing") + + logger.WarnCF("agent", "Reverted plan to interviewing: "+err.Error(), + + map[string]any{"agent_id": agent.ID}) + + rejectionMsg := "[System] Plan rejected: " + err.Error() + ". Fix and try again." + + agent.Sessions.AddMessage(opts.SessionKey, "user", rejectionMsg) + } else { + _ = agent.ContextBuilder.SetPlanStatus("review") + + al.reporter().ReportStateChange(opts.SessionKey, orch.AgentStatePlanReview, "") + + if !constants.IsInternalChannel(opts.Channel) { + planDisplay := agent.ContextBuilder.FormatPlanDisplay() + + _ = al.bus.PublishOutbound(ctx, bus.OutboundMessage{ + Channel: opts.Channel, + + ChatID: opts.ChatID, + + Content: planDisplay + "\n\nUse /plan start to approve, or continue chatting to refine.", + + SkipPlaceholder: true, + }) + } + } + + break + } + + if postStatus == "executing" && agent.ContextBuilder.GetTotalPhases() == 0 { + _ = agent.ContextBuilder.SetPlanStatus("interviewing") + + logger.WarnCF("agent", "Reverted plan to interviewing: no phases defined", + + map[string]any{"agent_id": agent.ID}) + + break + } + + if agent.ContextBuilder.IsPlanComplete() { + total := agent.ContextBuilder.GetTotalPhases() + + _ = agent.ContextBuilder.SetCurrentPhase(total) + + if preStatus != "completed" { + _ = agent.ContextBuilder.SetPlanStatus("completed") + + al.reporter().ReportStateChange(opts.SessionKey, orch.AgentStatePlanCompleted, "") + + // Deactivate worktree on plan completion + + commitMsg := "plan: " + agent.ContextBuilder.Memory().GetPlanTaskName() + + wtResult, _ := agent.DeactivateWorktree(opts.SessionKey, commitMsg, false) + + if !constants.IsInternalChannel(opts.Channel) { + msg := "\u2705 Plan completed!" + + if wtResult != nil && wtResult.CommitsAhead > 0 { + msg += fmt.Sprintf("\nBranch `%s` retained (%d commits). To merge: `git merge %s`", + + wtResult.Branch, wtResult.CommitsAhead, wtResult.Branch) + } + + _ = al.bus.PublishOutbound(ctx, bus.OutboundMessage{ + Channel: opts.Channel, + + ChatID: opts.ChatID, + + Content: msg, + + SkipPlaceholder: true, + }) + } + } + + break + } + + if agent.ContextBuilder.IsCurrentPhaseComplete() { + if phaseLoop >= maxPhaseTransitions { + logger.WarnCF("agent", "Max phase transitions reached, stopping", + + map[string]any{"agent_id": agent.ID, "transitions": phaseLoop}) + + break + } + + prev := agent.ContextBuilder.GetCurrentPhase() + + _ = agent.ContextBuilder.AdvancePhase() + + next := agent.ContextBuilder.GetCurrentPhase() + + if !constants.IsInternalChannel(opts.Channel) { + _ = al.bus.PublishOutbound(ctx, bus.OutboundMessage{ + Channel: opts.Channel, + + ChatID: opts.ChatID, + + Content: fmt.Sprintf("Phase %d complete. Moving to Phase %d.", prev, next), + + SkipPlaceholder: true, + }) + } + + al.notifyStateChange() + + continue + } + + break + } + + al.notifyStateChange() + + // 5b. Interview staleness detection: compare MEMORY.md size after iteration. + + if agent.ContextBuilder.GetPlanStatus() == "interviewing" { + postMemoryLen := len(agent.ContextBuilder.ReadMemory()) + + if postMemoryLen == preMemoryLen { + agent.interviewStaleCount++ + } else { + agent.interviewStaleCount = 0 + } + + agent.interviewMemoryLen = postMemoryLen + } else { + // Reset counter when not interviewing. + + agent.interviewStaleCount = 0 + } + + // 5c. Handle empty response + + if finalContent == "" { + finalContent = opts.DefaultResponse + } + + // 5d. Store result summary for task completion notification + task.Result = utils.Truncate(finalContent, 280) + + // 6. Save final assistant message to session (deferred write-behind) + + agent.Sessions.AddMessage(opts.SessionKey, "assistant", finalContent) + + agent.Sessions.MarkDirty(opts.SessionKey) + + // 7. Optional: summarization + + if opts.EnableSummary { + al.maybeSummarize(agent, opts.SessionKey, opts.Channel, opts.ChatID) + } + + // 8. Optional: send response via bus + + if opts.SendResponse { + _ = al.bus.PublishOutbound(ctx, bus.OutboundMessage{ + Channel: opts.Channel, + + ChatID: opts.ChatID, + + Content: finalContent, + + SkipPlaceholder: opts.SystemMessage, // suppress Telegram "Thinking..." for system messages + + }) + } + + // 9. Log response + + responsePreview := utils.Truncate(finalContent, 120) + logger.InfoCF("agent", fmt.Sprintf("Response: %s", responsePreview), + map[string]any{ + "agent_id": agent.ID, + + "session_key": opts.SessionKey, + + "iterations": iteration, + + "final_length": len(finalContent), + }) + + return finalContent, nil +} + +// cleanupHeartbeatWorktree handles worktree cleanup for background tasks. +// Waits for spawned subagents, auto-commits, and attempts fast-forward merge. +func (al *AgentLoop) cleanupHeartbeatWorktree(agent *AgentInstance, opts processOptions) { + if !opts.Background { + return + } + + if agent.SubagentMgr != nil { + agent.SubagentMgr.WaitAll(35 * time.Minute) // slightly above spawnTimeout + } + + wt := agent.GetWorktree(opts.SessionKey) + + if wt == nil { + return + } + + // 1. Auto-commit uncommitted changes in worktree + if git.HasUncommittedChanges(wt.Path) { + _ = git.AutoCommit(wt.Path, "heartbeat: auto-save") + } + + // 2. Check if there are unique commits worth merging + repoRoot := git.FindRepoRoot(agent.Workspace) + ahead := git.CommitsAhead(repoRoot, wt.BaseBranch, wt.Branch) + + if ahead > 0 && repoRoot != "" { + // 3. Try fast-forward merge into base branch + mr := git.MergeWorktreeBranch(repoRoot, wt) + + // 4. Notify based on merge result + if !constants.IsInternalChannel(opts.Channel) { + cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 5*time.Second) + + if mr.Merged { + _ = al.bus.PublishOutbound(cleanupCtx, bus.OutboundMessage{ + Channel: opts.Channel, + ChatID: opts.ChatID, + Content: fmt.Sprintf("Heartbeat: merged %d commit(s) to %s.", + ahead, wt.BaseBranch), + }) + } else if mr.Conflict { + _ = al.bus.PublishOutbound(cleanupCtx, bus.OutboundMessage{ + Channel: opts.Channel, + ChatID: opts.ChatID, + Content: fmt.Sprintf("Heartbeat: merge conflict on branch `%s` — manual merge needed.", + mr.Branch), + }) + } + + cleanupCancel() + } + } + + // 5. Dispose worktree (branch auto-deleted if merged, kept if conflict) + agent.DeactivateWorktree(opts.SessionKey, "", false) +} + +// publishTaskCompletion publishes the final task status on completion for +// background tasks, including the LLM response in the completion bubble. +func (al *AgentLoop) publishTaskCompletion( + task *activeTask, finalContent *string, opts processOptions, taskKey string, +) { + al.activeTasks.Delete(taskKey) + + if opts.TaskID == "" { + return + } + + elapsed := time.Since(task.StartedAt) + + completionMsg := fmt.Sprintf("\u2705 Task completed (%.1fs)", elapsed.Seconds()) + + // Determine the best content to show in the completion bubble. + // Priority: message tool content > finalContent > task.Result + task.mu.Lock() + msgContent := task.messageContent + task.mu.Unlock() + + var resultContent string + + switch { + case msgContent != "": + // The message tool already sent this to the user via the + // task bubble; re-include it so the completion doesn't erase it. + resultContent = msgContent + + case *finalContent != "" && *finalContent != defaultResponse && *finalContent != "HEARTBEAT_OK": + resultContent = *finalContent + + default: + summary := task.Result + if summary == "" { + summary = task.Description + } + resultContent = summary + } + + if resultContent != "" { + combined := completionMsg + "\n\n" + resultContent + + if len([]rune(combined)) <= 4096 { + completionMsg = combined + } else { + // Too long for one bubble: send header as task status, + // body as regular message (auto-split by SplitMessage). + doneCtx, doneCancel := context.WithTimeout(context.Background(), 5*time.Second) + + _ = al.bus.PublishOutbound(doneCtx, bus.OutboundMessage{ + Channel: opts.Channel, + ChatID: opts.ChatID, + Content: completionMsg, + IsTaskStatus: true, + TaskID: opts.TaskID, + Final: true, + }) + + _ = al.bus.PublishOutbound(doneCtx, bus.OutboundMessage{ + Channel: opts.Channel, + ChatID: opts.ChatID, + Content: resultContent, + }) + + doneCancel() + return + } + } + + doneCtx, doneCancel := context.WithTimeout(context.Background(), 5*time.Second) + + _ = al.bus.PublishOutbound(doneCtx, bus.OutboundMessage{ + Channel: opts.Channel, + ChatID: opts.ChatID, + Content: completionMsg, + IsTaskStatus: true, + TaskID: opts.TaskID, + Final: true, + }) + + doneCancel() +} + +// redirectMessageToolForTask redirects the message tool to publish as IsTaskStatus +// for background tasks that don't send a final response. +func (al *AgentLoop) redirectMessageToolForTask(agent *AgentInstance, task *activeTask, opts processOptions) { + tool, ok := agent.Tools.Get("message") + if !ok { + return + } + mt, ok := tool.(*tools.MessageTool) + if !ok { + return + } + taskID := opts.TaskID + + mt.SetSendCallback(func(channel, chatID, content string) error { + // Capture the message tool's content so the completion + // defer can include it instead of losing it to an overwrite. + if task != nil { + task.mu.Lock() + task.messageContent = content + task.mu.Unlock() + } + + pubCtx, pubCancel := context.WithTimeout(context.Background(), 5*time.Second) + defer pubCancel() + + return al.bus.PublishOutbound(pubCtx, bus.OutboundMessage{ + Channel: channel, + ChatID: chatID, + Content: content, + IsTaskStatus: true, + TaskID: taskID, + }) + }) +} + +// runLLMIteration executes the LLM call loop with tool handling using hooks. +func (al *AgentLoop) runLLMIteration( + ctx context.Context, + agent *AgentInstance, + messages []providers.Message, + opts processOptions, + task *activeTask, + planSnapshot string, +) (string, int, error) { + hooks := al.buildHooks(agent, opts, task, planSnapshot) + + iteration := 0 + var finalContent string + + for iteration < agent.MaxIterations { + iteration++ + + if msg := hooks.OnIterationStart(iteration); msg != "" { + messages = append(messages, providers.Message{Role: "user", Content: msg}) + } + + logger.DebugCF("agent", "LLM iteration", + map[string]any{ + "agent_id": agent.ID, + "iteration": iteration, + "max": agent.MaxIterations, + }) + + // Build tool definitions + providerToolDefs := hooks.FilterTools(agent.Tools.ToProviderDefs()) + + // Resolve model and candidates for this call + candidates := agent.Candidates + activeModel := agent.Model + if m, c := hooks.SelectModel(); m != "" { + activeModel = m + candidates = c + } + + // Log LLM request details + logger.DebugCF("agent", "LLM request", + map[string]any{ + "agent_id": agent.ID, + "iteration": iteration, + "model": activeModel, + "messages_count": len(messages), + "tools_count": len(providerToolDefs), + "max_tokens": agent.MaxTokens, + "temperature": agent.Temperature, + "system_prompt_len": len(messages[0].Content), + }) + logger.DebugCF("agent", "Full LLM request", + map[string]any{ + "iteration": iteration, + "messages_json": formatMessagesForLog(messages), + "tools_json": formatToolsForLog(providerToolDefs), + }) + + // Streaming setup + onChunk, streamCleanup := hooks.SetupStreaming() + + hooks.OnPreLLMCall() + + // Call LLM with retry + response, err := al.callLLMWithRetry(ctx, agent, &messages, opts, + providerToolDefs, candidates, activeModel, onChunk, iteration) + + // Streaming cleanup + if streamCleanup != nil { + onChunk = nil + streamCleanup() + } + + if err != nil { + logger.ErrorCF("agent", "LLM call failed", + map[string]any{ + "agent_id": agent.ID, + "iteration": iteration, + "model": activeModel, + "error": err.Error(), + }) + return "", iteration, fmt.Errorf("LLM call failed after retries: %w", err) + } + + // Record token usage + if response.Usage != nil && al.stats != nil { + al.stats.RecordUsage( + response.Usage.PromptTokens, + response.Usage.CompletionTokens, + response.Usage.TotalTokens, + ) + } + + go al.handleReasoning(ctx, response.Reasoning, opts.Channel, al.targetReasoningChannelID(opts.Channel)) + + logger.DebugCF("agent", "LLM response", + map[string]any{ + "agent_id": agent.ID, + "iteration": iteration, + "content_chars": len(response.Content), + "tool_calls": len(response.ToolCalls), + "reasoning": response.Reasoning, + "target_channel": al.targetReasoningChannelID(opts.Channel), + "channel": opts.Channel, + }) + + // Clean up response content + response = al.cleanLLMResponse(ctx, response, &messages, agent, iteration, + providerToolDefs, candidates, activeModel, onChunk) + + // No tool calls — check for plan nudge or return + if len(response.ToolCalls) == 0 { + if nudge, cont := hooks.OnNoToolCalls(response.Content, iteration); cont { + messages = append(messages, + providers.Message{Role: "assistant", Content: response.Content}, + providers.Message{Role: "user", Content: nudge}, + ) + continue + } + + finalContent = response.Content + if finalContent == "" && response.ReasoningContent != "" { + finalContent = response.ReasoningContent + } + logger.InfoCF("agent", "LLM response without tool calls (direct answer)", + map[string]any{ + "agent_id": agent.ID, + "iteration": iteration, + "content_chars": len(finalContent), + }) + break + } + + // Normalize and filter tool calls + normalizedToolCalls := make([]providers.ToolCall, 0, len(response.ToolCalls)) + for _, tc := range response.ToolCalls { + normalizedToolCalls = append(normalizedToolCalls, providers.NormalizeToolCall(tc)) + } + + filtered, rejMsg := hooks.FilterToolCalls(normalizedToolCalls) + if len(filtered) < len(normalizedToolCalls) && rejMsg != "" { + messages = append(messages, providers.Message{Role: "user", Content: rejMsg}) + } + normalizedToolCalls = filtered + if len(normalizedToolCalls) == 0 { + continue + } + + // Log tool calls + toolNames := make([]string, 0, len(normalizedToolCalls)) + for _, tc := range normalizedToolCalls { + toolNames = append(toolNames, tc.Name) + } + logger.InfoCF("agent", "LLM requested tool calls", + map[string]any{ + "agent_id": agent.ID, + "tools": toolNames, + "count": len(normalizedToolCalls), + "iteration": iteration, + }) + + hooks.OnToolsProcessed(ctx, iteration, normalizedToolCalls) + + // Build and save assistant message + assistantMsg := buildAssistantMessage(response, normalizedToolCalls) + messages = append(messages, assistantMsg) + agent.Sessions.AddFullMessage(opts.SessionKey, assistantMsg) + + // Execute tool calls and collect results + lastBlocker := al.executeToolCalls(ctx, agent, normalizedToolCalls, &messages, opts, hooks, iteration) + + hooks.InjectReminders(iteration, &messages, lastBlocker) + hooks.RefreshSystemPrompt(messages) + } + + // Force a final text response if max iterations exhausted + if finalContent == "" && iteration >= agent.MaxIterations { + finalContent = al.forceTextResponse(ctx, agent, messages) + } + + return finalContent, iteration, nil +} + +// executeToolCalls runs each tool call sequentially, publishes results, +// and returns the last blocker (error content) for reminder injection. +func (al *AgentLoop) executeToolCalls( + ctx context.Context, + agent *AgentInstance, + toolCalls []providers.ToolCall, + messages *[]providers.Message, + opts processOptions, + hooks iterationHooks, + iteration int, +) string { + var lastBlocker string + for _, tc := range toolCalls { + argsJSON, _ := json.Marshal(tc.Arguments) + argsPreview := utils.Truncate(string(argsJSON), 200) + logger.InfoCF("agent", fmt.Sprintf("Tool call: %s(%s)", tc.Name, argsPreview), + map[string]any{ + "agent_id": agent.ID, + "tool": tc.Name, + "iteration": iteration, + }) + + // Heartbeat lazy worktree: create worktree on first write-tool call + // Always use ai.Workspace (not GetPlanWorkDir) to avoid creating worktrees + // against stale project paths from previous plans. + if opts.Background && isWriteTool(tc.Name) && !agent.IsInWorktree(opts.SessionKey) { + taskName := "heartbeat-" + time.Now().Format("20060102") + if wt, wtErr := agent.ActivateWorktree(opts.SessionKey, taskName, agent.Workspace); wtErr == nil { + logger.InfoCF("agent", "Heartbeat worktree created", map[string]any{"branch": wt.Branch}) + } + } + + asyncCallback := hooks.OnPreToolExec(ctx, tc) + + toolStart := time.Now() + toolCtx := ctx + if wt := agent.GetWorktree(opts.SessionKey); wt != nil { + toolCtx = tools.WithWorkspaceOverride(toolCtx, wt.Path) + toolCtx = tools.WithWorktreeInfo(toolCtx, wt) + } + + toolResult := agent.Tools.ExecuteWithContext( + toolCtx, tc.Name, tc.Arguments, + opts.Channel, opts.ChatID, asyncCallback, + ) + toolDuration := time.Since(toolStart) + + hooks.OnToolExecDone(tc, toolResult, toolDuration) + + // Publish results to user + if !toolResult.Silent && toolResult.ForUser != "" && opts.SendResponse { + _ = al.bus.PublishOutbound(ctx, bus.OutboundMessage{ + Channel: opts.Channel, + ChatID: opts.ChatID, + Content: toolResult.ForUser, + }) + logger.DebugCF("agent", "Sent tool result to user", + map[string]any{"tool": tc.Name, "content_len": len(toolResult.ForUser)}) + } + + if len(toolResult.Media) > 0 && opts.SendResponse { + al.publishToolMedia(ctx, toolResult, opts) + } + + // Build tool result message + contentForLLM := toolResult.ForLLM + if contentForLLM == "" && toolResult.Err != nil { + contentForLLM = toolResult.Err.Error() + } + if toolResult.IsError || toolResult.Err != nil { + lastBlocker = contentForLLM + } + + toolResultMsg := providers.Message{ + Role: "tool", + Content: contentForLLM, + ToolCallID: tc.ID, + } + *messages = append(*messages, toolResultMsg) + agent.Sessions.AddFullMessage(opts.SessionKey, toolResultMsg) + } + return lastBlocker +} + +// forceTextResponse makes a final LLM call without tools when max iterations +// are exhausted, forcing a text response. +func (al *AgentLoop) forceTextResponse(ctx context.Context, agent *AgentInstance, messages []providers.Message) string { + logger.WarnCF("agent", "Max iterations reached, forcing final response without tools", + map[string]any{"agent_id": agent.ID}) + + forceResp, forceErr := agent.Provider.Chat(ctx, messages, nil, agent.Model, map[string]any{ + "max_tokens": agent.MaxTokens, + "temperature": agent.Temperature, + "prompt_cache_key": agent.ID, + }) + if forceErr != nil || forceResp.Content == "" { + return "" + } + content := utils.StripThinkBlocks(forceResp.Content) + if forceResp.Usage != nil && al.stats != nil { + al.stats.RecordUsage( + forceResp.Usage.PromptTokens, + forceResp.Usage.CompletionTokens, + forceResp.Usage.TotalTokens, + ) + } + return content +} + +// updateToolContexts updates the context for tools that need channel/chatID info. +func (al *AgentLoop) updateToolContexts(agent *AgentInstance, channel, chatID string) { + // Use ContextualTool interface instead of type assertions + + if tool, ok := agent.Tools.Get("message"); ok { + if mt, ok := tool.(tools.ContextualTool); ok { + mt.SetContext(channel, chatID) + } + } + + if tool, ok := agent.Tools.Get("spawn"); ok { + if st, ok := tool.(tools.ContextualTool); ok { + st.SetContext(channel, chatID) + } + } + + if tool, ok := agent.Tools.Get("subagent"); ok { + if st, ok := tool.(tools.ContextualTool); ok { + st.SetContext(channel, chatID) + } + } +} diff --git a/pkg/providers/azure/provider.go b/pkg/providers/azure/provider.go new file mode 100644 index 000000000..e0ddbbde4 --- /dev/null +++ b/pkg/providers/azure/provider.go @@ -0,0 +1,150 @@ +package azure + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "net/http" + "net/url" + "strings" + "time" + + "github.com/sipeed/picoclaw/pkg/providers/common" + "github.com/sipeed/picoclaw/pkg/providers/protocoltypes" +) + +type ( + LLMResponse = protocoltypes.LLMResponse + Message = protocoltypes.Message + ToolDefinition = protocoltypes.ToolDefinition +) + +const ( + // azureAPIVersion is the Azure OpenAI API version used for all requests. + azureAPIVersion = "2024-10-21" + defaultRequestTimeout = common.DefaultRequestTimeout +) + +// Provider implements the LLM provider interface for Azure OpenAI endpoints. +// It handles Azure-specific authentication (api-key header), URL construction +// (deployment-based), and request body formatting (max_completion_tokens, no model field). +type Provider struct { + apiKey string + apiBase string + httpClient *http.Client +} + +// Option configures the Azure Provider. +type Option func(*Provider) + +// WithRequestTimeout sets the HTTP request timeout. +func WithRequestTimeout(timeout time.Duration) Option { + return func(p *Provider) { + if timeout > 0 { + p.httpClient.Timeout = timeout + } + } +} + +// NewProvider creates a new Azure OpenAI provider. +func NewProvider(apiKey, apiBase, proxy string, opts ...Option) *Provider { + p := &Provider{ + apiKey: apiKey, + apiBase: strings.TrimRight(apiBase, "/"), + httpClient: common.NewHTTPClient(proxy), + } + + for _, opt := range opts { + if opt != nil { + opt(p) + } + } + + return p +} + +// NewProviderWithTimeout creates a new Azure OpenAI provider with a custom request timeout in seconds. +func NewProviderWithTimeout(apiKey, apiBase, proxy string, requestTimeoutSeconds int) *Provider { + return NewProvider( + apiKey, apiBase, proxy, + WithRequestTimeout(time.Duration(requestTimeoutSeconds)*time.Second), + ) +} + +// Chat sends a chat completion request to the Azure OpenAI endpoint. +// The model parameter is used as the Azure deployment name in the URL. +func (p *Provider) Chat( + ctx context.Context, + messages []Message, + tools []ToolDefinition, + model string, + options map[string]any, +) (*LLMResponse, error) { + if p.apiBase == "" { + return nil, fmt.Errorf("Azure API base not configured") + } + + // model is the deployment name for Azure OpenAI + deployment := model + + // Build Azure-specific URL safely using url.JoinPath and query encoding + // to prevent path traversal or query injection via deployment names. + base, err := url.JoinPath(p.apiBase, "openai/deployments", deployment, "chat/completions") + if err != nil { + return nil, fmt.Errorf("failed to build Azure request URL: %w", err) + } + requestURL := base + "?api-version=" + azureAPIVersion + + // Build request body — no "model" field (Azure infers from deployment URL) + requestBody := map[string]any{ + "messages": common.SerializeMessages(messages), + } + + if len(tools) > 0 { + requestBody["tools"] = tools + requestBody["tool_choice"] = "auto" + } + + // Azure OpenAI always uses max_completion_tokens + if maxTokens, ok := common.AsInt(options["max_tokens"]); ok { + requestBody["max_completion_tokens"] = maxTokens + } + + if temperature, ok := common.AsFloat(options["temperature"]); ok { + requestBody["temperature"] = temperature + } + + jsonData, err := json.Marshal(requestBody) + if err != nil { + return nil, fmt.Errorf("failed to marshal request: %w", err) + } + + req, err := http.NewRequestWithContext(ctx, "POST", requestURL, bytes.NewReader(jsonData)) + if err != nil { + return nil, fmt.Errorf("failed to create request: %w", err) + } + + // Azure uses api-key header instead of Authorization: Bearer + req.Header.Set("Content-Type", "application/json") + if p.apiKey != "" { + req.Header.Set("Api-Key", p.apiKey) + } + + resp, err := p.httpClient.Do(req) + if err != nil { + return nil, fmt.Errorf("failed to send request: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return nil, common.HandleErrorResponse(resp, p.apiBase) + } + + return common.ReadAndParseResponse(resp, p.apiBase) +} + +// GetDefaultModel returns an empty string as Azure deployments are user-configured. +func (p *Provider) GetDefaultModel() string { + return "" +} diff --git a/pkg/providers/azure/provider_test.go b/pkg/providers/azure/provider_test.go new file mode 100644 index 000000000..531b81296 --- /dev/null +++ b/pkg/providers/azure/provider_test.go @@ -0,0 +1,232 @@ +package azure + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + "time" +) + +// writeValidResponse writes a minimal valid Azure OpenAI chat completion response. +func writeValidResponse(w http.ResponseWriter) { + resp := map[string]any{ + "choices": []map[string]any{ + { + "message": map[string]any{"content": "ok"}, + "finish_reason": "stop", + }, + }, + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(resp) +} + +func TestProviderChat_AzureURLConstruction(t *testing.T) { + var capturedPath string + var capturedAPIVersion string + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + capturedPath = r.URL.Path + capturedAPIVersion = r.URL.Query().Get("api-version") + writeValidResponse(w) + })) + defer server.Close() + + p := NewProvider("test-key", server.URL, "") + _, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "my-gpt5-deployment", nil) + if err != nil { + t.Fatalf("Chat() error = %v", err) + } + + wantPath := "/openai/deployments/my-gpt5-deployment/chat/completions" + if capturedPath != wantPath { + t.Errorf("URL path = %q, want %q", capturedPath, wantPath) + } + if capturedAPIVersion != azureAPIVersion { + t.Errorf("api-version = %q, want %q", capturedAPIVersion, azureAPIVersion) + } +} + +func TestProviderChat_AzureAuthHeader(t *testing.T) { + var capturedAPIKey string + var capturedAuth string + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + capturedAPIKey = r.Header.Get("Api-Key") + capturedAuth = r.Header.Get("Authorization") + writeValidResponse(w) + })) + defer server.Close() + + p := NewProvider("test-azure-key", server.URL, "") + _, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "deployment", nil) + if err != nil { + t.Fatalf("Chat() error = %v", err) + } + + if capturedAPIKey != "test-azure-key" { + t.Errorf("api-key header = %q, want %q", capturedAPIKey, "test-azure-key") + } + if capturedAuth != "" { + t.Errorf("Authorization header should be empty, got %q", capturedAuth) + } +} + +func TestProviderChat_AzureOmitsModelFromBody(t *testing.T) { + var requestBody map[string]any + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + json.NewDecoder(r.Body).Decode(&requestBody) + writeValidResponse(w) + })) + defer server.Close() + + p := NewProvider("test-key", server.URL, "") + _, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "deployment", nil) + if err != nil { + t.Fatalf("Chat() error = %v", err) + } + + if _, exists := requestBody["model"]; exists { + t.Error("request body should not contain 'model' field for Azure OpenAI") + } +} + +func TestProviderChat_AzureUsesMaxCompletionTokens(t *testing.T) { + var requestBody map[string]any + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + json.NewDecoder(r.Body).Decode(&requestBody) + writeValidResponse(w) + })) + defer server.Close() + + p := NewProvider("test-key", server.URL, "") + _, err := p.Chat( + t.Context(), + []Message{{Role: "user", Content: "hi"}}, + nil, + "deployment", + map[string]any{"max_tokens": 2048}, + ) + if err != nil { + t.Fatalf("Chat() error = %v", err) + } + + if _, exists := requestBody["max_completion_tokens"]; !exists { + t.Error("request body should contain 'max_completion_tokens'") + } + if _, exists := requestBody["max_tokens"]; exists { + t.Error("request body should not contain 'max_tokens'") + } +} + +func TestProviderChat_AzureHTTPError(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Error(w, `{"error":"unauthorized"}`, http.StatusUnauthorized) + })) + defer server.Close() + + p := NewProvider("bad-key", server.URL, "") + _, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "deployment", nil) + if err == nil { + t.Fatal("expected error, got nil") + } +} + +func TestProviderChat_AzureParseToolCalls(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + resp := map[string]any{ + "choices": []map[string]any{ + { + "message": map[string]any{ + "content": "", + "tool_calls": []map[string]any{ + { + "id": "call_1", + "type": "function", + "function": map[string]any{ + "name": "get_weather", + "arguments": `{"city":"Seattle"}`, + }, + }, + }, + }, + "finish_reason": "tool_calls", + }, + }, + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(resp) + })) + defer server.Close() + + p := NewProvider("test-key", server.URL, "") + out, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "weather?"}}, nil, "deployment", nil) + if err != nil { + t.Fatalf("Chat() error = %v", err) + } + + if len(out.ToolCalls) != 1 { + t.Fatalf("len(ToolCalls) = %d, want 1", len(out.ToolCalls)) + } + if out.ToolCalls[0].Name != "get_weather" { + t.Errorf("ToolCalls[0].Name = %q, want %q", out.ToolCalls[0].Name, "get_weather") + } +} + +func TestProvider_AzureEmptyAPIBase(t *testing.T) { + p := NewProvider("test-key", "", "") + _, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "deployment", nil) + if err == nil { + t.Fatal("expected error for empty API base") + } +} + +func TestProvider_AzureRequestTimeoutDefault(t *testing.T) { + p := NewProvider("test-key", "https://example.com", "") + if p.httpClient.Timeout != defaultRequestTimeout { + t.Errorf("timeout = %v, want %v", p.httpClient.Timeout, defaultRequestTimeout) + } +} + +func TestProvider_AzureRequestTimeoutOverride(t *testing.T) { + p := NewProvider("test-key", "https://example.com", "", WithRequestTimeout(300*time.Second)) + if p.httpClient.Timeout != 300*time.Second { + t.Errorf("timeout = %v, want %v", p.httpClient.Timeout, 300*time.Second) + } +} + +func TestProvider_AzureNewProviderWithTimeout(t *testing.T) { + p := NewProviderWithTimeout("test-key", "https://example.com", "", 180) + if p.httpClient.Timeout != 180*time.Second { + t.Errorf("timeout = %v, want %v", p.httpClient.Timeout, 180*time.Second) + } +} + +func TestProviderChat_AzureDeploymentNameEscaped(t *testing.T) { + var capturedPath string + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + capturedPath = r.URL.RawPath // use RawPath to see percent-encoding + if capturedPath == "" { + capturedPath = r.URL.Path + } + writeValidResponse(w) + })) + defer server.Close() + + p := NewProvider("test-key", server.URL, "") + + // Deployment name with characters that could cause path injection + _, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "my deploy/../../admin", nil) + if err != nil { + t.Fatalf("Chat() error = %v", err) + } + + // The slash and special chars in the deployment name must be escaped, not treated as path separators + if capturedPath == "/openai/deployments/my deploy/../../admin/chat/completions" { + t.Fatal("deployment name was interpolated without escaping — path injection possible") + } +} diff --git a/pkg/providers/common/common.go b/pkg/providers/common/common.go new file mode 100644 index 000000000..23680a1bf --- /dev/null +++ b/pkg/providers/common/common.go @@ -0,0 +1,380 @@ +// PicoClaw - Ultra-lightweight personal AI agent +// License: MIT +// +// Copyright (c) 2026 PicoClaw contributors + +// Package common provides shared utilities used by multiple LLM provider +// implementations (openai_compat, azure, etc.). +package common + +import ( + "bufio" + "bytes" + "encoding/json" + "fmt" + "io" + "log" + "net/http" + "net/url" + "strings" + "time" + + "github.com/sipeed/picoclaw/pkg/providers/protocoltypes" +) + +// Re-export protocol types used across providers. +type ( + ToolCall = protocoltypes.ToolCall + FunctionCall = protocoltypes.FunctionCall + LLMResponse = protocoltypes.LLMResponse + UsageInfo = protocoltypes.UsageInfo + Message = protocoltypes.Message + ToolDefinition = protocoltypes.ToolDefinition + ToolFunctionDefinition = protocoltypes.ToolFunctionDefinition + ExtraContent = protocoltypes.ExtraContent + GoogleExtra = protocoltypes.GoogleExtra + ReasoningDetail = protocoltypes.ReasoningDetail +) + +const DefaultRequestTimeout = 120 * time.Second + +// NewHTTPClient creates an *http.Client with an optional proxy and the default timeout. +func NewHTTPClient(proxy string) *http.Client { + client := &http.Client{ + Timeout: DefaultRequestTimeout, + } + if proxy != "" { + parsed, err := url.Parse(proxy) + if err == nil { + // Preserve http.DefaultTransport settings (TLS, HTTP/2, timeouts, etc.) + if base, ok := http.DefaultTransport.(*http.Transport); ok { + tr := base.Clone() + tr.Proxy = http.ProxyURL(parsed) + client.Transport = tr + } else { + // Fallback: minimal transport if DefaultTransport is not *http.Transport. + client.Transport = &http.Transport{ + Proxy: http.ProxyURL(parsed), + } + } + } else { + log.Printf("common: invalid proxy URL %q: %v", proxy, err) + } + } + return client +} + +// --- Message serialization --- + +// openaiMessage is the wire-format message for OpenAI-compatible APIs. +// It mirrors protocoltypes.Message but omits SystemParts, which is an +// internal field that would be unknown to third-party endpoints. +type openaiMessage struct { + Role string `json:"role"` + Content string `json:"content"` + ReasoningContent string `json:"reasoning_content,omitempty"` + ToolCalls []ToolCall `json:"tool_calls,omitempty"` + ToolCallID string `json:"tool_call_id,omitempty"` +} + +// SerializeMessages converts internal Message structs to the OpenAI wire format. +// - Strips SystemParts (unknown to third-party endpoints) +// - Converts messages with Media to multipart content format (text + image_url parts) +// - Preserves ToolCallID, ToolCalls, and ReasoningContent for all messages +func SerializeMessages(messages []Message) []any { + out := make([]any, 0, len(messages)) + for _, m := range messages { + if len(m.Media) == 0 { + out = append(out, openaiMessage{ + Role: m.Role, + Content: m.Content, + ReasoningContent: m.ReasoningContent, + ToolCalls: m.ToolCalls, + ToolCallID: m.ToolCallID, + }) + continue + } + + // Multipart content format for messages with media + parts := make([]map[string]any, 0, 1+len(m.Media)) + if m.Content != "" { + parts = append(parts, map[string]any{ + "type": "text", + "text": m.Content, + }) + } + for _, mediaURL := range m.Media { + if strings.HasPrefix(mediaURL, "data:image/") { + parts = append(parts, map[string]any{ + "type": "image_url", + "image_url": map[string]any{ + "url": mediaURL, + }, + }) + } + } + + msg := map[string]any{ + "role": m.Role, + "content": parts, + } + if m.ToolCallID != "" { + msg["tool_call_id"] = m.ToolCallID + } + if len(m.ToolCalls) > 0 { + msg["tool_calls"] = m.ToolCalls + } + if m.ReasoningContent != "" { + msg["reasoning_content"] = m.ReasoningContent + } + out = append(out, msg) + } + return out +} + +// --- Response parsing --- + +// ParseResponse parses a JSON chat completion response body into an LLMResponse. +func ParseResponse(body io.Reader) (*LLMResponse, error) { + var apiResponse struct { + Choices []struct { + Message struct { + Content string `json:"content"` + ReasoningContent string `json:"reasoning_content"` + Reasoning string `json:"reasoning"` + ReasoningDetails []ReasoningDetail `json:"reasoning_details"` + ToolCalls []struct { + ID string `json:"id"` + Type string `json:"type"` + Function *struct { + Name string `json:"name"` + Arguments json.RawMessage `json:"arguments"` + } `json:"function"` + ExtraContent *struct { + Google *struct { + ThoughtSignature string `json:"thought_signature"` + } `json:"google"` + } `json:"extra_content"` + } `json:"tool_calls"` + } `json:"message"` + FinishReason string `json:"finish_reason"` + } `json:"choices"` + Usage *UsageInfo `json:"usage"` + } + + if err := json.NewDecoder(body).Decode(&apiResponse); err != nil { + return nil, fmt.Errorf("failed to decode response: %w", err) + } + + if len(apiResponse.Choices) == 0 { + return &LLMResponse{ + Content: "", + FinishReason: "stop", + }, nil + } + + choice := apiResponse.Choices[0] + toolCalls := make([]ToolCall, 0, len(choice.Message.ToolCalls)) + for _, tc := range choice.Message.ToolCalls { + arguments := make(map[string]any) + name := "" + + // Extract thought_signature from Gemini/Google-specific extra content + thoughtSignature := "" + if tc.ExtraContent != nil && tc.ExtraContent.Google != nil { + thoughtSignature = tc.ExtraContent.Google.ThoughtSignature + } + + if tc.Function != nil { + name = tc.Function.Name + arguments = DecodeToolCallArguments(tc.Function.Arguments, name) + } + + toolCall := ToolCall{ + ID: tc.ID, + Name: name, + Arguments: arguments, + ThoughtSignature: thoughtSignature, + } + + if thoughtSignature != "" { + toolCall.ExtraContent = &ExtraContent{ + Google: &GoogleExtra{ + ThoughtSignature: thoughtSignature, + }, + } + } + + toolCalls = append(toolCalls, toolCall) + } + + return &LLMResponse{ + Content: choice.Message.Content, + ReasoningContent: choice.Message.ReasoningContent, + Reasoning: choice.Message.Reasoning, + ReasoningDetails: choice.Message.ReasoningDetails, + ToolCalls: toolCalls, + FinishReason: choice.FinishReason, + Usage: apiResponse.Usage, + }, nil +} + +// DecodeToolCallArguments decodes a tool call's arguments from raw JSON. +func DecodeToolCallArguments(raw json.RawMessage, name string) map[string]any { + arguments := make(map[string]any) + raw = bytes.TrimSpace(raw) + if len(raw) == 0 || bytes.Equal(raw, []byte("null")) { + return arguments + } + + var decoded any + if err := json.Unmarshal(raw, &decoded); err != nil { + log.Printf("common: failed to decode tool call arguments payload for %q: %v", name, err) + arguments["raw"] = string(raw) + return arguments + } + + switch v := decoded.(type) { + case string: + if strings.TrimSpace(v) == "" { + return arguments + } + if err := json.Unmarshal([]byte(v), &arguments); err != nil { + log.Printf("common: failed to decode tool call arguments for %q: %v", name, err) + arguments["raw"] = v + } + return arguments + case map[string]any: + return v + default: + log.Printf("common: unsupported tool call arguments type for %q: %T", name, decoded) + arguments["raw"] = string(raw) + return arguments + } +} + +// --- HTTP response helpers --- + +// HandleErrorResponse reads a non-200 response body and returns an appropriate error. +func HandleErrorResponse(resp *http.Response, apiBase string) error { + contentType := resp.Header.Get("Content-Type") + body, readErr := io.ReadAll(io.LimitReader(resp.Body, 256)) + if readErr != nil { + return fmt.Errorf("failed to read response: %w", readErr) + } + if LooksLikeHTML(body, contentType) { + return WrapHTMLResponseError(resp.StatusCode, body, contentType, apiBase) + } + return fmt.Errorf( + "API request failed:\n Status: %d\n Body: %s", + resp.StatusCode, + ResponsePreview(body, 128), + ) +} + +// ReadAndParseResponse peeks at the response body to detect HTML errors, +// then parses the JSON response into an LLMResponse. +func ReadAndParseResponse(resp *http.Response, apiBase string) (*LLMResponse, error) { + contentType := resp.Header.Get("Content-Type") + reader := bufio.NewReader(resp.Body) + prefix, err := reader.Peek(256) + if err != nil && err != io.EOF && err != bufio.ErrBufferFull { + return nil, fmt.Errorf("failed to inspect response: %w", err) + } + if LooksLikeHTML(prefix, contentType) { + return nil, WrapHTMLResponseError(resp.StatusCode, prefix, contentType, apiBase) + } + out, err := ParseResponse(reader) + if err != nil { + return nil, fmt.Errorf("failed to parse JSON response: %w", err) + } + return out, nil +} + +// LooksLikeHTML checks if the response body appears to be HTML. +func LooksLikeHTML(body []byte, contentType string) bool { + contentType = strings.ToLower(strings.TrimSpace(contentType)) + if strings.Contains(contentType, "text/html") || strings.Contains(contentType, "application/xhtml+xml") { + return true + } + prefix := bytes.ToLower(leadingTrimmedPrefix(body, 128)) + return bytes.HasPrefix(prefix, []byte("" + } + if len(trimmed) <= maxLen { + return string(trimmed) + } + return string(trimmed[:maxLen]) + "..." +} + +func leadingTrimmedPrefix(body []byte, maxLen int) []byte { + i := 0 + for i < len(body) { + switch body[i] { + case ' ', '\t', '\n', '\r', '\f', '\v': + i++ + default: + end := i + maxLen + if end > len(body) { + end = len(body) + } + return body[i:end] + } + } + return nil +} + +// --- Numeric helpers --- + +// AsInt converts various numeric types to int. +func AsInt(v any) (int, bool) { + switch val := v.(type) { + case int: + return val, true + case int64: + return int(val), true + case float64: + return int(val), true + case float32: + return int(val), true + default: + return 0, false + } +} + +// AsFloat converts various numeric types to float64. +func AsFloat(v any) (float64, bool) { + switch val := v.(type) { + case float64: + return val, true + case float32: + return float64(val), true + case int: + return float64(val), true + case int64: + return float64(val), true + default: + return 0, false + } +} diff --git a/pkg/providers/common/common_test.go b/pkg/providers/common/common_test.go new file mode 100644 index 000000000..bb7e7434d --- /dev/null +++ b/pkg/providers/common/common_test.go @@ -0,0 +1,558 @@ +package common + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "testing" + + "github.com/sipeed/picoclaw/pkg/providers/protocoltypes" +) + +// --- NewHTTPClient tests --- + +func TestNewHTTPClient_DefaultTimeout(t *testing.T) { + client := NewHTTPClient("") + if client.Timeout != DefaultRequestTimeout { + t.Errorf("timeout = %v, want %v", client.Timeout, DefaultRequestTimeout) + } +} + +func TestNewHTTPClient_WithProxy(t *testing.T) { + client := NewHTTPClient("http://127.0.0.1:8080") + transport, ok := client.Transport.(*http.Transport) + if !ok || transport == nil { + t.Fatalf("expected http.Transport with proxy, got %T", client.Transport) + } + req := &http.Request{URL: &url.URL{Scheme: "https", Host: "api.example.com"}} + gotProxy, err := transport.Proxy(req) + if err != nil { + t.Fatalf("proxy function error: %v", err) + } + if gotProxy == nil || gotProxy.String() != "http://127.0.0.1:8080" { + t.Errorf("proxy = %v, want http://127.0.0.1:8080", gotProxy) + } +} + +func TestNewHTTPClient_NoProxy(t *testing.T) { + client := NewHTTPClient("") + if client.Transport != nil { + t.Errorf("expected nil transport without proxy, got %T", client.Transport) + } +} + +func TestNewHTTPClient_InvalidProxy(t *testing.T) { + // Should not panic, just log and return client without proxy + client := NewHTTPClient("://bad-url") + if client == nil { + t.Fatal("expected non-nil client even with invalid proxy") + } +} + +// --- SerializeMessages tests --- + +func TestSerializeMessages_PlainText(t *testing.T) { + messages := []Message{ + {Role: "user", Content: "hello"}, + {Role: "assistant", Content: "hi", ReasoningContent: "thinking..."}, + } + result := SerializeMessages(messages) + + data, _ := json.Marshal(result) + var msgs []map[string]any + json.Unmarshal(data, &msgs) + + if msgs[0]["content"] != "hello" { + t.Errorf("expected plain string content, got %v", msgs[0]["content"]) + } + if msgs[1]["reasoning_content"] != "thinking..." { + t.Errorf("reasoning_content not preserved, got %v", msgs[1]["reasoning_content"]) + } +} + +func TestSerializeMessages_WithMedia(t *testing.T) { + messages := []Message{ + {Role: "user", Content: "describe this", Media: []string{"data:image/png;base64,abc123"}}, + } + result := SerializeMessages(messages) + + data, _ := json.Marshal(result) + var msgs []map[string]any + json.Unmarshal(data, &msgs) + + content, ok := msgs[0]["content"].([]any) + if !ok { + t.Fatalf("expected array content for media message, got %T", msgs[0]["content"]) + } + if len(content) != 2 { + t.Fatalf("expected 2 content parts, got %d", len(content)) + } +} + +func TestSerializeMessages_MediaWithToolCallID(t *testing.T) { + messages := []Message{ + {Role: "tool", Content: "result", Media: []string{"data:image/png;base64,xyz"}, ToolCallID: "call_1"}, + } + result := SerializeMessages(messages) + + data, _ := json.Marshal(result) + var msgs []map[string]any + json.Unmarshal(data, &msgs) + + if msgs[0]["tool_call_id"] != "call_1" { + t.Errorf("tool_call_id not preserved, got %v", msgs[0]["tool_call_id"]) + } +} + +func TestSerializeMessages_StripsSystemParts(t *testing.T) { + messages := []Message{ + { + Role: "system", + Content: "you are helpful", + SystemParts: []protocoltypes.ContentBlock{ + {Type: "text", Text: "you are helpful"}, + }, + }, + } + result := SerializeMessages(messages) + + data, _ := json.Marshal(result) + if strings.Contains(string(data), "system_parts") { + t.Error("system_parts should not appear in serialized output") + } +} + +// --- ParseResponse tests --- + +func TestParseResponse_BasicContent(t *testing.T) { + body := `{"choices":[{"message":{"content":"hello world"},"finish_reason":"stop"}]}` + out, err := ParseResponse(strings.NewReader(body)) + if err != nil { + t.Fatalf("ParseResponse() error = %v", err) + } + if out.Content != "hello world" { + t.Errorf("Content = %q, want %q", out.Content, "hello world") + } + if out.FinishReason != "stop" { + t.Errorf("FinishReason = %q, want %q", out.FinishReason, "stop") + } +} + +func TestParseResponse_EmptyChoices(t *testing.T) { + body := `{"choices":[]}` + out, err := ParseResponse(strings.NewReader(body)) + if err != nil { + t.Fatalf("ParseResponse() error = %v", err) + } + if out.Content != "" { + t.Errorf("Content = %q, want empty", out.Content) + } + if out.FinishReason != "stop" { + t.Errorf("FinishReason = %q, want %q", out.FinishReason, "stop") + } +} + +func TestParseResponse_WithToolCalls(t *testing.T) { + body := `{"choices":[{"message":{"content":"","tool_calls":[{"id":"call_1","type":"function","function":{"name":"get_weather","arguments":"{\"city\":\"SF\"}"}}]},"finish_reason":"tool_calls"}]}` + out, err := ParseResponse(strings.NewReader(body)) + if err != nil { + t.Fatalf("ParseResponse() error = %v", err) + } + if len(out.ToolCalls) != 1 { + t.Fatalf("len(ToolCalls) = %d, want 1", len(out.ToolCalls)) + } + if out.ToolCalls[0].Name != "get_weather" { + t.Errorf("ToolCalls[0].Name = %q, want %q", out.ToolCalls[0].Name, "get_weather") + } + if out.ToolCalls[0].Arguments["city"] != "SF" { + t.Errorf("ToolCalls[0].Arguments[city] = %v, want SF", out.ToolCalls[0].Arguments["city"]) + } +} + +func TestParseResponse_WithUsage(t *testing.T) { + body := `{"choices":[{"message":{"content":"ok"},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15}}` + out, err := ParseResponse(strings.NewReader(body)) + if err != nil { + t.Fatalf("ParseResponse() error = %v", err) + } + if out.Usage == nil { + t.Fatal("Usage is nil") + } + if out.Usage.PromptTokens != 10 { + t.Errorf("PromptTokens = %d, want 10", out.Usage.PromptTokens) + } +} + +func TestParseResponse_WithReasoningContent(t *testing.T) { + body := `{"choices":[{"message":{"content":"2","reasoning_content":"Let me think... 1+1=2"},"finish_reason":"stop"}]}` + out, err := ParseResponse(strings.NewReader(body)) + if err != nil { + t.Fatalf("ParseResponse() error = %v", err) + } + if out.ReasoningContent != "Let me think... 1+1=2" { + t.Errorf("ReasoningContent = %q, want %q", out.ReasoningContent, "Let me think... 1+1=2") + } +} + +func TestParseResponse_InvalidJSON(t *testing.T) { + _, err := ParseResponse(strings.NewReader("not json")) + if err == nil { + t.Fatal("expected error for invalid JSON") + } +} + +// --- DecodeToolCallArguments tests --- + +func TestDecodeToolCallArguments_ObjectJSON(t *testing.T) { + raw := json.RawMessage(`{"city":"Seattle","units":"metric"}`) + args := DecodeToolCallArguments(raw, "test") + if args["city"] != "Seattle" { + t.Errorf("city = %v, want Seattle", args["city"]) + } + if args["units"] != "metric" { + t.Errorf("units = %v, want metric", args["units"]) + } +} + +func TestDecodeToolCallArguments_StringJSON(t *testing.T) { + raw := json.RawMessage(`"{\"city\":\"SF\"}"`) + args := DecodeToolCallArguments(raw, "test") + if args["city"] != "SF" { + t.Errorf("city = %v, want SF", args["city"]) + } +} + +func TestDecodeToolCallArguments_EmptyInput(t *testing.T) { + args := DecodeToolCallArguments(nil, "test") + if len(args) != 0 { + t.Errorf("expected empty map, got %v", args) + } +} + +func TestDecodeToolCallArguments_NullInput(t *testing.T) { + args := DecodeToolCallArguments(json.RawMessage(`null`), "test") + if len(args) != 0 { + t.Errorf("expected empty map, got %v", args) + } +} + +func TestDecodeToolCallArguments_InvalidJSON(t *testing.T) { + args := DecodeToolCallArguments(json.RawMessage(`not-json`), "test") + if _, ok := args["raw"]; !ok { + t.Error("expected 'raw' fallback key for invalid JSON") + } +} + +func TestDecodeToolCallArguments_EmptyStringJSON(t *testing.T) { + args := DecodeToolCallArguments(json.RawMessage(`" "`), "test") + if len(args) != 0 { + t.Errorf("expected empty map for whitespace string, got %v", args) + } +} + +// --- HandleErrorResponse tests --- + +func TestHandleErrorResponse_JSONError(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + w.Write([]byte(`{"error":"bad request"}`)) + })) + defer server.Close() + + resp, err := http.Get(server.URL) + if err != nil { + t.Fatalf("http.Get() error = %v", err) + } + defer resp.Body.Close() + err = HandleErrorResponse(resp, server.URL) + if err == nil { + t.Fatal("expected error") + } + if !strings.Contains(err.Error(), "400") { + t.Errorf("error should contain status code, got %v", err) + } + if strings.Contains(err.Error(), "HTML") { + t.Errorf("should not mention HTML for JSON error, got %v", err) + } +} + +func TestHandleErrorResponse_HTMLError(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/html") + w.WriteHeader(http.StatusBadGateway) + w.Write([]byte("bad gateway")) + })) + defer server.Close() + + resp, err := http.Get(server.URL) + if err != nil { + t.Fatalf("http.Get() error = %v", err) + } + defer resp.Body.Close() + err = HandleErrorResponse(resp, server.URL) + if err == nil { + t.Fatal("expected error") + } + if !strings.Contains(err.Error(), "HTML instead of JSON") { + t.Errorf("expected HTML error message, got %v", err) + } +} + +// --- ReadAndParseResponse tests --- + +func TestReadAndParseResponse_ValidJSON(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"choices":[{"message":{"content":"ok"},"finish_reason":"stop"}]}`)) + })) + defer server.Close() + + resp, err := http.Get(server.URL) + if err != nil { + t.Fatalf("http.Get() error = %v", err) + } + defer resp.Body.Close() + out, err := ReadAndParseResponse(resp, server.URL) + if err != nil { + t.Fatalf("ReadAndParseResponse() error = %v", err) + } + if out.Content != "ok" { + t.Errorf("Content = %q, want %q", out.Content, "ok") + } +} + +func TestReadAndParseResponse_HTMLResponse(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/html") + w.Write([]byte("login page")) + })) + defer server.Close() + + resp, err := http.Get(server.URL) + if err != nil { + t.Fatalf("http.Get() error = %v", err) + } + defer resp.Body.Close() + _, err = ReadAndParseResponse(resp, server.URL) + if err == nil { + t.Fatal("expected error for HTML response") + } + if !strings.Contains(err.Error(), "HTML instead of JSON") { + t.Errorf("expected HTML error, got %v", err) + } +} + +// --- LooksLikeHTML tests --- + +func TestLooksLikeHTML_ContentTypeHTML(t *testing.T) { + if !LooksLikeHTML(nil, "text/html; charset=utf-8") { + t.Error("expected true for text/html content type") + } +} + +func TestLooksLikeHTML_ContentTypeXHTML(t *testing.T) { + if !LooksLikeHTML(nil, "application/xhtml+xml") { + t.Error("expected true for xhtml content type") + } +} + +func TestLooksLikeHTML_BodyPrefix(t *testing.T) { + tests := []struct { + name string + body string + }{ + {"doctype", ""}, + {"html tag", ""}, + {"head tag", ""}, + {"body tag", "<body>content"}, + {"whitespace before", " \n\t<!DOCTYPE html>"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if !LooksLikeHTML([]byte(tt.body), "application/json") { + t.Errorf("expected true for body %q", tt.body) + } + }) + } +} + +func TestLooksLikeHTML_NotHTML(t *testing.T) { + if LooksLikeHTML([]byte(`{"error":"bad"}`), "application/json") { + t.Error("expected false for JSON body") + } +} + +// --- ResponsePreview tests --- + +func TestResponsePreview_Short(t *testing.T) { + got := ResponsePreview([]byte("hello"), 128) + if got != "hello" { + t.Errorf("got %q, want %q", got, "hello") + } +} + +func TestResponsePreview_Truncated(t *testing.T) { + body := strings.Repeat("a", 200) + got := ResponsePreview([]byte(body), 128) + if len(got) != 131 { // 128 + "..." + t.Errorf("len = %d, want 131", len(got)) + } + if !strings.HasSuffix(got, "...") { + t.Error("expected ... suffix") + } +} + +func TestResponsePreview_Empty(t *testing.T) { + got := ResponsePreview([]byte(""), 128) + if got != "<empty>" { + t.Errorf("got %q, want %q", got, "<empty>") + } +} + +func TestResponsePreview_Whitespace(t *testing.T) { + got := ResponsePreview([]byte(" \n\t "), 128) + if got != "<empty>" { + t.Errorf("got %q, want %q for whitespace-only body", got, "<empty>") + } +} + +// --- AsInt tests --- + +func TestAsInt(t *testing.T) { + tests := []struct { + name string + val any + want int + ok bool + }{ + {"int", 42, 42, true}, + {"int64", int64(99), 99, true}, + {"float64", float64(512), 512, true}, + {"float32", float32(256), 256, true}, + {"string", "nope", 0, false}, + {"nil", nil, 0, false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, ok := AsInt(tt.val) + if ok != tt.ok || got != tt.want { + t.Errorf("AsInt(%v) = (%d, %v), want (%d, %v)", tt.val, got, ok, tt.want, tt.ok) + } + }) + } +} + +// --- AsFloat tests --- + +func TestAsFloat(t *testing.T) { + tests := []struct { + name string + val any + want float64 + ok bool + }{ + {"float64", float64(0.7), 0.7, true}, + {"float32", float32(0.5), float64(float32(0.5)), true}, + {"int", 1, 1.0, true}, + {"int64", int64(100), 100.0, true}, + {"string", "nope", 0, false}, + {"nil", nil, 0, false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, ok := AsFloat(tt.val) + if ok != tt.ok || got != tt.want { + t.Errorf("AsFloat(%v) = (%f, %v), want (%f, %v)", tt.val, got, ok, tt.want, tt.ok) + } + }) + } +} + +// --- WrapHTMLResponseError tests --- + +func TestWrapHTMLResponseError(t *testing.T) { + err := WrapHTMLResponseError(502, []byte("<html>bad</html>"), "text/html", "https://api.example.com") + if err == nil { + t.Fatal("expected error") + } + msg := err.Error() + if !strings.Contains(msg, "502") { + t.Errorf("expected status code in error, got %v", msg) + } + if !strings.Contains(msg, "https://api.example.com") { + t.Errorf("expected api base in error, got %v", msg) + } + if !strings.Contains(msg, "HTML instead of JSON") { + t.Errorf("expected HTML mention in error, got %v", msg) + } +} + +// --- HandleErrorResponse with read failure --- + +func TestHandleErrorResponse_EmptyBody(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusInternalServerError) + // empty body + })) + defer server.Close() + + resp, err := http.Get(server.URL) + if err != nil { + t.Fatalf("http.Get() error = %v", err) + } + defer resp.Body.Close() + err = HandleErrorResponse(resp, server.URL) + if err == nil { + t.Fatal("expected error") + } + if !strings.Contains(err.Error(), "500") { + t.Errorf("expected status code, got %v", err) + } +} + +// --- ReadAndParseResponse with invalid JSON --- + +func TestReadAndParseResponse_InvalidJSON(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.Write([]byte("not valid json")) + })) + defer server.Close() + + resp, err := http.Get(server.URL) + if err != nil { + t.Fatalf("http.Get() error = %v", err) + } + defer resp.Body.Close() + _, err = ReadAndParseResponse(resp, server.URL) + if err == nil { + t.Fatal("expected error for invalid JSON") + } +} + +// --- ParseResponse with thought_signature (Google/Gemini) --- + +func TestParseResponse_WithThoughtSignature(t *testing.T) { + body := `{"choices":[{"message":{"content":"","tool_calls":[{"id":"call_1","type":"function","function":{"name":"test_tool","arguments":"{}"},"extra_content":{"google":{"thought_signature":"sig123"}}}]},"finish_reason":"tool_calls"}]}` + out, err := ParseResponse(strings.NewReader(body)) + if err != nil { + t.Fatalf("ParseResponse() error = %v", err) + } + if len(out.ToolCalls) != 1 { + t.Fatalf("len(ToolCalls) = %d, want 1", len(out.ToolCalls)) + } + if out.ToolCalls[0].ThoughtSignature != "sig123" { + t.Errorf("ThoughtSignature = %q, want %q", out.ToolCalls[0].ThoughtSignature, "sig123") + } + if out.ToolCalls[0].ExtraContent == nil || out.ToolCalls[0].ExtraContent.Google == nil { + t.Fatal("ExtraContent.Google is nil") + } + if out.ToolCalls[0].ExtraContent.Google.ThoughtSignature != "sig123" { + t.Errorf("ExtraContent.Google.ThoughtSignature = %q, want %q", + out.ToolCalls[0].ExtraContent.Google.ThoughtSignature, "sig123") + } +}