Merge pull request #34 from dj-oyu/refactor/conflict-isolation
refactor: isolate fork changes into _ext.go files to reduce upstream merge conflicts
This commit is contained in:
commit
fe7ebe21ad
18 changed files with 2136 additions and 1984 deletions
|
|
@ -16,134 +16,15 @@ 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 {
|
||||
contextBuilderExt // fork-specific fields (see context_ext.go)
|
||||
|
||||
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
|
||||
|
||||
|
|
@ -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,
|
||||
}
|
||||
}
|
||||
|
|
|
|||
68
pkg/agent/context_ext.go
Normal file
68
pkg/agent/context_ext.go
Normal file
|
|
@ -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,
|
||||
}
|
||||
}
|
||||
117
pkg/agent/context_orch.go
Normal file
117
pkg/agent/context_orch.go
Normal file
|
|
@ -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`
|
||||
79
pkg/agent/context_plan.go
Normal file
79
pkg/agent/context_plan.go
Normal file
|
|
@ -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()
|
||||
}
|
||||
|
|
@ -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: <workspace>/.worktrees/<branch-basename>/
|
||||
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 {
|
||||
|
|
|
|||
122
pkg/agent/instance_ext.go
Normal file
122
pkg/agent/instance_ext.go
Normal file
|
|
@ -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: <workspace>/.worktrees/<branch-basename>/
|
||||
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 ""
|
||||
}
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
46
pkg/agent/loop_ext.go
Normal file
46
pkg/agent/loop_ext.go
Normal file
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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,6 +76,8 @@ type channelWorker struct {
|
|||
}
|
||||
|
||||
type Manager struct {
|
||||
managerExt // fork-specific fields (see manager_ext.go)
|
||||
|
||||
channels map[string]Channel
|
||||
workers map[string]*channelWorker
|
||||
bus *bus.MessageBus
|
||||
|
|
@ -103,9 +88,6 @@ type Manager struct {
|
|||
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
|
||||
}
|
||||
|
||||
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()
|
||||
|
|
|
|||
286
pkg/channels/manager_ext.go
Normal file
286
pkg/channels/manager_ext.go
Normal file
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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")
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
573
pkg/tools/shell_bg.go
Normal file
573
pkg/tools/shell_bg.go
Normal file
|
|
@ -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
|
||||
}
|
||||
26
pkg/tools/shell_ext.go
Normal file
26
pkg/tools/shell_ext.go
Normal file
|
|
@ -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
|
||||
}
|
||||
82
pkg/tools/shell_net.go
Normal file
82
pkg/tools/shell_net.go
Normal file
|
|
@ -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()
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
269
pkg/tools/subagent_container.go
Normal file
269
pkg/tools/subagent_container.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
389
pkg/tools/subagent_preset.go
Normal file
389
pkg/tools/subagent_preset.go
Normal file
|
|
@ -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, ",")
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue