feat(agent): add session action tools and background process control
This commit is contained in:
parent
8a92c0dc93
commit
2943255c07
12 changed files with 2320 additions and 31 deletions
|
|
@ -51,12 +51,16 @@ func NewAgentInstance(
|
||||||
toolsRegistry.Register(tools.NewReadFileTool(workspace, restrict))
|
toolsRegistry.Register(tools.NewReadFileTool(workspace, restrict))
|
||||||
toolsRegistry.Register(tools.NewWriteFileTool(workspace, restrict))
|
toolsRegistry.Register(tools.NewWriteFileTool(workspace, restrict))
|
||||||
toolsRegistry.Register(tools.NewListDirTool(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.NewEditFileTool(workspace, restrict))
|
||||||
toolsRegistry.Register(tools.NewAppendFileTool(workspace, restrict))
|
toolsRegistry.Register(tools.NewAppendFileTool(workspace, restrict))
|
||||||
|
|
||||||
sessionsDir := filepath.Join(workspace, "sessions")
|
sessionsDir := filepath.Join(workspace, "sessions")
|
||||||
sessionsManager := session.NewSessionManager(sessionsDir)
|
sessionsManager := session.NewSessionManager(sessionsDir)
|
||||||
|
toolsRegistry.Register(tools.NewSessionsListTool(sessionsManager))
|
||||||
|
toolsRegistry.Register(tools.NewSessionsHistoryTool(sessionsManager))
|
||||||
|
|
||||||
contextBuilder := NewContextBuilder(workspace)
|
contextBuilder := NewContextBuilder(workspace)
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -55,9 +55,6 @@ type processOptions struct {
|
||||||
func NewAgentLoop(cfg *config.Config, msgBus *bus.MessageBus, provider providers.LLMProvider) *AgentLoop {
|
func NewAgentLoop(cfg *config.Config, msgBus *bus.MessageBus, provider providers.LLMProvider) *AgentLoop {
|
||||||
registry := NewAgentRegistry(cfg, provider)
|
registry := NewAgentRegistry(cfg, provider)
|
||||||
|
|
||||||
// Register shared tools to all agents
|
|
||||||
registerSharedTools(cfg, msgBus, registry, provider)
|
|
||||||
|
|
||||||
// Set up shared fallback chain
|
// Set up shared fallback chain
|
||||||
cooldown := providers.NewCooldownTracker()
|
cooldown := providers.NewCooldownTracker()
|
||||||
fallbackChain := providers.NewFallbackChain(cooldown)
|
fallbackChain := providers.NewFallbackChain(cooldown)
|
||||||
|
|
@ -69,7 +66,7 @@ func NewAgentLoop(cfg *config.Config, msgBus *bus.MessageBus, provider providers
|
||||||
stateManager = state.NewManager(defaultAgent.Workspace)
|
stateManager = state.NewManager(defaultAgent.Workspace)
|
||||||
}
|
}
|
||||||
|
|
||||||
return &AgentLoop{
|
al := &AgentLoop{
|
||||||
bus: msgBus,
|
bus: msgBus,
|
||||||
cfg: cfg,
|
cfg: cfg,
|
||||||
registry: registry,
|
registry: registry,
|
||||||
|
|
@ -77,6 +74,11 @@ func NewAgentLoop(cfg *config.Config, msgBus *bus.MessageBus, provider providers
|
||||||
summarizing: sync.Map{},
|
summarizing: sync.Map{},
|
||||||
fallback: fallbackChain,
|
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).
|
// registerSharedTools registers tools that are shared across all agents (web, message, spawn).
|
||||||
|
|
@ -85,6 +87,7 @@ func registerSharedTools(
|
||||||
msgBus *bus.MessageBus,
|
msgBus *bus.MessageBus,
|
||||||
registry *AgentRegistry,
|
registry *AgentRegistry,
|
||||||
provider providers.LLMProvider,
|
provider providers.LLMProvider,
|
||||||
|
sessionsExecutor tools.SessionsSendExecutor,
|
||||||
) {
|
) {
|
||||||
for _, agentID := range registry.ListAgentIDs() {
|
for _, agentID := range registry.ListAgentIDs() {
|
||||||
agent, ok := registry.GetAgent(agentID)
|
agent, ok := registry.GetAgent(agentID)
|
||||||
|
|
@ -140,15 +143,30 @@ func registerSharedTools(
|
||||||
agent.Tools.Register(tools.NewFindSkillsTool(registryMgr, searchCache))
|
agent.Tools.Register(tools.NewFindSkillsTool(registryMgr, searchCache))
|
||||||
agent.Tools.Register(tools.NewInstallSkillTool(registryMgr, agent.Workspace))
|
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 := tools.NewSubagentManager(provider, agent.Model, agent.Workspace, msgBus)
|
||||||
subagentManager.SetLLMOptions(agent.MaxTokens, agent.Temperature)
|
subagentManager.SetLLMOptions(agent.MaxTokens, agent.Temperature)
|
||||||
spawnTool := tools.NewSpawnTool(subagentManager)
|
spawnTool := tools.NewSpawnTool(subagentManager)
|
||||||
|
sessionsSpawnTool := tools.NewSessionsSpawnTool(subagentManager)
|
||||||
currentAgentID := agentID
|
currentAgentID := agentID
|
||||||
spawnTool.SetAllowlistChecker(func(targetAgentID string) bool {
|
allowlist := func(targetAgentID string) bool {
|
||||||
return registry.CanSpawnSubagent(currentAgentID, targetAgentID)
|
return registry.CanSpawnSubagent(currentAgentID, targetAgentID)
|
||||||
})
|
}
|
||||||
|
spawnTool.SetAllowlistChecker(allowlist)
|
||||||
|
sessionsSpawnTool.SetAllowlistChecker(allowlist)
|
||||||
agent.Tools.Register(spawnTool)
|
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)
|
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.
|
// ProcessHeartbeat processes a heartbeat request without session history.
|
||||||
// Each heartbeat is independent and doesn't accumulate context.
|
// Each heartbeat is independent and doesn't accumulate context.
|
||||||
func (al *AgentLoop) ProcessHeartbeat(ctx context.Context, content, channel, chatID string) (string, error) {
|
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)
|
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.
|
// maybeSummarize triggers summarization if the session history exceeds thresholds.
|
||||||
|
|
|
||||||
|
|
@ -4,6 +4,7 @@ import (
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
"os"
|
"os"
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
|
"sort"
|
||||||
"strings"
|
"strings"
|
||||||
"sync"
|
"sync"
|
||||||
"time"
|
"time"
|
||||||
|
|
@ -280,3 +281,58 @@ func (sm *SessionManager) SetHistory(key string, history []providers.Message) {
|
||||||
session.Updated = time.Now()
|
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
|
||||||
|
}
|
||||||
|
|
|
||||||
|
|
@ -4,6 +4,7 @@ import (
|
||||||
"os"
|
"os"
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
"testing"
|
"testing"
|
||||||
|
"time"
|
||||||
)
|
)
|
||||||
|
|
||||||
func TestSanitizeFilename(t *testing.T) {
|
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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
|
||||||
710
pkg/tools/process.go
Normal file
710
pkg/tools/process.go
Normal file
|
|
@ -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))
|
||||||
|
}
|
||||||
168
pkg/tools/process_test.go
Normal file
168
pkg/tools/process_test.go
Normal file
|
|
@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
445
pkg/tools/sessions.go
Normal file
445
pkg/tools/sessions.go
Normal file
|
|
@ -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")
|
||||||
|
}
|
||||||
|
}
|
||||||
239
pkg/tools/sessions_actions.go
Normal file
239
pkg/tools/sessions_actions.go
Normal file
|
|
@ -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))
|
||||||
|
}
|
||||||
138
pkg/tools/sessions_actions_test.go
Normal file
138
pkg/tools/sessions_actions_test.go
Normal file
|
|
@ -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")
|
||||||
|
}
|
||||||
|
}
|
||||||
131
pkg/tools/sessions_test.go
Normal file
131
pkg/tools/sessions_test.go
Normal file
|
|
@ -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")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -3,8 +3,10 @@ package tools
|
||||||
import (
|
import (
|
||||||
"bytes"
|
"bytes"
|
||||||
"context"
|
"context"
|
||||||
|
"encoding/json"
|
||||||
"errors"
|
"errors"
|
||||||
"fmt"
|
"fmt"
|
||||||
|
"io"
|
||||||
"os"
|
"os"
|
||||||
"os/exec"
|
"os/exec"
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
|
|
@ -22,6 +24,7 @@ type ExecTool struct {
|
||||||
denyPatterns []*regexp.Regexp
|
denyPatterns []*regexp.Regexp
|
||||||
allowPatterns []*regexp.Regexp
|
allowPatterns []*regexp.Regexp
|
||||||
restrictToWorkspace bool
|
restrictToWorkspace bool
|
||||||
|
processes *ProcessManager
|
||||||
}
|
}
|
||||||
|
|
||||||
var defaultDenyPatterns = []*regexp.Regexp{
|
var defaultDenyPatterns = []*regexp.Regexp{
|
||||||
|
|
@ -106,6 +109,7 @@ func NewExecToolWithConfig(workingDir string, restrict bool, config *config.Conf
|
||||||
denyPatterns: denyPatterns,
|
denyPatterns: denyPatterns,
|
||||||
allowPatterns: nil,
|
allowPatterns: nil,
|
||||||
restrictToWorkspace: restrict,
|
restrictToWorkspace: restrict,
|
||||||
|
processes: NewProcessManager(defaultProcessMaxOutputChars),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -129,6 +133,20 @@ func (t *ExecTool) Parameters() map[string]any {
|
||||||
"type": "string",
|
"type": "string",
|
||||||
"description": "Optional working directory for the command",
|
"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"},
|
"required": []string{"command"},
|
||||||
}
|
}
|
||||||
|
|
@ -164,22 +182,56 @@ func (t *ExecTool) Execute(ctx context.Context, args map[string]any) *ToolResult
|
||||||
return ErrorResult(guardError)
|
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 cmdCtx context.Context
|
||||||
var cancel context.CancelFunc
|
var cancel context.CancelFunc
|
||||||
if t.timeout > 0 {
|
if timeout > 0 {
|
||||||
cmdCtx, cancel = context.WithTimeout(ctx, t.timeout)
|
cmdCtx, cancel = context.WithTimeout(ctx, timeout)
|
||||||
} else {
|
} else {
|
||||||
cmdCtx, cancel = context.WithCancel(ctx)
|
cmdCtx, cancel = context.WithCancel(ctx)
|
||||||
}
|
}
|
||||||
defer cancel()
|
defer cancel()
|
||||||
|
|
||||||
var cmd *exec.Cmd
|
cmd := shellCommand(cmdCtx, command)
|
||||||
if runtime.GOOS == "windows" {
|
|
||||||
cmd = exec.CommandContext(cmdCtx, "powershell", "-NoProfile", "-NonInteractive", "-Command", command)
|
|
||||||
} else {
|
|
||||||
cmd = exec.CommandContext(cmdCtx, "sh", "-c", command)
|
|
||||||
}
|
|
||||||
if cwd != "" {
|
if cwd != "" {
|
||||||
cmd.Dir = cwd
|
cmd.Dir = cwd
|
||||||
}
|
}
|
||||||
|
|
@ -221,7 +273,7 @@ func (t *ExecTool) Execute(ctx context.Context, args map[string]any) *ToolResult
|
||||||
|
|
||||||
if err != nil {
|
if err != nil {
|
||||||
if errors.Is(cmdCtx.Err(), context.DeadlineExceeded) {
|
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{
|
return &ToolResult{
|
||||||
ForLLM: msg,
|
ForLLM: msg,
|
||||||
ForUser: 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)
|
output += fmt.Sprintf("\nExit code: %v", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
if output == "" {
|
output = truncateExecOutput(output)
|
||||||
output = "(no output)"
|
|
||||||
}
|
|
||||||
|
|
||||||
maxLen := 10000
|
|
||||||
if len(output) > maxLen {
|
|
||||||
output = output[:maxLen] + fmt.Sprintf("\n... (truncated, %d more chars)", len(output)-maxLen)
|
|
||||||
}
|
|
||||||
|
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return &ToolResult{
|
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 {
|
func (t *ExecTool) guardCommand(command, cwd string) string {
|
||||||
cmd := strings.TrimSpace(command)
|
cmd := strings.TrimSpace(command)
|
||||||
lower := strings.ToLower(cmd)
|
lower := strings.ToLower(cmd)
|
||||||
|
|
@ -315,6 +583,10 @@ func (t *ExecTool) SetTimeout(timeout time.Duration) {
|
||||||
t.timeout = timeout
|
t.timeout = timeout
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (t *ExecTool) ProcessManager() *ProcessManager {
|
||||||
|
return t.processes
|
||||||
|
}
|
||||||
|
|
||||||
func (t *ExecTool) SetRestrictToWorkspace(restrict bool) {
|
func (t *ExecTool) SetRestrictToWorkspace(restrict bool) {
|
||||||
t.restrictToWorkspace = restrict
|
t.restrictToWorkspace = restrict
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -85,6 +85,23 @@ func (sm *SubagentManager) Spawn(
|
||||||
task, label, agentID, originChannel, originChatID string,
|
task, label, agentID, originChannel, originChatID string,
|
||||||
callback AsyncCallback,
|
callback AsyncCallback,
|
||||||
) (string, error) {
|
) (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()
|
sm.mu.Lock()
|
||||||
defer sm.mu.Unlock()
|
defer sm.mu.Unlock()
|
||||||
|
|
||||||
|
|
@ -103,13 +120,11 @@ func (sm *SubagentManager) Spawn(
|
||||||
}
|
}
|
||||||
sm.tasks[taskID] = subagentTask
|
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)
|
go sm.runTask(ctx, subagentTask, callback)
|
||||||
|
|
||||||
if label != "" {
|
snapshot := *subagentTask
|
||||||
return fmt.Sprintf("Spawned subagent '%s' for task: %s", label, task), nil
|
return &snapshot, nil
|
||||||
}
|
|
||||||
return fmt.Sprintf("Spawned subagent for task: %s", task), nil
|
|
||||||
}
|
}
|
||||||
|
|
||||||
func (sm *SubagentManager) runTask(ctx context.Context, task *SubagentTask, callback AsyncCallback) {
|
func (sm *SubagentManager) runTask(ctx context.Context, task *SubagentTask, callback AsyncCallback) {
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue