diff --git a/pkg/agent/instance.go b/pkg/agent/instance.go index a6fd365c7..d5172d806 100644 --- a/pkg/agent/instance.go +++ b/pkg/agent/instance.go @@ -51,12 +51,16 @@ func NewAgentInstance( toolsRegistry.Register(tools.NewReadFileTool(workspace, restrict)) toolsRegistry.Register(tools.NewWriteFileTool(workspace, restrict)) toolsRegistry.Register(tools.NewListDirTool(workspace, restrict)) - toolsRegistry.Register(tools.NewExecToolWithConfig(workspace, restrict, cfg)) + execTool := tools.NewExecToolWithConfig(workspace, restrict, cfg) + toolsRegistry.Register(execTool) + toolsRegistry.Register(tools.NewProcessTool(execTool.ProcessManager())) toolsRegistry.Register(tools.NewEditFileTool(workspace, restrict)) toolsRegistry.Register(tools.NewAppendFileTool(workspace, restrict)) sessionsDir := filepath.Join(workspace, "sessions") sessionsManager := session.NewSessionManager(sessionsDir) + toolsRegistry.Register(tools.NewSessionsListTool(sessionsManager)) + toolsRegistry.Register(tools.NewSessionsHistoryTool(sessionsManager)) contextBuilder := NewContextBuilder(workspace) diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 693f2227b..ad88fda8f 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -55,9 +55,6 @@ type processOptions struct { func NewAgentLoop(cfg *config.Config, msgBus *bus.MessageBus, provider providers.LLMProvider) *AgentLoop { registry := NewAgentRegistry(cfg, provider) - // Register shared tools to all agents - registerSharedTools(cfg, msgBus, registry, provider) - // Set up shared fallback chain cooldown := providers.NewCooldownTracker() fallbackChain := providers.NewFallbackChain(cooldown) @@ -69,7 +66,7 @@ func NewAgentLoop(cfg *config.Config, msgBus *bus.MessageBus, provider providers stateManager = state.NewManager(defaultAgent.Workspace) } - return &AgentLoop{ + al := &AgentLoop{ bus: msgBus, cfg: cfg, registry: registry, @@ -77,6 +74,11 @@ func NewAgentLoop(cfg *config.Config, msgBus *bus.MessageBus, provider providers summarizing: sync.Map{}, fallback: fallbackChain, } + + // Register shared tools to all agents. + registerSharedTools(cfg, msgBus, registry, provider, al) + + return al } // registerSharedTools registers tools that are shared across all agents (web, message, spawn). @@ -85,6 +87,7 @@ func registerSharedTools( msgBus *bus.MessageBus, registry *AgentRegistry, provider providers.LLMProvider, + sessionsExecutor tools.SessionsSendExecutor, ) { for _, agentID := range registry.ListAgentIDs() { agent, ok := registry.GetAgent(agentID) @@ -140,15 +143,30 @@ func registerSharedTools( agent.Tools.Register(tools.NewFindSkillsTool(registryMgr, searchCache)) agent.Tools.Register(tools.NewInstallSkillTool(registryMgr, agent.Workspace)) - // Spawn tool with allowlist checker + // Spawn/session tools with allowlist checker. subagentManager := tools.NewSubagentManager(provider, agent.Model, agent.Workspace, msgBus) subagentManager.SetLLMOptions(agent.MaxTokens, agent.Temperature) spawnTool := tools.NewSpawnTool(subagentManager) + sessionsSpawnTool := tools.NewSessionsSpawnTool(subagentManager) currentAgentID := agentID - spawnTool.SetAllowlistChecker(func(targetAgentID string) bool { + allowlist := func(targetAgentID string) bool { return registry.CanSpawnSubagent(currentAgentID, targetAgentID) - }) + } + spawnTool.SetAllowlistChecker(allowlist) + sessionsSpawnTool.SetAllowlistChecker(allowlist) agent.Tools.Register(spawnTool) + agent.Tools.Register(sessionsSpawnTool) + + if sessionsExecutor != nil { + agent.Tools.Register(tools.NewSessionsSendTool(sessionsExecutor)) + } else { + logger.WarnCF("agent", "sessions_send tool disabled: executor unavailable", map[string]any{ + "agent_id": currentAgentID, + }) + } + + // Update context builder with the complete tools registry + agent.ContextBuilder.SetToolsRegistry(agent.Tools) } } @@ -251,6 +269,45 @@ func (al *AgentLoop) ProcessDirectWithChannel( return al.processMessage(ctx, msg) } +// ProcessSessionMessage injects a message into a specific session key directly. +// Unlike ProcessDirectWithChannel, this bypasses route-derived session rewriting. +func (al *AgentLoop) ProcessSessionMessage( + ctx context.Context, + content, sessionKey, channel, chatID string, +) (string, error) { + key := strings.TrimSpace(sessionKey) + if key == "" { + return "", fmt.Errorf("sessionKey is required") + } + + targetAgent := al.registry.GetDefaultAgent() + if parsed := routing.ParseAgentSessionKey(strings.ToLower(key)); parsed != nil { + if agent, ok := al.registry.GetAgent(parsed.AgentID); ok { + targetAgent = agent + } + } + if targetAgent == nil { + return "", fmt.Errorf("no agent available for session %q", key) + } + + if strings.TrimSpace(channel) == "" { + channel = "system" + } + if strings.TrimSpace(chatID) == "" { + chatID = "sessions-send" + } + + return al.runAgentLoop(ctx, targetAgent, processOptions{ + SessionKey: key, + Channel: channel, + ChatID: chatID, + UserMessage: content, + DefaultResponse: "I've completed processing but have no response to give.", + EnableSummary: true, + SendResponse: false, + }) +} + // ProcessHeartbeat processes a heartbeat request without session history. // Each heartbeat is independent and doesn't accumulate context. func (al *AgentLoop) ProcessHeartbeat(ctx context.Context, content, channel, chatID string) (string, error) { @@ -746,6 +803,16 @@ func (al *AgentLoop) updateToolContexts(agent *AgentInstance, channel, chatID st st.SetContext(channel, chatID) } } + if tool, ok := agent.Tools.Get("sessions_send"); ok { + if st, ok := tool.(tools.ContextualTool); ok { + st.SetContext(channel, chatID) + } + } + if tool, ok := agent.Tools.Get("sessions_spawn"); ok { + if st, ok := tool.(tools.ContextualTool); ok { + st.SetContext(channel, chatID) + } + } } // maybeSummarize triggers summarization if the session history exceeds thresholds. diff --git a/pkg/session/manager.go b/pkg/session/manager.go index 08f0b0ad2..aad54ec4c 100644 --- a/pkg/session/manager.go +++ b/pkg/session/manager.go @@ -4,6 +4,7 @@ import ( "encoding/json" "os" "path/filepath" + "sort" "strings" "sync" "time" @@ -280,3 +281,58 @@ func (sm *SessionManager) SetHistory(key string, history []providers.Message) { session.Updated = time.Now() } } + +// GetSessionSnapshot returns a deep-copied snapshot of a session by key. +func (sm *SessionManager) GetSessionSnapshot(key string) (*Session, bool) { + sm.mu.RLock() + defer sm.mu.RUnlock() + + stored, ok := sm.sessions[key] + if !ok { + return nil, false + } + + snapshot := Session{ + Key: stored.Key, + Summary: stored.Summary, + Created: stored.Created, + Updated: stored.Updated, + } + if len(stored.Messages) > 0 { + snapshot.Messages = make([]providers.Message, len(stored.Messages)) + copy(snapshot.Messages, stored.Messages) + } else { + snapshot.Messages = []providers.Message{} + } + + return &snapshot, true +} + +// ListSessionSnapshots returns deep-copied snapshots of all sessions, sorted by Updated descending. +func (sm *SessionManager) ListSessionSnapshots() []Session { + sm.mu.RLock() + defer sm.mu.RUnlock() + + snapshots := make([]Session, 0, len(sm.sessions)) + for _, stored := range sm.sessions { + snapshot := Session{ + Key: stored.Key, + Summary: stored.Summary, + Created: stored.Created, + Updated: stored.Updated, + } + if len(stored.Messages) > 0 { + snapshot.Messages = make([]providers.Message, len(stored.Messages)) + copy(snapshot.Messages, stored.Messages) + } else { + snapshot.Messages = []providers.Message{} + } + snapshots = append(snapshots, snapshot) + } + + sort.Slice(snapshots, func(i, j int) bool { + return snapshots[i].Updated.After(snapshots[j].Updated) + }) + + return snapshots +} diff --git a/pkg/session/manager_test.go b/pkg/session/manager_test.go index 5ef5f4349..0800fc0b1 100644 --- a/pkg/session/manager_test.go +++ b/pkg/session/manager_test.go @@ -4,6 +4,7 @@ import ( "os" "path/filepath" "testing" + "time" ) func TestSanitizeFilename(t *testing.T) { @@ -72,3 +73,46 @@ func TestSave_RejectsPathTraversal(t *testing.T) { } } } + +func TestGetSessionSnapshot_IsDeepCopy(t *testing.T) { + sm := NewSessionManager(t.TempDir()) + key := "agent:main:main" + sm.AddMessage(key, "user", "hello") + + snapshot, ok := sm.GetSessionSnapshot(key) + if !ok { + t.Fatalf("expected snapshot for key %q", key) + } + if len(snapshot.Messages) != 1 { + t.Fatalf("expected 1 message in snapshot, got %d", len(snapshot.Messages)) + } + + // Mutate returned snapshot; internal state should remain unchanged. + snapshot.Messages[0].Content = "mutated" + history := sm.GetHistory(key) + if history[0].Content != "hello" { + t.Fatalf("snapshot mutation leaked into manager state, got %q", history[0].Content) + } +} + +func TestListSessionSnapshots_SortedByUpdatedDesc(t *testing.T) { + sm := NewSessionManager(t.TempDir()) + sm.AddMessage("session-a", "user", "old") + time.Sleep(10 * time.Millisecond) + sm.AddMessage("session-b", "user", "new") + + snapshots := sm.ListSessionSnapshots() + if len(snapshots) != 2 { + t.Fatalf("expected 2 snapshots, got %d", len(snapshots)) + } + if snapshots[0].Key != "session-b" { + t.Fatalf("expected newest session first, got %q", snapshots[0].Key) + } + + // Ensure returned slices are copies. + snapshots[0].Messages[0].Content = "changed" + history := sm.GetHistory("session-b") + if history[0].Content != "new" { + t.Fatalf("snapshot mutation leaked into manager state, got %q", history[0].Content) + } +} diff --git a/pkg/tools/process.go b/pkg/tools/process.go new file mode 100644 index 000000000..02f4f473f --- /dev/null +++ b/pkg/tools/process.go @@ -0,0 +1,710 @@ +package tools + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "os/exec" + "sort" + "strings" + "sync" + "sync/atomic" + "time" +) + +const ( + defaultProcessMaxOutputChars = 30000 + defaultProcessMaxPendingChars = 12000 + defaultProcessLogTailLines = 200 +) + +var ( + ErrProcessSessionNotFound = errors.New("process session not found") + ErrProcessSessionRunning = errors.New("process session is still running") +) + +type processSession struct { + ID string + Command string + CWD string + + StartedAt time.Time + UpdatedAt time.Time + EndedAt time.Time + + PID int + Status string + ExitCode *int + ExitError string + Truncated bool + + output string + pending string + cmd *exec.Cmd + stdin io.WriteCloser + cancel func() + killRequested bool + + notify chan struct{} + mu sync.Mutex +} + +type ProcessSessionSnapshot struct { + SessionID string `json:"session_id"` + Status string `json:"status"` + Command string `json:"command"` + CWD string `json:"cwd,omitempty"` + PID int `json:"pid,omitempty"` + + StartedAt string `json:"started_at"` + UpdatedAt string `json:"updated_at"` + EndedAt string `json:"ended_at,omitempty"` + + ExitCode *int `json:"exit_code,omitempty"` + ExitError string `json:"exit_error,omitempty"` + Truncated bool `json:"truncated,omitempty"` +} + +type ProcessPollResult struct { + Session ProcessSessionSnapshot `json:"session"` + Output string `json:"output,omitempty"` + TimedOut bool `json:"timed_out,omitempty"` +} + +type ProcessLogResult struct { + Session ProcessSessionSnapshot `json:"session"` + TotalLines int `json:"total_lines"` + Offset int `json:"offset"` + Limit int `json:"limit,omitempty"` + Lines []string `json:"lines"` + Output string `json:"output"` +} + +type ProcessManager struct { + mu sync.RWMutex + sessions map[string]*processSession + nextID atomic.Uint64 + maxOutputChars int + maxPendingChars int +} + +func NewProcessManager(maxOutputChars int) *ProcessManager { + if maxOutputChars <= 0 { + maxOutputChars = defaultProcessMaxOutputChars + } + + maxPendingChars := defaultProcessMaxPendingChars + if maxPendingChars > maxOutputChars { + maxPendingChars = maxOutputChars + } + + return &ProcessManager{ + sessions: make(map[string]*processSession), + maxOutputChars: maxOutputChars, + maxPendingChars: maxPendingChars, + } +} + +func (pm *ProcessManager) StartSession( + command, cwd string, + cmd *exec.Cmd, + stdin io.WriteCloser, + cancel func(), +) string { + id := fmt.Sprintf("proc-%d", pm.nextID.Add(1)) + now := time.Now() + + pid := 0 + if cmd != nil && cmd.Process != nil { + pid = cmd.Process.Pid + } + + session := &processSession{ + ID: id, + Command: command, + CWD: cwd, + StartedAt: now, + UpdatedAt: now, + PID: pid, + Status: "running", + cmd: cmd, + stdin: stdin, + cancel: cancel, + notify: make(chan struct{}, 1), + } + + pm.mu.Lock() + pm.sessions[id] = session + pm.mu.Unlock() + + return id +} + +func (pm *ProcessManager) AppendOutput(sessionID, chunk string) { + if chunk == "" { + return + } + + session, ok := pm.getSession(sessionID) + if !ok { + return + } + + session.mu.Lock() + defer session.mu.Unlock() + + session.output, session.Truncated = appendWithCap( + session.output, + chunk, + pm.maxOutputChars, + session.Truncated, + ) + session.pending, session.Truncated = appendWithCap( + session.pending, + chunk, + pm.maxPendingChars, + session.Truncated, + ) + session.UpdatedAt = time.Now() + session.signalNotifyLocked() +} + +func (pm *ProcessManager) MarkExited(sessionID string, waitErr error, timedOut bool) { + session, ok := pm.getSession(sessionID) + if !ok { + return + } + + session.mu.Lock() + defer session.mu.Unlock() + + now := time.Now() + session.UpdatedAt = now + session.EndedAt = now + + switch { + case timedOut: + session.Status = "timeout" + case session.killRequested: + session.Status = "killed" + case waitErr != nil: + session.Status = "failed" + default: + session.Status = "completed" + } + + if waitErr != nil { + session.ExitError = waitErr.Error() + } + if code, ok := extractExitCode(waitErr); ok { + session.ExitCode = &code + } else if waitErr == nil { + code := 0 + session.ExitCode = &code + } + + if session.stdin != nil { + _ = session.stdin.Close() + session.stdin = nil + } + if session.cancel != nil { + session.cancel() + session.cancel = nil + } + session.cmd = nil + session.signalNotifyLocked() +} + +func (pm *ProcessManager) GetSnapshot(sessionID string) (ProcessSessionSnapshot, bool) { + session, ok := pm.getSession(sessionID) + if !ok { + return ProcessSessionSnapshot{}, false + } + + session.mu.Lock() + defer session.mu.Unlock() + return session.snapshotLocked(), true +} + +func (pm *ProcessManager) ListSnapshots() []ProcessSessionSnapshot { + pm.mu.RLock() + sessions := make([]*processSession, 0, len(pm.sessions)) + for _, session := range pm.sessions { + sessions = append(sessions, session) + } + pm.mu.RUnlock() + + out := make([]ProcessSessionSnapshot, 0, len(sessions)) + for _, session := range sessions { + session.mu.Lock() + out = append(out, session.snapshotLocked()) + session.mu.Unlock() + } + + sort.Slice(out, func(i, j int) bool { + return out[i].StartedAt > out[j].StartedAt + }) + return out +} + +func (pm *ProcessManager) Poll(sessionID string, timeout time.Duration) (ProcessPollResult, error) { + session, ok := pm.getSession(sessionID) + if !ok { + return ProcessPollResult{}, ErrProcessSessionNotFound + } + if timeout < 0 { + timeout = 0 + } + + deadline := time.Now().Add(timeout) + for { + output, snapshot := session.drainPending() + if output != "" || snapshot.Status != "running" || timeout == 0 { + return ProcessPollResult{ + Session: snapshot, + Output: output, + }, nil + } + + waitFor := time.Until(deadline) + if waitFor <= 0 { + return ProcessPollResult{ + Session: snapshot, + TimedOut: true, + }, nil + } + + select { + case <-session.notify: + // loop and re-check + case <-time.After(waitFor): + snapshot := session.snapshot() + return ProcessPollResult{ + Session: snapshot, + TimedOut: true, + }, nil + } + } +} + +func (pm *ProcessManager) Log( + sessionID string, + offset, limit int, + useDefaultTail bool, +) (ProcessLogResult, error) { + session, ok := pm.getSession(sessionID) + if !ok { + return ProcessLogResult{}, ErrProcessSessionNotFound + } + + output, snapshot := session.outputSnapshot() + lines := normalizeOutputLines(output) + total := len(lines) + + start := offset + if start < 0 { + start = 0 + } + if start > total { + start = total + } + + end := total + effectiveLimit := limit + if useDefaultTail { + if total > defaultProcessLogTailLines { + start = total - defaultProcessLogTailLines + } else { + start = 0 + } + end = total + effectiveLimit = defaultProcessLogTailLines + } else if effectiveLimit > 0 { + if start+effectiveLimit < end { + end = start + effectiveLimit + } + } + + window := []string{} + if start < end { + window = lines[start:end] + } + + return ProcessLogResult{ + Session: snapshot, + TotalLines: total, + Offset: start, + Limit: effectiveLimit, + Lines: window, + Output: strings.Join(window, "\n"), + }, nil +} + +func (pm *ProcessManager) Write(sessionID, data string, eof bool) error { + session, ok := pm.getSession(sessionID) + if !ok { + return ErrProcessSessionNotFound + } + + session.mu.Lock() + if session.Status != "running" { + session.mu.Unlock() + return fmt.Errorf("session %s is not running", sessionID) + } + stdin := session.stdin + session.mu.Unlock() + + if stdin == nil { + return fmt.Errorf("session %s has no writable stdin", sessionID) + } + + if data != "" { + if _, err := io.WriteString(stdin, data); err != nil { + return err + } + } + if eof { + return stdin.Close() + } + return nil +} + +func (pm *ProcessManager) Kill(sessionID string) (bool, error) { + session, ok := pm.getSession(sessionID) + if !ok { + return false, ErrProcessSessionNotFound + } + + session.mu.Lock() + if session.Status != "running" { + session.mu.Unlock() + return false, nil + } + session.killRequested = true + cmd := session.cmd + cancel := session.cancel + session.mu.Unlock() + + if cancel != nil { + cancel() + } + if cmd != nil { + _ = terminateProcessTree(cmd) + } + return true, nil +} + +func (pm *ProcessManager) Clear(sessionID string) error { + session, ok := pm.getSession(sessionID) + if !ok { + return ErrProcessSessionNotFound + } + + session.mu.Lock() + running := session.Status == "running" + session.mu.Unlock() + if running { + return ErrProcessSessionRunning + } + + pm.mu.Lock() + delete(pm.sessions, sessionID) + pm.mu.Unlock() + return nil +} + +func (pm *ProcessManager) Remove(sessionID string) (bool, error) { + session, ok := pm.getSession(sessionID) + if !ok { + return false, ErrProcessSessionNotFound + } + + session.mu.Lock() + running := session.Status == "running" + session.mu.Unlock() + if running { + _, err := pm.Kill(sessionID) + return false, err + } + + pm.mu.Lock() + delete(pm.sessions, sessionID) + pm.mu.Unlock() + return true, nil +} + +func (pm *ProcessManager) getSession(sessionID string) (*processSession, bool) { + pm.mu.RLock() + defer pm.mu.RUnlock() + session, ok := pm.sessions[sessionID] + return session, ok +} + +func (s *processSession) snapshot() ProcessSessionSnapshot { + s.mu.Lock() + defer s.mu.Unlock() + return s.snapshotLocked() +} + +func (s *processSession) drainPending() (string, ProcessSessionSnapshot) { + s.mu.Lock() + defer s.mu.Unlock() + + pending := s.pending + s.pending = "" + return pending, s.snapshotLocked() +} + +func (s *processSession) outputSnapshot() (string, ProcessSessionSnapshot) { + s.mu.Lock() + defer s.mu.Unlock() + return s.output, s.snapshotLocked() +} + +func (s *processSession) snapshotLocked() ProcessSessionSnapshot { + snapshot := ProcessSessionSnapshot{ + SessionID: s.ID, + Status: s.Status, + Command: s.Command, + CWD: s.CWD, + PID: s.PID, + StartedAt: s.StartedAt.Format(time.RFC3339), + UpdatedAt: s.UpdatedAt.Format(time.RFC3339), + ExitCode: s.ExitCode, + ExitError: s.ExitError, + Truncated: s.Truncated, + } + if !s.EndedAt.IsZero() { + snapshot.EndedAt = s.EndedAt.Format(time.RFC3339) + } + return snapshot +} + +func (s *processSession) signalNotifyLocked() { + select { + case s.notify <- struct{}{}: + default: + } +} + +func appendWithCap(current, appendChunk string, max int, truncated bool) (string, bool) { + if max <= 0 { + return current + appendChunk, truncated + } + + combined := current + appendChunk + if len(combined) <= max { + return combined, truncated + } + + return combined[len(combined)-max:], true +} + +func normalizeOutputLines(output string) []string { + if output == "" { + return []string{} + } + normalized := strings.ReplaceAll(output, "\r\n", "\n") + normalized = strings.TrimSuffix(normalized, "\n") + if normalized == "" { + return []string{} + } + return strings.Split(normalized, "\n") +} + +func extractExitCode(err error) (int, bool) { + if err == nil { + return 0, false + } + var exitErr *exec.ExitError + if errors.As(err, &exitErr) { + return exitErr.ExitCode(), true + } + return 0, false +} + +type ProcessTool struct { + processes *ProcessManager +} + +func NewProcessTool(processes *ProcessManager) *ProcessTool { + return &ProcessTool{processes: processes} +} + +func (t *ProcessTool) Name() string { + return "process" +} + +func (t *ProcessTool) Description() string { + return "Manage background exec sessions: list, poll, log, write, kill, clear, remove." +} + +func (t *ProcessTool) Parameters() map[string]any { + return map[string]any{ + "type": "object", + "properties": map[string]any{ + "action": map[string]any{ + "type": "string", + "enum": []string{"list", "poll", "log", "write", "kill", "clear", "remove"}, + "description": "Process action", + }, + "session_id": map[string]any{ + "type": "string", + "description": "Session ID returned by exec background mode", + }, + "data": map[string]any{ + "type": "string", + "description": "Input data for write action", + }, + "eof": map[string]any{ + "type": "boolean", + "description": "Close stdin after write", + }, + "offset": map[string]any{ + "type": "integer", + "description": "Log line offset", + "minimum": 0.0, + }, + "limit": map[string]any{ + "type": "integer", + "description": "Max log lines to return", + "minimum": 0.0, + }, + "timeout_ms": map[string]any{ + "type": "integer", + "description": "Poll wait timeout in milliseconds", + "minimum": 0.0, + }, + }, + "required": []string{"action"}, + } +} + +func (t *ProcessTool) Execute(_ context.Context, args map[string]any) *ToolResult { + if t.processes == nil { + return ErrorResult("process manager not configured") + } + + action, ok := getStringArg(args, "action") + if !ok || strings.TrimSpace(action) == "" { + return ErrorResult("action is required") + } + action = strings.ToLower(strings.TrimSpace(action)) + + switch action { + case "list": + sessions := t.processes.ListSnapshots() + return marshalSilentJSON(map[string]any{ + "count": len(sessions), + "sessions": sessions, + }) + case "poll": + sessionID, ok := getStringArg(args, "session_id") + if !ok || strings.TrimSpace(sessionID) == "" { + return ErrorResult("session_id is required for poll") + } + timeoutMS, err := parseOptionalIntArg(args, "timeout_ms", 0, 0, 5*60*1000) + if err != nil { + return ErrorResult(err.Error()) + } + result, err := t.processes.Poll(strings.TrimSpace(sessionID), time.Duration(timeoutMS)*time.Millisecond) + if err != nil { + return ErrorResult(err.Error()) + } + return marshalSilentJSON(result) + case "log": + sessionID, ok := getStringArg(args, "session_id") + if !ok || strings.TrimSpace(sessionID) == "" { + return ErrorResult("session_id is required for log") + } + offset, err := parseOptionalIntArg(args, "offset", 0, 0, 1_000_000) + if err != nil { + return ErrorResult(err.Error()) + } + limit, err := parseOptionalIntArg(args, "limit", 0, 0, 1_000_000) + if err != nil { + return ErrorResult(err.Error()) + } + _, hasOffset := args["offset"] + _, hasLimit := args["limit"] + useDefaultTail := !hasOffset && !hasLimit + result, err := t.processes.Log(strings.TrimSpace(sessionID), offset, limit, useDefaultTail) + if err != nil { + return ErrorResult(err.Error()) + } + return marshalSilentJSON(result) + case "write": + sessionID, ok := getStringArg(args, "session_id") + if !ok || strings.TrimSpace(sessionID) == "" { + return ErrorResult("session_id is required for write") + } + data, _ := getStringArg(args, "data") + eof, err := parseBoolArg(args, "eof", false) + if err != nil { + return ErrorResult(err.Error()) + } + if err := t.processes.Write(strings.TrimSpace(sessionID), data, eof); err != nil { + return ErrorResult(err.Error()) + } + return marshalSilentJSON(map[string]any{ + "status": "ok", + "action": "write", + "session_id": strings.TrimSpace(sessionID), + }) + case "kill": + sessionID, ok := getStringArg(args, "session_id") + if !ok || strings.TrimSpace(sessionID) == "" { + return ErrorResult("session_id is required for kill") + } + signaled, err := t.processes.Kill(strings.TrimSpace(sessionID)) + if err != nil { + return ErrorResult(err.Error()) + } + return marshalSilentJSON(map[string]any{ + "status": "ok", + "action": "kill", + "session_id": strings.TrimSpace(sessionID), + "kill_signal": signaled, + }) + case "clear": + sessionID, ok := getStringArg(args, "session_id") + if !ok || strings.TrimSpace(sessionID) == "" { + return ErrorResult("session_id is required for clear") + } + if err := t.processes.Clear(strings.TrimSpace(sessionID)); err != nil { + return ErrorResult(err.Error()) + } + return marshalSilentJSON(map[string]any{ + "status": "ok", + "action": "clear", + "session_id": strings.TrimSpace(sessionID), + }) + case "remove": + sessionID, ok := getStringArg(args, "session_id") + if !ok || strings.TrimSpace(sessionID) == "" { + return ErrorResult("session_id is required for remove") + } + removed, err := t.processes.Remove(strings.TrimSpace(sessionID)) + if err != nil { + return ErrorResult(err.Error()) + } + return marshalSilentJSON(map[string]any{ + "status": "ok", + "action": "remove", + "session_id": strings.TrimSpace(sessionID), + "removed": removed, + }) + default: + return ErrorResult(fmt.Sprintf("unknown action: %s", action)) + } +} + +func marshalSilentJSON(payload any) *ToolResult { + data, err := json.MarshalIndent(payload, "", " ") + if err != nil { + return ErrorResult(fmt.Sprintf("failed to encode process payload: %v", err)) + } + return SilentResult(string(data)) +} diff --git a/pkg/tools/process_test.go b/pkg/tools/process_test.go new file mode 100644 index 000000000..4a8325737 --- /dev/null +++ b/pkg/tools/process_test.go @@ -0,0 +1,168 @@ +package tools + +import ( + "context" + "encoding/json" + "strings" + "testing" +) + +func decodeSessionID(t *testing.T, payload string) string { + t.Helper() + var out struct { + SessionID string `json:"session_id"` + } + if err := json.Unmarshal([]byte(payload), &out); err != nil { + t.Fatalf("failed to decode exec payload: %v", err) + } + if out.SessionID == "" { + t.Fatalf("missing session_id in payload: %s", payload) + } + return out.SessionID +} + +func TestExecBackground_WithProcessPoll(t *testing.T) { + execTool := NewExecTool(t.TempDir(), false) + processTool := NewProcessTool(execTool.ProcessManager()) + + start := execTool.Execute(context.Background(), map[string]any{ + "command": "sleep 0.2; echo done", + "background": true, + }) + if start.IsError { + t.Fatalf("exec background failed: %s", start.ForLLM) + } + sessionID := decodeSessionID(t, start.ForLLM) + + poll := processTool.Execute(context.Background(), map[string]any{ + "action": "poll", + "session_id": sessionID, + "timeout_ms": 3000, + }) + if poll.IsError { + t.Fatalf("process poll failed: %s", poll.ForLLM) + } + + var payload struct { + Output string `json:"output"` + Session struct { + Status string `json:"status"` + } `json:"session"` + } + if err := json.Unmarshal([]byte(poll.ForLLM), &payload); err != nil { + t.Fatalf("failed to decode poll payload: %v", err) + } + if payload.Session.Status == "running" { + poll = processTool.Execute(context.Background(), map[string]any{ + "action": "poll", + "session_id": sessionID, + "timeout_ms": 2000, + }) + if poll.IsError { + t.Fatalf("second poll failed: %s", poll.ForLLM) + } + if err := json.Unmarshal([]byte(poll.ForLLM), &payload); err != nil { + t.Fatalf("failed to decode second poll payload: %v", err) + } + } + if !strings.Contains(payload.Output, "done") && payload.Session.Status != "completed" { + t.Fatalf("expected output to contain done or completed status, got: status=%q output=%q", payload.Session.Status, payload.Output) + } + if payload.Session.Status != "completed" { + t.Fatalf("expected completed status, got %q", payload.Session.Status) + } +} + +func TestProcessTool_KillAndRemove(t *testing.T) { + execTool := NewExecTool(t.TempDir(), false) + processTool := NewProcessTool(execTool.ProcessManager()) + + start := execTool.Execute(context.Background(), map[string]any{ + "command": "sleep 60", + "background": true, + }) + if start.IsError { + t.Fatalf("exec background failed: %s", start.ForLLM) + } + sessionID := decodeSessionID(t, start.ForLLM) + + kill := processTool.Execute(context.Background(), map[string]any{ + "action": "kill", + "session_id": sessionID, + }) + if kill.IsError { + t.Fatalf("process kill failed: %s", kill.ForLLM) + } + + poll := processTool.Execute(context.Background(), map[string]any{ + "action": "poll", + "session_id": sessionID, + "timeout_ms": 3000, + }) + if poll.IsError { + t.Fatalf("process poll after kill failed: %s", poll.ForLLM) + } + + var pollPayload struct { + Session struct { + Status string `json:"status"` + } `json:"session"` + } + if err := json.Unmarshal([]byte(poll.ForLLM), &pollPayload); err != nil { + t.Fatalf("failed to decode poll payload: %v", err) + } + if pollPayload.Session.Status == "running" { + t.Fatalf("expected killed session to stop running") + } + + remove := processTool.Execute(context.Background(), map[string]any{ + "action": "remove", + "session_id": sessionID, + }) + if remove.IsError { + t.Fatalf("process remove failed: %s", remove.ForLLM) + } +} + +func TestProcessTool_Write(t *testing.T) { + execTool := NewExecTool(t.TempDir(), false) + processTool := NewProcessTool(execTool.ProcessManager()) + + start := execTool.Execute(context.Background(), map[string]any{ + "command": "cat", + "background": true, + }) + if start.IsError { + t.Fatalf("exec background failed: %s", start.ForLLM) + } + sessionID := decodeSessionID(t, start.ForLLM) + + write := processTool.Execute(context.Background(), map[string]any{ + "action": "write", + "session_id": sessionID, + "data": "ping\n", + "eof": true, + }) + if write.IsError { + t.Fatalf("process write failed: %s", write.ForLLM) + } + + poll := processTool.Execute(context.Background(), map[string]any{ + "action": "poll", + "session_id": sessionID, + "timeout_ms": 3000, + }) + if poll.IsError { + t.Fatalf("poll after write failed: %s", poll.ForLLM) + } + + var payload struct { + Output string `json:"output"` + } + if err := json.Unmarshal([]byte(poll.ForLLM), &payload); err != nil { + t.Fatalf("failed to decode poll payload: %v", err) + } + if !strings.Contains(payload.Output, "ping") { + t.Fatalf("expected output to contain ping, got %q", payload.Output) + } +} diff --git a/pkg/tools/sessions.go b/pkg/tools/sessions.go new file mode 100644 index 000000000..d27c2381f --- /dev/null +++ b/pkg/tools/sessions.go @@ -0,0 +1,445 @@ +package tools + +import ( + "context" + "encoding/json" + "fmt" + "math" + "strings" + "time" + + "github.com/sipeed/picoclaw/pkg/providers" + "github.com/sipeed/picoclaw/pkg/session" +) + +const ( + defaultSessionsListLimit = 50 + maxSessionsListLimit = 200 + defaultSessionsHistorySize = 200 + maxSessionsHistorySize = 1000 + maxSessionPreviewMessages = 50 + maxMessagePreviewChars = 500 +) + +type SessionsListTool struct { + sessions *session.SessionManager +} + +func NewSessionsListTool(sm *session.SessionManager) *SessionsListTool { + return &SessionsListTool{sessions: sm} +} + +func (t *SessionsListTool) Name() string { + return "sessions_list" +} + +func (t *SessionsListTool) Description() string { + return "List known conversation sessions with metadata. Useful for debugging, navigation, and context inspection." +} + +func (t *SessionsListTool) Parameters() map[string]any { + return map[string]any{ + "type": "object", + "properties": map[string]any{ + "kinds": map[string]any{ + "type": "array", + "items": map[string]any{ + "type": "string", + }, + "description": "Optional session kind filter (e.g. main, direct, group, cron, subagent)", + }, + "limit": map[string]any{ + "type": "integer", + "description": "Maximum number of sessions to return (default 50, max 200)", + "minimum": 1.0, + "maximum": 200.0, + }, + "active_minutes": map[string]any{ + "type": "integer", + "description": "Only include sessions updated within N minutes", + "minimum": 1.0, + }, + "message_limit": map[string]any{ + "type": "integer", + "description": "Include up to N recent preview messages per session (default 0, max 50)", + "minimum": 0.0, + "maximum": 50.0, + }, + "include_tools": map[string]any{ + "type": "boolean", + "description": "When message_limit > 0, include tool messages in preview (default false)", + }, + }, + } +} + +func (t *SessionsListTool) Execute(_ context.Context, args map[string]any) *ToolResult { + if t.sessions == nil { + return ErrorResult("session manager not configured") + } + + limit, err := parseIntArg(args, "limit", defaultSessionsListLimit, 1, maxSessionsListLimit) + if err != nil { + return ErrorResult(err.Error()) + } + + activeMinutes, err := parseOptionalIntArg(args, "active_minutes", 0, 0, 365*24*60) + if err != nil { + return ErrorResult(err.Error()) + } + + messageLimit, err := parseIntArg(args, "message_limit", 0, 0, maxSessionPreviewMessages) + if err != nil { + return ErrorResult(err.Error()) + } + + includeTools, err := parseBoolArg(args, "include_tools", false) + if err != nil { + return ErrorResult(err.Error()) + } + + kinds, err := parseStringSliceArg(args, "kinds") + if err != nil { + return ErrorResult(err.Error()) + } + + allowedKinds := make(map[string]struct{}, len(kinds)) + for _, k := range kinds { + allowedKinds[strings.ToLower(strings.TrimSpace(k))] = struct{}{} + } + + now := time.Now() + snapshots := t.sessions.ListSessionSnapshots() + outSessions := make([]sessionListItem, 0, min(limit, len(snapshots))) + + for _, s := range snapshots { + kind := classifySessionKind(s.Key) + + if len(allowedKinds) > 0 { + if _, ok := allowedKinds[kind]; !ok { + continue + } + } + + if activeMinutes > 0 && now.Sub(s.Updated) > time.Duration(activeMinutes)*time.Minute { + continue + } + + item := sessionListItem{ + Key: s.Key, + Kind: kind, + CreatedAt: s.Created.Format(time.RFC3339), + UpdatedAt: s.Updated.Format(time.RFC3339), + MessageCount: len(s.Messages), + } + if strings.TrimSpace(s.Summary) != "" { + item.Summary = s.Summary + } + if messageLimit > 0 { + item.Messages = tailMessages(s.Messages, messageLimit, includeTools) + } + + outSessions = append(outSessions, item) + if len(outSessions) >= limit { + break + } + } + + payload := sessionsListOutput{ + Count: len(outSessions), + Sessions: outSessions, + } + + data, err := json.MarshalIndent(payload, "", " ") + if err != nil { + return ErrorResult(fmt.Sprintf("failed to encode sessions list: %v", err)) + } + return SilentResult(string(data)) +} + +type SessionsHistoryTool struct { + sessions *session.SessionManager +} + +func NewSessionsHistoryTool(sm *session.SessionManager) *SessionsHistoryTool { + return &SessionsHistoryTool{sessions: sm} +} + +func (t *SessionsHistoryTool) Name() string { + return "sessions_history" +} + +func (t *SessionsHistoryTool) Description() string { + return "Get full or partial message history for one session key." +} + +func (t *SessionsHistoryTool) Parameters() map[string]any { + return map[string]any{ + "type": "object", + "properties": map[string]any{ + "session_key": map[string]any{ + "type": "string", + "description": "Session key returned by sessions_list", + }, + "limit": map[string]any{ + "type": "integer", + "description": "Maximum messages to return from the tail (default 200, max 1000)", + "minimum": 1.0, + "maximum": 1000.0, + }, + "include_tools": map[string]any{ + "type": "boolean", + "description": "Include tool-role messages (default false)", + }, + }, + "required": []string{"session_key"}, + } +} + +func (t *SessionsHistoryTool) Execute(_ context.Context, args map[string]any) *ToolResult { + if t.sessions == nil { + return ErrorResult("session manager not configured") + } + + key, ok := getStringArg(args, "session_key") + if !ok { + return ErrorResult("session_key is required") + } + key = strings.TrimSpace(key) + if key == "" { + return ErrorResult("session_key is required") + } + + limit, err := parseIntArg(args, "limit", defaultSessionsHistorySize, 1, maxSessionsHistorySize) + if err != nil { + return ErrorResult(err.Error()) + } + + includeTools, err := parseBoolArg(args, "include_tools", false) + if err != nil { + return ErrorResult(err.Error()) + } + + snapshot, ok := t.sessions.GetSessionSnapshot(key) + if !ok { + return ErrorResult(fmt.Sprintf("session %q not found", key)) + } + + messages := tailMessages(snapshot.Messages, limit, includeTools) + payload := sessionHistoryOutput{ + SessionKey: snapshot.Key, + Kind: classifySessionKind(snapshot.Key), + CreatedAt: snapshot.Created.Format(time.RFC3339), + UpdatedAt: snapshot.Updated.Format(time.RFC3339), + MessageCount: len(messages), + Messages: messages, + } + if strings.TrimSpace(snapshot.Summary) != "" { + payload.Summary = snapshot.Summary + } + + data, err := json.MarshalIndent(payload, "", " ") + if err != nil { + return ErrorResult(fmt.Sprintf("failed to encode session history: %v", err)) + } + return SilentResult(string(data)) +} + +type sessionListItem struct { + Key string `json:"key"` + Kind string `json:"kind"` + CreatedAt string `json:"created_at"` + UpdatedAt string `json:"updated_at"` + MessageCount int `json:"message_count"` + Summary string `json:"summary,omitempty"` + Messages []providers.Message `json:"messages,omitempty"` +} + +type sessionsListOutput struct { + Count int `json:"count"` + Sessions []sessionListItem `json:"sessions"` +} + +type sessionHistoryOutput struct { + SessionKey string `json:"session_key"` + Kind string `json:"kind"` + CreatedAt string `json:"created_at"` + UpdatedAt string `json:"updated_at"` + MessageCount int `json:"message_count"` + Summary string `json:"summary,omitempty"` + Messages []providers.Message `json:"messages"` +} + +func classifySessionKind(key string) string { + k := strings.ToLower(strings.TrimSpace(key)) + switch { + case k == "": + return "other" + case strings.HasPrefix(k, "agent:") && strings.HasSuffix(k, ":main"): + return "main" + case strings.Contains(k, ":subagent:"): + return "subagent" + case strings.Contains(k, ":group:"): + return "group" + case strings.Contains(k, ":channel:"): + return "channel" + case strings.Contains(k, ":direct:"): + return "direct" + case strings.HasPrefix(k, "cron:") || strings.HasPrefix(k, "cron-"): + return "cron" + case strings.HasPrefix(k, "hook:"): + return "hook" + case strings.HasPrefix(k, "node-"): + return "node" + case k == "heartbeat": + return "heartbeat" + default: + return "other" + } +} + +func tailMessages(messages []providers.Message, limit int, includeTools bool) []providers.Message { + if limit <= 0 || len(messages) == 0 { + return nil + } + + selected := make([]providers.Message, 0, min(limit, len(messages))) + for i := len(messages) - 1; i >= 0 && len(selected) < limit; i-- { + msg := messages[i] + if !includeTools && msg.Role == "tool" { + continue + } + if len(msg.Content) > maxMessagePreviewChars { + msg.Content = msg.Content[:maxMessagePreviewChars] + "...(truncated)" + } + selected = append(selected, msg) + } + + // Reverse back to chronological order. + for i, j := 0, len(selected)-1; i < j; i, j = i+1, j-1 { + selected[i], selected[j] = selected[j], selected[i] + } + return selected +} + +func parseIntArg(args map[string]any, key string, defaultVal, minVal, maxVal int) (int, error) { + val, exists := args[key] + if !exists { + return defaultVal, nil + } + + n, err := toInt(val) + if err != nil { + return 0, fmt.Errorf("%s must be an integer", key) + } + if n < minVal || n > maxVal { + return 0, fmt.Errorf("%s must be between %d and %d", key, minVal, maxVal) + } + return n, nil +} + +func parseOptionalIntArg(args map[string]any, key string, defaultVal, minVal, maxVal int) (int, error) { + val, exists := args[key] + if !exists { + return defaultVal, nil + } + + n, err := toInt(val) + if err != nil { + return 0, fmt.Errorf("%s must be an integer", key) + } + if n < minVal || n > maxVal { + return 0, fmt.Errorf("%s must be between %d and %d", key, minVal, maxVal) + } + return n, nil +} + +func parseBoolArg(args map[string]any, key string, defaultVal bool) (bool, error) { + val, exists := args[key] + if !exists { + return defaultVal, nil + } + b, ok := val.(bool) + if !ok { + return false, fmt.Errorf("%s must be a boolean", key) + } + return b, nil +} + +func parseStringSliceArg(args map[string]any, key string) ([]string, error) { + val, exists := args[key] + if !exists { + return nil, nil + } + + switch v := val.(type) { + case []string: + out := make([]string, 0, len(v)) + for _, s := range v { + if trimmed := strings.TrimSpace(s); trimmed != "" { + out = append(out, trimmed) + } + } + return out, nil + case []any: + out := make([]string, 0, len(v)) + for _, item := range v { + s, ok := item.(string) + if !ok { + return nil, fmt.Errorf("%s must be an array of strings", key) + } + if trimmed := strings.TrimSpace(s); trimmed != "" { + out = append(out, trimmed) + } + } + return out, nil + default: + return nil, fmt.Errorf("%s must be an array of strings", key) + } +} + +func getStringArg(args map[string]any, key string) (string, bool) { + v, ok := args[key] + if !ok { + return "", false + } + s, ok := v.(string) + return s, ok +} + +func toInt(v any) (int, error) { + switch n := v.(type) { + case int: + return n, nil + case int8: + return int(n), nil + case int16: + return int(n), nil + case int32: + return int(n), nil + case int64: + return int(n), nil + case uint: + return int(n), nil + case uint8: + return int(n), nil + case uint16: + return int(n), nil + case uint32: + return int(n), nil + case uint64: + return int(n), nil + case float64: + if n != math.Trunc(n) { + return 0, fmt.Errorf("not an integer") + } + return int(n), nil + case float32: + if float64(n) != math.Trunc(float64(n)) { + return 0, fmt.Errorf("not an integer") + } + return int(n), nil + default: + return 0, fmt.Errorf("not an integer") + } +} diff --git a/pkg/tools/sessions_actions.go b/pkg/tools/sessions_actions.go new file mode 100644 index 000000000..64dd6e0a1 --- /dev/null +++ b/pkg/tools/sessions_actions.go @@ -0,0 +1,239 @@ +package tools + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "strings" + "time" +) + +const ( + defaultSessionsSendTimeoutSeconds = 30 + maxSessionsSendTimeoutSeconds = 3600 +) + +// SessionsSendExecutor executes a message inside a target session key. +type SessionsSendExecutor interface { + ProcessSessionMessage(ctx context.Context, content, sessionKey, channel, chatID string) (string, error) +} + +type SessionsSendTool struct { + executor SessionsSendExecutor + channel string + chatID string +} + +func NewSessionsSendTool(executor SessionsSendExecutor) *SessionsSendTool { + return &SessionsSendTool{ + executor: executor, + channel: "system", + chatID: "sessions-send", + } +} + +func (t *SessionsSendTool) Name() string { + return "sessions_send" +} + +func (t *SessionsSendTool) Description() string { + return "Send a message into another session and return the target session's assistant reply." +} + +func (t *SessionsSendTool) Parameters() map[string]any { + return map[string]any{ + "type": "object", + "properties": map[string]any{ + "session_key": map[string]any{ + "type": "string", + "description": "Target session key (from sessions_list)", + }, + "message": map[string]any{ + "type": "string", + "description": "Message to send into the target session", + }, + "timeout_seconds": map[string]any{ + "type": "integer", + "description": "Max wait time for target response (default 30, max 3600, 0 = no timeout)", + "minimum": 0.0, + "maximum": 3600.0, + }, + }, + "required": []string{"session_key", "message"}, + } +} + +func (t *SessionsSendTool) SetContext(channel, chatID string) { + if strings.TrimSpace(channel) != "" { + t.channel = channel + } + if strings.TrimSpace(chatID) != "" { + t.chatID = chatID + } +} + +func (t *SessionsSendTool) Execute(ctx context.Context, args map[string]any) *ToolResult { + if t.executor == nil { + return ErrorResult("sessions executor not configured") + } + + sessionKey, ok := getStringArg(args, "session_key") + if !ok || strings.TrimSpace(sessionKey) == "" { + return ErrorResult("session_key is required") + } + message, ok := getStringArg(args, "message") + if !ok || strings.TrimSpace(message) == "" { + return ErrorResult("message is required") + } + + timeoutSeconds, err := parseOptionalIntArg( + args, + "timeout_seconds", + defaultSessionsSendTimeoutSeconds, + 0, + maxSessionsSendTimeoutSeconds, + ) + if err != nil { + return ErrorResult(err.Error()) + } + + execCtx := ctx + cancel := func() {} + if timeoutSeconds > 0 { + execCtx, cancel = context.WithTimeout(ctx, time.Duration(timeoutSeconds)*time.Second) + } + defer cancel() + + reply, execErr := t.executor.ProcessSessionMessage( + execCtx, + message, + strings.TrimSpace(sessionKey), + "system", + "sessions-send", + ) + + payload := map[string]any{ + "session_key": strings.TrimSpace(sessionKey), + } + + switch { + case execErr == nil: + payload["status"] = "ok" + payload["reply"] = reply + case errors.Is(execErr, context.DeadlineExceeded) || errors.Is(execCtx.Err(), context.DeadlineExceeded): + payload["status"] = "timeout" + payload["error"] = execErr.Error() + default: + payload["status"] = "error" + payload["error"] = execErr.Error() + } + + data, marshalErr := json.MarshalIndent(payload, "", " ") + if marshalErr != nil { + return ErrorResult(fmt.Sprintf("failed to encode sessions_send payload: %v", marshalErr)) + } + if execErr != nil && payload["status"] == "error" { + return ErrorResult(string(data)) + } + return SilentResult(string(data)) +} + +type SessionsSpawnTool struct { + manager *SubagentManager + originChannel string + originChatID string + allowlistCheck func(targetAgentID string) bool +} + +func NewSessionsSpawnTool(manager *SubagentManager) *SessionsSpawnTool { + return &SessionsSpawnTool{ + manager: manager, + originChannel: "cli", + originChatID: "direct", + } +} + +func (t *SessionsSpawnTool) Name() string { + return "sessions_spawn" +} + +func (t *SessionsSpawnTool) Description() string { + return "Spawn a background subagent task and return task metadata." +} + +func (t *SessionsSpawnTool) Parameters() map[string]any { + return map[string]any{ + "type": "object", + "properties": map[string]any{ + "task": map[string]any{ + "type": "string", + "description": "Task content for the spawned subagent", + }, + "label": map[string]any{ + "type": "string", + "description": "Optional short label for tracking", + }, + "agent_id": map[string]any{ + "type": "string", + "description": "Optional target agent id", + }, + }, + "required": []string{"task"}, + } +} + +func (t *SessionsSpawnTool) SetContext(channel, chatID string) { + t.originChannel = channel + t.originChatID = chatID +} + +func (t *SessionsSpawnTool) SetAllowlistChecker(check func(targetAgentID string) bool) { + t.allowlistCheck = check +} + +func (t *SessionsSpawnTool) Execute(ctx context.Context, args map[string]any) *ToolResult { + if t.manager == nil { + return ErrorResult("Subagent manager not configured") + } + + task, ok := getStringArg(args, "task") + if !ok || strings.TrimSpace(task) == "" { + return ErrorResult("task is required") + } + + label, _ := getStringArg(args, "label") + label = strings.TrimSpace(label) + agentID, _ := getStringArg(args, "agent_id") + agentID = strings.TrimSpace(agentID) + + if agentID != "" && t.allowlistCheck != nil && !t.allowlistCheck(agentID) { + return ErrorResult(fmt.Sprintf("not allowed to spawn agent '%s'", agentID)) + } + + taskInfo, err := t.manager.SpawnTask( + ctx, + strings.TrimSpace(task), + label, + agentID, + t.originChannel, + t.originChatID, + nil, + ) + if err != nil { + return ErrorResult(fmt.Sprintf("failed to spawn subagent: %v", err)) + } + + payload := map[string]any{ + "status": "accepted", + "task_id": taskInfo.ID, + "task": taskInfo.Task, + "label": taskInfo.Label, + "agent_id": taskInfo.AgentID, + } + data, marshalErr := json.MarshalIndent(payload, "", " ") + if marshalErr != nil { + return ErrorResult(fmt.Sprintf("failed to encode sessions_spawn payload: %v", marshalErr)) + } + return SilentResult(string(data)) +} diff --git a/pkg/tools/sessions_actions_test.go b/pkg/tools/sessions_actions_test.go new file mode 100644 index 000000000..d14fe15af --- /dev/null +++ b/pkg/tools/sessions_actions_test.go @@ -0,0 +1,138 @@ +package tools + +import ( + "context" + "encoding/json" + "testing" + "time" +) + +type stubSessionsSendExecutor struct { + reply string + err error + called bool + content string + key string + channel string + chatID string +} + +func (s *stubSessionsSendExecutor) ProcessSessionMessage( + ctx context.Context, + content, sessionKey, channel, chatID string, +) (string, error) { + s.called = true + s.content = content + s.key = sessionKey + s.channel = channel + s.chatID = chatID + if s.err != nil { + return "", s.err + } + return s.reply, nil +} + +func TestSessionsSendTool_Success(t *testing.T) { + exec := &stubSessionsSendExecutor{reply: "target reply"} + tool := NewSessionsSendTool(exec) + tool.SetContext("cli", "direct") + + result := tool.Execute(context.Background(), map[string]any{ + "session_key": "agent:main:main", + "message": "hello", + }) + if result.IsError { + t.Fatalf("sessions_send returned error: %s", result.ForLLM) + } + if !result.Silent { + t.Fatalf("sessions_send should be silent") + } + if !exec.called { + t.Fatalf("expected executor to be called") + } + + var payload struct { + Status string `json:"status"` + SessionKey string `json:"session_key"` + Reply string `json:"reply"` + } + if err := json.Unmarshal([]byte(result.ForLLM), &payload); err != nil { + t.Fatalf("failed to decode payload: %v", err) + } + if payload.Status != "ok" { + t.Fatalf("expected status ok, got %q", payload.Status) + } + if payload.Reply != "target reply" { + t.Fatalf("expected reply %q, got %q", "target reply", payload.Reply) + } +} + +func TestSessionsSendTool_TimeoutFromContext(t *testing.T) { + exec := &stubSessionsSendExecutor{ + err: context.DeadlineExceeded, + } + tool := NewSessionsSendTool(exec) + + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancel() + + result := tool.Execute(ctx, map[string]any{ + "session_key": "agent:main:main", + "message": "hello", + }) + if result.IsError { + t.Fatalf("timeout should be reported as structured silent result, got error: %s", result.ForLLM) + } + + var payload struct { + Status string `json:"status"` + } + if err := json.Unmarshal([]byte(result.ForLLM), &payload); err != nil { + t.Fatalf("failed to decode payload: %v", err) + } + if payload.Status != "timeout" { + t.Fatalf("expected timeout status, got %q", payload.Status) + } +} + +func TestSessionsSpawnTool_AcceptedAndAllowlist(t *testing.T) { + manager := NewSubagentManager(&MockLLMProvider{}, "test-model", t.TempDir(), nil) + tool := NewSessionsSpawnTool(manager) + tool.SetAllowlistChecker(func(targetAgentID string) bool { + return targetAgentID != "blocked" + }) + + result := tool.Execute(context.Background(), map[string]any{ + "task": "do something", + "label": "unit", + "agent_id": "worker", + }) + if result.IsError { + t.Fatalf("sessions_spawn returned error: %s", result.ForLLM) + } + if !result.Silent { + t.Fatalf("sessions_spawn should be silent") + } + + var payload struct { + Status string `json:"status"` + TaskID string `json:"task_id"` + } + if err := json.Unmarshal([]byte(result.ForLLM), &payload); err != nil { + t.Fatalf("failed to decode payload: %v", err) + } + if payload.Status != "accepted" { + t.Fatalf("expected status accepted, got %q", payload.Status) + } + if payload.TaskID == "" { + t.Fatalf("expected non-empty task_id") + } + + denied := tool.Execute(context.Background(), map[string]any{ + "task": "do something else", + "agent_id": "blocked", + }) + if !denied.IsError { + t.Fatalf("expected allowlist failure") + } +} diff --git a/pkg/tools/sessions_test.go b/pkg/tools/sessions_test.go new file mode 100644 index 000000000..e9f23e854 --- /dev/null +++ b/pkg/tools/sessions_test.go @@ -0,0 +1,131 @@ +package tools + +import ( + "context" + "encoding/json" + "testing" + + "github.com/sipeed/picoclaw/pkg/providers" + "github.com/sipeed/picoclaw/pkg/session" +) + +func TestSessionsListTool_BasicAndKindFilter(t *testing.T) { + sm := session.NewSessionManager(t.TempDir()) + sm.AddMessage("agent:main:main", "user", "hello") + sm.AddMessage("cron:daily-report", "assistant", "done") + + tool := NewSessionsListTool(sm) + + // Basic list + result := tool.Execute(context.Background(), map[string]any{"limit": 10}) + if result.IsError { + t.Fatalf("sessions_list returned error: %s", result.ForLLM) + } + if !result.Silent { + t.Fatalf("sessions_list should be silent") + } + + var payload struct { + Count int `json:"count"` + Sessions []struct { + Key string `json:"key"` + Kind string `json:"kind"` + } `json:"sessions"` + } + if err := json.Unmarshal([]byte(result.ForLLM), &payload); err != nil { + t.Fatalf("failed to decode sessions_list payload: %v", err) + } + if payload.Count != 2 { + t.Fatalf("expected 2 sessions, got %d", payload.Count) + } + + kindByKey := map[string]string{} + for _, s := range payload.Sessions { + kindByKey[s.Key] = s.Kind + } + if kindByKey["agent:main:main"] != "main" { + t.Fatalf("expected main kind for agent:main:main, got %q", kindByKey["agent:main:main"]) + } + if kindByKey["cron:daily-report"] != "cron" { + t.Fatalf("expected cron kind for cron:daily-report, got %q", kindByKey["cron:daily-report"]) + } + + // Kind filter + filtered := tool.Execute(context.Background(), map[string]any{ + "kinds": []any{"main"}, + "limit": 10, + }) + if filtered.IsError { + t.Fatalf("sessions_list (filtered) returned error: %s", filtered.ForLLM) + } + if err := json.Unmarshal([]byte(filtered.ForLLM), &payload); err != nil { + t.Fatalf("failed to decode filtered payload: %v", err) + } + if payload.Count != 1 { + t.Fatalf("expected 1 filtered session, got %d", payload.Count) + } + if payload.Sessions[0].Key != "agent:main:main" { + t.Fatalf("unexpected filtered session key: %q", payload.Sessions[0].Key) + } +} + +func TestSessionsHistoryTool_IncludeToolsToggle(t *testing.T) { + sm := session.NewSessionManager(t.TempDir()) + key := "agent:main:main" + sm.AddMessage(key, "user", "first") + sm.AddFullMessage(key, providers.Message{Role: "tool", Content: "tool-output", ToolCallID: "tc-1"}) + sm.AddMessage(key, "assistant", "second") + + tool := NewSessionsHistoryTool(sm) + + withoutTools := tool.Execute(context.Background(), map[string]any{ + "session_key": key, + "limit": 10, + }) + if withoutTools.IsError { + t.Fatalf("sessions_history returned error: %s", withoutTools.ForLLM) + } + + var payload struct { + MessageCount int `json:"message_count"` + Messages []providers.Message `json:"messages"` + } + if err := json.Unmarshal([]byte(withoutTools.ForLLM), &payload); err != nil { + t.Fatalf("decode sessions_history payload failed: %v", err) + } + if payload.MessageCount != 2 { + t.Fatalf("expected 2 messages without tools, got %d", payload.MessageCount) + } + for _, msg := range payload.Messages { + if msg.Role == "tool" { + t.Fatalf("tool message should be excluded by default") + } + } + + withTools := tool.Execute(context.Background(), map[string]any{ + "session_key": key, + "limit": 10, + "include_tools": true, + }) + if withTools.IsError { + t.Fatalf("sessions_history(include_tools=true) returned error: %s", withTools.ForLLM) + } + if err := json.Unmarshal([]byte(withTools.ForLLM), &payload); err != nil { + t.Fatalf("decode sessions_history payload failed: %v", err) + } + if payload.MessageCount != 3 { + t.Fatalf("expected 3 messages with tools, got %d", payload.MessageCount) + } +} + +func TestSessionsHistoryTool_NotFound(t *testing.T) { + sm := session.NewSessionManager(t.TempDir()) + tool := NewSessionsHistoryTool(sm) + + result := tool.Execute(context.Background(), map[string]any{ + "session_key": "does-not-exist", + }) + if !result.IsError { + t.Fatalf("expected error for missing session") + } +} diff --git a/pkg/tools/shell.go b/pkg/tools/shell.go index ad1664b5b..5b3cf8952 100644 --- a/pkg/tools/shell.go +++ b/pkg/tools/shell.go @@ -3,8 +3,10 @@ package tools import ( "bytes" "context" + "encoding/json" "errors" "fmt" + "io" "os" "os/exec" "path/filepath" @@ -22,6 +24,7 @@ type ExecTool struct { denyPatterns []*regexp.Regexp allowPatterns []*regexp.Regexp restrictToWorkspace bool + processes *ProcessManager } var defaultDenyPatterns = []*regexp.Regexp{ @@ -106,6 +109,7 @@ func NewExecToolWithConfig(workingDir string, restrict bool, config *config.Conf denyPatterns: denyPatterns, allowPatterns: nil, restrictToWorkspace: restrict, + processes: NewProcessManager(defaultProcessMaxOutputChars), } } @@ -129,6 +133,20 @@ func (t *ExecTool) Parameters() map[string]any { "type": "string", "description": "Optional working directory for the command", }, + "background": map[string]any{ + "type": "boolean", + "description": "Start command in background and manage it via process tool", + }, + "yield_ms": map[string]any{ + "type": "integer", + "description": "Wait this many milliseconds before returning running status", + "minimum": 0.0, + }, + "timeout_seconds": map[string]any{ + "type": "integer", + "description": "Override command timeout in seconds (0 disables timeout)", + "minimum": 0.0, + }, }, "required": []string{"command"}, } @@ -164,22 +182,56 @@ func (t *ExecTool) Execute(ctx context.Context, args map[string]any) *ToolResult return ErrorResult(guardError) } - // timeout == 0 means no timeout + background, err := parseBoolArg(args, "background", false) + if err != nil { + return ErrorResult(err.Error()) + } + + yieldMS, err := parseOptionalIntArg(args, "yield_ms", 0, 0, 60*60*1000) + if err != nil { + return ErrorResult(err.Error()) + } + + timeoutSeconds, hasTimeoutOverride, err := readOptionalIntArg(args, "timeout_seconds", 0, 24*60*60) + if err != nil { + return ErrorResult(err.Error()) + } + + timeout := t.timeout + if hasTimeoutOverride { + if timeoutSeconds == 0 { + timeout = 0 + } else { + timeout = time.Duration(timeoutSeconds) * time.Second + } + } + + if !background && yieldMS <= 0 { + return t.executeSync(ctx, command, cwd, timeout) + } + + return t.executeManaged( + ctx, + command, + cwd, + background, + time.Duration(yieldMS)*time.Millisecond, + timeout, + ) +} + +func (t *ExecTool) executeSync(ctx context.Context, command, cwd string, timeout time.Duration) *ToolResult { + // timeout == 0 means no timeout. var cmdCtx context.Context var cancel context.CancelFunc - if t.timeout > 0 { - cmdCtx, cancel = context.WithTimeout(ctx, t.timeout) + if timeout > 0 { + cmdCtx, cancel = context.WithTimeout(ctx, timeout) } else { cmdCtx, cancel = context.WithCancel(ctx) } defer cancel() - var cmd *exec.Cmd - if runtime.GOOS == "windows" { - cmd = exec.CommandContext(cmdCtx, "powershell", "-NoProfile", "-NonInteractive", "-Command", command) - } else { - cmd = exec.CommandContext(cmdCtx, "sh", "-c", command) - } + cmd := shellCommand(cmdCtx, command) if cwd != "" { cmd.Dir = cwd } @@ -221,7 +273,7 @@ func (t *ExecTool) Execute(ctx context.Context, args map[string]any) *ToolResult if err != nil { if errors.Is(cmdCtx.Err(), context.DeadlineExceeded) { - msg := fmt.Sprintf("Command timed out after %v", t.timeout) + msg := fmt.Sprintf("Command timed out after %v", timeout) return &ToolResult{ ForLLM: msg, ForUser: msg, @@ -231,14 +283,7 @@ func (t *ExecTool) Execute(ctx context.Context, args map[string]any) *ToolResult output += fmt.Sprintf("\nExit code: %v", err) } - if output == "" { - output = "(no output)" - } - - maxLen := 10000 - if len(output) > maxLen { - output = output[:maxLen] + fmt.Sprintf("\n... (truncated, %d more chars)", len(output)-maxLen) - } + output = truncateExecOutput(output) if err != nil { return &ToolResult{ @@ -255,6 +300,229 @@ func (t *ExecTool) Execute(ctx context.Context, args map[string]any) *ToolResult } } +func (t *ExecTool) executeManaged( + ctx context.Context, + command, cwd string, + background bool, + yield time.Duration, + timeout time.Duration, +) *ToolResult { + if t.processes == nil { + return t.executeSync(ctx, command, cwd, timeout) + } + + baseCtx := context.Background() + var cmdCtx context.Context + var cancel context.CancelFunc + if timeout > 0 { + cmdCtx, cancel = context.WithTimeout(baseCtx, timeout) + } else { + cmdCtx, cancel = context.WithCancel(baseCtx) + } + + cmd := shellCommand(cmdCtx, command) + if cwd != "" { + cmd.Dir = cwd + } + prepareCommandForTermination(cmd) + + stdoutPipe, err := cmd.StdoutPipe() + if err != nil { + cancel() + return ErrorResult(fmt.Sprintf("failed to attach stdout: %v", err)) + } + stderrPipe, err := cmd.StderrPipe() + if err != nil { + cancel() + return ErrorResult(fmt.Sprintf("failed to attach stderr: %v", err)) + } + stdinPipe, err := cmd.StdinPipe() + if err != nil { + cancel() + return ErrorResult(fmt.Sprintf("failed to attach stdin: %v", err)) + } + + if err := cmd.Start(); err != nil { + cancel() + return ErrorResult(fmt.Sprintf("failed to start command: %v", err)) + } + + sessionID := t.processes.StartSession(command, cwd, cmd, stdinPipe, cancel) + done := make(chan struct{}) + go t.watchManagedCommand(sessionID, cmdCtx, cmd, stdoutPipe, stderrPipe, done) + + if background { + return t.runningSessionResult(sessionID) + } + + if yield <= 0 { + yield = 10 * time.Second + } + timer := time.NewTimer(yield) + defer timer.Stop() + + select { + case <-done: + pollResult, err := t.processes.Poll(sessionID, 0) + if err != nil { + return ErrorResult(fmt.Sprintf("failed to read managed command result: %v", err)) + } + return formatManagedCompletion(pollResult, timeout) + case <-timer.C: + return t.runningSessionResult(sessionID) + case <-ctx.Done(): + _, _ = t.processes.Kill(sessionID) + return ErrorResult(fmt.Sprintf("command canceled: %v", ctx.Err())) + } +} + +func (t *ExecTool) watchManagedCommand( + sessionID string, + cmdCtx context.Context, + cmd *exec.Cmd, + stdoutPipe, stderrPipe io.ReadCloser, + done chan<- struct{}, +) { + defer close(done) + + stdoutDone := make(chan struct{}) + stderrDone := make(chan struct{}) + go func() { + t.streamManagedOutput(sessionID, stdoutPipe, false) + close(stdoutDone) + }() + go func() { + t.streamManagedOutput(sessionID, stderrPipe, true) + close(stderrDone) + }() + + waitDone := make(chan error, 1) + go func() { + waitDone <- cmd.Wait() + }() + + var waitErr error + select { + case waitErr = <-waitDone: + case <-cmdCtx.Done(): + _ = terminateProcessTree(cmd) + select { + case waitErr = <-waitDone: + case <-time.After(2 * time.Second): + if cmd.Process != nil { + _ = cmd.Process.Kill() + } + waitErr = <-waitDone + } + } + + <-stdoutDone + <-stderrDone + + t.processes.MarkExited(sessionID, waitErr, errors.Is(cmdCtx.Err(), context.DeadlineExceeded)) +} + +func (t *ExecTool) streamManagedOutput(sessionID string, reader io.ReadCloser, stderr bool) { + defer reader.Close() + + buf := make([]byte, 4096) + wroteStderrHeader := false + for { + n, err := reader.Read(buf) + if n > 0 { + if stderr && !wroteStderrHeader { + t.processes.AppendOutput(sessionID, "\nSTDERR:\n") + wroteStderrHeader = true + } + t.processes.AppendOutput(sessionID, string(buf[:n])) + } + if err != nil { + return + } + } +} + +func (t *ExecTool) runningSessionResult(sessionID string) *ToolResult { + payload := map[string]any{ + "status": "running", + "session_id": sessionID, + } + if snap, ok := t.processes.GetSnapshot(sessionID); ok { + payload["pid"] = snap.PID + payload["started_at"] = snap.StartedAt + payload["command"] = snap.Command + } + + data, err := json.MarshalIndent(payload, "", " ") + if err != nil { + return ErrorResult(fmt.Sprintf("failed to encode background exec result: %v", err)) + } + return SilentResult(string(data)) +} + +func formatManagedCompletion(result ProcessPollResult, timeout time.Duration) *ToolResult { + output := truncateExecOutput(result.Output) + status := strings.ToLower(result.Session.Status) + if status == "" { + status = "completed" + } + + switch status { + case "completed": + return UserResult(output) + case "timeout": + msg := fmt.Sprintf("Command timed out after %v", timeout) + if output != "(no output)" { + msg = output + "\n" + msg + } + return ErrorResult(msg) + default: + if result.Session.ExitError != "" && !strings.Contains(output, result.Session.ExitError) { + output += "\nExit code: " + result.Session.ExitError + } + return &ToolResult{ + ForLLM: output, + ForUser: output, + IsError: true, + } + } +} + +func shellCommand(ctx context.Context, command string) *exec.Cmd { + if runtime.GOOS == "windows" { + return exec.CommandContext(ctx, "powershell", "-NoProfile", "-NonInteractive", "-Command", command) + } + return exec.CommandContext(ctx, "sh", "-c", command) +} + +func truncateExecOutput(output string) string { + if output == "" { + output = "(no output)" + } + + const maxLen = 10000 + if len(output) > maxLen { + output = output[:maxLen] + fmt.Sprintf("\n... (truncated, %d more chars)", len(output)-maxLen) + } + return output +} + +func readOptionalIntArg(args map[string]any, key string, minVal, maxVal int) (int, bool, error) { + raw, exists := args[key] + if !exists { + return 0, false, nil + } + + n, err := toInt(raw) + if err != nil { + return 0, true, fmt.Errorf("%s must be an integer", key) + } + if n < minVal || n > maxVal { + return 0, true, fmt.Errorf("%s must be between %d and %d", key, minVal, maxVal) + } + return n, true, nil +} + func (t *ExecTool) guardCommand(command, cwd string) string { cmd := strings.TrimSpace(command) lower := strings.ToLower(cmd) @@ -315,6 +583,10 @@ func (t *ExecTool) SetTimeout(timeout time.Duration) { t.timeout = timeout } +func (t *ExecTool) ProcessManager() *ProcessManager { + return t.processes +} + func (t *ExecTool) SetRestrictToWorkspace(restrict bool) { t.restrictToWorkspace = restrict } diff --git a/pkg/tools/subagent.go b/pkg/tools/subagent.go index ad371a649..960a41e5e 100644 --- a/pkg/tools/subagent.go +++ b/pkg/tools/subagent.go @@ -85,6 +85,23 @@ func (sm *SubagentManager) Spawn( task, label, agentID, originChannel, originChatID string, callback AsyncCallback, ) (string, error) { + subagentTask, err := sm.SpawnTask(ctx, task, label, agentID, originChannel, originChatID, callback) + if err != nil { + return "", err + } + + if label != "" { + return fmt.Sprintf("Spawned subagent '%s' for task: %s (id: %s)", label, task, subagentTask.ID), nil + } + return fmt.Sprintf("Spawned subagent for task: %s (id: %s)", task, subagentTask.ID), nil +} + +// SpawnTask starts a background subagent task and returns an immutable snapshot of the created task. +func (sm *SubagentManager) SpawnTask( + ctx context.Context, + task, label, agentID, originChannel, originChatID string, + callback AsyncCallback, +) (*SubagentTask, error) { sm.mu.Lock() defer sm.mu.Unlock() @@ -103,13 +120,11 @@ func (sm *SubagentManager) Spawn( } sm.tasks[taskID] = subagentTask - // Start task in background with context cancellation support + // Start task in background with context cancellation support. go sm.runTask(ctx, subagentTask, callback) - if label != "" { - return fmt.Sprintf("Spawned subagent '%s' for task: %s", label, task), nil - } - return fmt.Sprintf("Spawned subagent for task: %s", task), nil + snapshot := *subagentTask + return &snapshot, nil } func (sm *SubagentManager) runTask(ctx context.Context, task *SubagentTask, callback AsyncCallback) {