automagik-genie
Version:
Self-evolving AI agent orchestration framework with Model Context Protocol support
251 lines (250 loc) • 10.2 kB
JavaScript
;
/**
* Task Monitor - WebSocket-based task completion monitoring
*
* Provides event-driven task monitoring using Forge's normalized-logs WebSocket endpoint.
* Replaces polling-based approach with real-time status updates.
*/
var __importDefault = (this && this.__importDefault) || function (mod) {
return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.monitorTaskCompletion = monitorTaskCompletion;
const ws_1 = __importDefault(require("ws"));
// @ts-ignore - legacy JS Forge client does not ship TypeScript typings
const forge_client_js_1 = require("../../../src/lib/forge-client.js");
/**
* Monitor task completion using WebSocket connection
*
* Connects to Forge's normalized-logs WebSocket endpoint and listens for
* task completion events. More efficient than polling.
*
* @param options Task monitoring configuration
* @returns Promise resolving to task result
*/
async function monitorTaskCompletion(options) {
const { attemptId, baseUrl, onLog, onStatus, timeout = 300000, token, taskUrl: providedTaskUrl } = options;
const startTime = Date.now();
const authToken = token ?? process.env.FORGE_TOKEN;
const client = new forge_client_js_1.ForgeClient(baseUrl, authToken);
const attemptMeta = await getTaskAttemptSafe(client, attemptId);
const taskUrl = providedTaskUrl || deriveTaskUrl(baseUrl, attemptMeta, attemptId);
const processWaitMs = Math.min(Math.max(timeout - 5000, 60000), timeout);
const processId = await waitForLatestProcessId(client, attemptId, processWaitMs);
if (!processId) {
return {
status: 'failed',
output: '',
error: 'No execution process found for task attempt (executor did not start in time)',
duration_ms: Date.now() - startTime,
task_url: taskUrl
};
}
const wsUrl = client.getNormalizedLogsStreamUrl(processId);
return new Promise((resolve, reject) => {
let ws = null;
let timeoutId = null;
let outputBuffer = [];
let lastStatus = 'running';
let settled = false;
// Cleanup function
const cleanup = () => {
if (timeoutId)
clearTimeout(timeoutId);
timeoutId = null;
if (ws) {
ws.removeAllListeners();
if (ws.readyState === ws_1.default.OPEN || ws.readyState === ws_1.default.CONNECTING) {
ws.close();
}
}
};
const resolveOnce = (result) => {
if (settled)
return;
settled = true;
cleanup();
resolve(result);
};
// Timeout handler
timeoutId = setTimeout(() => {
resolveOnce({
status: 'timeout',
output: outputBuffer.join('\n'),
error: 'Task did not complete within timeout',
duration_ms: Date.now() - startTime,
task_url: taskUrl
});
}, timeout);
// Create WebSocket connection
try {
ws = new ws_1.default(wsUrl, {
headers: {
'User-Agent': 'Genie-TaskMonitor/1.0',
...(authToken ? { Authorization: `Bearer ${authToken}` } : {})
}
});
}
catch (error) {
cleanup();
reject(new Error(`Failed to create WebSocket: ${error}`));
return;
}
// Connection opened
ws.on('open', () => {
// Connection established, waiting for messages
});
// Message received
ws.on('message', (data) => {
try {
const message = JSON.parse(data.toString());
// Handle different message types
if (message.type === 'log' && message.content) {
const logLine = message.content;
outputBuffer.push(logLine);
if (onLog)
onLog(logLine);
}
if (message.type === 'status' && message.status) {
lastStatus = message.status;
if (onStatus)
onStatus(message.status);
// Check for completion
if (message.status === 'completed' || message.status === 'success') {
resolveOnce({
status: 'completed',
output: outputBuffer.join('\n') || '(no output)',
duration_ms: Date.now() - startTime,
task_url: taskUrl
});
}
else if (message.status === 'failed' || message.status === 'error') {
resolveOnce({
status: 'failed',
output: outputBuffer.join('\n') || '(no output)',
error: message.error || 'Task execution failed',
duration_ms: Date.now() - startTime,
task_url: taskUrl
});
}
}
// Handle finished flag from normalized log stream
if (message.finished === true) {
lastStatus = 'finished';
// WebSocket will close after this, and close handler will query actual status
}
// Handle raw log messages (fallback)
if (typeof message === 'string') {
outputBuffer.push(message);
if (onLog)
onLog(message);
}
}
catch (error) {
// If parsing fails, treat as plain text
const textMessage = data.toString();
outputBuffer.push(textMessage);
if (onLog)
onLog(textMessage);
}
});
// Error handler
ws.on('error', (error) => {
resolveOnce({
status: 'failed',
output: outputBuffer.join('\n'),
error: `WebSocket error: ${error.message}`,
duration_ms: Date.now() - startTime,
task_url: taskUrl
});
});
// Connection closed
ws.on('close', async (code, reason) => {
// If we received the 'finished' flag, query actual task status
if (lastStatus === 'finished') {
try {
const attempt = await client.getTaskAttempt(attemptId);
const actualStatus = attempt?.status || 'unknown';
if (actualStatus === 'completed' || actualStatus === 'success') {
resolveOnce({
status: 'completed',
output: outputBuffer.join('\n') || '(no output)',
duration_ms: Date.now() - startTime,
task_url: taskUrl
});
return;
}
else if (actualStatus === 'failed' || actualStatus === 'error') {
resolveOnce({
status: 'failed',
output: outputBuffer.join('\n') || '(no output)',
error: 'Task execution failed',
duration_ms: Date.now() - startTime,
task_url: taskUrl
});
return;
}
}
catch (error) {
// Fall through to default handling if query fails
}
}
const closedStatus = lastStatus === 'completed' || lastStatus === 'success'
? 'completed'
: lastStatus === 'failed' || lastStatus === 'error'
? 'failed'
: 'failed';
const errorMessage = closedStatus === 'failed' && (!lastStatus || lastStatus === 'running')
? `Connection closed (code ${code}) before receiving final status${reason?.length ? `: ${reason.toString()}` : ''}`
: undefined;
resolveOnce({
status: closedStatus,
output: outputBuffer.join('\n') || '(no output)',
error: errorMessage,
duration_ms: Date.now() - startTime,
task_url: taskUrl
});
});
});
}
async function getTaskAttemptSafe(client, attemptId) {
try {
return await client.getTaskAttempt(attemptId);
}
catch {
return null;
}
}
function deriveTaskUrl(baseUrl, attempt, attemptId) {
if (attempt) {
const projectId = attempt.project_id || attempt.projectId;
const taskId = attempt.task_id || attempt.taskId;
if (projectId && taskId) {
const cleanBase = baseUrl.replace(/\/$/, '');
return `${cleanBase}/projects/${projectId}/tasks/${taskId}/attempts/${attemptId}?view=diffs`;
}
}
return `${baseUrl.replace(/\/$/, '')}/tasks/${attemptId}`;
}
async function waitForLatestProcessId(client, attemptId, timeoutMs = 60000, pollIntervalMs = 1000) {
const start = Date.now();
while (Date.now() - start < timeoutMs) {
try {
const processes = await client.listExecutionProcesses(attemptId);
if (Array.isArray(processes) && processes.length > 0) {
const latest = processes[processes.length - 1];
if (latest?.id) {
return latest.id;
}
}
}
catch {
// Ignore temporary errors – Forge may still be booting the executor
}
await sleep(pollIntervalMs);
}
return null;
}
function sleep(durationMs) {
return new Promise(resolve => setTimeout(resolve, durationMs));
}