diff --git a/apps/mobile/src/app/task/[id].tsx b/apps/mobile/src/app/task/[id].tsx index 35d95ead55..c06bd0da15 100644 --- a/apps/mobile/src/app/task/[id].tsx +++ b/apps/mobile/src/app/task/[id].tsx @@ -1,4 +1,8 @@ import { Text } from "@components/text"; +import type { + CloudTaskQueuedMessage, + CloudTaskQueueMoveDirection, +} from "@posthog/core/sessions/cloudTaskQueue"; import { countUserMessages, getSessionActivityPhase, @@ -45,17 +49,18 @@ import { useQueuedCount, useToggleMessagingMode, } from "@/features/tasks/hooks/useMessagingMode"; +import { useTaskMessageQueue } from "@/features/tasks/hooks/useTaskMessageQueue"; import { taskKeys } from "@/features/tasks/hooks/useTasks"; -import { - type MoveDirection, - type QueuedMessage, - useMessageQueueStore, -} from "@/features/tasks/stores/messageQueueStore"; +import { taskMessageQueue } from "@/features/tasks/lib/taskMessageQueue"; +import { taskSessionActions } from "@/features/tasks/services/taskSessionService"; import { pendingTaskPromptStoreApi, usePendingTaskPrompt, } from "@/features/tasks/stores/pendingTaskPromptStore"; -import { useTaskSessionStore } from "@/features/tasks/stores/taskSessionStore"; +import { + getTaskSession, + useTaskSessionStore, +} from "@/features/tasks/stores/taskSessionStore"; import { useTaskStore } from "@/features/tasks/stores/taskStore"; import { confirmStopRun } from "@/features/tasks/utils/archiveGuard"; import { useScreenInsets } from "@/hooks/useScreenInsets"; @@ -69,6 +74,7 @@ import { getPostHogApiClient } from "@/lib/posthogApiClient"; import { useThemeColors } from "@/lib/theme"; const log = logger.scope("task-detail"); +type QueuedMessage = CloudTaskQueuedMessage; function getFirstParam(value?: string | string[]): string | undefined { return Array.isArray(value) ? value[0] : value; @@ -107,12 +113,13 @@ export default function TaskDetailScreen() { sendInterrupting, sendPermissionResponse, setConfigOption, - getSessionForTask, - setFocusedTaskId, steerQueuedMessage, flushQueuedMessagesIfIdle, stopRun, - } = useTaskSessionStore(); + } = taskSessionActions; + const setFocusedTaskId = useTaskSessionStore( + (state) => state.setFocusedTaskId, + ); useEffect(() => { if (!taskId) return; @@ -127,7 +134,7 @@ export default function TaskDetailScreen() { // Cleared when the screen unmounts. Matches the desktop super-property. useActiveTaskAnalyticsContext(task?.signal_report ?? null); - const session = taskId ? getSessionForTask(taskId) : undefined; + const session = taskId ? getTaskSession(taskId) : undefined; // Optimistic echo set by the new-task screen (or the terminal-resume path // below) so the user's prompt appears in the thread immediately, before @@ -176,9 +183,7 @@ export default function TaskDetailScreen() { const messagingMode = useMessagingMode(taskId); const queuedCount = useQueuedCount(taskId); - const editingQueuedId = useMessageQueueStore((s) => - taskId ? s.editingByTaskId[taskId] : undefined, - ); + const { editingId: editingQueuedId } = useTaskMessageQueue(taskId); const toggleMessagingMode = useToggleMessagingMode(taskId); const analytics = useAnalytics(); @@ -369,11 +374,13 @@ export default function TaskDetailScreen() { // Saving an in-place edit: overwrite the queued message and release the // drain hold. If the turn already ended while editing, flush now — the // turn-end drain won't fire again on its own. - const queue = useMessageQueueStore.getState(); - const editingId = queue.editingByTaskId[taskId]; + const editingId = taskMessageQueue.getSnapshot().editingByTaskId[taskId]; if (editingId) { - queue.update(taskId, editingId, { content: text, attachments }); - queue.clearEditing(taskId); + taskMessageQueue.update(taskId, editingId, { + content: text, + attachments, + }); + taskMessageQueue.clearEditing(taskId); flushQueuedMessagesIfIdle(taskId); return; } @@ -395,7 +402,7 @@ export default function TaskDetailScreen() { // Steer interrupts the turn and resends right away. if (session?.isPromptPending) { if (messagingMode === "queue") { - useMessageQueueStore.getState().enqueue(taskId, text, attachments); + taskMessageQueue.enqueue(taskId, text, attachments); return; } sendInterrupting(taskId, text, attachments) @@ -448,7 +455,7 @@ export default function TaskDetailScreen() { const handleEditQueued = useCallback( (message: QueuedMessage) => { if (!taskId) return; - useMessageQueueStore.getState().setEditing(taskId, message.id); + taskMessageQueue.setEditing(taskId, message.id); setRestoredDraft({ text: message.content, attachments: message.attachments, @@ -459,16 +466,16 @@ export default function TaskDetailScreen() { const handleCancelEdit = useCallback(() => { if (!taskId) return; - useMessageQueueStore.getState().clearEditing(taskId); + taskMessageQueue.clearEditing(taskId); setRestoredDraft({ text: "", attachments: [] }); flushQueuedMessagesIfIdle(taskId); }, [taskId, flushQueuedMessagesIfIdle]); const handleMoveQueued = useCallback( - (message: QueuedMessage, direction: MoveDirection) => { + (message: QueuedMessage, direction: CloudTaskQueueMoveDirection) => { if (!taskId) return; Haptics.selectionAsync(); - useMessageQueueStore.getState().move(taskId, message.id, direction); + taskMessageQueue.move(taskId, message.id, direction); }, [taskId], ); @@ -477,8 +484,8 @@ export default function TaskDetailScreen() { (message: QueuedMessage) => { if (!taskId) return; const wasEditing = - useMessageQueueStore.getState().editingByTaskId[taskId] === message.id; - useMessageQueueStore.getState().remove(taskId, message.id); + taskMessageQueue.getSnapshot().editingByTaskId[taskId] === message.id; + taskMessageQueue.remove(taskId, message.id); if (wasEditing) setRestoredDraft({ text: "", attachments: [] }); }, [taskId], @@ -525,7 +532,7 @@ export default function TaskDetailScreen() { const handleStopRun = useCallback(() => { if (!taskId) return; confirmStopRun(() => { - const promptsSent = countUserMessages(getSessionForTask(taskId)?.events); + const promptsSent = countUserMessages(getTaskSession(taskId)?.events); stopRun(taskId) .then((ok) => { if (ok) { @@ -543,7 +550,7 @@ export default function TaskDetailScreen() { }) .catch(() => {}); }); - }, [taskId, stopRun, analytics, getSessionForTask]); + }, [taskId, stopRun, analytics]); const canStopRun = !!task && diff --git a/apps/mobile/src/features/tasks/api.ts b/apps/mobile/src/features/tasks/api.ts index 2408f688c6..56180264de 100644 --- a/apps/mobile/src/features/tasks/api.ts +++ b/apps/mobile/src/features/tasks/api.ts @@ -1,9 +1,4 @@ -import type { - Adapter, - StoredLogEntry, - Task, - TaskRun, -} from "@posthog/shared"; +import type { Adapter, StoredLogEntry, Task, TaskRun } from "@posthog/shared"; import { fetch } from "expo/fetch"; import { authedFetch, diff --git a/apps/mobile/src/features/tasks/components/TaskSessionView.tsx b/apps/mobile/src/features/tasks/components/TaskSessionView.tsx index fb4c23ed13..70c2611826 100644 --- a/apps/mobile/src/features/tasks/components/TaskSessionView.tsx +++ b/apps/mobile/src/features/tasks/components/TaskSessionView.tsx @@ -1,3 +1,7 @@ +import type { + PortableSessionEvent as SessionEvent, + PortableSessionNotification as SessionNotification, +} from "@posthog/core/sessions/portableSessionEvents"; import { ArrowDown, Brain, @@ -25,8 +29,6 @@ import { useThemeColors } from "@/lib/theme"; import type { CloudPendingPermissionRequest, PlanEntry, - SessionEvent, - SessionNotification, SessionNotificationAttachment, } from "../types"; import { PlanApprovalCard } from "./PlanApprovalCard"; diff --git a/apps/mobile/src/features/tasks/composer/QueuedMessagesDock.tsx b/apps/mobile/src/features/tasks/composer/QueuedMessagesDock.tsx index 7b17729167..ae361d4330 100644 --- a/apps/mobile/src/features/tasks/composer/QueuedMessagesDock.tsx +++ b/apps/mobile/src/features/tasks/composer/QueuedMessagesDock.tsx @@ -1,4 +1,8 @@ import { Text } from "@components/text"; +import type { + CloudTaskQueuedMessage, + CloudTaskQueueMoveDirection, +} from "@posthog/core/sessions/cloudTaskQueue"; import { CaretDown, CaretUp, @@ -12,11 +16,10 @@ import { type ReactNode, useState } from "react"; import { Pressable, View } from "react-native"; import { SheetContainer } from "@/components/SheetContainer"; import { useThemeColors } from "@/lib/theme"; -import { - type MoveDirection, - type QueuedMessage, - useMessageQueueStore, -} from "../stores/messageQueueStore"; +import { useTaskMessageQueue } from "../hooks/useTaskMessageQueue"; +import type { PendingAttachment } from "./attachments/types"; + +type QueuedMessage = CloudTaskQueuedMessage; interface QueuedMessagesDockProps { taskId: string; @@ -24,7 +27,10 @@ interface QueuedMessagesDockProps { onSteer: (message: QueuedMessage) => void; onEdit: (message: QueuedMessage) => void; onDiscard: (message: QueuedMessage) => void; - onMove: (message: QueuedMessage, direction: MoveDirection) => void; + onMove: ( + message: QueuedMessage, + direction: CloudTaskQueueMoveDirection, + ) => void; } function previewText(message: QueuedMessage): string { @@ -42,11 +48,10 @@ export function QueuedMessagesDock({ onMove, }: QueuedMessagesDockProps) { const themeColors = useThemeColors(); - const queued = useMessageQueueStore((s) => s.queuesByTaskId[taskId]); - const editingId = useMessageQueueStore((s) => s.editingByTaskId[taskId]); + const { messages: queued, editingId } = useTaskMessageQueue(taskId); const [activeId, setActiveId] = useState(null); - if (!queued || queued.length === 0) return null; + if (queued.length === 0) return null; const active = queued.find((m) => m.id === activeId) ?? null; return ( diff --git a/apps/mobile/src/features/tasks/hooks/useMessagingMode.ts b/apps/mobile/src/features/tasks/hooks/useMessagingMode.ts index 9e6cae9277..4be6801b02 100644 --- a/apps/mobile/src/features/tasks/hooks/useMessagingMode.ts +++ b/apps/mobile/src/features/tasks/hooks/useMessagingMode.ts @@ -1,10 +1,10 @@ import { useCallback } from "react"; -import { useMessageQueueStore } from "../stores/messageQueueStore"; +import { taskSessionActions } from "../services/taskSessionService"; import { type MessagingMode, useMessagingModeStore, } from "../stores/messagingModeStore"; -import { useTaskSessionStore } from "../stores/taskSessionStore"; +import { useTaskMessageQueue } from "./useTaskMessageQueue"; /** Effective mode for a task: per-task override, else the global default. */ export function useMessagingMode(taskId: string | undefined): MessagingMode { @@ -12,7 +12,7 @@ export function useMessagingMode(taskId: string | undefined): MessagingMode { } export function useQueuedCount(taskId: string | undefined): number { - return useMessageQueueStore((s) => (taskId ? s.getQueue(taskId).length : 0)); + return useTaskMessageQueue(taskId ?? "").messages.length; } /** @@ -27,7 +27,7 @@ export function useToggleMessagingMode(taskId: string | undefined): () => void { const next: MessagingMode = mode === "steer" ? "queue" : "steer"; useMessagingModeStore.getState().setMode(taskId, next); if (next === "steer") { - void useTaskSessionStore.getState().flushQueuedMessages(taskId); + void taskSessionActions.flushQueuedMessages(taskId); } }, [taskId, mode]); } diff --git a/apps/mobile/src/features/tasks/hooks/useTaskMessageQueue.ts b/apps/mobile/src/features/tasks/hooks/useTaskMessageQueue.ts new file mode 100644 index 0000000000..c3d4833f05 --- /dev/null +++ b/apps/mobile/src/features/tasks/hooks/useTaskMessageQueue.ts @@ -0,0 +1,20 @@ +import type { CloudTaskQueuedMessage } from "@posthog/core/sessions/cloudTaskQueue"; +import { useSyncExternalStore } from "react"; +import type { PendingAttachment } from "../composer/attachments/types"; +import { taskMessageQueue } from "../lib/taskMessageQueue"; + +interface TaskMessageQueueSelection { + messages: readonly CloudTaskQueuedMessage[]; + editingId: string | undefined; +} + +export function useTaskMessageQueue(taskId: string): TaskMessageQueueSelection { + const snapshot = useSyncExternalStore( + taskMessageQueue.subscribe, + taskMessageQueue.getSnapshot, + ); + return { + messages: snapshot.queuesByTaskId[taskId] ?? [], + editingId: snapshot.editingByTaskId[taskId], + }; +} diff --git a/apps/mobile/src/features/tasks/index.ts b/apps/mobile/src/features/tasks/index.ts index 7da4db747e..2bb6820f00 100644 --- a/apps/mobile/src/features/tasks/index.ts +++ b/apps/mobile/src/features/tasks/index.ts @@ -16,10 +16,7 @@ export { useUpdateTask, } from "./hooks/useTasks"; // Stores -export { - type TaskSession, - useTaskSessionStore, -} from "./stores/taskSessionStore"; +export { useTaskSessionStore } from "./stores/taskSessionStore"; export { useTaskStore } from "./stores/taskStore"; // Types diff --git a/apps/mobile/src/features/tasks/lib/cloudTaskStream.test.ts b/apps/mobile/src/features/tasks/lib/cloudTaskStream.test.ts new file mode 100644 index 0000000000..a3fc53a610 --- /dev/null +++ b/apps/mobile/src/features/tasks/lib/cloudTaskStream.test.ts @@ -0,0 +1,54 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; + +const mocks = vi.hoisted(() => ({ + engine: { + off: vi.fn(), + on: vi.fn(), + reconnectIfDisconnected: vi.fn(), + unwatch: vi.fn(), + watch: vi.fn(), + }, +})); + +vi.mock("@posthog/core/cloud-task/cloud-task-engine", () => ({ + createCloudTaskEngine: () => mocks.engine, +})); + +vi.mock("@posthog/core/cloud-task/schemas", () => ({ + CloudTaskEvent: { Update: "cloud-task-update" }, +})); + +vi.mock("expo/fetch", () => ({ fetch: vi.fn() })); + +vi.mock("@/lib/api", () => ({ + authedFetch: vi.fn(), + getBaseUrl: () => "https://app.posthog.test", + getProjectId: () => 42, +})); + +vi.mock("@/lib/logger", () => ({ + logger: { scope: vi.fn() }, +})); + +import { watchCloudTask } from "./cloudTaskStream"; + +describe("watchCloudTask", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("asks the shared engine to reconnect only when disconnected", () => { + const handle = watchCloudTask({ + taskId: "task-1", + runId: "run-1", + onUpdate: vi.fn(), + }); + + handle.reconnectIfDisconnected(); + + expect(mocks.engine.reconnectIfDisconnected).toHaveBeenCalledWith( + "task-1", + "run-1", + ); + }); +}); diff --git a/apps/mobile/src/features/tasks/lib/cloudTaskStream.ts b/apps/mobile/src/features/tasks/lib/cloudTaskStream.ts index 7ec96e3851..ac761b6bc9 100644 --- a/apps/mobile/src/features/tasks/lib/cloudTaskStream.ts +++ b/apps/mobile/src/features/tasks/lib/cloudTaskStream.ts @@ -1,122 +1,18 @@ import { - type CloudTaskUpdatePayload, - isTerminalStatus, - type StoredLogEntry, - type TaskRun, - type TaskRunStatus, -} from "@posthog/shared"; + type CloudTaskEngine, + type CloudTaskFetch, + createCloudTaskEngine, +} from "@posthog/core/cloud-task/cloud-task-engine"; +import { CloudTaskEvent } from "@posthog/core/cloud-task/schemas"; +import type { CloudTaskUpdatePayload } from "@posthog/shared"; import { fetch } from "expo/fetch"; -import { createTimeoutSignal } from "@/lib/api"; -import { logger } from "@/lib/logger"; -import { - fetchSessionLogs, - getTaskRun, - HttpError, - streamCloudTask, -} from "../api"; import { - isKeepaliveEvent, - isPermissionRequestEvent, - isSseErrorEvent, - isTaskRunStateEvent, - type TaskRunStateEvent, -} from "../types"; -import { parseSessionLogs } from "../utils/parseSessionLogs"; -import { type SseEvent, SseEventParser } from "./sseParser"; - -const log = logger.scope("cloud-task-stream"); - -const MAX_SSE_RECONNECT_ATTEMPTS = 5; -const SSE_RECONNECT_BASE_DELAY_MS = 2_000; -const SSE_RECONNECT_MAX_DELAY_MS = 30_000; -const EVENT_BATCH_FLUSH_MS = 16; -const EVENT_BATCH_MAX_SIZE = 50; -const SESSION_LOG_PAGE_LIMIT = 5_000; - -interface CloudTaskConnectionError { - title: string; - message: string; - retryable: boolean; - autoRetry?: boolean; -} - -class CloudTaskStreamError extends Error { - constructor( - message: string, - public readonly details: CloudTaskConnectionError, - public readonly status?: number, - ) { - super(message); - this.name = "CloudTaskStreamError"; - } -} - -function createStreamStatusError(status: number): CloudTaskStreamError { - switch (status) { - case 401: - return new CloudTaskStreamError( - "Cloud authentication expired", - { - title: "Cloud authentication expired", - message: "Please reauthenticate and retry the cloud run stream.", - retryable: true, - autoRetry: false, - }, - status, - ); - case 403: - return new CloudTaskStreamError( - "Cloud access denied", - { - title: "Cloud access denied", - message: - "You no longer have access to this cloud run. Reauthenticate and retry.", - retryable: true, - autoRetry: false, - }, - status, - ); - case 404: - return new CloudTaskStreamError( - "Cloud run not found", - { - title: "Cloud run not found", - message: - "This cloud run could not be found. It may have been deleted or moved.", - retryable: false, - autoRetry: false, - }, - status, - ); - case 406: - return new CloudTaskStreamError( - "Cloud stream unavailable", - { - title: "Cloud stream unavailable", - message: - "The backend rejected the live stream request. Restart the backend and retry.", - retryable: true, - autoRetry: false, - }, - status, - ); - default: - return new CloudTaskStreamError( - `Stream request failed with status ${status}`, - { - title: "Cloud stream failed", - message: `The cloud stream request failed with status ${status}. Retry to reconnect.`, - retryable: true, - autoRetry: true, - }, - status, - ); - } -} - -function shouldFailWatcherForFetchStatus(status: number): boolean { - return status === 401 || status === 403 || status === 404; -} + authedFetch, + type FetchInit, + getBaseUrl, + getProjectId, +} from "@/lib/api"; +import { logger } from "@/lib/logger"; export interface WatchCloudTaskOptions { taskId: string; @@ -129,786 +25,78 @@ export interface WatchCloudTaskHandle { reconnectIfDisconnected: () => void; } -interface WatcherState { - taskId: string; - runId: string; - onUpdate: (update: CloudTaskUpdatePayload) => void; - stopped: boolean; - sseAbortController: AbortController | null; - reconnectTimeoutId: ReturnType | null; - batchFlushTimeoutId: ReturnType | null; - pendingLogEntries: StoredLogEntry[]; - totalEntryCount: number; - reconnectAttempts: number; - lastEventId: string | null; - lastStatus: TaskRunStatus | null; - lastStage: string | null; - lastOutput: Record | null; - lastErrorMessage: string | null; - lastBranch: string | null; - lastStatusUpdatedAt: string | null; - isBootstrapping: boolean; - hasEmittedSnapshot: boolean; - bufferedLogBatches: StoredLogEntry[][]; - failed: boolean; - needsPostBootstrapReconnect: boolean; - needsStopAfterBootstrap: boolean; -} - -export function watchCloudTask( - options: WatchCloudTaskOptions, -): WatchCloudTaskHandle { - const watcher: WatcherState = { - taskId: options.taskId, - runId: options.runId, - onUpdate: options.onUpdate, - stopped: false, - sseAbortController: null, - reconnectTimeoutId: null, - batchFlushTimeoutId: null, - pendingLogEntries: [], - totalEntryCount: 0, - reconnectAttempts: 0, - lastEventId: null, - lastStatus: null, - lastStage: null, - lastOutput: null, - lastErrorMessage: null, - lastBranch: null, - lastStatusUpdatedAt: null, - isBootstrapping: false, - hasEmittedSnapshot: false, - bufferedLogBatches: [], - failed: false, - needsPostBootstrapReconnect: false, - needsStopAfterBootstrap: false, - }; - - void bootstrapWatcher(watcher); - - return { - stop: () => stopWatcher(watcher), - reconnectIfDisconnected: () => { - if ( - watcher.stopped || - watcher.failed || - isTerminalStatus(watcher.lastStatus) - ) { - return; - } - if (watcher.sseAbortController || watcher.reconnectTimeoutId) { - return; - } - log.debug("Force reconnect after suspension", { runId: watcher.runId }); - watcher.reconnectAttempts = 0; - void connectSse(watcher, { - startLatest: !watcher.lastEventId, - }); +const mobileCloudTaskAnalytics = { + initialize: () => {}, + track: () => {}, + identify: () => {}, + setCurrentUserId: () => {}, + getCurrentUserId: () => null, + getOrCreateSessionId: () => "mobile-cloud-task", + resetUser: () => {}, + captureException: () => {}, + flush: async () => {}, + shutdown: async () => {}, +}; + +let cloudTaskEngine: CloudTaskEngine | null = null; + +function getCloudTaskEngine(): CloudTaskEngine { + if (cloudTaskEngine) { + return cloudTaskEngine; + } + + cloudTaskEngine = createCloudTaskEngine({ + auth: { + authenticatedFetch: (url, init) => + authedFetch(url, init as FetchInit | undefined), + getCloudContext: async () => ({ + apiHost: getBaseUrl(), + teamId: getProjectId(), + }), }, - }; -} - -function stopWatcher(watcher: WatcherState): void { - if (watcher.stopped) return; - watcher.stopped = true; - - watcher.sseAbortController?.abort(); - watcher.sseAbortController = null; - - if (watcher.reconnectTimeoutId) { - clearTimeout(watcher.reconnectTimeoutId); - watcher.reconnectTimeoutId = null; - } - - if (watcher.batchFlushTimeoutId) { - clearTimeout(watcher.batchFlushTimeoutId); - watcher.batchFlushTimeoutId = null; - } - - // Drop any unflushed batches; the consumer is gone. - watcher.pendingLogEntries = []; - watcher.bufferedLogBatches = []; -} - -async function bootstrapWatcher(watcher: WatcherState): Promise { - if (watcher.stopped) return; - - watcher.failed = false; - watcher.needsPostBootstrapReconnect = false; - watcher.needsStopAfterBootstrap = false; - - const run = await fetchTaskRunState(watcher); - if (watcher.stopped || watcher.failed) return; - - if (!run) { - failWatcher(watcher, { - title: "Failed to load cloud run", - message: "Could not fetch the cloud run state. Retry to reconnect.", - retryable: true, - }); - return; - } - - applyTaskRunState(watcher, run); - - if (isTerminalStatus(run.status)) { - const historicalEntries = await fetchHistoricalEntries(watcher, run); - if (watcher.stopped || watcher.failed) return; - if (!historicalEntries) { - failWatcher(watcher, { - title: "Failed to load task history", - message: - "Could not load the persisted cloud task logs. Retry to reconnect.", - retryable: true, - }); - return; - } - - watcher.totalEntryCount = historicalEntries.length; - watcher.hasEmittedSnapshot = true; - emitSnapshot(watcher, historicalEntries); - stopWatcher(watcher); - return; - } - - watcher.isBootstrapping = true; - watcher.bufferedLogBatches = []; - void connectSse(watcher, { startLatest: true }); - - const historicalEntries = await fetchHistoricalEntries(watcher, run); - if (watcher.stopped || watcher.failed) return; - if (!historicalEntries) { - failWatcher(watcher, { - title: "Failed to load cloud run history", - message: - "Could not load the existing cloud run logs. Retry to reconnect.", - retryable: true, - }); - return; - } - - // Flush any pending live entries into the bootstrap buffer before snapshot. - flushLogBatch(watcher); - - watcher.totalEntryCount = historicalEntries.length; - watcher.hasEmittedSnapshot = true; - emitSnapshot(watcher, historicalEntries); - - watcher.isBootstrapping = false; - drainBufferedLogBatches(watcher, historicalEntries); - - if (watcher.failed) return; - - if (watcher.needsStopAfterBootstrap || isTerminalStatus(watcher.lastStatus)) { - watcher.needsStopAfterBootstrap = false; - stopWatcher(watcher); - return; - } - - if (watcher.needsPostBootstrapReconnect) { - watcher.needsPostBootstrapReconnect = false; - scheduleReconnect(watcher, undefined, { countAttempt: false }); - } - - void verifyPostBootstrapStatus(watcher); -} - -async function verifyPostBootstrapStatus(watcher: WatcherState): Promise { - if (watcher.stopped) return; - if (isTerminalStatus(watcher.lastStatus)) return; - - const run = await fetchTaskRunState(watcher); - if (watcher.stopped || !run) return; - - if (!applyTaskRunState(watcher, run)) return; - if (isTerminalStatus(watcher.lastStatus)) return; - - emitStatus(watcher); -} - -async function connectSse( - watcher: WatcherState, - options?: { startLatest?: boolean }, -): Promise { - if (watcher.stopped) return; - - const controller = new AbortController(); - watcher.sseAbortController = controller; - - const parser = new SseEventParser(); - const decoder = new TextDecoder(); - - try { - const response = await streamCloudTask(watcher.taskId, watcher.runId, { - lastEventId: watcher.lastEventId, - startLatest: options?.startLatest, - signal: controller.signal, - }); - - if (!response.ok) { - throw createStreamStatusError(response.status); - } - - if (!response.body) { - throw new Error("Stream response did not include a body"); - } - - const reader = response.body.getReader(); - - while (true) { - const { done, value } = await reader.read(); - if (done) { - break; - } - - if (!value) { - continue; - } - - const chunk = decoder.decode(value, { stream: true }); - const events = parser.parse(chunk); - for (const event of events) { - handleSseEvent(watcher, event); - if (watcher.failed) return; - } - } - - const trailingEvents = parser.parse(decoder.decode()); - for (const event of trailingEvents) { - handleSseEvent(watcher, event); - if (watcher.failed) return; - } - - flushLogBatch(watcher); - - if (controller.signal.aborted) { - return; - } - - await handleStreamCompletion(watcher, { reconnectIfNonTerminal: true }); - } catch (error) { - flushLogBatch(watcher); - - if (controller.signal.aborted) { - return; - } - - if ( - error instanceof CloudTaskStreamError && - error.details.autoRetry === false - ) { - failWatcher(watcher, error.details); - return; - } - - const errorMessage = - error instanceof Error ? error.message : "Unknown stream error"; - log.warn("Cloud task stream error", { - runId: watcher.runId, - error: errorMessage, - }); - await handleStreamCompletion(watcher, { - reconnectIfNonTerminal: true, - reconnectError: error, - countReconnectAttempt: true, - }); - } finally { - if (watcher.sseAbortController === controller) { - watcher.sseAbortController = null; - } - } -} - -function handleSseEvent(watcher: WatcherState, event: SseEvent): void { - if (watcher.failed || watcher.stopped) return; - - if (event.id) { - watcher.lastEventId = event.id; - } - - if (event.event === "error") { - const message = isSseErrorEvent(event.data) - ? event.data.error - : "Unknown stream error"; - throw new Error(message); - } - - if (event.event === "keepalive" || isKeepaliveEvent(event.data)) { - return; - } - - watcher.reconnectAttempts = 0; - - if (isTaskRunStateEvent(event.data)) { - if (applyTaskRunState(watcher, event.data)) { - if (!watcher.isBootstrapping && !isTerminalStatus(watcher.lastStatus)) { - emitStatus(watcher); - } - } - return; - } - - if (isPermissionRequestEvent(event.data)) { - watcher.onUpdate({ - taskId: watcher.taskId, - runId: watcher.runId, - kind: "permission_request", - requestId: event.data.requestId, - toolCall: event.data.toolCall, - options: event.data.options, - }); - return; - } - - // StoredLogEntry always has a string `type`. Anything else is a server - // event the mobile client doesn't understand yet — drop it instead of - // forwarding a malformed entry to convertStoredEntriesToEvents. - if ( - typeof event.data !== "object" || - event.data === null || - typeof (event.data as { type?: unknown }).type !== "string" - ) { - log.warn("Skipping unrecognized SSE event", { - runId: watcher.runId, - eventName: event.event, - }); - return; - } - - watcher.pendingLogEntries.push(event.data as StoredLogEntry); - if (watcher.pendingLogEntries.length >= EVENT_BATCH_MAX_SIZE) { - flushLogBatch(watcher); - return; - } - - if (!watcher.batchFlushTimeoutId) { - watcher.batchFlushTimeoutId = setTimeout(() => { - watcher.batchFlushTimeoutId = null; - flushLogBatch(watcher); - }, EVENT_BATCH_FLUSH_MS); - } -} - -function flushLogBatch(watcher: WatcherState): void { - if (watcher.pendingLogEntries.length === 0) return; - - if (watcher.batchFlushTimeoutId) { - clearTimeout(watcher.batchFlushTimeoutId); - watcher.batchFlushTimeoutId = null; - } - - const entries = watcher.pendingLogEntries; - watcher.pendingLogEntries = []; - - if (watcher.isBootstrapping) { - watcher.bufferedLogBatches.push(entries); - return; - } - - watcher.totalEntryCount += entries.length; - watcher.onUpdate({ - taskId: watcher.taskId, - runId: watcher.runId, - kind: "logs", - newEntries: entries, - totalEntryCount: watcher.totalEntryCount, - }); -} - -function drainBufferedLogBatches( - watcher: WatcherState, - historicalEntries: StoredLogEntry[], -): void { - if (watcher.bufferedLogBatches.length === 0) return; - - // Content-based dedup because SSE IDs (Redis stream IDs) don't exist in - // the S3-backed historical entries — the JSON payload is the only shared key. - const historicalCounts = new Map(); - for (const entry of historicalEntries) { - const serialized = JSON.stringify(entry); - historicalCounts.set( - serialized, - (historicalCounts.get(serialized) ?? 0) + 1, - ); - } - - for (const entries of watcher.bufferedLogBatches) { - const dedupedEntries = entries.filter((entry) => { - const serialized = JSON.stringify(entry); - const remaining = historicalCounts.get(serialized) ?? 0; - if (remaining <= 0) return true; - historicalCounts.set(serialized, remaining - 1); - return false; - }); - - if (dedupedEntries.length === 0) continue; - - watcher.totalEntryCount += dedupedEntries.length; - watcher.onUpdate({ - taskId: watcher.taskId, - runId: watcher.runId, - kind: "logs", - newEntries: dedupedEntries, - totalEntryCount: watcher.totalEntryCount, - }); - } - - watcher.bufferedLogBatches = []; -} - -function emitSnapshot(watcher: WatcherState, entries: StoredLogEntry[]): void { - watcher.onUpdate({ - taskId: watcher.taskId, - runId: watcher.runId, - kind: "snapshot", - newEntries: entries, - totalEntryCount: watcher.totalEntryCount, - status: watcher.lastStatus ?? undefined, - stage: watcher.lastStage, - output: watcher.lastOutput, - errorMessage: watcher.lastErrorMessage, - branch: watcher.lastBranch, - }); -} - -function emitStatus(watcher: WatcherState): void { - watcher.onUpdate({ - taskId: watcher.taskId, - runId: watcher.runId, - kind: "status", - status: watcher.lastStatus ?? undefined, - stage: watcher.lastStage, - output: watcher.lastOutput, - errorMessage: watcher.lastErrorMessage, - branch: watcher.lastBranch, - }); -} - -function failWatcher( - watcher: WatcherState, - error: CloudTaskConnectionError, -): void { - if (watcher.stopped) return; - - watcher.failed = true; - watcher.isBootstrapping = false; - watcher.pendingLogEntries = []; - watcher.bufferedLogBatches = []; - - if (watcher.reconnectTimeoutId) { - clearTimeout(watcher.reconnectTimeoutId); - watcher.reconnectTimeoutId = null; - } - - if (watcher.batchFlushTimeoutId) { - clearTimeout(watcher.batchFlushTimeoutId); - watcher.batchFlushTimeoutId = null; - } - - watcher.sseAbortController?.abort(); - watcher.sseAbortController = null; - - watcher.onUpdate({ - taskId: watcher.taskId, - runId: watcher.runId, - kind: "error", - errorTitle: error.title, - errorMessage: error.message, - retryable: error.retryable, + analytics: mobileCloudTaskAnalytics, + logger, + streamFetch: fetch as CloudTaskFetch, }); -} - -function scheduleReconnect( - watcher: WatcherState, - error?: unknown, - options: { countAttempt?: boolean } = {}, -): void { - if ( - watcher.stopped || - watcher.failed || - isTerminalStatus(watcher.lastStatus) - ) { - return; - } - - if (watcher.reconnectTimeoutId) { - clearTimeout(watcher.reconnectTimeoutId); - } - - const countAttempt = options.countAttempt ?? true; - if (countAttempt) { - watcher.reconnectAttempts += 1; - } else { - watcher.reconnectAttempts = 0; - } - - if (watcher.reconnectAttempts > MAX_SSE_RECONNECT_ATTEMPTS) { - const details = - error instanceof CloudTaskStreamError - ? error.details - : { - title: "Cloud stream disconnected", - message: - "Lost connection to the cloud run stream. Retry to reconnect.", - retryable: true, - }; - failWatcher(watcher, details); - return; - } - - const delay = Math.min( - SSE_RECONNECT_BASE_DELAY_MS * - 2 ** Math.max(watcher.reconnectAttempts - 1, 0), - SSE_RECONNECT_MAX_DELAY_MS, - ); - - watcher.reconnectTimeoutId = setTimeout(() => { - if (watcher.stopped) return; - watcher.reconnectTimeoutId = null; - void connectSse(watcher, { - startLatest: watcher.isBootstrapping || watcher.hasEmittedSnapshot, - }); - }, delay); -} - -async function handleStreamCompletion( - watcher: WatcherState, - options: { - reconnectIfNonTerminal: boolean; - reconnectError?: unknown; - countReconnectAttempt?: boolean; - }, -): Promise { - if (watcher.stopped) return; - - const { reconnectIfNonTerminal } = options; - const run = await fetchTaskRunState(watcher); - if (watcher.stopped || watcher.failed) return; - - if (watcher.isBootstrapping) { - if (!run) { - watcher.needsPostBootstrapReconnect = true; - return; - } - - applyTaskRunState(watcher, run); - if (isTerminalStatus(watcher.lastStatus) || !reconnectIfNonTerminal) { - watcher.needsStopAfterBootstrap = true; - } else { - watcher.needsPostBootstrapReconnect = true; - } - return; - } - - if (!run) { - scheduleReconnect( - watcher, - new CloudTaskStreamError("Failed to fetch terminal cloud run state", { - title: "Cloud run state unavailable", - message: - "Could not fetch the latest cloud run state after the stream ended. Retry to reconnect.", - retryable: true, - }), - ); - return; - } - - const stateChanged = applyTaskRunState(watcher, run); - - if (!isTerminalStatus(watcher.lastStatus) && reconnectIfNonTerminal) { - if (stateChanged) { - emitStatus(watcher); - } - log.warn("Cloud task stream ended before terminal status", { - runId: watcher.runId, - status: watcher.lastStatus, - }); - scheduleReconnect(watcher, options.reconnectError, { - countAttempt: options.countReconnectAttempt ?? false, - }); - return; - } - - emitStatus(watcher); - stopWatcher(watcher); -} - -function applyTaskRunState( - watcher: WatcherState, - run: - | Pick< - TaskRun, - | "status" - | "stage" - | "output" - | "error_message" - | "branch" - | "updated_at" - > - | TaskRunStateEvent, -): boolean { - const updatedAt = run.updated_at ?? null; - if ( - updatedAt && - watcher.lastStatusUpdatedAt && - Date.parse(updatedAt) <= Date.parse(watcher.lastStatusUpdatedAt) - ) { - return false; - } - - const nextStatus = run.status ?? watcher.lastStatus; - const nextStage = run.stage ?? null; - const nextOutput = run.output ?? null; - const nextErrorMessage = run.error_message ?? null; - const nextBranch = run.branch ?? null; - - const changed = - nextStatus !== watcher.lastStatus || - nextStage !== watcher.lastStage || - JSON.stringify(nextOutput) !== JSON.stringify(watcher.lastOutput) || - nextErrorMessage !== watcher.lastErrorMessage || - nextBranch !== watcher.lastBranch; - watcher.lastStatus = nextStatus ?? null; - watcher.lastStage = nextStage; - watcher.lastOutput = nextOutput; - watcher.lastErrorMessage = nextErrorMessage; - watcher.lastBranch = nextBranch; - if (updatedAt) { - watcher.lastStatusUpdatedAt = updatedAt; - } - - return changed; -} - -async function fetchTaskRunState( - watcher: WatcherState, -): Promise { - try { - return await getTaskRun(watcher.taskId, watcher.runId); - } catch (error) { - if (error instanceof HttpError) { - log.warn("Cloud task status fetch failed", { - runId: watcher.runId, - status: error.status, - }); - if (shouldFailWatcherForFetchStatus(error.status)) { - failWatcher(watcher, createStreamStatusError(error.status).details); - } - return null; - } - log.warn("Cloud task status fetch error", { - runId: watcher.runId, - error, - }); - return null; - } + return cloudTaskEngine; } -/** - * Loads the historical log entries for the run, mirroring the desktop's - * dual-source strategy: - * 1. Try the paginated `session_logs/` API — the live source while a run - * is active. For older / archived runs this can come back empty even - * though the canonical log exists on S3. - * 2. Fall back to the run's presigned `log_url` (S3 NDJSON), which is the - * canonical archive for completed runs. - * - * Returns `null` only when both sources fail outright (so the bootstrap can - * surface a retryable error). An empty paginated result is treated as "no - * data yet" and falls through to S3 — if S3 also has nothing we return the - * empty array so the snapshot can still flip the session to `"connected"`. - */ -async function fetchHistoricalEntries( - watcher: WatcherState, - run: TaskRun, -): Promise { - const paginated = await fetchAllSessionLogs(watcher); - if (watcher.stopped || watcher.failed) return null; - if (paginated && paginated.length > 0) return paginated; - - if (run.log_url) { - const s3Entries = await fetchS3LogEntries(watcher, run.log_url); - if (watcher.stopped || watcher.failed) return null; - if (s3Entries && s3Entries.length > 0) return s3Entries; - } - - // Both sources returned no rows. Prefer the paginated result (which is - // `[]` rather than `null`) so the caller can still emit an empty snapshot - // and the session flips to `"connected"` instead of hanging on loading. - return paginated ?? null; -} - -async function fetchS3LogEntries( - watcher: WatcherState, - logUrl: string, -): Promise { - try { - const response = await fetch(logUrl, { - signal: createTimeoutSignal(15_000), - }); - if (response.status === 404) { - // No archived log yet for this run — not an error, just no data. - return []; - } - if (!response.ok) { - log.warn("S3 session log fetch returned non-OK", { - runId: watcher.runId, - status: response.status, - }); - return null; +export function watchCloudTask({ + taskId, + runId, + onUpdate, +}: WatchCloudTaskOptions): WatchCloudTaskHandle { + const engine = getCloudTaskEngine(); + const listener = (update: CloudTaskUpdatePayload): void => { + if (update.taskId === taskId && update.runId === runId) { + onUpdate(update); } - const content = await response.text(); - if (!content.trim()) return []; - return parseSessionLogs(content).rawEntries; - } catch (error) { - log.warn("S3 session log fetch failed", { - runId: watcher.runId, - error, - }); - return null; - } -} + }; -async function fetchAllSessionLogs( - watcher: WatcherState, -): Promise { - const entries: StoredLogEntry[] = []; - let offset = 0; + engine.on(CloudTaskEvent.Update, listener); + engine.watch({ + taskId, + runId, + apiHost: getBaseUrl(), + teamId: getProjectId(), + }); - while (true) { - if (watcher.stopped || watcher.failed) return null; - try { - const page = await fetchSessionLogs(watcher.taskId, watcher.runId, { - limit: SESSION_LOG_PAGE_LIMIT, - offset, - }); + let stopped = false; - for (const entry of page.entries) { - entries.push(entry); - } - if (!page.hasMore || page.entries.length === 0) { - return entries; + return { + stop: () => { + if (stopped) { + return; } - offset += page.entries.length; - } catch (error) { - if (error instanceof HttpError) { - log.warn("Cloud task session logs fetch failed", { - runId: watcher.runId, - status: error.status, - offset, - }); - if (shouldFailWatcherForFetchStatus(error.status)) { - failWatcher(watcher, createStreamStatusError(error.status).details); - } - return null; + stopped = true; + engine.off(CloudTaskEvent.Update, listener); + engine.unwatch(taskId, runId); + }, + reconnectIfDisconnected: () => { + if (!stopped) { + engine.reconnectIfDisconnected(taskId, runId); } - log.warn("Cloud task session logs fetch error", { - runId: watcher.runId, - offset, - error, - }); - return null; - } - } + }, + }; } diff --git a/apps/mobile/src/features/tasks/lib/sseParser.ts b/apps/mobile/src/features/tasks/lib/sseParser.ts deleted file mode 100644 index 4c626fd657..0000000000 --- a/apps/mobile/src/features/tasks/lib/sseParser.ts +++ /dev/null @@ -1,89 +0,0 @@ -import { logger } from "@/lib/logger"; - -const log = logger.scope("sse-parser"); - -export interface SseEvent { - event?: string; - id?: string; - data: unknown; -} - -export class SseEventParser { - private buffer = ""; - private currentEventName: string | null = null; - private currentEventId: string | null = null; - private currentData: string[] = []; - - parse(chunk: string): SseEvent[] { - this.buffer += chunk; - const lines = this.buffer.split("\n"); - this.buffer = lines.pop() || ""; - - const events: SseEvent[] = []; - - for (const rawLine of lines) { - const line = rawLine.endsWith("\r") ? rawLine.slice(0, -1) : rawLine; - - if (line === "") { - const event = this.flushEvent(); - if (event) { - events.push(event); - } - continue; - } - - if (line.startsWith(":")) { - continue; - } - - if (line.startsWith("event:")) { - this.currentEventName = line.slice(6).trim() || null; - continue; - } - - if (line.startsWith("id:")) { - this.currentEventId = line.slice(3).trim() || null; - continue; - } - - if (line.startsWith("data:")) { - this.currentData.push(line.slice(5).trimStart()); - } - } - - return events; - } - - reset(): void { - this.buffer = ""; - this.currentEventName = null; - this.currentEventId = null; - this.currentData = []; - } - - private flushEvent(): SseEvent | null { - if (this.currentData.length === 0) { - this.currentEventName = null; - this.currentEventId = null; - return null; - } - - const rawData = this.currentData.join("\n"); - this.currentData = []; - - try { - const data = JSON.parse(rawData); - return { - event: this.currentEventName ?? undefined, - id: this.currentEventId ?? undefined, - data, - }; - } catch { - log.warn("SSE event JSON parse failure", { rawData }); - return null; - } finally { - this.currentEventName = null; - this.currentEventId = null; - } - } -} diff --git a/apps/mobile/src/features/tasks/lib/taskMessageQueue.ts b/apps/mobile/src/features/tasks/lib/taskMessageQueue.ts new file mode 100644 index 0000000000..7ef0234f22 --- /dev/null +++ b/apps/mobile/src/features/tasks/lib/taskMessageQueue.ts @@ -0,0 +1,8 @@ +import { CloudTaskQueue } from "@posthog/core/sessions/cloudTaskQueue"; +import type { PendingAttachment } from "../composer/attachments/types"; + +let queueId = 0; + +export const taskMessageQueue = new CloudTaskQueue({ + createId: () => `queue-${++queueId}`, +}); diff --git a/apps/mobile/src/features/tasks/services/taskSessionService.ts b/apps/mobile/src/features/tasks/services/taskSessionService.ts new file mode 100644 index 0000000000..d51cf22e7a --- /dev/null +++ b/apps/mobile/src/features/tasks/services/taskSessionService.ts @@ -0,0 +1,248 @@ +import { combineCloudTaskQueuedMessages } from "@posthog/core/sessions/cloudTaskQueue"; +import { + type CloudTaskSessionNotificationKind, + CloudTaskSessionService, + type CloudTaskSessionTask, +} from "@posthog/core/sessions/cloudTaskSessionService"; +import { serializeCloudPrompt, type Task } from "@posthog/shared"; +import * as Haptics from "expo-haptics"; +import { AppState } from "react-native"; +import { presentLocalNotification } from "@/features/notifications/lib/notifications"; +import { usePreferencesStore } from "@/features/preferences/stores/preferencesStore"; +import { logger } from "@/lib/logger"; +import { getPostHogApiClient } from "@/lib/posthogApiClient"; +import { + CloudCommandError, + cancelRun, + runTaskInCloud, + sendCloudCommand, +} from "../api"; +import { buildCloudPromptBlocks } from "../composer/attachments/buildCloudPrompt"; +import type { PendingAttachment } from "../composer/attachments/types"; +import { watchCloudTask } from "../lib/cloudTaskStream"; +import { taskMessageQueue } from "../lib/taskMessageQueue"; +import { useAttachmentEchoStore } from "../stores/attachmentEchoStore"; +import { + taskSessionStatePort, + useTaskSessionStore, +} from "../stores/taskSessionStore"; +import { useTaskStore } from "../stores/taskStore"; +import type { SessionNotificationAttachment } from "../types"; +import { playbackRateForTaskDuration } from "../utils/playbackRate"; +import { playCompletionSound } from "../utils/sounds"; + +const log = logger.scope("task-session-service"); +const NOTIFICATION_DEDUP_WINDOW_MS = 30_000; +const lastNotificationAt = new Map(); + +function toSessionTask(task: Task): CloudTaskSessionTask { + const run = task.latest_run; + const permissionMode = run?.state?.initial_permission_mode; + return { + id: task.id, + title: task.title, + latestRun: run + ? { + id: run.id, + branch: run.branch, + reasoningEffort: run.reasoning_effort, + initialPermissionMode: + typeof permissionMode === "string" ? permissionMode : undefined, + } + : undefined, + }; +} + +function completionPlaybackRate(promptStartedAt?: number): number { + if ( + !usePreferencesStore.getState().scaleSoundWithTaskLength || + promptStartedAt == null + ) { + return 1; + } + return playbackRateForTaskDuration(Date.now() - promptStartedAt); +} + +function presentTaskNotification(args: { + taskId: string; + taskRunId: string; + kind: CloudTaskSessionNotificationKind; +}): void { + if (!usePreferencesStore.getState().pushNotificationsEnabled) return; + const state = useTaskSessionStore.getState(); + const session = state.sessions[args.taskRunId]; + if (!session || state.focusedTaskId === args.taskId) return; + + const now = Date.now(); + const previous = lastNotificationAt.get(args.taskId); + if (previous && now - previous < NOTIFICATION_DEDUP_WINDOW_MS) return; + lastNotificationAt.set(args.taskId, now); + + const title = session.taskTitle ?? "PostHog Code"; + const body = + args.kind === "awaiting_user_input" + ? `"${title}" needs your input` + : args.kind === "task_failed" + ? `"${title}" failed` + : `"${title}" finished`; + + void presentLocalNotification({ + title: "PostHog Code", + body, + data: { taskId: args.taskId, taskRunId: args.taskRunId }, + }); +} + +function reinjectAttachmentEchoes( + taskRunId: string, + events: Parameters< + NonNullable< + ConstructorParameters< + typeof CloudTaskSessionService + >[0]["prompts"]["reinjectSnapshotAttachments"] + > + >[1], +): void { + const echoes = useAttachmentEchoStore.getState().getEchoes(taskRunId); + let echoIndex = 0; + for (const event of events) { + if (echoIndex >= echoes.length) return; + if (event.type !== "session_update") { + continue; + } + const update = event.notification.update; + if (update?.sessionUpdate !== "user_message_chunk") continue; + if (update.attachments?.length) { + echoIndex += 1; + continue; + } + const echo = echoes[echoIndex]; + echoIndex += 1; + if (echo.text === (update.content?.text ?? "")) { + update.attachments = echo.attachments; + } + } +} + +export const taskSessionService = + new CloudTaskSessionService({ + state: taskSessionStatePort, + api: { + getTask: async (taskId) => + toSessionTask(await getPostHogApiClient().getTask(taskId)), + runTask: async (taskId, options) => { + if (!options) return toSessionTask(await runTaskInCloud(taskId)); + return toSessionTask( + await runTaskInCloud(taskId, { + branch: options.branch, + runtimeAdapter: "claude", + reasoningEffort: options.reasoningEffort, + initialPermissionMode: options.initialPermissionMode, + rtkEnabled: options.rtkEnabled, + resumeFromRunId: options.resumeFromRunId, + pendingUserMessage: options.pendingUserMessage, + }), + ); + }, + sendCommand: async (taskId, runId, command, payload) => { + await sendCloudCommand(taskId, runId, command, payload); + }, + cancelRun: async (taskId, runId) => { + await cancelRun(taskId, runId); + }, + classifyCommandError: (error) => { + if (error instanceof CloudCommandError) { + if (error.isSandboxInactive()) return { kind: "sandbox_inactive" }; + if ([502, 503, 504].includes(error.status)) + return { kind: "transient" }; + } + return { kind: "other" }; + }, + }, + watchers: { + create: (args) => watchCloudTask({ ...args }), + }, + queue: { + get: (taskId) => taskMessageQueue.getQueue(taskId), + drain: (taskId) => taskMessageQueue.drain(taskId, { stopAtEdited: true }), + prepend: (taskId, messages) => taskMessageQueue.prepend(taskId, messages), + remove: (taskId, messageId) => taskMessageQueue.remove(taskId, messageId), + combine: combineCloudTaskQueuedMessages, + }, + prompts: { + prepare: async (prompt, attachments) => { + const eventAttachments: SessionNotificationAttachment[] = + attachments.map(({ kind, uri, fileName, mimeType }) => ({ + kind, + uri, + fileName, + mimeType, + })); + return { + wirePayload: + attachments.length > 0 + ? serializeCloudPrompt( + await buildCloudPromptBlocks(prompt, attachments), + ) + : prompt, + eventAttachments, + }; + }, + reinjectSnapshotAttachments: reinjectAttachmentEchoes, + recordAttachmentEcho: (taskRunId, prompt, attachments) => + useAttachmentEchoStore + .getState() + .recordEcho(taskRunId, prompt, attachments), + }, + time: { + now: Date.now, + defer: (callback) => setTimeout(callback, 0), + }, + logger: log, + effects: { + onCompletion: ({ promptStartedAt }) => { + if (!usePreferencesStore.getState().pingsEnabled) return; + void playCompletionSound( + undefined, + undefined, + completionPlaybackRate(promptStartedAt), + ); + void Haptics.notificationAsync( + Haptics.NotificationFeedbackType.Success, + ); + }, + onNotification: presentTaskNotification, + }, + preferences: { + getComposerConfig: (taskId) => + useTaskStore.getState().composerConfigByTaskId[taskId], + isRtkEnabled: () => usePreferencesStore.getState().rtkEnabledCloud, + }, + }); + +export function connectToTask(task: Task): Promise { + return taskSessionService.connect(toSessionTask(task)); +} + +export const taskSessionActions = { + connectToTask, + disconnectFromTask: (taskId: string) => taskSessionService.disconnect(taskId), + sendPrompt: taskSessionService.sendPrompt.bind(taskSessionService), + sendPermissionResponse: + taskSessionService.sendPermissionResponse.bind(taskSessionService), + cancelPrompt: taskSessionService.cancelPrompt.bind(taskSessionService), + stopRun: taskSessionService.stopRun.bind(taskSessionService), + sendInterrupting: + taskSessionService.sendInterrupting.bind(taskSessionService), + flushQueuedMessages: + taskSessionService.flushQueuedMessages.bind(taskSessionService), + flushQueuedMessagesIfIdle: + taskSessionService.flushQueuedMessagesIfIdle.bind(taskSessionService), + steerQueuedMessage: + taskSessionService.steerQueuedMessage.bind(taskSessionService), + setConfigOption: taskSessionService.setConfigOption.bind(taskSessionService), +}; + +AppState.addEventListener("change", (nextState) => { + if (nextState === "active") taskSessionService.reconnectWatchers(); +}); diff --git a/apps/mobile/src/features/tasks/stores/messageQueueStore.test.ts b/apps/mobile/src/features/tasks/stores/messageQueueStore.test.ts deleted file mode 100644 index 53c1d4f3dc..0000000000 --- a/apps/mobile/src/features/tasks/stores/messageQueueStore.test.ts +++ /dev/null @@ -1,332 +0,0 @@ -import { beforeEach, describe, expect, it } from "vitest"; -import type { PendingAttachment } from "../composer/attachments/types"; -import { - combineQueuedMessages, - type QueuedMessage, - useMessageQueueStore, -} from "./messageQueueStore"; - -function image(id: string): PendingAttachment { - return { - kind: "image", - id, - uri: `file://${id}.png`, - fileName: `${id}.png`, - mimeType: "image/png", - }; -} - -describe("messageQueueStore", () => { - beforeEach(() => { - useMessageQueueStore.setState( - { queuesByTaskId: {}, editingByTaskId: {} }, - false, - ); - }); - - it("enqueues messages in FIFO order", () => { - const { enqueue, getQueue } = useMessageQueueStore.getState(); - enqueue("t1", "first", []); - enqueue("t1", "second", []); - enqueue("t1", "third", []); - expect(getQueue("t1").map((m) => m.content)).toEqual([ - "first", - "second", - "third", - ]); - }); - - it("keeps separate queues per task", () => { - const { enqueue, getQueue } = useMessageQueueStore.getState(); - enqueue("t1", "a", []); - enqueue("t2", "b", []); - expect(getQueue("t1").map((m) => m.content)).toEqual(["a"]); - expect(getQueue("t2").map((m) => m.content)).toEqual(["b"]); - }); - - it("drains the queue in order and clears it", () => { - const { enqueue, drain, getQueue } = useMessageQueueStore.getState(); - enqueue("t1", "a", []); - enqueue("t1", "b", []); - expect(drain("t1").map((m) => m.content)).toEqual(["a", "b"]); - expect(getQueue("t1")).toEqual([]); - // A second drain on the emptied queue is a no-op. - expect(drain("t1")).toEqual([]); - }); - - it("prepends restored messages at the head (failed-flush rollback)", () => { - const { enqueue, drain, prepend, getQueue } = - useMessageQueueStore.getState(); - enqueue("t1", "a", []); - enqueue("t1", "b", []); - const drained = drain("t1"); - - // A new message arrives while the flush is in flight. - enqueue("t1", "c", []); - // Flush failed — the drained messages go back ahead of the newcomer. - prepend("t1", drained); - - expect(getQueue("t1").map((m) => m.content)).toEqual(["a", "b", "c"]); - }); - - it.each([ - { - name: "removes exactly the targeted message", - contents: ["a", "b", "c"], - removeIndex: 1, - expected: ["a", "c"], - }, - { - name: "clears the entry once the last message is removed", - contents: ["only"], - removeIndex: 0, - expected: [], - }, - ])("$name", ({ contents, removeIndex, expected }) => { - const { enqueue, remove, getQueue } = useMessageQueueStore.getState(); - for (const content of contents) enqueue("t1", content, []); - remove("t1", getQueue("t1")[removeIndex].id); - expect(getQueue("t1").map((m) => m.content)).toEqual(expected); - }); - - it("ignores removal of an unknown id", () => { - const { enqueue, remove, getQueue } = useMessageQueueStore.getState(); - enqueue("t1", "a", []); - remove("t1", "nope"); - expect(getQueue("t1").map((m) => m.content)).toEqual(["a"]); - }); -}); - -describe("move", () => { - beforeEach(() => { - useMessageQueueStore.setState( - { queuesByTaskId: {}, editingByTaskId: {} }, - false, - ); - }); - - function seed(contents: string[]) { - const { enqueue, getQueue } = useMessageQueueStore.getState(); - for (const content of contents) enqueue("t1", content, []); - return getQueue("t1"); - } - - it("moves a message down one slot", () => { - const [a] = seed(["a", "b", "c"]); - useMessageQueueStore.getState().move("t1", a.id, "down"); - expect( - useMessageQueueStore - .getState() - .getQueue("t1") - .map((m) => m.content), - ).toEqual(["b", "a", "c"]); - }); - - it("moves a message up one slot", () => { - const queue = seed(["a", "b", "c"]); - useMessageQueueStore.getState().move("t1", queue[2].id, "up"); - expect( - useMessageQueueStore - .getState() - .getQueue("t1") - .map((m) => m.content), - ).toEqual(["a", "c", "b"]); - }); - - it.each([ - { name: "up at the head", index: 0, direction: "up" as const }, - { name: "down at the tail", index: 2, direction: "down" as const }, - ])("is a no-op moving $name", ({ index, direction }) => { - const queue = seed(["a", "b", "c"]); - useMessageQueueStore.getState().move("t1", queue[index].id, direction); - expect( - useMessageQueueStore - .getState() - .getQueue("t1") - .map((m) => m.content), - ).toEqual(["a", "b", "c"]); - }); - - it("ignores an unknown id", () => { - seed(["a", "b"]); - useMessageQueueStore.getState().move("t1", "nope", "up"); - expect( - useMessageQueueStore - .getState() - .getQueue("t1") - .map((m) => m.content), - ).toEqual(["a", "b"]); - }); -}); - -describe("edit in place", () => { - beforeEach(() => { - useMessageQueueStore.setState( - { queuesByTaskId: {}, editingByTaskId: {} }, - false, - ); - }); - - it("updates content and attachments in place, keeping id and position", () => { - const { enqueue, update, getQueue } = useMessageQueueStore.getState(); - enqueue("t1", "a", []); - enqueue("t1", "b", []); - const target = getQueue("t1")[0]; - - update("t1", target.id, { content: "edited", attachments: [image("x")] }); - - const queue = useMessageQueueStore.getState().getQueue("t1"); - expect(queue.map((m) => m.id)).toEqual([target.id, getQueue("t1")[1].id]); - expect(queue[0].content).toBe("edited"); - expect(queue[0].attachments.map((att) => att.id)).toEqual(["x"]); - expect(queue[1].content).toBe("b"); - }); - - it("is a no-op when the target id is gone", () => { - const { enqueue, update, getQueue } = useMessageQueueStore.getState(); - enqueue("t1", "a", []); - update("t1", "missing", { content: "edited", attachments: [] }); - expect(getQueue("t1")[0].content).toBe("a"); - }); - - it("set/clear stores and releases the edit hold", () => { - const { enqueue, getQueue, setEditing, clearEditing } = - useMessageQueueStore.getState(); - enqueue("t1", "a", []); - const id = getQueue("t1")[0].id; - - setEditing("t1", id); - expect(useMessageQueueStore.getState().editingByTaskId.t1).toBe(id); - - clearEditing("t1"); - expect(useMessageQueueStore.getState().editingByTaskId.t1).toBeUndefined(); - }); - - it("clears the edit hold when the edited message is removed", () => { - const { enqueue, getQueue, setEditing, remove } = - useMessageQueueStore.getState(); - enqueue("t1", "a", []); - enqueue("t1", "b", []); - const id = getQueue("t1")[0].id; - setEditing("t1", id); - - remove("t1", id); - - expect(useMessageQueueStore.getState().editingByTaskId.t1).toBeUndefined(); - expect( - useMessageQueueStore - .getState() - .getQueue("t1") - .map((m) => m.content), - ).toEqual(["b"]); - }); - - it("keeps the edit hold when a different message is removed", () => { - const { enqueue, getQueue, setEditing, remove } = - useMessageQueueStore.getState(); - enqueue("t1", "a", []); - enqueue("t1", "b", []); - const editingId = getQueue("t1")[1].id; - setEditing("t1", editingId); - - remove("t1", getQueue("t1")[0].id); - - expect(useMessageQueueStore.getState().editingByTaskId.t1).toBe(editingId); - }); -}); - -describe("drain boundary (stopAtEdited)", () => { - beforeEach(() => { - useMessageQueueStore.setState( - { queuesByTaskId: {}, editingByTaskId: {} }, - false, - ); - }); - - function seed(contents: string[]) { - const { enqueue, getQueue } = useMessageQueueStore.getState(); - for (const content of contents) enqueue("t1", content, []); - return getQueue("t1"); - } - - it("drains only the messages before the edited one", () => { - const queue = seed(["a", "b", "c"]); - useMessageQueueStore.getState().setEditing("t1", queue[1].id); - - const drained = useMessageQueueStore - .getState() - .drain("t1", { stopAtEdited: true }); - - expect(drained.map((m) => m.content)).toEqual(["a"]); - // The edited message and everything after it stay queued. - expect( - useMessageQueueStore - .getState() - .getQueue("t1") - .map((m) => m.content), - ).toEqual(["b", "c"]); - }); - - it("drains nothing when the head message is being edited", () => { - const queue = seed(["a", "b"]); - useMessageQueueStore.getState().setEditing("t1", queue[0].id); - - const drained = useMessageQueueStore - .getState() - .drain("t1", { stopAtEdited: true }); - - expect(drained).toEqual([]); - expect( - useMessageQueueStore - .getState() - .getQueue("t1") - .map((m) => m.content), - ).toEqual(["a", "b"]); - }); - - it("drains the whole queue when nothing is being edited", () => { - seed(["a", "b", "c"]); - const drained = useMessageQueueStore - .getState() - .drain("t1", { stopAtEdited: true }); - expect(drained.map((m) => m.content)).toEqual(["a", "b", "c"]); - expect(useMessageQueueStore.getState().getQueue("t1")).toEqual([]); - }); - - it("ignores the boundary without stopAtEdited, even mid-edit", () => { - const queue = seed(["a", "b", "c"]); - useMessageQueueStore.getState().setEditing("t1", queue[1].id); - - const drained = useMessageQueueStore.getState().drain("t1"); - - expect(drained.map((m) => m.content)).toEqual(["a", "b", "c"]); - expect(useMessageQueueStore.getState().getQueue("t1")).toEqual([]); - }); -}); - -describe("combineQueuedMessages", () => { - function msg( - content: string, - attachments: PendingAttachment[], - ): QueuedMessage { - return { id: content, content, attachments }; - } - - it("joins text in order with a blank line and concatenates attachments", () => { - const result = combineQueuedMessages([ - msg("first", [image("one")]), - msg("second", []), - msg("third", [image("two"), image("three")]), - ]); - expect(result.text).toBe("first\n\nsecond\n\nthird"); - expect(result.attachments.map((a) => a.id)).toEqual([ - "one", - "two", - "three", - ]); - }); - - it("handles an empty list", () => { - expect(combineQueuedMessages([])).toEqual({ text: "", attachments: [] }); - }); -}); diff --git a/apps/mobile/src/features/tasks/stores/messageQueueStore.ts b/apps/mobile/src/features/tasks/stores/messageQueueStore.ts deleted file mode 100644 index f676cb326f..0000000000 --- a/apps/mobile/src/features/tasks/stores/messageQueueStore.ts +++ /dev/null @@ -1,193 +0,0 @@ -import { create } from "zustand"; -import type { PendingAttachment } from "../composer/attachments/types"; - -export interface QueuedMessage { - id: string; - content: string; - attachments: PendingAttachment[]; -} - -const EMPTY: QueuedMessage[] = []; - -export type MoveDirection = "up" | "down"; - -let queueIdCounter = 0; -function nextQueueId(): string { - queueIdCounter += 1; - return `queue-${queueIdCounter}`; -} - -// A message being edited in place acts as a drain boundary: only the messages -// queued before it may auto-send; it and everything after stay put until the -// edit is saved or cancelled. Returns the full length when nothing is being -// edited (or the edited message has already left the queue). Mirrors -// `sendableQueuePrefixLength` in @posthog/shared, reimplemented here because -// mobile's queue rows (`{id, content, attachments}`) don't match the shared -// `AgentSession` message shape. -function sendablePrefixLength( - queue: QueuedMessage[], - editingId: string | undefined, -): number { - if (!editingId) return queue.length; - const index = queue.findIndex((m) => m.id === editingId); - return index === -1 ? queue.length : index; -} - -interface MessageQueueState { - queuesByTaskId: Record; - /** Per-task id of the queued message currently open in the composer. */ - editingByTaskId: Record; - enqueue: ( - taskId: string, - content: string, - attachments: PendingAttachment[], - ) => void; - /** - * Remove and return queued messages from the head, in FIFO order. With - * `stopAtEdited`, stops at the in-place edit boundary so the edited message - * and everything after it stay queued. - */ - drain: ( - taskId: string, - options?: { stopAtEdited?: boolean }, - ) => QueuedMessage[]; - /** Restore messages at the head of the queue, e.g. after a failed flush. */ - prepend: (taskId: string, messages: QueuedMessage[]) => void; - /** Drop a single queued message by id. Clears the edit hold if it targeted it. */ - remove: (taskId: string, messageId: string) => void; - /** Reorder a queued message one slot up or down; the send order follows. */ - move: (taskId: string, messageId: string, direction: MoveDirection) => void; - /** Replace a queued message's content/attachments in place, keeping position. */ - update: ( - taskId: string, - messageId: string, - patch: { content: string; attachments: PendingAttachment[] }, - ) => void; - /** Mark a queued message as being edited in the composer (drain boundary). */ - setEditing: (taskId: string, messageId: string) => void; - /** Release the in-place edit hold. */ - clearEditing: (taskId: string) => void; - getQueue: (taskId: string) => QueuedMessage[]; -} - -export const useMessageQueueStore = create((set, get) => ({ - queuesByTaskId: {}, - editingByTaskId: {}, - enqueue: (taskId, content, attachments) => - set((state) => ({ - queuesByTaskId: { - ...state.queuesByTaskId, - [taskId]: [ - ...(state.queuesByTaskId[taskId] ?? []), - { id: nextQueueId(), content, attachments }, - ], - }, - })), - drain: (taskId, options) => { - const queued = get().queuesByTaskId[taskId] ?? EMPTY; - if (queued.length === 0) return EMPTY; - const cutoff = options?.stopAtEdited - ? sendablePrefixLength(queued, get().editingByTaskId[taskId]) - : queued.length; - if (cutoff === 0) return EMPTY; - const drained = queued.slice(0, cutoff); - const rest = queued.slice(cutoff); - set((state) => { - if (rest.length === 0) { - const { [taskId]: _drained, ...others } = state.queuesByTaskId; - return { queuesByTaskId: others }; - } - return { - queuesByTaskId: { ...state.queuesByTaskId, [taskId]: rest }, - }; - }); - return drained; - }, - prepend: (taskId, messages) => - set((state) => ({ - queuesByTaskId: { - ...state.queuesByTaskId, - [taskId]: [...messages, ...(state.queuesByTaskId[taskId] ?? [])], - }, - })), - remove: (taskId, messageId) => - set((state) => { - const queue = state.queuesByTaskId[taskId]; - if (!queue) return state; - const next = queue.filter((m) => m.id !== messageId); - if (next.length === queue.length) return state; - const editingByTaskId = - state.editingByTaskId[taskId] === messageId - ? omit(state.editingByTaskId, taskId) - : state.editingByTaskId; - if (next.length === 0) { - const { [taskId]: _emptied, ...rest } = state.queuesByTaskId; - return { queuesByTaskId: rest, editingByTaskId }; - } - return { - queuesByTaskId: { ...state.queuesByTaskId, [taskId]: next }, - editingByTaskId, - }; - }), - move: (taskId, messageId, direction) => - set((state) => { - const queue = state.queuesByTaskId[taskId]; - if (!queue) return state; - const from = queue.findIndex((m) => m.id === messageId); - if (from === -1) return state; - const to = direction === "up" ? from - 1 : from + 1; - if (to < 0 || to >= queue.length) return state; - const next = [...queue]; - const [moved] = next.splice(from, 1); - next.splice(to, 0, moved); - return { - queuesByTaskId: { ...state.queuesByTaskId, [taskId]: next }, - }; - }), - update: (taskId, messageId, patch) => - set((state) => { - const queue = state.queuesByTaskId[taskId]; - if (!queue?.some((m) => m.id === messageId)) return state; - const next = queue.map((m) => - m.id === messageId - ? { ...m, content: patch.content, attachments: patch.attachments } - : m, - ); - return { - queuesByTaskId: { ...state.queuesByTaskId, [taskId]: next }, - }; - }), - setEditing: (taskId, messageId) => - set((state) => ({ - editingByTaskId: { ...state.editingByTaskId, [taskId]: messageId }, - })), - clearEditing: (taskId) => - set((state) => - taskId in state.editingByTaskId - ? { editingByTaskId: omit(state.editingByTaskId, taskId) } - : state, - ), - getQueue: (taskId) => get().queuesByTaskId[taskId] ?? EMPTY, -})); - -function omit>( - record: T, - key: string, -): Record { - const { [key]: _removed, ...rest } = record; - return rest; -} - -/** - * Combine buffered messages into a single prompt, preserving the order they - * were typed: texts join with a blank line, attachments concatenate. - */ -export function combineQueuedMessages(messages: QueuedMessage[]): { - text: string; - attachments: PendingAttachment[]; -} { - return { - text: messages.map((m) => m.content).join("\n\n"), - attachments: messages.flatMap((m) => m.attachments), - }; -} diff --git a/apps/mobile/src/features/tasks/stores/taskSessionStore.test.ts b/apps/mobile/src/features/tasks/stores/taskSessionStore.test.ts index 6b4fa9152a..839de96009 100644 --- a/apps/mobile/src/features/tasks/stores/taskSessionStore.test.ts +++ b/apps/mobile/src/features/tasks/stores/taskSessionStore.test.ts @@ -1,350 +1,52 @@ -import type { - CloudTaskUpdatePayload, - StoredLogEntry, - Task, - TaskRun, -} from "@posthog/shared"; -import { beforeEach, describe, expect, it, vi } from "vitest"; - -const { mockGetTask } = vi.hoisted(() => ({ mockGetTask: vi.fn() })); - -vi.mock("expo-haptics", () => ({ - impactAsync: vi.fn(), - notificationAsync: vi.fn(), - ImpactFeedbackStyle: { Light: "light", Medium: "medium" }, - NotificationFeedbackType: { Success: "success" }, -})); -vi.mock("../lib/cloudTaskStream", () => ({ watchCloudTask: vi.fn() })); -vi.mock("../composer/attachments/buildCloudPrompt", () => ({ - buildCloudPromptBlocks: vi.fn(() => Promise.resolve([])), -})); -vi.mock("../utils/sounds", () => ({ - playCompletionSound: vi.fn(() => Promise.resolve()), -})); -vi.mock("@/features/notifications/lib/notifications", () => ({ - presentLocalNotification: vi.fn(() => Promise.resolve()), -})); -vi.mock("../api", () => ({ - CloudCommandError: class CloudCommandError extends Error {}, - runTaskInCloud: vi.fn(), - sendCloudCommand: vi.fn(), -})); - -vi.mock("@/lib/posthogApiClient", () => ({ - getPostHogApiClient: () => ({ getTask: mockGetTask }), -})); - -import { usePreferencesStore } from "@/features/preferences/stores/preferencesStore"; -import { runTaskInCloud } from "../api"; -import { useMessageQueueStore } from "./messageQueueStore"; -import { type TaskSession, useTaskSessionStore } from "./taskSessionStore"; -import { useTaskStore } from "./taskStore"; - -function seedSession(overrides: Partial = {}): void { - const session: TaskSession = { +import type { CloudTaskSession } from "@posthog/core/sessions/cloudTaskSessionService"; +import { beforeEach, describe, expect, it } from "vitest"; +import { + getTaskSession, + taskSessionStatePort, + useTaskSessionStore, +} from "./taskSessionStore"; + +function session(overrides: Partial = {}): CloudTaskSession { + return { + taskId: "task-1", taskRunId: "run-1", - taskId: "t1", events: [], status: "connected", - isPromptPending: true, + isPromptPending: false, ...overrides, }; - useTaskSessionStore.setState({ sessions: { "run-1": session } }); } -describe("steerQueuedMessage", () => { - beforeEach(() => { - useMessageQueueStore.setState({ queuesByTaskId: {} }, false); - useTaskSessionStore.setState({ sessions: {} }); - }); - - it("removes the message and resends it as a steer", async () => { - seedSession(); - const sendInterrupting = vi.fn(() => Promise.resolve()); - useTaskSessionStore.setState({ sendInterrupting }); - - useMessageQueueStore.getState().enqueue("t1", "first", []); - useMessageQueueStore.getState().enqueue("t1", "second", []); - const target = useMessageQueueStore.getState().getQueue("t1")[0]; - - await useTaskSessionStore.getState().steerQueuedMessage("t1", target.id); - - expect(sendInterrupting).toHaveBeenCalledWith("t1", "first", []); - expect( - useMessageQueueStore - .getState() - .getQueue("t1") - .map((m) => m.content), - ).toEqual(["second"]); - }); - - it("rolls the message back onto the head when the resend fails", async () => { - seedSession(); - const sendInterrupting = vi.fn(() => Promise.reject(new Error("boom"))); - useTaskSessionStore.setState({ sendInterrupting }); - - useMessageQueueStore.getState().enqueue("t1", "first", []); - useMessageQueueStore.getState().enqueue("t1", "second", []); - const target = useMessageQueueStore.getState().getQueue("t1")[0]; - - await expect( - useTaskSessionStore.getState().steerQueuedMessage("t1", target.id), - ).rejects.toThrow("boom"); - - expect( - useMessageQueueStore - .getState() - .getQueue("t1") - .map((m) => m.content), - ).toEqual(["first", "second"]); - }); - - it("no-ops while the session is compacting", async () => { - seedSession({ isCompacting: true }); - const sendInterrupting = vi.fn(() => Promise.resolve()); - useTaskSessionStore.setState({ sendInterrupting }); - - useMessageQueueStore.getState().enqueue("t1", "first", []); - const target = useMessageQueueStore.getState().getQueue("t1")[0]; - - await useTaskSessionStore.getState().steerQueuedMessage("t1", target.id); - - expect(sendInterrupting).not.toHaveBeenCalled(); - expect( - useMessageQueueStore - .getState() - .getQueue("t1") - .map((m) => m.content), - ).toEqual(["first"]); - }); - - it("no-ops for an unknown message id", async () => { - seedSession(); - const sendInterrupting = vi.fn(() => Promise.resolve()); - useTaskSessionStore.setState({ sendInterrupting }); - - useMessageQueueStore.getState().enqueue("t1", "first", []); - - await useTaskSessionStore.getState().steerQueuedMessage("t1", "missing"); - - expect(sendInterrupting).not.toHaveBeenCalled(); - expect(useMessageQueueStore.getState().getQueue("t1")).toHaveLength(1); - }); - - it("no-ops when no turn is running", async () => { - seedSession({ isPromptPending: false }); - const sendInterrupting = vi.fn(() => Promise.resolve()); - useTaskSessionStore.setState({ sendInterrupting }); - - useMessageQueueStore.getState().enqueue("t1", "first", []); - const target = useMessageQueueStore.getState().getQueue("t1")[0]; - - await useTaskSessionStore.getState().steerQueuedMessage("t1", target.id); - - expect(sendInterrupting).not.toHaveBeenCalled(); - expect(useMessageQueueStore.getState().getQueue("t1")).toHaveLength(1); - }); -}); - -describe("flushQueuedMessagesIfIdle", () => { - beforeEach(() => { - useMessageQueueStore.setState( - { queuesByTaskId: {}, editingByTaskId: {} }, - false, - ); - useTaskSessionStore.setState({ sessions: {} }); - }); - - it("sends the queue when the agent is idle", async () => { - seedSession({ isPromptPending: false }); - const sendInterrupting = vi.fn(() => Promise.resolve()); - useTaskSessionStore.setState({ sendInterrupting }); - useMessageQueueStore.getState().enqueue("t1", "a", []); - useMessageQueueStore.getState().enqueue("t1", "b", []); - - useTaskSessionStore.getState().flushQueuedMessagesIfIdle("t1"); - await vi.waitFor(() => expect(sendInterrupting).toHaveBeenCalled()); - - expect(sendInterrupting).toHaveBeenCalledWith("t1", "a\n\nb", []); - expect(useMessageQueueStore.getState().getQueue("t1")).toEqual([]); - }); - - it("sends only the messages before the one being edited", async () => { - seedSession({ isPromptPending: false }); - const sendInterrupting = vi.fn(() => Promise.resolve()); - useTaskSessionStore.setState({ sendInterrupting }); - useMessageQueueStore.getState().enqueue("t1", "a", []); - useMessageQueueStore.getState().enqueue("t1", "b", []); - useMessageQueueStore.getState().enqueue("t1", "c", []); - const edited = useMessageQueueStore.getState().getQueue("t1")[1]; - useMessageQueueStore.getState().setEditing("t1", edited.id); - - useTaskSessionStore.getState().flushQueuedMessagesIfIdle("t1"); - await vi.waitFor(() => expect(sendInterrupting).toHaveBeenCalled()); - - expect(sendInterrupting).toHaveBeenCalledWith("t1", "a", []); - expect( - useMessageQueueStore - .getState() - .getQueue("t1") - .map((m) => m.content), - ).toEqual(["b", "c"]); - }); - - it.each([ - { name: "a turn is running", overrides: { isPromptPending: true } }, - { name: "the run is terminal", overrides: { terminalStatus: "completed" } }, - { name: "the agent is compacting", overrides: { isCompacting: true } }, - ] as const)("no-ops when $name", async ({ overrides }) => { - seedSession({ isPromptPending: false, ...overrides }); - const sendInterrupting = vi.fn(() => Promise.resolve()); - useTaskSessionStore.setState({ sendInterrupting }); - useMessageQueueStore.getState().enqueue("t1", "a", []); - - useTaskSessionStore.getState().flushQueuedMessagesIfIdle("t1"); - await Promise.resolve(); - - expect(sendInterrupting).not.toHaveBeenCalled(); - expect(useMessageQueueStore.getState().getQueue("t1")).toHaveLength(1); - }); - - it("no-ops when the queue is empty", async () => { - seedSession({ isPromptPending: false }); - const sendInterrupting = vi.fn(() => Promise.resolve()); - useTaskSessionStore.setState({ sendInterrupting }); - - useTaskSessionStore.getState().flushQueuedMessagesIfIdle("t1"); - await Promise.resolve(); - - expect(sendInterrupting).not.toHaveBeenCalled(); - }); -}); - -describe("_resumeCloudRun", () => { - const mockRunTaskInCloud = vi.mocked(runTaskInCloud); - - function previousTask(latestRun: Partial): Task { - return { - id: "t1", - latest_run: { id: "prev-run", ...latestRun } as TaskRun, - } as Task; - } - +describe("taskSessionStatePort", () => { beforeEach(() => { - vi.clearAllMocks(); - useTaskStore.setState({ composerConfigByTaskId: {} }); - useTaskSessionStore.setState({ sessions: {} }); - usePreferencesStore.setState({ rtkEnabledCloud: true }); - mockRunTaskInCloud.mockResolvedValue({ - id: "t1", - latest_run: { id: "new-run" }, - } as Task); + useTaskSessionStore.setState({ sessions: {}, focusedTaskId: null }); }); - it("forwards the previous run's effort and permission mode", async () => { - mockGetTask.mockResolvedValue( - previousTask({ - branch: "feature", - reasoning_effort: "low", - state: { initial_permission_mode: "acceptEdits" }, - }), - ); - - await useTaskSessionStore - .getState() - ._resumeCloudRun("t1", "prev-run", "hi"); - - expect(mockRunTaskInCloud).toHaveBeenCalledWith("t1", { - branch: "feature", - resumeFromRunId: "prev-run", - pendingUserMessage: "hi", - reasoningEffort: "low", - initialPermissionMode: "acceptEdits", - rtkEnabled: true, - }); - }); - - it("forwards the rtk compression opt-out so resume preserves it", async () => { - usePreferencesStore.setState({ rtkEnabledCloud: false }); - mockGetTask.mockResolvedValue(previousTask({ branch: "feature" })); - - await useTaskSessionStore - .getState() - ._resumeCloudRun("t1", "prev-run", "hi"); + it("indexes sessions by both task and run", () => { + taskSessionStatePort.set(session()); - expect(mockRunTaskInCloud).toHaveBeenCalledWith( - "t1", - expect.objectContaining({ rtkEnabled: false }), - ); + expect(getTaskSession("task-1")?.taskRunId).toBe("run-1"); }); - it("prefers the composer's current selection over the previous run", async () => { - useTaskStore.setState({ - composerConfigByTaskId: { t1: { mode: "plan", reasoning: "max" } }, - }); - mockGetTask.mockResolvedValue( - previousTask({ - branch: null, - reasoning_effort: "low", - state: { initial_permission_mode: "acceptEdits" }, - }), - ); + it("updates one session without replacing the others", () => { + taskSessionStatePort.set(session()); + taskSessionStatePort.set(session({ taskId: "task-2", taskRunId: "run-2" })); - await useTaskSessionStore - .getState() - ._resumeCloudRun("t1", "prev-run", "hi"); + taskSessionStatePort.update("run-1", (current) => ({ + ...current, + isPromptPending: true, + })); - expect(mockRunTaskInCloud).toHaveBeenCalledWith( - "t1", - expect.objectContaining({ - reasoningEffort: "max", - initialPermissionMode: "plan", - }), + expect(useTaskSessionStore.getState().sessions["run-2"]?.taskId).toBe( + "task-2", ); }); -}); -describe("compaction tracking from the log stream", () => { - beforeEach(() => { - useTaskSessionStore.setState({ sessions: {} }); - }); + it("removes sessions by run id", () => { + taskSessionStatePort.set(session()); - function statusEntry(isComplete: boolean): StoredLogEntry { - return { - type: "notification", - notification: { - method: "_posthog/status", - params: { status: "compacting", isComplete }, - }, - }; - } + taskSessionStatePort.remove("run-1"); - function logsUpdate(entries: StoredLogEntry[]): CloudTaskUpdatePayload { - return { - kind: "logs", - taskId: "t1", - runId: "run-1", - newEntries: entries, - totalEntryCount: entries.length, - }; - } - - it("sets isCompacting on a compacting status and clears it on the boundary", () => { - seedSession({ isCompacting: false }); - const store = useTaskSessionStore.getState(); - - store._handleCloudUpdate("run-1", logsUpdate([statusEntry(false)])); - expect(store.getSessionForTask("t1")?.isCompacting).toBe(true); - - store._handleCloudUpdate( - "run-1", - logsUpdate([ - { - type: "notification", - notification: { method: "_posthog/compact_boundary" }, - }, - ]), - ); - expect(store.getSessionForTask("t1")?.isCompacting).toBe(false); + expect(getTaskSession("task-1")).toBeUndefined(); }); }); diff --git a/apps/mobile/src/features/tasks/stores/taskSessionStore.ts b/apps/mobile/src/features/tasks/stores/taskSessionStore.ts index 3a9adcbdad..4bc02d26e9 100644 --- a/apps/mobile/src/features/tasks/stores/taskSessionStore.ts +++ b/apps/mobile/src/features/tasks/stores/taskSessionStore.ts @@ -1,1285 +1,47 @@ -import { - type CloudTaskUpdatePayload, - isTerminalStatus, - type StoredLogEntry, - type Task, -} from "@posthog/shared"; -import * as Haptics from "expo-haptics"; -import { AppState } from "react-native"; -import { create } from "zustand"; -import { presentLocalNotification } from "@/features/notifications/lib/notifications"; -import { usePreferencesStore } from "@/features/preferences/stores/preferencesStore"; -import { logger } from "@/lib/logger"; -import { getPostHogApiClient } from "@/lib/posthogApiClient"; -import { - CloudCommandError, - cancelRun, - runTaskInCloud, - sendCloudCommand, -} from "../api"; -import { buildCloudPromptBlocks } from "../composer/attachments/buildCloudPrompt"; -import { serializeCloudPrompt } from "../composer/attachments/cloudPrompt"; -import type { PendingAttachment } from "../composer/attachments/types"; -import { - type WatchCloudTaskHandle, - watchCloudTask, -} from "../lib/cloudTaskStream"; import type { - CloudPendingPermissionRequest, - SessionEvent, - SessionNotification, - SessionNotificationAttachment, -} from "../types"; -import { convertStoredEntriesToEvents } from "../utils/parseSessionLogs"; -import { playbackRateForTaskDuration } from "../utils/playbackRate"; -import { playCompletionSound } from "../utils/sounds"; -import { useAttachmentEchoStore } from "./attachmentEchoStore"; -import { - combineQueuedMessages, - useMessageQueueStore, -} from "./messageQueueStore"; -import { useTaskStore } from "./taskStore"; - -const log = logger.scope("task-session-store"); - -function completionPlaybackRate(promptStartedAt?: number): number { - if ( - !usePreferencesStore.getState().scaleSoundWithTaskLength || - promptStartedAt == null - ) { - return 1; - } - return playbackRateForTaskDuration(Date.now() - promptStartedAt); -} - -// Match historical `user_message_chunk` events (text-only, as the cloud -// stores them) against locally-cached attachment echoes by position+text. -// Echoes are written in send-order; we walk user messages in receive-order -// and zip them up. Drift (text mismatch at the same index) is treated as a -// no-op rather than a misattribution. -function reinjectAttachmentEchoes( - taskRunId: string, - events: SessionEvent[], -): void { - const echoes = useAttachmentEchoStore.getState().getEchoes(taskRunId); - if (echoes.length === 0) return; - - let echoIdx = 0; - for (const event of events) { - if (echoIdx >= echoes.length) return; - if (event.type !== "session_update") continue; - const update = event.notification?.update; - if (update?.sessionUpdate !== "user_message_chunk") continue; - if (update.attachments && update.attachments.length > 0) { - echoIdx++; - continue; - } - const echo = echoes[echoIdx]; - echoIdx++; - if (echo.text === (update.content?.text ?? "")) { - update.attachments = echo.attachments; - } - } -} - -type LocalNotificationKind = - | "turn_complete" - | "awaiting_user_input" - | "task_failed"; - -// Per-task cooldown so a noisy stream of terminal/awaiting events doesn't -// fire a burst of identical banners. Keyed by taskId, value is the epoch ms -// of the most recent notification for that task. -const NOTIFICATION_DEDUP_WINDOW_MS = 30_000; -const lastNotificationAt = new Map(); - -// TODO: server-side device presence. Today we can only suppress notifications -// when *this* device is foregrounded on the task (`focusedTaskId` check -// below). That leaves the cross-device case uncovered — e.g. desktop is open -// on the same task, server fans the push to every registered token, mobile -// still rings. Once the `/api/projects/{team_id}/tasks/{task_id}/presence/` -// beacon endpoint lands in posthog/posthog, add: -// 1. A stable per-install device_id (probably derive from pushTokenStore). -// 2. POST presence every ~30s while a task screen is mounted AND AppState -// is "active". -// 3. DELETE presence on screen blur or AppState → "background"/"inactive". -// The server will then drop pushes to devices with non-expired presence for -// the target task, so this client-side maybePresentLocalNotification stays -// as-is (it's only the OS fanout path we're improving). - -function maybePresentLocalNotification(args: { - taskRunId: string; - kind: LocalNotificationKind; -}): void { - if (!usePreferencesStore.getState().pushNotificationsEnabled) return; - - const storeState = useTaskSessionStore.getState(); - const session = storeState.sessions[args.taskRunId]; - if (!session) return; - - // Skip when the user is actively viewing this task — the UI already - // surfaces what changed; an OS banner would be redundant noise. - if (storeState.focusedTaskId === session.taskId) return; - - // Dedup: skip if we just notified about this task. - const now = Date.now(); - const previous = lastNotificationAt.get(session.taskId); - if (previous && now - previous < NOTIFICATION_DEDUP_WINDOW_MS) return; - lastNotificationAt.set(session.taskId, now); - - const title = session.taskTitle ?? "PostHog Code"; - let body: string; - switch (args.kind) { - case "awaiting_user_input": - body = `"${title}" needs your input`; - break; - case "task_failed": - body = `"${title}" failed`; - break; - default: - body = `"${title}" finished`; - break; - } - - presentLocalNotification({ - title: "PostHog Code", - body, - data: { taskId: session.taskId, taskRunId: session.taskRunId }, - }).catch(() => {}); -} - -// Session-update kinds that count as "the agent produced visible output" — -// once we've seen one of these the connecting/thinking indicator should clear. -const VISIBLE_AGENT_SESSION_UPDATES = new Set([ - "agent_message_chunk", - "agent_message", - "agent_thought_chunk", - "tool_call", - "tool_call_update", -]); - -// Notification methods that mark the end of an agent turn — clearing -// isPromptPending so the composer unblocks. -const TURN_END_METHODS = new Set([ - "_posthog/turn_complete", - "_posthog/task_complete", - "_posthog/error", - "_posthog/awaiting_user_input", -]); - -interface BatchAnalysis { - hasTurnEnd: boolean; - hasAwaitingUserInput: boolean; - hasTurnCompleted: boolean; - hasTurnFailed: boolean; - hasVisibleAgentOutput: boolean; - externalUserMessageCount: number; - agentMessageFinalized: boolean; - // Latest compaction state seen in the batch (undefined = no change). - compacting?: boolean; -} - -function analyzeEntries( - entries: StoredLogEntry[], - localUserEchoes: Set, -): BatchAnalysis { - let hasTurnEnd = false; - let hasAwaitingUserInput = false; - let hasTurnCompleted = false; - let hasTurnFailed = false; - let hasVisibleAgentOutput = false; - let externalUserMessageCount = 0; - let agentMessageFinalized = false; - let compacting: boolean | undefined; - - for (const entry of entries) { - const method = entry.notification?.method; - if (method && TURN_END_METHODS.has(method)) { - hasTurnEnd = true; - if (method === "_posthog/awaiting_user_input") { - hasAwaitingUserInput = true; - } - if ( - method === "_posthog/turn_complete" || - method === "_posthog/task_complete" - ) { - hasTurnCompleted = true; - } - if (method === "_posthog/error") { - hasTurnFailed = true; - } - } - - if (method === "_posthog/status") { - const params = entry.notification?.params as - | { status?: string; isComplete?: boolean } - | undefined; - if (params?.status === "compacting") { - compacting = !params.isComplete; - } - } - if (method === "_posthog/compact_boundary") { - compacting = false; - } - - if ( - entry.type === "notification" && - method === "session/update" && - entry.notification?.params - ) { - const params = entry.notification.params as SessionNotification; - const sessionUpdate = params.update?.sessionUpdate; - if (sessionUpdate && VISIBLE_AGENT_SESSION_UPDATES.has(sessionUpdate)) { - hasVisibleAgentOutput = true; - } - if (sessionUpdate === "agent_message") { - agentMessageFinalized = true; - } - if (sessionUpdate === "user_message_chunk") { - const text = params.update?.content?.text; - if (text && !localUserEchoes.has(text)) { - externalUserMessageCount += 1; - } - } - } - } - - return { - hasTurnEnd, - hasAwaitingUserInput, - hasTurnCompleted, - hasTurnFailed, - hasVisibleAgentOutput, - externalUserMessageCount, - agentMessageFinalized, - compacting, - }; -} - -// Strip user_message_chunk entries whose text matches a pending local echo -// (one match per echo). The echo set is mutated so each echo only cancels -// one canonical copy. -function dedupAgainstLocalEchoes( - entries: StoredLogEntry[], - localUserEchoes: Set, -): StoredLogEntry[] { - if (localUserEchoes.size === 0) return entries; - const result: StoredLogEntry[] = []; - for (const entry of entries) { - if ( - entry.type === "notification" && - entry.notification?.method === "session/update" - ) { - const params = entry.notification?.params as SessionNotification; - const sessionUpdate = params?.update?.sessionUpdate; - if (sessionUpdate === "user_message_chunk") { - const text = params?.update?.content?.text; - if (text && localUserEchoes.has(text)) { - localUserEchoes.delete(text); - continue; - } - } - } - result.push(entry); - } - return result; -} - -export interface TaskSession { - taskRunId: string; - taskId: string; - taskTitle?: string; - events: SessionEvent[]; - status: "connecting" | "connected" | "disconnected" | "error"; - isPromptPending: boolean; - // Content of user prompts echoed locally (before the agent writes them to - // the log). Used to dedup the canonical copy against the echo. - localUserEchoes?: Set; - // Terminal backend status for this run, populated by status updates so the - // UI can surface "Run failed" / "Run completed". - terminalStatus?: "failed" | "completed"; - lastError?: string | null; - // True when the user initiated work (new task, sendPrompt, resume) and - // we should play a sound when control returns. False when reconnecting - // to an already-running task to avoid spurious pings. - awaitingPing?: boolean; - // Timestamp when the current prompt started on this device. Used to scale - // the completion sound's playback rate by how long the turn ran. - promptStartedAt?: number; - // True after a user prompt is sent, cleared when the first piece of - // agent output (tool call, message, etc.) arrives. - awaitingAgentOutput?: boolean; - // Timestamp of the last new event received. Used to detect stale local - // sessions (desktop stopped syncing). - lastEventAt?: number; - // Maps toolCallId → cloud requestId for routing permission responses. The - // cloud's permission_response command requires the requestId it generated - // when emitting the original permission_request SSE event; we capture it - // here so the response can be routed back to the awaiting tool call. - cloudPermissionRequestIds?: Record; - pendingPermissions?: Record; - // True while the agent is compacting context. Steering cancels and resends - // the running turn, which would abort an in-flight compaction, so queued - // messages are held until compaction ends. - isCompacting?: boolean; - // True once the user has requested the whole run be stopped, until the run - // reaches a terminal status. Hides the Stop control so it can't be tapped - // twice while the cancel is in flight. - stopRequested?: boolean; -} + CloudTaskSession, + CloudTaskSessionStatePort, +} from "@posthog/core/sessions/cloudTaskSessionService"; +import { create } from "zustand"; -interface TaskSessionStore { - sessions: Record; +interface TaskSessionState { + sessions: Record; focusedTaskId: string | null; - - setFocusedTaskId: (taskId: string | null) => void; - - connectToTask: (task: Task) => Promise; - disconnectFromTask: (taskId: string) => void; - sendPrompt: ( - taskId: string, - prompt: string, - attachments?: PendingAttachment[], - ) => Promise; - sendPermissionResponse: ( - taskId: string, - args: { - toolCallId: string; - optionId: string; - answers?: Record; - customInput?: string; - displayText: string; - }, - ) => Promise; - cancelPrompt: (taskId: string) => Promise; - /** Cancel the whole cloud run. Optimistically marks the session stop-requested - * and reverts on failure. Returns false if there is no session or the API fails. */ - stopRun: (taskId: string) => Promise; - /** Send a prompt now, interrupting the running turn first if one is live. */ - sendInterrupting: ( - taskId: string, - prompt: string, - attachments?: PendingAttachment[], - ) => Promise; - flushQueuedMessages: (taskId: string) => Promise; - /** Flush the queue only if the agent is idle. Used after an in-place edit is - * saved or cancelled: the turn may have ended while the user was editing, so - * nothing else would trigger the turn-end drain. A no-op mid-turn. */ - flushQueuedMessagesIfIdle: (taskId: string) => void; - /** Drop one queued message and resend it now as a steer (interrupt + resend). */ - steerQueuedMessage: (taskId: string, messageId: string) => Promise; - setConfigOption: ( - taskId: string, - configId: string, - value: string, - ) => Promise; - getSessionForTask: (taskId: string) => TaskSession | undefined; - - _handleCloudUpdate: ( - taskRunId: string, - update: CloudTaskUpdatePayload, - ) => void; - _startWatcher: (taskRunId: string, taskId: string) => void; - _stopWatcher: (taskRunId: string) => void; - _resumeCloudRun: ( - taskId: string, - previousRunId: string, - prompt: string, - ) => Promise; + setFocusedTaskId(taskId: string | null): void; } -const watchHandles = new Map(); -const connectAttempts = new Set(); -// Guards against a turn-end batch and a mode toggle racing to flush the same -// queue twice. -const flushingTasks = new Set(); - -function mapTerminalStatus( - status: string | undefined | null, -): "completed" | "failed" | undefined { - if (status === "completed") return "completed"; - if (status === "failed" || status === "cancelled") return "failed"; - return undefined; -} - -export const useTaskSessionStore = create((set, get) => ({ +export const useTaskSessionStore = create((set) => ({ sessions: {}, focusedTaskId: null, + setFocusedTaskId: (focusedTaskId) => set({ focusedTaskId }), +})); - setFocusedTaskId: (taskId) => set({ focusedTaskId: taskId }), - - connectToTask: async (task: Task) => { - const taskId = task.id; - const latestRunId = task.latest_run?.id; - - if (connectAttempts.has(taskId)) { - log.debug("Connection already in progress", { taskId }); - return; - } - - const existing = get().getSessionForTask(taskId); - if (existing && existing.status === "connected") { - log.debug("Already connected to task", { taskId }); - return; - } - - connectAttempts.add(taskId); - - try { - let runId = latestRunId; - let awaitingPing = false; - - if (!runId) { - log.debug("Task has no run yet, starting cloud run", { taskId }); - const updatedTask = await runTaskInCloud(taskId); - runId = updatedTask.latest_run?.id; - if (!runId) { - log.error("Failed to start cloud run"); - return; - } - awaitingPing = true; - } - - set((state) => ({ - sessions: { - ...state.sessions, - [runId]: { - taskRunId: runId, - taskId, - taskTitle: task.title, - events: [], - status: "connecting", - // Assume the run is working until the bootstrap snapshot tells - // us otherwise — the SSE watcher will refine these fields. - isPromptPending: true, - awaitingPing, - promptStartedAt: awaitingPing ? Date.now() : undefined, - awaitingAgentOutput: true, - }, - }, - })); - - get()._startWatcher(runId, taskId); - log.debug("Started SSE watcher", { taskId, runId }); - } catch (error) { - log.error("Failed to connect to task", error); - } finally { - connectAttempts.delete(taskId); - } - }, - - disconnectFromTask: (taskId: string) => { - const session = get().getSessionForTask(taskId); - if (!session) return; - - get()._stopWatcher(session.taskRunId); - - set((state) => { - const { [session.taskRunId]: _, ...rest } = state.sessions; - return { sessions: rest }; - }); - log.debug("Disconnected from task", { taskId }); - }, - - sendPrompt: async ( - taskId: string, - prompt: string, - attachments: PendingAttachment[] = [], - ) => { - const session = get().getSessionForTask(taskId); - if (!session) { - throw new Error("No active session for task"); - } - - // The local echo always shows the plain prompt text in the chat. When - // attachments are present we send a structured cloud-prompt blob on the - // wire (`__twig_cloud_prompt_v1__:…`) so the agent receives the image - // and resource blocks alongside the text. - const wirePayload = - attachments.length > 0 - ? serializeCloudPrompt( - await buildCloudPromptBlocks(prompt, attachments), - ) - : prompt; - - const ts = Date.now(); - const echoAttachments: SessionNotificationAttachment[] = - attachments.length > 0 - ? attachments.map((a) => ({ - kind: a.kind, - uri: a.uri, - fileName: a.fileName, - mimeType: a.mimeType, - })) - : []; - const userEvent: SessionEvent = { - type: "session_update", - ts, - notification: { - update: { - sessionUpdate: "user_message_chunk", - content: { type: "text", text: prompt }, - attachments: echoAttachments.length > 0 ? echoAttachments : undefined, - }, - }, - }; - if (echoAttachments.length > 0) { - useAttachmentEchoStore - .getState() - .recordEcho(session.taskRunId, prompt, echoAttachments); - } - - set((state) => { - const current = state.sessions[session.taskRunId]; - const nextLocalEchoes = new Set(current.localUserEchoes ?? []); - nextLocalEchoes.add(prompt); - return { - sessions: { - ...state.sessions, - [session.taskRunId]: { - ...current, - events: [...current.events, userEvent], - localUserEchoes: nextLocalEchoes, - isPromptPending: true, - awaitingPing: true, - promptStartedAt: ts, - awaitingAgentOutput: true, - }, - }, - }; - }); - - try { - await sendCloudCommand(taskId, session.taskRunId, "user_message", { - content: wirePayload, - }); - log.debug("Sent cloud command user_message", { - taskId, - runId: session.taskRunId, - }); - } catch (err) { - if ( - err instanceof CloudCommandError && - (err.status === 504 || err.status === 502 || err.status === 503) - ) { - log.warn("Transient server error sending prompt, rolling back", { - status: err.status, - taskId, - }); - set((state) => { - const current = state.sessions[session.taskRunId]; - if (!current) return state; - const nextLocalEchoes = new Set(current.localUserEchoes ?? []); - nextLocalEchoes.delete(prompt); - return { - sessions: { - ...state.sessions, - [session.taskRunId]: { - ...current, - events: current.events.filter((e) => e !== userEvent), - localUserEchoes: nextLocalEchoes, - isPromptPending: false, - }, - }, - }; - }); - throw err; - } - - let rollbackError: unknown = err; - if (err instanceof CloudCommandError && err.isSandboxInactive()) { - log.info("Sandbox inactive, creating resume run", { - taskId, - previousRunId: session.taskRunId, - }); - try { - await get()._resumeCloudRun(taskId, session.taskRunId, wirePayload); - return; - } catch (resumeErr) { - log.error("Failed to resume cloud run", resumeErr); - rollbackError = resumeErr; - } - } - - set((state) => { - const current = state.sessions[session.taskRunId]; - if (!current) return state; - const nextLocalEchoes = new Set(current.localUserEchoes ?? []); - nextLocalEchoes.delete(prompt); - return { - sessions: { - ...state.sessions, - [session.taskRunId]: { - ...current, - events: current.events.filter((e) => e !== userEvent), - localUserEchoes: nextLocalEchoes, - isPromptPending: false, - }, - }, - }; - }); - throw rollbackError; - } - }, - - sendPermissionResponse: async (taskId, args) => { - const session = get().getSessionForTask(taskId); - if (!session) { - throw new Error("No active session for task"); - } - - const ts = Date.now(); - const userEvent: SessionEvent = { - type: "session_update", - ts, - notification: { - update: { - sessionUpdate: "user_message_chunk", - content: { type: "text", text: args.displayText }, - }, - }, - }; - - set((state) => { - const current = state.sessions[session.taskRunId]; - if (!current) return state; - const nextLocalEchoes = new Set(current.localUserEchoes ?? []); - nextLocalEchoes.add(args.displayText); - return { - sessions: { - ...state.sessions, - [session.taskRunId]: { - ...current, - events: [...current.events, userEvent], - localUserEchoes: nextLocalEchoes, - isPromptPending: true, - awaitingPing: true, - promptStartedAt: ts, - awaitingAgentOutput: true, - }, - }, - }; - }); - - // The cloud command requires the requestId it generated when emitting - // the permission_request SSE event — toolCallId alone is not sufficient - // for routing the response back to the awaiting tool call. - const cloudRequestId = - session.cloudPermissionRequestIds?.[args.toolCallId] ?? - session.pendingPermissions?.[args.toolCallId]?.requestId; - - set((state) => { - const current = state.sessions[session.taskRunId]; - const currentPermission = current?.pendingPermissions?.[args.toolCallId]; - if (!current || !currentPermission) return state; - return { - sessions: { - ...state.sessions, - [session.taskRunId]: { - ...current, - pendingPermissions: { - ...(current.pendingPermissions ?? {}), - [args.toolCallId]: { - ...currentPermission, - response: { - optionId: args.optionId, - displayText: args.displayText, - ...(args.answers ? { answers: args.answers } : {}), - ...(args.customInput - ? { customInput: args.customInput } - : {}), - }, - }, - }, - }, - }, - }; - }); - - try { - await sendCloudCommand(taskId, session.taskRunId, "permission_response", { - ...(cloudRequestId ? { requestId: cloudRequestId } : {}), - toolCallId: args.toolCallId, - optionId: args.optionId, - ...(args.answers ? { answers: args.answers } : {}), - ...(args.customInput ? { customInput: args.customInput } : {}), - }); - log.debug("Sent permission_response", { - taskId, - runId: session.taskRunId, - toolCallId: args.toolCallId, - requestId: cloudRequestId, - }); - - // One-shot: drop the mapping once we've responded so we don't reuse - // it accidentally. - if (cloudRequestId) { - set((state) => { - const current = state.sessions[session.taskRunId]; - if (!current?.cloudPermissionRequestIds) return state; - const next = { ...current.cloudPermissionRequestIds }; - delete next[args.toolCallId]; - return { - sessions: { - ...state.sessions, - [session.taskRunId]: { - ...current, - cloudPermissionRequestIds: next, - }, - }, - }; - }); - } - } catch (err) { - log.error("Failed to send permission_response", err); - set((state) => { - const current = state.sessions[session.taskRunId]; - if (!current) return state; - const nextLocalEchoes = new Set(current.localUserEchoes ?? []); - nextLocalEchoes.delete(args.displayText); - const currentPermission = current.pendingPermissions?.[args.toolCallId]; - return { - sessions: { - ...state.sessions, - [session.taskRunId]: { - ...current, - events: current.events.filter((e) => e !== userEvent), - localUserEchoes: nextLocalEchoes, - pendingPermissions: currentPermission - ? { - ...(current.pendingPermissions ?? {}), - [args.toolCallId]: { - ...currentPermission, - response: undefined, - }, - } - : current.pendingPermissions, - isPromptPending: false, - }, - }, - }; - }); - throw err; - } - }, - - setConfigOption: async (taskId, configId, value) => { - const session = get().getSessionForTask(taskId); - if (!session || session.terminalStatus) return; - - try { - await sendCloudCommand(taskId, session.taskRunId, "set_config_option", { - configId, - value, - }); - log.debug("Sent set_config_option", { - taskId, - runId: session.taskRunId, - configId, - value, - }); - } catch (err) { - log.warn("Failed to send set_config_option", { - taskId, - configId, - error: err, - }); - throw err; - } - }, - - cancelPrompt: async (taskId: string) => { - const session = get().getSessionForTask(taskId); - if (!session) return false; - - try { - await sendCloudCommand(taskId, session.taskRunId, "cancel"); - log.debug("Sent cancel command", { - taskId, - runId: session.taskRunId, - }); - - set((state) => ({ - sessions: { - ...state.sessions, - [session.taskRunId]: { - ...state.sessions[session.taskRunId], - isPromptPending: false, - awaitingPing: false, - promptStartedAt: undefined, - awaitingAgentOutput: false, - }, - }, - })); - return true; - } catch (error) { - log.error("Failed to send cancel request", error); - return false; - } - }, - - stopRun: async (taskId: string) => { - const session = get().getSessionForTask(taskId); - if (!session) return false; - const runId = session.taskRunId; - - const previous = { - stopRequested: session.stopRequested, - isPromptPending: session.isPromptPending, - }; - set((state) => ({ - sessions: { - ...state.sessions, - [runId]: { - ...state.sessions[runId], - stopRequested: true, - isPromptPending: false, - }, - }, - })); - - try { - await cancelRun(taskId, runId); - return true; - } catch (error) { - log.error("Failed to stop cloud run", error); - set((state) => { - const current = state.sessions[runId]; - if (!current) return state; - return { - sessions: { - ...state.sessions, - [runId]: { ...current, ...previous }, - }, - }; - }); - return false; - } - }, - - sendInterrupting: async (taskId, prompt, attachments) => { - // The cloud has no mid-turn inject, so steering interrupts the running - // turn and resends as a fresh prompt. - if (get().getSessionForTask(taskId)?.isPromptPending) { - await get().cancelPrompt(taskId); - } - await get().sendPrompt(taskId, prompt, attachments); - }, - - flushQueuedMessages: async (taskId: string) => { - if (flushingTasks.has(taskId)) return; - flushingTasks.add(taskId); - try { - const drained = useMessageQueueStore - .getState() - .drain(taskId, { stopAtEdited: true }); - if (drained.length === 0) return; - - const { text, attachments } = combineQueuedMessages(drained); - try { - await get().sendInterrupting(taskId, text, attachments); - } catch (err) { - log.warn("Failed to flush queued messages, restoring queue", { - taskId, - error: err, - }); - useMessageQueueStore.getState().prepend(taskId, drained); - } - } finally { - flushingTasks.delete(taskId); - } - }, - - flushQueuedMessagesIfIdle: (taskId: string) => { - const session = get().getSessionForTask(taskId); - if ( - session?.status === "connected" && - !session.isPromptPending && - !session.terminalStatus && - !session.isCompacting && - useMessageQueueStore.getState().getQueue(taskId).length > 0 - ) { - get() - .flushQueuedMessages(taskId) - .catch((err) => log.warn("Queue flush failed", err)); - } - }, - - steerQueuedMessage: async (taskId: string, messageId: string) => { - const session = get().getSessionForTask(taskId); - // Steering only makes sense against a live turn. Mid-compaction it would - // abort the compaction; with no turn running there is nothing to interrupt - // and the message drains via the normal turn-end flush. - if (!session || !session.isPromptPending || session.isCompacting) return; - - const message = useMessageQueueStore - .getState() - .getQueue(taskId) - .find((m) => m.id === messageId); - if (!message) return; - - useMessageQueueStore.getState().remove(taskId, messageId); - try { - await get().sendInterrupting( - taskId, - message.content, - message.attachments, - ); - } catch (err) { - // Restore at the head so a failed steer never silently drops the message. - useMessageQueueStore.getState().prepend(taskId, [message]); - throw err; - } - }, - - getSessionForTask: (taskId: string) => { - return Object.values(get().sessions).find((s) => s.taskId === taskId); - }, - - _startWatcher: (taskRunId: string, taskId: string) => { - if (watchHandles.has(taskRunId)) return; - - const handle = watchCloudTask({ - taskId, - runId: taskRunId, - onUpdate: (update) => get()._handleCloudUpdate(taskRunId, update), - }); - watchHandles.set(taskRunId, handle); - }, - - _stopWatcher: (taskRunId: string) => { - const handle = watchHandles.get(taskRunId); - if (handle) { - handle.stop(); - watchHandles.delete(taskRunId); - log.debug("Stopped SSE watcher", { taskRunId }); - } - }, - - _handleCloudUpdate: (taskRunId: string, update: CloudTaskUpdatePayload) => { - if (update.kind === "error") { - set((state) => { - const current = state.sessions[taskRunId]; - if (!current) return state; - return { - sessions: { - ...state.sessions, - [taskRunId]: { - ...current, - status: "error", - isPromptPending: false, - lastError: update.errorMessage, - }, - }, - }; - }); - return; - } - - if (update.kind === "permission_request") { - // The tool_call UI itself comes from the `session/update` log stream; - // this SSE-only payload exists so we can capture the cloud-side - // requestId required to route a permission_response back to the - // correct pending tool call. - const toolCallId = update.toolCall?.toolCallId; - if (toolCallId && update.requestId) { - set((state) => { - const current = state.sessions[taskRunId]; - if (!current) return state; - return { - sessions: { - ...state.sessions, - [taskRunId]: { - ...current, - cloudPermissionRequestIds: { - ...(current.cloudPermissionRequestIds ?? {}), - [toolCallId]: update.requestId, - }, - pendingPermissions: { - ...(current.pendingPermissions ?? {}), - [toolCallId]: { - requestId: update.requestId, - toolCall: update.toolCall, - options: update.options, - }, - }, - }, - }, - }; - }); - } - return; - } - - if (update.kind === "snapshot" || update.kind === "logs") { - const isSnapshot = update.kind === "snapshot"; - - // Snapshot replaces all events; drop pending echoes since the snapshot - // already includes the canonical copies. - const existing = get().sessions[taskRunId]; - const echoSet = isSnapshot - ? new Set() - : new Set(existing?.localUserEchoes ?? []); - - const dedupedEntries = isSnapshot - ? update.newEntries - : dedupAgainstLocalEchoes(update.newEntries, echoSet); - - const events = convertStoredEntriesToEvents(dedupedEntries); - // Snapshots are S3-backed and lose attachment metadata; reattach from - // the local echo store so historical user messages keep their images. - if (isSnapshot) { - reinjectAttachmentEchoes(taskRunId, events); - } - - const analysis = analyzeEntries( - dedupedEntries, - isSnapshot ? new Set() : echoSet, - ); - - const wasAwaitingPing = existing?.awaitingPing ?? false; - const wasPromptPending = existing?.isPromptPending ?? false; - - set((state) => { - const current = state.sessions[taskRunId]; - if (!current) return state; - - let nextIsPromptPending = current.isPromptPending; - if (analysis.externalUserMessageCount > 0) nextIsPromptPending = true; - if (analysis.hasTurnEnd || analysis.agentMessageFinalized) { - nextIsPromptPending = false; - } - - // Snapshots replay historical content — we don't mutate awaitingPing - // based on history, otherwise turn-end markers inside an existing - // run's snapshot would clear the user's pending ping before the - // status block has a chance to fire its (more specific, e.g. - // "task_failed") notification. The status block below is the - // canonical owner of awaitingPing for terminal snapshots. - // - // awaitingPing is only ever set by this-device actions (sendPrompt, - // sendPermissionResponse, fresh runs, resumes). External user - // messages — i.e. another device chatting in the same task — must - // NOT arm it; otherwise mobile would fire notifications for desktop - // activity. Clearing on turn-end / finalized agent message stays. - let nextAwaitingPing = current.awaitingPing; - if ( - !isSnapshot && - (analysis.hasTurnEnd || analysis.agentMessageFinalized) - ) { - nextAwaitingPing = false; - } - - const nextAwaitingAgentOutput = - current.awaitingAgentOutput && !analysis.hasVisibleAgentOutput; - - const nextEvents = isSnapshot - ? events - : events.length > 0 - ? [...current.events, ...events] - : current.events; - - return { - sessions: { - ...state.sessions, - [taskRunId]: { - ...current, - events: nextEvents, - status: "connected", - isPromptPending: nextIsPromptPending, - awaitingPing: nextAwaitingPing, - awaitingAgentOutput: nextAwaitingAgentOutput, - isCompacting: analysis.compacting ?? current.isCompacting, - localUserEchoes: echoSet.size > 0 ? echoSet : undefined, - lastEventAt: events.length > 0 ? Date.now() : current.lastEventAt, - }, - }, - }; - }); - - // Live `logs` deltas fire pings for three turn-boundary cases: - // * agent is blocked on the user (_posthog/awaiting_user_input) - // * agent finished its turn (_posthog/turn_complete / task_complete) - // * agent errored out the turn (_posthog/error) - // The terminal-status block below can't be relied on for these: the - // turn-end log entry arrives first and clears `awaitingPing`, so by - // the time status terminal fires its `preState.awaitingPing` is - // already false. Status-only termination (sandbox killed without a - // turn-end log) still falls through to the status block. Snapshots - // are historical replay — never ping for those. - const shouldPingForAwaitingInput = - !isSnapshot && wasAwaitingPing && analysis.hasAwaitingUserInput; - const shouldPingForTurnComplete = - !isSnapshot && - wasAwaitingPing && - analysis.hasTurnCompleted && - !analysis.hasAwaitingUserInput; - const shouldPingForTurnFailed = - !isSnapshot && - wasAwaitingPing && - analysis.hasTurnFailed && - !analysis.hasAwaitingUserInput && - !analysis.hasTurnCompleted; - const shouldPingNow = - shouldPingForAwaitingInput || - shouldPingForTurnComplete || - shouldPingForTurnFailed; - if (shouldPingNow && usePreferencesStore.getState().pingsEnabled) { - playCompletionSound( - undefined, - undefined, - completionPlaybackRate(existing?.promptStartedAt), - ).catch(() => {}); - Haptics.notificationAsync(Haptics.NotificationFeedbackType.Success); - } - if (shouldPingForAwaitingInput) { - maybePresentLocalNotification({ - taskRunId, - kind: "awaiting_user_input", - }); - } else if (shouldPingForTurnComplete) { - maybePresentLocalNotification({ - taskRunId, - kind: "turn_complete", - }); - } else if (shouldPingForTurnFailed) { - maybePresentLocalNotification({ - taskRunId, - kind: "task_failed", - }); - } - - // Turn just ended on a live delta — drain any queued messages the user - // buffered while it was running. Deferred a tick so the store state is - // committed before the flush reads it. - const after = get().sessions[taskRunId]; - if ( - !isSnapshot && - wasPromptPending && - after && - !after.isPromptPending && - after.status === "connected" && - useMessageQueueStore.getState().getQueue(after.taskId).length > 0 - ) { - const flushTaskId = after.taskId; - setTimeout(() => { - get() - .flushQueuedMessages(flushTaskId) - .catch((err) => log.warn("Queue flush failed", err)); - }, 0); - } - } - - if (update.kind === "status" || update.kind === "snapshot") { - if (isTerminalStatus(update.status)) { - const preState = get().sessions[taskRunId]; - const shouldPing = preState?.awaitingPing ?? false; - const terminal = mapTerminalStatus(update.status); - set((state) => { - const current = state.sessions[taskRunId]; - if (!current) return state; - return { - sessions: { - ...state.sessions, - [taskRunId]: { - ...current, - isPromptPending: false, - terminalStatus: terminal, - lastError: update.errorMessage ?? null, - awaitingPing: false, - }, - }, - }; - }); - if (shouldPing && usePreferencesStore.getState().pingsEnabled) { - playCompletionSound( - undefined, - undefined, - completionPlaybackRate(preState?.promptStartedAt), - ).catch(() => {}); - Haptics.notificationAsync(Haptics.NotificationFeedbackType.Success); - } - if (shouldPing) { - maybePresentLocalNotification({ - taskRunId, - kind: terminal === "failed" ? "task_failed" : "turn_complete", - }); - } - } - } - }, - - _resumeCloudRun: async ( - taskId: string, - previousRunId: string, - prompt: string, - ) => { - const freshTask = await getPostHogApiClient().getTask(taskId); - const previousRun = freshTask.latest_run; - const previousBranch = previousRun?.branch ?? null; - - const composerConfig = - useTaskStore.getState().composerConfigByTaskId[taskId]; - const previousPermissionMode = previousRun?.state?.initial_permission_mode; - const reasoningEffort = - composerConfig?.reasoning ?? previousRun?.reasoning_effort ?? undefined; - const initialPermissionMode = - composerConfig?.mode ?? - (typeof previousPermissionMode === "string" - ? previousPermissionMode - : undefined); - - const updatedTask = await runTaskInCloud(taskId, { - branch: previousBranch, - resumeFromRunId: previousRunId, - pendingUserMessage: prompt, - reasoningEffort, - initialPermissionMode, - rtkEnabled: usePreferencesStore.getState().rtkEnabledCloud, - }); - - const newRun = updatedTask.latest_run; - if (!newRun?.id) { - throw new Error("Resume run was created but has no id"); - } - - get()._stopWatcher(previousRunId); - - set((state) => { - const previousSession = state.sessions[previousRunId]; - if (!previousSession) return state; - const { [previousRunId]: _old, ...rest } = state.sessions; +export const taskSessionStatePort: CloudTaskSessionStatePort = { + getByTaskId: (taskId) => + Object.values(useTaskSessionStore.getState().sessions).find( + (session) => session.taskId === taskId, + ), + getByRunId: (runId) => useTaskSessionStore.getState().sessions[runId], + set: (session) => + useTaskSessionStore.setState((state) => ({ + sessions: { ...state.sessions, [session.taskRunId]: session }, + })), + update: (runId, updater) => + useTaskSessionStore.setState((state) => { + const session = state.sessions[runId]; + if (!session) return state; return { - sessions: { - ...rest, - [newRun.id]: { - ...previousSession, - taskRunId: newRun.id, - status: "connecting", - isPromptPending: true, - awaitingPing: true, - promptStartedAt: Date.now(), - awaitingAgentOutput: true, - }, - }, + sessions: { ...state.sessions, [runId]: updater(session) }, }; - }); - - get()._startWatcher(newRun.id, taskId); - log.debug("Swapped to resume run", { - taskId, - previousRunId, - newRunId: newRun.id, - }); - }, -})); - -// When the app returns from background, iOS may have killed the SSE -// connection. Nudge every active watcher to reconnect so the stream resumes -// with Last-Event-ID. -AppState.addEventListener("change", (nextState) => { - if (nextState !== "active") return; - for (const handle of watchHandles.values()) { - handle.reconnectIfDisconnected(); - } -}); + }), + remove: (runId) => + useTaskSessionStore.setState((state) => { + if (!(runId in state.sessions)) return state; + const { [runId]: _removed, ...sessions } = state.sessions; + return { sessions }; + }), +}; + +export function getTaskSession(taskId: string): CloudTaskSession | undefined { + return taskSessionStatePort.getByTaskId(taskId); +} diff --git a/packages/core/src/cloud-task/cloud-task-engine.ts b/packages/core/src/cloud-task/cloud-task-engine.ts new file mode 100644 index 0000000000..006b0ab6d2 --- /dev/null +++ b/packages/core/src/cloud-task/cloud-task-engine.ts @@ -0,0 +1,2250 @@ +import type { RootLogger, ScopedLogger } from "@posthog/di/logger"; +import type { IAnalytics } from "@posthog/platform/analytics"; +import { + type CloudTaskPermissionRequestUpdate, + isTerminalStatus, + mcpToolKey, + posthogToolMeta, + type StoredLogEntry, + serializeError, + type TaskRunStatus, + TypedEventEmitter, +} from "@posthog/shared"; +import { ANALYTICS_EVENTS } from "@posthog/shared/analytics-events"; +import type { ICloudTaskAuth, McpRelayExecutor } from "./identifiers"; +import { + CloudTaskEvent, + type CloudTaskEvents, + type SendCommandInput, + type SendCommandOutput, + type StopInput, + type StopOutput, + type WatchInput, +} from "./schemas"; +import { type SseEvent, SseEventParser } from "./sse-parser"; + +// Reconnect backoff: flat base delay for the first SSE_RECONNECT_FLAT_ATTEMPTS attempts, then +// exponential up to the cap (0.5, 0.5, 0.5, 1, 2, 4, 8, 16, 30s), spanning ~60s before giving up. +const MAX_SSE_RECONNECT_ATTEMPTS = 9; +const MAX_CUMULATIVE_RECONNECT_ATTEMPTS = 30; +const SSE_RECONNECT_BASE_DELAY_MS = 500; +const SSE_RECONNECT_FLAT_ATTEMPTS = 3; +const SSE_RECONNECT_MAX_DELAY_MS = 30_000; +const SSE_HEALTHY_CONNECTION_MS = 60_000; +const EVENT_BATCH_FLUSH_MS = 16; +const EVENT_BATCH_MAX_SIZE = 50; +const SESSION_LOG_PAGE_LIMIT = 5_000; +const MAX_HANDLED_RELAY_REQUEST_IDS = 1_000; +const MCP_RELAY_METHODS_WITHOUT_APPROVAL = new Set([ + "initialize", + "notifications/initialized", + "ping", + "tools/list", + "prompts/list", + "resources/list", + "resources/templates/list", +]); + +// Authoritative end-of-stream sentinel, matched on the SSE event name (event.event, not data.type). +// The client stops on it without consulting run status. +const STREAM_END_EVENT_NAME = "stream-end"; + +interface SessionLogsPage { + entries: StoredLogEntry[]; + hasMore: boolean; +} + +interface CloudTaskConnectionError { + title: string; + message: string; + retryable: boolean; + autoRetry?: boolean; +} + +class CloudTaskStreamError extends Error { + constructor( + message: string, + public readonly details: CloudTaskConnectionError, + public readonly status?: number, + ) { + super(message); + this.name = "CloudTaskStreamError"; + } +} + +class BackendStreamError extends Error { + constructor(message: string) { + super(message); + this.name = "BackendStreamError"; + } +} + +interface TaskRunResponse { + id: string; + status: TaskRunStatus; + stage?: string | null; + output?: Record | null; + state?: Record | null; + error_message?: string | null; + branch?: string | null; + updated_at?: string; + completed_at?: string | null; +} + +interface TaskRunStateEvent { + type: "task_run_state"; + status?: TaskRunStatus; + stage?: string | null; + output?: Record | null; + state?: Record | null; + error_message?: string | null; + branch?: string | null; + updated_at?: string | null; + completed_at?: string | null; +} + +// Which endpoint a connection reads from. Event ids are only meaningful within their issuing leg. +type StreamLeg = "proxy" | "django"; + +interface WatcherState { + taskId: string; + runId: string; + apiHost: string; + teamId: number; + subscriberCount: number; + sseAbortController: AbortController | null; + reconnectTimeoutId: ReturnType | null; + batchFlushTimeoutId: ReturnType | null; + pendingLogEntries: StoredLogEntry[]; + totalEntryCount: number; + /** On resume the renderer already holds the prior conversation; start live- + * only (no bootstrap fetch/snapshot) seeded at this count so the in-flight + * turn can't collide with a re-fetched snapshot. Null on non-resume watches. */ + resumeFromEntryCount: number | null; + reconnectAttempts: number; + streamErrorAttempts: number; + cumulativeReconnectAttempts: number; + lastEventId: string | null; + // Leg that issued lastEventId, and the leg of the connection currently being read. + lastEventIdLeg: StreamLeg | null; + streamLeg: StreamLeg | null; + // Ids of log entries already ingested on the current leg. The durable stream + // re-sends the tail by id on reconnect/replay, so dropping a seen id here is + // what stops a re-delivered entry (e.g. a `turn_complete`) from being counted + // and emitted twice. Cleared on a leg switch, where the id space changes. + seenEventIds: Set; + lastStatus: TaskRunStatus | null; + lastStage: string | null; + lastOutput: Record | null; + lastErrorMessage: string | null; + lastBranch: string | null; + lastSandboxAlive: boolean | null; + lastStatusUpdatedAt: string | null; + connStartedAt: number; + connSentLastEventId: string | null; + connDataEventsReceived: number; + isBootstrapping: boolean; + hasEmittedSnapshot: boolean; + bufferedLogBatches: StoredLogEntry[][]; + // Live entries emitted since the last snapshot, retained so a re-subscribe snapshot can reconcile + // entries the server has not persisted yet. emitCurrentSnapshot trims this to the still-missing + // set; with no re-subscribe it holds the run's emitted entries until the watch ends. + emittedLogEntries: StoredLogEntry[]; + failed: boolean; + needsPostBootstrapReconnect: boolean; + needsStopAfterBootstrap: boolean; + streamEnded: boolean; + // Consumes one automatic re-bootstrap recovery; re-armed by a data event or healthy connection. + selfHealAttempted: boolean; + // Both streamBaseUrl and streamReadToken non-null => read via the agent-proxy; either null => Django. + streamTargetResolved: boolean; + streamBaseUrl: string | null; + streamReadToken: string | null; + // True once stream_token resolved. False for old servers (404), which fall back to status polling. + durableStreamEnabled: boolean; +} + +function watcherKey(taskId: string, runId: string): string { + return `${taskId}:${runId}`; +} + +function isTaskRunStateEvent(data: unknown): data is TaskRunStateEvent { + return ( + typeof data === "object" && + data !== null && + (data as { type?: string }).type === "task_run_state" + ); +} + +interface SseErrorEventData { + error: string; +} + +function isSseErrorEvent(data: unknown): data is SseErrorEventData { + return ( + typeof data === "object" && + data !== null && + "error" in data && + typeof (data as SseErrorEventData).error === "string" + ); +} + +interface PermissionRequestEventData { + type: "permission_request"; + requestId: string; + toolCall: CloudTaskPermissionRequestUpdate["toolCall"]; + options: CloudTaskPermissionRequestUpdate["options"]; +} + +function isPermissionRequestEvent( + data: unknown, +): data is PermissionRequestEventData { + return ( + typeof data === "object" && + data !== null && + (data as { type?: string }).type === "permission_request" && + typeof (data as { requestId?: string }).requestId === "string" + ); +} + +interface McpRequestEventData { + type: "mcp_request"; + requestId: string; + server: string; + payload: Record; + expiresAt: string; +} + +function isMcpRequestEvent(data: unknown): data is McpRequestEventData { + if (typeof data !== "object" || data === null) return false; + const candidate = data as Partial; + return ( + candidate.type === "mcp_request" && + typeof candidate.requestId === "string" && + typeof candidate.server === "string" && + typeof candidate.payload === "object" && + candidate.payload !== null + ); +} + +/** Prefix marking a desktop-issued relay approval prompt, so `sendCommand` can + * resolve its response locally instead of POSTing it to the sandbox. */ +const RELAY_APPROVAL_REQUEST_PREFIX = "relay-approval:"; + +const RELAY_KEY_SEPARATOR = ""; + +function relayApprovalKey( + runId: string, + server: string, + kind: "method" | "tool", + name: string, +): string { + return [runId, server, kind, name].join(RELAY_KEY_SEPARATOR); +} + +interface RelayApprovalRequest { + approvalKey: string; + title: string; + toolName: string; + rawInput: Record; + mcp: { server: string; tool: string }; +} + +function relayApprovalRequest( + runId: string, + server: string, + payload: Record, +): RelayApprovalRequest | null { + const method = + typeof payload.method === "string" ? payload.method : "unknown"; + if (MCP_RELAY_METHODS_WITHOUT_APPROVAL.has(method)) return null; + + const params = + payload.params && typeof payload.params === "object" + ? (payload.params as Record) + : {}; + + if (method === "tools/call") { + const tool = typeof params.name === "string" ? params.name : "unknown"; + const args = + params.arguments && typeof params.arguments === "object" + ? (params.arguments as Record) + : {}; + const toolName = mcpToolKey({ server, tool }); + return { + approvalKey: relayApprovalKey(runId, server, "tool", tool), + title: `The agent wants to call ${tool} (${server}) on your machine`, + toolName, + rawInput: { ...args, toolName }, + mcp: { server, tool }, + }; + } + + const toolName = `mcp:${server}:${method}`; + return { + approvalKey: relayApprovalKey(runId, server, "method", method), + title: `The agent wants to send ${method} to ${server} on your machine`, + toolName, + rawInput: { method, params }, + mcp: { server, tool: method }, + }; +} + +function isKeepaliveEvent(event: SseEvent): boolean { + return ( + event.event === "keepalive" || + (typeof event.data === "object" && + event.data !== null && + "type" in event.data && + event.data.type === "keepalive") + ); +} + +function createStreamStatusError(status: number): CloudTaskStreamError { + switch (status) { + case 401: + return new CloudTaskStreamError( + "Cloud authentication expired", + { + title: "Cloud authentication expired", + message: "Please reauthenticate and retry the cloud run stream.", + retryable: true, + autoRetry: false, + }, + status, + ); + case 403: + return new CloudTaskStreamError( + "Cloud access denied", + { + title: "Cloud access denied", + message: + "You no longer have access to this cloud run. Reauthenticate and retry.", + retryable: true, + autoRetry: false, + }, + status, + ); + case 404: + return new CloudTaskStreamError( + "Cloud run not found", + { + title: "Cloud run not found", + message: + "This cloud run could not be found. It may have been deleted or moved.", + retryable: false, + autoRetry: false, + }, + status, + ); + case 406: + return new CloudTaskStreamError( + "Cloud stream unavailable", + { + title: "Cloud stream unavailable", + message: + "The backend rejected the live stream request. Restart the backend and retry.", + retryable: true, + autoRetry: false, + }, + status, + ); + default: + return new CloudTaskStreamError( + `Stream request failed with status ${status}`, + { + title: "Cloud stream failed", + message: `The cloud stream request failed with status ${status}. Retry to reconnect.`, + retryable: true, + autoRetry: true, + }, + status, + ); + } +} + +function shouldFailWatcherForFetchStatus(status: number): boolean { + return status === 401 || status === 403 || status === 404; +} + +// 5xx and 429 are momentary: the stream-token endpoint exists but is briefly unavailable, so the +// target stays unresolved and the next reconnect retries instead of caching a Django fallback. +function isTransientStreamTargetStatus(status: number): boolean { + return status >= 500 || status === 429; +} + +// Content-based frequency map keyed by the serialized entry. SSE ids are absent from persisted +// (historical) entries, so the payload itself is the identity used to dedup live against historical. +function buildEntryFrequencyMap( + entries: StoredLogEntry[], +): Map { + const counts = new Map(); + for (const entry of entries) { + const serialized = JSON.stringify(entry); + counts.set(serialized, (counts.get(serialized) ?? 0) + 1); + } + return counts; +} + +// Keeps only entries absent from counts, consuming one occurrence per match so a payload present N +// times in the reference set is suppressed at most N times. Mutates counts. +function filterEntriesNotInFrequencyMap( + entries: StoredLogEntry[], + counts: Map, +): StoredLogEntry[] { + return entries.filter((entry) => { + const serialized = JSON.stringify(entry); + const remaining = counts.get(serialized) ?? 0; + if (remaining <= 0) { + return true; + } + counts.set(serialized, remaining - 1); + return false; + }); +} + +function extractSandboxAlive( + state: Record | null | undefined, +): boolean | null | undefined { + if (!state || !Object.hasOwn(state, "sandbox_alive")) { + return undefined; + } + + const sandboxAlive = state.sandbox_alive; + return typeof sandboxAlive === "boolean" ? sandboxAlive : null; +} + +function sandboxAlivePayload(watcher: { lastSandboxAlive: boolean | null }): { + sandboxAlive?: boolean | null; +} { + return watcher.lastSandboxAlive === null + ? {} + : { sandboxAlive: watcher.lastSandboxAlive }; +} + +export interface CloudTaskEngineDependencies { + auth: ICloudTaskAuth; + analytics: IAnalytics; + logger: RootLogger; + mcpRelayExecutor?: McpRelayExecutor | null; + streamFetch?: CloudTaskFetch; +} + +export type CloudTaskFetch = ( + input: string | URL | Request, + init?: RequestInit, +) => Promise; + +export function createCloudTaskEngine( + dependencies: CloudTaskEngineDependencies, +): CloudTaskEngine { + return new CloudTaskEngine(dependencies); +} + +export class CloudTaskEngine extends TypedEventEmitter { + private watchers = new Map(); + private readonly log: ScopedLogger; + private readonly auth: ICloudTaskAuth; + private readonly analytics: IAnalytics; + private readonly mcpRelayExecutor: McpRelayExecutor | null; + private readonly streamFetch: CloudTaskFetch; + + constructor({ + auth, + analytics, + logger, + mcpRelayExecutor = null, + streamFetch = globalThis.fetch.bind(globalThis), + }: CloudTaskEngineDependencies) { + super(); + this.auth = auth; + this.analytics = analytics; + this.mcpRelayExecutor = mcpRelayExecutor; + this.streamFetch = streamFetch; + this.log = logger.scope("cloud-task"); + } + + /** + * Relay-designated server names per run (docs/cloud-mcp-relay.md). + * In-memory by design: only the client that created a run in this app + * session may execute relay requests for it; requests for undesignated + * runs or names are dropped. + */ + private readonly relayDesignations = new Map>(); + /** requestId dedupe — the event stream is at-least-once and replays on reconnect. */ + private readonly handledRelayRequestIds = new Set(); + private readonly handledRelayRequestOrder: string[] = []; + + /** Sensitive relay requests require desktop-owned approval. */ + private readonly relayAlwaysApprovals = new Set(); + /** Desktop-issued relay approval prompts awaiting a task-view answer. */ + private readonly pendingLocalRelayPrompts = new Map< + string, + { + runId: string; + resolve: (outcome: { + optionId: string | null; + customInput?: string; + }) => void; + } + >(); + + designateRelayedMcpServers(runId: string, servers: string[]): void { + if (servers.length === 0) return; + this.relayDesignations.set(runId, new Set(servers)); + this.log.info("Designated relayed MCP servers for run", { + runId, + servers, + }); + } + + private markRelayRequestHandled(requestId: string): void { + this.handledRelayRequestIds.add(requestId); + this.handledRelayRequestOrder.push(requestId); + if (this.handledRelayRequestOrder.length > MAX_HANDLED_RELAY_REQUEST_IDS) { + const evicted = this.handledRelayRequestOrder.shift(); + if (evicted) this.handledRelayRequestIds.delete(evicted); + } + } + + private async handleMcpRelayRequest( + watcher: WatcherState, + data: McpRequestEventData, + ): Promise { + if (!this.mcpRelayExecutor) return; + const designated = this.relayDesignations.get(watcher.runId); + if (!designated?.has(data.server)) { + // Not created by this client, or a name the run never declared. + return; + } + if (this.handledRelayRequestIds.has(data.requestId)) return; + this.markRelayRequestHandled(data.requestId); + + const expiresAt = Date.parse(data.expiresAt); + if (this.relayRequestExpired(expiresAt)) { + this.log.info("Dropping expired MCP relay request", { + runId: watcher.runId, + server: data.server, + requestId: data.requestId, + }); + return; + } + + const approvalRequest = relayApprovalRequest( + watcher.runId, + data.server, + data.payload, + ); + if (approvalRequest) { + const approval = await this.ensureRelayRequestApproval( + watcher, + approvalRequest, + expiresAt, + ); + if (!approval.approved) { + // Expired prompts get no response: the sandbox has already timed the + // request out, and a late mcp_response would be rejected as unknown. + if (!approval.expired) { + await this.sendRelayResponse(watcher, data, { + error: { code: -32000, message: approval.message }, + }); + } + return; + } + if (this.relayRequestExpired(expiresAt)) return; + } + + let execution: { + payload?: Record; + error?: { code: number; message: string }; + }; + try { + execution = await this.mcpRelayExecutor.execute( + watcher.runId, + data.server, + data.payload, + ); + } catch (error) { + execution = { + error: { + code: -32000, + message: + error instanceof Error + ? error.message + : "MCP relay execution failed", + }, + }; + } + + // Fire-and-forget notifications produce no response payload or error. + if (!execution.payload && !execution.error) return; + + await this.sendRelayResponse(watcher, data, execution); + } + + private relayRequestExpired(expiresAt: number): boolean { + return Number.isFinite(expiresAt) && expiresAt < Date.now(); + } + + private async sendRelayResponse( + watcher: WatcherState, + data: McpRequestEventData, + execution: { + payload?: Record; + error?: { code: number; message: string }; + }, + ): Promise { + try { + await this.sendCommand({ + taskId: watcher.taskId, + runId: watcher.runId, + apiHost: watcher.apiHost, + teamId: watcher.teamId, + method: "mcp_response", + params: { + requestId: data.requestId, + server: data.server, + ...(execution.payload + ? { payload: execution.payload } + : { error: execution.error }), + }, + }); + } catch (error) { + // The sandbox times the request out on its own; nothing to unwind here. + this.log.warn("Failed to deliver mcp_response command", { + runId: watcher.runId, + requestId: data.requestId, + error: serializeError(error), + }); + } + } + + private async ensureRelayRequestApproval( + watcher: WatcherState, + request: RelayApprovalRequest, + expiresAt: number, + ): Promise< + { approved: true } | { approved: false; expired: boolean; message: string } + > { + const { runId } = watcher; + if (this.relayAlwaysApprovals.has(request.approvalKey)) { + return { approved: true }; + } + + const requestId = `${RELAY_APPROVAL_REQUEST_PREFIX}${globalThis.crypto.randomUUID()}`; + this.emit(CloudTaskEvent.Update, { + taskId: watcher.taskId, + runId, + kind: "permission_request" as const, + requestId, + toolCall: { + toolCallId: requestId, + title: request.title, + kind: "other", + rawInput: request.rawInput, + _meta: posthogToolMeta({ + toolName: request.toolName, + mcp: request.mcp, + }), + }, + options: [ + { kind: "allow_once", name: "Yes", optionId: "allow" }, + { + kind: "allow_always", + name: "Yes, always allow", + optionId: "allow_always", + }, + { + kind: "reject_once", + name: "Type here to tell the agent what to do differently", + optionId: "reject", + _meta: { customInput: true }, + }, + ], + }); + + const outcome = await new Promise<{ + optionId: string | null; + customInput?: string; + }>((resolve) => { + this.pendingLocalRelayPrompts.set(requestId, { runId, resolve }); + // The sandbox abandons the request at expiresAt; keep waiting any longer + // and an approval would execute a call whose result nothing consumes. + const waitMs = Number.isFinite(expiresAt) + ? Math.max(0, expiresAt - Date.now()) + : 60_000; + const timer = setTimeout(() => { + if (this.pendingLocalRelayPrompts.delete(requestId)) { + resolve({ optionId: null }); + } + }, waitMs); + timer.unref?.(); + }); + + if (outcome.optionId === "allow_always") { + this.relayAlwaysApprovals.add(request.approvalKey); + return { approved: true }; + } + if (outcome.optionId === "allow") return { approved: true }; + if (outcome.optionId === null) { + return { + approved: false, + expired: true, + message: "The user did not respond in time.", + }; + } + return { + approved: false, + expired: false, + message: outcome.customInput + ? `The user denied this MCP request: ${outcome.customInput}` + : "The user denied this MCP request.", + }; + } + + /** Drop a terminal run's relay approval state and abandon its open prompts. */ + private evictRelayApprovalState(runId: string): void { + const prefix = `${runId}${RELAY_KEY_SEPARATOR}`; + for (const key of [...this.relayAlwaysApprovals]) { + if (key.startsWith(prefix)) this.relayAlwaysApprovals.delete(key); + } + for (const [requestId, prompt] of [...this.pendingLocalRelayPrompts]) { + if (prompt.runId !== runId) continue; + this.pendingLocalRelayPrompts.delete(requestId); + prompt.resolve({ optionId: null }); + } + } + + watch(input: WatchInput): void { + const key = watcherKey(input.taskId, input.runId); + + const existing = this.watchers.get(key); + if (existing) { + existing.subscriberCount++; + this.log.info("Cloud task watcher subscriber added", { + key, + subscribers: existing.subscriberCount, + }); + void this.emitCurrentSnapshot(key); + return; + } + + this.startWatcher(input, 1); + } + + unwatch(taskId: string, runId: string): void { + const key = watcherKey(taskId, runId); + const watcher = this.watchers.get(key); + if (!watcher) { + return; + } + + watcher.subscriberCount--; + if (watcher.subscriberCount <= 0) { + this.stopWatcher(key); + } else { + this.log.info("Cloud task watcher subscriber removed", { + key, + subscribers: watcher.subscriberCount, + }); + } + } + + async retry(taskId: string, runId: string): Promise { + const key = watcherKey(taskId, runId); + const watcher = this.watchers.get(key); + if (!watcher) return; + + if (watcher.reconnectTimeoutId) { + clearTimeout(watcher.reconnectTimeoutId); + watcher.reconnectTimeoutId = null; + } + + watcher.sseAbortController?.abort(); + watcher.sseAbortController = null; + + if (watcher.batchFlushTimeoutId) { + clearTimeout(watcher.batchFlushTimeoutId); + watcher.batchFlushTimeoutId = null; + } + + this.log.info("Retrying cloud task watcher", { + key, + hasSnapshot: watcher.hasEmittedSnapshot, + }); + + // Start over from scratch: a poisoned resume position loops straight back into the same + // failure, so re-bootstrap to re-resolve the read leg and emit a fresh snapshot. + this.resetWatcherForRebootstrap(watcher); + void this.bootstrapWatcher(key); + } + + reconnectIfDisconnected(taskId: string, runId: string): void { + const key = watcherKey(taskId, runId); + const watcher = this.watchers.get(key); + if ( + !watcher || + watcher.sseAbortController || + watcher.reconnectTimeoutId || + watcher.isBootstrapping || + isTerminalStatus(watcher.lastStatus) + ) { + return; + } + + void this.connectSse(key); + } + + // Resets a watcher to its pre-bootstrap state so bootstrapWatcher can rebuild it from server truth. + private resetWatcherForRebootstrap(watcher: WatcherState): void { + watcher.reconnectAttempts = 0; + watcher.streamErrorAttempts = 0; + watcher.cumulativeReconnectAttempts = 0; + watcher.failed = false; + watcher.pendingLogEntries = []; + watcher.bufferedLogBatches = []; + watcher.needsPostBootstrapReconnect = false; + watcher.needsStopAfterBootstrap = false; + watcher.streamEnded = false; + watcher.selfHealAttempted = false; + watcher.lastEventId = null; + watcher.lastEventIdLeg = null; + watcher.streamLeg = null; + // The rebuild re-resolves the read leg, so a retained id could false-match a + // different entry on the next connection — and the leg-switch clear in + // connectSse can't catch it, since lastEventId was just nulled. The re-fetched + // snapshot re-delivers history, so no dedup state is lost. + watcher.seenEventIds.clear(); + watcher.totalEntryCount = 0; + watcher.isBootstrapping = false; + watcher.streamTargetResolved = false; + watcher.streamBaseUrl = null; + watcher.streamReadToken = null; + watcher.durableStreamEnabled = false; + } + + async sendCommand(input: SendCommandInput): Promise { + if (input.method === "permission_response") { + const params = input.params ?? {}; + const requestId = + typeof params.requestId === "string" ? params.requestId : null; + if (requestId?.startsWith(RELAY_APPROVAL_REQUEST_PREFIX)) { + // A desktop-issued relay approval: resolve it locally — the sandbox + // never saw this prompt, so there is nothing to POST. + const pending = this.pendingLocalRelayPrompts.get(requestId); + this.pendingLocalRelayPrompts.delete(requestId); + pending?.resolve({ + optionId: + typeof params.optionId === "string" ? params.optionId : null, + customInput: + typeof params.customInput === "string" + ? params.customInput + : undefined, + }); + return { success: true }; + } + } + + const url = `${input.apiHost}/api/projects/${input.teamId}/tasks/${input.taskId}/runs/${input.runId}/command/`; + const body = { + jsonrpc: "2.0", + method: input.method, + params: input.params ?? {}, + id: `posthog-code-${Date.now()}`, + }; + + try { + const response = await this.auth.authenticatedFetch(url, { + method: "POST", + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify(body), + }); + + if (!response.ok) { + const errorText = await response.text().catch(() => ""); + let errorMessage = `Command failed with status ${response.status}`; + try { + const errorJson = JSON.parse(errorText); + if (errorJson.error?.message) { + errorMessage = errorJson.error.message; + } else if (errorJson.error) { + errorMessage = + typeof errorJson.error === "string" + ? errorJson.error + : JSON.stringify(errorJson.error); + } + } catch { + if (errorText) errorMessage = errorText; + } + + this.log.warn("Cloud task command failed", { + taskId: input.taskId, + runId: input.runId, + method: input.method, + status: response.status, + error: errorMessage, + }); + return { success: false, error: errorMessage }; + } + + const data = (await response.json()) as { + error?: { message?: string }; + result?: unknown; + }; + + if (data.error) { + this.log.warn("Cloud task command returned error", { + taskId: input.taskId, + method: input.method, + error: data.error, + }); + return { + success: false, + error: data.error.message ?? JSON.stringify(data.error), + }; + } + + this.log.info("Cloud task command sent", { + taskId: input.taskId, + runId: input.runId, + method: input.method, + }); + + return { success: true, result: data.result }; + } catch (error) { + const errorMessage = + error instanceof Error ? error.message : "Unknown error"; + this.log.error("Cloud task command error", { + taskId: input.taskId, + method: input.method, + error: errorMessage, + }); + return { success: false, error: errorMessage }; + } + } + + async stop(input: StopInput): Promise { + try { + const context = await this.auth.getCloudContext(); + if (!context) { + return { success: false, error: "No active cloud project" }; + } + const url = `${context.apiHost}/api/projects/${context.teamId}/tasks/${input.taskId}/runs/${input.runId}/cancel/`; + const response = await this.auth.authenticatedFetch(url, { + method: "POST", + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify(input.reason ? { reason: input.reason } : {}), + }); + + if (!response.ok) { + const errorText = await response.text().catch(() => ""); + let errorMessage = `Stop failed with status ${response.status}`; + try { + const errorJson = JSON.parse(errorText) as { error?: unknown }; + if (typeof errorJson.error === "string" && errorJson.error) { + errorMessage = errorJson.error; + } + } catch { + if (errorText) errorMessage = errorText; + } + + this.log.warn("Cloud run stop failed", { + taskId: input.taskId, + runId: input.runId, + status: response.status, + error: errorMessage, + }); + return { + success: false, + error: errorMessage, + retryable: response.status === 503 || response.status >= 500, + }; + } + + const data = (await response.json()) as { status?: string }; + this.log.info("Cloud run stop accepted", { + taskId: input.taskId, + runId: input.runId, + runStatus: data.status, + }); + return { success: true, runStatus: data.status }; + } catch (error) { + const errorMessage = + error instanceof Error ? error.message : "Unknown error"; + this.log.error("Cloud run stop error", { + taskId: input.taskId, + runId: input.runId, + error: errorMessage, + }); + return { success: false, error: errorMessage, retryable: true }; + } + } + + unwatchAll(): void { + for (const key of [...this.watchers.keys()]) { + this.stopWatcher(key); + } + } + + private startWatcher(input: WatchInput, subscriberCount: number): void { + const key = watcherKey(input.taskId, input.runId); + + const watcher: WatcherState = { + taskId: input.taskId, + runId: input.runId, + apiHost: input.apiHost, + teamId: input.teamId, + subscriberCount, + sseAbortController: null, + reconnectTimeoutId: null, + batchFlushTimeoutId: null, + pendingLogEntries: [], + totalEntryCount: 0, + resumeFromEntryCount: input.resumeFromEntryCount ?? null, + reconnectAttempts: 0, + streamErrorAttempts: 0, + cumulativeReconnectAttempts: 0, + lastEventId: null, + lastEventIdLeg: null, + streamLeg: null, + seenEventIds: new Set(), + lastStatus: null, + lastStage: null, + lastOutput: null, + lastErrorMessage: null, + lastBranch: null, + lastSandboxAlive: null, + lastStatusUpdatedAt: null, + connStartedAt: 0, + connSentLastEventId: null, + connDataEventsReceived: 0, + isBootstrapping: false, + hasEmittedSnapshot: false, + bufferedLogBatches: [], + emittedLogEntries: [], + failed: false, + needsPostBootstrapReconnect: false, + needsStopAfterBootstrap: false, + streamEnded: false, + selfHealAttempted: false, + streamTargetResolved: false, + streamBaseUrl: null, + streamReadToken: null, + durableStreamEnabled: false, + }; + + this.watchers.set(key, watcher); + this.log.info("Cloud task watcher started", { key }); + void this.bootstrapWatcher(key); + } + + private stopWatcher(key: string): void { + const watcher = this.watchers.get(key); + if (!watcher) return; + + if (this.relayDesignations.has(watcher.runId)) { + // No watcher → no relay events → nothing executes; release the run's + // live server connections (stdio children included). They reopen + // lazily if the run is watched again. + void this.mcpRelayExecutor?.closeRun?.(watcher.runId).catch(() => {}); + } + + watcher.sseAbortController?.abort(); + + if (watcher.reconnectTimeoutId) { + clearTimeout(watcher.reconnectTimeoutId); + watcher.reconnectTimeoutId = null; + } + + if (watcher.batchFlushTimeoutId) { + clearTimeout(watcher.batchFlushTimeoutId); + watcher.batchFlushTimeoutId = null; + } + + this.flushLogBatch(key); + this.watchers.delete(key); + this.log.info("Cloud task watcher stopped", { key }); + } + + private async bootstrapWatcher(key: string): Promise { + const watcher = this.watchers.get(key); + if (!watcher) return; + + watcher.failed = false; + watcher.needsPostBootstrapReconnect = false; + watcher.needsStopAfterBootstrap = false; + + const run = await this.fetchTaskRun(watcher); + const currentWatcher = this.watchers.get(key); + if (!currentWatcher || currentWatcher !== watcher) return; + if (watcher.failed) return; + + if (!run) { + this.failWatcher(key, { + title: "Failed to load cloud run", + message: "Could not fetch the cloud run state. Retry to reconnect.", + retryable: true, + }); + return; + } + + this.applyTaskRunState(watcher, run); + + if ( + !isTerminalStatus(run.status) && + watcher.resumeFromEntryCount !== null + ) { + watcher.totalEntryCount = watcher.resumeFromEntryCount; + watcher.hasEmittedSnapshot = true; + watcher.isBootstrapping = false; + void this.connectSse(key, { startLatest: true }); + return; + } + + if (isTerminalStatus(run.status)) { + const historicalEntries = await this.fetchAllSessionLogs(watcher); + const terminalWatcher = this.watchers.get(key); + if (!terminalWatcher || terminalWatcher !== watcher) return; + if (watcher.failed) return; + if (!historicalEntries) { + this.failWatcher(key, { + title: "Failed to load task history", + message: + "Could not load the persisted cloud task logs. Retry to reconnect.", + retryable: true, + }); + return; + } + + watcher.totalEntryCount = historicalEntries.length; + watcher.hasEmittedSnapshot = true; + this.emit(CloudTaskEvent.Update, { + taskId: watcher.taskId, + runId: watcher.runId, + kind: "snapshot", + newEntries: historicalEntries, + totalEntryCount: watcher.totalEntryCount, + status: watcher.lastStatus ?? undefined, + stage: watcher.lastStage, + output: watcher.lastOutput, + errorMessage: watcher.lastErrorMessage, + branch: watcher.lastBranch, + ...sandboxAlivePayload(watcher), + }); + this.stopWatcher(key); + return; + } + + watcher.isBootstrapping = true; + watcher.bufferedLogBatches = []; + void this.connectSse(key, { startLatest: true }); + + const historicalEntries = await this.fetchAllSessionLogs(watcher); + const bootstrappingWatcher = this.watchers.get(key); + if (!bootstrappingWatcher || bootstrappingWatcher !== watcher) return; + if (watcher.failed) return; + if (!historicalEntries) { + this.failWatcher(key, { + title: "Failed to load cloud run history", + message: + "Could not load the existing cloud run logs. Retry to reconnect.", + retryable: true, + }); + return; + } + + // Flush any pending live entries into the bootstrap buffer before snapshot. + this.flushLogBatch(key); + + watcher.totalEntryCount = historicalEntries.length; + watcher.hasEmittedSnapshot = true; + + this.emit(CloudTaskEvent.Update, { + taskId: watcher.taskId, + runId: watcher.runId, + kind: "snapshot", + newEntries: historicalEntries, + totalEntryCount: watcher.totalEntryCount, + status: watcher.lastStatus ?? undefined, + stage: watcher.lastStage, + output: watcher.lastOutput, + errorMessage: watcher.lastErrorMessage, + branch: watcher.lastBranch, + ...sandboxAlivePayload(watcher), + }); + + watcher.isBootstrapping = false; + this.drainBufferedLogBatches(key, historicalEntries); + + if (watcher.failed) { + return; + } + + if (watcher.needsStopAfterBootstrap) { + watcher.needsStopAfterBootstrap = false; + await this.finalizeWatcherStop(key); + return; + } + + if (watcher.needsPostBootstrapReconnect) { + watcher.needsPostBootstrapReconnect = false; + this.scheduleReconnect(key, undefined, { countAttempt: false }); + } + + void this.verifyPostBootstrapStatus(key); + } + + private async verifyPostBootstrapStatus(key: string): Promise { + const watcher = this.watchers.get(key); + if (!watcher) return; + if (isTerminalStatus(watcher.lastStatus)) return; + + const run = await this.fetchTaskRun(watcher); + const currentWatcher = this.watchers.get(key); + if (!currentWatcher || currentWatcher !== watcher) return; + if (!run) return; + + if (!this.applyTaskRunState(watcher, run)) return; + if (isTerminalStatus(watcher.lastStatus)) return; + + this.emitStatusUpdate(watcher); + } + + private emitStatusUpdate(watcher: WatcherState): void { + this.emit(CloudTaskEvent.Update, { + taskId: watcher.taskId, + runId: watcher.runId, + kind: "status", + status: watcher.lastStatus ?? undefined, + stage: watcher.lastStage, + output: watcher.lastOutput, + errorMessage: watcher.lastErrorMessage, + branch: watcher.lastBranch, + ...sandboxAlivePayload(watcher), + }); + } + + private async connectSse( + key: string, + options?: { startLatest?: boolean }, + ): Promise { + const watcher = this.watchers.get(key); + if (!watcher) return; + + const controller = new AbortController(); + watcher.sseAbortController = controller; + + watcher.connStartedAt = 0; + watcher.connDataEventsReceived = 0; + + // Resolve the read target once (proxy URL + token, or Django), reused across reconnects. + if (!watcher.streamTargetResolved) { + await this.resolveStreamTarget(watcher); + const resolvedWatcher = this.watchers.get(key); + if ( + !resolvedWatcher || + resolvedWatcher !== watcher || + controller.signal.aborted + ) { + return; + } + } + + const usingProxy = Boolean( + watcher.streamBaseUrl && watcher.streamReadToken, + ); + const base = usingProxy + ? watcher.streamBaseUrl?.replace(/\/+$/, "") + : watcher.apiHost; + const leg: StreamLeg = usingProxy ? "proxy" : "django"; + // Proxy and Django id spaces are unrelated, so drop the resume position on a leg switch and + // let start=latest plus the next snapshot cover the gap. + if (watcher.lastEventId && watcher.lastEventIdLeg !== leg) { + this.log.info("Cloud task stream leg changed, dropping resume position", { + key, + from: watcher.lastEventIdLeg, + to: leg, + }); + watcher.lastEventId = null; + watcher.lastEventIdLeg = null; + // Proxy and Django ids are unrelated, so a retained id could false-match a + // different entry on the new leg. Drop them; the snapshot covers the gap. + watcher.seenEventIds.clear(); + } + watcher.streamLeg = leg; + + // Captured after the leg-switch drop so they reflect what this connection actually sends. + watcher.connSentLastEventId = watcher.lastEventId; + const startLatest = Boolean(options?.startLatest && !watcher.lastEventId); + const url = new URL( + usingProxy + ? `${base}/v1/runs/${encodeURIComponent(watcher.runId)}/stream` + : `${base}/api/projects/${watcher.teamId}/tasks/${encodeURIComponent( + watcher.taskId, + )}/runs/${encodeURIComponent(watcher.runId)}/stream/`, + ); + if (startLatest) { + url.searchParams.set("start", "latest"); + } + const headers: Record = { + Accept: "text/event-stream", + }; + if (watcher.lastEventId) { + headers["Last-Event-ID"] = watcher.lastEventId; + } + if (usingProxy) { + headers.Authorization = `Bearer ${watcher.streamReadToken}`; + } + + // Info so every stream attempt is visible in the logs; Bearer token redacted. + this.log.info(`Opening cloud task stream via ${leg}: ${url.toString()}`, { + key, + leg, + usingProxy, + durableStream: watcher.durableStreamEnabled, + method: "GET", + streamUrl: url.toString(), + lastEventId: watcher.lastEventId, + startLatest, + headers: usingProxy + ? { ...headers, Authorization: "Bearer " } + : headers, + }); + + const parser = new SseEventParser((message, data) => + this.log.warn(message, data), + ); + const decoder = new TextDecoder(); + + // Track how long the body stayed open so healthy long-lived connections cut by churn + // aren't penalized as failed reconnects (see SSE_HEALTHY_CONNECTION_MS). + let connectedAt = 0; + let streamWasEstablished = false; + let bytesReceived = 0; + let eventsReceived = 0; + + try { + // The proxy authenticates with the run-scoped Bearer token; the Django leg uses the session. + const response = usingProxy + ? await this.streamFetch(url.toString(), { + method: "GET", + headers, + signal: controller.signal, + }) + : await this.auth.authenticatedFetch(url.toString(), { + method: "GET", + headers, + signal: controller.signal, + }); + + this.log.info( + `Cloud task stream response ${response.status} ${ + response.ok ? "ok" : "FAILED" + } via ${leg}`, + { + key, + leg, + status: response.status, + ok: response.ok, + streamUrl: url.toString(), + }, + ); + + if (!response.ok) { + throw createStreamStatusError(response.status); + } + + if (!response.body) { + throw new Error("Stream response did not include a body"); + } + + connectedAt = Date.now(); + streamWasEstablished = true; + watcher.connStartedAt = connectedAt; + + this.log.info(`Cloud task SSE connected via ${leg}: ${url.toString()}`, { + key, + leg, + streamUrl: url.toString(), + sentLastEventId: watcher.connSentLastEventId, + startLatest, + status: response.status, + server: response.headers.get("server"), + via: response.headers.get("via"), + cfRay: response.headers.get("cf-ray"), + cfCacheStatus: response.headers.get("cf-cache-status"), + xAccelBuffering: response.headers.get("x-accel-buffering"), + contentType: response.headers.get("content-type"), + requestId: response.headers.get("x-request-id"), + }); + + const reader = response.body.getReader(); + + while (true) { + const { done, value } = await reader.read(); + if (done) { + break; + } + + if (!value) { + continue; + } + + bytesReceived += value.byteLength; + const chunk = decoder.decode(value, { stream: true }); + const events = parser.parse(chunk); + for (const event of events) { + eventsReceived += 1; + const backendError = this.handleSseEvent(key, event); + if (backendError) { + throw backendError; + } + } + } + + const trailingEvents = parser.parse(decoder.decode()); + for (const event of trailingEvents) { + const backendError = this.handleSseEvent(key, event); + if (backendError) { + throw backendError; + } + } + + this.flushLogBatch(key); + + if (controller.signal.aborted) { + return; + } + + this.log.info("Cloud task stream closed cleanly", { + key, + connectionDurationMs: Date.now() - connectedAt, + bytesReceived, + eventsReceived, + dataEventsReceived: watcher.connDataEventsReceived, + lastEventId: watcher.lastEventId, + }); + + // A long-lived clean close is healthy churn, not a loop: clear the cumulative budget so an + // idle run can ride out proxy timeout cycles, while instant-EOF loops still exhaust it. + const completedWatcher = this.watchers.get(key); + if ( + completedWatcher && + streamWasEstablished && + Date.now() - connectedAt >= SSE_HEALTHY_CONNECTION_MS + ) { + completedWatcher.cumulativeReconnectAttempts = 0; + completedWatcher.selfHealAttempted = false; + } + + await this.handleStreamCompletion(key, { reconnectOnDisconnect: true }); + } catch (error) { + this.flushLogBatch(key); + + if (controller.signal.aborted) { + return; + } + + // Proxy-leg 401: the read token expired or its signing key rotated. Re-resolve to mint a + // fresh token (or route back to Django) instead of failing. Django-leg 401 stays fatal below. + const unauthorizedWatcher = this.watchers.get(key); + if ( + error instanceof CloudTaskStreamError && + error.status === 401 && + unauthorizedWatcher?.streamBaseUrl + ) { + // Keep durableStreamEnabled set: clearing it would route this disconnect through legacy + // status polling, which can stop the watch on a terminal status before stream-end arrives. + // The next connectSse re-resolves the target and resolveStreamTarget re-derives durability. + unauthorizedWatcher.streamTargetResolved = false; + unauthorizedWatcher.streamBaseUrl = null; + unauthorizedWatcher.streamReadToken = null; + this.log.info("Cloud task stream proxy token rejected, re-resolving", { + key, + }); + await this.handleStreamCompletion(key, { + reconnectOnDisconnect: true, + reconnectError: error, + countReconnectAttempt: true, + }); + return; + } + + if ( + error instanceof CloudTaskStreamError && + error.details.autoRetry === false + ) { + this.failWatcher(key, error.details); + return; + } + + const errorMessage = + error instanceof Error ? error.message : "Unknown stream error"; + + const isBackendError = error instanceof BackendStreamError; + const wasHealthyStream = + !isBackendError && + streamWasEstablished && + Date.now() - connectedAt >= SSE_HEALTHY_CONNECTION_MS; + + const errorWatcher = this.watchers.get(key); + if (errorWatcher) { + if (isBackendError) { + errorWatcher.streamErrorAttempts += 1; + } else if (wasHealthyStream) { + errorWatcher.streamErrorAttempts = 0; + // A healthy-length connection proves timeout cycling, not a loop. + errorWatcher.cumulativeReconnectAttempts = 0; + errorWatcher.selfHealAttempted = false; + } + } + + this.log.warn("Cloud task stream error", { + key, + leg, + streamUrl: url.toString(), + error: errorMessage, + errorDetail: serializeError(error), + wasHealthyStream, + isBackendError, + streamWasEstablished, + connectionDurationMs: streamWasEstablished + ? Date.now() - connectedAt + : 0, + bytesReceived, + eventsReceived, + dataEventsReceived: errorWatcher?.connDataEventsReceived ?? 0, + lastEventId: errorWatcher?.lastEventId ?? null, + reconnectAttempts: errorWatcher?.reconnectAttempts ?? 0, + streamErrorAttempts: errorWatcher?.streamErrorAttempts ?? 0, + cumulativeReconnectAttempts: + errorWatcher?.cumulativeReconnectAttempts ?? 0, + }); + await this.handleStreamCompletion(key, { + reconnectOnDisconnect: true, + reconnectError: error, + countReconnectAttempt: !isBackendError && !wasHealthyStream, + }); + } finally { + const currentWatcher = this.watchers.get(key); + if (currentWatcher?.sseAbortController === controller) { + currentWatcher.sseAbortController = null; + } + } + } + + // Returns a BackendStreamError when the stream carries an error event so the caller can throw at + // the read site; returns null otherwise. It does not throw, so a single event cannot unwind the + // reader loop unexpectedly. + private handleSseEvent( + key: string, + event: SseEvent, + ): BackendStreamError | null { + const watcher = this.watchers.get(key); + if (!watcher || watcher.failed) return null; + + if (event.id) { + watcher.lastEventId = event.id; + watcher.lastEventIdLeg = watcher.streamLeg; + } + + if (event.event === "error") { + const message = isSseErrorEvent(event.data) + ? event.data.error + : "Unknown stream error"; + return new BackendStreamError(message); + } + + if (event.event === STREAM_END_EVENT_NAME) { + // The run's stream is durably complete. Mark it so completion stops instead + // of reconnecting, independent of run status. The connection will close + // naturally (clean EOF) right after this sentinel. + watcher.streamEnded = true; + return null; + } + + // A keepalive or real event proves the transport recovered. A keepalive does not clear the + // backend-error budget, which only a real data event below resets. + watcher.reconnectAttempts = 0; + + if (isKeepaliveEvent(event)) { + return null; + } + + // A real data event proves the stream materialized; clear the remaining budgets and re-arm self-heal. + watcher.streamErrorAttempts = 0; + watcher.cumulativeReconnectAttempts = 0; + watcher.selfHealAttempted = false; + + watcher.connDataEventsReceived += 1; + if (watcher.connDataEventsReceived === 1 && watcher.connSentLastEventId) { + this.log.info("Cloud task SSE resumed", { + key, + resumedFrom: watcher.connSentLastEventId, + firstEventIdAfterResume: event.id ?? null, + }); + } + + if (isTaskRunStateEvent(event.data)) { + if (this.applyTaskRunState(watcher, event.data)) { + if (!watcher.isBootstrapping && !isTerminalStatus(watcher.lastStatus)) { + this.emit(CloudTaskEvent.Update, { + taskId: watcher.taskId, + runId: watcher.runId, + kind: "status", + status: watcher.lastStatus ?? undefined, + stage: watcher.lastStage, + output: watcher.lastOutput, + errorMessage: watcher.lastErrorMessage, + branch: watcher.lastBranch, + ...sandboxAlivePayload(watcher), + }); + } + } + return null; + } + + // Drop a re-delivered event by its stream id. The durable stream re-sends + // the tail on reconnect/replay: each re-sent log entry would otherwise be + // counted as a new entry (advancing totalEntryCount past the renderer's + // processedLineCount guard) and emitted again — the root cause of duplicate + // transcript entries and back-to-back completion notifications — and a + // re-sent permission_request frame would re-surface an already-answered + // question as a fresh pending card. Events without an id (legacy servers) + // fall through and are handled downstream. + const eventId = event.id; + if (eventId !== undefined) { + if (watcher.seenEventIds.has(eventId)) { + return null; + } + watcher.seenEventIds.add(eventId); + } + + if (isMcpRequestEvent(event.data)) { + void this.handleMcpRelayRequest(watcher, event.data); + return null; + } + + if (isPermissionRequestEvent(event.data)) { + this.emit(CloudTaskEvent.Update, { + taskId: watcher.taskId, + runId: watcher.runId, + kind: "permission_request" as const, + requestId: event.data.requestId, + toolCall: event.data.toolCall, + options: event.data.options, + }); + return null; + } + + watcher.pendingLogEntries.push(event.data as StoredLogEntry); + if (watcher.pendingLogEntries.length >= EVENT_BATCH_MAX_SIZE) { + this.flushLogBatch(key); + return null; + } + + if (!watcher.batchFlushTimeoutId) { + watcher.batchFlushTimeoutId = setTimeout(() => { + watcher.batchFlushTimeoutId = null; + this.flushLogBatch(key); + }, EVENT_BATCH_FLUSH_MS); + } + + return null; + } + + private flushLogBatch(key: string): void { + const watcher = this.watchers.get(key); + if (!watcher || watcher.pendingLogEntries.length === 0) return; + + if (watcher.batchFlushTimeoutId) { + clearTimeout(watcher.batchFlushTimeoutId); + watcher.batchFlushTimeoutId = null; + } + + const entries = watcher.pendingLogEntries; + watcher.pendingLogEntries = []; + + if (watcher.isBootstrapping) { + watcher.bufferedLogBatches.push(entries); + return; + } + + watcher.totalEntryCount += entries.length; + this.rememberEmittedLogEntries(watcher, entries); + + this.emit(CloudTaskEvent.Update, { + taskId: watcher.taskId, + runId: watcher.runId, + kind: "logs", + newEntries: entries, + totalEntryCount: watcher.totalEntryCount, + }); + } + + private drainBufferedLogBatches( + key: string, + historicalEntries: StoredLogEntry[], + ): void { + const watcher = this.watchers.get(key); + if (!watcher || watcher.bufferedLogBatches.length === 0) return; + + const historicalCounts = buildEntryFrequencyMap(historicalEntries); + + for (const entries of watcher.bufferedLogBatches) { + const dedupedEntries = filterEntriesNotInFrequencyMap( + entries, + historicalCounts, + ); + + if (dedupedEntries.length === 0) { + continue; + } + + watcher.totalEntryCount += dedupedEntries.length; + this.rememberEmittedLogEntries(watcher, dedupedEntries); + this.emit(CloudTaskEvent.Update, { + taskId: watcher.taskId, + runId: watcher.runId, + kind: "logs", + newEntries: dedupedEntries, + totalEntryCount: watcher.totalEntryCount, + }); + } + + watcher.bufferedLogBatches = []; + } + + private rememberEmittedLogEntries( + watcher: WatcherState, + entries: StoredLogEntry[], + ): void { + watcher.emittedLogEntries.push(...entries); + } + + private mergeHistoricalAndEmittedEntries( + historicalEntries: StoredLogEntry[], + emittedEntries: StoredLogEntry[], + ): { + snapshotEntries: StoredLogEntry[]; + missingEmittedEntries: StoredLogEntry[]; + } { + if (emittedEntries.length === 0) { + return { snapshotEntries: historicalEntries, missingEmittedEntries: [] }; + } + + const historicalCounts = buildEntryFrequencyMap(historicalEntries); + const missingEmittedEntries = filterEntriesNotInFrequencyMap( + emittedEntries, + historicalCounts, + ); + + return { + snapshotEntries: [...historicalEntries, ...missingEmittedEntries], + missingEmittedEntries, + }; + } + + private async emitCurrentSnapshot(key: string): Promise { + const watcher = this.watchers.get(key); + if (!watcher || watcher.failed) return; + + const historicalEntries = await this.fetchAllSessionLogs(watcher); + const currentWatcher = this.watchers.get(key); + if (!currentWatcher || currentWatcher !== watcher || watcher.failed) { + return; + } + + if (!historicalEntries) { + this.log.warn("Cloud task snapshot replay failed", { + taskId: watcher.taskId, + runId: watcher.runId, + }); + return; + } + + const { snapshotEntries, missingEmittedEntries } = + this.mergeHistoricalAndEmittedEntries( + historicalEntries, + watcher.emittedLogEntries, + ); + watcher.emittedLogEntries = missingEmittedEntries; + if (snapshotEntries.length > watcher.totalEntryCount) { + watcher.totalEntryCount = snapshotEntries.length; + } + + this.emit(CloudTaskEvent.Update, { + taskId: watcher.taskId, + runId: watcher.runId, + kind: "snapshot", + newEntries: snapshotEntries, + totalEntryCount: snapshotEntries.length, + status: watcher.lastStatus ?? undefined, + stage: watcher.lastStage, + output: watcher.lastOutput, + errorMessage: watcher.lastErrorMessage, + branch: watcher.lastBranch, + ...sandboxAlivePayload(watcher), + }); + } + + private failWatcher(key: string, error: CloudTaskConnectionError): void { + const watcher = this.watchers.get(key); + if (!watcher) return; + + this.log.warn("Cloud task watcher failed", { + key, + errorTitle: error.title, + retryable: error.retryable, + status: watcher.lastStatus, + wasBootstrapping: watcher.isBootstrapping, + reconnectAttempts: watcher.reconnectAttempts, + cumulativeReconnectAttempts: watcher.cumulativeReconnectAttempts, + totalEntryCount: watcher.totalEntryCount, + lastEventId: watcher.lastEventId, + }); + + this.analytics.track(ANALYTICS_EVENTS.CLOUD_STREAM_DISCONNECTED, { + task_id: watcher.taskId, + run_id: watcher.runId, + team_id: watcher.teamId, + error_title: error.title, + retryable: error.retryable, + reconnect_attempts: watcher.reconnectAttempts, + stream_error_attempts: watcher.streamErrorAttempts, + cumulative_reconnect_attempts: watcher.cumulativeReconnectAttempts, + was_bootstrapping: watcher.isBootstrapping, + }); + + watcher.failed = true; + watcher.isBootstrapping = false; + watcher.pendingLogEntries = []; + watcher.bufferedLogBatches = []; + + if (watcher.reconnectTimeoutId) { + clearTimeout(watcher.reconnectTimeoutId); + watcher.reconnectTimeoutId = null; + } + + if (watcher.batchFlushTimeoutId) { + clearTimeout(watcher.batchFlushTimeoutId); + watcher.batchFlushTimeoutId = null; + } + + watcher.sseAbortController?.abort(); + watcher.sseAbortController = null; + + this.emit(CloudTaskEvent.Update, { + taskId: watcher.taskId, + runId: watcher.runId, + kind: "error", + errorTitle: error.title, + errorMessage: error.message, + retryable: error.retryable, + }); + } + + private scheduleReconnect( + key: string, + error?: unknown, + options: { countAttempt?: boolean } = {}, + ): void { + const watcher = this.watchers.get(key); + // Status-unaware: the loop only stops on the stream-end sentinel or budget exhaustion below. + if (!watcher || watcher.failed) { + return; + } + + if (watcher.reconnectTimeoutId) { + clearTimeout(watcher.reconnectTimeoutId); + } + + // Bounds runaway loops that clean-EOF (countAttempt=false) and dodge reconnectAttempts. + watcher.cumulativeReconnectAttempts += 1; + const countAttempt = options.countAttempt ?? true; + if (countAttempt) { + watcher.reconnectAttempts += 1; + } + + if ( + watcher.cumulativeReconnectAttempts > MAX_CUMULATIVE_RECONNECT_ATTEMPTS + ) { + // A poisoned resume position burns the budget without an error frame. Rebuild once from + // scratch (the app-restart recovery) before failing; if it loops straight back, fail for real. + if (!watcher.selfHealAttempted) { + watcher.reconnectTimeoutId = null; + this.log.warn( + "Cloud task stream looping without events, re-bootstrapping", + { key }, + ); + this.resetWatcherForRebootstrap(watcher); + // Set after the reset (which clears it): consumes the single allowed self-heal so a + // straight-back loop fails next time instead of re-bootstrapping forever. + watcher.selfHealAttempted = true; + void this.bootstrapWatcher(key); + return; + } + this.failWatcher(key, { + title: "Cloud run unreachable", + message: + "Could not maintain a connection to the cloud run after many attempts. Click retry once the issue is resolved.", + retryable: true, + }); + return; + } + + // Fail once either budget (transport reconnect or backend stream-error) is exhausted. + const attemptCount = Math.max( + watcher.reconnectAttempts, + watcher.streamErrorAttempts, + ); + if (attemptCount > MAX_SSE_RECONNECT_ATTEMPTS) { + const details = + error instanceof CloudTaskStreamError + ? error.details + : { + title: "Cloud stream disconnected", + message: + "Lost connection to the cloud run stream. Retry to reconnect.", + retryable: true, + }; + this.failWatcher(key, details); + return; + } + + const backoffAttempts = + error instanceof BackendStreamError + ? watcher.streamErrorAttempts + : watcher.reconnectAttempts; + const delay = Math.min( + SSE_RECONNECT_BASE_DELAY_MS * + 2 ** Math.max(backoffAttempts - SSE_RECONNECT_FLAT_ATTEMPTS, 0), + SSE_RECONNECT_MAX_DELAY_MS, + ); + + watcher.reconnectTimeoutId = setTimeout(() => { + const currentWatcher = this.watchers.get(key); + if (!currentWatcher) return; + currentWatcher.reconnectTimeoutId = null; + void this.connectSse(key, { + startLatest: + currentWatcher.isBootstrapping || currentWatcher.hasEmittedSnapshot, + }); + }, delay); + } + + private async handleStreamCompletion( + key: string, + options: { + reconnectOnDisconnect: boolean; + reconnectError?: unknown; + countReconnectAttempt?: boolean; + }, + ): Promise { + const watcher = this.watchers.get(key); + if (!watcher) return; + if (watcher.failed) return; + + const { reconnectOnDisconnect } = options; + + // Bootstrap owns the snapshot lifecycle: stopping mid-bootstrap would discard the backlog and + // buffered live entries. Record intent and let bootstrap finish. + if (watcher.isBootstrapping) { + if (watcher.streamEnded || !reconnectOnDisconnect) { + watcher.needsStopAfterBootstrap = true; + } else { + watcher.needsPostBootstrapReconnect = true; + } + return; + } + + // The stream-end sentinel is the only signal that ends a durable watch. Any disconnect without + // it is transport churn to reconnect through; status is tracked for display only, never to stop. + if (watcher.streamEnded) { + await this.finalizeWatcherStop(key); + return; + } + + // Legacy mode (old server): no sentinel, so poll run status on disconnect to decide stop vs + // reconnect. The reconnect budgets keep the new semantics, so self-heal stays active here too. + if (!watcher.durableStreamEnabled && reconnectOnDisconnect) { + const run = await this.fetchTaskRun(watcher); + const legacyWatcher = this.watchers.get(key); + if (!legacyWatcher || legacyWatcher !== watcher) return; + if (watcher.failed) return; + + if (run) { + this.applyTaskRunState(watcher, run); + } + if (isTerminalStatus(watcher.lastStatus)) { + this.emitStatusUpdate(watcher); + this.stopWatcher(key); + return; + } + if (run) { + this.emitStatusUpdate(watcher); + } + this.scheduleReconnect(key, options.reconnectError, { + countAttempt: options.countReconnectAttempt ?? false, + }); + return; + } + + // All callers pass reconnectOnDisconnect, and durable watches only stop via the stream-end + // sentinel or a terminal legacy poll (both handled above); any other disconnect reconnects. + if (reconnectOnDisconnect) { + this.scheduleReconnect(key, options.reconnectError, { + countAttempt: options.countReconnectAttempt ?? false, + }); + } + } + + // Stops a watcher whose stream is durably complete. Repairs the displayed status if the stream + // ended non-terminal (dropped final frame); the poll never decides whether to stop. + private async finalizeWatcherStop(key: string): Promise { + const watcher = this.watchers.get(key); + if (!watcher) return; + + if (!isTerminalStatus(watcher.lastStatus)) { + const run = await this.fetchTaskRun(watcher); + const currentWatcher = this.watchers.get(key); + if (!currentWatcher || currentWatcher !== watcher) return; + if (run) { + this.applyTaskRunState(watcher, run); + } + } + + this.emitStatusUpdate(watcher); + this.stopWatcher(key); + } + + private applyTaskRunState( + watcher: WatcherState, + run: + | Pick< + TaskRunResponse, + | "status" + | "stage" + | "output" + | "state" + | "error_message" + | "branch" + | "updated_at" + > + | TaskRunStateEvent, + ): boolean { + const updatedAt = run.updated_at ?? null; + if ( + updatedAt && + watcher.lastStatusUpdatedAt && + Date.parse(updatedAt) <= Date.parse(watcher.lastStatusUpdatedAt) + ) { + return false; + } + + const nextStatus = run.status ?? watcher.lastStatus; + const nextStage = run.stage ?? null; + const nextOutput = run.output ?? null; + const nextErrorMessage = run.error_message ?? null; + const nextBranch = run.branch ?? null; + const sandboxAlive = extractSandboxAlive(run.state); + const nextSandboxAlive = + sandboxAlive === undefined ? watcher.lastSandboxAlive : sandboxAlive; + + const changed = + nextStatus !== watcher.lastStatus || + nextStage !== watcher.lastStage || + JSON.stringify(nextOutput) !== JSON.stringify(watcher.lastOutput) || + nextErrorMessage !== watcher.lastErrorMessage || + nextBranch !== watcher.lastBranch || + nextSandboxAlive !== watcher.lastSandboxAlive; + + watcher.lastStatus = nextStatus ?? null; + watcher.lastStage = nextStage; + watcher.lastOutput = nextOutput; + watcher.lastErrorMessage = nextErrorMessage; + watcher.lastBranch = nextBranch; + watcher.lastSandboxAlive = nextSandboxAlive; + if (updatedAt) { + watcher.lastStatusUpdatedAt = updatedAt; + } + + // A terminal run gets no further relay requests; drop its designation and + // approval state so the maps don't grow for the lifetime of the app session. + if (isTerminalStatus(watcher.lastStatus)) { + this.relayDesignations.delete(watcher.runId); + this.evictRelayApprovalState(watcher.runId); + } + + return changed; + } + + private async fetchSessionLogsPage( + watcher: WatcherState, + offset: number, + ): Promise { + const url = new URL( + `${watcher.apiHost}/api/projects/${watcher.teamId}/tasks/${watcher.taskId}/runs/${watcher.runId}/session_logs/`, + ); + url.searchParams.set("limit", SESSION_LOG_PAGE_LIMIT.toString()); + url.searchParams.set("offset", offset.toString()); + + try { + const authedResponse = await this.auth.authenticatedFetch( + url.toString(), + { + method: "GET", + }, + ); + + if (!authedResponse.ok) { + this.log.warn("Cloud task session logs fetch failed", { + status: authedResponse.status, + taskId: watcher.taskId, + runId: watcher.runId, + offset, + }); + if (shouldFailWatcherForFetchStatus(authedResponse.status)) { + this.failWatcher( + watcherKey(watcher.taskId, watcher.runId), + createStreamStatusError(authedResponse.status).details, + ); + } + return null; + } + + const raw = await authedResponse.text(); + return { + entries: JSON.parse(raw) as StoredLogEntry[], + hasMore: authedResponse.headers.get("X-Has-More") === "true", + }; + } catch (error) { + this.log.warn("Cloud task session logs fetch error", { + taskId: watcher.taskId, + runId: watcher.runId, + offset, + error, + }); + return null; + } + } + + private async fetchAllSessionLogs( + watcher: WatcherState, + ): Promise { + const entries: StoredLogEntry[] = []; + let offset = 0; + + while (true) { + const page = await this.fetchSessionLogsPage(watcher, offset); + if (!page) { + return null; + } + + entries.push(...page.entries); + if (!page.hasMore || page.entries.length === 0) { + return entries; + } + + offset += page.entries.length; + } + } + + private async resolveStreamTarget(watcher: WatcherState): Promise { + const url = `${watcher.apiHost}/api/projects/${watcher.teamId}/tasks/${watcher.taskId}/runs/${watcher.runId}/stream_token/`; + try { + const response = await this.auth.authenticatedFetch(url, { + method: "GET", + }); + if (!response.ok) { + watcher.streamBaseUrl = null; + watcher.streamReadToken = null; + if (isTransientStreamTargetStatus(response.status)) { + // Transient: read from Django this round but leave the target unresolved so the next + // reconnect retries durable resolution instead of pinning the run to status polling. + this.log.warn("Cloud task stream target temporarily unavailable", { + taskId: watcher.taskId, + runId: watcher.runId, + status: response.status, + }); + return; + } + // Refused, or an old server without the endpoint: read from Django with status polling. + watcher.durableStreamEnabled = false; + watcher.streamTargetResolved = true; + this.log.info("Cloud task stream reading from API host", { + taskId: watcher.taskId, + runId: watcher.runId, + status: response.status, + }); + return; + } + const data = (await response.json()) as { + token?: string; + stream_base_url?: string | null; + }; + watcher.streamReadToken = data.token ?? null; + watcher.streamBaseUrl = data.stream_base_url ?? null; + // The endpoint resolving at all opts this watcher into the status-unaware contract; + // old servers 404 above and stay on legacy status polling. + watcher.durableStreamEnabled = true; + watcher.streamTargetResolved = true; + this.log.info("Cloud task stream target resolved", { + taskId: watcher.taskId, + runId: watcher.runId, + streamBaseUrl: watcher.streamBaseUrl, + hasToken: Boolean(watcher.streamReadToken), + durableStream: watcher.durableStreamEnabled, + }); + } catch (error) { + // Transient failure: leave unresolved so the next reconnect retries and falls back to Django. + watcher.streamBaseUrl = null; + watcher.streamReadToken = null; + this.log.warn("Cloud task stream target resolution failed", { + taskId: watcher.taskId, + runId: watcher.runId, + error, + }); + } + } + + private async fetchTaskRun( + watcher: WatcherState, + ): Promise { + const url = `${watcher.apiHost}/api/projects/${watcher.teamId}/tasks/${watcher.taskId}/runs/${watcher.runId}/`; + + try { + const authedResponse = await this.auth.authenticatedFetch(url, { + method: "GET", + }); + + if (!authedResponse.ok) { + this.log.warn("Cloud task status fetch failed", { + status: authedResponse.status, + taskId: watcher.taskId, + runId: watcher.runId, + }); + if (shouldFailWatcherForFetchStatus(authedResponse.status)) { + this.failWatcher( + watcherKey(watcher.taskId, watcher.runId), + createStreamStatusError(authedResponse.status).details, + ); + } + return null; + } + + return (await authedResponse.json()) as TaskRunResponse; + } catch (error) { + this.log.warn("Cloud task status fetch error", { + taskId: watcher.taskId, + runId: watcher.runId, + error, + }); + return null; + } + } +} diff --git a/packages/core/src/cloud-task/cloud-task-service.test.ts b/packages/core/src/cloud-task/cloud-task-service.test.ts new file mode 100644 index 0000000000..644f9dd7cb --- /dev/null +++ b/packages/core/src/cloud-task/cloud-task-service.test.ts @@ -0,0 +1,27 @@ +import { describe, expect, it, vi } from "vitest"; +import { CloudTaskService } from "./cloud-task"; +import { CloudTaskEngine } from "./cloud-task-engine"; + +describe("CloudTaskService", () => { + it("preserves the injectable service API as a thin engine wrapper", () => { + const scopedLog = { + debug: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + }; + const service = new CloudTaskService( + { + authenticatedFetch: vi.fn(), + getCloudContext: vi.fn(), + }, + { track: vi.fn() } as never, + { ...scopedLog, scope: vi.fn(() => scopedLog) }, + ); + + expect(service).toBeInstanceOf(CloudTaskEngine); + expect(service.watch).toBeTypeOf("function"); + expect(service.retry).toBeTypeOf("function"); + expect(service.unwatchAll).toBeTypeOf("function"); + }); +}); diff --git a/packages/core/src/cloud-task/cloud-task.test.ts b/packages/core/src/cloud-task/cloud-task.test.ts index a7fbf37c66..4178bd3d78 100644 --- a/packages/core/src/cloud-task/cloud-task.test.ts +++ b/packages/core/src/cloud-task/cloud-task.test.ts @@ -5,14 +5,17 @@ const mockNetFetch = vi.hoisted(() => vi.fn()); const mockStreamFetch = vi.hoisted(() => vi.fn()); const mockStreamTokenFetch = vi.hoisted(() => vi.fn()); -// The service now uses global fetch for BOTH authenticated API calls (JSON) -// and SSE streaming. The two used to be distinct (net.fetch vs global fetch). // Route by URL: /stream_token/ → token mock (read-leg resolution), the stream leg // (Django /stream/ or proxy /v1/runs/:run/stream) → stream mock, everything else → API mock. // The token mock has a Django-path default so existing fixtures (which never set it) are untouched. const fetchRouter = vi.hoisted(() => - vi.fn((input: string | Request, init?: RequestInit) => { - const url = typeof input === "string" ? input : input.url; + vi.fn((input: string | URL | Request, init?: RequestInit) => { + const url = + typeof input === "string" + ? input + : input instanceof URL + ? input.toString() + : input.url; const impl = url.includes("/stream_token/") ? mockStreamTokenFetch : /\/stream(\/|\?|$)/.test(url) @@ -22,7 +25,10 @@ const fetchRouter = vi.hoisted(() => }), ); -import { CloudTaskService } from "./cloud-task"; +import { + type CloudTaskEngine, + createCloudTaskEngine, +} from "./cloud-task-engine"; const mockAuthService = { authenticatedFetch: vi.fn(), @@ -86,8 +92,8 @@ async function waitFor( } } -describe("CloudTaskService", () => { - let service: CloudTaskService; +describe("CloudTaskEngine", () => { + let service: CloudTaskEngine; beforeEach(() => { const scopedLog = { @@ -98,11 +104,12 @@ describe("CloudTaskService", () => { }; const loggerMock = { ...scopedLog, scope: vi.fn(() => scopedLog) }; const analyticsMock = { track: vi.fn() }; - service = new CloudTaskService( - mockAuthService as never, - analyticsMock as never, - loggerMock, - ); + service = createCloudTaskEngine({ + auth: mockAuthService as never, + analytics: analyticsMock as never, + logger: loggerMock, + streamFetch: fetchRouter, + }); mockNetFetch.mockReset(); mockStreamFetch.mockReset(); mockStreamTokenFetch.mockReset(); @@ -3077,8 +3084,8 @@ describe("CloudTaskService", () => { }); }); -describe("CloudTaskService MCP relay", () => { - let relayService: CloudTaskService; +describe("CloudTaskEngine MCP relay", () => { + let relayService: CloudTaskEngine; let mcpRelayExecutor: { execute: ReturnType; closeRun: ReturnType; @@ -3099,12 +3106,12 @@ describe("CloudTaskService MCP relay", () => { })), closeRun: vi.fn(async () => {}), }; - relayService = new CloudTaskService( - mockAuthService as never, - analyticsMock as never, - loggerMock, - mcpRelayExecutor as never, - ); + relayService = createCloudTaskEngine({ + auth: mockAuthService as never, + analytics: analyticsMock as never, + logger: loggerMock, + mcpRelayExecutor: mcpRelayExecutor as never, + }); mockNetFetch.mockReset(); mockStreamFetch.mockReset(); diff --git a/packages/core/src/cloud-task/cloud-task.ts b/packages/core/src/cloud-task/cloud-task.ts index 22b5a1ddab..1e2003d85a 100644 --- a/packages/core/src/cloud-task/cloud-task.ts +++ b/packages/core/src/cloud-task/cloud-task.ts @@ -1,2226 +1,35 @@ -import { - ROOT_LOGGER, - type RootLogger, - type ScopedLogger, -} from "@posthog/di/logger"; +import { ROOT_LOGGER, type RootLogger } from "@posthog/di/logger"; import { ANALYTICS_SERVICE, type IAnalytics, } from "@posthog/platform/analytics"; -import type { StoredLogEntry } from "@posthog/shared"; -import { - mcpToolKey, - posthogToolMeta, - serializeError, - TypedEventEmitter, -} from "@posthog/shared"; -import { ANALYTICS_EVENTS } from "@posthog/shared/analytics-events"; import { inject, injectable, optional, preDestroy } from "inversify"; -import type { CloudTaskPermissionRequestUpdate } from "./cloud-task-types"; +import { CloudTaskEngine } from "./cloud-task-engine"; import { CLOUD_TASK_AUTH, type ICloudTaskAuth, MCP_RELAY_EXECUTOR, type McpRelayExecutor, } from "./identifiers"; -import { - CloudTaskEvent, - type CloudTaskEvents, - isTerminalStatus, - type SendCommandInput, - type SendCommandOutput, - type StopInput, - type StopOutput, - type TaskRunStatus, - type WatchInput, -} from "./schemas"; -import { type SseEvent, SseEventParser } from "./sse-parser"; - -// Reconnect backoff: flat base delay for the first SSE_RECONNECT_FLAT_ATTEMPTS attempts, then -// exponential up to the cap (0.5, 0.5, 0.5, 1, 2, 4, 8, 16, 30s), spanning ~60s before giving up. -const MAX_SSE_RECONNECT_ATTEMPTS = 9; -const MAX_CUMULATIVE_RECONNECT_ATTEMPTS = 30; -const SSE_RECONNECT_BASE_DELAY_MS = 500; -const SSE_RECONNECT_FLAT_ATTEMPTS = 3; -const SSE_RECONNECT_MAX_DELAY_MS = 30_000; -const SSE_HEALTHY_CONNECTION_MS = 60_000; -const EVENT_BATCH_FLUSH_MS = 16; -const EVENT_BATCH_MAX_SIZE = 50; -const SESSION_LOG_PAGE_LIMIT = 5_000; -const MAX_HANDLED_RELAY_REQUEST_IDS = 1_000; -const MCP_RELAY_METHODS_WITHOUT_APPROVAL = new Set([ - "initialize", - "notifications/initialized", - "ping", - "tools/list", - "prompts/list", - "resources/list", - "resources/templates/list", -]); - -// Authoritative end-of-stream sentinel, matched on the SSE event name (event.event, not data.type). -// The client stops on it without consulting run status. -const STREAM_END_EVENT_NAME = "stream-end"; - -interface SessionLogsPage { - entries: StoredLogEntry[]; - hasMore: boolean; -} - -interface CloudTaskConnectionError { - title: string; - message: string; - retryable: boolean; - autoRetry?: boolean; -} - -class CloudTaskStreamError extends Error { - constructor( - message: string, - public readonly details: CloudTaskConnectionError, - public readonly status?: number, - ) { - super(message); - this.name = "CloudTaskStreamError"; - } -} - -class BackendStreamError extends Error { - constructor(message: string) { - super(message); - this.name = "BackendStreamError"; - } -} - -interface TaskRunResponse { - id: string; - status: TaskRunStatus; - stage?: string | null; - output?: Record | null; - state?: Record | null; - error_message?: string | null; - branch?: string | null; - updated_at?: string; - completed_at?: string | null; -} - -interface TaskRunStateEvent { - type: "task_run_state"; - status?: TaskRunStatus; - stage?: string | null; - output?: Record | null; - state?: Record | null; - error_message?: string | null; - branch?: string | null; - updated_at?: string | null; - completed_at?: string | null; -} - -// Which endpoint a connection reads from. Event ids are only meaningful within their issuing leg. -type StreamLeg = "proxy" | "django"; - -interface WatcherState { - taskId: string; - runId: string; - apiHost: string; - teamId: number; - subscriberCount: number; - sseAbortController: AbortController | null; - reconnectTimeoutId: ReturnType | null; - batchFlushTimeoutId: ReturnType | null; - pendingLogEntries: StoredLogEntry[]; - totalEntryCount: number; - /** On resume the renderer already holds the prior conversation; start live- - * only (no bootstrap fetch/snapshot) seeded at this count so the in-flight - * turn can't collide with a re-fetched snapshot. Null on non-resume watches. */ - resumeFromEntryCount: number | null; - reconnectAttempts: number; - streamErrorAttempts: number; - cumulativeReconnectAttempts: number; - lastEventId: string | null; - // Leg that issued lastEventId, and the leg of the connection currently being read. - lastEventIdLeg: StreamLeg | null; - streamLeg: StreamLeg | null; - // Ids of log entries already ingested on the current leg. The durable stream - // re-sends the tail by id on reconnect/replay, so dropping a seen id here is - // what stops a re-delivered entry (e.g. a `turn_complete`) from being counted - // and emitted twice. Cleared on a leg switch, where the id space changes. - seenEventIds: Set; - lastStatus: TaskRunStatus | null; - lastStage: string | null; - lastOutput: Record | null; - lastErrorMessage: string | null; - lastBranch: string | null; - lastSandboxAlive: boolean | null; - lastStatusUpdatedAt: string | null; - connStartedAt: number; - connSentLastEventId: string | null; - connDataEventsReceived: number; - isBootstrapping: boolean; - hasEmittedSnapshot: boolean; - bufferedLogBatches: StoredLogEntry[][]; - // Live entries emitted since the last snapshot, retained so a re-subscribe snapshot can reconcile - // entries the server has not persisted yet. emitCurrentSnapshot trims this to the still-missing - // set; with no re-subscribe it holds the run's emitted entries until the watch ends. - emittedLogEntries: StoredLogEntry[]; - failed: boolean; - needsPostBootstrapReconnect: boolean; - needsStopAfterBootstrap: boolean; - streamEnded: boolean; - // Consumes one automatic re-bootstrap recovery; re-armed by a data event or healthy connection. - selfHealAttempted: boolean; - // Both streamBaseUrl and streamReadToken non-null => read via the agent-proxy; either null => Django. - streamTargetResolved: boolean; - streamBaseUrl: string | null; - streamReadToken: string | null; - // True once stream_token resolved. False for old servers (404), which fall back to status polling. - durableStreamEnabled: boolean; -} - -function watcherKey(taskId: string, runId: string): string { - return `${taskId}:${runId}`; -} - -function isTaskRunStateEvent(data: unknown): data is TaskRunStateEvent { - return ( - typeof data === "object" && - data !== null && - (data as { type?: string }).type === "task_run_state" - ); -} - -interface SseErrorEventData { - error: string; -} - -function isSseErrorEvent(data: unknown): data is SseErrorEventData { - return ( - typeof data === "object" && - data !== null && - "error" in data && - typeof (data as SseErrorEventData).error === "string" - ); -} - -interface PermissionRequestEventData { - type: "permission_request"; - requestId: string; - toolCall: CloudTaskPermissionRequestUpdate["toolCall"]; - options: CloudTaskPermissionRequestUpdate["options"]; -} - -function isPermissionRequestEvent( - data: unknown, -): data is PermissionRequestEventData { - return ( - typeof data === "object" && - data !== null && - (data as { type?: string }).type === "permission_request" && - typeof (data as { requestId?: string }).requestId === "string" - ); -} - -interface McpRequestEventData { - type: "mcp_request"; - requestId: string; - server: string; - payload: Record; - expiresAt: string; -} - -function isMcpRequestEvent(data: unknown): data is McpRequestEventData { - if (typeof data !== "object" || data === null) return false; - const candidate = data as Partial; - return ( - candidate.type === "mcp_request" && - typeof candidate.requestId === "string" && - typeof candidate.server === "string" && - typeof candidate.payload === "object" && - candidate.payload !== null - ); -} - -/** Prefix marking a desktop-issued relay approval prompt, so `sendCommand` can - * resolve its response locally instead of POSTing it to the sandbox. */ -const RELAY_APPROVAL_REQUEST_PREFIX = "relay-approval:"; - -const RELAY_KEY_SEPARATOR = ""; - -function relayApprovalKey( - runId: string, - server: string, - kind: "method" | "tool", - name: string, -): string { - return [runId, server, kind, name].join(RELAY_KEY_SEPARATOR); -} - -interface RelayApprovalRequest { - approvalKey: string; - title: string; - toolName: string; - rawInput: Record; - mcp: { server: string; tool: string }; -} - -function relayApprovalRequest( - runId: string, - server: string, - payload: Record, -): RelayApprovalRequest | null { - const method = - typeof payload.method === "string" ? payload.method : "unknown"; - if (MCP_RELAY_METHODS_WITHOUT_APPROVAL.has(method)) return null; - - const params = - payload.params && typeof payload.params === "object" - ? (payload.params as Record) - : {}; - - if (method === "tools/call") { - const tool = typeof params.name === "string" ? params.name : "unknown"; - const args = - params.arguments && typeof params.arguments === "object" - ? (params.arguments as Record) - : {}; - const toolName = mcpToolKey({ server, tool }); - return { - approvalKey: relayApprovalKey(runId, server, "tool", tool), - title: `The agent wants to call ${tool} (${server}) on your machine`, - toolName, - rawInput: { ...args, toolName }, - mcp: { server, tool }, - }; - } - - const toolName = `mcp:${server}:${method}`; - return { - approvalKey: relayApprovalKey(runId, server, "method", method), - title: `The agent wants to send ${method} to ${server} on your machine`, - toolName, - rawInput: { method, params }, - mcp: { server, tool: method }, - }; -} - -function isKeepaliveEvent(event: SseEvent): boolean { - return ( - event.event === "keepalive" || - (typeof event.data === "object" && - event.data !== null && - "type" in event.data && - event.data.type === "keepalive") - ); -} - -function createStreamStatusError(status: number): CloudTaskStreamError { - switch (status) { - case 401: - return new CloudTaskStreamError( - "Cloud authentication expired", - { - title: "Cloud authentication expired", - message: "Please reauthenticate and retry the cloud run stream.", - retryable: true, - autoRetry: false, - }, - status, - ); - case 403: - return new CloudTaskStreamError( - "Cloud access denied", - { - title: "Cloud access denied", - message: - "You no longer have access to this cloud run. Reauthenticate and retry.", - retryable: true, - autoRetry: false, - }, - status, - ); - case 404: - return new CloudTaskStreamError( - "Cloud run not found", - { - title: "Cloud run not found", - message: - "This cloud run could not be found. It may have been deleted or moved.", - retryable: false, - autoRetry: false, - }, - status, - ); - case 406: - return new CloudTaskStreamError( - "Cloud stream unavailable", - { - title: "Cloud stream unavailable", - message: - "The backend rejected the live stream request. Restart the backend and retry.", - retryable: true, - autoRetry: false, - }, - status, - ); - default: - return new CloudTaskStreamError( - `Stream request failed with status ${status}`, - { - title: "Cloud stream failed", - message: `The cloud stream request failed with status ${status}. Retry to reconnect.`, - retryable: true, - autoRetry: true, - }, - status, - ); - } -} - -function shouldFailWatcherForFetchStatus(status: number): boolean { - return status === 401 || status === 403 || status === 404; -} - -// 5xx and 429 are momentary: the stream-token endpoint exists but is briefly unavailable, so the -// target stays unresolved and the next reconnect retries instead of caching a Django fallback. -function isTransientStreamTargetStatus(status: number): boolean { - return status >= 500 || status === 429; -} - -// Content-based frequency map keyed by the serialized entry. SSE ids are absent from persisted -// (historical) entries, so the payload itself is the identity used to dedup live against historical. -function buildEntryFrequencyMap( - entries: StoredLogEntry[], -): Map { - const counts = new Map(); - for (const entry of entries) { - const serialized = JSON.stringify(entry); - counts.set(serialized, (counts.get(serialized) ?? 0) + 1); - } - return counts; -} - -// Keeps only entries absent from counts, consuming one occurrence per match so a payload present N -// times in the reference set is suppressed at most N times. Mutates counts. -function filterEntriesNotInFrequencyMap( - entries: StoredLogEntry[], - counts: Map, -): StoredLogEntry[] { - return entries.filter((entry) => { - const serialized = JSON.stringify(entry); - const remaining = counts.get(serialized) ?? 0; - if (remaining <= 0) { - return true; - } - counts.set(serialized, remaining - 1); - return false; - }); -} - -function extractSandboxAlive( - state: Record | null | undefined, -): boolean | null | undefined { - if (!state || !Object.hasOwn(state, "sandbox_alive")) { - return undefined; - } - - const sandboxAlive = state.sandbox_alive; - return typeof sandboxAlive === "boolean" ? sandboxAlive : null; -} - -function sandboxAlivePayload(watcher: { lastSandboxAlive: boolean | null }): { - sandboxAlive?: boolean | null; -} { - return watcher.lastSandboxAlive === null - ? {} - : { sandboxAlive: watcher.lastSandboxAlive }; -} @injectable() -export class CloudTaskService extends TypedEventEmitter { - private watchers = new Map(); - private readonly log: ScopedLogger; - +export class CloudTaskService extends CloudTaskEngine { constructor( @inject(CLOUD_TASK_AUTH) - private readonly auth: ICloudTaskAuth, + auth: ICloudTaskAuth, @inject(ANALYTICS_SERVICE) - private readonly analytics: IAnalytics, + analytics: IAnalytics, @inject(ROOT_LOGGER) logger: RootLogger, @inject(MCP_RELAY_EXECUTOR) @optional() - private readonly mcpRelayExecutor: McpRelayExecutor | null = null, + mcpRelayExecutor: McpRelayExecutor | null = null, ) { - super(); - this.log = logger.scope("cloud-task"); - } - - /** - * Relay-designated server names per run (docs/cloud-mcp-relay.md). - * In-memory by design: only the client that created a run in this app - * session may execute relay requests for it; requests for undesignated - * runs or names are dropped. - */ - private readonly relayDesignations = new Map>(); - /** requestId dedupe — the event stream is at-least-once and replays on reconnect. */ - private readonly handledRelayRequestIds = new Set(); - private readonly handledRelayRequestOrder: string[] = []; - - /** Sensitive relay requests require desktop-owned approval. */ - private readonly relayAlwaysApprovals = new Set(); - /** Desktop-issued relay approval prompts awaiting a task-view answer. */ - private readonly pendingLocalRelayPrompts = new Map< - string, - { - runId: string; - resolve: (outcome: { - optionId: string | null; - customInput?: string; - }) => void; - } - >(); - - designateRelayedMcpServers(runId: string, servers: string[]): void { - if (servers.length === 0) return; - this.relayDesignations.set(runId, new Set(servers)); - this.log.info("Designated relayed MCP servers for run", { - runId, - servers, - }); - } - - private markRelayRequestHandled(requestId: string): void { - this.handledRelayRequestIds.add(requestId); - this.handledRelayRequestOrder.push(requestId); - if (this.handledRelayRequestOrder.length > MAX_HANDLED_RELAY_REQUEST_IDS) { - const evicted = this.handledRelayRequestOrder.shift(); - if (evicted) this.handledRelayRequestIds.delete(evicted); - } - } - - private async handleMcpRelayRequest( - watcher: WatcherState, - data: McpRequestEventData, - ): Promise { - if (!this.mcpRelayExecutor) return; - const designated = this.relayDesignations.get(watcher.runId); - if (!designated?.has(data.server)) { - // Not created by this client, or a name the run never declared. - return; - } - if (this.handledRelayRequestIds.has(data.requestId)) return; - this.markRelayRequestHandled(data.requestId); - - const expiresAt = Date.parse(data.expiresAt); - if (this.relayRequestExpired(expiresAt)) { - this.log.info("Dropping expired MCP relay request", { - runId: watcher.runId, - server: data.server, - requestId: data.requestId, - }); - return; - } - - const approvalRequest = relayApprovalRequest( - watcher.runId, - data.server, - data.payload, - ); - if (approvalRequest) { - const approval = await this.ensureRelayRequestApproval( - watcher, - approvalRequest, - expiresAt, - ); - if (!approval.approved) { - // Expired prompts get no response: the sandbox has already timed the - // request out, and a late mcp_response would be rejected as unknown. - if (!approval.expired) { - await this.sendRelayResponse(watcher, data, { - error: { code: -32000, message: approval.message }, - }); - } - return; - } - if (this.relayRequestExpired(expiresAt)) return; - } - - let execution: { - payload?: Record; - error?: { code: number; message: string }; - }; - try { - execution = await this.mcpRelayExecutor.execute( - watcher.runId, - data.server, - data.payload, - ); - } catch (error) { - execution = { - error: { - code: -32000, - message: - error instanceof Error - ? error.message - : "MCP relay execution failed", - }, - }; - } - - // Fire-and-forget notifications produce no response payload or error. - if (!execution.payload && !execution.error) return; - - await this.sendRelayResponse(watcher, data, execution); - } - - private relayRequestExpired(expiresAt: number): boolean { - return Number.isFinite(expiresAt) && expiresAt < Date.now(); - } - - private async sendRelayResponse( - watcher: WatcherState, - data: McpRequestEventData, - execution: { - payload?: Record; - error?: { code: number; message: string }; - }, - ): Promise { - try { - await this.sendCommand({ - taskId: watcher.taskId, - runId: watcher.runId, - apiHost: watcher.apiHost, - teamId: watcher.teamId, - method: "mcp_response", - params: { - requestId: data.requestId, - server: data.server, - ...(execution.payload - ? { payload: execution.payload } - : { error: execution.error }), - }, - }); - } catch (error) { - // The sandbox times the request out on its own; nothing to unwind here. - this.log.warn("Failed to deliver mcp_response command", { - runId: watcher.runId, - requestId: data.requestId, - error: serializeError(error), - }); - } - } - - private async ensureRelayRequestApproval( - watcher: WatcherState, - request: RelayApprovalRequest, - expiresAt: number, - ): Promise< - { approved: true } | { approved: false; expired: boolean; message: string } - > { - const { runId } = watcher; - if (this.relayAlwaysApprovals.has(request.approvalKey)) { - return { approved: true }; - } - - const requestId = `${RELAY_APPROVAL_REQUEST_PREFIX}${globalThis.crypto.randomUUID()}`; - this.emit(CloudTaskEvent.Update, { - taskId: watcher.taskId, - runId, - kind: "permission_request" as const, - requestId, - toolCall: { - toolCallId: requestId, - title: request.title, - kind: "other", - rawInput: request.rawInput, - _meta: posthogToolMeta({ - toolName: request.toolName, - mcp: request.mcp, - }), - }, - options: [ - { kind: "allow_once", name: "Yes", optionId: "allow" }, - { - kind: "allow_always", - name: "Yes, always allow", - optionId: "allow_always", - }, - { - kind: "reject_once", - name: "Type here to tell the agent what to do differently", - optionId: "reject", - _meta: { customInput: true }, - }, - ], - }); - - const outcome = await new Promise<{ - optionId: string | null; - customInput?: string; - }>((resolve) => { - this.pendingLocalRelayPrompts.set(requestId, { runId, resolve }); - // The sandbox abandons the request at expiresAt; keep waiting any longer - // and an approval would execute a call whose result nothing consumes. - const waitMs = Number.isFinite(expiresAt) - ? Math.max(0, expiresAt - Date.now()) - : 60_000; - const timer = setTimeout(() => { - if (this.pendingLocalRelayPrompts.delete(requestId)) { - resolve({ optionId: null }); - } - }, waitMs); - timer.unref?.(); - }); - - if (outcome.optionId === "allow_always") { - this.relayAlwaysApprovals.add(request.approvalKey); - return { approved: true }; - } - if (outcome.optionId === "allow") return { approved: true }; - if (outcome.optionId === null) { - return { - approved: false, - expired: true, - message: "The user did not respond in time.", - }; - } - return { - approved: false, - expired: false, - message: outcome.customInput - ? `The user denied this MCP request: ${outcome.customInput}` - : "The user denied this MCP request.", - }; - } - - /** Drop a terminal run's relay approval state and abandon its open prompts. */ - private evictRelayApprovalState(runId: string): void { - const prefix = `${runId}${RELAY_KEY_SEPARATOR}`; - for (const key of [...this.relayAlwaysApprovals]) { - if (key.startsWith(prefix)) this.relayAlwaysApprovals.delete(key); - } - for (const [requestId, prompt] of [...this.pendingLocalRelayPrompts]) { - if (prompt.runId !== runId) continue; - this.pendingLocalRelayPrompts.delete(requestId); - prompt.resolve({ optionId: null }); - } - } - - watch(input: WatchInput): void { - const key = watcherKey(input.taskId, input.runId); - - const existing = this.watchers.get(key); - if (existing) { - existing.subscriberCount++; - this.log.info("Cloud task watcher subscriber added", { - key, - subscribers: existing.subscriberCount, - }); - void this.emitCurrentSnapshot(key); - return; - } - - this.startWatcher(input, 1); - } - - unwatch(taskId: string, runId: string): void { - const key = watcherKey(taskId, runId); - const watcher = this.watchers.get(key); - if (!watcher) { - return; - } - - watcher.subscriberCount--; - if (watcher.subscriberCount <= 0) { - this.stopWatcher(key); - } else { - this.log.info("Cloud task watcher subscriber removed", { - key, - subscribers: watcher.subscriberCount, - }); - } - } - - async retry(taskId: string, runId: string): Promise { - const key = watcherKey(taskId, runId); - const watcher = this.watchers.get(key); - if (!watcher) return; - - if (watcher.reconnectTimeoutId) { - clearTimeout(watcher.reconnectTimeoutId); - watcher.reconnectTimeoutId = null; - } - - watcher.sseAbortController?.abort(); - watcher.sseAbortController = null; - - if (watcher.batchFlushTimeoutId) { - clearTimeout(watcher.batchFlushTimeoutId); - watcher.batchFlushTimeoutId = null; - } - - this.log.info("Retrying cloud task watcher", { - key, - hasSnapshot: watcher.hasEmittedSnapshot, - }); - - // Start over from scratch: a poisoned resume position loops straight back into the same - // failure, so re-bootstrap to re-resolve the read leg and emit a fresh snapshot. - this.resetWatcherForRebootstrap(watcher); - void this.bootstrapWatcher(key); - } - - // Resets a watcher to its pre-bootstrap state so bootstrapWatcher can rebuild it from server truth. - private resetWatcherForRebootstrap(watcher: WatcherState): void { - watcher.reconnectAttempts = 0; - watcher.streamErrorAttempts = 0; - watcher.cumulativeReconnectAttempts = 0; - watcher.failed = false; - watcher.pendingLogEntries = []; - watcher.bufferedLogBatches = []; - watcher.needsPostBootstrapReconnect = false; - watcher.needsStopAfterBootstrap = false; - watcher.streamEnded = false; - watcher.selfHealAttempted = false; - watcher.lastEventId = null; - watcher.lastEventIdLeg = null; - watcher.streamLeg = null; - // The rebuild re-resolves the read leg, so a retained id could false-match a - // different entry on the next connection — and the leg-switch clear in - // connectSse can't catch it, since lastEventId was just nulled. The re-fetched - // snapshot re-delivers history, so no dedup state is lost. - watcher.seenEventIds.clear(); - watcher.totalEntryCount = 0; - watcher.isBootstrapping = false; - watcher.streamTargetResolved = false; - watcher.streamBaseUrl = null; - watcher.streamReadToken = null; - watcher.durableStreamEnabled = false; - } - - async sendCommand(input: SendCommandInput): Promise { - if (input.method === "permission_response") { - const params = input.params ?? {}; - const requestId = - typeof params.requestId === "string" ? params.requestId : null; - if (requestId?.startsWith(RELAY_APPROVAL_REQUEST_PREFIX)) { - // A desktop-issued relay approval: resolve it locally — the sandbox - // never saw this prompt, so there is nothing to POST. - const pending = this.pendingLocalRelayPrompts.get(requestId); - this.pendingLocalRelayPrompts.delete(requestId); - pending?.resolve({ - optionId: - typeof params.optionId === "string" ? params.optionId : null, - customInput: - typeof params.customInput === "string" - ? params.customInput - : undefined, - }); - return { success: true }; - } - } - - const url = `${input.apiHost}/api/projects/${input.teamId}/tasks/${input.taskId}/runs/${input.runId}/command/`; - const body = { - jsonrpc: "2.0", - method: input.method, - params: input.params ?? {}, - id: `posthog-code-${Date.now()}`, - }; - - try { - const response = await this.auth.authenticatedFetch(url, { - method: "POST", - headers: { - "Content-Type": "application/json", - }, - body: JSON.stringify(body), - }); - - if (!response.ok) { - const errorText = await response.text().catch(() => ""); - let errorMessage = `Command failed with status ${response.status}`; - try { - const errorJson = JSON.parse(errorText); - if (errorJson.error?.message) { - errorMessage = errorJson.error.message; - } else if (errorJson.error) { - errorMessage = - typeof errorJson.error === "string" - ? errorJson.error - : JSON.stringify(errorJson.error); - } - } catch { - if (errorText) errorMessage = errorText; - } - - this.log.warn("Cloud task command failed", { - taskId: input.taskId, - runId: input.runId, - method: input.method, - status: response.status, - error: errorMessage, - }); - return { success: false, error: errorMessage }; - } - - const data = (await response.json()) as { - error?: { message?: string }; - result?: unknown; - }; - - if (data.error) { - this.log.warn("Cloud task command returned error", { - taskId: input.taskId, - method: input.method, - error: data.error, - }); - return { - success: false, - error: data.error.message ?? JSON.stringify(data.error), - }; - } - - this.log.info("Cloud task command sent", { - taskId: input.taskId, - runId: input.runId, - method: input.method, - }); - - return { success: true, result: data.result }; - } catch (error) { - const errorMessage = - error instanceof Error ? error.message : "Unknown error"; - this.log.error("Cloud task command error", { - taskId: input.taskId, - method: input.method, - error: errorMessage, - }); - return { success: false, error: errorMessage }; - } - } - - async stop(input: StopInput): Promise { - try { - const context = await this.auth.getCloudContext(); - if (!context) { - return { success: false, error: "No active cloud project" }; - } - const url = `${context.apiHost}/api/projects/${context.teamId}/tasks/${input.taskId}/runs/${input.runId}/cancel/`; - const response = await this.auth.authenticatedFetch(url, { - method: "POST", - headers: { - "Content-Type": "application/json", - }, - body: JSON.stringify(input.reason ? { reason: input.reason } : {}), - }); - - if (!response.ok) { - const errorText = await response.text().catch(() => ""); - let errorMessage = `Stop failed with status ${response.status}`; - try { - const errorJson = JSON.parse(errorText) as { error?: unknown }; - if (typeof errorJson.error === "string" && errorJson.error) { - errorMessage = errorJson.error; - } - } catch { - if (errorText) errorMessage = errorText; - } - - this.log.warn("Cloud run stop failed", { - taskId: input.taskId, - runId: input.runId, - status: response.status, - error: errorMessage, - }); - return { - success: false, - error: errorMessage, - retryable: response.status === 503 || response.status >= 500, - }; - } - - const data = (await response.json()) as { status?: string }; - this.log.info("Cloud run stop accepted", { - taskId: input.taskId, - runId: input.runId, - runStatus: data.status, - }); - return { success: true, runStatus: data.status }; - } catch (error) { - const errorMessage = - error instanceof Error ? error.message : "Unknown error"; - this.log.error("Cloud run stop error", { - taskId: input.taskId, - runId: input.runId, - error: errorMessage, - }); - return { success: false, error: errorMessage, retryable: true }; - } + super({ auth, analytics, logger, mcpRelayExecutor }); } @preDestroy() - unwatchAll(): void { - for (const key of [...this.watchers.keys()]) { - this.stopWatcher(key); - } - } - - private startWatcher(input: WatchInput, subscriberCount: number): void { - const key = watcherKey(input.taskId, input.runId); - - const watcher: WatcherState = { - taskId: input.taskId, - runId: input.runId, - apiHost: input.apiHost, - teamId: input.teamId, - subscriberCount, - sseAbortController: null, - reconnectTimeoutId: null, - batchFlushTimeoutId: null, - pendingLogEntries: [], - totalEntryCount: 0, - resumeFromEntryCount: input.resumeFromEntryCount ?? null, - reconnectAttempts: 0, - streamErrorAttempts: 0, - cumulativeReconnectAttempts: 0, - lastEventId: null, - lastEventIdLeg: null, - streamLeg: null, - seenEventIds: new Set(), - lastStatus: null, - lastStage: null, - lastOutput: null, - lastErrorMessage: null, - lastBranch: null, - lastSandboxAlive: null, - lastStatusUpdatedAt: null, - connStartedAt: 0, - connSentLastEventId: null, - connDataEventsReceived: 0, - isBootstrapping: false, - hasEmittedSnapshot: false, - bufferedLogBatches: [], - emittedLogEntries: [], - failed: false, - needsPostBootstrapReconnect: false, - needsStopAfterBootstrap: false, - streamEnded: false, - selfHealAttempted: false, - streamTargetResolved: false, - streamBaseUrl: null, - streamReadToken: null, - durableStreamEnabled: false, - }; - - this.watchers.set(key, watcher); - this.log.info("Cloud task watcher started", { key }); - void this.bootstrapWatcher(key); - } - - private stopWatcher(key: string): void { - const watcher = this.watchers.get(key); - if (!watcher) return; - - if (this.relayDesignations.has(watcher.runId)) { - // No watcher → no relay events → nothing executes; release the run's - // live server connections (stdio children included). They reopen - // lazily if the run is watched again. - void this.mcpRelayExecutor?.closeRun?.(watcher.runId).catch(() => {}); - } - - watcher.sseAbortController?.abort(); - - if (watcher.reconnectTimeoutId) { - clearTimeout(watcher.reconnectTimeoutId); - watcher.reconnectTimeoutId = null; - } - - if (watcher.batchFlushTimeoutId) { - clearTimeout(watcher.batchFlushTimeoutId); - watcher.batchFlushTimeoutId = null; - } - - this.flushLogBatch(key); - this.watchers.delete(key); - this.log.info("Cloud task watcher stopped", { key }); - } - - private async bootstrapWatcher(key: string): Promise { - const watcher = this.watchers.get(key); - if (!watcher) return; - - watcher.failed = false; - watcher.needsPostBootstrapReconnect = false; - watcher.needsStopAfterBootstrap = false; - - const run = await this.fetchTaskRun(watcher); - const currentWatcher = this.watchers.get(key); - if (!currentWatcher || currentWatcher !== watcher) return; - if (watcher.failed) return; - - if (!run) { - this.failWatcher(key, { - title: "Failed to load cloud run", - message: "Could not fetch the cloud run state. Retry to reconnect.", - retryable: true, - }); - return; - } - - this.applyTaskRunState(watcher, run); - - if ( - !isTerminalStatus(run.status) && - watcher.resumeFromEntryCount !== null - ) { - watcher.totalEntryCount = watcher.resumeFromEntryCount; - watcher.hasEmittedSnapshot = true; - watcher.isBootstrapping = false; - void this.connectSse(key, { startLatest: true }); - return; - } - - if (isTerminalStatus(run.status)) { - const historicalEntries = await this.fetchAllSessionLogs(watcher); - const terminalWatcher = this.watchers.get(key); - if (!terminalWatcher || terminalWatcher !== watcher) return; - if (watcher.failed) return; - if (!historicalEntries) { - this.failWatcher(key, { - title: "Failed to load task history", - message: - "Could not load the persisted cloud task logs. Retry to reconnect.", - retryable: true, - }); - return; - } - - watcher.totalEntryCount = historicalEntries.length; - watcher.hasEmittedSnapshot = true; - this.emit(CloudTaskEvent.Update, { - taskId: watcher.taskId, - runId: watcher.runId, - kind: "snapshot", - newEntries: historicalEntries, - totalEntryCount: watcher.totalEntryCount, - status: watcher.lastStatus ?? undefined, - stage: watcher.lastStage, - output: watcher.lastOutput, - errorMessage: watcher.lastErrorMessage, - branch: watcher.lastBranch, - ...sandboxAlivePayload(watcher), - }); - this.stopWatcher(key); - return; - } - - watcher.isBootstrapping = true; - watcher.bufferedLogBatches = []; - void this.connectSse(key, { startLatest: true }); - - const historicalEntries = await this.fetchAllSessionLogs(watcher); - const bootstrappingWatcher = this.watchers.get(key); - if (!bootstrappingWatcher || bootstrappingWatcher !== watcher) return; - if (watcher.failed) return; - if (!historicalEntries) { - this.failWatcher(key, { - title: "Failed to load cloud run history", - message: - "Could not load the existing cloud run logs. Retry to reconnect.", - retryable: true, - }); - return; - } - - // Flush any pending live entries into the bootstrap buffer before snapshot. - this.flushLogBatch(key); - - watcher.totalEntryCount = historicalEntries.length; - watcher.hasEmittedSnapshot = true; - - this.emit(CloudTaskEvent.Update, { - taskId: watcher.taskId, - runId: watcher.runId, - kind: "snapshot", - newEntries: historicalEntries, - totalEntryCount: watcher.totalEntryCount, - status: watcher.lastStatus ?? undefined, - stage: watcher.lastStage, - output: watcher.lastOutput, - errorMessage: watcher.lastErrorMessage, - branch: watcher.lastBranch, - ...sandboxAlivePayload(watcher), - }); - - watcher.isBootstrapping = false; - this.drainBufferedLogBatches(key, historicalEntries); - - if (watcher.failed) { - return; - } - - if (watcher.needsStopAfterBootstrap) { - watcher.needsStopAfterBootstrap = false; - await this.finalizeWatcherStop(key); - return; - } - - if (watcher.needsPostBootstrapReconnect) { - watcher.needsPostBootstrapReconnect = false; - this.scheduleReconnect(key, undefined, { countAttempt: false }); - } - - void this.verifyPostBootstrapStatus(key); - } - - private async verifyPostBootstrapStatus(key: string): Promise { - const watcher = this.watchers.get(key); - if (!watcher) return; - if (isTerminalStatus(watcher.lastStatus)) return; - - const run = await this.fetchTaskRun(watcher); - const currentWatcher = this.watchers.get(key); - if (!currentWatcher || currentWatcher !== watcher) return; - if (!run) return; - - if (!this.applyTaskRunState(watcher, run)) return; - if (isTerminalStatus(watcher.lastStatus)) return; - - this.emitStatusUpdate(watcher); - } - - private emitStatusUpdate(watcher: WatcherState): void { - this.emit(CloudTaskEvent.Update, { - taskId: watcher.taskId, - runId: watcher.runId, - kind: "status", - status: watcher.lastStatus ?? undefined, - stage: watcher.lastStage, - output: watcher.lastOutput, - errorMessage: watcher.lastErrorMessage, - branch: watcher.lastBranch, - ...sandboxAlivePayload(watcher), - }); - } - - private async connectSse( - key: string, - options?: { startLatest?: boolean }, - ): Promise { - const watcher = this.watchers.get(key); - if (!watcher) return; - - const controller = new AbortController(); - watcher.sseAbortController = controller; - - watcher.connStartedAt = 0; - watcher.connDataEventsReceived = 0; - - // Resolve the read target once (proxy URL + token, or Django), reused across reconnects. - if (!watcher.streamTargetResolved) { - await this.resolveStreamTarget(watcher); - const resolvedWatcher = this.watchers.get(key); - if ( - !resolvedWatcher || - resolvedWatcher !== watcher || - controller.signal.aborted - ) { - return; - } - } - - const usingProxy = Boolean( - watcher.streamBaseUrl && watcher.streamReadToken, - ); - const base = usingProxy - ? watcher.streamBaseUrl?.replace(/\/+$/, "") - : watcher.apiHost; - const leg: StreamLeg = usingProxy ? "proxy" : "django"; - // Proxy and Django id spaces are unrelated, so drop the resume position on a leg switch and - // let start=latest plus the next snapshot cover the gap. - if (watcher.lastEventId && watcher.lastEventIdLeg !== leg) { - this.log.info("Cloud task stream leg changed, dropping resume position", { - key, - from: watcher.lastEventIdLeg, - to: leg, - }); - watcher.lastEventId = null; - watcher.lastEventIdLeg = null; - // Proxy and Django ids are unrelated, so a retained id could false-match a - // different entry on the new leg. Drop them; the snapshot covers the gap. - watcher.seenEventIds.clear(); - } - watcher.streamLeg = leg; - - // Captured after the leg-switch drop so they reflect what this connection actually sends. - watcher.connSentLastEventId = watcher.lastEventId; - const startLatest = Boolean(options?.startLatest && !watcher.lastEventId); - const url = new URL( - usingProxy - ? `${base}/v1/runs/${encodeURIComponent(watcher.runId)}/stream` - : `${base}/api/projects/${watcher.teamId}/tasks/${encodeURIComponent( - watcher.taskId, - )}/runs/${encodeURIComponent(watcher.runId)}/stream/`, - ); - if (startLatest) { - url.searchParams.set("start", "latest"); - } - const headers: Record = { - Accept: "text/event-stream", - }; - if (watcher.lastEventId) { - headers["Last-Event-ID"] = watcher.lastEventId; - } - if (usingProxy) { - headers.Authorization = `Bearer ${watcher.streamReadToken}`; - } - - // Info so every stream attempt is visible in the logs; Bearer token redacted. - this.log.info(`Opening cloud task stream via ${leg}: ${url.toString()}`, { - key, - leg, - usingProxy, - durableStream: watcher.durableStreamEnabled, - method: "GET", - streamUrl: url.toString(), - lastEventId: watcher.lastEventId, - startLatest, - headers: usingProxy - ? { ...headers, Authorization: "Bearer " } - : headers, - }); - - const parser = new SseEventParser((message, data) => - this.log.warn(message, data), - ); - const decoder = new TextDecoder(); - - // Track how long the body stayed open so healthy long-lived connections cut by churn - // aren't penalized as failed reconnects (see SSE_HEALTHY_CONNECTION_MS). - let connectedAt = 0; - let streamWasEstablished = false; - let bytesReceived = 0; - let eventsReceived = 0; - - try { - // The proxy authenticates with the run-scoped Bearer token; the Django leg uses the session. - const response = usingProxy - ? await fetch(url.toString(), { - method: "GET", - headers, - signal: controller.signal, - }) - : await this.auth.authenticatedFetch(url.toString(), { - method: "GET", - headers, - signal: controller.signal, - }); - - this.log.info( - `Cloud task stream response ${response.status} ${ - response.ok ? "ok" : "FAILED" - } via ${leg}`, - { - key, - leg, - status: response.status, - ok: response.ok, - streamUrl: url.toString(), - }, - ); - - if (!response.ok) { - throw createStreamStatusError(response.status); - } - - if (!response.body) { - throw new Error("Stream response did not include a body"); - } - - connectedAt = Date.now(); - streamWasEstablished = true; - watcher.connStartedAt = connectedAt; - - this.log.info(`Cloud task SSE connected via ${leg}: ${url.toString()}`, { - key, - leg, - streamUrl: url.toString(), - sentLastEventId: watcher.connSentLastEventId, - startLatest, - status: response.status, - server: response.headers.get("server"), - via: response.headers.get("via"), - cfRay: response.headers.get("cf-ray"), - cfCacheStatus: response.headers.get("cf-cache-status"), - xAccelBuffering: response.headers.get("x-accel-buffering"), - contentType: response.headers.get("content-type"), - requestId: response.headers.get("x-request-id"), - }); - - const reader = response.body.getReader(); - - while (true) { - const { done, value } = await reader.read(); - if (done) { - break; - } - - if (!value) { - continue; - } - - bytesReceived += value.byteLength; - const chunk = decoder.decode(value, { stream: true }); - const events = parser.parse(chunk); - for (const event of events) { - eventsReceived += 1; - const backendError = this.handleSseEvent(key, event); - if (backendError) { - throw backendError; - } - } - } - - const trailingEvents = parser.parse(decoder.decode()); - for (const event of trailingEvents) { - const backendError = this.handleSseEvent(key, event); - if (backendError) { - throw backendError; - } - } - - this.flushLogBatch(key); - - if (controller.signal.aborted) { - return; - } - - this.log.info("Cloud task stream closed cleanly", { - key, - connectionDurationMs: Date.now() - connectedAt, - bytesReceived, - eventsReceived, - dataEventsReceived: watcher.connDataEventsReceived, - lastEventId: watcher.lastEventId, - }); - - // A long-lived clean close is healthy churn, not a loop: clear the cumulative budget so an - // idle run can ride out proxy timeout cycles, while instant-EOF loops still exhaust it. - const completedWatcher = this.watchers.get(key); - if ( - completedWatcher && - streamWasEstablished && - Date.now() - connectedAt >= SSE_HEALTHY_CONNECTION_MS - ) { - completedWatcher.cumulativeReconnectAttempts = 0; - completedWatcher.selfHealAttempted = false; - } - - await this.handleStreamCompletion(key, { reconnectOnDisconnect: true }); - } catch (error) { - this.flushLogBatch(key); - - if (controller.signal.aborted) { - return; - } - - // Proxy-leg 401: the read token expired or its signing key rotated. Re-resolve to mint a - // fresh token (or route back to Django) instead of failing. Django-leg 401 stays fatal below. - const unauthorizedWatcher = this.watchers.get(key); - if ( - error instanceof CloudTaskStreamError && - error.status === 401 && - unauthorizedWatcher?.streamBaseUrl - ) { - // Keep durableStreamEnabled set: clearing it would route this disconnect through legacy - // status polling, which can stop the watch on a terminal status before stream-end arrives. - // The next connectSse re-resolves the target and resolveStreamTarget re-derives durability. - unauthorizedWatcher.streamTargetResolved = false; - unauthorizedWatcher.streamBaseUrl = null; - unauthorizedWatcher.streamReadToken = null; - this.log.info("Cloud task stream proxy token rejected, re-resolving", { - key, - }); - await this.handleStreamCompletion(key, { - reconnectOnDisconnect: true, - reconnectError: error, - countReconnectAttempt: true, - }); - return; - } - - if ( - error instanceof CloudTaskStreamError && - error.details.autoRetry === false - ) { - this.failWatcher(key, error.details); - return; - } - - const errorMessage = - error instanceof Error ? error.message : "Unknown stream error"; - - const isBackendError = error instanceof BackendStreamError; - const wasHealthyStream = - !isBackendError && - streamWasEstablished && - Date.now() - connectedAt >= SSE_HEALTHY_CONNECTION_MS; - - const errorWatcher = this.watchers.get(key); - if (errorWatcher) { - if (isBackendError) { - errorWatcher.streamErrorAttempts += 1; - } else if (wasHealthyStream) { - errorWatcher.streamErrorAttempts = 0; - // A healthy-length connection proves timeout cycling, not a loop. - errorWatcher.cumulativeReconnectAttempts = 0; - errorWatcher.selfHealAttempted = false; - } - } - - this.log.warn("Cloud task stream error", { - key, - leg, - streamUrl: url.toString(), - error: errorMessage, - errorDetail: serializeError(error), - wasHealthyStream, - isBackendError, - streamWasEstablished, - connectionDurationMs: streamWasEstablished - ? Date.now() - connectedAt - : 0, - bytesReceived, - eventsReceived, - dataEventsReceived: errorWatcher?.connDataEventsReceived ?? 0, - lastEventId: errorWatcher?.lastEventId ?? null, - reconnectAttempts: errorWatcher?.reconnectAttempts ?? 0, - streamErrorAttempts: errorWatcher?.streamErrorAttempts ?? 0, - cumulativeReconnectAttempts: - errorWatcher?.cumulativeReconnectAttempts ?? 0, - }); - await this.handleStreamCompletion(key, { - reconnectOnDisconnect: true, - reconnectError: error, - countReconnectAttempt: !isBackendError && !wasHealthyStream, - }); - } finally { - const currentWatcher = this.watchers.get(key); - if (currentWatcher?.sseAbortController === controller) { - currentWatcher.sseAbortController = null; - } - } - } - - // Returns a BackendStreamError when the stream carries an error event so the caller can throw at - // the read site; returns null otherwise. It does not throw, so a single event cannot unwind the - // reader loop unexpectedly. - private handleSseEvent( - key: string, - event: SseEvent, - ): BackendStreamError | null { - const watcher = this.watchers.get(key); - if (!watcher || watcher.failed) return null; - - if (event.id) { - watcher.lastEventId = event.id; - watcher.lastEventIdLeg = watcher.streamLeg; - } - - if (event.event === "error") { - const message = isSseErrorEvent(event.data) - ? event.data.error - : "Unknown stream error"; - return new BackendStreamError(message); - } - - if (event.event === STREAM_END_EVENT_NAME) { - // The run's stream is durably complete. Mark it so completion stops instead - // of reconnecting, independent of run status. The connection will close - // naturally (clean EOF) right after this sentinel. - watcher.streamEnded = true; - return null; - } - - // A keepalive or real event proves the transport recovered. A keepalive does not clear the - // backend-error budget, which only a real data event below resets. - watcher.reconnectAttempts = 0; - - if (isKeepaliveEvent(event)) { - return null; - } - - // A real data event proves the stream materialized; clear the remaining budgets and re-arm self-heal. - watcher.streamErrorAttempts = 0; - watcher.cumulativeReconnectAttempts = 0; - watcher.selfHealAttempted = false; - - watcher.connDataEventsReceived += 1; - if (watcher.connDataEventsReceived === 1 && watcher.connSentLastEventId) { - this.log.info("Cloud task SSE resumed", { - key, - resumedFrom: watcher.connSentLastEventId, - firstEventIdAfterResume: event.id ?? null, - }); - } - - if (isTaskRunStateEvent(event.data)) { - if (this.applyTaskRunState(watcher, event.data)) { - if (!watcher.isBootstrapping && !isTerminalStatus(watcher.lastStatus)) { - this.emit(CloudTaskEvent.Update, { - taskId: watcher.taskId, - runId: watcher.runId, - kind: "status", - status: watcher.lastStatus ?? undefined, - stage: watcher.lastStage, - output: watcher.lastOutput, - errorMessage: watcher.lastErrorMessage, - branch: watcher.lastBranch, - ...sandboxAlivePayload(watcher), - }); - } - } - return null; - } - - // Drop a re-delivered event by its stream id. The durable stream re-sends - // the tail on reconnect/replay: each re-sent log entry would otherwise be - // counted as a new entry (advancing totalEntryCount past the renderer's - // processedLineCount guard) and emitted again — the root cause of duplicate - // transcript entries and back-to-back completion notifications — and a - // re-sent permission_request frame would re-surface an already-answered - // question as a fresh pending card. Events without an id (legacy servers) - // fall through and are handled downstream. - const eventId = event.id; - if (eventId !== undefined) { - if (watcher.seenEventIds.has(eventId)) { - return null; - } - watcher.seenEventIds.add(eventId); - } - - if (isMcpRequestEvent(event.data)) { - void this.handleMcpRelayRequest(watcher, event.data); - return null; - } - - if (isPermissionRequestEvent(event.data)) { - this.emit(CloudTaskEvent.Update, { - taskId: watcher.taskId, - runId: watcher.runId, - kind: "permission_request" as const, - requestId: event.data.requestId, - toolCall: event.data.toolCall, - options: event.data.options, - }); - return null; - } - - watcher.pendingLogEntries.push(event.data as StoredLogEntry); - if (watcher.pendingLogEntries.length >= EVENT_BATCH_MAX_SIZE) { - this.flushLogBatch(key); - return null; - } - - if (!watcher.batchFlushTimeoutId) { - watcher.batchFlushTimeoutId = setTimeout(() => { - watcher.batchFlushTimeoutId = null; - this.flushLogBatch(key); - }, EVENT_BATCH_FLUSH_MS); - } - - return null; - } - - private flushLogBatch(key: string): void { - const watcher = this.watchers.get(key); - if (!watcher || watcher.pendingLogEntries.length === 0) return; - - if (watcher.batchFlushTimeoutId) { - clearTimeout(watcher.batchFlushTimeoutId); - watcher.batchFlushTimeoutId = null; - } - - const entries = watcher.pendingLogEntries; - watcher.pendingLogEntries = []; - - if (watcher.isBootstrapping) { - watcher.bufferedLogBatches.push(entries); - return; - } - - watcher.totalEntryCount += entries.length; - this.rememberEmittedLogEntries(watcher, entries); - - this.emit(CloudTaskEvent.Update, { - taskId: watcher.taskId, - runId: watcher.runId, - kind: "logs", - newEntries: entries, - totalEntryCount: watcher.totalEntryCount, - }); - } - - private drainBufferedLogBatches( - key: string, - historicalEntries: StoredLogEntry[], - ): void { - const watcher = this.watchers.get(key); - if (!watcher || watcher.bufferedLogBatches.length === 0) return; - - const historicalCounts = buildEntryFrequencyMap(historicalEntries); - - for (const entries of watcher.bufferedLogBatches) { - const dedupedEntries = filterEntriesNotInFrequencyMap( - entries, - historicalCounts, - ); - - if (dedupedEntries.length === 0) { - continue; - } - - watcher.totalEntryCount += dedupedEntries.length; - this.rememberEmittedLogEntries(watcher, dedupedEntries); - this.emit(CloudTaskEvent.Update, { - taskId: watcher.taskId, - runId: watcher.runId, - kind: "logs", - newEntries: dedupedEntries, - totalEntryCount: watcher.totalEntryCount, - }); - } - - watcher.bufferedLogBatches = []; - } - - private rememberEmittedLogEntries( - watcher: WatcherState, - entries: StoredLogEntry[], - ): void { - watcher.emittedLogEntries.push(...entries); - } - - private mergeHistoricalAndEmittedEntries( - historicalEntries: StoredLogEntry[], - emittedEntries: StoredLogEntry[], - ): { - snapshotEntries: StoredLogEntry[]; - missingEmittedEntries: StoredLogEntry[]; - } { - if (emittedEntries.length === 0) { - return { snapshotEntries: historicalEntries, missingEmittedEntries: [] }; - } - - const historicalCounts = buildEntryFrequencyMap(historicalEntries); - const missingEmittedEntries = filterEntriesNotInFrequencyMap( - emittedEntries, - historicalCounts, - ); - - return { - snapshotEntries: [...historicalEntries, ...missingEmittedEntries], - missingEmittedEntries, - }; - } - - private async emitCurrentSnapshot(key: string): Promise { - const watcher = this.watchers.get(key); - if (!watcher || watcher.failed) return; - - const historicalEntries = await this.fetchAllSessionLogs(watcher); - const currentWatcher = this.watchers.get(key); - if (!currentWatcher || currentWatcher !== watcher || watcher.failed) { - return; - } - - if (!historicalEntries) { - this.log.warn("Cloud task snapshot replay failed", { - taskId: watcher.taskId, - runId: watcher.runId, - }); - return; - } - - const { snapshotEntries, missingEmittedEntries } = - this.mergeHistoricalAndEmittedEntries( - historicalEntries, - watcher.emittedLogEntries, - ); - watcher.emittedLogEntries = missingEmittedEntries; - if (snapshotEntries.length > watcher.totalEntryCount) { - watcher.totalEntryCount = snapshotEntries.length; - } - - this.emit(CloudTaskEvent.Update, { - taskId: watcher.taskId, - runId: watcher.runId, - kind: "snapshot", - newEntries: snapshotEntries, - totalEntryCount: snapshotEntries.length, - status: watcher.lastStatus ?? undefined, - stage: watcher.lastStage, - output: watcher.lastOutput, - errorMessage: watcher.lastErrorMessage, - branch: watcher.lastBranch, - ...sandboxAlivePayload(watcher), - }); - } - - private failWatcher(key: string, error: CloudTaskConnectionError): void { - const watcher = this.watchers.get(key); - if (!watcher) return; - - this.log.warn("Cloud task watcher failed", { - key, - errorTitle: error.title, - retryable: error.retryable, - status: watcher.lastStatus, - wasBootstrapping: watcher.isBootstrapping, - reconnectAttempts: watcher.reconnectAttempts, - cumulativeReconnectAttempts: watcher.cumulativeReconnectAttempts, - totalEntryCount: watcher.totalEntryCount, - lastEventId: watcher.lastEventId, - }); - - this.analytics.track(ANALYTICS_EVENTS.CLOUD_STREAM_DISCONNECTED, { - task_id: watcher.taskId, - run_id: watcher.runId, - team_id: watcher.teamId, - error_title: error.title, - retryable: error.retryable, - reconnect_attempts: watcher.reconnectAttempts, - stream_error_attempts: watcher.streamErrorAttempts, - cumulative_reconnect_attempts: watcher.cumulativeReconnectAttempts, - was_bootstrapping: watcher.isBootstrapping, - }); - - watcher.failed = true; - watcher.isBootstrapping = false; - watcher.pendingLogEntries = []; - watcher.bufferedLogBatches = []; - - if (watcher.reconnectTimeoutId) { - clearTimeout(watcher.reconnectTimeoutId); - watcher.reconnectTimeoutId = null; - } - - if (watcher.batchFlushTimeoutId) { - clearTimeout(watcher.batchFlushTimeoutId); - watcher.batchFlushTimeoutId = null; - } - - watcher.sseAbortController?.abort(); - watcher.sseAbortController = null; - - this.emit(CloudTaskEvent.Update, { - taskId: watcher.taskId, - runId: watcher.runId, - kind: "error", - errorTitle: error.title, - errorMessage: error.message, - retryable: error.retryable, - }); - } - - private scheduleReconnect( - key: string, - error?: unknown, - options: { countAttempt?: boolean } = {}, - ): void { - const watcher = this.watchers.get(key); - // Status-unaware: the loop only stops on the stream-end sentinel or budget exhaustion below. - if (!watcher || watcher.failed) { - return; - } - - if (watcher.reconnectTimeoutId) { - clearTimeout(watcher.reconnectTimeoutId); - } - - // Bounds runaway loops that clean-EOF (countAttempt=false) and dodge reconnectAttempts. - watcher.cumulativeReconnectAttempts += 1; - const countAttempt = options.countAttempt ?? true; - if (countAttempt) { - watcher.reconnectAttempts += 1; - } - - if ( - watcher.cumulativeReconnectAttempts > MAX_CUMULATIVE_RECONNECT_ATTEMPTS - ) { - // A poisoned resume position burns the budget without an error frame. Rebuild once from - // scratch (the app-restart recovery) before failing; if it loops straight back, fail for real. - if (!watcher.selfHealAttempted) { - watcher.reconnectTimeoutId = null; - this.log.warn( - "Cloud task stream looping without events, re-bootstrapping", - { key }, - ); - this.resetWatcherForRebootstrap(watcher); - // Set after the reset (which clears it): consumes the single allowed self-heal so a - // straight-back loop fails next time instead of re-bootstrapping forever. - watcher.selfHealAttempted = true; - void this.bootstrapWatcher(key); - return; - } - this.failWatcher(key, { - title: "Cloud run unreachable", - message: - "Could not maintain a connection to the cloud run after many attempts. Click retry once the issue is resolved.", - retryable: true, - }); - return; - } - - // Fail once either budget (transport reconnect or backend stream-error) is exhausted. - const attemptCount = Math.max( - watcher.reconnectAttempts, - watcher.streamErrorAttempts, - ); - if (attemptCount > MAX_SSE_RECONNECT_ATTEMPTS) { - const details = - error instanceof CloudTaskStreamError - ? error.details - : { - title: "Cloud stream disconnected", - message: - "Lost connection to the cloud run stream. Retry to reconnect.", - retryable: true, - }; - this.failWatcher(key, details); - return; - } - - const backoffAttempts = - error instanceof BackendStreamError - ? watcher.streamErrorAttempts - : watcher.reconnectAttempts; - const delay = Math.min( - SSE_RECONNECT_BASE_DELAY_MS * - 2 ** Math.max(backoffAttempts - SSE_RECONNECT_FLAT_ATTEMPTS, 0), - SSE_RECONNECT_MAX_DELAY_MS, - ); - - watcher.reconnectTimeoutId = setTimeout(() => { - const currentWatcher = this.watchers.get(key); - if (!currentWatcher) return; - currentWatcher.reconnectTimeoutId = null; - void this.connectSse(key, { - startLatest: - currentWatcher.isBootstrapping || currentWatcher.hasEmittedSnapshot, - }); - }, delay); - } - - private async handleStreamCompletion( - key: string, - options: { - reconnectOnDisconnect: boolean; - reconnectError?: unknown; - countReconnectAttempt?: boolean; - }, - ): Promise { - const watcher = this.watchers.get(key); - if (!watcher) return; - if (watcher.failed) return; - - const { reconnectOnDisconnect } = options; - - // Bootstrap owns the snapshot lifecycle: stopping mid-bootstrap would discard the backlog and - // buffered live entries. Record intent and let bootstrap finish. - if (watcher.isBootstrapping) { - if (watcher.streamEnded || !reconnectOnDisconnect) { - watcher.needsStopAfterBootstrap = true; - } else { - watcher.needsPostBootstrapReconnect = true; - } - return; - } - - // The stream-end sentinel is the only signal that ends a durable watch. Any disconnect without - // it is transport churn to reconnect through; status is tracked for display only, never to stop. - if (watcher.streamEnded) { - await this.finalizeWatcherStop(key); - return; - } - - // Legacy mode (old server): no sentinel, so poll run status on disconnect to decide stop vs - // reconnect. The reconnect budgets keep the new semantics, so self-heal stays active here too. - if (!watcher.durableStreamEnabled && reconnectOnDisconnect) { - const run = await this.fetchTaskRun(watcher); - const legacyWatcher = this.watchers.get(key); - if (!legacyWatcher || legacyWatcher !== watcher) return; - if (watcher.failed) return; - - if (run) { - this.applyTaskRunState(watcher, run); - } - if (isTerminalStatus(watcher.lastStatus)) { - this.emitStatusUpdate(watcher); - this.stopWatcher(key); - return; - } - if (run) { - this.emitStatusUpdate(watcher); - } - this.scheduleReconnect(key, options.reconnectError, { - countAttempt: options.countReconnectAttempt ?? false, - }); - return; - } - - // All callers pass reconnectOnDisconnect, and durable watches only stop via the stream-end - // sentinel or a terminal legacy poll (both handled above); any other disconnect reconnects. - if (reconnectOnDisconnect) { - this.scheduleReconnect(key, options.reconnectError, { - countAttempt: options.countReconnectAttempt ?? false, - }); - } - } - - // Stops a watcher whose stream is durably complete. Repairs the displayed status if the stream - // ended non-terminal (dropped final frame); the poll never decides whether to stop. - private async finalizeWatcherStop(key: string): Promise { - const watcher = this.watchers.get(key); - if (!watcher) return; - - if (!isTerminalStatus(watcher.lastStatus)) { - const run = await this.fetchTaskRun(watcher); - const currentWatcher = this.watchers.get(key); - if (!currentWatcher || currentWatcher !== watcher) return; - if (run) { - this.applyTaskRunState(watcher, run); - } - } - - this.emitStatusUpdate(watcher); - this.stopWatcher(key); - } - - private applyTaskRunState( - watcher: WatcherState, - run: - | Pick< - TaskRunResponse, - | "status" - | "stage" - | "output" - | "state" - | "error_message" - | "branch" - | "updated_at" - > - | TaskRunStateEvent, - ): boolean { - const updatedAt = run.updated_at ?? null; - if ( - updatedAt && - watcher.lastStatusUpdatedAt && - Date.parse(updatedAt) <= Date.parse(watcher.lastStatusUpdatedAt) - ) { - return false; - } - - const nextStatus = run.status ?? watcher.lastStatus; - const nextStage = run.stage ?? null; - const nextOutput = run.output ?? null; - const nextErrorMessage = run.error_message ?? null; - const nextBranch = run.branch ?? null; - const sandboxAlive = extractSandboxAlive(run.state); - const nextSandboxAlive = - sandboxAlive === undefined ? watcher.lastSandboxAlive : sandboxAlive; - - const changed = - nextStatus !== watcher.lastStatus || - nextStage !== watcher.lastStage || - JSON.stringify(nextOutput) !== JSON.stringify(watcher.lastOutput) || - nextErrorMessage !== watcher.lastErrorMessage || - nextBranch !== watcher.lastBranch || - nextSandboxAlive !== watcher.lastSandboxAlive; - - watcher.lastStatus = nextStatus ?? null; - watcher.lastStage = nextStage; - watcher.lastOutput = nextOutput; - watcher.lastErrorMessage = nextErrorMessage; - watcher.lastBranch = nextBranch; - watcher.lastSandboxAlive = nextSandboxAlive; - if (updatedAt) { - watcher.lastStatusUpdatedAt = updatedAt; - } - - // A terminal run gets no further relay requests; drop its designation and - // approval state so the maps don't grow for the lifetime of the app session. - if (isTerminalStatus(watcher.lastStatus)) { - this.relayDesignations.delete(watcher.runId); - this.evictRelayApprovalState(watcher.runId); - } - - return changed; - } - - private async fetchSessionLogsPage( - watcher: WatcherState, - offset: number, - ): Promise { - const url = new URL( - `${watcher.apiHost}/api/projects/${watcher.teamId}/tasks/${watcher.taskId}/runs/${watcher.runId}/session_logs/`, - ); - url.searchParams.set("limit", SESSION_LOG_PAGE_LIMIT.toString()); - url.searchParams.set("offset", offset.toString()); - - try { - const authedResponse = await this.auth.authenticatedFetch( - url.toString(), - { - method: "GET", - }, - ); - - if (!authedResponse.ok) { - this.log.warn("Cloud task session logs fetch failed", { - status: authedResponse.status, - taskId: watcher.taskId, - runId: watcher.runId, - offset, - }); - if (shouldFailWatcherForFetchStatus(authedResponse.status)) { - this.failWatcher( - watcherKey(watcher.taskId, watcher.runId), - createStreamStatusError(authedResponse.status).details, - ); - } - return null; - } - - const raw = await authedResponse.text(); - return { - entries: JSON.parse(raw) as StoredLogEntry[], - hasMore: authedResponse.headers.get("X-Has-More") === "true", - }; - } catch (error) { - this.log.warn("Cloud task session logs fetch error", { - taskId: watcher.taskId, - runId: watcher.runId, - offset, - error, - }); - return null; - } - } - - private async fetchAllSessionLogs( - watcher: WatcherState, - ): Promise { - const entries: StoredLogEntry[] = []; - let offset = 0; - - while (true) { - const page = await this.fetchSessionLogsPage(watcher, offset); - if (!page) { - return null; - } - - entries.push(...page.entries); - if (!page.hasMore || page.entries.length === 0) { - return entries; - } - - offset += page.entries.length; - } - } - - private async resolveStreamTarget(watcher: WatcherState): Promise { - const url = `${watcher.apiHost}/api/projects/${watcher.teamId}/tasks/${watcher.taskId}/runs/${watcher.runId}/stream_token/`; - try { - const response = await this.auth.authenticatedFetch(url, { - method: "GET", - }); - if (!response.ok) { - watcher.streamBaseUrl = null; - watcher.streamReadToken = null; - if (isTransientStreamTargetStatus(response.status)) { - // Transient: read from Django this round but leave the target unresolved so the next - // reconnect retries durable resolution instead of pinning the run to status polling. - this.log.warn("Cloud task stream target temporarily unavailable", { - taskId: watcher.taskId, - runId: watcher.runId, - status: response.status, - }); - return; - } - // Refused, or an old server without the endpoint: read from Django with status polling. - watcher.durableStreamEnabled = false; - watcher.streamTargetResolved = true; - this.log.info("Cloud task stream reading from API host", { - taskId: watcher.taskId, - runId: watcher.runId, - status: response.status, - }); - return; - } - const data = (await response.json()) as { - token?: string; - stream_base_url?: string | null; - }; - watcher.streamReadToken = data.token ?? null; - watcher.streamBaseUrl = data.stream_base_url ?? null; - // The endpoint resolving at all opts this watcher into the status-unaware contract; - // old servers 404 above and stay on legacy status polling. - watcher.durableStreamEnabled = true; - watcher.streamTargetResolved = true; - this.log.info("Cloud task stream target resolved", { - taskId: watcher.taskId, - runId: watcher.runId, - streamBaseUrl: watcher.streamBaseUrl, - hasToken: Boolean(watcher.streamReadToken), - durableStream: watcher.durableStreamEnabled, - }); - } catch (error) { - // Transient failure: leave unresolved so the next reconnect retries and falls back to Django. - watcher.streamBaseUrl = null; - watcher.streamReadToken = null; - this.log.warn("Cloud task stream target resolution failed", { - taskId: watcher.taskId, - runId: watcher.runId, - error, - }); - } - } - - private async fetchTaskRun( - watcher: WatcherState, - ): Promise { - const url = `${watcher.apiHost}/api/projects/${watcher.teamId}/tasks/${watcher.taskId}/runs/${watcher.runId}/`; - - try { - const authedResponse = await this.auth.authenticatedFetch(url, { - method: "GET", - }); - - if (!authedResponse.ok) { - this.log.warn("Cloud task status fetch failed", { - status: authedResponse.status, - taskId: watcher.taskId, - runId: watcher.runId, - }); - if (shouldFailWatcherForFetchStatus(authedResponse.status)) { - this.failWatcher( - watcherKey(watcher.taskId, watcher.runId), - createStreamStatusError(authedResponse.status).details, - ); - } - return null; - } - - return (await authedResponse.json()) as TaskRunResponse; - } catch (error) { - this.log.warn("Cloud task status fetch error", { - taskId: watcher.taskId, - runId: watcher.runId, - error, - }); - return null; - } + override unwatchAll(): void { + super.unwatchAll(); } } diff --git a/packages/core/src/cloud-task/schemas.ts b/packages/core/src/cloud-task/schemas.ts index d694e52141..b8c03eb202 100644 --- a/packages/core/src/cloud-task/schemas.ts +++ b/packages/core/src/cloud-task/schemas.ts @@ -1,20 +1,12 @@ -import type { TaskRunStatus } from "@posthog/shared"; +import type { CloudTaskUpdatePayload } from "@posthog/shared"; import { z } from "zod"; -import type { CloudTaskUpdatePayload } from "./cloud-task-types"; -export type { CloudTaskUpdatePayload, TaskRunStatus }; - -export const TERMINAL_STATUSES = ["completed", "failed", "cancelled"] as const; - -export function isTerminalStatus( - status: TaskRunStatus | string | null | undefined, -): boolean { - return ( - status !== null && - status !== undefined && - TERMINAL_STATUSES.includes(status as (typeof TERMINAL_STATUSES)[number]) - ); -} +export { + type CloudTaskUpdatePayload, + isTerminalStatus, + type TaskRunStatus, + TERMINAL_STATUSES, +} from "@posthog/shared"; // --- Events --- diff --git a/packages/core/src/sessions/cloudTaskCommandController.test.ts b/packages/core/src/sessions/cloudTaskCommandController.test.ts new file mode 100644 index 0000000000..a800a43747 --- /dev/null +++ b/packages/core/src/sessions/cloudTaskCommandController.test.ts @@ -0,0 +1,126 @@ +import { describe, expect, it, vi } from "vitest"; +import { + CloudTaskCommandController, + type CloudTaskCommandTransport, +} from "./cloudTaskCommandController"; + +function createHarness(): { + controller: CloudTaskCommandController; + transport: CloudTaskCommandTransport; +} { + const transport: CloudTaskCommandTransport = { + sendCommand: vi.fn(async () => {}), + stopRun: vi.fn(async () => {}), + }; + return { + controller: new CloudTaskCommandController(transport), + transport, + }; +} + +describe("CloudTaskCommandController", () => { + it("dispatches user messages with canonical parameters", async () => { + const { controller, transport } = createHarness(); + + await controller.sendUserMessage( + { taskId: "task-1", taskRunId: "run-1" }, + "encoded prompt", + ); + + expect(transport.sendCommand).toHaveBeenCalledWith( + { taskId: "task-1", taskRunId: "run-1" }, + "user_message", + { content: "encoded prompt" }, + ); + }); + + it("dispatches permission responses with canonical parameters", async () => { + const { controller, transport } = createHarness(); + + await controller.respondToPermission( + { taskId: "task-1", taskRunId: "run-1" }, + { + requestId: "request-1", + optionId: "allow", + answers: { reason: "approved" }, + customInput: "Proceed", + }, + ); + + expect(transport.sendCommand).toHaveBeenCalledWith( + { taskId: "task-1", taskRunId: "run-1" }, + "permission_response", + { + requestId: "request-1", + optionId: "allow", + answers: { reason: "approved" }, + customInput: "Proceed", + }, + ); + }); + + it("dispatches prompt cancellation to the active cloud run", async () => { + const { controller, transport } = createHarness(); + + await controller.cancelPrompt({ taskId: "task-1", taskRunId: "run-1" }); + + expect(transport.sendCommand).toHaveBeenCalledWith( + { taskId: "task-1", taskRunId: "run-1" }, + "cancel", + ); + }); + + it("dispatches configuration changes with canonical parameters", async () => { + const { controller, transport } = createHarness(); + + await controller.setConfigOption( + { taskId: "task-1", taskRunId: "run-1" }, + "model", + "claude-sonnet-4-5", + ); + + expect(transport.sendCommand).toHaveBeenCalledWith( + { taskId: "task-1", taskRunId: "run-1" }, + "set_config_option", + { configId: "model", value: "claude-sonnet-4-5" }, + ); + }); + + it("preserves transport failures for service-specific handling", async () => { + const { controller, transport } = createHarness(); + vi.mocked(transport.sendCommand).mockRejectedValueOnce(new Error("failed")); + + await expect( + controller.cancelPrompt({ taskId: "task-1", taskRunId: "run-1" }), + ).rejects.toThrow("failed"); + }); + + it("preserves command results for callers", async () => { + const commandResult = { success: true, result: { accepted: true } }; + const transport: CloudTaskCommandTransport = { + sendCommand: vi.fn(async () => commandResult), + stopRun: vi.fn(async () => undefined), + }; + const controller = new CloudTaskCommandController(transport); + + await expect( + controller.sendUserMessage( + { taskId: "task-1", taskRunId: "run-1" }, + "hello", + ), + ).resolves.toBe(commandResult); + }); + + it("dispatches stop requests and preserves their results", async () => { + const stopResult = { status: "cancelling" }; + const transport: CloudTaskCommandTransport = { + sendCommand: vi.fn(async () => {}), + stopRun: vi.fn(async () => stopResult), + }; + const controller = new CloudTaskCommandController(transport); + const target = { taskId: "task-1", taskRunId: "run-1" }; + + await expect(controller.stopRun(target)).resolves.toBe(stopResult); + expect(transport.stopRun).toHaveBeenCalledWith(target); + }); +}); diff --git a/packages/core/src/sessions/cloudTaskCommandController.ts b/packages/core/src/sessions/cloudTaskCommandController.ts new file mode 100644 index 0000000000..bb1e2cf004 --- /dev/null +++ b/packages/core/src/sessions/cloudTaskCommandController.ts @@ -0,0 +1,80 @@ +export interface CloudTaskCommandTarget { + taskId: string; + taskRunId: string; +} + +export type CloudTaskCommandMethod = + | "user_message" + | "permission_response" + | "cancel" + | "set_config_option"; + +export interface CloudTaskPermissionResponseCommand { + requestId: string; + optionId: string; + answers?: Record; + customInput?: string; +} + +export interface CloudTaskCommandTransport< + CommandResult = unknown, + StopResult = unknown, +> { + sendCommand( + target: CloudTaskCommandTarget, + method: CloudTaskCommandMethod, + params?: Record, + ): Promise; + stopRun(target: CloudTaskCommandTarget): Promise; +} + +export class CloudTaskCommandController< + CommandResult = unknown, + StopResult = unknown, +> { + constructor( + private readonly transport: CloudTaskCommandTransport< + CommandResult, + StopResult + >, + ) {} + + sendUserMessage( + target: CloudTaskCommandTarget, + command: string | Record, + ): Promise { + return this.transport.sendCommand( + target, + "user_message", + typeof command === "string" ? { content: command } : command, + ); + } + + respondToPermission( + target: CloudTaskCommandTarget, + response: CloudTaskPermissionResponseCommand, + ): Promise { + return this.transport.sendCommand(target, "permission_response", { + ...response, + }); + } + + cancelPrompt(target: CloudTaskCommandTarget): Promise { + return this.transport.sendCommand(target, "cancel"); + } + + setConfigOption( + target: CloudTaskCommandTarget, + configId: string, + value: string, + ): Promise { + return this.transport.sendCommand(target, "set_config_option", { + configId, + value, + }); + } + + stopRun(target: CloudTaskCommandTarget): Promise { + return this.transport.stopRun(target); + } +} diff --git a/packages/core/src/sessions/cloudTaskQueue.test.ts b/packages/core/src/sessions/cloudTaskQueue.test.ts new file mode 100644 index 0000000000..93967d67c9 --- /dev/null +++ b/packages/core/src/sessions/cloudTaskQueue.test.ts @@ -0,0 +1,261 @@ +import { describe, expect, it, vi } from "vitest"; +import { + CloudTaskQueue, + type CloudTaskQueuedMessage, + combineCloudTaskQueuedMessages, +} from "./cloudTaskQueue"; + +interface Attachment { + uri: string; +} + +function createQueue(): CloudTaskQueue { + let nextId = 0; + let now = 100; + return new CloudTaskQueue({ + createId: () => `message-${++nextId}`, + now: () => ++now, + }); +} + +function message( + id: string, + content: string, + attachments: Attachment[] = [], +): CloudTaskQueuedMessage { + return { id, content, attachments, queuedAt: 1 }; +} + +function ids(queue: CloudTaskQueue, taskId = "task-1"): string[] { + return queue.getQueue(taskId).map(({ id }) => id); +} + +describe("CloudTaskQueue", () => { + it("enqueues messages in FIFO order with host-provided metadata", () => { + const queue = createQueue(); + + const first = queue.enqueue("task-1", "first", [{ uri: "one" }]); + const second = queue.enqueue("task-1", "second", []); + + expect(queue.getQueue("task-1")).toEqual([first, second]); + expect(first).toEqual({ + id: "message-1", + content: "first", + attachments: [{ uri: "one" }], + queuedAt: 101, + }); + }); + + it("drains all messages in FIFO order by default", () => { + const queue = createQueue(); + queue.enqueue("task-1", "first", []); + queue.enqueue("task-1", "second", []); + + expect(queue.drain("task-1").map(({ content }) => content)).toEqual([ + "first", + "second", + ]); + expect(queue.getQueue("task-1")).toEqual([]); + }); + + it("stops a drain before the message being edited", () => { + const queue = createQueue(); + const first = queue.enqueue("task-1", "first", []); + const edited = queue.enqueue("task-1", "edited", []); + const last = queue.enqueue("task-1", "last", []); + queue.setEditing("task-1", edited.id); + + expect(queue.drain("task-1", { stopAtEdited: true })).toEqual([first]); + expect(queue.getQueue("task-1")).toEqual([edited, last]); + }); + + it("drains nothing when the head message is being edited", () => { + const queue = createQueue(); + const first = queue.enqueue("task-1", "first", []); + queue.enqueue("task-1", "second", []); + queue.setEditing("task-1", first.id); + + expect(queue.drain("task-1", { stopAtEdited: true })).toEqual([]); + expect(ids(queue)).toEqual(["message-1", "message-2"]); + }); + + it("ignores a stale edit boundary and drains the full queue", () => { + const queue = createQueue(); + queue.enqueue("task-1", "first", []); + queue.setEditing("task-1", "missing"); + + expect(queue.drain("task-1", { stopAtEdited: true })).toHaveLength(1); + expect(queue.getQueue("task-1")).toEqual([]); + }); + + it("prepends restored messages without changing their order", () => { + const queue = createQueue(); + queue.enqueue("task-1", "last", []); + + queue.prepend("task-1", [message("a", "first"), message("b", "second")]); + + expect(ids(queue)).toEqual(["a", "b", "message-1"]); + }); + + it("removes a message and releases its edit boundary", () => { + const queue = createQueue(); + queue.enqueue("task-1", "first", []); + const edited = queue.enqueue("task-1", "edited", []); + queue.enqueue("task-1", "last", []); + queue.setEditing("task-1", edited.id); + + queue.remove("task-1", edited.id); + + expect(queue.drain("task-1", { stopAtEdited: true })).toHaveLength(2); + }); + + it.each([ + ["up" as const, "message-2", ["message-2", "message-1", "message-3"]], + ["down" as const, "message-2", ["message-1", "message-3", "message-2"]], + ])("moves a message one position %s", (direction, messageId, expected) => { + const queue = createQueue(); + queue.enqueue("task-1", "first", []); + queue.enqueue("task-1", "second", []); + queue.enqueue("task-1", "third", []); + + queue.move("task-1", messageId, direction); + + expect(ids(queue)).toEqual(expected); + }); + + it.each([ + ["up" as const, "message-1"], + ["down" as const, "message-3"], + ["up" as const, "missing"], + ])("does not move a message beyond the queue boundary", (direction, id) => { + const queue = createQueue(); + queue.enqueue("task-1", "first", []); + queue.enqueue("task-1", "second", []); + queue.enqueue("task-1", "third", []); + + queue.move("task-1", id, direction); + + expect(ids(queue)).toEqual(["message-1", "message-2", "message-3"]); + }); + + it("updates content and generic attachments without changing position", () => { + const queue = createQueue(); + const first = queue.enqueue("task-1", "first", [{ uri: "old" }]); + queue.enqueue("task-1", "second", []); + + queue.update("task-1", first.id, { + content: "edited", + attachments: [{ uri: "new" }], + }); + + expect(queue.getQueue("task-1")).toEqual([ + { ...first, content: "edited", attachments: [{ uri: "new" }] }, + { + id: "message-2", + content: "second", + attachments: [], + queuedAt: 102, + }, + ]); + }); + + it("clears an edit boundary so the whole queue can drain", () => { + const queue = createQueue(); + const first = queue.enqueue("task-1", "first", []); + queue.enqueue("task-1", "second", []); + queue.setEditing("task-1", first.id); + + queue.clearEditing("task-1"); + + expect(queue.drain("task-1", { stopAtEdited: true })).toHaveLength(2); + }); + + it("keeps queues and edit boundaries isolated by task", () => { + const queue = createQueue(); + const first = queue.enqueue("task-1", "first", []); + queue.enqueue("task-2", "second", []); + queue.setEditing("task-1", first.id); + + expect(queue.drain("task-1", { stopAtEdited: true })).toEqual([]); + expect(queue.drain("task-2", { stopAtEdited: true })).toHaveLength(1); + }); + + it("exposes stable snapshots and notifies subscribers after changes", () => { + const queue = createQueue(); + const initialSnapshot = queue.getSnapshot(); + const snapshots: ReturnType[] = []; + const unsubscribe = queue.subscribe(() => + snapshots.push(queue.getSnapshot()), + ); + + const queued = queue.enqueue("task-1", "first", [{ uri: "image" }]); + queue.setEditing("task-1", queued.id); + + expect(queue.getSnapshot()).toBe(snapshots[1]); + expect(queue.getSnapshot()).not.toBe(initialSnapshot); + expect(snapshots).toEqual([ + { + queuesByTaskId: { "task-1": [queued] }, + editingByTaskId: {}, + }, + { + queuesByTaskId: { "task-1": [queued] }, + editingByTaskId: { "task-1": queued.id }, + }, + ]); + + unsubscribe(); + queue.clearEditing("task-1"); + expect(snapshots).toHaveLength(2); + }); + + it("does not notify subscribers for no-op operations", () => { + const queue = createQueue(); + const listener = vi.fn(); + queue.subscribe(listener); + + queue.drain("missing"); + queue.prepend("task-1", []); + queue.remove("task-1", "missing"); + queue.move("task-1", "missing", "up"); + queue.update("task-1", "missing", { content: "edited", attachments: [] }); + queue.clearEditing("task-1"); + + expect(listener).not.toHaveBeenCalled(); + }); + + it("does not mutate an earlier snapshot when a message is updated", () => { + const queue = createQueue(); + const queued = queue.enqueue("task-1", "first", [{ uri: "old" }]); + const earlierSnapshot = queue.getSnapshot(); + + queue.update("task-1", queued.id, { + content: "edited", + attachments: [{ uri: "new" }], + }); + + expect(earlierSnapshot.queuesByTaskId["task-1"]?.[0]).toEqual(queued); + expect(queue.getSnapshot().queuesByTaskId["task-1"]?.[0]).toEqual({ + ...queued, + content: "edited", + attachments: [{ uri: "new" }], + }); + }); +}); + +describe("combineCloudTaskQueuedMessages", () => { + it("joins content in order and preserves generic attachments", () => { + const image = { uri: "image" }; + const file = { uri: "file" }; + + expect( + combineCloudTaskQueuedMessages([ + message("a", "first", [image]), + message("b", "second", [file]), + ]), + ).toEqual({ + text: "first\n\nsecond", + attachments: [image, file], + }); + }); +}); diff --git a/packages/core/src/sessions/cloudTaskQueue.ts b/packages/core/src/sessions/cloudTaskQueue.ts new file mode 100644 index 0000000000..a610f90f65 --- /dev/null +++ b/packages/core/src/sessions/cloudTaskQueue.ts @@ -0,0 +1,204 @@ +import { type QueuedMessage, sendableQueuePrefixLength } from "@posthog/shared"; + +export interface CloudTaskQueuedMessage extends QueuedMessage { + attachments: Attachment[]; +} + +export type CloudTaskQueueMoveDirection = "up" | "down"; + +export interface CloudTaskQueueOptions { + createId: () => string; + now?: () => number; +} + +export interface CloudTaskQueueDrainOptions { + stopAtEdited?: boolean; +} + +export interface CloudTaskQueueMessagePatch { + content: string; + attachments: readonly Attachment[]; +} + +export interface CloudTaskQueueSnapshot { + queuesByTaskId: Readonly< + Record[]> + >; + editingByTaskId: Readonly>; +} + +export type CloudTaskQueueListener = () => void; + +const EMPTY_QUEUE: readonly never[] = []; + +export class CloudTaskQueue { + private readonly queuesByTaskId = new Map< + string, + CloudTaskQueuedMessage[] + >(); + private readonly editingByTaskId = new Map(); + private readonly listeners = new Set(); + private readonly createId: () => string; + private readonly now: () => number; + private snapshot: CloudTaskQueueSnapshot = { + queuesByTaskId: {}, + editingByTaskId: {}, + }; + + constructor(options: CloudTaskQueueOptions) { + this.createId = options.createId; + this.now = options.now ?? Date.now; + } + + enqueue( + taskId: string, + content: string, + attachments: readonly Attachment[], + ): CloudTaskQueuedMessage { + const message = { + id: this.createId(), + content, + attachments: [...attachments], + queuedAt: this.now(), + }; + const queue = this.queuesByTaskId.get(taskId); + if (queue) { + queue.push(message); + } else { + this.queuesByTaskId.set(taskId, [message]); + } + this.publish(); + return message; + } + + drain( + taskId: string, + options?: CloudTaskQueueDrainOptions, + ): CloudTaskQueuedMessage[] { + const queue = this.queuesByTaskId.get(taskId); + if (!queue || queue.length === 0) return []; + + const cutoff = options?.stopAtEdited + ? sendableQueuePrefixLength({ + messageQueue: queue, + editingQueuedId: this.editingByTaskId.get(taskId), + }) + : queue.length; + if (cutoff === 0) return []; + + const drained = queue.splice(0, cutoff); + if (queue.length === 0) this.queuesByTaskId.delete(taskId); + this.publish(); + return drained; + } + + prepend( + taskId: string, + messages: readonly CloudTaskQueuedMessage[], + ): void { + if (messages.length === 0) return; + const queue = this.queuesByTaskId.get(taskId) ?? []; + this.queuesByTaskId.set(taskId, [ + ...messages.map((message) => ({ + ...message, + attachments: [...message.attachments], + })), + ...queue, + ]); + this.publish(); + } + + remove(taskId: string, messageId: string): void { + const queue = this.queuesByTaskId.get(taskId); + if (!queue) return; + + const index = queue.findIndex((message) => message.id === messageId); + if (index === -1) return; + + queue.splice(index, 1); + if (queue.length === 0) this.queuesByTaskId.delete(taskId); + if (this.editingByTaskId.get(taskId) === messageId) { + this.editingByTaskId.delete(taskId); + } + this.publish(); + } + + move( + taskId: string, + messageId: string, + direction: CloudTaskQueueMoveDirection, + ): void { + const queue = this.queuesByTaskId.get(taskId); + if (!queue) return; + + const from = queue.findIndex((message) => message.id === messageId); + if (from === -1) return; + + const to = direction === "up" ? from - 1 : from + 1; + if (to < 0 || to >= queue.length) return; + + const [message] = queue.splice(from, 1); + queue.splice(to, 0, message); + this.publish(); + } + + update( + taskId: string, + messageId: string, + patch: CloudTaskQueueMessagePatch, + ): void { + const queue = this.queuesByTaskId.get(taskId); + if (!queue) return; + const index = queue.findIndex((message) => message.id === messageId); + if (index === -1) return; + + queue[index] = { + ...queue[index], + content: patch.content, + attachments: [...patch.attachments], + }; + this.publish(); + } + + setEditing(taskId: string, messageId: string): void { + if (this.editingByTaskId.get(taskId) === messageId) return; + this.editingByTaskId.set(taskId, messageId); + this.publish(); + } + + clearEditing(taskId: string): void { + if (!this.editingByTaskId.delete(taskId)) return; + this.publish(); + } + + getQueue(taskId: string): readonly CloudTaskQueuedMessage[] { + return this.snapshot.queuesByTaskId[taskId] ?? EMPTY_QUEUE; + } + + readonly getSnapshot = (): CloudTaskQueueSnapshot => + this.snapshot; + + readonly subscribe = (listener: CloudTaskQueueListener): (() => void) => { + this.listeners.add(listener); + return () => this.listeners.delete(listener); + }; + + private publish(): void { + this.snapshot = { + queuesByTaskId: Object.fromEntries( + [...this.queuesByTaskId].map(([taskId, queue]) => [taskId, [...queue]]), + ), + editingByTaskId: Object.fromEntries(this.editingByTaskId), + }; + for (const listener of this.listeners) listener(); + } +} + +export function combineCloudTaskQueuedMessages( + messages: readonly CloudTaskQueuedMessage[], +): { text: string; attachments: Attachment[] } { + return { + text: messages.map((message) => message.content).join("\n\n"), + attachments: messages.flatMap((message) => message.attachments), + }; +} diff --git a/packages/core/src/sessions/cloudTaskRunLifecycle.test.ts b/packages/core/src/sessions/cloudTaskRunLifecycle.test.ts new file mode 100644 index 0000000000..2137e40ef4 --- /dev/null +++ b/packages/core/src/sessions/cloudTaskRunLifecycle.test.ts @@ -0,0 +1,115 @@ +import { describe, expect, it, vi } from "vitest"; +import { CloudTaskCommandController } from "./cloudTaskCommandController"; +import { + CloudTaskRunLifecycle, + type CloudTaskRunSession, +} from "./cloudTaskRunLifecycle"; + +function createHarness() { + const sendCommand = vi.fn(async () => undefined); + const stopRun = vi.fn(async () => undefined); + const update = vi.fn(); + let current: CloudTaskRunSession | undefined = session; + const lifecycle = new CloudTaskRunLifecycle( + new CloudTaskCommandController({ sendCommand, stopRun }), + { + get: () => current, + update: (taskRunId, patch) => { + update(taskRunId, patch); + if (current?.taskRunId === taskRunId) + current = { ...current, ...patch }; + }, + }, + ); + return { + lifecycle, + stopRun, + update, + setCurrent: (next: CloudTaskRunSession) => { + current = next; + }, + getCurrent: () => current, + }; +} + +const session: CloudTaskRunSession = { + taskRunId: "run-1", + stopRequested: false, + isPromptPending: true, + promptStartedAt: 123, + activityVersion: 1, +}; + +describe("CloudTaskRunLifecycle", () => { + it("marks a mounted session stopping before dispatch", async () => { + const { lifecycle, stopRun, update } = createHarness(); + + await lifecycle.stopRun({ taskId: "task-1", taskRunId: "run-1" }, session); + + expect(update).toHaveBeenCalledWith("run-1", { + stopRequested: true, + isPromptPending: false, + promptStartedAt: null, + }); + expect(stopRun).toHaveBeenCalledWith({ + taskId: "task-1", + taskRunId: "run-1", + }); + }); + + it("restores prompt state when stopping fails", async () => { + const { lifecycle, stopRun, update } = createHarness(); + stopRun.mockRejectedValueOnce(new Error("failed")); + + await expect( + lifecycle.stopRun({ taskId: "task-1", taskRunId: "run-1" }, session), + ).rejects.toThrow("failed"); + + expect(update).toHaveBeenLastCalledWith("run-1", { + stopRequested: false, + isPromptPending: true, + promptStartedAt: 123, + }); + }); + + it("stops an unmounted run without writing session state", async () => { + const { lifecycle, update } = createHarness(); + + await lifecycle.stopRun({ taskId: "task-1", taskRunId: "run-1" }); + + expect(update).not.toHaveBeenCalled(); + }); + + it("does not restore stale prompt state after newer activity", async () => { + let rejectStop: (error: Error) => void = () => undefined; + const { lifecycle, stopRun, setCurrent, getCurrent } = createHarness(); + stopRun.mockImplementationOnce( + () => + new Promise((_, reject) => { + rejectStop = reject; + }), + ); + + const stopping = lifecycle.stopRun( + { taskId: "task-1", taskRunId: "run-1" }, + session, + ); + setCurrent({ + ...session, + stopRequested: true, + isPromptPending: false, + promptStartedAt: null, + activityVersion: 2, + }); + rejectStop(new Error("failed")); + + await expect(stopping).rejects.toThrow("failed"); + expect(getCurrent()).toEqual({ + ...session, + stopRequested: true, + isPromptPending: false, + promptStartedAt: null, + activityVersion: 2, + }); + }); +}); diff --git a/packages/core/src/sessions/cloudTaskRunLifecycle.ts b/packages/core/src/sessions/cloudTaskRunLifecycle.ts new file mode 100644 index 0000000000..6dff3ae271 --- /dev/null +++ b/packages/core/src/sessions/cloudTaskRunLifecycle.ts @@ -0,0 +1,64 @@ +import type { + CloudTaskCommandController, + CloudTaskCommandTarget, +} from "./cloudTaskCommandController"; + +export interface CloudTaskRunSession { + taskRunId: string; + stopRequested?: boolean; + isPromptPending: boolean; + promptStartedAt?: number | null; + activityVersion: string | number; +} + +export interface CloudTaskRunStatePort { + get(taskRunId: string): CloudTaskRunSession | undefined; + update( + taskRunId: string, + patch: { + stopRequested?: boolean; + isPromptPending: boolean; + promptStartedAt?: number | null; + }, + ): void; +} + +export class CloudTaskRunLifecycle { + constructor( + private readonly commands: CloudTaskCommandController, + private readonly state: CloudTaskRunStatePort, + ) {} + + async stopRun( + target: CloudTaskCommandTarget, + session?: CloudTaskRunSession, + ): Promise { + if (session) { + this.state.update(session.taskRunId, { + stopRequested: true, + isPromptPending: false, + promptStartedAt: null, + }); + } + + try { + await this.commands.stopRun(target); + } catch (error) { + const current = session ? this.state.get(session.taskRunId) : undefined; + if ( + session && + current?.stopRequested === true && + current.isPromptPending === false && + current.promptStartedAt == null && + current.activityVersion === session.activityVersion + ) { + this.state.update(session.taskRunId, { + stopRequested: session.stopRequested, + isPromptPending: session.isPromptPending, + promptStartedAt: session.promptStartedAt, + }); + } + throw error; + } + } +} diff --git a/packages/core/src/sessions/cloudTaskSessionService.test.ts b/packages/core/src/sessions/cloudTaskSessionService.test.ts new file mode 100644 index 0000000000..f433275275 --- /dev/null +++ b/packages/core/src/sessions/cloudTaskSessionService.test.ts @@ -0,0 +1,421 @@ +import type { CloudTaskUpdatePayload, StoredLogEntry } from "@posthog/shared"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import type { CloudTaskQueuedMessage } from "./cloudTaskQueue"; +import { + type CloudTaskSession, + CloudTaskSessionService, + type CloudTaskSessionServicePorts, + type CloudTaskWatcherHandle, +} from "./cloudTaskSessionService"; + +interface TestAttachment { + name: string; +} + +function session(overrides: Partial = {}): CloudTaskSession { + return { + taskRunId: "run-1", + taskId: "task-1", + events: [], + status: "connected", + isPromptPending: false, + ...overrides, + }; +} + +function logEntry(method: string, params?: unknown): StoredLogEntry { + return { + type: "notification", + notification: { method, ...(params === undefined ? {} : { params }) }, + }; +} + +function createHarness(initialSessions: CloudTaskSession[] = [session()]): { + service: CloudTaskSessionService; + ports: CloudTaskSessionServicePorts; + sessions: Map; + queue: CloudTaskQueuedMessage[]; + watcher: CloudTaskWatcherHandle; + update: (payload: CloudTaskUpdatePayload) => void; +} { + const sessions = new Map( + initialSessions.map((current) => [current.taskRunId, current]), + ); + const queue: CloudTaskQueuedMessage[] = []; + let onUpdate: (payload: CloudTaskUpdatePayload) => void = () => {}; + const watcher = { + stop: vi.fn(), + reconnectIfDisconnected: vi.fn(), + }; + const ports: CloudTaskSessionServicePorts = { + state: { + getByTaskId: (taskId) => + [...sessions.values()].find((current) => current.taskId === taskId), + getByRunId: (runId) => sessions.get(runId), + set: (current) => sessions.set(current.taskRunId, current), + update: (runId, updater) => { + const current = sessions.get(runId); + if (current) sessions.set(runId, updater(current)); + }, + remove: (runId) => { + sessions.delete(runId); + }, + }, + api: { + getTask: vi.fn(async () => ({ + id: "task-1", + latestRun: { id: "run-1", branch: "main" }, + })), + runTask: vi.fn(async () => ({ + id: "task-1", + latestRun: { id: "run-2" }, + })), + sendCommand: vi.fn(async () => {}), + cancelRun: vi.fn(async () => {}), + classifyCommandError: vi.fn(() => ({ kind: "other" as const })), + }, + watchers: { + create: vi.fn((args) => { + onUpdate = args.onUpdate; + return watcher; + }), + }, + queue: { + get: () => queue, + drain: () => queue.splice(0, queue.length), + prepend: (_taskId, messages) => queue.unshift(...messages), + remove: (_taskId, messageId) => { + const index = queue.findIndex((message) => message.id === messageId); + if (index >= 0) queue.splice(index, 1); + }, + combine: (messages) => ({ + text: messages.map((message) => message.content).join("\n"), + attachments: messages.flatMap((message) => message.attachments), + }), + }, + prompts: { + prepare: vi.fn(async (prompt) => ({ wirePayload: `wire:${prompt}` })), + }, + time: { + now: vi.fn(() => 1000), + defer: vi.fn((callback) => callback()), + }, + logger: { + debug: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + }, + effects: { + onCompletion: vi.fn(), + onNotification: vi.fn(), + }, + preferences: { + getComposerConfig: vi.fn(() => undefined), + isRtkEnabled: vi.fn(() => true), + }, + }; + return { + service: new CloudTaskSessionService(ports), + ports, + sessions, + queue, + watcher, + update: (payload) => onUpdate(payload), + }; +} + +describe("CloudTaskSessionService", () => { + beforeEach(() => { + vi.restoreAllMocks(); + }); + + it("creates a missing run and starts its watcher", async () => { + const harness = createHarness([]); + + await harness.service.connect({ id: "task-1", title: "Task" }); + + expect(harness.ports.api.runTask).toHaveBeenCalledWith("task-1"); + expect(harness.sessions.get("run-2")).toMatchObject({ + taskId: "task-1", + taskTitle: "Task", + status: "connecting", + isPromptPending: true, + awaitingPing: true, + promptStartedAt: 1000, + }); + expect(harness.ports.watchers.create).toHaveBeenCalledWith( + expect.objectContaining({ taskId: "task-1", runId: "run-2" }), + ); + }); + + it("sends a local echo and rolls it back when the command fails", async () => { + const harness = createHarness(); + vi.mocked(harness.ports.api.sendCommand).mockRejectedValueOnce( + new Error("offline"), + ); + + await expect(harness.service.sendPrompt("task-1", "hello")).rejects.toThrow( + "offline", + ); + + expect(harness.sessions.get("run-1")).toMatchObject({ + events: [], + isPromptPending: false, + }); + expect(harness.ports.api.sendCommand).toHaveBeenCalledWith( + "task-1", + "run-1", + "user_message", + { content: "wire:hello" }, + ); + }); + + it("resumes an inactive sandbox with the current composer configuration", async () => { + const harness = createHarness([session({ status: "disconnected" })]); + await harness.service.connect({ + id: "task-1", + latestRun: { id: "run-1" }, + }); + vi.mocked(harness.ports.api.sendCommand).mockRejectedValueOnce( + new Error("inactive"), + ); + vi.mocked(harness.ports.api.classifyCommandError).mockReturnValueOnce({ + kind: "sandbox_inactive", + }); + vi.mocked(harness.ports.preferences.getComposerConfig).mockReturnValueOnce({ + reasoning: "high", + mode: "plan", + }); + + await harness.service.sendPrompt("task-1", "continue"); + + expect(harness.ports.api.runTask).toHaveBeenCalledWith("task-1", { + branch: "main", + resumeFromRunId: "run-1", + pendingUserMessage: "wire:continue", + reasoningEffort: "high", + initialPermissionMode: "plan", + rtkEnabled: true, + }); + expect(harness.watcher.stop).toHaveBeenCalledOnce(); + expect(harness.sessions.has("run-1")).toBe(false); + expect(harness.sessions.get("run-2")).toMatchObject({ + taskRunId: "run-2", + status: "connecting", + isPromptPending: true, + awaitingPing: true, + }); + }); + + it("routes permission responses by cloud request id and clears it", async () => { + const harness = createHarness([ + session({ + cloudPermissionRequestIds: { "tool-1": "request-1" }, + pendingPermissions: { + "tool-1": { + requestId: "request-1", + toolCall: { toolCallId: "tool-1", title: "Run", kind: "execute" }, + options: [], + }, + }, + }), + ]); + + await harness.service.sendPermissionResponse("task-1", { + toolCallId: "tool-1", + optionId: "allow", + displayText: "Allowed", + }); + + expect(harness.ports.api.sendCommand).toHaveBeenCalledWith( + "task-1", + "run-1", + "permission_response", + { + requestId: "request-1", + optionId: "allow", + }, + ); + expect(harness.sessions.get("run-1")?.cloudPermissionRequestIds).toEqual( + {}, + ); + }); + + it("dispatches configuration changes without mutating portable state", async () => { + const currentSession = session(); + const harness = createHarness([currentSession]); + + await harness.service.setConfigOption("task-1", "model", "sonnet"); + + expect(harness.ports.api.sendCommand).toHaveBeenCalledWith( + "task-1", + "run-1", + "set_config_option", + { configId: "model", value: "sonnet" }, + ); + expect(harness.sessions.get("run-1")).toEqual(currentSession); + }); + + it("retains prompt state when cancellation dispatch fails", async () => { + const harness = createHarness([ + session({ isPromptPending: true, awaitingPing: true }), + ]); + vi.mocked(harness.ports.api.sendCommand).mockRejectedValueOnce( + new Error("failed"), + ); + + await expect(harness.service.cancelPrompt("task-1")).resolves.toBe(false); + + expect(harness.sessions.get("run-1")).toMatchObject({ + isPromptPending: true, + awaitingPing: true, + }); + expect(harness.ports.logger.error).toHaveBeenCalledWith( + "Failed to cancel cloud prompt", + expect.any(Error), + ); + }); + + it("restores optimistic stop state when run cancellation fails", async () => { + const harness = createHarness([session({ isPromptPending: true })]); + vi.mocked(harness.ports.api.cancelRun).mockRejectedValueOnce( + new Error("failed"), + ); + + await expect(harness.service.stopRun("task-1")).resolves.toBe(false); + + expect(harness.sessions.get("run-1")).toMatchObject({ + stopRequested: undefined, + isPromptPending: true, + }); + }); + + it("dispatches stop requests to the active cloud run", async () => { + const harness = createHarness([session({ isPromptPending: true })]); + + await expect(harness.service.stopRun("task-1")).resolves.toBe(true); + + expect(harness.ports.api.cancelRun).toHaveBeenCalledWith("task-1", "run-1"); + }); + + it("restores drained messages when queue delivery fails", async () => { + const harness = createHarness(); + harness.queue.push({ + id: "message-1", + content: "queued", + attachments: [], + queuedAt: 1000, + }); + vi.mocked(harness.ports.api.sendCommand).mockRejectedValueOnce( + new Error("failed"), + ); + + await harness.service.flushQueuedMessages("task-1"); + + expect(harness.queue.map((message) => message.id)).toEqual(["message-1"]); + }); + + it("tracks compaction and blocks steering until compaction ends", async () => { + const harness = createHarness([session({ isPromptPending: true })]); + harness.queue.push({ + id: "message-1", + content: "queued", + attachments: [], + queuedAt: 1000, + }); + + harness.service.handleUpdate("run-1", { + kind: "logs", + taskId: "task-1", + runId: "run-1", + newEntries: [ + logEntry("_posthog/status", { + status: "compacting", + isComplete: false, + }), + ], + totalEntryCount: 1, + }); + await harness.service.steerQueuedMessage("task-1", "message-1"); + + expect(harness.sessions.get("run-1")?.isCompacting).toBe(true); + expect(harness.queue).toHaveLength(1); + expect(harness.ports.api.sendCommand).not.toHaveBeenCalled(); + }); + + it("notifies and flushes the queue when a live turn ends", () => { + const harness = createHarness([ + session({ + isPromptPending: true, + awaitingPing: true, + promptStartedAt: 500, + }), + ]); + harness.queue.push({ + id: "message-1", + content: "next", + attachments: [], + queuedAt: 1000, + }); + + harness.service.handleUpdate("run-1", { + kind: "logs", + taskId: "task-1", + runId: "run-1", + newEntries: [logEntry("_posthog/turn_complete")], + totalEntryCount: 1, + }); + + expect(harness.ports.effects.onCompletion).toHaveBeenCalledWith({ + taskId: "task-1", + taskRunId: "run-1", + promptStartedAt: 500, + }); + expect(harness.ports.effects.onNotification).toHaveBeenCalledWith({ + taskId: "task-1", + taskRunId: "run-1", + kind: "turn_complete", + }); + expect(harness.ports.time.defer).toHaveBeenCalledOnce(); + }); + + it("records terminal failure and emits a status-only completion", () => { + const harness = createHarness([ + session({ isPromptPending: true, awaitingPing: true }), + ]); + + harness.service.handleUpdate("run-1", { + kind: "status", + taskId: "task-1", + runId: "run-1", + status: "failed", + errorMessage: "boom", + }); + + expect(harness.sessions.get("run-1")).toMatchObject({ + terminalStatus: "failed", + isPromptPending: false, + awaitingPing: false, + lastError: "boom", + }); + expect(harness.ports.effects.onNotification).toHaveBeenCalledWith( + expect.objectContaining({ kind: "task_failed" }), + ); + }); + + it("reconnects every active watcher and stops it on disconnect", async () => { + const harness = createHarness([]); + await harness.service.connect({ + id: "task-1", + latestRun: { id: "run-1" }, + }); + + harness.service.reconnectWatchers(); + harness.service.disconnect("task-1"); + + expect(harness.watcher.reconnectIfDisconnected).toHaveBeenCalledOnce(); + expect(harness.watcher.stop).toHaveBeenCalledOnce(); + expect(harness.sessions.size).toBe(0); + }); +}); diff --git a/packages/core/src/sessions/cloudTaskSessionService.ts b/packages/core/src/sessions/cloudTaskSessionService.ts new file mode 100644 index 0000000000..f2f6fd25de --- /dev/null +++ b/packages/core/src/sessions/cloudTaskSessionService.ts @@ -0,0 +1,942 @@ +import type { + CloudPermissionOption, + CloudTaskUpdatePayload, + StoredLogEntry, +} from "@posthog/shared"; +import { isTerminalStatus } from "@posthog/shared"; +import { CloudTaskCommandController } from "./cloudTaskCommandController"; +import type { CloudTaskQueuedMessage } from "./cloudTaskQueue"; +import { CloudTaskRunLifecycle } from "./cloudTaskRunLifecycle"; +import { + convertStoredEntriesToPortableSessionEvents, + type PortableSessionEvent, + type PortableSessionNotification, +} from "./portableSessionEvents"; + +export type CloudTaskSessionConnectionStatus = + | "connecting" + | "connected" + | "disconnected" + | "error"; + +export type CloudTaskSessionTerminalStatus = "completed" | "failed"; + +export type CloudTaskSessionNotificationKind = + | "turn_complete" + | "awaiting_user_input" + | "task_failed"; + +export interface CloudTaskSessionPermissionRequest { + requestId: string; + toolCall: { + toolCallId: string; + title: string; + kind: string; + content?: unknown[]; + rawInput?: Record; + _meta?: Record; + }; + options: CloudPermissionOption[]; + response?: { + optionId: string; + displayText: string; + answers?: Record; + customInput?: string; + }; +} + +export interface CloudTaskSession { + taskRunId: string; + taskId: string; + taskTitle?: string; + events: PortableSessionEvent[]; + status: CloudTaskSessionConnectionStatus; + isPromptPending: boolean; + localUserEchoes?: Set; + terminalStatus?: CloudTaskSessionTerminalStatus; + lastError?: string | null; + awaitingPing?: boolean; + promptStartedAt?: number; + awaitingAgentOutput?: boolean; + lastEventAt?: number; + cloudPermissionRequestIds?: Record; + pendingPermissions?: Record; + isCompacting?: boolean; + stopRequested?: boolean; +} + +export interface CloudTaskSessionTask { + id: string; + title?: string; + latestRun?: CloudTaskSessionRun; +} + +export interface CloudTaskSessionRun { + id: string; + branch?: string | null; + reasoningEffort?: string | null; + initialPermissionMode?: string; +} + +export interface CloudTaskRunOptions { + branch?: string | null; + resumeFromRunId?: string; + pendingUserMessage?: string; + reasoningEffort?: string; + initialPermissionMode?: string; + rtkEnabled?: boolean; +} + +export interface CloudTaskComposerConfig { + reasoning?: string; + mode?: string; +} + +export interface CloudTaskPermissionResponse { + toolCallId: string; + optionId: string; + answers?: Record; + customInput?: string; + displayText: string; +} + +export interface CloudTaskPreparedPrompt { + wirePayload: string; + eventAttachments?: PortableSessionNotification["update"] extends infer Update + ? Update extends { attachments?: infer Attachments } + ? Attachments + : never + : never; +} + +export interface CloudTaskSessionStatePort { + getByTaskId(taskId: string): CloudTaskSession | undefined; + getByRunId(runId: string): CloudTaskSession | undefined; + set(session: CloudTaskSession): void; + update( + runId: string, + updater: (session: CloudTaskSession) => CloudTaskSession, + ): void; + remove(runId: string): void; +} + +export interface CloudTaskSessionApiPort { + getTask(taskId: string): Promise; + runTask( + taskId: string, + options?: CloudTaskRunOptions, + ): Promise; + sendCommand( + taskId: string, + runId: string, + command: + | "user_message" + | "permission_response" + | "set_config_option" + | "cancel", + payload?: Record, + ): Promise; + cancelRun(taskId: string, runId: string): Promise; + classifyCommandError( + error: unknown, + ): { kind: "sandbox_inactive" } | { kind: "transient" } | { kind: "other" }; +} + +export interface CloudTaskWatcherHandle { + stop(): void; + reconnectIfDisconnected(): void; +} + +export interface CloudTaskWatcherPort { + create(args: { + taskId: string; + runId: string; + onUpdate: (update: CloudTaskUpdatePayload) => void; + }): CloudTaskWatcherHandle; +} + +export interface CloudTaskQueuePort { + get(taskId: string): readonly CloudTaskQueuedMessage[]; + drain( + taskId: string, + options: { stopAtEdited: true }, + ): CloudTaskQueuedMessage[]; + prepend( + taskId: string, + messages: readonly CloudTaskQueuedMessage[], + ): void; + remove(taskId: string, messageId: string): void; + combine(messages: readonly CloudTaskQueuedMessage[]): { + text: string; + attachments: Attachment[]; + }; +} + +export interface CloudTaskPromptPort { + prepare( + prompt: string, + attachments: Attachment[], + ): Promise; + reinjectSnapshotAttachments?( + taskRunId: string, + events: PortableSessionEvent[], + ): void; + recordAttachmentEcho?( + taskRunId: string, + prompt: string, + attachments: NonNullable, + ): void; +} + +export interface CloudTaskSessionTimePort { + now(): number; + defer(callback: () => void): void; +} + +export interface CloudTaskSessionLogger { + debug(message: string, context?: unknown): void; + info(message: string, context?: unknown): void; + warn(message: string, context?: unknown): void; + error(message: string, error?: unknown): void; +} + +export interface CloudTaskSessionEffectsPort { + onCompletion(args: { + taskId: string; + taskRunId: string; + promptStartedAt?: number; + }): void | Promise; + onNotification(args: { + taskId: string; + taskRunId: string; + kind: CloudTaskSessionNotificationKind; + }): void | Promise; +} + +export interface CloudTaskSessionPreferencesPort { + getComposerConfig(taskId: string): CloudTaskComposerConfig | undefined; + isRtkEnabled(): boolean; +} + +export interface CloudTaskSessionServicePorts { + state: CloudTaskSessionStatePort; + api: CloudTaskSessionApiPort; + watchers: CloudTaskWatcherPort; + queue: CloudTaskQueuePort; + prompts: CloudTaskPromptPort; + time: CloudTaskSessionTimePort; + logger: CloudTaskSessionLogger; + effects: CloudTaskSessionEffectsPort; + preferences: CloudTaskSessionPreferencesPort; +} + +const visibleAgentSessionUpdates = new Set([ + "agent_message_chunk", + "agent_message", + "agent_thought_chunk", + "tool_call", + "tool_call_update", +]); + +const turnEndMethods = new Set([ + "_posthog/turn_complete", + "_posthog/task_complete", + "_posthog/error", + "_posthog/awaiting_user_input", +]); + +interface EntryAnalysis { + hasTurnEnd: boolean; + hasAwaitingUserInput: boolean; + hasTurnCompleted: boolean; + hasTurnFailed: boolean; + hasVisibleAgentOutput: boolean; + externalUserMessageCount: number; + agentMessageFinalized: boolean; + compacting?: boolean; +} + +export class CloudTaskSessionService { + private readonly watcherHandles = new Map(); + private readonly connectAttempts = new Set(); + private readonly flushingTasks = new Set(); + private readonly commandController: CloudTaskCommandController; + private readonly runLifecycle: CloudTaskRunLifecycle; + + constructor( + private readonly ports: CloudTaskSessionServicePorts, + ) { + this.commandController = new CloudTaskCommandController({ + sendCommand: (target, method, params) => + this.ports.api.sendCommand( + target.taskId, + target.taskRunId, + method, + params, + ), + stopRun: (target) => + this.ports.api.cancelRun(target.taskId, target.taskRunId), + }); + this.runLifecycle = new CloudTaskRunLifecycle(this.commandController, { + get: (taskRunId) => { + const session = this.ports.state.getByRunId(taskRunId); + return session + ? { + ...session, + activityVersion: `${session.events.length}:${session.lastEventAt ?? ""}:${session.terminalStatus ?? ""}:${session.status}`, + } + : undefined; + }, + update: (taskRunId, patch) => + this.ports.state.update(taskRunId, (current) => ({ + ...current, + ...patch, + promptStartedAt: + patch.promptStartedAt === null ? undefined : patch.promptStartedAt, + })), + }); + } + + async connect(task: CloudTaskSessionTask): Promise { + if (this.connectAttempts.has(task.id)) return; + if (this.ports.state.getByTaskId(task.id)?.status === "connected") return; + + this.connectAttempts.add(task.id); + try { + let runId = task.latestRun?.id; + let awaitingPing = false; + if (!runId) { + runId = (await this.ports.api.runTask(task.id)).latestRun?.id; + if (!runId) throw new Error("Cloud run was created without an id"); + awaitingPing = true; + } + + const now = this.ports.time.now(); + this.ports.state.set({ + taskRunId: runId, + taskId: task.id, + taskTitle: task.title, + events: [], + status: "connecting", + isPromptPending: true, + awaitingPing, + promptStartedAt: awaitingPing ? now : undefined, + awaitingAgentOutput: true, + }); + this.startWatcher(runId, task.id); + } catch (error) { + this.ports.logger.error("Failed to connect to cloud task", error); + throw error; + } finally { + this.connectAttempts.delete(task.id); + } + } + + disconnect(taskId: string): void { + const session = this.ports.state.getByTaskId(taskId); + if (!session) return; + this.stopWatcher(session.taskRunId); + this.ports.state.remove(session.taskRunId); + } + + async sendPrompt( + taskId: string, + prompt: string, + attachments: Attachment[] = [], + ): Promise { + const session = this.requireSession(taskId); + const prepared = await this.ports.prompts.prepare(prompt, attachments); + const userEvent = this.createUserEvent(prompt, prepared.eventAttachments); + + if (prepared.eventAttachments?.length) { + this.ports.prompts.recordAttachmentEcho?.( + session.taskRunId, + prompt, + prepared.eventAttachments, + ); + } + this.addLocalEcho(session.taskRunId, prompt, userEvent); + + try { + await this.commandController.sendUserMessage( + { taskId, taskRunId: session.taskRunId }, + prepared.wirePayload, + ); + } catch (error) { + const classification = this.ports.api.classifyCommandError(error); + if (classification.kind === "sandbox_inactive") { + try { + await this.resume(taskId, session.taskRunId, prepared.wirePayload); + return; + } catch (resumeError) { + this.rollbackLocalEcho(session.taskRunId, prompt, userEvent); + throw resumeError; + } + } + this.rollbackLocalEcho(session.taskRunId, prompt, userEvent); + throw error; + } + } + + async resume( + taskId: string, + previousRunId: string, + pendingUserMessage: string, + ): Promise { + const freshTask = await this.ports.api.getTask(taskId); + const previousRun = freshTask.latestRun; + const composerConfig = this.ports.preferences.getComposerConfig(taskId); + const updatedTask = await this.ports.api.runTask(taskId, { + branch: previousRun?.branch ?? null, + resumeFromRunId: previousRunId, + pendingUserMessage, + reasoningEffort: + composerConfig?.reasoning ?? previousRun?.reasoningEffort ?? undefined, + initialPermissionMode: + composerConfig?.mode ?? previousRun?.initialPermissionMode, + rtkEnabled: this.ports.preferences.isRtkEnabled(), + }); + const newRunId = updatedTask.latestRun?.id; + if (!newRunId) throw new Error("Resume run was created without an id"); + + const previousSession = this.ports.state.getByRunId(previousRunId); + if (!previousSession) throw new Error("No active session for previous run"); + + this.stopWatcher(previousRunId); + this.ports.state.remove(previousRunId); + this.ports.state.set({ + ...previousSession, + taskRunId: newRunId, + status: "connecting", + isPromptPending: true, + awaitingPing: true, + promptStartedAt: this.ports.time.now(), + awaitingAgentOutput: true, + }); + this.startWatcher(newRunId, taskId); + } + + async sendPermissionResponse( + taskId: string, + response: CloudTaskPermissionResponse, + ): Promise { + const session = this.requireSession(taskId); + const event = this.createUserEvent(response.displayText); + const requestId = session.cloudPermissionRequestIds?.[response.toolCallId]; + if (!requestId) { + throw new Error("No cloud permission request id found"); + } + + this.ports.state.update(session.taskRunId, (current) => ({ + ...this.withLocalEcho(current, response.displayText, event), + pendingPermissions: { + ...(current.pendingPermissions ?? {}), + ...(current.pendingPermissions?.[response.toolCallId] + ? { + [response.toolCallId]: { + ...current.pendingPermissions[response.toolCallId], + response: { + optionId: response.optionId, + displayText: response.displayText, + ...(response.answers ? { answers: response.answers } : {}), + ...(response.customInput + ? { customInput: response.customInput } + : {}), + }, + }, + } + : {}), + }, + })); + + try { + await this.commandController.respondToPermission( + { taskId, taskRunId: session.taskRunId }, + { + requestId, + optionId: response.optionId, + ...(response.answers ? { answers: response.answers } : {}), + ...(response.customInput + ? { customInput: response.customInput } + : {}), + }, + ); + this.ports.state.update(session.taskRunId, (current) => { + const requestIds = { ...(current.cloudPermissionRequestIds ?? {}) }; + delete requestIds[response.toolCallId]; + return { ...current, cloudPermissionRequestIds: requestIds }; + }); + } catch (error) { + this.ports.state.update(session.taskRunId, (current) => { + const permission = current.pendingPermissions?.[response.toolCallId]; + return { + ...this.withoutLocalEcho(current, response.displayText, event), + pendingPermissions: permission + ? { + ...(current.pendingPermissions ?? {}), + [response.toolCallId]: { ...permission, response: undefined }, + } + : current.pendingPermissions, + isPromptPending: false, + }; + }); + throw error; + } + } + + async setConfigOption( + taskId: string, + configId: string, + value: string, + ): Promise { + const session = this.ports.state.getByTaskId(taskId); + if (!session || session.terminalStatus) return; + await this.commandController.setConfigOption( + { taskId, taskRunId: session.taskRunId }, + configId, + value, + ); + } + + async cancelPrompt(taskId: string): Promise { + const session = this.ports.state.getByTaskId(taskId); + if (!session) return false; + try { + await this.commandController.cancelPrompt({ + taskId, + taskRunId: session.taskRunId, + }); + this.ports.state.update(session.taskRunId, (current) => ({ + ...current, + isPromptPending: false, + awaitingPing: false, + promptStartedAt: undefined, + awaitingAgentOutput: false, + })); + return true; + } catch (error) { + this.ports.logger.error("Failed to cancel cloud prompt", error); + return false; + } + } + + async stopRun(taskId: string): Promise { + const session = this.ports.state.getByTaskId(taskId); + if (!session) return false; + try { + await this.runLifecycle.stopRun( + { taskId, taskRunId: session.taskRunId }, + { + ...session, + activityVersion: `${session.events.length}:${session.lastEventAt ?? ""}:${session.terminalStatus ?? ""}:${session.status}`, + }, + ); + return true; + } catch (error) { + this.ports.logger.error("Failed to stop cloud run", error); + return false; + } + } + + async sendInterrupting( + taskId: string, + prompt: string, + attachments: Attachment[] = [], + ): Promise { + if (this.ports.state.getByTaskId(taskId)?.isPromptPending) { + await this.cancelPrompt(taskId); + } + await this.sendPrompt(taskId, prompt, attachments); + } + + async flushQueuedMessages(taskId: string): Promise { + if (this.flushingTasks.has(taskId)) return; + this.flushingTasks.add(taskId); + try { + const drained = this.ports.queue.drain(taskId, { stopAtEdited: true }); + if (drained.length === 0) return; + const combined = this.ports.queue.combine(drained); + try { + await this.sendInterrupting( + taskId, + combined.text, + combined.attachments, + ); + } catch (error) { + this.ports.queue.prepend(taskId, drained); + this.ports.logger.warn("Failed to flush cloud task queue", error); + } + } finally { + this.flushingTasks.delete(taskId); + } + } + + flushQueuedMessagesIfIdle(taskId: string): void { + const session = this.ports.state.getByTaskId(taskId); + if ( + session?.status === "connected" && + !session.isPromptPending && + !session.terminalStatus && + !session.isCompacting && + this.ports.queue.get(taskId).length > 0 + ) { + void this.flushQueuedMessages(taskId); + } + } + + async steerQueuedMessage(taskId: string, messageId: string): Promise { + const session = this.ports.state.getByTaskId(taskId); + if (!session?.isPromptPending || session.isCompacting) return; + const message = this.ports.queue + .get(taskId) + .find((candidate) => candidate.id === messageId); + if (!message) return; + + this.ports.queue.remove(taskId, messageId); + try { + await this.sendInterrupting(taskId, message.content, message.attachments); + } catch (error) { + this.ports.queue.prepend(taskId, [message]); + throw error; + } + } + + reconnectWatchers(): void { + for (const handle of this.watcherHandles.values()) { + handle.reconnectIfDisconnected(); + } + } + + handleUpdate(taskRunId: string, update: CloudTaskUpdatePayload): void { + if (update.kind === "error") { + this.ports.state.update(taskRunId, (current) => ({ + ...current, + status: "error", + isPromptPending: false, + lastError: update.errorMessage, + })); + return; + } + + if (update.kind === "permission_request") { + this.ports.state.update(taskRunId, (current) => ({ + ...current, + cloudPermissionRequestIds: { + ...(current.cloudPermissionRequestIds ?? {}), + [update.toolCall.toolCallId]: update.requestId, + }, + pendingPermissions: { + ...(current.pendingPermissions ?? {}), + [update.toolCall.toolCallId]: { + requestId: update.requestId, + toolCall: update.toolCall, + options: update.options, + }, + }, + })); + return; + } + + if (update.kind === "snapshot" || update.kind === "logs") { + this.handleLogUpdate(taskRunId, update); + } + + if ( + (update.kind === "status" || update.kind === "snapshot") && + update.status !== undefined && + isTerminalStatus(update.status) + ) { + this.handleTerminalUpdate(taskRunId, update.status, update.errorMessage); + } + } + + private startWatcher(taskRunId: string, taskId: string): void { + if (this.watcherHandles.has(taskRunId)) return; + this.watcherHandles.set( + taskRunId, + this.ports.watchers.create({ + taskId, + runId: taskRunId, + onUpdate: (update) => this.handleUpdate(taskRunId, update), + }), + ); + } + + private stopWatcher(taskRunId: string): void { + this.watcherHandles.get(taskRunId)?.stop(); + this.watcherHandles.delete(taskRunId); + } + + private handleLogUpdate( + taskRunId: string, + update: Extract, + ): void { + const isSnapshot = update.kind === "snapshot"; + const existing = this.ports.state.getByRunId(taskRunId); + if (!existing) return; + const echoSet = isSnapshot + ? new Set() + : new Set(existing.localUserEchoes ?? []); + const entries = isSnapshot + ? update.newEntries + : dedupAgainstLocalEchoes(update.newEntries, echoSet); + const events = convertStoredEntriesToPortableSessionEvents(entries); + if (isSnapshot) { + this.ports.prompts.reinjectSnapshotAttachments?.(taskRunId, events); + } + const analysis = analyzeEntries(entries, isSnapshot ? new Set() : echoSet); + const wasPromptPending = existing.isPromptPending; + + this.ports.state.update(taskRunId, (current) => { + let isPromptPending = current.isPromptPending; + if (analysis.externalUserMessageCount > 0) isPromptPending = true; + if (analysis.hasTurnEnd || analysis.agentMessageFinalized) { + isPromptPending = false; + } + return { + ...current, + events: isSnapshot ? events : [...current.events, ...events], + status: "connected", + isPromptPending, + awaitingPing: + !isSnapshot && (analysis.hasTurnEnd || analysis.agentMessageFinalized) + ? false + : current.awaitingPing, + awaitingAgentOutput: + current.awaitingAgentOutput && !analysis.hasVisibleAgentOutput, + isCompacting: analysis.compacting ?? current.isCompacting, + localUserEchoes: echoSet.size > 0 ? echoSet : undefined, + lastEventAt: + events.length > 0 ? this.ports.time.now() : current.lastEventAt, + }; + }); + + if (!isSnapshot && existing.awaitingPing) { + const kind = notificationKindForAnalysis(analysis); + if (kind) this.emitCompletion(existing, kind); + } + + const after = this.ports.state.getByRunId(taskRunId); + if ( + !isSnapshot && + wasPromptPending && + after && + !after.isPromptPending && + after.status === "connected" && + this.ports.queue.get(after.taskId).length > 0 + ) { + this.ports.time.defer(() => { + void this.flushQueuedMessages(after.taskId); + }); + } + } + + private handleTerminalUpdate( + taskRunId: string, + status: string, + errorMessage?: string | null, + ): void { + const session = this.ports.state.getByRunId(taskRunId); + if (!session) return; + const terminalStatus = mapTerminalStatus(status); + this.ports.state.update(taskRunId, (current) => ({ + ...current, + isPromptPending: false, + terminalStatus, + lastError: errorMessage ?? null, + awaitingPing: false, + })); + if (session.awaitingPing) { + this.emitCompletion( + session, + terminalStatus === "failed" ? "task_failed" : "turn_complete", + ); + } + } + + private emitCompletion( + session: CloudTaskSession, + kind: CloudTaskSessionNotificationKind, + ): void { + void this.ports.effects.onCompletion({ + taskId: session.taskId, + taskRunId: session.taskRunId, + promptStartedAt: session.promptStartedAt, + }); + void this.ports.effects.onNotification({ + taskId: session.taskId, + taskRunId: session.taskRunId, + kind, + }); + } + + private requireSession(taskId: string): CloudTaskSession { + const session = this.ports.state.getByTaskId(taskId); + if (!session) throw new Error("No active session for task"); + return session; + } + + private createUserEvent( + text: string, + attachments?: CloudTaskPreparedPrompt["eventAttachments"], + ): PortableSessionEvent { + return { + type: "session_update", + ts: this.ports.time.now(), + notification: { + update: { + sessionUpdate: "user_message_chunk", + content: { type: "text", text }, + ...(attachments?.length ? { attachments } : {}), + }, + }, + }; + } + + private addLocalEcho( + taskRunId: string, + text: string, + event: PortableSessionEvent, + ): void { + this.ports.state.update(taskRunId, (session) => + this.withLocalEcho(session, text, event), + ); + } + + private withLocalEcho( + session: CloudTaskSession, + text: string, + event: PortableSessionEvent, + ): CloudTaskSession { + const localUserEchoes = new Set(session.localUserEchoes ?? []); + localUserEchoes.add(text); + return { + ...session, + events: [...session.events, event], + localUserEchoes, + isPromptPending: true, + awaitingPing: true, + promptStartedAt: event.ts, + awaitingAgentOutput: true, + }; + } + + private rollbackLocalEcho( + taskRunId: string, + text: string, + event: PortableSessionEvent, + ): void { + this.ports.state.update(taskRunId, (session) => + this.withoutLocalEcho(session, text, event), + ); + } + + private withoutLocalEcho( + session: CloudTaskSession, + text: string, + event: PortableSessionEvent, + ): CloudTaskSession { + const localUserEchoes = new Set(session.localUserEchoes ?? []); + localUserEchoes.delete(text); + return { + ...session, + events: session.events.filter((candidate) => candidate !== event), + localUserEchoes, + isPromptPending: false, + }; + } +} + +function analyzeEntries( + entries: readonly StoredLogEntry[], + localUserEchoes: ReadonlySet, +): EntryAnalysis { + const analysis: EntryAnalysis = { + hasTurnEnd: false, + hasAwaitingUserInput: false, + hasTurnCompleted: false, + hasTurnFailed: false, + hasVisibleAgentOutput: false, + externalUserMessageCount: 0, + agentMessageFinalized: false, + }; + for (const entry of entries) { + const method = entry.notification?.method; + if (method && turnEndMethods.has(method)) { + analysis.hasTurnEnd = true; + analysis.hasAwaitingUserInput ||= + method === "_posthog/awaiting_user_input"; + analysis.hasTurnCompleted ||= + method === "_posthog/turn_complete" || + method === "_posthog/task_complete"; + analysis.hasTurnFailed ||= method === "_posthog/error"; + } + if (method === "_posthog/status") { + const params = entry.notification?.params as + | { status?: string; isComplete?: boolean } + | undefined; + if (params?.status === "compacting") + analysis.compacting = !params.isComplete; + } + if (method === "_posthog/compact_boundary") analysis.compacting = false; + if (entry.type !== "notification" || method !== "session/update") continue; + const notification = entry.notification?.params as + | PortableSessionNotification + | undefined; + const update = notification?.update; + if ( + update?.sessionUpdate && + visibleAgentSessionUpdates.has(update.sessionUpdate) + ) { + analysis.hasVisibleAgentOutput = true; + } + if (update?.sessionUpdate === "agent_message") { + analysis.agentMessageFinalized = true; + } + if ( + update?.sessionUpdate === "user_message_chunk" && + update.content?.text && + !localUserEchoes.has(update.content.text) + ) { + analysis.externalUserMessageCount += 1; + } + } + return analysis; +} + +function dedupAgainstLocalEchoes( + entries: readonly StoredLogEntry[], + localUserEchoes: Set, +): StoredLogEntry[] { + if (localUserEchoes.size === 0) return [...entries]; + return entries.filter((entry) => { + if ( + entry.type !== "notification" || + entry.notification?.method !== "session/update" + ) { + return true; + } + const notification = entry.notification.params as + | PortableSessionNotification + | undefined; + const update = notification?.update; + const text = update?.content?.text; + if (update?.sessionUpdate !== "user_message_chunk" || !text) return true; + if (!localUserEchoes.has(text)) return true; + localUserEchoes.delete(text); + return false; + }); +} + +function notificationKindForAnalysis( + analysis: EntryAnalysis, +): CloudTaskSessionNotificationKind | undefined { + if (analysis.hasAwaitingUserInput) return "awaiting_user_input"; + if (analysis.hasTurnCompleted) return "turn_complete"; + if (analysis.hasTurnFailed) return "task_failed"; + return undefined; +} + +function mapTerminalStatus(status: string): CloudTaskSessionTerminalStatus { + return status === "completed" ? "completed" : "failed"; +} diff --git a/packages/core/src/sessions/cloudTaskUpdateNotifications.test.ts b/packages/core/src/sessions/cloudTaskUpdateNotifications.test.ts index 5ad145a0c0..64ea2e9b24 100644 --- a/packages/core/src/sessions/cloudTaskUpdateNotifications.test.ts +++ b/packages/core/src/sessions/cloudTaskUpdateNotifications.test.ts @@ -127,6 +127,7 @@ function createHarness() { const notifyPermissionRequest = vi.fn(); const enqueueSpeech = vi.fn(); const markActivity = vi.fn(); + const sendCommand = vi.fn().mockResolvedValue({ success: true }); const noopLog = { info: vi.fn(), warn: vi.fn(), @@ -141,6 +142,23 @@ function createHarness() { notifyPermissionRequest, enqueueSpeech, taskViewedApi: { markActivity }, + track: vi.fn(), + buildPermissionToolMetadata: vi.fn().mockReturnValue({}), + fetchAuthState: vi.fn().mockResolvedValue({ + status: "authenticated", + bootstrapComplete: true, + cloudRegion: "us", + currentProjectId: 1, + }), + createAuthenticatedClient: vi.fn().mockReturnValue({ + getTaskRun: vi.fn().mockResolvedValue({ + status: "in_progress", + stage: null, + output: null, + error_message: null, + log_url: "https://example.com/log", + }), + }), getPersistedConfigOptions: () => undefined, setPersistedConfigOptions: vi.fn(), adapterStore: { @@ -161,6 +179,7 @@ function createHarness() { readLocalLogs: { query: vi.fn().mockResolvedValue("") }, }, cloudTask: { + sendCommand: { mutate: sendCommand }, onUpdate: { subscribe: ( _input: unknown, @@ -186,6 +205,9 @@ function createHarness() { notifyPermissionRequest, enqueueSpeech, markActivity, + sendCommand, + service, + sessions, }; } @@ -321,4 +343,81 @@ describe("cloud task update notifications", () => { }); expect(harness.notifyPermissionRequest).toHaveBeenCalledTimes(2); }); + + it("restores a rejected permission response so it can be retried", async () => { + const harness = createHarness(); + harness.sendUpdate({ + taskId: TASK_ID, + runId: RUN_ID, + kind: "permission_request", + requestId: "r1", + toolCall: { toolCallId: "t1", title: "Run command", kind: "execute" }, + options: [], + }); + const session = harness.sessions[RUN_ID]; + const permission = session?.pendingPermissions.get("t1"); + if (permission) permission.receivedAt = Date.now() - 1_000; + if (session) session.pausedDurationMs = 25; + harness.sendCommand.mockResolvedValueOnce({ + success: false, + error: "sandbox rejected response", + }); + + await harness.service.respondToPermission(TASK_ID, "t1", "allow"); + + expect(harness.sessions[RUN_ID]?.pendingPermissions.has("t1")).toBe(true); + expect(harness.sessions[RUN_ID]?.pausedDurationMs).toBe(25); + + await harness.service.respondToPermission(TASK_ID, "t1", "allow"); + expect(harness.sendCommand).toHaveBeenLastCalledWith( + expect.objectContaining({ + method: "permission_response", + params: { requestId: "r1", optionId: "allow" }, + }), + ); + }); + + it("does not restore a stale permission over a newer request", async () => { + const harness = createHarness(); + harness.sendUpdate({ + taskId: TASK_ID, + runId: RUN_ID, + kind: "permission_request", + requestId: "r1", + toolCall: { toolCallId: "t1", title: "Run command", kind: "execute" }, + options: [], + }); + let rejectResponse: (error: Error) => void = () => undefined; + harness.sendCommand.mockImplementationOnce( + () => + new Promise((_, reject) => { + rejectResponse = reject; + }), + ); + + const response = harness.service.respondToPermission( + TASK_ID, + "t1", + "allow", + ); + await vi.waitFor(() => expect(harness.sendCommand).toHaveBeenCalled()); + harness.sendUpdate({ + taskId: TASK_ID, + runId: RUN_ID, + kind: "permission_request", + requestId: "r2", + toolCall: { toolCallId: "t1", title: "Run command", kind: "execute" }, + options: [], + }); + rejectResponse(new Error("failed")); + await response; + + await harness.service.respondToPermission(TASK_ID, "t1", "allow"); + expect(harness.sendCommand).toHaveBeenLastCalledWith( + expect.objectContaining({ + method: "permission_response", + params: { requestId: "r2", optionId: "allow" }, + }), + ); + }); }); diff --git a/packages/core/src/sessions/sessionLogs.test.ts b/packages/core/src/sessions/sessionLogs.test.ts index 76ca883d78..860fbc4ef5 100644 --- a/packages/core/src/sessions/sessionLogs.test.ts +++ b/packages/core/src/sessions/sessionLogs.test.ts @@ -17,6 +17,27 @@ describe("parseSessionLogContent", () => { expect(result.adapter).toBeUndefined(); }); + it("returns an empty result for blank content", () => { + expect(parseSessionLogContent(" \n ")).toEqual({ + notifications: [], + rawEntries: [], + totalLineCount: 0, + parseFailureCount: 0, + }); + }); + + it("collects session update notifications", () => { + const params = { + update: { sessionUpdate: "agent_message_chunk" }, + }; + const content = JSON.stringify({ + type: "notification", + notification: { method: "session/update", params }, + }); + + expect(parseSessionLogContent(content).notifications).toEqual([params]); + }); + it("extracts sessionId and adapter from a posthog/sdk_session notification", () => { const content = JSON.stringify({ type: "notification", diff --git a/packages/core/src/sessions/sessionLogs.ts b/packages/core/src/sessions/sessionLogs.ts index 974cd7e2e0..c2c11be659 100644 --- a/packages/core/src/sessions/sessionLogs.ts +++ b/packages/core/src/sessions/sessionLogs.ts @@ -1,6 +1,8 @@ import type { Adapter, StoredLogEntry } from "@posthog/shared"; +import type { PortableSessionNotification } from "./portableSessionEvents"; export interface ParsedSessionLogs { + notifications: PortableSessionNotification[]; rawEntries: StoredLogEntry[]; totalLineCount: number; parseFailureCount: number; @@ -12,6 +14,16 @@ export function parseSessionLogContent( content: string, options: { onParseError?: (line: string) => void } = {}, ): ParsedSessionLogs { + if (!content.trim()) { + return { + notifications: [], + rawEntries: [], + totalLineCount: 0, + parseFailureCount: 0, + }; + } + + const notifications: PortableSessionNotification[] = []; const rawEntries: StoredLogEntry[] = []; let sessionId: string | undefined; let adapter: Adapter | undefined; @@ -23,6 +35,16 @@ export function parseSessionLogContent( const stored = JSON.parse(line) as StoredLogEntry; rawEntries.push(stored); + if ( + stored.type === "notification" && + stored.notification?.method === "session/update" && + stored.notification.params + ) { + notifications.push( + stored.notification.params as PortableSessionNotification, + ); + } + if ( stored.type === "notification" && stored.notification?.method?.endsWith("posthog/sdk_session") @@ -43,6 +65,7 @@ export function parseSessionLogContent( } return { + notifications, rawEntries, totalLineCount: lines.length, parseFailureCount, diff --git a/packages/core/src/sessions/sessionService.ts b/packages/core/src/sessions/sessionService.ts index 2879033ee8..27340ae03e 100644 --- a/packages/core/src/sessions/sessionService.ts +++ b/packages/core/src/sessions/sessionService.ts @@ -70,6 +70,11 @@ import { buildCloudDefaultConfigOptions, extractLatestConfigOptionsFromEntries, } from "./cloudSessionConfig"; +import { + CloudTaskCommandController, + type CloudTaskPermissionResponseCommand, +} from "./cloudTaskCommandController"; +import { CloudTaskRunLifecycle } from "./cloudTaskRunLifecycle"; import { computeAutoRetryFinalState, OFFLINE_SESSION_MESSAGE, @@ -143,6 +148,8 @@ class GitHubAuthorizationRequiredForCloudHandoffError extends Error { } } +class CloudTaskCommandRejectedError extends Error {} + type TrpcMutation = { mutate: (input?: any) => Promise }; type TrpcQuery = { query: (input?: any) => Promise }; type TrpcSubscription = { @@ -1611,8 +1618,59 @@ export class SessionService { * the placeholder with this richer text instead, then drop it once consumed. */ private initialCloudOptimisticPrompt = new Map(); + private readonly cloudTaskCommandController: CloudTaskCommandController; + private readonly cloudTaskRunLifecycle: CloudTaskRunLifecycle; constructor(private readonly d: SessionServiceDeps) { + this.cloudTaskCommandController = new CloudTaskCommandController({ + sendCommand: async (target, method, params) => { + const auth = await this.getCloudCommandAuth(); + if (!auth) { + throw new Error("No cloud auth credentials available"); + } + const result = await this.d.trpc.cloudTask.sendCommand.mutate({ + taskId: target.taskId, + runId: target.taskRunId, + apiHost: auth.apiHost, + teamId: auth.teamId, + method, + params, + }); + if (!result.success) { + throw new CloudTaskCommandRejectedError( + result.error ?? "Failed to send cloud command", + ); + } + return result.result; + }, + stopRun: async (target) => { + const result = await this.d.trpc.cloudTask.stop.mutate({ + taskId: target.taskId, + runId: target.taskRunId, + }); + if (!result.success) { + throw new CloudTaskCommandRejectedError( + result.error ?? "Failed to stop cloud run", + ); + } + }, + }); + this.cloudTaskRunLifecycle = new CloudTaskRunLifecycle( + this.cloudTaskCommandController, + { + get: (taskRunId) => { + const session = this.getSessionByRunId(taskRunId); + return session + ? { + ...session, + activityVersion: `${session.events.length}:${session.processedLineCount ?? ""}:${session.cloudStatus ?? ""}:${session.status}`, + } + : undefined; + }, + update: (taskRunId, patch) => + this.d.store.updateSession(taskRunId, patch), + }, + ); this.cloudRunIdleTracker = new CloudRunIdleTracker(); this.cloudLogGapReconciler = new CloudLogGapReconciler({ fetchLogs: (logUrl, taskRunId, minEntryCount) => @@ -4115,22 +4173,13 @@ export class SessionService { }); try { - const result = await this.d.trpc.cloudTask.sendCommand.mutate({ - taskId: session.taskId, - runId: session.taskRunId, - apiHost: cloudCommandAuth.apiHost, - teamId: cloudCommandAuth.teamId, - method: "user_message", - params, - }); - - if (!result.success) { - throw new Error(result.error ?? "Failed to send cloud command"); - } - - const commandResult = result.result as - | { queued?: boolean; steered?: boolean; stopReason?: string } - | undefined; + const commandResult = + (await this.cloudTaskCommandController.sendUserMessage( + { taskId: session.taskId, taskRunId: session.taskRunId }, + params, + )) as + | { queued?: boolean; steered?: boolean; stopReason?: string } + | undefined; const stopReason = commandResult?.queued ? "queued" : (commandResult?.stopReason ?? "end_turn"); @@ -4444,26 +4493,14 @@ export class SessionService { return false; } - const auth = await this.getCloudCommandAuth(); - if (!auth) { - this.d.log.error("No auth for cloud cancel"); - return false; - } - - try { - const result = await this.d.trpc.cloudTask.sendCommand.mutate({ - taskId: session.taskId, - runId: session.taskRunId, - apiHost: auth.apiHost, - teamId: auth.teamId, - method: "cancel", - }); - + const trackCancellation = (): void => { const durationSeconds = Math.round( (Date.now() - session.startedAt) / 1000, ); const promptCount = session.events.filter( - (e) => "method" in e.message && e.message.method === "session/prompt", + (event) => + "method" in event.message && + event.message.method === "session/prompt", ).length; this.d.track(ANALYTICS_EVENTS.TASK_RUN_CANCELLED, { task_id: session.taskId, @@ -4471,14 +4508,23 @@ export class SessionService { duration_seconds: durationSeconds, prompts_sent: promptCount, }); + }; - if (!result.success) { - this.d.log.warn("Cloud cancel command failed", { error: result.error }); - return false; - } - + try { + await this.cloudTaskCommandController.cancelPrompt({ + taskId: session.taskId, + taskRunId: session.taskRunId, + }); + trackCancellation(); return true; } catch (error) { + if (error instanceof CloudTaskCommandRejectedError) { + trackCancellation(); + this.d.log.warn("Cloud cancel command failed", { + error: error.message, + }); + return false; + } this.d.log.error("Failed to cancel cloud prompt", error); return false; } @@ -4520,40 +4566,17 @@ export class SessionService { const matchingSession = session?.isCloud && session.taskRunId === taskRunId ? session : undefined; - const previousPromptState = matchingSession - ? { - isPromptPending: matchingSession.isPromptPending, - promptStartedAt: matchingSession.promptStartedAt, - } - : undefined; - if (matchingSession) { - this.d.store.updateSession(matchingSession.taskRunId, { - stopRequested: true, - isPromptPending: false, - promptStartedAt: null, - }); - } - try { - const result = await this.d.trpc.cloudTask.stop.mutate({ - taskId, - runId: taskRunId, - }); - - if (!result.success) { - if (matchingSession) { - this.d.store.updateSession(matchingSession.taskRunId, { - stopRequested: false, - ...previousPromptState, - }); - } - this.d.log.warn("Cloud run stop failed", { - taskId, - error: result.error, - retryable: result.retryable, - }); - return false; - } + await this.cloudTaskRunLifecycle.stopRun( + { taskId, taskRunId }, + matchingSession + ? { + ...matchingSession, + stopRequested: matchingSession.stopRequested ?? false, + activityVersion: `${matchingSession.events.length}:${matchingSession.processedLineCount ?? ""}:${matchingSession.cloudStatus ?? ""}:${matchingSession.status}`, + } + : undefined, + ); const durationSeconds = matchingSession ? Math.round((Date.now() - matchingSession.startedAt) / 1000) @@ -4572,13 +4595,14 @@ export class SessionService { return true; } catch (error) { - if (matchingSession) { - this.d.store.updateSession(matchingSession.taskRunId, { - stopRequested: false, - ...previousPromptState, + if (error instanceof CloudTaskCommandRejectedError) { + this.d.log.warn("Cloud run stop failed", { + taskId, + error: error.message, }); + } else { + this.d.log.error("Failed to stop cloud run", error); } - this.d.log.error("Failed to stop cloud run", error); return false; } } @@ -4595,27 +4619,14 @@ export class SessionService { }; } - /** - * Send a command to the cloud agent server via the backend proxy. - * Handles auth lookup and throws if credentials are unavailable. - */ - private async sendCloudCommand( + private async sendCloudPermissionResponse( session: AgentSession, - method: "permission_response" | "set_config_option", - params: Record, + params: CloudTaskPermissionResponseCommand, ): Promise { - const auth = await this.getCloudCommandAuth(); - if (!auth) { - throw new Error("No cloud auth credentials available"); - } - await this.d.trpc.cloudTask.sendCommand.mutate({ - taskId: session.taskId, - runId: session.taskRunId, - apiHost: auth.apiHost, - teamId: auth.teamId, - method, + await this.cloudTaskCommandController.respondToPermission( + { taskId: session.taskId, taskRunId: session.taskRunId }, params, - }); + ); } private async refreshCloudRunStatus( @@ -4752,19 +4763,28 @@ export class SessionService { // --- Permissions --- - private resolvePermission(session: AgentSession, toolCallId: string): void { + private resolvePermission( + session: AgentSession, + toolCallId: string, + ): { + previousPausedDurationMs: number; + resolvedPausedDurationMs: number; + } { const permission = session.pendingPermissions.get(toolCallId); const newPermissions = new Map(session.pendingPermissions); newPermissions.delete(toolCallId); this.d.store.setPendingPermissions(session.taskRunId, newPermissions); + const previousPausedDurationMs = session.pausedDurationMs ?? 0; + const resolvedPausedDurationMs = permission?.receivedAt + ? previousPausedDurationMs + (Date.now() - permission.receivedAt) + : previousPausedDurationMs; if (permission?.receivedAt) { this.d.store.updateSession(session.taskRunId, { - pausedDurationMs: - (session.pausedDurationMs ?? 0) + - (Date.now() - permission.receivedAt), + pausedDurationMs: resolvedPausedDurationMs, }); } + return { previousPausedDurationMs, resolvedPausedDurationMs }; } /** @@ -4790,7 +4810,7 @@ export class SessionService { }); const cloudRequestId = this.cloudPermissionRequestIds.get(toolCallId); - this.resolvePermission(session, toolCallId); + const permissionResolution = this.resolvePermission(session, toolCallId); try { const refreshedCloudStatus = @@ -4820,7 +4840,7 @@ export class SessionService { if (session.isCloud && cloudRequestId) { this.cloudPermissionRequestIds.delete(toolCallId); try { - await this.sendCloudCommand(session, "permission_response", { + await this.sendCloudPermissionResponse(session, { requestId: cloudRequestId, optionId, customInput, @@ -4865,6 +4885,28 @@ export class SessionService { hasCustomInput: !!customInput, }); } catch (error) { + const latestSession = this.d.store.getSessionByTaskId(taskId); + if ( + latestSession?.pausedDurationMs === + permissionResolution.resolvedPausedDurationMs + ) { + this.d.store.updateSession(latestSession.taskRunId, { + pausedDurationMs: permissionResolution.previousPausedDurationMs, + }); + } + const currentCloudRequestId = + this.cloudPermissionRequestIds.get(toolCallId); + if (permission && latestSession && !currentCloudRequestId) { + const restoredPermissions = new Map(latestSession.pendingPermissions); + restoredPermissions.set(toolCallId, permission); + this.d.store.setPendingPermissions( + latestSession.taskRunId, + restoredPermissions, + ); + } + if (cloudRequestId && !currentCloudRequestId) { + this.cloudPermissionRequestIds.set(toolCallId, cloudRequestId); + } this.d.log.error("Failed to respond to permission", { taskId, toolCallId, @@ -4912,7 +4954,7 @@ export class SessionService { } if (session.isCloud && cloudRequestId) { this.cloudPermissionRequestIds.delete(toolCallId); - await this.sendCloudCommand(session, "permission_response", { + await this.sendCloudPermissionResponse(session, { requestId: cloudRequestId, optionId: "reject_with_feedback", customInput: "User cancelled the permission request.", @@ -4990,10 +5032,11 @@ export class SessionService { try { if (session.isCloud) { - await this.sendCloudCommand(session, "set_config_option", { + await this.cloudTaskCommandController.setConfigOption( + { taskId: session.taskId, taskRunId: session.taskRunId }, configId, value, - }); + ); } else { await this.d.trpc.agent.setConfigOption.mutate({ sessionId: session.taskRunId, @@ -7199,6 +7242,7 @@ export class SessionService { options: { minEntryCount?: number } = {}, ): Promise { const empty: ParsedSessionLogs = { + notifications: [], rawEntries: [], totalLineCount: 0, parseFailureCount: 0, diff --git a/packages/core/src/sessions/sessionServiceStopCloudRun.test.ts b/packages/core/src/sessions/sessionServiceStopCloudRun.test.ts index ba1d1a127f..4dde0281f2 100644 --- a/packages/core/src/sessions/sessionServiceStopCloudRun.test.ts +++ b/packages/core/src/sessions/sessionServiceStopCloudRun.test.ts @@ -45,12 +45,17 @@ function createHarness( } : { id: "run-1", environment: "cloud", status: "in_progress" }, ) { - const updateSession = vi.fn(); + const updateSession = vi.fn( + (_taskRunId: string, updates: Partial) => { + if (session) Object.assign(session, updates); + }, + ); const track = vi.fn(); const deps = { store: { getSessionByTaskId: (taskId: string) => session?.taskId === taskId ? session : undefined, + getSessions: () => (session ? { [session.taskRunId]: session } : {}), updateSession, }, log: { info: vi.fn(), warn: vi.fn(), error: vi.fn(), debug: vi.fn() }, diff --git a/packages/core/vitest.config.ts b/packages/core/vitest.config.ts new file mode 100644 index 0000000000..bed14b24e0 --- /dev/null +++ b/packages/core/vitest.config.ts @@ -0,0 +1,23 @@ +import { defineConfig } from "vitest/config"; +import { trunkTestOptions } from "../../vitest.config.base"; + +export default defineConfig({ + oxc: false, + esbuild: { + tsconfigRaw: { + compilerOptions: { + experimentalDecorators: true, + target: "ES2022", + useDefineForClassFields: false, + verbatimModuleSyntax: true, + }, + }, + }, + test: { + globals: true, + ...trunkTestOptions, + environment: "node", + include: ["src/**/*.test.ts", "src/**/*.test.tsx"], + exclude: ["**/node_modules/**", "**/dist/**"], + }, +});