feat: implement pass_mgr, pcexec, and safe-restart modules
- Add pass_mgr Go binary with AES-256-GCM encryption - Add pcexec TypeScript tool with password sanitization - Add safe-restart module with state machine and API - Add slash command handler with cooldown support - Update README with usage documentation
This commit is contained in:
root
100
safe-restart/src/api.ts
Normal file
100
safe-restart/src/api.ts
Normal file
@@ -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<void> {
|
||||
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);
|
||||
});
|
||||
});
|
||||
}
|
||||
4
safe-restart/src/index.ts
Normal file
4
safe-restart/src/index.ts
Normal file
@@ -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';
|
||||
288
safe-restart/src/safe-restart.ts
Normal file
288
safe-restart/src/safe-restart.ts
Normal file
@@ -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<void>;
|
||||
}
|
||||
|
||||
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<SafeRestartResult> {
|
||||
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<boolean> {
|
||||
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<void> {
|
||||
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,
|
||||
};
|
||||
},
|
||||
};
|
||||
}
|
||||
182
safe-restart/src/slash-commands.ts
Normal file
182
safe-restart/src/slash-commands.ts
Normal file
@@ -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<void>;
|
||||
}
|
||||
|
||||
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<void>;
|
||||
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<void> {
|
||||
// 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<void> {
|
||||
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<void> {
|
||||
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<void> {
|
||||
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;
|
||||
}
|
||||
}
|
||||
396
safe-restart/src/status-manager.ts
Normal file
396
safe-restart/src/status-manager.ts
Normal file
@@ -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<string, AgentStatus> = 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 '在忙，无法应答';
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user