Skip to content

Commit f901af8

Browse files
liuyliuy
authored
feat(tools): add exec tool enhancement with background execution and PTY support (#1752)
- Unified exec tool with actions: run/list/poll/read/write/send-keys/kill - PTY support using creack/pty library - Process session management with background execution - Process group kill for cleaning up child processes - Session cleanup: 30-minute TTL for old sessions - Output buffer: 100MB limit with truncation Actions: - run: execute command (sync or background) - list: list all sessions - poll: check session status - read: read session output - write: send input to session stdin - send-keys: send special keys (up, down, ctrl-c, enter, etc.) - kill: terminate session Tests: - PTY: allowed commands, write/read, poll, kill, process group kill - Non-PTY: background execution, list, read, write, poll, kill, process group kill - Session management: add/get/remove/list/cleanup
1 parent 6148ccc commit f901af8

File tree

11 files changed

+2082
-31
lines changed

11 files changed

+2082
-31
lines changed

‎go.mod‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,13 @@ module github.com/sipeed/picoclaw
33
go 1.25.8
44

55
require (
6-
github.com/BurntSushi/toml v1.6.0
76
fyne.io/systray v1.12.0
7+
github.com/BurntSushi/toml v1.6.0
88
github.com/adhocore/gronx v1.19.6
99
github.com/anthropics/anthropic-sdk-go v1.26.0
1010
github.com/bwmarrin/discordgo v0.29.0
1111
github.com/caarlos0/env/v11 v11.4.0
12+
github.com/creack/pty v1.1.9
1213
github.com/ergochat/irc-go v0.6.0
1314
github.com/ergochat/readline v0.1.3
1415
github.com/gdamore/tcell/v2 v2.13.8

‎go.sum‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ github.com/coder/websocket v1.8.14 h1:9L0p0iKiNOibykf283eHkKUHHrpG7f65OE3BhhO7v9
3737
github.com/coder/websocket v1.8.14/go.mod h1:NX3SzP+inril6yawo5CQXx8+fk145lPDC6pumgx0mVg=
3838
github.com/coreos/go-systemd/v22 v22.5.0/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc=
3939
github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g=
40+
github.com/creack/pty v1.1.9 h1:uDmaGzcdjhF4i/plgjmEsriH11Y0o7RKapEf/LDaM3w=
4041
github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
4142
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
4243
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=

‎pkg/agent/instance_test.go‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -236,8 +236,9 @@ func TestNewAgentInstance_AllowsMediaTempDirForReadListAndExec(t *testing.T) {
236236
t.Fatal("exec tool not registered")
237237
}
238238
execResult := execTool.Execute(context.Background(), map[string]any{
239-
"command": "cat " + filepath.Base(mediaPath),
240-
"working_dir": mediaDir,
239+
"action": "run",
240+
"command": "cat " + filepath.Base(mediaPath),
241+
"cwd": mediaDir,
241242
})
242243
if execResult.IsError {
243244
t.Fatalf("exec should allow media temp dir, got: %s", execResult.ForLLM)

‎pkg/tools/session.go‎

Lines changed: 252 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,252 @@
1+
package tools
2+
3+
import (
4+
"bytes"
5+
"errors"
6+
"io"
7+
"os"
8+
"sync"
9+
"time"
10+
11+
"github.com/google/uuid"
12+
)
13+
14+
const maxOutputBufferSize = 100 * 1024 * 1024 // 100MB
15+
16+
const outputTruncateMarker = "\n... [output truncated, exceeded 100MB]\n"
17+
18+
// PtyKeyMode represents arrow key encoding mode for PTY sessions.
19+
// Programs send smkx/rmkx sequences to switch between CSI and SS3 modes.
20+
type PtyKeyMode uint8
21+
22+
const (
23+
PtyKeyModeCSI PtyKeyMode = iota // triggered by rmkx (\x1b[?1l)
24+
PtyKeyModeSS3 // triggered by smkx (\x1b[?1h)
25+
)
26+
27+
const PtyKeyModeNotFound PtyKeyMode = 255
28+
29+
var (
30+
ErrSessionNotFound = errors.New("session not found")
31+
ErrSessionDone = errors.New("session already completed")
32+
ErrPTYNotSupported = errors.New("PTY is not supported on this platform")
33+
ErrNoStdin = errors.New("no stdin available")
34+
)
35+
36+
type ProcessSession struct {
37+
mu sync.Mutex
38+
ID string
39+
PID int
40+
Command string
41+
PTY bool
42+
Background bool
43+
StartTime int64
44+
ExitCode int
45+
Status string
46+
stdinWriter io.Writer
47+
stdoutPipe io.Reader
48+
outputBuffer *bytes.Buffer
49+
outputTruncated bool
50+
ptyMaster *os.File
51+
52+
// ptyKeyMode tracks arrow key encoding mode (CSI vs SS3)
53+
ptyKeyMode PtyKeyMode
54+
}
55+
56+
func (s *ProcessSession) IsDone() bool {
57+
s.mu.Lock()
58+
defer s.mu.Unlock()
59+
return s.Status == "done" || s.Status == "exited"
60+
}
61+
62+
func (s *ProcessSession) GetPtyKeyMode() PtyKeyMode {
63+
s.mu.Lock()
64+
defer s.mu.Unlock()
65+
return s.ptyKeyMode
66+
}
67+
68+
func (s *ProcessSession) SetPtyKeyMode(mode PtyKeyMode) {
69+
s.mu.Lock()
70+
defer s.mu.Unlock()
71+
s.ptyKeyMode = mode
72+
}
73+
74+
func (s *ProcessSession) GetStatus() string {
75+
s.mu.Lock()
76+
defer s.mu.Unlock()
77+
return s.Status
78+
}
79+
80+
func (s *ProcessSession) SetStatus(status string) {
81+
s.mu.Lock()
82+
defer s.mu.Unlock()
83+
s.Status = status
84+
}
85+
86+
func (s *ProcessSession) GetExitCode() int {
87+
s.mu.Lock()
88+
defer s.mu.Unlock()
89+
return s.ExitCode
90+
}
91+
92+
func (s *ProcessSession) SetExitCode(code int) {
93+
s.mu.Lock()
94+
defer s.mu.Unlock()
95+
s.ExitCode = code
96+
}
97+
98+
func (s *ProcessSession) killProcess() error {
99+
s.mu.Lock()
100+
defer s.mu.Unlock()
101+
102+
if s.Status != "running" {
103+
return ErrSessionDone
104+
}
105+
106+
pid := s.PID
107+
if pid <= 0 {
108+
return ErrSessionNotFound
109+
}
110+
111+
if err := killProcessGroup(pid); err != nil {
112+
return err
113+
}
114+
115+
s.Status = "done"
116+
s.ExitCode = -1
117+
return nil
118+
}
119+
120+
func (s *ProcessSession) Kill() error {
121+
return s.killProcess()
122+
}
123+
124+
func (s *ProcessSession) Write(data string) error {
125+
s.mu.Lock()
126+
defer s.mu.Unlock()
127+
128+
if s.Status != "running" {
129+
return ErrSessionDone
130+
}
131+
132+
var writer io.Writer
133+
if s.PTY && s.ptyMaster != nil {
134+
writer = s.ptyMaster
135+
} else if s.stdinWriter != nil {
136+
writer = s.stdinWriter
137+
} else {
138+
return ErrNoStdin
139+
}
140+
141+
_, err := writer.Write([]byte(data))
142+
return err
143+
}
144+
145+
func (s *ProcessSession) Read() string {
146+
s.mu.Lock()
147+
defer s.mu.Unlock()
148+
149+
if s.outputBuffer.Len() == 0 {
150+
return ""
151+
}
152+
153+
data := s.outputBuffer.String()
154+
s.outputBuffer.Reset()
155+
return data
156+
}
157+
158+
func (s *ProcessSession) ToSessionInfo() SessionInfo {
159+
s.mu.Lock()
160+
defer s.mu.Unlock()
161+
162+
return SessionInfo{
163+
ID: s.ID,
164+
Command: s.Command,
165+
Status: s.Status,
166+
PID: s.PID,
167+
StartedAt: s.StartTime,
168+
}
169+
}
170+
171+
type SessionManager struct {
172+
mu sync.RWMutex
173+
sessions map[string]*ProcessSession
174+
}
175+
176+
func NewSessionManager() *SessionManager {
177+
sm := &SessionManager{
178+
sessions: make(map[string]*ProcessSession),
179+
}
180+
181+
// Start cleaner goroutine - runs every 5 minutes, cleans up sessions done for >30 minutes
182+
go func() {
183+
ticker := time.NewTicker(5 * time.Minute)
184+
defer ticker.Stop()
185+
for range ticker.C {
186+
sm.cleanupOldSessions()
187+
}
188+
}()
189+
190+
return sm
191+
}
192+
193+
// cleanupOldSessions removes sessions that are done and older than 30 minutes
194+
func (sm *SessionManager) cleanupOldSessions() {
195+
sm.mu.Lock()
196+
defer sm.mu.Unlock()
197+
198+
cutoff := time.Now().Add(-30 * time.Minute)
199+
for id, session := range sm.sessions {
200+
if session.IsDone() && session.StartTime < cutoff.Unix() {
201+
delete(sm.sessions, id)
202+
}
203+
}
204+
}
205+
206+
func (sm *SessionManager) Add(session *ProcessSession) {
207+
sm.mu.Lock()
208+
defer sm.mu.Unlock()
209+
sm.sessions[session.ID] = session
210+
}
211+
212+
func (sm *SessionManager) Get(sessionID string) (*ProcessSession, error) {
213+
sm.mu.RLock()
214+
defer sm.mu.RUnlock()
215+
216+
session, ok := sm.sessions[sessionID]
217+
if !ok {
218+
return nil, ErrSessionNotFound
219+
}
220+
221+
return session, nil
222+
}
223+
224+
func (sm *SessionManager) Remove(sessionID string) {
225+
sm.mu.Lock()
226+
defer sm.mu.Unlock()
227+
delete(sm.sessions, sessionID)
228+
}
229+
230+
func (sm *SessionManager) List() []SessionInfo {
231+
sm.mu.RLock()
232+
defer sm.mu.RUnlock()
233+
234+
result := make([]SessionInfo, 0, len(sm.sessions))
235+
for _, session := range sm.sessions {
236+
result = append(result, session.ToSessionInfo())
237+
}
238+
239+
return result
240+
}
241+
242+
func generateSessionID() string {
243+
return uuid.New().String()[:8]
244+
}
245+
246+
type SessionInfo struct {
247+
ID string `json:"id"`
248+
Command string `json:"command"`
249+
Status string `json:"status"`
250+
PID int `json:"pid"`
251+
StartedAt int64 `json:"startedAt"`
252+
}

‎pkg/tools/session_process_unix.go‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
//go:build !windows
2+
3+
package tools
4+
5+
import (
6+
"syscall"
7+
)
8+
9+
func killProcessGroup(pid int) error {
10+
if err := syscall.Kill(-pid, syscall.SIGKILL); err != nil {
11+
_ = syscall.Kill(pid, syscall.SIGKILL)
12+
}
13+
return nil
14+
}

‎pkg/tools/session_process_windows.go‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
//go:build windows
2+
3+
package tools
4+
5+
import (
6+
"os/exec"
7+
"strconv"
8+
)
9+
10+
func killProcessGroup(pid int) error {
11+
_ = exec.Command("taskkill", "/T", "/F", "/PID", strconv.Itoa(pid)).Run()
12+
return nil
13+
}

0 commit comments

Comments
 (0)