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