diff --git a/pkg/agent/context.go b/pkg/agent/context.go index 1680b5eaa..cfa40e69a 100644 --- a/pkg/agent/context.go +++ b/pkg/agent/context.go @@ -16,136 +16,17 @@ import ( "github.com/sipeed/picoclaw/pkg/logger" "github.com/sipeed/picoclaw/pkg/providers" "github.com/sipeed/picoclaw/pkg/skills" - "github.com/sipeed/picoclaw/pkg/tools" "github.com/sipeed/picoclaw/pkg/utils" ) -const orchestrationGuidance = `## Orchestration - - - -You are the conductor, not the performer. **Your primary job is to delegate, not to implement.** - - - -### spawn (non-blocking) — DEFAULT choice - -Returns immediately. Use for any task that can run independently. - -Call the spawn tool with JSON arguments like this: - - - -Tool: spawn - -Arguments: {"task": "Examine pkg/auth/ and report middleware pattern", "preset": "scout", "label": "auth-scout"} - - - -Tool: spawn - -Arguments: {"task": "Implement rate limiter in pkg/ratelimit/ with tests", "preset": "coder", "label": "rate-limiter"} - - - -### subagent (blocking) — only when you need the answer NOW - -Blocks until the subagent finishes. Use only when you cannot proceed without the result. - -Does not take a preset — it runs with default tools. - - - -Tool: subagent - -Arguments: {"task": "Read pkg/config/config.go and list all SubagentsConfig fields", "label": "config-check"} - - - -### When to use which - -- spawn: parallel tasks, independent work, implementation, long analysis, >2 tool calls - -- subagent: you need the result before your next decision - -- inline: single quick tool call where delegation overhead is wasteful - - - -### Presets (for spawn only) - -| preset | role | can write | can exec | - -|--------|------|-----------|----------| - -| scout | explore, investigate | no | no | - -| analyst | analyze, run tests | no | go test/vet, git | - -| coder | implement + verify | yes (sandbox) | test/lint/fmt | - -| worker | build + install | yes (sandbox) | build/package mgr | - -| coordinator | orchestrate others | yes (sandbox) | general + spawn | - - - -### Parallel spawning - -Spawn multiple independent tasks at once — do NOT wait between them: - - - -Tool: spawn - -Arguments: {"task": "Analyze error handling patterns in pkg/providers/", "preset": "analyst", "label": "error-patterns"} - - - -Tool: spawn - -Arguments: {"task": "List all HTTP endpoints in pkg/miniapp/", "preset": "scout", "label": "endpoints"} - - - -After spawning, record the assignment in ## Orchestration > Delegated in MEMORY.md. - -When results come back, synthesize findings and decide the next fork. - - - -### Subagent escalation - -Deliberate subagents (coder/worker/coordinator) may ask you questions or submit plans for review. - -When a subagent question appears, respond with the appropriate tool: - -- answer_subagent: Answer a subagent's clarifying question - -- review_subagent_plan: Approve or reject a subagent's execution plan (decision: "approved" or rejection feedback) - - - -### Orchestration Memory - -Maintain these sections in MEMORY.md under ## Orchestration: - -- **Delegated**: Active subagent assignments (task ID, preset, description) - -- **Findings**: Synthesized results from completed subagents - -- **Decisions**: Key architectural/implementation decisions made during orchestration` - type ContextBuilder struct { - workspace string - workDir string // session-specific working directory (worktree or project subdir) - skillsLoader *skills.SkillsLoader - memory *MemoryStore - tools *tools.ToolRegistry // Direct reference to tool registry - peerNote string // set per-call from loop.go for peer session awareness - orchestrationEnabled bool // set from AgentLoop when --orchestration flag is used - toolDiscoveryBM25 bool - toolDiscoveryRegex bool + contextBuilderExt // fork-specific fields (see context_ext.go) + + workspace string + skillsLoader *skills.SkillsLoader + memory *MemoryStore + toolDiscoveryBM25 bool + toolDiscoveryRegex bool // Cache for system prompt to avoid rebuilding on every call. // This fixes issue #607: repeated reprocessing of the entire context. @@ -200,27 +81,6 @@ func NewContextBuilder(workspace string) *ContextBuilder { } } -// SetToolsRegistry sets the tools registry for dynamic tool summary generation. -func (cb *ContextBuilder) SetToolsRegistry(registry *tools.ToolRegistry) { - cb.tools = registry -} - -// SetWorkDir sets the session-specific working directory (e.g., worktree path -// or project subdirectory). Bootstrap files found here take priority over workspace. -func (cb *ContextBuilder) SetWorkDir(dir string) { - cb.workDir = dir -} - -// SetPeerNote sets the peer session awareness note for the current call. -func (cb *ContextBuilder) SetPeerNote(note string) { - cb.peerNote = note -} - -// SetOrchestrationEnabled sets whether orchestration is enabled. -func (cb *ContextBuilder) SetOrchestrationEnabled(enabled bool) { - cb.orchestrationEnabled = enabled -} - func (cb *ContextBuilder) getIdentity() string { workspacePath, _ := filepath.Abs(filepath.Join(cb.workspace)) toolDiscovery := cb.getDiscoveryRule() @@ -1056,114 +916,3 @@ func (cb *ContextBuilder) LoadSkill(name string) (string, bool) { func (cb *ContextBuilder) ListSkills() []skills.SkillInfo { return cb.skillsLoader.ListSkills() } - -// Memory returns the underlying MemoryStore for direct plan queries. -func (cb *ContextBuilder) Memory() *MemoryStore { - return cb.memory -} - -// ---------- Plan passthrough methods ---------- - -// ReadMemory reads the long-term memory (MEMORY.md). -func (cb *ContextBuilder) ReadMemory() string { - return cb.memory.ReadLongTerm() -} - -// WriteMemory writes content to the long-term memory file. -func (cb *ContextBuilder) WriteMemory(content string) error { - return cb.memory.WriteLongTerm(content) -} - -// ClearMemory removes the long-term memory file. -func (cb *ContextBuilder) ClearMemory() error { - return cb.memory.ClearLongTerm() -} - -// HasActivePlan returns true if MEMORY.md contains an active plan. -func (cb *ContextBuilder) HasActivePlan() bool { - return cb.memory.HasActivePlan() -} - -// GetPlanStatus returns the plan status: "interviewing", "executing", or "". -func (cb *ContextBuilder) GetPlanStatus() string { - return cb.memory.GetPlanStatus() -} - -// IsPlanComplete returns true if all steps in all phases are [x]. -func (cb *ContextBuilder) IsPlanComplete() bool { - return cb.memory.IsPlanComplete() -} - -// IsCurrentPhaseComplete returns true if all steps in the current phase are [x]. -func (cb *ContextBuilder) IsCurrentPhaseComplete() bool { - return cb.memory.IsCurrentPhaseComplete() -} - -// AdvancePhase increments the current phase number by 1. -func (cb *ContextBuilder) AdvancePhase() error { - return cb.memory.AdvancePhase() -} - -// SetCurrentPhase sets the current phase number to n. -func (cb *ContextBuilder) SetCurrentPhase(n int) error { - return cb.memory.SetPhase(n) -} - -// GetCurrentPhase returns the current phase number. -func (cb *ContextBuilder) GetCurrentPhase() int { - return cb.memory.GetCurrentPhase() -} - -// GetTotalPhases returns the total number of phases in the plan. -func (cb *ContextBuilder) GetTotalPhases() int { - return cb.memory.GetTotalPhases() -} - -// FormatPlanDisplay returns a user-facing display of the full plan. -func (cb *ContextBuilder) FormatPlanDisplay() string { - return cb.memory.FormatPlanDisplay() -} - -// MarkStep marks a step as done in the specified phase. -func (cb *ContextBuilder) MarkStep(phase, step int) error { - return cb.memory.MarkStep(phase, step) -} - -// AddStep appends a new step to the given phase. -func (cb *ContextBuilder) AddStep(phase int, desc string) error { - return cb.memory.AddStep(phase, desc) -} - -// ValidatePlanStructure validates plan structure for interview->review transition. -func (cb *ContextBuilder) ValidatePlanStructure() error { - return cb.memory.ValidatePlanStructure() -} - -// SetPlanStatus sets the plan status. -func (cb *ContextBuilder) SetPlanStatus(status string) error { - return cb.memory.SetStatus(status) -} - -// GetPlanWorkDir returns the WorkDir from the plan metadata, or "". -func (cb *ContextBuilder) GetPlanWorkDir() string { - return cb.memory.GetPlanWorkDir() -} - -// GetPlanTaskName returns the task description from the plan metadata, or "". -func (cb *ContextBuilder) GetPlanTaskName() string { - return cb.memory.GetPlanTaskName() -} - -// GetSkillsInfo returns information about loaded skills. -func (cb *ContextBuilder) GetSkillsInfo() map[string]any { - allSkills := cb.skillsLoader.ListSkills() - skillNames := make([]string, 0, len(allSkills)) - for _, s := range allSkills { - skillNames = append(skillNames, s.Name) - } - return map[string]any{ - "total": len(allSkills), - "available": len(allSkills), - "names": skillNames, - } -} diff --git a/pkg/agent/context_ext.go b/pkg/agent/context_ext.go new file mode 100644 index 000000000..61b245488 --- /dev/null +++ b/pkg/agent/context_ext.go @@ -0,0 +1,68 @@ +package agent + +import "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. +// Upstream additions to ContextBuilder won't conflict with these fields. +type contextBuilderExt struct { + workDir string // session-specific working directory (worktree or project subdir) + tools *tools.ToolRegistry // Direct reference to tool registry + peerNote string // set per-call from loop.go for peer session awareness + orchestrationEnabled bool // set from AgentLoop when --orchestration flag is used +} + +// SetToolsRegistry sets the tools registry for dynamic tool summary generation. +func (cb *ContextBuilder) SetToolsRegistry(registry *tools.ToolRegistry) { + cb.tools = registry +} + +// SetWorkDir sets the session-specific working directory (e.g., worktree path +// or project subdirectory). Bootstrap files found here take priority over workspace. +func (cb *ContextBuilder) SetWorkDir(dir string) { + cb.workDir = dir +} + +// SetPeerNote sets the peer session awareness note for the current call. +func (cb *ContextBuilder) SetPeerNote(note string) { + cb.peerNote = note +} + +// SetOrchestrationEnabled sets whether orchestration is enabled. +func (cb *ContextBuilder) SetOrchestrationEnabled(enabled bool) { + cb.orchestrationEnabled = enabled +} + +// Memory returns the underlying MemoryStore for direct plan queries. +func (cb *ContextBuilder) Memory() *MemoryStore { + return cb.memory +} + +// ReadMemory reads the long-term memory (MEMORY.md). +func (cb *ContextBuilder) ReadMemory() string { + return cb.memory.ReadLongTerm() +} + +// WriteMemory writes content to the long-term memory file. +func (cb *ContextBuilder) WriteMemory(content string) error { + return cb.memory.WriteLongTerm(content) +} + +// ClearMemory removes the long-term memory file. +func (cb *ContextBuilder) ClearMemory() error { + return cb.memory.ClearLongTerm() +} + +// GetSkillsInfo returns information about loaded skills. +func (cb *ContextBuilder) GetSkillsInfo() map[string]any { + allSkills := cb.skillsLoader.ListSkills() + skillNames := make([]string, 0, len(allSkills)) + for _, s := range allSkills { + skillNames = append(skillNames, s.Name) + } + return map[string]any{ + "total": len(allSkills), + "available": len(allSkills), + "names": skillNames, + } +} diff --git a/pkg/agent/context_orch.go b/pkg/agent/context_orch.go new file mode 100644 index 000000000..4c16efcfe --- /dev/null +++ b/pkg/agent/context_orch.go @@ -0,0 +1,117 @@ +package agent + +const orchestrationGuidance = `## Orchestration + + + +You are the conductor, not the performer. **Your primary job is to delegate, not to implement.** + + + +### spawn (non-blocking) — DEFAULT choice + +Returns immediately. Use for any task that can run independently. + +Call the spawn tool with JSON arguments like this: + + + +Tool: spawn + +Arguments: {"task": "Examine pkg/auth/ and report middleware pattern", "preset": "scout", "label": "auth-scout"} + + + +Tool: spawn + +Arguments: {"task": "Implement rate limiter in pkg/ratelimit/ with tests", "preset": "coder", "label": "rate-limiter"} + + + +### subagent (blocking) — only when you need the answer NOW + +Blocks until the subagent finishes. Use only when you cannot proceed without the result. + +Does not take a preset — it runs with default tools. + + + +Tool: subagent + +Arguments: {"task": "Read pkg/config/config.go and list all SubagentsConfig fields", "label": "config-check"} + + + +### When to use which + +- spawn: parallel tasks, independent work, implementation, long analysis, >2 tool calls + +- subagent: you need the result before your next decision + +- inline: single quick tool call where delegation overhead is wasteful + + + +### Presets (for spawn only) + +| preset | role | can write | can exec | + +|--------|------|-----------|----------| + +| scout | explore, investigate | no | no | + +| analyst | analyze, run tests | no | go test/vet, git | + +| coder | implement + verify | yes (sandbox) | test/lint/fmt | + +| worker | build + install | yes (sandbox) | build/package mgr | + +| coordinator | orchestrate others | yes (sandbox) | general + spawn | + + + +### Parallel spawning + +Spawn multiple independent tasks at once — do NOT wait between them: + + + +Tool: spawn + +Arguments: {"task": "Analyze error handling patterns in pkg/providers/", "preset": "analyst", "label": "error-patterns"} + + + +Tool: spawn + +Arguments: {"task": "List all HTTP endpoints in pkg/miniapp/", "preset": "scout", "label": "endpoints"} + + + +After spawning, record the assignment in ## Orchestration > Delegated in MEMORY.md. + +When results come back, synthesize findings and decide the next fork. + + + +### Subagent escalation + +Deliberate subagents (coder/worker/coordinator) may ask you questions or submit plans for review. + +When a subagent question appears, respond with the appropriate tool: + +- answer_subagent: Answer a subagent's clarifying question + +- review_subagent_plan: Approve or reject a subagent's execution plan (decision: "approved" or rejection feedback) + + + +### Orchestration Memory + +Maintain these sections in MEMORY.md under ## Orchestration: + +- **Delegated**: Active subagent assignments (task ID, preset, description) + +- **Findings**: Synthesized results from completed subagents + +- **Decisions**: Key architectural/implementation decisions made during orchestration` diff --git a/pkg/agent/context_plan.go b/pkg/agent/context_plan.go new file mode 100644 index 000000000..a164a8333 --- /dev/null +++ b/pkg/agent/context_plan.go @@ -0,0 +1,79 @@ +package agent + +// ---------- Plan passthrough methods ---------- +// These delegate to MemoryStore and are separated to reduce upstream conflicts. + +// HasActivePlan returns true if MEMORY.md contains an active plan. +func (cb *ContextBuilder) HasActivePlan() bool { + return cb.memory.HasActivePlan() +} + +// GetPlanStatus returns the plan status: "interviewing", "executing", or "". +func (cb *ContextBuilder) GetPlanStatus() string { + return cb.memory.GetPlanStatus() +} + +// IsPlanComplete returns true if all steps in all phases are [x]. +func (cb *ContextBuilder) IsPlanComplete() bool { + return cb.memory.IsPlanComplete() +} + +// IsCurrentPhaseComplete returns true if all steps in the current phase are [x]. +func (cb *ContextBuilder) IsCurrentPhaseComplete() bool { + return cb.memory.IsCurrentPhaseComplete() +} + +// AdvancePhase increments the current phase number by 1. +func (cb *ContextBuilder) AdvancePhase() error { + return cb.memory.AdvancePhase() +} + +// SetCurrentPhase sets the current phase number to n. +func (cb *ContextBuilder) SetCurrentPhase(n int) error { + return cb.memory.SetPhase(n) +} + +// GetCurrentPhase returns the current phase number. +func (cb *ContextBuilder) GetCurrentPhase() int { + return cb.memory.GetCurrentPhase() +} + +// GetTotalPhases returns the total number of phases in the plan. +func (cb *ContextBuilder) GetTotalPhases() int { + return cb.memory.GetTotalPhases() +} + +// FormatPlanDisplay returns a user-facing display of the full plan. +func (cb *ContextBuilder) FormatPlanDisplay() string { + return cb.memory.FormatPlanDisplay() +} + +// MarkStep marks a step as done in the specified phase. +func (cb *ContextBuilder) MarkStep(phase, step int) error { + return cb.memory.MarkStep(phase, step) +} + +// AddStep appends a new step to the given phase. +func (cb *ContextBuilder) AddStep(phase int, desc string) error { + return cb.memory.AddStep(phase, desc) +} + +// ValidatePlanStructure validates plan structure for interview->review transition. +func (cb *ContextBuilder) ValidatePlanStructure() error { + return cb.memory.ValidatePlanStructure() +} + +// SetPlanStatus sets the plan status. +func (cb *ContextBuilder) SetPlanStatus(status string) error { + return cb.memory.SetStatus(status) +} + +// GetPlanWorkDir returns the WorkDir from the plan metadata, or "". +func (cb *ContextBuilder) GetPlanWorkDir() string { + return cb.memory.GetPlanWorkDir() +} + +// GetPlanTaskName returns the task description from the plan metadata, or "". +func (cb *ContextBuilder) GetPlanTaskName() string { + return cb.memory.GetPlanTaskName() +} diff --git a/pkg/agent/instance.go b/pkg/agent/instance.go index 809c89d49..0ac3d8894 100644 --- a/pkg/agent/instance.go +++ b/pkg/agent/instance.go @@ -7,7 +7,6 @@ import ( "path/filepath" "regexp" "strings" - "sync" "github.com/sipeed/picoclaw/pkg/config" "github.com/sipeed/picoclaw/pkg/git" @@ -20,6 +19,8 @@ import ( // AgentInstance represents a fully configured agent with its own workspace, // session manager, context builder, and tool registry. type AgentInstance struct { + instanceExt // fork-specific fields (see instance_ext.go) + ID string Name string Model string @@ -37,8 +38,6 @@ type AgentInstance struct { Sessions *session.LegacyAdapter ContextBuilder *ContextBuilder Tools *tools.ToolRegistry - Subagents *config.SubagentsConfig - SkillsFilter []string Candidates []providers.FallbackCandidate PlanModel string PlanFallbacks []string @@ -51,18 +50,6 @@ type AgentInstance struct { // LightCandidates holds the resolved provider candidates for the light model. // Pre-computed at agent creation to avoid repeated model_list lookups at runtime. LightCandidates []providers.FallbackCandidate - - // SubagentMgr is set during registerSharedTools when orchestration is enabled. - // Used by runAgentLoop to wait for spawned subagents before worktree cleanup. - SubagentMgr *tools.SubagentManager - - // Interview staleness tracking: consecutive turns where MEMORY.md was not updated. - interviewStaleCount int - interviewMemoryLen int - - // Per-session worktree isolation - worktrees map[string]*git.WorktreeInfo // sessionKey → worktree - worktreeMu sync.RWMutex } // Close releases resources held by the agent instance. @@ -295,6 +282,10 @@ func NewAgentInstance( } return &AgentInstance{ + instanceExt: instanceExt{ + Subagents: subagents, + SkillsFilter: skillsFilter, + }, ID: agentID, Name: agentName, Model: model, @@ -312,8 +303,6 @@ func NewAgentInstance( Sessions: sessionsManager, ContextBuilder: contextBuilder, Tools: toolsRegistry, - Subagents: subagents, - SkillsFilter: skillsFilter, Candidates: candidates, PlanModel: planModel, PlanFallbacks: planFallbacks, @@ -370,94 +359,6 @@ func resolvePlanFallbacks(agentCfg *config.AgentConfig, defaults *config.AgentDe return defaults.PlanModelFallbacks } -// ActivateWorktree creates a worktree for a session. -// projectDir is the git repository to create the worktree in. -// If empty, falls back to ai.Workspace. -// Worktree path: /.worktrees// -func (ai *AgentInstance) ActivateWorktree(sessionKey, taskName, projectDir string) (*git.WorktreeInfo, error) { - if projectDir == "" { - projectDir = ai.Workspace - } - - repoRoot := git.FindRepoRoot(projectDir) - if repoRoot == "" { - return nil, fmt.Errorf("directory is not a git repository: %s", projectDir) - } - - branchName := git.SanitizeBranchName(taskName) - baseName := git.BranchBaseName(branchName) - wtPath := filepath.Join(ai.Workspace, ".worktrees", baseName) - - wt, err := git.CreateWorktree(repoRoot, wtPath, branchName) - if err != nil { - return nil, err - } - - ai.worktreeMu.Lock() - if ai.worktrees == nil { - ai.worktrees = make(map[string]*git.WorktreeInfo) - } - ai.worktrees[sessionKey] = wt - ai.worktreeMu.Unlock() - - return wt, nil -} - -// DeactivateWorktree safe-disposes the session's worktree. -func (ai *AgentInstance) DeactivateWorktree(sessionKey, commitMsg string, discard bool) (*git.DisposeResult, error) { - ai.worktreeMu.Lock() - wt, ok := ai.worktrees[sessionKey] - if ok { - delete(ai.worktrees, sessionKey) - } - ai.worktreeMu.Unlock() - - if !ok || wt == nil { - return nil, nil - } - - repoRoot := git.FindRepoRoot(ai.Workspace) - if repoRoot == "" { - return nil, fmt.Errorf("workspace is not a git repository") - } - - // Even on discard, SafeDispose auto-commits first for safety - if commitMsg != "" && git.HasUncommittedChanges(wt.Path) { - _ = git.AutoCommit(wt.Path, commitMsg) - } - - result := git.SafeDispose(repoRoot, wt) - return &result, nil -} - -// GetWorktree returns the session's active worktree, or nil. -func (ai *AgentInstance) GetWorktree(sessionKey string) *git.WorktreeInfo { - ai.worktreeMu.RLock() - defer ai.worktreeMu.RUnlock() - return ai.worktrees[sessionKey] -} - -// IsInWorktree returns true if the session has an active worktree. -func (ai *AgentInstance) IsInWorktree(sessionKey string) bool { - return ai.GetWorktree(sessionKey) != nil -} - -// EffectiveWorkspace returns worktree path for session, or original Workspace. -func (ai *AgentInstance) EffectiveWorkspace(sessionKey string) string { - if wt := ai.GetWorktree(sessionKey); wt != nil { - return wt.Path - } - return ai.Workspace -} - -// GetWorktreeBranch returns the branch name for the session's worktree, or "". -func (ai *AgentInstance) GetWorktreeBranch(sessionKey string) string { - if wt := ai.GetWorktree(sessionKey); wt != nil { - return wt.Branch - } - return "" -} - func compilePatterns(patterns []string) []*regexp.Regexp { compiled := make([]*regexp.Regexp, 0, len(patterns)) for _, p := range patterns { diff --git a/pkg/agent/instance_ext.go b/pkg/agent/instance_ext.go new file mode 100644 index 000000000..10aa835e9 --- /dev/null +++ b/pkg/agent/instance_ext.go @@ -0,0 +1,122 @@ +package agent + +import ( + "fmt" + "sync" + + "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/git" + "github.com/sipeed/picoclaw/pkg/tools" +) + +// instanceExt holds fork-specific fields for AgentInstance. +// Embedded in AgentInstance so existing field access continues to work. +type instanceExt struct { + // SubagentMgr is set during registerSharedTools when orchestration is enabled. + // 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 + + // Per-session worktree isolation + worktrees map[string]*git.WorktreeInfo // sessionKey → worktree + worktreeMu sync.RWMutex +} + +// ActivateWorktree creates a worktree for a session. +// projectDir is the git repository to create the worktree in. +// If empty, falls back to ai.Workspace. +// Worktree path: /.worktrees// +func (ai *AgentInstance) ActivateWorktree(sessionKey, taskName, projectDir string) (*git.WorktreeInfo, error) { + if projectDir == "" { + projectDir = ai.Workspace + } + + repoRoot := git.FindRepoRoot(projectDir) + if repoRoot == "" { + return nil, fmt.Errorf("directory is not a git repository: %s", projectDir) + } + + branchName := git.SanitizeBranchName(taskName) + baseName := git.BranchBaseName(branchName) + wtPath := ai.worktreePath(baseName) + + wt, err := git.CreateWorktree(repoRoot, wtPath, branchName) + if err != nil { + return nil, err + } + + ai.worktreeMu.Lock() + if ai.worktrees == nil { + ai.worktrees = make(map[string]*git.WorktreeInfo) + } + ai.worktrees[sessionKey] = wt + ai.worktreeMu.Unlock() + + return wt, nil +} + +// worktreePath returns the standard path for a worktree under the workspace. +func (ai *AgentInstance) worktreePath(baseName string) string { + return ai.Workspace + "/.worktrees/" + baseName +} + +// DeactivateWorktree safe-disposes the session's worktree. +func (ai *AgentInstance) DeactivateWorktree(sessionKey, commitMsg string, discard bool) (*git.DisposeResult, error) { + ai.worktreeMu.Lock() + wt, ok := ai.worktrees[sessionKey] + if ok { + delete(ai.worktrees, sessionKey) + } + ai.worktreeMu.Unlock() + + if !ok || wt == nil { + return nil, nil + } + + repoRoot := git.FindRepoRoot(ai.Workspace) + if repoRoot == "" { + return nil, fmt.Errorf("workspace is not a git repository") + } + + // Even on discard, SafeDispose auto-commits first for safety + if commitMsg != "" && git.HasUncommittedChanges(wt.Path) { + _ = git.AutoCommit(wt.Path, commitMsg) + } + + result := git.SafeDispose(repoRoot, wt) + return &result, nil +} + +// GetWorktree returns the session's active worktree, or nil. +func (ai *AgentInstance) GetWorktree(sessionKey string) *git.WorktreeInfo { + ai.worktreeMu.RLock() + defer ai.worktreeMu.RUnlock() + return ai.worktrees[sessionKey] +} + +// IsInWorktree returns true if the session has an active worktree. +func (ai *AgentInstance) IsInWorktree(sessionKey string) bool { + return ai.GetWorktree(sessionKey) != nil +} + +// EffectiveWorkspace returns worktree path for session, or original Workspace. +func (ai *AgentInstance) EffectiveWorkspace(sessionKey string) string { + if wt := ai.GetWorktree(sessionKey); wt != nil { + return wt.Path + } + return ai.Workspace +} + +// GetWorktreeBranch returns the branch name for the session's worktree, or "". +func (ai *AgentInstance) GetWorktreeBranch(sessionKey string) string { + if wt := ai.GetWorktree(sessionKey); wt != nil { + return wt.Branch + } + return "" +} diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 259f2cd28..2117e2f39 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -39,6 +39,8 @@ import ( ) type AgentLoop struct { + loopExt // fork-specific fields (see loop_ext.go) + bus *bus.MessageBus cfg *config.Config @@ -47,8 +49,6 @@ type AgentLoop struct { state *state.Manager - stats *stats.Tracker // nil when --stats not passed - running atomic.Bool summarizing sync.Map @@ -67,16 +67,6 @@ type AgentLoop struct { providerCache map[string]providers.LLMProvider - planStartPending bool // set by /plan start to trigger LLM execution - - planClearHistory bool // set by /plan start clear to wipe history on transition - - sessionLocks sync.Map // sessionKey → *sessionSemaphore - - activeTasks sync.Map // sessionKey → *activeTask - - sessions *SessionTracker - lastSystemPrompt atomic.Value // string — last system prompt sent to LLM promptDirty atomic.Bool // true = rebuild needed on next GetSystemPrompt read @@ -84,16 +74,6 @@ type AgentLoop struct { OnStateChange func() // called on plan/session/skills mutations OnUserMessage func() // called when a real user message is processed - - saveConfig func(*config.Config) error - - onHeartbeatThreadUpdate func(int) - - orchBroadcaster *orch.Broadcaster // nil when --orchestration not set - - orchReporter orch.AgentReporter // always non-nil (Noop when disabled) - - done chan struct{} // closed by Close() to stop background goroutines } // processOptions configures how a message is processed @@ -186,6 +166,14 @@ func NewAgentLoop( } al := &AgentLoop{ + loopExt: loopExt{ + stats: statsTracker, + sessions: NewSessionTracker(), + orchBroadcaster: orchBroadcaster, + orchReporter: orchReporter, + done: make(chan struct{}), + }, + bus: msgBus, cfg: cfg, @@ -194,22 +182,12 @@ func NewAgentLoop( state: stateManager, - stats: statsTracker, - summarizing: sync.Map{}, fallback: fallbackChain, providerCache: providerCache, - sessions: NewSessionTracker(), - - orchBroadcaster: orchBroadcaster, - - orchReporter: orchReporter, - - done: make(chan struct{}), - cmdRegistry: commands.NewRegistry(commands.BuiltinDefinitions()), } @@ -222,16 +200,6 @@ func NewAgentLoop( return al } -func (al *AgentLoop) SetConfigSaver(fn func(*config.Config) error) { - al.saveConfig = fn -} - -// SetHeartbeatThreadUpdater registers a callback to apply runtime heartbeat thread updates. - -func (al *AgentLoop) SetHeartbeatThreadUpdater(fn func(int)) { - al.onHeartbeatThreadUpdate = fn -} - // registerSharedTools registers tools that are shared across all agents (web, message, spawn). func registerSharedTools( cfg *config.Config, @@ -616,6 +584,16 @@ func (al *AgentLoop) llmWorker(ctx context.Context, queue <-chan bus.InboundMess return } + // Reset per-round message-tool state so a previous round's + // tool-sent flag does not suppress this round's response. + if defaultAgent := al.registry.GetDefaultAgent(); defaultAgent != nil { + if tool, ok := defaultAgent.Tools.Get("message"); ok { + if mt, ok := tool.(*tools.MessageTool); ok { + mt.ResetSentInRound() + } + } + } + response, err := al.processMessage(ctx, msg) if err != nil { response = fmt.Sprintf("Error processing message: %v", err) diff --git a/pkg/agent/loop_ext.go b/pkg/agent/loop_ext.go new file mode 100644 index 000000000..67d99a89b --- /dev/null +++ b/pkg/agent/loop_ext.go @@ -0,0 +1,46 @@ +package agent + +import ( + "sync" + + "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/orch" + "github.com/sipeed/picoclaw/pkg/stats" +) + +// loopExt holds fork-specific fields for AgentLoop. +// Embedded in AgentLoop so existing field access (al.stats, al.sessions, etc.) continues to work. +// Upstream additions to AgentLoop won't conflict with these fields. +type loopExt struct { + stats *stats.Tracker // nil when --stats not passed + + sessions *SessionTracker + + orchBroadcaster *orch.Broadcaster // nil when --orchestration not set + + orchReporter orch.AgentReporter // always non-nil (Noop when disabled) + + planStartPending bool // set by /plan start to trigger LLM execution + + planClearHistory bool // set by /plan start clear to wipe history on transition + + sessionLocks sync.Map // sessionKey → *sessionSemaphore + + activeTasks sync.Map // sessionKey → *activeTask + + done chan struct{} // closed by Close() to stop background goroutines + + saveConfig func(*config.Config) error + + onHeartbeatThreadUpdate func(int) +} + +// SetConfigSaver registers a callback to persist config changes. +func (al *AgentLoop) SetConfigSaver(fn func(*config.Config) error) { + al.saveConfig = fn +} + +// SetHeartbeatThreadUpdater registers a callback to apply runtime heartbeat thread updates. +func (al *AgentLoop) SetHeartbeatThreadUpdater(fn func(int)) { + al.onHeartbeatThreadUpdate = fn +} diff --git a/pkg/channels/manager.go b/pkg/channels/manager.go index 631f0aa1c..8e663ff86 100644 --- a/pkg/channels/manager.go +++ b/pkg/channels/manager.go @@ -10,9 +10,7 @@ import ( "context" "errors" "fmt" - "hash/fnv" "math" - "strings" "sync" "time" @@ -37,14 +35,6 @@ const ( janitorInterval = 10 * time.Second typingStopTTL = 5 * time.Minute placeholderTTL = 10 * time.Minute - statusMsgTTL = 5 * time.Minute - taskMsgTTL = 30 * time.Minute - - // statusEditInterval is the minimum interval between EditMessage calls - // for the same status/task bubble. EditMessage APIs are more rate-sensitive - // than SendMessageDraft, so we throttle edits to avoid "(edited)" flicker - // and API rate limit errors. Draft-based channels bypass this throttle. - statusEditInterval = 500 * time.Millisecond ) // typingEntry wraps a typing stop function with a creation timestamp for TTL eviction. @@ -65,13 +55,6 @@ type placeholderEntry struct { createdAt time.Time } -// statusMsgEntry tracks a status or task message ID for later editing. -type statusMsgEntry struct { - messageID string - draftID int // non-zero when using draft-based streaming - createdAt time.Time -} - // channelRateConfig maps channel name to per-second rate limit. var channelRateConfig = map[string]float64{ "telegram": 20, @@ -93,19 +76,18 @@ type channelWorker struct { } type Manager struct { - channels map[string]Channel - workers map[string]*channelWorker - bus *bus.MessageBus - config *config.Config - mediaStore media.MediaStore - dispatchTask *asyncTask - mu sync.RWMutex - placeholders sync.Map // "channel:chatID" → placeholderEntry - typingStops sync.Map // "channel:chatID" → typingEntry - reactionUndos sync.Map // "channel:chatID" → reactionEntry - statusMsgIDs sync.Map // "channel:chatID" → statusMsgEntry (streaming preview) - taskMsgIDs sync.Map // "channel:chatID:taskID" → statusMsgEntry (background task status) - statusEditTimes sync.Map // key → time.Time — last EditMessage time for throttling + managerExt // fork-specific fields (see manager_ext.go) + + channels map[string]Channel + workers map[string]*channelWorker + bus *bus.MessageBus + config *config.Config + mediaStore media.MediaStore + dispatchTask *asyncTask + mu sync.RWMutex + placeholders sync.Map // "channel:chatID" → placeholderEntry + typingStops sync.Map // "channel:chatID" → typingEntry + reactionUndos sync.Map // "channel:chatID" → reactionEntry } type asyncTask struct { @@ -551,235 +533,6 @@ func (m *Manager) runWorker(ctx context.Context, name string, w *channelWorker) } } -// handleStatusSend processes IsStatus messages (streaming previews). -// It reuses an existing placeholder or tracked status message, or sends a new -// one via SendWithID so subsequent status updates edit the same bubble. -// For channels implementing DraftSender (e.g. Telegram private chats), -// sendMessageDraft is preferred as it avoids the "(edited)" indicator. -// If the channel doesn't support editing, the message is silently dropped. -func (m *Manager) handleStatusSend(ctx context.Context, name string, w *channelWorker, msg bus.OutboundMessage) { - if err := w.limiter.Wait(ctx); err != nil { - return - } - - key := name + ":" + msg.ChatID - - // 0. Draft-based streaming (preferred for supported channels) - if drafter, ok := w.ch.(DraftSender); ok { - var did int - if v, loaded := m.statusMsgIDs.Load(key); loaded { - if entry, ok := v.(statusMsgEntry); ok && entry.draftID != 0 { - did = entry.draftID - } - } - if did == 0 { - did = generateDraftID(key) - } - if err := drafter.SendDraft(ctx, msg.ChatID, did, msg.Content); err == nil { - // Track draft only after successful send. If draft fails (e.g. group - // main thread), keep existing messageID entry so fallback edits can - // reuse the same status bubble instead of creating duplicates. - m.statusMsgIDs.Store(key, statusMsgEntry{ - draftID: did, - createdAt: time.Now(), - }) - return - } - // Draft failed — fall through to edit-based approach - } - - // Edit-based path: throttle to statusEditInterval per key to avoid - // API rate limit errors and "(edited)" flicker. - if v, loaded := m.statusEditTimes.Load(key); loaded { - if t, ok := v.(time.Time); ok && time.Since(t) < statusEditInterval { - return // too recent, skip this update - } - } - - // 1. Try editing an existing placeholder - if v, loaded := m.placeholders.Load(key); loaded { - if entry, ok := v.(placeholderEntry); ok && entry.id != "" { - if editor, ok := w.ch.(MessageEditor); ok { - if err := editor.EditMessage(ctx, msg.ChatID, entry.id, msg.Content); err == nil { - m.statusEditTimes.Store(key, time.Now()) - return - } - } - } - } - - // 2. Try editing a previously tracked status message - if v, loaded := m.statusMsgIDs.Load(key); loaded { - if entry, ok := v.(statusMsgEntry); ok && entry.messageID != "" { - if editor, ok := w.ch.(MessageEditor); ok { - if err := editor.EditMessage(ctx, msg.ChatID, entry.messageID, msg.Content); err == nil { - m.statusEditTimes.Store(key, time.Now()) - return - } - } - } - } - - // 3. Send new message via SendWithID and track it - if sender, ok := w.ch.(MessageSenderWithID); ok { - if msgID, err := sender.SendWithID(ctx, msg.ChatID, msg.Content); err == nil && msgID != "" { - m.statusMsgIDs.Store(key, statusMsgEntry{ - messageID: msgID, - createdAt: time.Now(), - }) - return - } - } - - // 4. Channel doesn't support SendWithID or editing — drop silently -} - -func taskStatusKey(channel, chatID, taskID string) string { - if taskID == "" { - return "" - } - if channel == "" || chatID == "" { - return taskID - } - return channel + ":" + chatID + ":" + taskID -} - -// handleTaskStatusSend processes IsTaskStatus messages (background task status). -// It reuses a previously tracked task message, or sends a new one via SendWithID. -// For channels implementing DraftSender, sendMessageDraft is used to avoid "(edited)". -// If the channel doesn't support editing, falls back to regular Send. -func (m *Manager) handleTaskStatusSend(ctx context.Context, name string, w *channelWorker, msg bus.OutboundMessage) { - if err := w.limiter.Wait(ctx); err != nil { - return - } - - taskKey := taskStatusKey(name, msg.ChatID, msg.TaskID) - - // Final message: reuse the existing bubble when possible to avoid - // duplicate messages. If a permanent message (messageID) is tracked, - // edit it in-place. If a draft (draftID) is tracked, update it with - // the completion content (the draft persists in Telegram and serves - // as the visible message; sending a separate permanent message would - // create a duplicate). - if msg.Final { - v, loaded := m.taskMsgIDs.LoadAndDelete(taskKey) - m.statusEditTimes.Delete(taskKey) - - if loaded { - if entry, ok := v.(statusMsgEntry); ok { - // Path A: a permanent message exists — edit it in-place. - if entry.messageID != "" { - if editor, ok := w.ch.(MessageEditor); ok { - if err := editor.EditMessage(ctx, msg.ChatID, entry.messageID, msg.Content); err == nil { - return - } - } - // Edit failed — fall through to send a new message. - } - - // Path B: a draft exists — update it with the final - // content. Drafts persist visibly in Telegram, so do NOT - // send a separate permanent message (that causes duplicates). - if entry.draftID != 0 { - if drafter, ok := w.ch.(DraftSender); ok { - if err := drafter.SendDraft(ctx, msg.ChatID, entry.draftID, msg.Content); err == nil { - return - } - } - // Draft update failed — fall through to send permanent. - } - } - } - - // No existing bubble to reuse — send a new permanent message. - if sender, ok := w.ch.(MessageSenderWithID); ok { - if msgID, err := sender.SendWithID(ctx, msg.ChatID, msg.Content); err == nil && msgID != "" { - return - } - } - _ = w.ch.Send(ctx, msg) - return - } - - // 0. Draft-based streaming (preferred for supported channels) - if drafter, ok := w.ch.(DraftSender); ok && taskKey != "" { - var did int - if v, loaded := m.taskMsgIDs.Load(taskKey); loaded { - if entry, ok := v.(statusMsgEntry); ok && entry.draftID != 0 { - did = entry.draftID - } - } - if did == 0 { - did = generateDraftID(taskKey) - } - if err := drafter.SendDraft(ctx, msg.ChatID, did, msg.Content); err == nil { - // Track draft only after successful send to avoid clobbering an - // existing messageID entry when drafts are unsupported. - m.taskMsgIDs.Store(taskKey, statusMsgEntry{ - draftID: did, - createdAt: time.Now(), - }) - return - } - // Draft failed — fall through to edit-based approach - } - - // Edit-based path: throttle to statusEditInterval per task key. - if taskKey != "" { - if v, loaded := m.statusEditTimes.Load(taskKey); loaded { - if t, ok := v.(time.Time); ok && time.Since(t) < statusEditInterval { - return - } - } - } - - // 1. Try editing an existing task message - if taskKey != "" { - if v, loaded := m.taskMsgIDs.Load(taskKey); loaded { - if entry, ok := v.(statusMsgEntry); ok && entry.messageID != "" { - if editor, ok := w.ch.(MessageEditor); ok { - if err := editor.EditMessage(ctx, msg.ChatID, entry.messageID, msg.Content); err == nil { - m.statusEditTimes.Store(taskKey, time.Now()) - return - } - } - } - } - } - - // 2. Send new message via SendWithID and track it - if sender, ok := w.ch.(MessageSenderWithID); ok { - if msgID, err := sender.SendWithID(ctx, msg.ChatID, msg.Content); err == nil && msgID != "" { - if taskKey != "" { - m.taskMsgIDs.Store(taskKey, statusMsgEntry{ - messageID: msgID, - createdAt: time.Now(), - }) - } - return - } - } - - // 3. Fallback: regular Send (for channels without SendWithID) - _ = w.ch.Send(ctx, msg) -} - -// generateDraftID produces a stable non-zero int from a key string. -// The same key always maps to the same draft ID so successive calls -// animate the same Telegram draft bubble. -func generateDraftID(key string) int { - h := fnv.New32a() - h.Write([]byte(key)) - v := int(h.Sum32()) - if v == 0 { - v = 1 // draftID must be non-zero - } - if v < 0 { - v = -v - } - return v -} - // sendWithRetry sends a message through the channel with rate limiting and // retry logic. It classifies errors to determine the retry strategy: // - ErrNotRunning / ErrSendFailed: permanent, no retry @@ -1073,26 +826,6 @@ func (m *Manager) runTTLJanitor(ctx context.Context) { } } -// PromoteStatusToTask moves the tracked streaming status message for the given -// channel:chatID key into the task message map under channel:chatID:taskID. This allows the -// next IsTaskStatus publish to edit the streaming bubble instead of creating a -// new message. Returns true if a status message was found and promoted. -func (m *Manager) PromoteStatusToTask(statusKey, taskID string) bool { - v, loaded := m.statusMsgIDs.LoadAndDelete(statusKey) - if !loaded { - return false - } - - parts := strings.SplitN(statusKey, ":", 2) - if len(parts) == 2 { - m.taskMsgIDs.Store(taskStatusKey(parts[0], parts[1], taskID), v) - return true - } - - m.taskMsgIDs.Store(taskID, v) - return true -} - func (m *Manager) GetChannel(name string) (Channel, bool) { m.mu.RLock() defer m.mu.RUnlock() diff --git a/pkg/channels/manager_ext.go b/pkg/channels/manager_ext.go new file mode 100644 index 000000000..0c84cd07f --- /dev/null +++ b/pkg/channels/manager_ext.go @@ -0,0 +1,286 @@ +package channels + +import ( + "context" + "hash/fnv" + "strings" + "sync" + "time" + + "github.com/sipeed/picoclaw/pkg/bus" +) + +const ( + statusMsgTTL = 5 * time.Minute + taskMsgTTL = 30 * time.Minute + + // statusEditInterval is the minimum interval between EditMessage calls + // for the same status/task bubble. EditMessage APIs are more rate-sensitive + // than SendMessageDraft, so we throttle edits to avoid "(edited)" flicker + // and API rate limit errors. Draft-based channels bypass this throttle. + statusEditInterval = 500 * time.Millisecond +) + +// statusMsgEntry tracks a status or task message ID for later editing. +type statusMsgEntry struct { + messageID string + draftID int // non-zero when using draft-based streaming + createdAt time.Time +} + +// managerExt holds fork-specific fields for Manager. +// Embedded in Manager so existing field access continues to work. +type managerExt struct { + statusMsgIDs sync.Map // "channel:chatID" → statusMsgEntry (streaming preview) + taskMsgIDs sync.Map // "channel:chatID:taskID" → statusMsgEntry (background task status) + statusEditTimes sync.Map // key → time.Time — last EditMessage time for throttling +} + +// handleStatusSend processes IsStatus messages (streaming previews). +// It reuses an existing placeholder or tracked status message, or sends a new +// one via SendWithID so subsequent status updates edit the same bubble. +// For channels implementing DraftSender (e.g. Telegram private chats), +// sendMessageDraft is preferred as it avoids the "(edited)" indicator. +// If the channel doesn't support editing, the message is silently dropped. +func (m *Manager) handleStatusSend(ctx context.Context, name string, w *channelWorker, msg bus.OutboundMessage) { + if err := w.limiter.Wait(ctx); err != nil { + return + } + + key := name + ":" + msg.ChatID + + // 0. Draft-based streaming (preferred for supported channels) + if drafter, ok := w.ch.(DraftSender); ok { + var did int + if v, loaded := m.statusMsgIDs.Load(key); loaded { + if entry, ok := v.(statusMsgEntry); ok && entry.draftID != 0 { + did = entry.draftID + } + } + if did == 0 { + did = generateDraftID(key) + } + if err := drafter.SendDraft(ctx, msg.ChatID, did, msg.Content); err == nil { + // Track draft only after successful send. If draft fails (e.g. group + // main thread), keep existing messageID entry so fallback edits can + // reuse the same status bubble instead of creating duplicates. + m.statusMsgIDs.Store(key, statusMsgEntry{ + draftID: did, + createdAt: time.Now(), + }) + return + } + // Draft failed — fall through to edit-based approach + } + + // Edit-based path: throttle to statusEditInterval per key to avoid + // API rate limit errors and "(edited)" flicker. + if v, loaded := m.statusEditTimes.Load(key); loaded { + if t, ok := v.(time.Time); ok && time.Since(t) < statusEditInterval { + return // too recent, skip this update + } + } + + // 1. Try editing an existing placeholder + if v, loaded := m.placeholders.Load(key); loaded { + if entry, ok := v.(placeholderEntry); ok && entry.id != "" { + if editor, ok := w.ch.(MessageEditor); ok { + if err := editor.EditMessage(ctx, msg.ChatID, entry.id, msg.Content); err == nil { + m.statusEditTimes.Store(key, time.Now()) + return + } + } + } + } + + // 2. Try editing a previously tracked status message + if v, loaded := m.statusMsgIDs.Load(key); loaded { + if entry, ok := v.(statusMsgEntry); ok && entry.messageID != "" { + if editor, ok := w.ch.(MessageEditor); ok { + if err := editor.EditMessage(ctx, msg.ChatID, entry.messageID, msg.Content); err == nil { + m.statusEditTimes.Store(key, time.Now()) + return + } + } + } + } + + // 3. Send new message via SendWithID and track it + if sender, ok := w.ch.(MessageSenderWithID); ok { + if msgID, err := sender.SendWithID(ctx, msg.ChatID, msg.Content); err == nil && msgID != "" { + m.statusMsgIDs.Store(key, statusMsgEntry{ + messageID: msgID, + createdAt: time.Now(), + }) + return + } + } + + // 4. Channel doesn't support SendWithID or editing — drop silently +} + +func taskStatusKey(channel, chatID, taskID string) string { + if taskID == "" { + return "" + } + if channel == "" || chatID == "" { + return taskID + } + return channel + ":" + chatID + ":" + taskID +} + +// handleTaskStatusSend processes IsTaskStatus messages (background task status). +// It reuses a previously tracked task message, or sends a new one via SendWithID. +// For channels implementing DraftSender, sendMessageDraft is used to avoid "(edited)". +// If the channel doesn't support editing, falls back to regular Send. +func (m *Manager) handleTaskStatusSend(ctx context.Context, name string, w *channelWorker, msg bus.OutboundMessage) { + if err := w.limiter.Wait(ctx); err != nil { + return + } + + taskKey := taskStatusKey(name, msg.ChatID, msg.TaskID) + + // Final message: reuse the existing bubble when possible to avoid + // duplicate messages. If a permanent message (messageID) is tracked, + // edit it in-place. If a draft (draftID) is tracked, update it with + // the completion content (the draft persists in Telegram and serves + // as the visible message; sending a separate permanent message would + // create a duplicate). + if msg.Final { + v, loaded := m.taskMsgIDs.LoadAndDelete(taskKey) + m.statusEditTimes.Delete(taskKey) + + if loaded { + if entry, ok := v.(statusMsgEntry); ok { + // Path A: a permanent message exists — edit it in-place. + if entry.messageID != "" { + if editor, ok := w.ch.(MessageEditor); ok { + if err := editor.EditMessage(ctx, msg.ChatID, entry.messageID, msg.Content); err == nil { + return + } + } + // Edit failed — fall through to send a new message. + } + + // Path B: a draft exists — update it with the final + // content. Drafts persist visibly in Telegram, so do NOT + // send a separate permanent message (that causes duplicates). + if entry.draftID != 0 { + if drafter, ok := w.ch.(DraftSender); ok { + if err := drafter.SendDraft(ctx, msg.ChatID, entry.draftID, msg.Content); err == nil { + return + } + } + // Draft update failed — fall through to send permanent. + } + } + } + + // No existing bubble to reuse — send a new permanent message. + if sender, ok := w.ch.(MessageSenderWithID); ok { + if msgID, err := sender.SendWithID(ctx, msg.ChatID, msg.Content); err == nil && msgID != "" { + return + } + } + _ = w.ch.Send(ctx, msg) + return + } + + // 0. Draft-based streaming (preferred for supported channels) + if drafter, ok := w.ch.(DraftSender); ok && taskKey != "" { + var did int + if v, loaded := m.taskMsgIDs.Load(taskKey); loaded { + if entry, ok := v.(statusMsgEntry); ok && entry.draftID != 0 { + did = entry.draftID + } + } + if did == 0 { + did = generateDraftID(taskKey) + } + if err := drafter.SendDraft(ctx, msg.ChatID, did, msg.Content); err == nil { + // Track draft only after successful send to avoid clobbering an + // existing messageID entry when drafts are unsupported. + m.taskMsgIDs.Store(taskKey, statusMsgEntry{ + draftID: did, + createdAt: time.Now(), + }) + return + } + // Draft failed — fall through to edit-based approach + } + + // Edit-based path: throttle to statusEditInterval per task key. + if taskKey != "" { + if v, loaded := m.statusEditTimes.Load(taskKey); loaded { + if t, ok := v.(time.Time); ok && time.Since(t) < statusEditInterval { + return + } + } + } + + // 1. Try editing an existing task message + if taskKey != "" { + if v, loaded := m.taskMsgIDs.Load(taskKey); loaded { + if entry, ok := v.(statusMsgEntry); ok && entry.messageID != "" { + if editor, ok := w.ch.(MessageEditor); ok { + if err := editor.EditMessage(ctx, msg.ChatID, entry.messageID, msg.Content); err == nil { + m.statusEditTimes.Store(taskKey, time.Now()) + return + } + } + } + } + } + + // 2. Send new message via SendWithID and track it + if sender, ok := w.ch.(MessageSenderWithID); ok { + if msgID, err := sender.SendWithID(ctx, msg.ChatID, msg.Content); err == nil && msgID != "" { + if taskKey != "" { + m.taskMsgIDs.Store(taskKey, statusMsgEntry{ + messageID: msgID, + createdAt: time.Now(), + }) + } + return + } + } + + // 3. Fallback: regular Send (for channels without SendWithID) + _ = w.ch.Send(ctx, msg) +} + +// generateDraftID produces a stable non-zero int from a key string. +// The same key always maps to the same draft ID so successive calls +// animate the same Telegram draft bubble. +func generateDraftID(key string) int { + h := fnv.New32a() + h.Write([]byte(key)) + v := int(h.Sum32()) + if v == 0 { + v = 1 // draftID must be non-zero + } + if v < 0 { + v = -v + } + return v +} + +// PromoteStatusToTask moves the tracked streaming status message for the given +// channel:chatID key into the task message map under channel:chatID:taskID. This allows the +// next IsTaskStatus publish to edit the streaming bubble instead of creating a +// new message. Returns true if a status message was found and promoted. +func (m *Manager) PromoteStatusToTask(statusKey, taskID string) bool { + v, loaded := m.statusMsgIDs.LoadAndDelete(statusKey) + if !loaded { + return false + } + + parts := strings.SplitN(statusKey, ":", 2) + if len(parts) == 2 { + m.taskMsgIDs.Store(taskStatusKey(parts[0], parts[1], taskID), v) + return true + } + + m.taskMsgIDs.Store(taskID, v) + return true +} diff --git a/pkg/tools/message_test.go b/pkg/tools/message_test.go index 05630972e..abd34448a 100644 --- a/pkg/tools/message_test.go +++ b/pkg/tools/message_test.go @@ -252,3 +252,29 @@ func TestMessageTool_Parameters(t *testing.T) { t.Error("Expected chat_id type to be 'string'") } } + +func TestMessageTool_ResetSentInRound(t *testing.T) { + tool := NewMessageTool() + tool.SetSendCallback(func(channel, chatID, content string) error { + return nil + }) + + ctx := WithToolContext(context.Background(), "ch", "cid") + + // First round: tool sends a message + tool.Execute(ctx, map[string]any{"content": "hello"}) + if !tool.HasSentInRound() { + t.Fatal("expected sentInRound=true after Execute") + } + + // Reset for second round + tool.ResetSentInRound() + if tool.HasSentInRound() { + t.Fatal("expected sentInRound=false after ResetSentInRound") + } + + // Second round: tool is NOT used (direct answer) → flag stays false + if tool.HasSentInRound() { + t.Error("expected sentInRound=false when tool was not used in this round") + } +} diff --git a/pkg/tools/shell.go b/pkg/tools/shell.go index e2c8a5c11..16bb82b64 100644 --- a/pkg/tools/shell.go +++ b/pkg/tools/shell.go @@ -5,191 +5,32 @@ import ( "context" "errors" "fmt" - "io" - "net" - "net/url" "os" "os/exec" "path/filepath" "regexp" "runtime" - "sort" "strings" - "sync" "time" "github.com/sipeed/picoclaw/pkg/config" "github.com/sipeed/picoclaw/pkg/constants" ) -const ( - bgMaxLifetime = 45 * time.Minute - - bgRingBufSize = 32 * 1024 // 32KB - - bgInitCapture = 3 * time.Second - - bgMaxProcesses = 10 -) - -// ringBuffer is a thread-safe circular buffer that retains the most recent bytes. - -type ringBuffer struct { - mu sync.Mutex - - buf []byte - - size int -} - -func newRingBuffer(size int) *ringBuffer { - return &ringBuffer{size: size} -} - -// Write appends data to the ring buffer, dropping oldest bytes if capacity is exceeded. - -func (rb *ringBuffer) Write(p []byte) (int, error) { - rb.mu.Lock() - - defer rb.mu.Unlock() - - rb.buf = append(rb.buf, p...) - - if len(rb.buf) > rb.size { - rb.buf = rb.buf[len(rb.buf)-rb.size:] - } - - return len(p), nil -} - -// String returns the current buffer contents. - -func (rb *ringBuffer) String() string { - rb.mu.Lock() - - defer rb.mu.Unlock() - - return string(rb.buf) -} - -// Lines returns the last n lines from the buffer. - -func (rb *ringBuffer) Lines(n int) []string { - rb.mu.Lock() - - defer rb.mu.Unlock() - - if len(rb.buf) == 0 { - return nil - } - - all := strings.Split(string(rb.buf), "\n") - - // Remove trailing empty element from final newline - - if len(all) > 0 && all[len(all)-1] == "" { - all = all[:len(all)-1] - } - - if n <= 0 || n >= len(all) { - return all - } - - return all[len(all)-n:] -} - -// Match checks if any line in the buffer matches the given regex pattern. - -// Returns the first matching line, or empty string if no match. - -func (rb *ringBuffer) Match(pattern *regexp.Regexp) string { - rb.mu.Lock() - - defer rb.mu.Unlock() - - for _, line := range strings.Split(string(rb.buf), "\n") { - if pattern.MatchString(line) { - return line - } - } - - return "" -} - -// Len returns the current number of bytes in the buffer. - -func (rb *ringBuffer) Len() int { - rb.mu.Lock() - - defer rb.mu.Unlock() - - return len(rb.buf) -} - -// bgProcess represents a background process managed by ExecTool. - -type bgProcess struct { - id string - - command string - - cmd *exec.Cmd - - pid int - - startedAt time.Time - - output *ringBuffer - - done chan struct{} // closed when process exits - - exitErr error - - cancel context.CancelFunc // cancels the monitor goroutine -} - -// isRunning returns true if the process has not yet exited. - -func (bp *bgProcess) isRunning() bool { - select { - case <-bp.done: - - return false - - default: - - return true - } -} - type ExecTool struct { + execToolExt // fork-specific fields (see shell_ext.go) + workingDir string timeout time.Duration denyPatterns []*regexp.Regexp - allowRules [][]string // pre-split command prefix allowlist - customAllowPatterns []*regexp.Regexp restrictToWorkspace bool - localNetOnly bool // restrict curl/wget to localhost + RFC 1918 - allowRemote bool - - // Background process management - - bgMu sync.Mutex - - bgProcesses map[string]*bgProcess - - bgNextID int - - bgShutdown context.CancelFunc // cancels all bg monitor goroutines - - bgCtx context.Context } var ( @@ -351,25 +192,23 @@ func NewExecToolWithConfig(workingDir string, restrict bool, config *config.Conf bgCtx, bgCancel := context.WithCancel(context.Background()) return &ExecTool{ + execToolExt: execToolExt{ + bgProcesses: make(map[string]*bgProcess), + bgCtx: bgCtx, + bgShutdown: bgCancel, + }, + workingDir: workingDir, timeout: timeout, denyPatterns: denyPatterns, - allowRules: nil, - customAllowPatterns: customAllowPatterns, restrictToWorkspace: restrict, allowRemote: allowRemote, - - bgProcesses: make(map[string]*bgProcess), - - bgCtx: bgCtx, - - bgShutdown: bgCancel, }, nil } @@ -613,403 +452,6 @@ func (t *ExecTool) executeSync(ctx context.Context, command, cwd string) *ToolRe } } -// executeBg starts a background process and returns immediately. - -func (t *ExecTool) executeBg(command, cwd string) *ToolResult { - t.bgMu.Lock() - - // Check max processes limit - - running := 0 - - for _, bp := range t.bgProcesses { - if bp.isRunning() { - running++ - } - } - - if running >= bgMaxProcesses { - t.bgMu.Unlock() - - return ErrorResult( - - fmt.Sprintf("maximum background processes reached (%d). Kill an existing one first.", bgMaxProcesses), - ) - } - - t.bgNextID++ - - id := fmt.Sprintf("bg-%d", t.bgNextID) - - t.bgMu.Unlock() - - var cmd *exec.Cmd - - if runtime.GOOS == "windows" { - cmd = exec.Command("powershell", "-NoProfile", "-NonInteractive", "-Command", command) - } else { - cmd = exec.Command("sh", "-c", command) - } - - if cwd != "" { - cmd.Dir = cwd - } - - prepareCommandForTermination(cmd) - - output := newRingBuffer(bgRingBufSize) - - // Use pipes to capture output - - stdoutPipe, err := cmd.StdoutPipe() - if err != nil { - return ErrorResult(fmt.Sprintf("failed to create stdout pipe: %v", err)) - } - - stderrPipe, err := cmd.StderrPipe() - if err != nil { - return ErrorResult(fmt.Sprintf("failed to create stderr pipe: %v", err)) - } - - if err := cmd.Start(); err != nil { - return ErrorResult(fmt.Sprintf("failed to start background command: %v", err)) - } - - monitorCtx, monitorCancel := context.WithCancel(t.bgCtx) - - bp := &bgProcess{ - id: id, - - command: command, - - cmd: cmd, - - pid: cmd.Process.Pid, - - startedAt: time.Now(), - - output: output, - - done: make(chan struct{}), - - cancel: monitorCancel, - } - - t.bgMu.Lock() - - t.bgProcesses[id] = bp - - t.bgMu.Unlock() - - // io.Copy goroutines: pipe stdout/stderr into ring buffer - - go io.Copy(output, stdoutPipe) - - go io.Copy(output, stderrPipe) - - // cmd.Wait goroutine - - waitDone := make(chan error, 1) - - go func() { - waitDone <- cmd.Wait() - }() - - // Monitor goroutine: handles lifetime timer, process exit, and shutdown - - go func() { - lifetime := time.NewTimer(getBgMaxLifetime()) - - defer lifetime.Stop() - - select { - case err := <-waitDone: - - // Process exited naturally - - bp.exitErr = err - - close(bp.done) - - case <-lifetime.C: - - // Max lifetime exceeded — kill - - _ = terminateProcessTree(cmd) - - select { - case err := <-waitDone: - - bp.exitErr = err - - case <-time.After(2 * time.Second): - - if cmd.Process != nil { - _ = cmd.Process.Kill() - } - - bp.exitErr = <-waitDone - } - - close(bp.done) - - case <-monitorCtx.Done(): - - // Shutdown or explicit kill via cancel - - _ = terminateProcessTree(cmd) - - select { - case err := <-waitDone: - - bp.exitErr = err - - case <-time.After(2 * time.Second): - - if cmd.Process != nil { - _ = cmd.Process.Kill() - } - - bp.exitErr = <-waitDone - } - - select { - case <-bp.done: - - default: - - close(bp.done) - } - } - }() - - // Capture initial output (wait up to bgInitCapture) - - time.Sleep(bgInitCapture) - - initialOutput := output.String() - - var sb strings.Builder - - fmt.Fprintf(&sb, "Background process started.\n") - - fmt.Fprintf(&sb, " id: %s\n", id) - - fmt.Fprintf(&sb, " pid: %d\n", bp.pid) - - fmt.Fprintf(&sb, " cmd: %s\n", command) - - fmt.Fprintf(&sb, " max lifetime: %s\n", getBgMaxLifetime()) - - if initialOutput != "" { - fmt.Fprintf(&sb, "\nInitial output:\n%s", initialOutput) - } - - return &ToolResult{ - ForLLM: sb.String(), - - ForUser: fmt.Sprintf("Background process %s (pid=%d) started: %s", id, bp.pid, command), - } -} - -// handleBgAction handles bg_action=output and bg_action=kill. - -func (t *ExecTool) handleBgAction(action, bgID string) *ToolResult { - if bgID == "" { - return ErrorResult("bg_id is required for bg_action") - } - - t.bgMu.Lock() - - bp, ok := t.bgProcesses[bgID] - - t.bgMu.Unlock() - - if !ok { - return ErrorResult(fmt.Sprintf("background process %q not found", bgID)) - } - - switch action { - case "output": - - return t.bgOutput(bp) - - case "kill": - - return t.bgKill(bp) - - default: - - return ErrorResult(fmt.Sprintf("unknown bg_action %q (use 'output' or 'kill')", action)) - } -} - -func (t *ExecTool) bgOutput(bp *bgProcess) *ToolResult { - var sb strings.Builder - - fmt.Fprintf(&sb, "[%s] pid=%d %s\n", bp.id, bp.pid, bp.command) - - if bp.isRunning() { - uptime := time.Since(bp.startedAt).Truncate(time.Second) - - fmt.Fprintf(&sb, "Status: running (uptime: %s, max: %s)\n", uptime, getBgMaxLifetime()) - } else { - ran := time.Since(bp.startedAt).Truncate(time.Second) - - if bp.exitErr != nil { - fmt.Fprintf(&sb, "Status: exited with error (ran: %s): %v\n", ran, bp.exitErr) - } else { - fmt.Fprintf(&sb, "Status: exited=0 (ran: %s)\n", ran) - } - } - - output := bp.output.String() - - if output == "" { - fmt.Fprintf(&sb, "\n(no output)") - } else { - fmt.Fprintf(&sb, "\nOutput:\n%s", output) - } - - return &ToolResult{ - ForLLM: sb.String(), - - ForUser: sb.String(), - } -} - -func (t *ExecTool) bgKill(bp *bgProcess) *ToolResult { - if bp.isRunning() { - bp.cancel() // triggers monitor goroutine cleanup - - // Wait for process to actually exit - - select { - case <-bp.done: - - case <-time.After(5 * time.Second): - } - } - - t.bgMu.Lock() - - delete(t.bgProcesses, bp.id) - - t.bgMu.Unlock() - - msg := fmt.Sprintf("Background process %s (pid=%d) terminated: %s", bp.id, bp.pid, bp.command) - - return &ToolResult{ - ForLLM: msg, - - ForUser: msg, - } -} - -// BgProcesses returns a snapshot of background processes for use by bg_monitor. - -func (t *ExecTool) BgProcesses() map[string]*bgProcess { - t.bgMu.Lock() - - defer t.bgMu.Unlock() - - snapshot := make(map[string]*bgProcess, len(t.bgProcesses)) - - for k, v := range t.bgProcesses { - snapshot[k] = v - } - - return snapshot -} - -// RuntimeStatus implements StatusProvider for system prompt injection. - -func (t *ExecTool) RuntimeStatus() string { - t.bgMu.Lock() - - defer t.bgMu.Unlock() - - if len(t.bgProcesses) == 0 { - return "" - } - - // Sort by ID for stable output - - ids := make([]string, 0, len(t.bgProcesses)) - - for id := range t.bgProcesses { - ids = append(ids, id) - } - - sort.Strings(ids) - - var sb strings.Builder - - sb.WriteString("## Background Processes\n\n") - - for _, id := range ids { - bp := t.bgProcesses[id] - - if bp.isRunning() { - uptime := time.Since(bp.startedAt).Truncate(time.Second) - - fmt.Fprintf(&sb, " [%s] pid=%d running (uptime: %s, max: %s) %s\n", - - id, bp.pid, uptime, getBgMaxLifetime(), bp.command) - } else { - ran := time.Since(bp.startedAt).Truncate(time.Second) - - if bp.exitErr != nil { - fmt.Fprintf(&sb, " [%s] pid=%d exited=err (ran: %s) %s\n", - - id, bp.pid, ran, bp.command) - } else { - fmt.Fprintf(&sb, " [%s] pid=%d exited=0 (ran: %s) %s\n", - - id, bp.pid, ran, bp.command) - } - } - } - - sb.WriteString("\nUse exec with bg_action=\"output\" / \"kill\" and bg_id to manage.\n") - - sb.WriteString("Use bg_monitor for list/watch/tail operations.") - - return sb.String() -} - -// Shutdown terminates all background processes. Call on application exit. - -func (t *ExecTool) Shutdown() { - t.bgShutdown() // cancel all monitor goroutines - - t.bgMu.Lock() - - procs := make([]*bgProcess, 0, len(t.bgProcesses)) - - for _, bp := range t.bgProcesses { - procs = append(procs, bp) - } - - t.bgMu.Unlock() - - // Wait for all processes to exit - - for _, bp := range procs { - select { - case <-bp.done: - - case <-time.After(5 * time.Second): - - // Force kill if still running - - if bp.cmd.Process != nil { - _ = bp.cmd.Process.Kill() - } - } - } -} - func (t *ExecTool) guardCommand(command, cwd string) string { cmd := strings.TrimSpace(command) lower := strings.ToLower(cmd) @@ -1256,97 +698,3 @@ func matchAllowRules(cmd string, rules [][]string) bool { return false } - -func (t *ExecTool) SetLocalNetOnly(v bool) { - t.localNetOnly = v -} - -// isCurlOrWget reports whether command is a curl or wget invocation. - -func isCurlOrWget(command string) bool { - fields := strings.Fields(command) - - if len(fields) == 0 { - return false - } - - base := filepath.Base(fields[0]) - - return base == "curl" || base == "wget" -} - -// checkCurlLocalNet validates that all http/https URLs in a curl/wget command - -// target localhost or RFC 1918 private addresses. - -// Returns an error message string, or empty string if the command is allowed. - -func checkCurlLocalNet(command string) string { - for _, token := range strings.Fields(command) { - token = strings.Trim(token, "\"'") - - if !strings.HasPrefix(token, "http://") && !strings.HasPrefix(token, "https://") { - continue - } - - u, err := url.Parse(token) - if err != nil { - continue - } - - host := u.Hostname() - - if !isLocalHost(host) { - return fmt.Sprintf( - - "Command blocked by safety guard "+ - - "(curl/wget is restricted to localhost and private network; %q is a public address)", - - host, - ) - } - } - - return "" -} - -// isLocalHost reports whether host is localhost or a loopback/RFC 1918 private IP. - -// DNS resolution is intentionally avoided to prevent DNS rebinding attacks. - -func isLocalHost(host string) bool { - if strings.EqualFold(host, "localhost") { - return true - } - - ip := net.ParseIP(host) - - if ip == nil { - return false - } - - return ip.IsLoopback() || ip.IsPrivate() -} - -// SetBgMaxLifetimeForTest overrides bgMaxLifetime for testing purposes. - -// This is exposed only for tests; the returned function restores the original value. - -var bgMaxLifetimeOverride time.Duration - -func SetBgMaxLifetimeForTest(d time.Duration) func() { - old := bgMaxLifetimeOverride - - bgMaxLifetimeOverride = d - - return func() { bgMaxLifetimeOverride = old } -} - -func getBgMaxLifetime() time.Duration { - if bgMaxLifetimeOverride > 0 { - return bgMaxLifetimeOverride - } - - return bgMaxLifetime -} diff --git a/pkg/tools/shell_bg.go b/pkg/tools/shell_bg.go new file mode 100644 index 000000000..19ec015fc --- /dev/null +++ b/pkg/tools/shell_bg.go @@ -0,0 +1,573 @@ +package tools + +import ( + "context" + "fmt" + "io" + "os/exec" + "regexp" + "runtime" + "sort" + "strings" + "sync" + "time" +) + +const ( + bgMaxLifetime = 45 * time.Minute + + bgRingBufSize = 32 * 1024 // 32KB + + bgInitCapture = 3 * time.Second + + bgMaxProcesses = 10 +) + +// ringBuffer is a thread-safe circular buffer that retains the most recent bytes. + +type ringBuffer struct { + mu sync.Mutex + + buf []byte + + size int +} + +func newRingBuffer(size int) *ringBuffer { + return &ringBuffer{size: size} +} + +// Write appends data to the ring buffer, dropping oldest bytes if capacity is exceeded. + +func (rb *ringBuffer) Write(p []byte) (int, error) { + rb.mu.Lock() + + defer rb.mu.Unlock() + + rb.buf = append(rb.buf, p...) + + if len(rb.buf) > rb.size { + rb.buf = rb.buf[len(rb.buf)-rb.size:] + } + + return len(p), nil +} + +// String returns the current buffer contents. + +func (rb *ringBuffer) String() string { + rb.mu.Lock() + + defer rb.mu.Unlock() + + return string(rb.buf) +} + +// Lines returns the last n lines from the buffer. + +func (rb *ringBuffer) Lines(n int) []string { + rb.mu.Lock() + + defer rb.mu.Unlock() + + if len(rb.buf) == 0 { + return nil + } + + all := strings.Split(string(rb.buf), "\n") + + // Remove trailing empty element from final newline + + if len(all) > 0 && all[len(all)-1] == "" { + all = all[:len(all)-1] + } + + if n <= 0 || n >= len(all) { + return all + } + + return all[len(all)-n:] +} + +// Match checks if any line in the buffer matches the given regex pattern. + +// Returns the first matching line, or empty string if no match. + +func (rb *ringBuffer) Match(pattern *regexp.Regexp) string { + rb.mu.Lock() + + defer rb.mu.Unlock() + + for _, line := range strings.Split(string(rb.buf), "\n") { + if pattern.MatchString(line) { + return line + } + } + + return "" +} + +// Len returns the current number of bytes in the buffer. + +func (rb *ringBuffer) Len() int { + rb.mu.Lock() + + defer rb.mu.Unlock() + + return len(rb.buf) +} + +// bgProcess represents a background process managed by ExecTool. + +type bgProcess struct { + id string + + command string + + cmd *exec.Cmd + + pid int + + startedAt time.Time + + output *ringBuffer + + done chan struct{} // closed when process exits + + exitErr error + + cancel context.CancelFunc // cancels the monitor goroutine +} + +// isRunning returns true if the process has not yet exited. + +func (bp *bgProcess) isRunning() bool { + select { + case <-bp.done: + + return false + + default: + + return true + } +} + +// executeBg starts a background process and returns immediately. + +func (t *ExecTool) executeBg(command, cwd string) *ToolResult { + t.bgMu.Lock() + + // Check max processes limit + + running := 0 + + for _, bp := range t.bgProcesses { + if bp.isRunning() { + running++ + } + } + + if running >= bgMaxProcesses { + t.bgMu.Unlock() + + return ErrorResult( + + fmt.Sprintf("maximum background processes reached (%d). Kill an existing one first.", bgMaxProcesses), + ) + } + + t.bgNextID++ + + id := fmt.Sprintf("bg-%d", t.bgNextID) + + t.bgMu.Unlock() + + var cmd *exec.Cmd + + if runtime.GOOS == "windows" { + cmd = exec.Command("powershell", "-NoProfile", "-NonInteractive", "-Command", command) + } else { + cmd = exec.Command("sh", "-c", command) + } + + if cwd != "" { + cmd.Dir = cwd + } + + prepareCommandForTermination(cmd) + + output := newRingBuffer(bgRingBufSize) + + // Use pipes to capture output + + stdoutPipe, err := cmd.StdoutPipe() + if err != nil { + return ErrorResult(fmt.Sprintf("failed to create stdout pipe: %v", err)) + } + + stderrPipe, err := cmd.StderrPipe() + if err != nil { + return ErrorResult(fmt.Sprintf("failed to create stderr pipe: %v", err)) + } + + if err := cmd.Start(); err != nil { + return ErrorResult(fmt.Sprintf("failed to start background command: %v", err)) + } + + monitorCtx, monitorCancel := context.WithCancel(t.bgCtx) + + bp := &bgProcess{ + id: id, + + command: command, + + cmd: cmd, + + pid: cmd.Process.Pid, + + startedAt: time.Now(), + + output: output, + + done: make(chan struct{}), + + cancel: monitorCancel, + } + + t.bgMu.Lock() + + t.bgProcesses[id] = bp + + t.bgMu.Unlock() + + // io.Copy goroutines: pipe stdout/stderr into ring buffer + + go io.Copy(output, stdoutPipe) + + go io.Copy(output, stderrPipe) + + // cmd.Wait goroutine + + waitDone := make(chan error, 1) + + go func() { + waitDone <- cmd.Wait() + }() + + // Monitor goroutine: handles lifetime timer, process exit, and shutdown + + go func() { + lifetime := time.NewTimer(getBgMaxLifetime()) + + defer lifetime.Stop() + + select { + case err := <-waitDone: + + // Process exited naturally + + bp.exitErr = err + + close(bp.done) + + case <-lifetime.C: + + // Max lifetime exceeded — kill + + _ = terminateProcessTree(cmd) + + select { + case err := <-waitDone: + + bp.exitErr = err + + case <-time.After(2 * time.Second): + + if cmd.Process != nil { + _ = cmd.Process.Kill() + } + + bp.exitErr = <-waitDone + } + + close(bp.done) + + case <-monitorCtx.Done(): + + // Shutdown or explicit kill via cancel + + _ = terminateProcessTree(cmd) + + select { + case err := <-waitDone: + + bp.exitErr = err + + case <-time.After(2 * time.Second): + + if cmd.Process != nil { + _ = cmd.Process.Kill() + } + + bp.exitErr = <-waitDone + } + + select { + case <-bp.done: + + default: + + close(bp.done) + } + } + }() + + // Capture initial output (wait up to bgInitCapture) + + time.Sleep(bgInitCapture) + + initialOutput := output.String() + + var sb strings.Builder + + fmt.Fprintf(&sb, "Background process started.\n") + + fmt.Fprintf(&sb, " id: %s\n", id) + + fmt.Fprintf(&sb, " pid: %d\n", bp.pid) + + fmt.Fprintf(&sb, " cmd: %s\n", command) + + fmt.Fprintf(&sb, " max lifetime: %s\n", getBgMaxLifetime()) + + if initialOutput != "" { + fmt.Fprintf(&sb, "\nInitial output:\n%s", initialOutput) + } + + return &ToolResult{ + ForLLM: sb.String(), + + ForUser: fmt.Sprintf("Background process %s (pid=%d) started: %s", id, bp.pid, command), + } +} + +// handleBgAction handles bg_action=output and bg_action=kill. + +func (t *ExecTool) handleBgAction(action, bgID string) *ToolResult { + if bgID == "" { + return ErrorResult("bg_id is required for bg_action") + } + + t.bgMu.Lock() + + bp, ok := t.bgProcesses[bgID] + + t.bgMu.Unlock() + + if !ok { + return ErrorResult(fmt.Sprintf("background process %q not found", bgID)) + } + + switch action { + case "output": + + return t.bgOutput(bp) + + case "kill": + + return t.bgKill(bp) + + default: + + return ErrorResult(fmt.Sprintf("unknown bg_action %q (use 'output' or 'kill')", action)) + } +} + +func (t *ExecTool) bgOutput(bp *bgProcess) *ToolResult { + var sb strings.Builder + + fmt.Fprintf(&sb, "[%s] pid=%d %s\n", bp.id, bp.pid, bp.command) + + if bp.isRunning() { + uptime := time.Since(bp.startedAt).Truncate(time.Second) + + fmt.Fprintf(&sb, "Status: running (uptime: %s, max: %s)\n", uptime, getBgMaxLifetime()) + } else { + ran := time.Since(bp.startedAt).Truncate(time.Second) + + if bp.exitErr != nil { + fmt.Fprintf(&sb, "Status: exited with error (ran: %s): %v\n", ran, bp.exitErr) + } else { + fmt.Fprintf(&sb, "Status: exited=0 (ran: %s)\n", ran) + } + } + + output := bp.output.String() + + if output == "" { + fmt.Fprintf(&sb, "\n(no output)") + } else { + fmt.Fprintf(&sb, "\nOutput:\n%s", output) + } + + return &ToolResult{ + ForLLM: sb.String(), + + ForUser: sb.String(), + } +} + +func (t *ExecTool) bgKill(bp *bgProcess) *ToolResult { + if bp.isRunning() { + bp.cancel() // triggers monitor goroutine cleanup + + // Wait for process to actually exit + + select { + case <-bp.done: + + case <-time.After(5 * time.Second): + } + } + + t.bgMu.Lock() + + delete(t.bgProcesses, bp.id) + + t.bgMu.Unlock() + + msg := fmt.Sprintf("Background process %s (pid=%d) terminated: %s", bp.id, bp.pid, bp.command) + + return &ToolResult{ + ForLLM: msg, + + ForUser: msg, + } +} + +// BgProcesses returns a snapshot of background processes for use by bg_monitor. + +func (t *ExecTool) BgProcesses() map[string]*bgProcess { + t.bgMu.Lock() + + defer t.bgMu.Unlock() + + snapshot := make(map[string]*bgProcess, len(t.bgProcesses)) + + for k, v := range t.bgProcesses { + snapshot[k] = v + } + + return snapshot +} + +// RuntimeStatus implements StatusProvider for system prompt injection. + +func (t *ExecTool) RuntimeStatus() string { + t.bgMu.Lock() + + defer t.bgMu.Unlock() + + if len(t.bgProcesses) == 0 { + return "" + } + + // Sort by ID for stable output + + ids := make([]string, 0, len(t.bgProcesses)) + + for id := range t.bgProcesses { + ids = append(ids, id) + } + + sort.Strings(ids) + + var sb strings.Builder + + sb.WriteString("## Background Processes\n\n") + + for _, id := range ids { + bp := t.bgProcesses[id] + + if bp.isRunning() { + uptime := time.Since(bp.startedAt).Truncate(time.Second) + + fmt.Fprintf(&sb, " [%s] pid=%d running (uptime: %s, max: %s) %s\n", + + id, bp.pid, uptime, getBgMaxLifetime(), bp.command) + } else { + ran := time.Since(bp.startedAt).Truncate(time.Second) + + if bp.exitErr != nil { + fmt.Fprintf(&sb, " [%s] pid=%d exited=err (ran: %s) %s\n", + + id, bp.pid, ran, bp.command) + } else { + fmt.Fprintf(&sb, " [%s] pid=%d exited=0 (ran: %s) %s\n", + + id, bp.pid, ran, bp.command) + } + } + } + + sb.WriteString("\nUse exec with bg_action=\"output\" / \"kill\" and bg_id to manage.\n") + + sb.WriteString("Use bg_monitor for list/watch/tail operations.") + + return sb.String() +} + +// Shutdown terminates all background processes. Call on application exit. + +func (t *ExecTool) Shutdown() { + t.bgShutdown() // cancel all monitor goroutines + + t.bgMu.Lock() + + procs := make([]*bgProcess, 0, len(t.bgProcesses)) + + for _, bp := range t.bgProcesses { + procs = append(procs, bp) + } + + t.bgMu.Unlock() + + // Wait for all processes to exit + + for _, bp := range procs { + select { + case <-bp.done: + + case <-time.After(5 * time.Second): + + // Force kill if still running + + if bp.cmd.Process != nil { + _ = bp.cmd.Process.Kill() + } + } + } +} + +// SetBgMaxLifetimeForTest overrides bgMaxLifetime for testing purposes. + +// This is exposed only for tests; the returned function restores the original value. + +var bgMaxLifetimeOverride time.Duration + +func SetBgMaxLifetimeForTest(d time.Duration) func() { + old := bgMaxLifetimeOverride + + bgMaxLifetimeOverride = d + + return func() { bgMaxLifetimeOverride = old } +} + +func getBgMaxLifetime() time.Duration { + if bgMaxLifetimeOverride > 0 { + return bgMaxLifetimeOverride + } + + return bgMaxLifetime +} diff --git a/pkg/tools/shell_ext.go b/pkg/tools/shell_ext.go new file mode 100644 index 000000000..e3b034148 --- /dev/null +++ b/pkg/tools/shell_ext.go @@ -0,0 +1,26 @@ +package tools + +import ( + "context" + "sync" +) + +// execToolExt holds fork-specific fields for ExecTool. +// Embedded in ExecTool so existing field access (t.bgProcesses, etc.) continues to work. +type execToolExt struct { + allowRules [][]string // pre-split command prefix allowlist + + localNetOnly bool // restrict curl/wget to localhost + RFC 1918 + + // Background process management + + bgMu sync.Mutex + + bgProcesses map[string]*bgProcess + + bgNextID int + + bgShutdown context.CancelFunc // cancels all bg monitor goroutines + + bgCtx context.Context +} diff --git a/pkg/tools/shell_net.go b/pkg/tools/shell_net.go new file mode 100644 index 000000000..34f3397a6 --- /dev/null +++ b/pkg/tools/shell_net.go @@ -0,0 +1,82 @@ +package tools + +import ( + "fmt" + "net" + "net/url" + "path/filepath" + "strings" +) + +// SetLocalNetOnly restricts curl/wget to localhost and RFC 1918 private addresses. +func (t *ExecTool) SetLocalNetOnly(v bool) { + t.localNetOnly = v +} + +// isCurlOrWget reports whether command is a curl or wget invocation. + +func isCurlOrWget(command string) bool { + fields := strings.Fields(command) + + if len(fields) == 0 { + return false + } + + base := filepath.Base(fields[0]) + + return base == "curl" || base == "wget" +} + +// checkCurlLocalNet validates that all http/https URLs in a curl/wget command + +// target localhost or RFC 1918 private addresses. + +// Returns an error message string, or empty string if the command is allowed. + +func checkCurlLocalNet(command string) string { + for _, token := range strings.Fields(command) { + token = strings.Trim(token, "\"'") + + if !strings.HasPrefix(token, "http://") && !strings.HasPrefix(token, "https://") { + continue + } + + u, err := url.Parse(token) + if err != nil { + continue + } + + host := u.Hostname() + + if !isLocalHost(host) { + return fmt.Sprintf( + + "Command blocked by safety guard "+ + + "(curl/wget is restricted to localhost and private network; %q is a public address)", + + host, + ) + } + } + + return "" +} + +// isLocalHost reports whether host is localhost or a loopback/RFC 1918 private IP. + +// DNS resolution is intentionally avoided to prevent DNS rebinding attacks. + +func isLocalHost(host string) bool { + if strings.EqualFold(host, "localhost") { + return true + } + + ip := net.ParseIP(host) + + if ip == nil { + return false + } + + return ip.IsLoopback() || ip.IsPrivate() +} diff --git a/pkg/tools/subagent.go b/pkg/tools/subagent.go index 09b3a11ab..8ce9d961a 100644 --- a/pkg/tools/subagent.go +++ b/pkg/tools/subagent.go @@ -3,11 +3,7 @@ package tools import ( "context" "fmt" - "os" - "path/filepath" - "sort" "strconv" - "strings" "sync" "time" @@ -23,73 +19,6 @@ import ( const spawnTimeout = 30 * time.Minute -// ContainerMessage is sent from a subagent to the conductor via outCh. - -type ContainerMessage struct { - Type string // "question" or "plan_review" - - Content string - - TaskID string -} - -// isDeliberatePreset returns true for presets that use the deliberate - -// (clarifying → review → executing) workflow with escalation channels. - -func isDeliberatePreset(p Preset) bool { - switch p { - case PresetCoder, PresetWorker, PresetCoordinator: - - return true - } - - return false -} - -// SubagentPlanState represents the deliberate workflow phase. - -type SubagentPlanState int - -const ( - PlanNone SubagentPlanState = iota // Not a deliberate preset - - PlanClarifying // Gathering info, asking questions - - PlanReview // Plan submitted, awaiting approval - - PlanExecuting // Plan approved, executing - - PlanCompleted // Done - -) - -// String returns a human-readable label for the plan state. - -func (s SubagentPlanState) String() string { - switch s { - case PlanClarifying: - - return "clarifying" - - case PlanReview: - - return "review" - - case PlanExecuting: - - return "executing" - - case PlanCompleted: - - return "completed" - - default: - - return "none" - } -} - type SubagentTask struct { ID string @@ -337,62 +266,6 @@ func (sm *SubagentManager) runTask(ctx context.Context, task *SubagentTask, pres } } -// clarifyingSystemPrompt returns the system prompt for the clarifying phase. - -func clarifyingSystemPrompt() string { - return `You are a deliberate subagent in the CLARIFYING phase. - -Your job is to understand the task fully before acting. You MUST: - -1. Read relevant files and gather context using your tools. - -2. If anything is unclear, use ask_conductor to ask the conductor. - -3. When you have a clear plan, use submit_plan with a goal and steps. - - - -Do NOT execute any changes yet. Only investigate and plan. - -Available escalation tools: ask_conductor, submit_plan.` -} - -// executingSystemPrompt returns the system prompt for the executing phase. - -func executingSystemPrompt() string { - return `You are a deliberate subagent in the EXECUTING phase. Your plan was approved. - -Execute the plan steps methodically. Use all available tools to complete the work. - -After completing, provide a clear summary of what was done and how it was verified. - - - -If you encounter a blocker, use ask_conductor to escalate.` -} - -// exploratorySystemPrompt returns the system prompt for exploratory presets. - -func exploratorySystemPrompt(p Preset) string { - switch p { - case PresetScout, PresetAnalyst: - - return `You are an exploratory subagent. Investigate the task and report your findings. - -Use your best judgment when encountering ambiguity. Use tools as needed. - -Return clear findings and observations.` - - default: - - return `You are a subagent. Complete the given task independently and report the result. - -You have access to tools - use them as needed to complete your task. - -After completing the task, provide a clear summary of what was done.` - } -} - // getLLMOptions returns the LLM options snapshot under read lock. func (sm *SubagentManager) getLLMOptions() map[string]any { @@ -546,402 +419,6 @@ func (sm *SubagentManager) finishTask( } } -// setPlanState updates task's plan state in memory and records status in session DAG. - -func (sm *SubagentManager) setPlanState(task *SubagentTask, state SubagentPlanState) { - task.PlanState = state - - if sm.recorder != nil { - subKey := routing.BuildSubagentSessionKey(task.ID) - - _ = sm.recorder.RecordCompletion(subKey, state.String(), "") - } -} - -// runExploratoryTask runs a single-phase tool loop for exploratory presets. - -func (sm *SubagentManager) runExploratoryTask( - ctx context.Context, - task *SubagentTask, - preset Preset, - callback AsyncCallback, -) { - systemPrompt := buildSubagentSystemPrompt(exploratorySystemPrompt(preset), sm.workspace) - - messages := []providers.Message{ - {Role: "system", Content: systemPrompt}, - - {Role: "user", Content: task.Task}, - } - - select { - case <-ctx.Done(): - - sm.mu.Lock() - - task.Status = "canceled" - - task.Result = "Task canceled before execution" - - sm.mu.Unlock() - - return - - default: - } - - sm.mu.RLock() - - reg := sm.tools - - if IsValidPreset(preset) { - reg = sm.buildPresetRegistry(preset, sm.workspace, task) - } - - maxIter := sm.maxIterations - - sm.mu.RUnlock() - - sm.reporter.ReportConversation("conductor", task.ID, task.Task) - - loopResult, err := RunToolLoop(ctx, ToolLoopConfig{ - Provider: sm.provider, - - Model: sm.defaultModel, - - Tools: reg, - - MaxIterations: maxIter, - - LLMOptions: sm.getLLMOptions(), - - Reporter: sm.reporter, - - AgentID: task.ID, - }, messages, task.OriginChannel, task.OriginChatID) - - sm.finishTask(ctx, task, messages, loopResult, err, callback) -} - -// runDeliberateTask runs the clarifying → review → executing workflow. - -func (sm *SubagentManager) runDeliberateTask( - ctx context.Context, - task *SubagentTask, - preset Preset, - callback AsyncCallback, -) { - select { - case <-ctx.Done(): - - sm.mu.Lock() - - task.Status = "canceled" - - task.Result = "Task canceled before execution" - - sm.mu.Unlock() - - return - - default: - } - - sm.mu.RLock() - - reg := sm.buildPresetRegistry(preset, sm.workspace, task) - - maxIter := sm.maxIterations - - sm.mu.RUnlock() - - sm.reporter.ReportConversation("conductor", task.ID, task.Task) - - sm.setPlanState(task, PlanClarifying) - - // Phase 1: Clarifying — subagent gathers info and submits a plan. - - clarifyMsgs := []providers.Message{ - {Role: "system", Content: buildSubagentSystemPrompt(clarifyingSystemPrompt(), sm.workspace)}, - - {Role: "user", Content: task.Task}, - } - - clarifyResult, err := RunToolLoop(ctx, ToolLoopConfig{ - Provider: sm.provider, - - Model: sm.defaultModel, - - Tools: reg, - - MaxIterations: maxIter, - - LLMOptions: sm.getLLMOptions(), - - Reporter: sm.reporter, - - AgentID: task.ID, - }, clarifyMsgs, task.OriginChannel, task.OriginChatID) - if err != nil { - sm.finishTask(ctx, task, clarifyMsgs, nil, err, callback) - - return - } - - // After clarifying, the subagent should have used submit_plan. - - // If it didn't produce a plan, treat the clarifying result as direct completion. - - if task.PlanGoal == "" { - sm.finishTask(ctx, task, clarifyMsgs, clarifyResult, nil, callback) - - return - } - - // Phase 2: Executing — plan was approved, now execute it. - - sm.setPlanState(task, PlanExecuting) - - executeMsgs := []providers.Message{ - {Role: "system", Content: buildSubagentSystemPrompt(executingSystemPrompt(), sm.workspace)}, - - {Role: "user", Content: fmt.Sprintf("Execute the approved plan:\nGoal: %s\nSteps:\n%s", - - task.PlanGoal, formatPlanSteps(task.PlanSteps))}, - } - - execResult, err := RunToolLoop(ctx, ToolLoopConfig{ - Provider: sm.provider, - - Model: sm.defaultModel, - - Tools: reg, - - MaxIterations: maxIter * 2, // Executing gets more iterations - - LLMOptions: sm.getLLMOptions(), - - Reporter: sm.reporter, - - AgentID: task.ID, - }, executeMsgs, task.OriginChannel, task.OriginChatID) - - sm.finishTask(ctx, task, executeMsgs, execResult, err, callback) -} - -// formatPlanSteps formats plan steps as a numbered list. - -func formatPlanSteps(steps []string) string { - var b strings.Builder - - for i, step := range steps { - fmt.Fprintf(&b, "%d. %s\n", i+1, step) - } - - return b.String() -} - -// buildPresetRegistry constructs a ToolRegistry for the given preset with appropriate restrictions. - -// If task is non-nil and has escalation channels, ask_conductor and submit_plan are registered. - -func (sm *SubagentManager) buildPresetRegistry(preset Preset, writeRoot string, task ...*SubagentTask) *ToolRegistry { - registry := NewToolRegistry() - - config := SandboxConfigForPreset(preset, writeRoot) - - readRoot := writeRoot - - if readRoot == "" { - readRoot = sm.workspace - } - - // Register read_file and list_dir with restrict=true - - if config.AllowedTools["read_file"] { - registry.Register(NewReadFileTool(readRoot, true, 0)) - } - - if config.AllowedTools["list_dir"] { - registry.Register(NewListDirTool(readRoot, true)) - } - - // Register write tools only if allowed and writeRoot is set - - if config.AllowedTools["write_file"] && writeRoot != "" { - registry.Register(NewWriteFileTool(writeRoot, true)) - - registry.Register(NewEditFileTool(writeRoot, true)) - - registry.Register(NewAppendFileTool(writeRoot, true)) - } - - // Register exec and bg_monitor if allowed. - - // Each subagent gets its own ExecTool to avoid mutating the shared instance's - - // allowRules (which would leak sandbox restrictions to the conductor). - - if config.AllowedTools["exec"] { - execWorkDir := writeRoot - - if execWorkDir == "" { - execWorkDir = sm.workspace - } - - execTool, err := NewExecTool(execWorkDir, true) - if err != nil { - // exec disabled for this subagent; skip registration - - return registry - } - - if config.ExecPolicy != nil { - execTool.SetAllowRules(config.ExecPolicy.AllowRules) - - execTool.SetLocalNetOnly(config.ExecPolicy.LocalNetOnly) - } - - registry.Register(execTool) - - if config.AllowedTools["bg_monitor"] { - registry.Register(NewBgMonitorTool(execTool)) - } - } - - // Register git tools (worktree-safe push and PR creation) - - if config.AllowedTools["git_push"] { - registry.Register(NewGitPushTool()) - } - - if config.AllowedTools["create_pr"] { - registry.Register(NewCreatePRTool()) - } - - // Register web tools - - if config.AllowedTools["web_search"] { - webSearchTool, _ := NewWebSearchTool(sm.webSearchOpts) - - if webSearchTool != nil { - registry.Register(webSearchTool) - } - } - - if config.AllowedTools["web_fetch"] { - if fetchTool, err := NewWebFetchTool(50000); err == nil { - registry.Register(fetchTool) - } - } - - // Register message tool (always available) - - registry.Register(NewMessageTool()) - - // Register spawn tool only for coordinator preset - - if config.AllowedTools["spawn"] && preset == PresetCoordinator { - spawnTool := NewSpawnTool(sm) - - registry.Register(spawnTool) - } - - // Register escalation tools for deliberate presets with channels. - - if len(task) > 0 && task[0] != nil && task[0].outCh != nil { - t := task[0] - - subKey := "subagent:" + t.ID - - registry.Register(NewAskConductorTool( - - t.ID, sm.conductorSessionKey, subKey, - - t.outCh, t.inCh, sm.recorder, - )) - - submitPlan := NewSubmitPlanTool( - - t.ID, sm.conductorSessionKey, subKey, - - t.outCh, t.inCh, sm.recorder, - ) - - submitPlan.SetPlanCallback(func(goal string, steps []string) { - t.PlanGoal = goal - - t.PlanSteps = steps - }) - - registry.Register(submitPlan) - } - - return registry -} - -// PendingQuestions drains all outCh channels and returns pending container messages. - -// Non-blocking: reads all available messages without waiting. - -func (sm *SubagentManager) PendingQuestions() []ContainerMessage { - sm.mu.RLock() - - defer sm.mu.RUnlock() - - var msgs []ContainerMessage - - for _, task := range sm.tasks { - if task.outCh == nil { - continue - } - - for { - select { - case msg := <-task.outCh: - - msgs = append(msgs, msg) - - default: - - goto nextTask - } - } - - nextTask: - } - - return msgs -} - -// AnswerQuestion sends an answer to a subagent's inCh (non-blocking). - -func (sm *SubagentManager) AnswerQuestion(taskID, answer string) error { - sm.mu.RLock() - - task, ok := sm.tasks[taskID] - - sm.mu.RUnlock() - - if !ok { - return fmt.Errorf("task %q not found", taskID) - } - - if task.inCh == nil { - return fmt.Errorf("task %q has no escalation channel", taskID) - } - - select { - case task.inCh <- answer: - - return nil - - default: - - return fmt.Errorf("task %q answer channel full", taskID) - } -} - // WaitAll blocks until all spawned subagent goroutines have finished // or the timeout expires. Returns true if all goroutines finished, @@ -1152,120 +629,3 @@ func (t *SubagentTool) Execute(ctx context.Context, args map[string]any) *ToolRe Async: false, } } - -// formatToolStats formats a tool stats map as a compact string: "exec:3,read_file:5". - -// Keys are sorted alphabetically for deterministic output. - -func formatToolStats(stats map[string]int) string { - keys := make([]string, 0, len(stats)) - - for k := range stats { - keys = append(keys, k) - } - - sort.Strings(keys) - - parts := make([]string, 0, len(keys)) - - for _, k := range keys { - parts = append(parts, k+":"+strconv.Itoa(stats[k])) - } - - return strings.Join(parts, ",") -} - -// extractPlanContext reads MEMORY.md from the workspace and extracts relevant - -// sections (Task, Context, Commands) to provide as subagent environment. - -func extractPlanContext(workspace string) string { - memPath := filepath.Join(workspace, "memory", "MEMORY.md") - - data, err := os.ReadFile(memPath) - if err != nil { - return "" - } - - content := string(data) - - var sections []string - - // Extract key sections by header. - - for _, header := range []string{"## Context", "## Commands", "## Orchestration"} { - if section := extractSection(content, header); section != "" { - sections = append(sections, section) - } - } - - // Also extract the task line from the header block. - - for _, line := range strings.Split(content, "\n") { - if strings.HasPrefix(line, "> Task:") { - sections = append([]string{strings.TrimSpace(line)}, sections...) - - break - } - } - - if len(sections) == 0 { - return "" - } - - return strings.Join(sections, "\n\n") -} - -// extractSection extracts a markdown section by header (including its content - -// until the next section of the same or higher level). - -func extractSection(content, header string) string { - idx := strings.Index(content, header) - - if idx < 0 { - return "" - } - - // Determine header level. - - level := 0 - - for _, c := range header { - if c == '#' { - level++ - } else { - break - } - } - - start := idx - - rest := content[idx+len(header):] - - // Find next section at same or higher level. - - nextHeader := "\n" + strings.Repeat("#", level) + " " - - end := strings.Index(rest, nextHeader) - - if end < 0 { - return strings.TrimSpace(content[start:]) - } - - return strings.TrimSpace(content[start : start+len(header)+end]) -} - -// buildSubagentSystemPrompt builds an enriched system prompt for a subagent - -// by combining the base prompt with environment context from MEMORY.md. - -func buildSubagentSystemPrompt(basePrompt, workspace string) string { - envContext := extractPlanContext(workspace) - - if envContext == "" { - return basePrompt - } - - return basePrompt + "\n\n## Environment Context\n\n" + envContext -} diff --git a/pkg/tools/subagent_container.go b/pkg/tools/subagent_container.go new file mode 100644 index 000000000..b68c963ce --- /dev/null +++ b/pkg/tools/subagent_container.go @@ -0,0 +1,269 @@ +package tools + +import ( + "context" + "fmt" + "strings" + + "github.com/sipeed/picoclaw/pkg/providers" + "github.com/sipeed/picoclaw/pkg/routing" +) + +// ContainerMessage is sent from a subagent to the conductor via outCh. + +type ContainerMessage struct { + Type string // "question" or "plan_review" + + Content string + + TaskID string +} + +// isDeliberatePreset returns true for presets that use the deliberate + +// (clarifying → review → executing) workflow with escalation channels. + +func isDeliberatePreset(p Preset) bool { + switch p { + case PresetCoder, PresetWorker, PresetCoordinator: + + return true + } + + return false +} + +// SubagentPlanState represents the deliberate workflow phase. + +type SubagentPlanState int + +const ( + PlanNone SubagentPlanState = iota // Not a deliberate preset + + PlanClarifying // Gathering info, asking questions + + PlanReview // Plan submitted, awaiting approval + + PlanExecuting // Plan approved, executing + + PlanCompleted // Done + +) + +// String returns a human-readable label for the plan state. + +func (s SubagentPlanState) String() string { + switch s { + case PlanClarifying: + + return "clarifying" + + case PlanReview: + + return "review" + + case PlanExecuting: + + return "executing" + + case PlanCompleted: + + return "completed" + + default: + + return "none" + } +} + +// setPlanState updates task's plan state in memory and records status in session DAG. + +func (sm *SubagentManager) setPlanState(task *SubagentTask, state SubagentPlanState) { + task.PlanState = state + + if sm.recorder != nil { + subKey := routing.BuildSubagentSessionKey(task.ID) + + _ = sm.recorder.RecordCompletion(subKey, state.String(), "") + } +} + +// runDeliberateTask runs the clarifying → review → executing workflow. + +func (sm *SubagentManager) runDeliberateTask( + ctx context.Context, + task *SubagentTask, + preset Preset, + callback AsyncCallback, +) { + select { + case <-ctx.Done(): + + sm.mu.Lock() + + task.Status = "canceled" + + task.Result = "Task canceled before execution" + + sm.mu.Unlock() + + return + + default: + } + + sm.mu.RLock() + + reg := sm.buildPresetRegistry(preset, sm.workspace, task) + + maxIter := sm.maxIterations + + sm.mu.RUnlock() + + sm.reporter.ReportConversation("conductor", task.ID, task.Task) + + sm.setPlanState(task, PlanClarifying) + + // Phase 1: Clarifying — subagent gathers info and submits a plan. + + clarifyMsgs := []providers.Message{ + {Role: "system", Content: buildSubagentSystemPrompt(clarifyingSystemPrompt(), sm.workspace)}, + + {Role: "user", Content: task.Task}, + } + + clarifyResult, err := RunToolLoop(ctx, ToolLoopConfig{ + Provider: sm.provider, + + Model: sm.defaultModel, + + Tools: reg, + + MaxIterations: maxIter, + + LLMOptions: sm.getLLMOptions(), + + Reporter: sm.reporter, + + AgentID: task.ID, + }, clarifyMsgs, task.OriginChannel, task.OriginChatID) + if err != nil { + sm.finishTask(ctx, task, clarifyMsgs, nil, err, callback) + + return + } + + // After clarifying, the subagent should have used submit_plan. + + // If it didn't produce a plan, treat the clarifying result as direct completion. + + if task.PlanGoal == "" { + sm.finishTask(ctx, task, clarifyMsgs, clarifyResult, nil, callback) + + return + } + + // Phase 2: Executing — plan was approved, now execute it. + + sm.setPlanState(task, PlanExecuting) + + executeMsgs := []providers.Message{ + {Role: "system", Content: buildSubagentSystemPrompt(executingSystemPrompt(), sm.workspace)}, + + {Role: "user", Content: fmt.Sprintf("Execute the approved plan:\nGoal: %s\nSteps:\n%s", + + task.PlanGoal, formatPlanSteps(task.PlanSteps))}, + } + + execResult, err := RunToolLoop(ctx, ToolLoopConfig{ + Provider: sm.provider, + + Model: sm.defaultModel, + + Tools: reg, + + MaxIterations: maxIter * 2, // Executing gets more iterations + + LLMOptions: sm.getLLMOptions(), + + Reporter: sm.reporter, + + AgentID: task.ID, + }, executeMsgs, task.OriginChannel, task.OriginChatID) + + sm.finishTask(ctx, task, executeMsgs, execResult, err, callback) +} + +// formatPlanSteps formats plan steps as a numbered list. + +func formatPlanSteps(steps []string) string { + var b strings.Builder + + for i, step := range steps { + fmt.Fprintf(&b, "%d. %s\n", i+1, step) + } + + return b.String() +} + +// PendingQuestions drains all outCh channels and returns pending container messages. + +// Non-blocking: reads all available messages without waiting. + +func (sm *SubagentManager) PendingQuestions() []ContainerMessage { + sm.mu.RLock() + + defer sm.mu.RUnlock() + + var msgs []ContainerMessage + + for _, task := range sm.tasks { + if task.outCh == nil { + continue + } + + for { + select { + case msg := <-task.outCh: + + msgs = append(msgs, msg) + + default: + + goto nextTask + } + } + + nextTask: + } + + return msgs +} + +// AnswerQuestion sends an answer to a subagent's inCh (non-blocking). + +func (sm *SubagentManager) AnswerQuestion(taskID, answer string) error { + sm.mu.RLock() + + task, ok := sm.tasks[taskID] + + sm.mu.RUnlock() + + if !ok { + return fmt.Errorf("task %q not found", taskID) + } + + if task.inCh == nil { + return fmt.Errorf("task %q has no escalation channel", taskID) + } + + select { + case task.inCh <- answer: + + return nil + + default: + + return fmt.Errorf("task %q answer channel full", taskID) + } +} diff --git a/pkg/tools/subagent_preset.go b/pkg/tools/subagent_preset.go new file mode 100644 index 000000000..e9354865a --- /dev/null +++ b/pkg/tools/subagent_preset.go @@ -0,0 +1,389 @@ +package tools + +import ( + "context" + "os" + "path/filepath" + "sort" + "strconv" + "strings" + + "github.com/sipeed/picoclaw/pkg/providers" +) + +// clarifyingSystemPrompt returns the system prompt for the clarifying phase. + +func clarifyingSystemPrompt() string { + return `You are a deliberate subagent in the CLARIFYING phase. + +Your job is to understand the task fully before acting. You MUST: + +1. Read relevant files and gather context using your tools. + +2. If anything is unclear, use ask_conductor to ask the conductor. + +3. When you have a clear plan, use submit_plan with a goal and steps. + + + +Do NOT execute any changes yet. Only investigate and plan. + +Available escalation tools: ask_conductor, submit_plan.` +} + +// executingSystemPrompt returns the system prompt for the executing phase. + +func executingSystemPrompt() string { + return `You are a deliberate subagent in the EXECUTING phase. Your plan was approved. + +Execute the plan steps methodically. Use all available tools to complete the work. + +After completing, provide a clear summary of what was done and how it was verified. + + + +If you encounter a blocker, use ask_conductor to escalate.` +} + +// exploratorySystemPrompt returns the system prompt for exploratory presets. + +func exploratorySystemPrompt(p Preset) string { + switch p { + case PresetScout, PresetAnalyst: + + return `You are an exploratory subagent. Investigate the task and report your findings. + +Use your best judgment when encountering ambiguity. Use tools as needed. + +Return clear findings and observations.` + + default: + + return `You are a subagent. Complete the given task independently and report the result. + +You have access to tools - use them as needed to complete your task. + +After completing the task, provide a clear summary of what was done.` + } +} + +// runExploratoryTask runs a single-phase tool loop for exploratory presets. + +func (sm *SubagentManager) runExploratoryTask( + ctx context.Context, + task *SubagentTask, + preset Preset, + callback AsyncCallback, +) { + systemPrompt := buildSubagentSystemPrompt(exploratorySystemPrompt(preset), sm.workspace) + + messages := []providers.Message{ + {Role: "system", Content: systemPrompt}, + + {Role: "user", Content: task.Task}, + } + + select { + case <-ctx.Done(): + + sm.mu.Lock() + + task.Status = "canceled" + + task.Result = "Task canceled before execution" + + sm.mu.Unlock() + + return + + default: + } + + sm.mu.RLock() + + reg := sm.tools + + if IsValidPreset(preset) { + reg = sm.buildPresetRegistry(preset, sm.workspace, task) + } + + maxIter := sm.maxIterations + + sm.mu.RUnlock() + + sm.reporter.ReportConversation("conductor", task.ID, task.Task) + + loopResult, err := RunToolLoop(ctx, ToolLoopConfig{ + Provider: sm.provider, + + Model: sm.defaultModel, + + Tools: reg, + + MaxIterations: maxIter, + + LLMOptions: sm.getLLMOptions(), + + Reporter: sm.reporter, + + AgentID: task.ID, + }, messages, task.OriginChannel, task.OriginChatID) + + sm.finishTask(ctx, task, messages, loopResult, err, callback) +} + +// buildPresetRegistry constructs a ToolRegistry for the given preset with appropriate restrictions. + +// If task is non-nil and has escalation channels, ask_conductor and submit_plan are registered. + +func (sm *SubagentManager) buildPresetRegistry(preset Preset, writeRoot string, task ...*SubagentTask) *ToolRegistry { + registry := NewToolRegistry() + + config := SandboxConfigForPreset(preset, writeRoot) + + readRoot := writeRoot + + if readRoot == "" { + readRoot = sm.workspace + } + + // Register read_file and list_dir with restrict=true + + if config.AllowedTools["read_file"] { + registry.Register(NewReadFileTool(readRoot, true, 0)) + } + + if config.AllowedTools["list_dir"] { + registry.Register(NewListDirTool(readRoot, true)) + } + + // Register write tools only if allowed and writeRoot is set + + if config.AllowedTools["write_file"] && writeRoot != "" { + registry.Register(NewWriteFileTool(writeRoot, true)) + + registry.Register(NewEditFileTool(writeRoot, true)) + + registry.Register(NewAppendFileTool(writeRoot, true)) + } + + // Register exec and bg_monitor if allowed. + + // Each subagent gets its own ExecTool to avoid mutating the shared instance's + + // allowRules (which would leak sandbox restrictions to the conductor). + + if config.AllowedTools["exec"] { + execWorkDir := writeRoot + + if execWorkDir == "" { + execWorkDir = sm.workspace + } + + execTool, err := NewExecTool(execWorkDir, true) + if err != nil { + // exec disabled for this subagent; skip registration + + return registry + } + + if config.ExecPolicy != nil { + execTool.SetAllowRules(config.ExecPolicy.AllowRules) + + execTool.SetLocalNetOnly(config.ExecPolicy.LocalNetOnly) + } + + registry.Register(execTool) + + if config.AllowedTools["bg_monitor"] { + registry.Register(NewBgMonitorTool(execTool)) + } + } + + // Register git tools (worktree-safe push and PR creation) + + if config.AllowedTools["git_push"] { + registry.Register(NewGitPushTool()) + } + + if config.AllowedTools["create_pr"] { + registry.Register(NewCreatePRTool()) + } + + // Register web tools + + if config.AllowedTools["web_search"] { + webSearchTool, _ := NewWebSearchTool(sm.webSearchOpts) + + if webSearchTool != nil { + registry.Register(webSearchTool) + } + } + + if config.AllowedTools["web_fetch"] { + if fetchTool, err := NewWebFetchTool(50000); err == nil { + registry.Register(fetchTool) + } + } + + // Register message tool (always available) + + registry.Register(NewMessageTool()) + + // Register spawn tool only for coordinator preset + + if config.AllowedTools["spawn"] && preset == PresetCoordinator { + spawnTool := NewSpawnTool(sm) + + registry.Register(spawnTool) + } + + // Register escalation tools for deliberate presets with channels. + + if len(task) > 0 && task[0] != nil && task[0].outCh != nil { + t := task[0] + + subKey := "subagent:" + t.ID + + registry.Register(NewAskConductorTool( + + t.ID, sm.conductorSessionKey, subKey, + + t.outCh, t.inCh, sm.recorder, + )) + + submitPlan := NewSubmitPlanTool( + + t.ID, sm.conductorSessionKey, subKey, + + t.outCh, t.inCh, sm.recorder, + ) + + submitPlan.SetPlanCallback(func(goal string, steps []string) { + t.PlanGoal = goal + + t.PlanSteps = steps + }) + + registry.Register(submitPlan) + } + + return registry +} + +// extractPlanContext reads MEMORY.md from the workspace and extracts relevant + +// sections (Task, Context, Commands) to provide as subagent environment. + +func extractPlanContext(workspace string) string { + memPath := filepath.Join(workspace, "memory", "MEMORY.md") + + data, err := os.ReadFile(memPath) + if err != nil { + return "" + } + + content := string(data) + + var sections []string + + // Extract key sections by header. + + for _, header := range []string{"## Context", "## Commands", "## Orchestration"} { + if section := extractSection(content, header); section != "" { + sections = append(sections, section) + } + } + + // Also extract the task line from the header block. + + for _, line := range strings.Split(content, "\n") { + if strings.HasPrefix(line, "> Task:") { + sections = append([]string{strings.TrimSpace(line)}, sections...) + + break + } + } + + if len(sections) == 0 { + return "" + } + + return strings.Join(sections, "\n\n") +} + +// extractSection extracts a markdown section by header (including its content + +// until the next section of the same or higher level). + +func extractSection(content, header string) string { + idx := strings.Index(content, header) + + if idx < 0 { + return "" + } + + // Determine header level. + + level := 0 + + for _, c := range header { + if c == '#' { + level++ + } else { + break + } + } + + start := idx + + rest := content[idx+len(header):] + + // Find next section at same or higher level. + + nextHeader := "\n" + strings.Repeat("#", level) + " " + + end := strings.Index(rest, nextHeader) + + if end < 0 { + return strings.TrimSpace(content[start:]) + } + + return strings.TrimSpace(content[start : start+len(header)+end]) +} + +// buildSubagentSystemPrompt builds an enriched system prompt for a subagent + +// by combining the base prompt with environment context from MEMORY.md. + +func buildSubagentSystemPrompt(basePrompt, workspace string) string { + envContext := extractPlanContext(workspace) + + if envContext == "" { + return basePrompt + } + + return basePrompt + "\n\n## Environment Context\n\n" + envContext +} + +// formatToolStats formats a tool stats map as a compact string: "exec:3,read_file:5". + +// Keys are sorted alphabetically for deterministic output. + +func formatToolStats(stats map[string]int) string { + keys := make([]string, 0, len(stats)) + + for k := range stats { + keys = append(keys, k) + } + + sort.Strings(keys) + + parts := make([]string, 0, len(keys)) + + for _, k := range keys { + parts = append(parts, k+":"+strconv.Itoa(stats[k])) + } + + return strings.Join(parts, ",") +}