diff --git a/README.md b/README.md new file mode 100644 index 0000000..11cf0dd --- /dev/null +++ b/README.md @@ -0,0 +1,127 @@ +# PaddedCell + +OpenClaw 插件:安全密码管理 + 安全执行 + 安全重启 + +## 功能模块 + +### 1. pass_mgr - 密码管理二进制 (Go) + +使用 AES-256-GCM 加密,每个 agent 基于公私钥进行加/解密。 + +```bash +# 初始化 +pass_mgr admin init [--key-path ] + +# 获取密码 +pass_mgr get [--username] + +# 生成密码 (agent 可用) +pass_mgr generate [--username ] + +# 设置密码 (仅人类) +pass_mgr set [--username ] + +# 删除密码 +pass_mgr unset + +# 轮换密码 +pass_mgr rotate +``` + +**安全特性:** +- Agent 无法执行 `set` 操作(通过环境变量检测) +- 未初始化前所有操作报错 +- Admin 密码泄露检测(监控 message/tool calling) + +### 2. pcexec - 安全执行工具 (TypeScript) + +与 OpenClaw 原生 exec 接口一致,自动处理 pass_mgr get 并脱敏输出。 + +```typescript +import { pcexec } from 'pcexec'; + +const result = await pcexec('echo $(pass_mgr get mypassword)', { + cwd: '/workspace', + timeout: 30000, +}); +// result.stdout 中密码会被替换为 ###### +``` + +### 3. safe-restart - 安全重启模块 (TypeScript) + +提供 agent 状态管理和协调重启。 + +**Agent 状态:** +- `idle` - 空闲 +- `busy` - 处理消息中 +- `focus` - 专注模式(工作流) +- `freeze` - 冻结(不接受新消息) +- `pre-freeze` - 准备冻结 +- `pre-freeze-focus` - 准备冻结(专注模式) + +**API:** +- `POST /query-restart` - 查询重启就绪状态 +- `POST /restart-result` - 报告重启结果 +- `GET /status` - 获取所有状态 + +**Slash 命令:** +``` +/padded-cell-ctrl status +/padded-cell-ctrl enable pass-mgr|safe-restart +/padded-cell-ctrl disable pass-mgr|safe-restart +``` + +## 目录结构 + +``` +PaddedCell/ +├── pass_mgr/ # Go 密码管理二进制 +│ ├── src/ +│ │ └── main.go +│ └── go.mod +├── pcexec/ # TypeScript 安全执行工具 +│ ├── src/ +│ │ └── index.ts +│ ├── package.json +│ └── tsconfig.json +├── safe-restart/ # TypeScript 安全重启模块 +│ ├── src/ +│ │ ├── index.ts +│ │ ├── status-manager.ts +│ │ ├── api.ts +│ │ ├── safe-restart.ts +│ │ └── slash-commands.ts +│ ├── package.json +│ └── tsconfig.json +├── docs/ # 文档 +├── PROJECT_PLAN.md # 项目计划 +├── AGENT_TASKS.md # 任务清单 +└── README.md +``` + +## 开发 + +```bash +# Clone +git clone https://git.hangman-lab.top/nav/PaddedCell.git +cd PaddedCell + +# Build pass_mgr (需要 Go) +cd pass_mgr +go build -o pass_mgr src/main.go + +# Build pcexec (需要 Node.js) +cd ../pcexec +npm install +npm run build + +# Build safe-restart (需要 Node.js) +cd ../safe-restart +npm install +npm run build +``` + +## 开发分支 + +- `main` - 主分支 +- `dev/zhi` - 当前开发分支 diff --git a/pass_mgr/go.mod b/pass_mgr/go.mod new file mode 100644 index 0000000..47a663c --- /dev/null +++ b/pass_mgr/go.mod @@ -0,0 +1,12 @@ +module pass_mgr + +go 1.22 + +require ( + github.com/spf13/cobra v1.8.0 +) + +require ( + github.com/inconshreveable/mousetrap v1.1.0 // indirect + github.com/spf13/pflag v1.0.5 // indirect +) diff --git a/pass_mgr/src/main.go b/pass_mgr/src/main.go new file mode 100644 index 0000000..7c196e4 --- /dev/null +++ b/pass_mgr/src/main.go @@ -0,0 +1,498 @@ +package main + +import ( + "crypto/aes" + "crypto/cipher" + "crypto/rand" + "crypto/sha256" + "encoding/base64" + "encoding/json" + "fmt" + "io" + "os" + "path/filepath" + "strings" + "time" + + "github.com/spf13/cobra" +) + +const ( + DefaultAlgorithm = "AES-256-GCM" + AdminKeyDir = ".pass_mgr" + AdminKeyFile = ".priv" + SecretsDirName = ".secrets" +) + +// EncryptedData represents the structure of encrypted password file +type EncryptedData struct { + Algorithm string `json:"algorithm"` + Nonce string `json:"nonce"` + Data string `json:"data"` + User string `json:"user,omitempty"` +} + +// Config holds admin key configuration +type Config struct { + KeyHash string `json:"key_hash"` + Algorithm string `json:"algorithm"` +} + +var ( + workspaceDir string + agentID string + username string +) + +func main() { + rootCmd := &cobra.Command{ + Use: "pass_mgr", + Short: "Password manager for OpenClaw agents", + Long: `A secure password management tool using AES-256-GCM encryption.`, + } + + // Get environment variables + workspaceDir = os.Getenv("AGENT_WORKSPACE") + agentID = os.Getenv("AGENT_ID") + + // Commands + rootCmd.AddCommand(getCmd()) + rootCmd.AddCommand(generateCmd()) + rootCmd.AddCommand(unsetCmd()) + rootCmd.AddCommand(rotateCmd()) + rootCmd.AddCommand(adminInitCmd()) + rootCmd.AddCommand(setCmd()) + + if err := rootCmd.Execute(); err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + os.Exit(1) + } +} + +func getCmd() *cobra.Command { + return &cobra.Command{ + Use: "get [key]", + Short: "Get password for a key", + Args: cobra.ExactArgs(1), + Run: func(cmd *cobra.Command, args []string) { + key := args[0] + password, user, err := getPassword(key) + if err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + os.Exit(1) + } + if username { + fmt.Println(user) + } else { + fmt.Println(password) + } + }, + } +} + +func generateCmd() *cobra.Command { + var user string + cmd := &cobra.Command{ + Use: "generate [key]", + Short: "Generate a new password", + Args: cobra.ExactArgs(1), + Run: func(cmd *cobra.Command, args []string) { + key := args[0] + // Check if agent is trying to set password + if os.Getenv("AGENT") != "" || os.Getenv("AGENT_WORKSPACE") != "" { + fmt.Fprintln(os.Stderr, "Error: Agents cannot set passwords. Use generate instead.") + os.Exit(1) + } + password, err := generatePassword(32) + if err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + os.Exit(1) + } + if err := setPassword(key, user, password); err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + os.Exit(1) + } + fmt.Println(password) + }, + } + cmd.Flags().StringVar(&user, "username", "", "Username associated with the password") + return cmd +} + +func unsetCmd() *cobra.Command { + return &cobra.Command{ + Use: "unset [key]", + Short: "Remove a password", + Args: cobra.ExactArgs(1), + Run: func(cmd *cobra.Command, args []string) { + key := args[0] + if err := removePassword(key); err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + os.Exit(1) + } + }, + } +} + +func rotateCmd() *cobra.Command { + return &cobra.Command{ + Use: "rotate [key]", + Short: "Rotate password for a key", + Args: cobra.ExactArgs(1), + Run: func(cmd *cobra.Command, args []string) { + key := args[0] + // Check if initialized + if !isInitialized() { + fmt.Fprintln(os.Stderr, "Error: pass_mgr not initialized. Run 'pass_mgr admin init' first.") + os.Exit(1) + } + + // Get current user if exists + _, user, err := getPassword(key) + if err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + os.Exit(1) + } + + // Generate new password + newPassword, err := generatePassword(32) + if err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + os.Exit(1) + } + + if err := setPassword(key, user, newPassword); err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + os.Exit(1) + } + fmt.Println(newPassword) + }, + } +} + +func adminInitCmd() *cobra.Command { + var keyPath string + cmd := &cobra.Command{ + Use: "admin init", + Short: "Initialize pass_mgr with admin key", + Run: func(cmd *cobra.Command, args []string) { + if err := initAdmin(keyPath); err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + os.Exit(1) + } + fmt.Println("pass_mgr initialized successfully") + }, + } + cmd.Flags().StringVar(&keyPath, "key-path", "", "Path to admin key file (optional)") + return cmd +} + +func setCmd() *cobra.Command { + var user string + cmd := &cobra.Command{ + Use: "set [key] [password]", + Short: "Set password (admin only)", + Args: cobra.ExactArgs(2), + Run: func(cmd *cobra.Command, args []string) { + // Check if agent is trying to set password + if os.Getenv("AGENT") != "" || os.Getenv("AGENT_WORKSPACE") != "" { + fmt.Fprintln(os.Stderr, "Error: Agents cannot set passwords. Only humans can use 'set'.") + os.Exit(1) + } + + key := args[0] + password := args[1] + + if err := setPassword(key, user, password); err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + os.Exit(1) + } + }, + } + cmd.Flags().StringVar(&user, "username", "", "Username associated with the password") + return cmd +} + +// Helper functions + +func getHomeDir() string { + home, err := os.UserHomeDir() + if err != nil { + return "." + } + return home +} + +func getAdminKeyPath() string { + return filepath.Join(getHomeDir(), AdminKeyDir, AdminKeyFile) +} + +func getConfigPath() string { + return filepath.Join(getHomeDir(), AdminKeyDir, "config.json") +} + +func isInitialized() bool { + _, err := os.Stat(getConfigPath()) + return err == nil +} + +func loadAdminKey() ([]byte, error) { + keyPath := getAdminKeyPath() + key, err := os.ReadFile(keyPath) + if err != nil { + return nil, fmt.Errorf("failed to load admin key: %w", err) + } + // Hash the key to get 32 bytes for AES-256 + hash := sha256.Sum256(key) + return hash[:], nil +} + +func initAdmin(keyPath string) error { + homeDir := getHomeDir() + adminDir := filepath.Join(homeDir, AdminKeyDir) + + // Create admin directory + if err := os.MkdirAll(adminDir, 0700); err != nil { + return fmt.Errorf("failed to create admin directory: %w", err) + } + + var key []byte + if keyPath != "" { + // Read provided key + var err error + key, err = os.ReadFile(keyPath) + if err != nil { + return fmt.Errorf("failed to read key file: %w", err) + } + } else { + // Generate new key + key = make([]byte, 32) + if _, err := rand.Read(key); err != nil { + return fmt.Errorf("failed to generate key: %w", err) + } + } + + // Save key + keyFile := filepath.Join(adminDir, AdminKeyFile) + if err := os.WriteFile(keyFile, key, 0600); err != nil { + return fmt.Errorf("failed to save key: %w", err) + } + + // Save config + config := Config{ + KeyHash: fmt.Sprintf("%x", sha256.Sum256(key)), + Algorithm: DefaultAlgorithm, + } + configData, _ := json.MarshalIndent(config, "", " ") + configPath := filepath.Join(adminDir, "config.json") + if err := os.WriteFile(configPath, configData, 0600); err != nil { + return fmt.Errorf("failed to save config: %w", err) + } + + return nil +} + +func getSecretsDir() string { + if workspaceDir != "" && agentID != "" { + return filepath.Join(workspaceDir, SecretsDirName, agentID) + } + // Fallback to home directory + return filepath.Join(getHomeDir(), SecretsDirName, "default") +} + +func getPasswordFilePath(key string) string { + return filepath.Join(getSecretsDir(), key+".gpg") +} + +func encrypt(plaintext []byte, key []byte) (*EncryptedData, error) { + block, err := aes.NewCipher(key) + if err != nil { + return nil, err + } + + gcm, err := cipher.NewGCM(block) + if err != nil { + return nil, err + } + + nonce := make([]byte, gcm.NonceSize()) + if _, err := io.ReadFull(rand.Reader, nonce); err != nil { + return nil, err + } + + ciphertext := gcm.Seal(nonce, nonce, plaintext, nil) + + return &EncryptedData{ + Algorithm: DefaultAlgorithm, + Nonce: base64.StdEncoding.EncodeToString(nonce), + Data: base64.StdEncoding.EncodeToString(ciphertext[gcm.NonceSize():]), + }, nil +} + +func decrypt(data *EncryptedData, key []byte) ([]byte, error) { + ciphertext, err := base64.StdEncoding.DecodeString(data.Data) + if err != nil { + return nil, err + } + + nonce, err := base64.StdEncoding.DecodeString(data.Nonce) + if err != nil { + return nil, err + } + + block, err := aes.NewCipher(key) + if err != nil { + return nil, err + } + + gcm, err := cipher.NewGCM(block) + if err != nil { + return nil, err + } + + plaintext, err := gcm.Open(nil, nonce, ciphertext, nil) + if err != nil { + return nil, err + } + + return plaintext, nil +} + +func setPassword(key, user, password string) error { + if !isInitialized() { + return fmt.Errorf("pass_mgr not initialized. Run 'pass_mgr admin init' first") + } + + adminKey, err := loadAdminKey() + if err != nil { + return err + } + + // Create secrets directory + secretsDir := getSecretsDir() + if err := os.MkdirAll(secretsDir, 0700); err != nil { + return fmt.Errorf("failed to create secrets directory: %w", err) + } + + // Encrypt password + data := map[string]string{ + "password": password, + "user": user, + } + plaintext, _ := json.Marshal(data) + + encrypted, err := encrypt(plaintext, adminKey) + if err != nil { + return fmt.Errorf("failed to encrypt: %w", err) + } + encrypted.User = user + + // Save to file + filePath := getPasswordFilePath(key) + fileData, _ := json.MarshalIndent(encrypted, "", " ") + if err := os.WriteFile(filePath, fileData, 0600); err != nil { + return fmt.Errorf("failed to save password: %w", err) + } + + return nil +} + +func getPassword(key string) (string, string, error) { + if !isInitialized() { + return "", "", fmt.Errorf("pass_mgr not initialized. Run 'pass_mgr admin init' first") + } + + adminKey, err := loadAdminKey() + if err != nil { + return "", "", err + } + + filePath := getPasswordFilePath(key) + fileData, err := os.ReadFile(filePath) + if err != nil { + return "", "", fmt.Errorf("password not found: %w", err) + } + + var encrypted EncryptedData + if err := json.Unmarshal(fileData, &encrypted); err != nil { + return "", "", fmt.Errorf("failed to parse password file: %w", err) + } + + plaintext, err := decrypt(&encrypted, adminKey) + if err != nil { + return "", "", fmt.Errorf("failed to decrypt: %w", err) + } + + var data map[string]string + if err := json.Unmarshal(plaintext, &data); err != nil { + return "", "", fmt.Errorf("failed to parse decrypted data: %w", err) + } + + return data["password"], data["user"], nil +} + +func removePassword(key string) error { + if !isInitialized() { + return fmt.Errorf("pass_mgr not initialized. Run 'pass_mgr admin init' first") + } + + filePath := getPasswordFilePath(key) + if err := os.Remove(filePath); err != nil { + return fmt.Errorf("failed to remove password: %w", err) + } + return nil +} + +func generatePassword(length int) (string, error) { + const charset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!@#$%^&*" + password := make([]byte, length) + for i := range password { + randomByte := make([]byte, 1) + if _, err := rand.Read(randomByte); err != nil { + return "", err + } + password[i] = charset[int(randomByte[0])%len(charset)] + } + return string(password), nil +} + +// CheckForAdminLeak checks if admin password appears in message/tool calling +func CheckForAdminLeak(content string) bool { + // This is a placeholder - actual implementation should check against actual admin password + // This function should be called by the plugin to monitor messages + configPath := getConfigPath() + if _, err := os.Stat(configPath); err != nil { + return false + } + + // TODO: Implement actual leak detection + // For now, just check if content contains common patterns + return strings.Contains(content, "admin") && strings.Contains(content, "password") +} + +// ResetOnLeak resets pass_mgr to uninitialized state and logs security breach +func ResetOnLeak() error { + configPath := getConfigPath() + + // Remove config (but keep key file for potential recovery) + if err := os.Remove(configPath); err != nil { + return err + } + + // Log security breach + logPath := filepath.Join(getHomeDir(), AdminKeyDir, "security_breach.log") + logEntry := fmt.Sprintf("[%s] CRITICAL: Admin password leaked! pass_mgr reset to uninitialized state.\n", + time.Now().Format(time.RFC3339)) + + f, err := os.OpenFile(logPath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0600) + if err != nil { + return err + } + defer f.Close() + + if _, err := f.WriteString(logEntry); err != nil { + return err + } + + return nil +} diff --git a/pcexec/package.json b/pcexec/package.json new file mode 100644 index 0000000..f9416bf --- /dev/null +++ b/pcexec/package.json @@ -0,0 +1,21 @@ +{ + "name": "pcexec", + "version": "0.1.0", + "description": "Safe exec wrapper for OpenClaw with password sanitization", + "main": "dist/index.js", + "types": "dist/index.d.ts", + "scripts": { + "build": "tsc", + "dev": "tsc --watch", + "test": "jest" + }, + "dependencies": { + "@types/node": "^20.0.0" + }, + "devDependencies": { + "typescript": "^5.0.0", + "jest": "^29.0.0", + "@types/jest": "^29.0.0", + "ts-jest": "^29.0.0" + } +} diff --git a/pcexec/src/index.ts b/pcexec/src/index.ts new file mode 100644 index 0000000..4df3424 --- /dev/null +++ b/pcexec/src/index.ts @@ -0,0 +1,346 @@ +import { spawn, SpawnOptions } from 'child_process'; +import { promisify } from 'util'; + +const execAsync = promisify(require('child_process').exec); + +export interface PcExecOptions { + /** Current working directory */ + cwd?: string; + /** Environment variables */ + env?: Record; + /** Timeout in milliseconds */ + timeout?: number; + /** Maximum buffer size for stdout/stderr */ + maxBuffer?: number; + /** Kill signal */ + killSignal?: NodeJS.Signals; + /** Shell to use */ + shell?: string | boolean; + /** UID to run as */ + uid?: number; + /** GID to run as */ + gid?: number; + /** Window style (Windows only) */ + windowsHide?: boolean; +} + +export interface PcExecResult { + /** Standard output */ + stdout: string; + /** Standard error */ + stderr: string; + /** Exit code */ + exitCode: number; + /** Command that was executed */ + command: string; +} + +export interface PcExecError extends Error { + /** Exit code */ + code?: number; + /** Signal that terminated the process */ + signal?: string; + /** Standard output */ + stdout: string; + /** Standard error */ + stderr: string; + /** Killed by timeout */ + killed?: boolean; +} + +/** + * Extract pass_mgr get commands from a command string + * Supports formats like: + * - $(pass_mgr get key) + * - `pass_mgr get key` + * - pass_mgr get key (direct invocation) + */ +function extractPassMgrGets(command: string): Array<{ key: string; fullMatch: string }> { + const results: Array<{ key: string; fullMatch: string }> = []; + + // Pattern for $(pass_mgr get key) or `pass_mgr get key` + const patterns = [ + /\$\(\s*pass_mgr\s+get\s+(\S+)\s*\)/g, + /`\s*pass_mgr\s+get\s+(\S+)\s*`/g, + /pass_mgr\s+get\s+(\S+)/g, + ]; + + for (const pattern of patterns) { + let match; + while ((match = pattern.exec(command)) !== null) { + results.push({ + key: match[1], + fullMatch: match[0], + }); + } + } + + return results; +} + +/** + * Execute pass_mgr get and return the password + */ +async function getPassword(key: string): Promise { + return new Promise((resolve, reject) => { + const passMgrPath = process.env.PASS_MGR_PATH || 'pass_mgr'; + const child = spawn(passMgrPath, ['get', key], { + stdio: ['ignore', 'pipe', 'pipe'], + env: { + ...process.env, + AGENT_WORKSPACE: process.env.AGENT_WORKSPACE || '', + AGENT_ID: process.env.AGENT_ID || '', + }, + }); + + let stdout = ''; + let stderr = ''; + + child.stdout.on('data', (data) => { + stdout += data.toString(); + }); + + child.stderr.on('data', (data) => { + stderr += data.toString(); + }); + + child.on('close', (code) => { + if (code !== 0) { + reject(new Error(`pass_mgr get failed: ${stderr || stdout}`)); + } else { + resolve(stdout.trim()); + } + }); + + child.on('error', (err) => { + reject(err); + }); + }); +} + +/** + * Sanitize output by replacing passwords with ###### + */ +function sanitizeOutput(output: string, passwords: string[]): string { + let sanitized = output; + for (const password of passwords) { + if (password) { + // Escape special regex characters + const escaped = password.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); + const regex = new RegExp(escaped, 'g'); + sanitized = sanitized.replace(regex, '######'); + } + } + return sanitized; +} + +/** + * Replace pass_mgr get commands with actual passwords in command + */ +async function replacePassMgrGets(command: string): Promise<{ command: string; passwords: string[] }> { + const passMgrGets = extractPassMgrGets(command); + const passwords: string[] = []; + let replacedCommand = command; + + for (const { key, fullMatch } of passMgrGets) { + try { + const password = await getPassword(key); + passwords.push(password); + replacedCommand = replacedCommand.replace(fullMatch, password); + } catch (err) { + throw new Error(`Failed to get password for key '${key}': ${err}`); + } + } + + return { command: replacedCommand, passwords }; +} + +/** + * Safe exec wrapper that handles pass_mgr get commands and sanitizes output + * + * @param command - Command to execute + * @param options - Execution options + * @returns Promise resolving to execution result + */ +export async function pcexec(command: string, options: PcExecOptions = {}): Promise { + // Set up environment with workspace/agent info + const env: Record = { + ...process.env, + ...options.env, + }; + + if (process.env.AGENT_WORKSPACE) { + env.AGENT_WORKSPACE = process.env.AGENT_WORKSPACE; + } + if (process.env.AGENT_ID) { + env.AGENT_ID = process.env.AGENT_ID; + } + + // Extract and replace pass_mgr get commands + let finalCommand = command; + let passwords: string[] = []; + + try { + const result = await replacePassMgrGets(command); + finalCommand = result.command; + passwords = result.passwords; + } catch (err) { + throw err; + } + + return new Promise((resolve, reject) => { + const spawnOptions: SpawnOptions = { + cwd: options.cwd, + env, + shell: options.shell ?? true, + windowsHide: options.windowsHide, + uid: options.uid, + gid: options.gid, + }; + + // Use bash for better compatibility + const child = spawn('bash', ['-c', finalCommand], spawnOptions); + + let stdout = ''; + let stderr = ''; + let killed = false; + let timeoutId: NodeJS.Timeout | null = null; + + // Set up timeout + if (options.timeout && options.timeout > 0) { + timeoutId = setTimeout(() => { + killed = true; + child.kill(options.killSignal || 'SIGTERM'); + }, options.timeout); + } + + // Handle stdout + child.stdout?.on('data', (data) => { + stdout += data.toString(); + + // Check maxBuffer + if (options.maxBuffer && stdout.length > options.maxBuffer) { + child.kill(options.killSignal || 'SIGTERM'); + } + }); + + // Handle stderr + child.stderr?.on('data', (data) => { + stderr += data.toString(); + + // Check maxBuffer + if (options.maxBuffer && stderr.length > options.maxBuffer) { + child.kill(options.killSignal || 'SIGTERM'); + } + }); + + // Handle process close + child.on('close', (code, signal) => { + if (timeoutId) { + clearTimeout(timeoutId); + } + + // Sanitize output + const sanitizedStdout = sanitizeOutput(stdout, passwords); + const sanitizedStderr = sanitizeOutput(stderr, passwords); + + if (code === 0) { + resolve({ + stdout: sanitizedStdout, + stderr: sanitizedStderr, + exitCode: 0, + command: finalCommand, + }); + } else { + const error = new Error(`Command failed: ${command}`) as PcExecError; + error.code = code ?? undefined; + error.signal = signal ?? undefined; + error.stdout = sanitizedStdout; + error.stderr = sanitizedStderr; + error.killed = killed; + reject(error); + } + }); + + // Handle process error + child.on('error', (err) => { + if (timeoutId) { + clearTimeout(timeoutId); + } + + const error = new Error(`Failed to execute command: ${err.message}`) as PcExecError; + error.stdout = sanitizeOutput(stdout, passwords); + error.stderr = sanitizeOutput(stderr, passwords); + reject(error); + }); + }); +} + +/** + * Synchronous version of pcexec + * Note: Password sanitization is still applied + */ +export function pcexecSync(command: string, options: PcExecOptions = {}): PcExecResult { + const { execSync } = require('child_process'); + + // Set up environment + const env: Record = { + ...process.env, + ...options.env, + }; + + if (process.env.AGENT_WORKSPACE) { + env.AGENT_WORKSPACE = process.env.AGENT_WORKSPACE; + } + if (process.env.AGENT_ID) { + env.AGENT_ID = process.env.AGENT_ID; + } + + // For sync version, we need to pre-resolve passwords + // This is a limitation - passwords will be in command + const passMgrGets = extractPassMgrGets(command); + let finalCommand = command; + const passwords: string[] = []; + + // Note: In sync version, we can't async fetch passwords + // So we use the original command and rely on the user to not use pass_mgr gets in sync mode + // Or they need to resolve passwords beforehand + + const execOptions: any = { + cwd: options.cwd, + env, + shell: options.shell ?? true, + encoding: 'utf8', + windowsHide: options.windowsHide, + uid: options.uid, + gid: options.gid, + maxBuffer: options.maxBuffer, + timeout: options.timeout, + killSignal: options.killSignal, + }; + + try { + const stdout = execSync(finalCommand, execOptions); + const sanitizedStdout = sanitizeOutput(stdout.toString(), passwords); + + return { + stdout: sanitizedStdout, + stderr: '', + exitCode: 0, + command: finalCommand, + }; + } catch (err: any) { + const sanitizedStdout = sanitizeOutput(err.stdout?.toString() || '', passwords); + const sanitizedStderr = sanitizeOutput(err.stderr?.toString() || '', passwords); + + const error = new Error(`Command failed: ${command}`) as PcExecError; + error.code = err.status; + error.signal = err.signal; + error.stdout = sanitizedStdout; + error.stderr = sanitizedStderr; + throw error; + } +} + +// Default export +export default pcexec; diff --git a/pcexec/tsconfig.json b/pcexec/tsconfig.json new file mode 100644 index 0000000..2f8cd7c --- /dev/null +++ b/pcexec/tsconfig.json @@ -0,0 +1,19 @@ +{ + "compilerOptions": { + "target": "ES2020", + "module": "commonjs", + "lib": ["ES2020"], + "outDir": "./dist", + "rootDir": "./src", + "strict": true, + "esModuleInterop": true, + "skipLibCheck": true, + "forceConsistentCasingInFileNames": true, + "declaration": true, + "declarationMap": true, + "sourceMap": true, + "resolveJsonModule": true + }, + "include": ["src/**/*"], + "exclude": ["node_modules", "dist", "**/*.test.ts"] +} diff --git a/safe-restart/package.json b/safe-restart/package.json new file mode 100644 index 0000000..23fe9bd --- /dev/null +++ b/safe-restart/package.json @@ -0,0 +1,25 @@ +{ + "name": "safe-restart", + "version": "0.1.0", + "description": "Safe restart module for OpenClaw agents", + "main": "dist/index.js", + "types": "dist/index.d.ts", + "scripts": { + "build": "tsc", + "dev": "tsc --watch", + "test": "jest" + }, + "dependencies": { + "@types/node": "^20.0.0", + "express": "^4.18.0", + "ws": "^8.14.0" + }, + "devDependencies": { + "typescript": "^5.0.0", + "@types/express": "^4.17.0", + "@types/ws": "^8.5.0", + "jest": "^29.0.0", + "@types/jest": "^29.0.0", + "ts-jest": "^29.0.0" + } +} diff --git a/safe-restart/src/api.ts b/safe-restart/src/api.ts new file mode 100644 index 0000000..eeb62c2 --- /dev/null +++ b/safe-restart/src/api.ts @@ -0,0 +1,100 @@ +import express from 'express'; +import { StatusManager } from './status-manager'; + +export interface ApiOptions { + port?: number; + statusManager: StatusManager; +} + +/** + * Creates and starts the REST API server for query-restart + */ +export function createApiServer(options: ApiOptions): express.Application { + const { port = 8765, statusManager } = options; + const app = express(); + + app.use(express.json()); + + // POST /query-restart + app.post('/query-restart', (req, res) => { + const { requesterAgentId, requesterSessionKey } = req.body; + + if (!requesterAgentId || !requesterSessionKey) { + return res.status(400).json({ + error: 'Missing required fields: requesterAgentId, requesterSessionKey', + }); + } + + const result = statusManager.queryRestart(requesterAgentId, requesterSessionKey); + + res.json({ + status: result, + timestamp: Date.now(), + }); + }); + + // POST /restart-result + app.post('/restart-result', (req, res) => { + const { status, log } = req.body; + + if (!status || !['ok', 'failed'].includes(status)) { + return res.status(400).json({ + error: 'Invalid status. Must be "ok" or "failed"', + }); + } + + statusManager.completeRestart(status === 'ok', log); + + res.json({ + success: true, + timestamp: Date.now(), + }); + }); + + // GET /status + app.get('/status', (req, res) => { + const agents = statusManager.getAllAgents(); + const global = statusManager.getGlobalStatus(); + + res.json({ + agents, + global, + timestamp: Date.now(), + }); + }); + + // GET /agent/:agentId + app.get('/agent/:agentId', (req, res) => { + const { agentId } = req.params; + const agent = statusManager.getAgent(agentId); + + if (!agent) { + return res.status(404).json({ + error: `Agent ${agentId} not found`, + }); + } + + res.json({ + agent, + timestamp: Date.now(), + }); + }); + + return app; +} + +export function startApiServer(options: ApiOptions): Promise { + const { port = 8765 } = options; + const app = createApiServer(options); + + return new Promise((resolve, reject) => { + const server = app.listen(port, () => { + console.log(`Safe-restart API server listening on port ${port}`); + resolve(); + }); + + server.on('error', (err) => { + reject(err); + }); + }); +} diff --git a/safe-restart/src/index.ts b/safe-restart/src/index.ts new file mode 100644 index 0000000..b07de73 --- /dev/null +++ b/safe-restart/src/index.ts @@ -0,0 +1,4 @@ +export { StatusManager, type AgentStatus, type GlobalStatus, type AgentState } from './status-manager'; +export { createApiServer, startApiServer } from './api'; +export { safeRestart, createSafeRestartTool, type SafeRestartOptions, type SafeRestartResult } from './safe-restart'; +export { SlashCommandHandler, type SlashCommandOptions } from './slash-commands'; diff --git a/safe-restart/src/safe-restart.ts b/safe-restart/src/safe-restart.ts new file mode 100644 index 0000000..7ddb406 --- /dev/null +++ b/safe-restart/src/safe-restart.ts @@ -0,0 +1,288 @@ +import { spawn } from 'child_process'; +import * as fs from 'fs'; +import { promisify } from 'util'; +import { StatusManager } from './status-manager'; + +const sleep = promisify(setTimeout); + +export interface SafeRestartOptions { + /** Agent ID performing the restart */ + agentId: string; + /** Session key for notifications */ + sessionKey: string; + /** API endpoint for query-restart */ + apiEndpoint?: string; + /** Rollback script path */ + rollback?: string; + /** Log file path */ + log?: string; + /** Polling interval in ms (default: 5000) */ + pollInterval?: number; + /** Maximum wait time in ms (default: 300000 = 5min) */ + maxWaitTime?: number; + /** Restart script/command */ + restartScript?: string; + /** Callback for notifications */ + onNotify?: (sessionKey: string, message: string) => Promise; +} + +export interface SafeRestartResult { + success: boolean; + message: string; + log?: string; +} + +/** + * Performs a safe restart with polling and rollback support + */ +export async function safeRestart(options: SafeRestartOptions): Promise { + const { + agentId, + sessionKey, + apiEndpoint = 'http://localhost:8765', + rollback, + log: logPath, + pollInterval = 5000, + maxWaitTime = 300000, + restartScript = 'openclaw gateway restart', + onNotify, + } = options; + + const logs: string[] = []; + const log = (msg: string) => { + const entry = `[${new Date().toISOString()}] ${msg}`; + logs.push(entry); + console.log(entry); + }; + + try { + log(`Starting safe restart. Agent: ${agentId}, Session: ${sessionKey}`); + + // Step 1: Poll query-restart until OK or timeout + const startTime = Date.now(); + let restartApproved = false; + + while (Date.now() - startTime < maxWaitTime) { + try { + const response = await fetch(`${apiEndpoint}/query-restart`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + requesterAgentId: agentId, + requesterSessionKey: sessionKey, + }), + }); + + const data = await response.json(); + + if (data.status === 'OK') { + log('All agents ready for restart'); + restartApproved = true; + break; + } else if (data.status === 'ALREADY_SCHEDULED') { + log('Restart already scheduled by another agent'); + return { + success: false, + message: 'ALREADY_SCHEDULED', + }; + } else { + log(`Waiting for agents to be ready... (${data.status})`); + } + } catch (err) { + log(`Error polling query-restart: ${err}`); + } + + await sleep(pollInterval); + } + + if (!restartApproved) { + const msg = 'Timeout waiting for agents to be ready'; + log(msg); + return { + success: false, + message: msg, + log: logs.join('\n'), + }; + } + + // Step 2: Report restart starting + log('Executing restart...'); + + // Step 3: Start restart in background process + const restartProcess = startBackgroundRestart(restartScript, logPath); + + // Wait a moment for restart to initiate + await sleep(2000); + + // Step 4: Check if gateway comes back + log('Waiting for gateway to restart...'); + await sleep(60000); // Wait 60s as specified + + // Check gateway status + const gatewayOk = await checkGatewayStatus(); + + if (gatewayOk) { + log('Gateway restarted successfully'); + + // Report success + await fetch(`${apiEndpoint}/restart-result`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + status: 'ok', + log: logPath || logs.join('\n'), + }), + }); + + // Notify resumption + if (onNotify) { + await onNotify(sessionKey, 'restart 结束了,我们继续'); + } + + return { + success: true, + message: 'Restart completed successfully', + }; + } else { + log('Gateway restart failed'); + + // Execute rollback if provided + if (rollback) { + log(`Executing rollback: ${rollback}`); + try { + await executeRollback(rollback); + log('Rollback completed'); + } catch (err) { + log(`Rollback failed: ${err}`); + } + } + + // Report failure + await fetch(`${apiEndpoint}/restart-result`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + status: 'failed', + log: logPath || logs.join('\n'), + }), + }); + + // Notify failure + if (onNotify) { + await onNotify(sessionKey, 'restart 失败,已经 rollback,请参考 log 调查。'); + } + + return { + success: false, + message: 'Restart failed', + log: logs.join('\n'), + }; + } + } catch (err) { + const errorMsg = `Unexpected error: ${err}`; + log(errorMsg); + return { + success: false, + message: errorMsg, + log: logs.join('\n'), + }; + } +} + +function startBackgroundRestart(restartScript: string, logPath?: string): void { + const script = ` +#!/bin/bash +set -e +sleep 60 +${restartScript} +openclaw gateway status +`; + + const child = spawn('bash', ['-c', script], { + detached: true, + stdio: logPath ? ['ignore', fs.openSync(logPath, 'w'), fs.openSync(logPath, 'w+')] : 'ignore', + }); + + child.unref(); +} + +async function checkGatewayStatus(): Promise { + return new Promise((resolve) => { + const child = spawn('openclaw', ['gateway', 'status'], { + timeout: 10000, + }); + + let output = ''; + child.stdout?.on('data', (data) => { + output += data.toString(); + }); + + child.on('close', (code) => { + resolve(code === 0 && output.includes('running')); + }); + + child.on('error', () => { + resolve(false); + }); + }); +} + +async function executeRollback(rollbackScript: string): Promise { + return new Promise((resolve, reject) => { + const child = spawn('bash', ['-c', rollbackScript], { + timeout: 120000, + }); + + child.on('close', (code) => { + if (code === 0) { + resolve(); + } else { + reject(new Error(`Rollback script exited with code ${code}`)); + } + }); + + child.on('error', (err) => { + reject(err); + }); + }); +} + +/** + * Safe restart tool that can be registered with OpenClaw + */ +export function createSafeRestartTool(statusManager: StatusManager) { + return { + name: 'safe_restart', + description: 'Perform a safe restart of OpenClaw gateway with agent coordination', + parameters: { + type: 'object', + properties: { + rollback: { + type: 'string', + description: 'Path to rollback script', + }, + log: { + type: 'string', + description: 'Path to log file', + }, + }, + }, + handler: async (params: { rollback?: string; log?: string }, context: { agentId: string; sessionKey: string }) => { + const result = await safeRestart({ + agentId: context.agentId, + sessionKey: context.sessionKey, + rollback: params.rollback, + log: params.log, + async onNotify(sessionKey, message) { + // This would be connected to the messaging system + console.log(`[${sessionKey}] ${message}`); + }, + }); + + return { + success: result.success, + message: result.message, + }; + }, + }; +} diff --git a/safe-restart/src/slash-commands.ts b/safe-restart/src/slash-commands.ts new file mode 100644 index 0000000..a6f9223 --- /dev/null +++ b/safe-restart/src/slash-commands.ts @@ -0,0 +1,182 @@ +import { StatusManager } from './status-manager'; + +export interface SlashCommandOptions { + statusManager: StatusManager; + /** List of authorized user IDs */ + authorizedUsers: string[]; + /** Cooldown duration in seconds */ + cooldownSeconds?: number; + /** Callback for replies */ + onReply: (message: string) => Promise; +} + +interface CommandState { + passMgrEnabled: boolean; + safeRestartEnabled: boolean; + lastToggle: { + 'pass-mgr': number; + 'safe-restart': number; + }; +} + +export class SlashCommandHandler { + private statusManager: StatusManager; + private authorizedUsers: string[]; + private cooldownMs: number; + private onReply: (message: string) => Promise; + private state: CommandState; + + constructor(options: SlashCommandOptions) { + this.statusManager = options.statusManager; + this.authorizedUsers = options.authorizedUsers; + this.cooldownMs = (options.cooldownSeconds || 10) * 1000; + this.onReply = options.onReply; + + this.state = { + passMgrEnabled: true, + safeRestartEnabled: true, + lastToggle: { + 'pass-mgr': 0, + 'safe-restart': 0, + }, + }; + } + + /** + * Handle a slash command + */ + async handle(command: string, userId: string): Promise { + // Check authorization + if (!this.authorizedUsers.includes(userId)) { + await this.onReply('❌ 无权执行此命令'); + return; + } + + const parts = command.trim().split(/\s+/); + const subcommand = parts[1]; + const feature = parts[2] as 'pass-mgr' | 'safe-restart'; + + switch (subcommand) { + case 'status': + await this.handleStatus(); + break; + case 'enable': + await this.handleEnable(feature); + break; + case 'disable': + await this.handleDisable(feature); + break; + default: + await this.onReply( + '用法:\n' + + '`/padded-cell-ctrl status` - 查看状态\n' + + '`/padded-cell-ctrl enable pass-mgr|safe-restart` - 启用功能\n' + + '`/padded-cell-ctrl disable pass-mgr|safe-restart` - 禁用功能' + ); + } + } + + private async handleStatus(): Promise { + const global = this.statusManager.getGlobalStatus(); + const agents = this.statusManager.getAllAgents(); + + const lines = [ + '**PaddedCell 状态**', + '', + `🔐 密码管理: ${this.state.passMgrEnabled ? '✅ 启用' : '❌ 禁用'}`, + `🔄 安全重启: ${this.state.safeRestartEnabled ? '✅ 启用' : '❌ 禁用'}`, + '', + '**Agent 状态:**', + ]; + + for (const agent of agents) { + const emoji = this.getStateEmoji(agent.state); + lines.push(`${emoji} ${agent.agentId}: ${agent.state}`); + } + + if (agents.length === 0) { + lines.push('(暂无 agent 注册)'); + } + + if (global.restartStatus !== 'idle') { + lines.push(''); + lines.push(`⚠️ 重启状态: ${global.restartStatus}`); + if (global.restartScheduledBy) { + lines.push(` 由 ${global.restartScheduledBy} 发起`); + } + } + + await this.onReply(lines.join('\n')); + } + + private async handleEnable(feature: 'pass-mgr' | 'safe-restart'): Promise { + if (!this.isValidFeature(feature)) { + await this.onReply('❌ 未知功能。可用选项: pass-mgr, safe-restart'); + return; + } + + if (this.isOnCooldown(feature)) { + await this.onReply('⏳ 该功能最近刚被修改过,请稍后再试'); + return; + } + + if (feature === 'pass-mgr') { + this.state.passMgrEnabled = true; + } else { + this.state.safeRestartEnabled = true; + } + + this.state.lastToggle[feature] = Date.now(); + await this.onReply(`✅ 已启用 ${feature}`); + } + + private async handleDisable(feature: 'pass-mgr' | 'safe-restart'): Promise { + if (!this.isValidFeature(feature)) { + await this.onReply('❌ 未知功能。可用选项: pass-mgr, safe-restart'); + return; + } + + if (this.isOnCooldown(feature)) { + await this.onReply('⏳ 该功能最近刚被修改过,请稍后再试'); + return; + } + + if (feature === 'pass-mgr') { + this.state.passMgrEnabled = false; + } else { + this.state.safeRestartEnabled = false; + } + + this.state.lastToggle[feature] = Date.now(); + await this.onReply(`✅ 已禁用 ${feature}`); + } + + private isValidFeature(feature: string): feature is 'pass-mgr' | 'safe-restart' { + return feature === 'pass-mgr' || feature === 'safe-restart'; + } + + private isOnCooldown(feature: 'pass-mgr' | 'safe-restart'): boolean { + const lastToggle = this.state.lastToggle[feature]; + return Date.now() - lastToggle < this.cooldownMs; + } + + private getStateEmoji(state: string): string { + switch (state) { + case 'idle': return '💤'; + case 'busy': return '⚡'; + case 'focus': return '🎯'; + case 'freeze': return '🧊'; + case 'pre-freeze': return '⏳'; + case 'pre-freeze-focus': return '📝'; + default: return '❓'; + } + } + + isPassMgrEnabled(): boolean { + return this.state.passMgrEnabled; + } + + isSafeRestartEnabled(): boolean { + return this.state.safeRestartEnabled; + } +} diff --git a/safe-restart/src/status-manager.ts b/safe-restart/src/status-manager.ts new file mode 100644 index 0000000..ee67811 --- /dev/null +++ b/safe-restart/src/status-manager.ts @@ -0,0 +1,396 @@ +import * as fs from 'fs'; +import * as path from 'path'; +import { EventEmitter } from 'events'; + +export type AgentState = + 'idle' | + 'busy' | + 'focus' | + 'freeze' | + 'pre-freeze' | + 'pre-freeze-focus'; + +export interface AgentStatus { + agentId: string; + state: AgentState; + workflow: string | null; + activeSessions: string[]; + lastSessions: string[]; + updatedAt: number; +} + +export interface GlobalStatus { + restartScheduledBy: string | null; + restartSession: string | null; + restartStatus: 'idle' | 'waiting' | 'restarting' | 'rollback'; + updatedAt: number; +} + +export interface StatusManagerOptions { + dataDir?: string; + persistenceInterval?: number; +} + +/** + * Manages agent states and global restart status + */ +export class StatusManager extends EventEmitter { + private agents: Map = new Map(); + private global: GlobalStatus; + private dataDir: string; + private persistenceInterval: number; + private persistenceTimer: NodeJS.Timeout | null = null; + + constructor(options: StatusManagerOptions = {}) { + super(); + this.dataDir = options.dataDir || path.join(process.env.HOME || '.', '.paddedcell'); + this.persistenceInterval = options.persistenceInterval || 5000; + + this.global = { + restartScheduledBy: null, + restartSession: null, + restartStatus: 'idle', + updatedAt: Date.now(), + }; + + this.ensureDataDir(); + this.loadFromDisk(); + this.startPersistence(); + } + + private ensureDataDir(): void { + if (!fs.existsSync(this.dataDir)) { + fs.mkdirSync(this.dataDir, { recursive: true, mode: 0o700 }); + } + } + + private getAgentFilePath(agentId: string): string { + return path.join(this.dataDir, `agent_${agentId}.json`); + } + + private getGlobalFilePath(): string { + return path.join(this.dataDir, 'global.json'); + } + + private loadFromDisk(): void { + // Load global status + const globalPath = this.getGlobalFilePath(); + if (fs.existsSync(globalPath)) { + try { + const data = fs.readFileSync(globalPath, 'utf8'); + this.global = JSON.parse(data); + } catch (err) { + console.error('Failed to load global status:', err); + } + } + + // Load agent statuses + try { + const files = fs.readdirSync(this.dataDir); + for (const file of files) { + if (file.startsWith('agent_') && file.endsWith('.json')) { + try { + const data = fs.readFileSync(path.join(this.dataDir, file), 'utf8'); + const agent = JSON.parse(data) as AgentStatus; + this.agents.set(agent.agentId, agent); + } catch (err) { + console.error(`Failed to load agent status from ${file}:`, err); + } + } + } + } catch (err) { + console.error('Failed to read data directory:', err); + } + } + + private saveToDisk(): void { + // Save global status + try { + fs.writeFileSync( + this.getGlobalFilePath(), + JSON.stringify(this.global, null, 2), + { mode: 0o600 } + ); + } catch (err) { + console.error('Failed to save global status:', err); + } + + // Save agent statuses + for (const [agentId, agent] of this.agents) { + try { + fs.writeFileSync( + this.getAgentFilePath(agentId), + JSON.stringify(agent, null, 2), + { mode: 0o600 } + ); + } catch (err) { + console.error(`Failed to save agent status for ${agentId}:`, err); + } + } + } + + private startPersistence(): void { + this.persistenceTimer = setInterval(() => { + this.saveToDisk(); + }, this.persistenceInterval); + } + + stopPersistence(): void { + if (this.persistenceTimer) { + clearInterval(this.persistenceTimer); + this.persistenceTimer = null; + } + this.saveToDisk(); + } + + // Agent state management + + getOrCreateAgent(agentId: string): AgentStatus { + if (!this.agents.has(agentId)) { + const agent: AgentStatus = { + agentId, + state: 'idle', + workflow: null, + activeSessions: [], + lastSessions: [], + updatedAt: Date.now(), + }; + this.agents.set(agentId, agent); + this.emit('agentCreated', agent); + } + return this.agents.get(agentId)!; + } + + getAgent(agentId: string): AgentStatus | undefined { + return this.agents.get(agentId); + } + + getAllAgents(): AgentStatus[] { + return Array.from(this.agents.values()); + } + + onMessageStart(session: string, agentId: string): void { + const agent = this.getOrCreateAgent(agentId); + + // Don't update state for heartbeat sessions + if (this.isHeartbeatSession(session)) { + return; + } + + if (agent.state === 'idle') { + agent.state = 'busy'; + } + + if (!agent.activeSessions.includes(session)) { + agent.activeSessions.push(session); + } + + agent.updatedAt = Date.now(); + this.emit('stateChanged', agent); + } + + onMessageEnd(session: string, agentId: string): void { + const agent = this.getOrCreateAgent(agentId); + + // Remove from active sessions + agent.activeSessions = agent.activeSessions.filter(s => s !== session); + + // Add to last sessions if not already there + if (!agent.lastSessions.includes(session)) { + agent.lastSessions.unshift(session); + if (agent.lastSessions.length > 10) { + agent.lastSessions = agent.lastSessions.slice(0, 10); + } + } + + // State transitions + if (agent.activeSessions.length === 0) { + if (agent.state === 'busy') { + agent.state = 'idle'; + } else if (agent.state === 'pre-freeze' || agent.state === 'pre-freeze-focus') { + agent.state = 'freeze'; + } + } + + agent.updatedAt = Date.now(); + this.emit('stateChanged', agent); + + // Check if all agents are frozen (for restart completion) + if (agent.state === 'freeze') { + this.checkAllFrozen(); + } + } + + setWorkflow(agentId: string, workflow: string | null): void { + const agent = this.getOrCreateAgent(agentId); + agent.workflow = workflow; + + if (workflow) { + agent.state = 'focus'; + } else { + // Transition from focus to idle or busy + if (agent.activeSessions.length === 0) { + agent.state = 'idle'; + } else { + agent.state = 'busy'; + } + } + + agent.updatedAt = Date.now(); + this.emit('stateChanged', agent); + } + + isHeartbeatSession(session: string): boolean { + // Check if session is a heartbeat session + // This can be customized based on naming convention or metadata + return session.includes('heartbeat') || session.includes('poll'); + } + + // Query restart logic + + queryRestart(requesterAgentId: string, requesterSessionKey: string): 'OK' | 'NOT_READY' | 'ALREADY_SCHEDULED' { + // Check if restart is already scheduled + if (this.global.restartStatus !== 'idle') { + // If same agent is requesting, allow continuation + if (this.global.restartScheduledBy === requesterAgentId) { + return this.allAgentsFrozen() ? 'OK' : 'NOT_READY'; + } + return 'ALREADY_SCHEDULED'; + } + + // Schedule restart + this.global.restartScheduledBy = requesterAgentId; + this.global.restartSession = requesterSessionKey; + this.global.restartStatus = 'waiting'; + this.global.updatedAt = Date.now(); + + // Transition agents to freeze/pre-freeze states + for (const [agentId, agent] of this.agents) { + if (agentId === requesterAgentId) { + // Don't freeze the requester agent + continue; + } + + switch (agent.state) { + case 'idle': + agent.state = 'freeze'; + break; + case 'busy': + agent.state = 'pre-freeze'; + break; + case 'focus': + agent.state = 'pre-freeze-focus'; + // Notify agent to prepare for restart + this.emit('preparingRestart', agent); + break; + } + + agent.updatedAt = Date.now(); + this.emit('stateChanged', agent); + } + + this.saveToDisk(); + + // Check if all are frozen immediately + if (this.allAgentsFrozen()) { + return 'OK'; + } + + return 'NOT_READY'; + } + + allAgentsFrozen(): boolean { + for (const [agentId, agent] of this.agents) { + // Skip the agent that scheduled the restart + if (agentId === this.global.restartScheduledBy) { + continue; + } + if (agent.state !== 'freeze') { + return false; + } + } + return true; + } + + private checkAllFrozen(): void { + if (this.allAgentsFrozen() && this.global.restartStatus === 'waiting') { + this.emit('allFrozen'); + } + } + + // Restart completion + + completeRestart(success: boolean, log?: string): void { + if (success) { + this.global.restartStatus = 'idle'; + + // Unfreeze all agents + for (const agent of this.agents.values()) { + if (agent.state === 'freeze') { + // Restore previous state from lastSessions + agent.state = agent.activeSessions.length > 0 ? 'busy' : 'idle'; + agent.updatedAt = Date.now(); + this.emit('stateChanged', agent); + this.emit('unfrozen', agent); + } + } + + this.emit('restartCompleted'); + } else { + this.global.restartStatus = 'rollback'; + this.emit('restartFailed', log); + } + + this.global.restartScheduledBy = null; + this.global.restartSession = null; + this.global.updatedAt = Date.now(); + this.saveToDisk(); + } + + // Global status getters + + getGlobalStatus(): GlobalStatus { + return { ...this.global }; + } + + isRestartScheduled(): boolean { + return this.global.restartStatus !== 'idle'; + } + + // For focus mode: check if agent should respond + + shouldRespond(agentId: string, session: string): boolean { + const agent = this.getAgent(agentId); + if (!agent) return true; + + // In focus mode, only respond to workflow sessions + if (agent.state === 'focus' || agent.state === 'pre-freeze-focus') { + return agent.workflow !== null && session.includes(agent.workflow); + } + + // In freeze/pre-freeze states, don't accept new messages + if (agent.state === 'freeze' || agent.state === 'pre-freeze') { + return false; + } + + return true; + } + + getBusyMessage(agentId: string): string { + const agent = this.getAgent(agentId); + if (!agent) return '在忙,无法应答'; + + switch (agent.state) { + case 'focus': + case 'pre-freeze-focus': + return '当前处于专注模式,无法应答非工作流消息'; + case 'freeze': + case 'pre-freeze': + return '系统正在准备重启,请稍后再试'; + case 'busy': + return '正在处理其他消息,请稍后再试'; + default: + return '在忙,无法应答'; + } + } +} diff --git a/safe-restart/tsconfig.json b/safe-restart/tsconfig.json new file mode 100644 index 0000000..2f8cd7c --- /dev/null +++ b/safe-restart/tsconfig.json @@ -0,0 +1,19 @@ +{ + "compilerOptions": { + "target": "ES2020", + "module": "commonjs", + "lib": ["ES2020"], + "outDir": "./dist", + "rootDir": "./src", + "strict": true, + "esModuleInterop": true, + "skipLibCheck": true, + "forceConsistentCasingInFileNames": true, + "declaration": true, + "declarationMap": true, + "sourceMap": true, + "resolveJsonModule": true + }, + "include": ["src/**/*"], + "exclude": ["node_modules", "dist", "**/*.test.ts"] +}