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
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