byterover-cli 1.7.2 → 2.0.0

package/dist/server/infra/process/task-router.js

@@ -0,0 +1,503 @@
+/**
+ * TaskRouter - Routes task and LLM events between clients and agents.
+ *
+ * Handles:
+ * - Task lifecycle: create → ack → started → completed/error/cancelled
+ * - LLM event routing: llmservice:* events from agent → client + project room
+ * - Grace period: keeps completed tasks briefly for late-arriving LLM events
+ * - Lifecycle hooks: extensible observer hooks (e.g. CurateLogHandler)
+ *
+ * Broadcasting: Task/LLM events are broadcast to project-scoped rooms
+ * (project:<sanitizedPath>:broadcast) so only clients in the same project
+ * receive them. Global events (auth, agent connect/disconnect) remain on
+ * the global broadcast channel.
+ *
+ * Consumed by TransportHandlers (orchestrator).
+ */
+import { AgentNotAvailableError, serializeTaskError } from '../../core/domain/errors/task-error.js';
+import { LlmEventNames, TransportLlmEventList, TransportTaskEventNames } from '../../core/domain/transport/schemas.js';
+import { transportLog } from '../../utils/process-logger.js';
+import { isValidTaskType } from '../../utils/type-guards.js';
+import { broadcastToProjectRoom } from './broadcast-utils.js';
+/**
+ * Grace period (in ms) to keep completed tasks in memory for late-arriving events.
+ * Prevents silent event drops when llmservice:* events arrive after task:completed.
+ */
+const TASK_CLEANUP_GRACE_PERIOD_MS = 5000;
+function hasTaskId(data) {
+    return typeof data === 'object' && data !== null && 'taskId' in data && typeof data.taskId === 'string';
+}
+export class TaskRouter {
+    agentPool;
+    /**
+     * Track recently completed tasks for grace period.
+     * Allows late-arriving llmservice:* events to be routed even after task:completed.
+     */
+    completedTasks = new Map();
+    getAgentForProject;
+    lifecycleHooks;
+    projectRegistry;
+    projectRouter;
+    resolveClientProjectPath;
+    /** Track active tasks */
+    tasks = new Map();
+    transport;
+    constructor(options) {
+        this.transport = options.transport;
+        this.agentPool = options.agentPool;
+        this.getAgentForProject = options.getAgentForProject;
+        this.lifecycleHooks = options.lifecycleHooks ?? [];
+        this.projectRegistry = options.projectRegistry;
+        this.projectRouter = options.projectRouter;
+        this.resolveClientProjectPath = options.resolveClientProjectPath;
+    }
+    clearTasks() {
+        this.tasks.clear();
+        this.completedTasks.clear();
+    }
+    /**
+     * Remove a task from tracking and send error to its client.
+     * Used by ConnectionCoordinator when an agent disconnects.
+     */
+    failTask(taskId, error) {
+        const task = this.tasks.get(taskId);
+        if (!task)
+            return;
+        this.transport.sendTo(task.clientId, TransportTaskEventNames.ERROR, { error, taskId });
+        broadcastToProjectRoom(this.projectRegistry, this.projectRouter, task.projectPath, TransportTaskEventNames.ERROR, { error, taskId }, task.clientId);
+        this.tasks.delete(taskId);
+        // Notify hooks (fire-and-forget)
+        this.notifyHooksError(taskId, error.message, task).catch(() => { });
+    }
+    getDebugState() {
+        return {
+            activeTasks: [...this.tasks.values()].map((t) => ({
+                clientId: t.clientId,
+                createdAt: t.createdAt,
+                projectPath: t.projectPath,
+                taskId: t.taskId,
+                type: t.type,
+            })),
+            completedTasks: [...this.completedTasks.entries()].map(([taskId, entry]) => ({
+                completedAt: entry.completedAt,
+                projectPath: entry.task.projectPath,
+                taskId,
+                type: entry.task.type,
+            })),
+        };
+    }
+    /**
+     * Returns all active tasks for a given project path.
+     * Used by ConnectionCoordinator to fail tasks on agent disconnect.
+     */
+    getTasksForProject(projectPath) {
+        const result = [];
+        for (const task of this.tasks.values()) {
+            if (projectPath === undefined) {
+                // No projectPath specified — only match tasks without a project
+                if (task.projectPath === undefined) {
+                    result.push(task);
+                }
+            }
+            else if (task.projectPath === projectPath || task.projectPath === undefined) {
+                // Specific project — match tasks for that project or unassigned tasks
+                result.push(task);
+            }
+        }
+        return result;
+    }
+    /**
+     * Register all task and LLM event handlers on the transport.
+     */
+    setup() {
+        // Task creation from clients
+        this.transport.onRequest(TransportTaskEventNames.CREATE, (data, clientId) => this.handleTaskCreate(data, clientId));
+        // Task cancellation from clients
+        this.transport.onRequest(TransportTaskEventNames.CANCEL, (data, clientId) => this.handleTaskCancel(data, clientId));
+        // Task lifecycle events from agent
+        this.transport.onRequest(TransportTaskEventNames.STARTED, (data) => {
+            this.handleTaskStarted(data);
+        });
+        this.transport.onRequest(TransportTaskEventNames.COMPLETED, (data) => {
+            this.handleTaskCompleted(data);
+        });
+        this.transport.onRequest(TransportTaskEventNames.ERROR, (data) => {
+            this.handleTaskError(data);
+        });
+        this.transport.onRequest(TransportTaskEventNames.CANCELLED, (data) => {
+            this.handleTaskCancelled(data);
+        });
+        // LLM events
+        for (const eventName of TransportLlmEventList) {
+            this.registerLlmEvent(eventName);
+        }
+    }
+    handleTaskCancel(data, _clientId) {
+        const { taskId } = data;
+        transportLog(`Task cancel requested: ${taskId}`);
+        const task = this.tasks.get(taskId);
+        if (!task) {
+            return { error: 'Task not found', success: false };
+        }
+        // If Agent connected for this task's project, forward cancel request
+        const agentId = this.getAgentForProject(task.projectPath);
+        if (agentId) {
+            this.transport.sendTo(agentId, TransportTaskEventNames.CANCEL, { taskId });
+            return { success: true };
+        }
+        // No Agent - cancel task locally and emit terminal event
+        transportLog(`No Agent connected, cancelling task locally: ${taskId}`);
+        this.transport.sendTo(task.clientId, TransportTaskEventNames.CANCELLED, { taskId });
+        broadcastToProjectRoom(this.projectRegistry, this.projectRouter, task.projectPath, TransportTaskEventNames.CANCELLED, { taskId }, task.clientId);
+        this.tasks.delete(taskId);
+        this.notifyHooksCancelled(taskId, task).catch(() => { });
+        return { success: true };
+    }
+    handleTaskCancelled(data) {
+        const { taskId } = data;
+        const task = this.tasks.get(taskId);
+        transportLog(`Task cancelled: ${taskId}`);
+        if (task) {
+            this.transport.sendTo(task.clientId, TransportTaskEventNames.CANCELLED, { taskId });
+        }
+        broadcastToProjectRoom(this.projectRegistry, this.projectRouter, task?.projectPath, TransportTaskEventNames.CANCELLED, { taskId }, task?.clientId);
+        this.moveToCompleted(taskId);
+        // Notify hooks (fire-and-forget)
+        if (task) {
+            this.notifyHooksCancelled(taskId, task).catch(() => { });
+        }
+    }
+    handleTaskCompleted(data) {
+        const { result, taskId } = data;
+        const task = this.tasks.get(taskId);
+        transportLog(`Task completed: ${taskId}`);
+        if (task) {
+            this.transport.sendTo(task.clientId, TransportTaskEventNames.COMPLETED, {
+                ...(task.logId ? { logId: task.logId } : {}),
+                result,
+                taskId,
+            });
+        }
+        broadcastToProjectRoom(this.projectRegistry, this.projectRouter, task?.projectPath, TransportTaskEventNames.COMPLETED, {
+            ...(task?.logId ? { logId: task.logId } : {}),
+            result,
+            taskId,
+        }, task?.clientId);
+        this.moveToCompleted(taskId);
+        // Notify pool so it can clear busy flag and drain queued tasks
+        if (task?.projectPath) {
+            this.agentPool?.notifyTaskCompleted(task.projectPath);
+        }
+        // Notify hooks (fire-and-forget)
+        if (task) {
+            this.notifyHooksCompleted(taskId, result, task).catch(() => { });
+        }
+    }
+    /**
+     * Handle task creation from a client.
+     *
+     * Ordering (critical for correctness):
+     * 1. Idempotency check
+     * 2. Early validation — on failure: send task:error, return. No task stored, no task:created, no hooks called.
+     * 3. Store task + send task:created synchronously (before any await)
+     * 4. Await lifecycle hooks → get logId
+     *    Note: task:ack is intentionally delayed until hooks resolve so logId can be included.
+     *    This reverses the old ordering (previously ack preceded created).
+     * 5. Send task:ack with logId
+     * 6. Submit to agentPool (fire-and-forget)
+     */
+    async handleTaskCreate(data, clientId) {
+        const { taskId } = data;
+        if (this.tasks.has(taskId)) {
+            // Idempotent — duplicate creation returns existing taskId (e.g. client retry)
+            return { taskId };
+        }
+        // ── Early validation: no hooks called if invalid ──────────────────────────
+        if (!this.agentPool) {
+            transportLog(`No AgentPool available, cannot process task ${taskId}`);
+            const error = serializeTaskError(new AgentNotAvailableError());
+            const projectPath = data.projectPath ?? this.resolveClientProjectPath?.(clientId) ?? data.clientCwd;
+            this.transport.sendTo(clientId, TransportTaskEventNames.ERROR, { error, taskId });
+            broadcastToProjectRoom(this.projectRegistry, this.projectRouter, projectPath, TransportTaskEventNames.ERROR, { error, taskId }, clientId);
+            return { taskId };
+        }
+        if (!isValidTaskType(data.type)) {
+            transportLog(`Invalid task type: ${data.type}`);
+            const error = serializeTaskError(new Error(`Invalid task type: ${data.type}`));
+            const projectPath = data.projectPath ?? this.resolveClientProjectPath?.(clientId) ?? data.clientCwd;
+            this.transport.sendTo(clientId, TransportTaskEventNames.ERROR, { error, taskId });
+            broadcastToProjectRoom(this.projectRegistry, this.projectRouter, projectPath, TransportTaskEventNames.ERROR, { error, taskId }, clientId);
+            return { taskId };
+        }
+        // ── Resolve projectPath & store task synchronously ────────────────────────
+        // Resolve projectPath: explicit field > client's registered projectPath > clientCwd.
+        // Client's registered projectPath is the walked-up project root (where .brv/ lives),
+        // while clientCwd is the raw working directory (may be a subdirectory).
+        const projectPath = data.projectPath ?? this.resolveClientProjectPath?.(clientId) ?? data.clientCwd;
+        transportLog(`Task accepted: ${taskId} (type=${data.type}, client=${clientId})`);
+        this.tasks.set(taskId, {
+            clientId,
+            content: data.content,
+            createdAt: Date.now(),
+            ...(data.clientCwd ? { clientCwd: data.clientCwd } : {}),
+            ...(data.files?.length ? { files: data.files } : {}),
+            ...(data.folderPath ? { folderPath: data.folderPath } : {}),
+            ...(projectPath ? { projectPath } : {}),
+            taskId,
+            type: data.type,
+        });
+        // ── Send task:created synchronously (before any await) ────────────────────
+        const createdPayload = {
+            content: data.content,
+            ...(data.clientCwd ? { clientCwd: data.clientCwd } : {}),
+            ...(data.files?.length ? { files: data.files } : {}),
+            taskId,
+            type: data.type,
+        };
+        this.transport.sendTo(clientId, TransportTaskEventNames.CREATED, createdPayload);
+        // Broadcast to other clients in the project room (exclude creator to avoid duplicate)
+        broadcastToProjectRoom(this.projectRegistry, this.projectRouter, projectPath, TransportTaskEventNames.CREATED, createdPayload, clientId);
+        // ── Await lifecycle hooks ─────────────────────────────────────────────────
+        const logId = await this.runCreateHooks(taskId);
+        const task = this.tasks.get(taskId);
+        if (task && logId) {
+            this.tasks.set(taskId, { ...task, logId });
+        }
+        // ── Send task:ack with logId ──────────────────────────────────────────────
+        this.transport.sendTo(clientId, TransportTaskEventNames.ACK, {
+            ...(logId ? { logId } : {}),
+            taskId,
+        });
+        // ── Submit to AgentPool (fire-and-forget) ─────────────────────────────────
+        const executeMsg = {
+            clientId,
+            content: data.content,
+            ...(data.clientCwd ? { clientCwd: data.clientCwd } : {}),
+            ...(data.files?.length ? { files: data.files } : {}),
+            ...(data.folderPath ? { folderPath: data.folderPath } : {}),
+            ...(projectPath ? { projectPath } : {}),
+            taskId,
+            type: data.type,
+        };
+        // eslint-disable-next-line no-void
+        void this.agentPool
+            .submitTask(executeMsg)
+            .then((submitResult) => {
+            if (!submitResult.success) {
+                transportLog(`AgentPool rejected task ${taskId}: ${submitResult.reason} — ${submitResult.message}`);
+                const error = serializeTaskError(new Error(submitResult.message));
+                const rejectedTask = this.tasks.get(taskId) ?? {
+                    clientId,
+                    content: data.content,
+                    createdAt: Date.now(),
+                    taskId,
+                    type: data.type,
+                };
+                this.tasks.delete(taskId);
+                this.transport.sendTo(clientId, TransportTaskEventNames.ERROR, {
+                    ...(rejectedTask.logId ? { logId: rejectedTask.logId } : {}),
+                    error,
+                    taskId,
+                });
+                broadcastToProjectRoom(this.projectRegistry, this.projectRouter, projectPath, TransportTaskEventNames.ERROR, {
+                    ...(rejectedTask.logId ? { logId: rejectedTask.logId } : {}),
+                    error,
+                    taskId,
+                }, clientId);
+                this.notifyHooksError(taskId, submitResult.message, rejectedTask).catch(() => { });
+            }
+        })
+            .catch((error_) => {
+            transportLog(`AgentPool.submitTask threw unexpectedly for task ${taskId}: ${error_ instanceof Error ? error_.message : String(error_)}`);
+            const error = serializeTaskError(error_ instanceof Error ? error_ : new Error(String(error_)));
+            const errorMsg = error_ instanceof Error ? error_.message : String(error_);
+            const thrownTask = this.tasks.get(taskId) ?? {
+                clientId,
+                content: data.content,
+                createdAt: Date.now(),
+                taskId,
+                type: data.type,
+            };
+            this.tasks.delete(taskId);
+            this.transport.sendTo(clientId, TransportTaskEventNames.ERROR, {
+                ...(thrownTask.logId ? { logId: thrownTask.logId } : {}),
+                error,
+                taskId,
+            });
+            broadcastToProjectRoom(this.projectRegistry, this.projectRouter, projectPath, TransportTaskEventNames.ERROR, {
+                ...(thrownTask.logId ? { logId: thrownTask.logId } : {}),
+                error,
+                taskId,
+            }, clientId);
+            this.notifyHooksError(taskId, errorMsg, thrownTask).catch(() => { });
+        });
+        return { ...(logId ? { logId } : {}), taskId };
+    }
+    handleTaskError(data) {
+        const { error, taskId } = data;
+        const task = this.tasks.get(taskId);
+        transportLog(`Task error: ${taskId} - [${error.code}] ${error.message}`);
+        if (task) {
+            this.transport.sendTo(task.clientId, TransportTaskEventNames.ERROR, {
+                ...(task.logId ? { logId: task.logId } : {}),
+                error,
+                taskId,
+            });
+        }
+        broadcastToProjectRoom(this.projectRegistry, this.projectRouter, task?.projectPath, TransportTaskEventNames.ERROR, {
+            ...(task?.logId ? { logId: task.logId } : {}),
+            error,
+            taskId,
+        }, task?.clientId);
+        this.moveToCompleted(taskId);
+        // Notify pool so it can clear busy flag and drain queued tasks
+        if (task?.projectPath) {
+            this.agentPool?.notifyTaskCompleted(task.projectPath);
+        }
+        // Notify hooks (fire-and-forget)
+        if (task) {
+            this.notifyHooksError(taskId, error.message, task).catch(() => { });
+        }
+    }
+    handleTaskStarted(data) {
+        const { taskId } = data;
+        const task = this.tasks.get(taskId);
+        if (task) {
+            this.transport.sendTo(task.clientId, TransportTaskEventNames.STARTED, { taskId });
+            broadcastToProjectRoom(this.projectRegistry, this.projectRouter, task.projectPath, TransportTaskEventNames.STARTED, {
+                content: task.content,
+                ...(task.clientCwd ? { clientCwd: task.clientCwd } : {}),
+                ...(task.files?.length ? { files: task.files } : {}),
+                taskId,
+                type: task.type,
+            }, task.clientId);
+        }
+        else {
+            // No task context — cannot determine project room, skip broadcast
+            transportLog(`Task started but no task context found: ${taskId}`);
+        }
+    }
+    /**
+     * Move a task to the completed tasks map with grace period cleanup.
+     */
+    moveToCompleted(taskId) {
+        const task = this.tasks.get(taskId);
+        if (task) {
+            this.completedTasks.set(taskId, { completedAt: Date.now(), task });
+            this.tasks.delete(taskId);
+            setTimeout(() => {
+                this.completedTasks.delete(taskId);
+            }, TASK_CLEANUP_GRACE_PERIOD_MS);
+        }
+    }
+    /**
+     * Notify all hooks of task cancellation.
+     * Each hook is called independently; errors are caught and logged.
+     * cleanup() is called for each hook after onTaskCancelled.
+     */
+    async notifyHooksCancelled(taskId, task) {
+        await Promise.allSettled(this.lifecycleHooks.map(async (hook) => {
+            try {
+                await hook.onTaskCancelled?.(taskId, task);
+            }
+            catch (error) {
+                transportLog(`LifecycleHook.onTaskCancelled error for ${taskId}: ${error instanceof Error ? error.message : String(error)}`);
+            }
+            finally {
+                hook.cleanup?.(taskId);
+            }
+        }));
+    }
+    /**
+     * Notify all hooks of task completion.
+     * Each hook is called independently; errors are caught and logged.
+     * cleanup() is called for each hook after onTaskCompleted.
+     */
+    async notifyHooksCompleted(taskId, result, task) {
+        await Promise.allSettled(this.lifecycleHooks.map(async (hook) => {
+            try {
+                await hook.onTaskCompleted?.(taskId, result, task);
+            }
+            catch (error) {
+                transportLog(`LifecycleHook.onTaskCompleted error for ${taskId}: ${error instanceof Error ? error.message : String(error)}`);
+            }
+            finally {
+                hook.cleanup?.(taskId);
+            }
+        }));
+    }
+    /**
+     * Notify all hooks of task error.
+     * Each hook is called independently; errors are caught and logged.
+     * cleanup() is called for each hook after onTaskError.
+     */
+    async notifyHooksError(taskId, errorMessage, task) {
+        await Promise.allSettled(this.lifecycleHooks.map(async (hook) => {
+            try {
+                await hook.onTaskError?.(taskId, errorMessage, task);
+            }
+            catch (error) {
+                transportLog(`LifecycleHook.onTaskError error for ${taskId}: ${error instanceof Error ? error.message : String(error)}`);
+            }
+            finally {
+                hook.cleanup?.(taskId);
+            }
+        }));
+    }
+    registerLlmEvent(eventName) {
+        this.transport.onRequest(eventName, (data) => {
+            if (!hasTaskId(data))
+                return;
+            this.routeLlmEvent(eventName, data);
+        });
+    }
+    /**
+     * Generic handler for routing LLM events from Agent to clients.
+     * Checks both active and recently completed tasks (within grace period).
+     * onToolResult hooks are called only for ACTIVE tasks (not grace-period).
+     */
+    routeLlmEvent(eventName, data) {
+        const { taskId, ...rest } = data;
+        const activeTask = this.tasks.get(taskId);
+        const task = activeTask ?? this.completedTasks.get(taskId)?.task;
+        if (!task) {
+            return;
+        }
+        // Notify onToolResult hooks only for active tasks
+        if (activeTask && eventName === LlmEventNames.TOOL_RESULT) {
+            for (const hook of this.lifecycleHooks) {
+                try {
+                    hook.onToolResult?.(taskId, data);
+                }
+                catch (error) {
+                    transportLog(`LifecycleHook.onToolResult error for ${taskId}: ${error instanceof Error ? error.message : String(error)}`);
+                }
+            }
+        }
+        this.transport.sendTo(task.clientId, eventName, { taskId, ...rest });
+        broadcastToProjectRoom(this.projectRegistry, this.projectRouter, task.projectPath, eventName, { taskId, ...rest }, task.clientId);
+    }
+    /**
+     * Run all onTaskCreate hooks and return the first logId.
+     * Each hook is called independently; errors are caught and logged.
+     */
+    async runCreateHooks(taskId) {
+        if (this.lifecycleHooks.length === 0)
+            return undefined;
+        const task = this.tasks.get(taskId);
+        if (!task)
+            return undefined;
+        const logIds = await Promise.all(this.lifecycleHooks.map(async (hook) => {
+            if (!hook.onTaskCreate)
+                return;
+            try {
+                const result = await hook.onTaskCreate(task);
+                return result?.logId;
+            }
+            catch (error) {
+                transportLog(`LifecycleHook.onTaskCreate error for ${taskId}: ${error instanceof Error ? error.message : String(error)}`);
+            }
+        }));
+        return logIds.find((id) => typeof id === 'string');
+    }
+}