mirror of
https://github.com/sipeed/picoclaw.git
synced 2026-06-12 18:08:54 +00:00
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
This commit is contained in:
Liu Yuan
@@ -0,0 +1,252 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"errors"
|
||||
"io"
|
||||
"os"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
const maxOutputBufferSize = 100 * 1024 * 1024 // 100MB
|
||||
|
||||
const outputTruncateMarker = "\n... [output truncated, exceeded 100MB]\n"
|
||||
|
||||
// PtyKeyMode represents arrow key encoding mode for PTY sessions.
|
||||
// Programs send smkx/rmkx sequences to switch between CSI and SS3 modes.
|
||||
type PtyKeyMode uint8
|
||||
|
||||
const (
|
||||
PtyKeyModeCSI PtyKeyMode = iota // triggered by rmkx (\x1b[?1l)
|
||||
PtyKeyModeSS3 // triggered by smkx (\x1b[?1h)
|
||||
)
|
||||
|
||||
const PtyKeyModeNotFound PtyKeyMode = 255
|
||||
|
||||
var (
|
||||
ErrSessionNotFound = errors.New("session not found")
|
||||
ErrSessionDone = errors.New("session already completed")
|
||||
ErrPTYNotSupported = errors.New("PTY is not supported on this platform")
|
||||
ErrNoStdin = errors.New("no stdin available")
|
||||
)
|
||||
|
||||
type ProcessSession struct {
|
||||
mu sync.Mutex
|
||||
ID string
|
||||
PID int
|
||||
Command string
|
||||
PTY bool
|
||||
Background bool
|
||||
StartTime int64
|
||||
ExitCode int
|
||||
Status string
|
||||
stdinWriter io.Writer
|
||||
stdoutPipe io.Reader
|
||||
outputBuffer *bytes.Buffer
|
||||
outputTruncated bool
|
||||
ptyMaster *os.File
|
||||
|
||||
// ptyKeyMode tracks arrow key encoding mode (CSI vs SS3)
|
||||
ptyKeyMode PtyKeyMode
|
||||
}
|
||||
|
||||
func (s *ProcessSession) IsDone() bool {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return s.Status == "done" || s.Status == "exited"
|
||||
}
|
||||
|
||||
func (s *ProcessSession) GetPtyKeyMode() PtyKeyMode {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return s.ptyKeyMode
|
||||
}
|
||||
|
||||
func (s *ProcessSession) SetPtyKeyMode(mode PtyKeyMode) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.ptyKeyMode = mode
|
||||
}
|
||||
|
||||
func (s *ProcessSession) GetStatus() string {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return s.Status
|
||||
}
|
||||
|
||||
func (s *ProcessSession) SetStatus(status string) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.Status = status
|
||||
}
|
||||
|
||||
func (s *ProcessSession) GetExitCode() int {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return s.ExitCode
|
||||
}
|
||||
|
||||
func (s *ProcessSession) SetExitCode(code int) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.ExitCode = code
|
||||
}
|
||||
|
||||
func (s *ProcessSession) killProcess() error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
if s.Status != "running" {
|
||||
return ErrSessionDone
|
||||
}
|
||||
|
||||
pid := s.PID
|
||||
if pid <= 0 {
|
||||
return ErrSessionNotFound
|
||||
}
|
||||
|
||||
if err := killProcessGroup(pid); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
s.Status = "done"
|
||||
s.ExitCode = -1
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *ProcessSession) Kill() error {
|
||||
return s.killProcess()
|
||||
}
|
||||
|
||||
func (s *ProcessSession) Write(data string) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
if s.Status != "running" {
|
||||
return ErrSessionDone
|
||||
}
|
||||
|
||||
var writer io.Writer
|
||||
if s.PTY && s.ptyMaster != nil {
|
||||
writer = s.ptyMaster
|
||||
} else if s.stdinWriter != nil {
|
||||
writer = s.stdinWriter
|
||||
} else {
|
||||
return ErrNoStdin
|
||||
}
|
||||
|
||||
_, err := writer.Write([]byte(data))
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *ProcessSession) Read() string {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
if s.outputBuffer.Len() == 0 {
|
||||
return ""
|
||||
}
|
||||
|
||||
data := s.outputBuffer.String()
|
||||
s.outputBuffer.Reset()
|
||||
return data
|
||||
}
|
||||
|
||||
func (s *ProcessSession) ToSessionInfo() SessionInfo {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
return SessionInfo{
|
||||
ID: s.ID,
|
||||
Command: s.Command,
|
||||
Status: s.Status,
|
||||
PID: s.PID,
|
||||
StartedAt: s.StartTime,
|
||||
}
|
||||
}
|
||||
|
||||
type SessionManager struct {
|
||||
mu sync.RWMutex
|
||||
sessions map[string]*ProcessSession
|
||||
}
|
||||
|
||||
func NewSessionManager() *SessionManager {
|
||||
sm := &SessionManager{
|
||||
sessions: make(map[string]*ProcessSession),
|
||||
}
|
||||
|
||||
// Start cleaner goroutine - runs every 5 minutes, cleans up sessions done for >30 minutes
|
||||
go func() {
|
||||
ticker := time.NewTicker(5 * time.Minute)
|
||||
defer ticker.Stop()
|
||||
for range ticker.C {
|
||||
sm.cleanupOldSessions()
|
||||
}
|
||||
}()
|
||||
|
||||
return sm
|
||||
}
|
||||
|
||||
// cleanupOldSessions removes sessions that are done and older than 30 minutes
|
||||
func (sm *SessionManager) cleanupOldSessions() {
|
||||
sm.mu.Lock()
|
||||
defer sm.mu.Unlock()
|
||||
|
||||
cutoff := time.Now().Add(-30 * time.Minute)
|
||||
for id, session := range sm.sessions {
|
||||
if session.IsDone() && session.StartTime < cutoff.Unix() {
|
||||
delete(sm.sessions, id)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (sm *SessionManager) Add(session *ProcessSession) {
|
||||
sm.mu.Lock()
|
||||
defer sm.mu.Unlock()
|
||||
sm.sessions[session.ID] = session
|
||||
}
|
||||
|
||||
func (sm *SessionManager) Get(sessionID string) (*ProcessSession, error) {
|
||||
sm.mu.RLock()
|
||||
defer sm.mu.RUnlock()
|
||||
|
||||
session, ok := sm.sessions[sessionID]
|
||||
if !ok {
|
||||
return nil, ErrSessionNotFound
|
||||
}
|
||||
|
||||
return session, nil
|
||||
}
|
||||
|
||||
func (sm *SessionManager) Remove(sessionID string) {
|
||||
sm.mu.Lock()
|
||||
defer sm.mu.Unlock()
|
||||
delete(sm.sessions, sessionID)
|
||||
}
|
||||
|
||||
func (sm *SessionManager) List() []SessionInfo {
|
||||
sm.mu.RLock()
|
||||
defer sm.mu.RUnlock()
|
||||
|
||||
result := make([]SessionInfo, 0, len(sm.sessions))
|
||||
for _, session := range sm.sessions {
|
||||
result = append(result, session.ToSessionInfo())
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
func generateSessionID() string {
|
||||
return uuid.New().String()[:8]
|
||||
}
|
||||
|
||||
type SessionInfo struct {
|
||||
ID string `json:"id"`
|
||||
Command string `json:"command"`
|
||||
Status string `json:"status"`
|
||||
PID int `json:"pid"`
|
||||
StartedAt int64 `json:"startedAt"`
|
||||
}
|
||||
Reference in New Issue
Block a user