diff --git a/packages/core/src/sessions/sessionService.ts b/packages/core/src/sessions/sessionService.ts index 751cdc7eca..1e5401647f 100644 --- a/packages/core/src/sessions/sessionService.ts +++ b/packages/core/src/sessions/sessionService.ts @@ -2235,15 +2235,25 @@ export class SessionService { prompt: string | ContentBlock[], options?: { steer?: boolean }, ): Promise<{ stopReason: string }> { - if (!this.d.getIsOnline()) { - throw new Error( - "No internet connection. Please check your connection and try again.", - ); - } - let session = this.d.store.getSessionByTaskId(taskId); if (!session) throw new Error("No active session for task"); + if (!this.d.getIsOnline()) { + // Cloud follow-ups queue (and persist) so an unstable connection can't + // drop them — the queue flushes when the stream reconnects. Local busy + // sessions fall through to the normal queue path below. Anything else + // genuinely can't proceed offline, so surface a clear error. + const canQueueOffline = + session.isCloud || + (session.status === "connected" && + (session.isPromptPending || session.isCompacting)); + if (!canQueueOffline) { + throw new Error( + "No internet connection. Please check your connection and try again.", + ); + } + } + // The /add-dir dialog mutates the per-task additional-directories list and // we re-read it during respawn below. Sending while it's open would race // and respawn with the pre-decision set, so block here. @@ -2390,22 +2400,26 @@ export class SessionService { /** * Send all queued messages as a single prompt. * Called internally when a turn completes and there are queued messages. - * Queue is cleared atomically before sending - if sending fails, messages are lost - * (this is acceptable since the user can re-type; avoiding complex retry logic). + * The queue is drained before sending, but rolled back via + * `prependQueuedMessages` if the send fails — so a transient failure doesn't + * silently drop the follow-ups (mirrors the cloud queue's rollback, and keeps + * the durable mirror in sync). */ private async sendQueuedMessages( taskId: string, ): Promise<{ stopReason: string }> { - const combinedText = this.d.store.dequeueMessagesAsText(taskId); - if (!combinedText) { + const drained = this.d.store.dequeueMessages(taskId); + if (drained.length === 0) { return { stopReason: "skipped" }; } + const combinedText = drained.map((message) => message.content).join("\n\n"); const session = this.d.store.getSessionByTaskId(taskId); if (!session) { - this.d.log.warn("No session found for queued messages, messages lost", { + this.d.store.prependQueuedMessages(taskId, drained); + this.d.log.warn("No session found for queued messages, re-queued", { taskId, - lostMessageLength: combinedText.length, + queued: drained.length, }); return { stopReason: "no_session" }; } @@ -2433,10 +2447,12 @@ export class SessionService { try { return await this.sendLocalPrompt(session, blocks, combinedText); } catch (error) { - // Log that queued messages were lost due to send failure - this.d.log.error("Failed to send queued messages, messages lost", { + // Roll the drain back so a transient send failure doesn't drop the + // follow-ups; the durable mirror re-syncs off the restored queue. + this.d.store.prependQueuedMessages(taskId, drained); + this.d.log.error("Failed to send queued messages, re-queued", { taskId, - lostMessageLength: combinedText.length, + queued: drained.length, error, }); throw error; @@ -2621,6 +2637,23 @@ export class SessionService { return { stopReason: "empty" }; } + // Offline: the run keeps executing in the cloud, so queue (and persist via + // the durable mirror) instead of attempting a doomed network send. The + // queue flushes when the main-process watcher resumes the stream and the + // agent is next idle. Don't retry the watch here — it would fail offline. + if (!this.d.getIsOnline()) { + this.d.store.enqueueMessage( + session.taskId, + transport.promptText, + normalizedPrompt, + ); + this.d.log.info("Cloud message queued (offline)", { + taskId: session.taskId, + cloudStatus: session.cloudStatus, + }); + return { stopReason: "queued" }; + } + if (isTerminalStatus(session.cloudStatus)) { // If the agent never booted (no `run_started`), resuming spins another // sandbox that hits the same provisioning failure — surface the error diff --git a/packages/ui/src/features/sessions/components/SessionView.tsx b/packages/ui/src/features/sessions/components/SessionView.tsx index 5c84aaefa6..1f29e9b694 100644 --- a/packages/ui/src/features/sessions/components/SessionView.tsx +++ b/packages/ui/src/features/sessions/components/SessionView.tsx @@ -312,13 +312,16 @@ export function SessionView({ const handleBeforeSubmit = useCallback( (text: string, clearEditor: () => void): boolean => { - if (!isOnline) { + // Cloud runs execute server-side, so an offline follow-up is queued (and + // persisted) rather than dropped — don't block it. Local sessions can't + // proceed offline. + if (!isOnline && !isCloud) { showOfflineToast(); return false; } return onBeforeSubmit ? onBeforeSubmit(text, clearEditor) : true; }, - [isOnline, onBeforeSubmit], + [isOnline, isCloud, onBeforeSubmit], ); const [isDraggingFile, setIsDraggingFile] = useState(false); @@ -657,10 +660,14 @@ export function SessionView({ placeholder="Type a message... @ to mention files, ! for bash mode, / for skills" disabled={!isRunning && !handoffInProgress} submitDisabledExternal={ - handoffInProgress || !isOnline + handoffInProgress || (!isOnline && !isCloud) } submitTooltipOverride={ - !isOnline ? "No internet connection" : undefined + !isOnline + ? isCloud + ? "Offline — your message will be queued and sent when you reconnect" + : "No internet connection" + : undefined } isLoading={!!isPromptPending} isActiveSession={isActiveSession} diff --git a/packages/ui/src/features/sessions/queuedMessagePersistence.ts b/packages/ui/src/features/sessions/queuedMessagePersistence.ts new file mode 100644 index 0000000000..597db9307b --- /dev/null +++ b/packages/ui/src/features/sessions/queuedMessagePersistence.ts @@ -0,0 +1,121 @@ +import { + sessionStore, + sessionStoreSetters, +} from "@posthog/core/sessions/sessionStore"; +import type { QueuedMessage } from "@posthog/shared"; +import { logger } from "@posthog/ui/shell/logger"; +import { queuedMessageStoreApi } from "./queuedMessageStore"; + +/** + * Keeps the durable {@link queuedMessageStoreApi} mirror in sync with each + * session's in-memory `messageQueue`, and rehydrates persisted follow-ups back + * into a session's core queue when it (re)appears for a task. + * + * A store subscription (rather than wrapping the injected setters) is used so + * EVERY queue mutation is captured — including the several UI call sites that + * mutate `sessionStoreSetters` directly (dock remove, steer/queue toggle, + * cancel-into-editor), which a service-seam wrapper would miss. + * + * Per task the flow is: on first observation, reconcile (merge persisted into + * core, dedup by id) BEFORE mirroring, so a freshly-created empty session can't + * wipe persisted messages that haven't been seeded yet. After reconciliation, + * queue changes mirror straight through (including drains → clear). Session + * removal un-reconciles the task so a later recreation re-seeds from disk. + */ + +const log = logger.scope("queued-messages"); + +const reconciledTasks = new Set(); +const lastQueueByTask = new Map(); +let unsubscribe: (() => void) | null = null; + +function reconcileTask(taskId: string): void { + void queuedMessageStoreApi + .whenHydrated() + .then(() => { + const session = sessionStoreSetters.getSessionByTaskId(taskId); + // The session vanished before hydration completed — leave the persisted + // queue intact so a later recreation can still seed from it. + if (!session) { + return; + } + + const persisted = queuedMessageStoreApi.get(taskId); + const current = session.messageQueue; + const existingIds = new Set(current.map((m) => m.id)); + const missing = persisted.filter((m) => !existingIds.has(m.id)); + + if (missing.length > 0) { + // Persisted follow-ups predate anything typed this session, so they + // belong at the head. The resulting store change re-enters the + // subscription (now reconciled) and mirrors the merged queue. + log.info("Rehydrating persisted queued messages", { + taskId, + count: missing.length, + }); + sessionStoreSetters.prependQueuedMessages(taskId, missing); + return; + } + + // Nothing to seed: make the mirror match the current queue. + lastQueueByTask.set(taskId, current); + queuedMessageStoreApi.set(taskId, current); + }) + .catch((error) => { + log.warn("Failed to reconcile persisted queued messages", { + taskId, + error, + }); + }); +} + +function handleSession(taskId: string, queue: QueuedMessage[]): void { + if (!reconciledTasks.has(taskId)) { + // Mark eagerly so repeat observations before the async reconcile finishes + // don't schedule it twice. + reconciledTasks.add(taskId); + lastQueueByTask.set(taskId, queue); + reconcileTask(taskId); + return; + } + + if (lastQueueByTask.get(taskId) === queue) { + return; + } + lastQueueByTask.set(taskId, queue); + queuedMessageStoreApi.set(taskId, queue); +} + +/** + * Starts mirroring the core session store into the durable queue store. Safe to + * call repeatedly; only the first call installs the subscription. + */ +export function startQueuedMessagePersistence(): void { + if (unsubscribe) { + return; + } + unsubscribe = sessionStore.subscribe((state) => { + const present = new Set(); + for (const session of Object.values(state.sessions)) { + present.add(session.taskId); + handleSession(session.taskId, session.messageQueue); + } + // A removed/evicted session un-reconciles its task so a later recreation + // re-seeds from disk instead of mirroring its empty starting queue. + for (const taskId of [...reconciledTasks]) { + if (!present.has(taskId)) { + reconciledTasks.delete(taskId); + lastQueueByTask.delete(taskId); + } + } + }); +} + +/** + * Drops all in-memory tracking. Pairs with `queuedMessageStoreApi.clearAll()` on + * logout / project switch so tracking doesn't leak across accounts. + */ +export function resetQueuedMessagePersistenceTracking(): void { + reconciledTasks.clear(); + lastQueueByTask.clear(); +} diff --git a/packages/ui/src/features/sessions/queuedMessageStore.test.ts b/packages/ui/src/features/sessions/queuedMessageStore.test.ts new file mode 100644 index 0000000000..e4686c6682 --- /dev/null +++ b/packages/ui/src/features/sessions/queuedMessageStore.test.ts @@ -0,0 +1,53 @@ +import type { QueuedMessage } from "@posthog/shared"; +import { beforeEach, describe, expect, it } from "vitest"; +import { + queuedMessageStoreApi, + useQueuedMessageStore, +} from "./queuedMessageStore"; + +function msg(id: string, queuedAt = 1): QueuedMessage { + return { id, content: `content-${id}`, queuedAt }; +} + +describe("queuedMessageStore", () => { + beforeEach(() => { + useQueuedMessageStore.setState({ byTaskId: {} }); + }); + + it("stores and reads a per-task queue", () => { + queuedMessageStoreApi.set("task-1", [msg("a"), msg("b")]); + expect(queuedMessageStoreApi.get("task-1").map((m) => m.id)).toEqual([ + "a", + "b", + ]); + expect(queuedMessageStoreApi.get("task-2")).toEqual([]); + }); + + it("drops the key when the queue is emptied", () => { + queuedMessageStoreApi.set("task-1", [msg("a")]); + queuedMessageStoreApi.set("task-1", []); + expect("task-1" in useQueuedMessageStore.getState().byTaskId).toBe(false); + }); + + it("caps a task's queue to the newest messages", () => { + const many = Array.from({ length: 25 }, (_, i) => msg(`m${i}`, i)); + queuedMessageStoreApi.set("task-1", many); + const kept = queuedMessageStoreApi.get("task-1"); + expect(kept).toHaveLength(20); + // Newest (tail) retained. + expect(kept[0].id).toBe("m5"); + expect(kept.at(-1)?.id).toBe("m24"); + }); + + it("clears a single task and clears all", () => { + queuedMessageStoreApi.set("task-1", [msg("a")]); + queuedMessageStoreApi.set("task-2", [msg("b")]); + + queuedMessageStoreApi.clear("task-1"); + expect(queuedMessageStoreApi.get("task-1")).toEqual([]); + expect(queuedMessageStoreApi.get("task-2").map((m) => m.id)).toEqual(["b"]); + + queuedMessageStoreApi.clearAll(); + expect(useQueuedMessageStore.getState().byTaskId).toEqual({}); + }); +}); diff --git a/packages/ui/src/features/sessions/queuedMessageStore.ts b/packages/ui/src/features/sessions/queuedMessageStore.ts new file mode 100644 index 0000000000..b62cc7f8a0 --- /dev/null +++ b/packages/ui/src/features/sessions/queuedMessageStore.ts @@ -0,0 +1,135 @@ +import type { QueuedMessage } from "@posthog/shared"; +import { logger } from "@posthog/ui/shell/logger"; +import { electronStorage } from "@posthog/ui/shell/rendererStorage"; +import { create } from "zustand"; +import { persist } from "zustand/middleware"; + +/** + * Durable mirror of each session's follow-up `messageQueue`, keyed by taskId, so + * queued follow-ups survive an app restart (and offline stretches) instead of + * dying with the in-memory core `sessionStore`. + * + * This holds COMMITTED, queued messages — distinct from `pendingTaskPromptStore` + * (unsent draft composer text for a not-yet-created task). The core store stays + * the source of truth; `queuedMessagePersistence` keeps this mirror in sync and + * rehydrates it back into core when a session (re)appears for a task. + */ + +const log = logger.scope("queued-messages"); + +const MAX_TASKS = 50; +const MAX_MESSAGES_PER_TASK = 20; + +function capMessages(messages: QueuedMessage[]): QueuedMessage[] { + if (messages.length <= MAX_MESSAGES_PER_TASK) { + return messages; + } + log.warn("Dropping oldest queued messages beyond per-task cap", { + dropped: messages.length - MAX_MESSAGES_PER_TASK, + }); + // Keep the newest (queue tail). + return messages.slice(-MAX_MESSAGES_PER_TASK); +} + +function newestQueuedAt(messages: QueuedMessage[]): number { + return messages.reduce((max, m) => Math.max(max, m.queuedAt), 0); +} + +function capTasks( + byTaskId: Record, +): Record { + const keys = Object.keys(byTaskId); + if (keys.length <= MAX_TASKS) { + return byTaskId; + } + const keptKeys = keys + .sort((a, b) => newestQueuedAt(byTaskId[b]) - newestQueuedAt(byTaskId[a])) + .slice(0, MAX_TASKS); + log.warn("Dropping oldest queued-message tasks beyond cap", { + dropped: keys.length - keptKeys.length, + }); + const kept: Record = {}; + for (const key of keptKeys) { + kept[key] = byTaskId[key]; + } + return kept; +} + +interface QueuedMessageStore { + byTaskId: Record; + _hasHydrated: boolean; + setHasHydrated: (hydrated: boolean) => void; + set: (taskId: string, messages: QueuedMessage[]) => void; + clear: (taskId: string) => void; + clearAll: () => void; +} + +export const useQueuedMessageStore = create()( + persist( + (set) => ({ + byTaskId: {}, + _hasHydrated: false, + setHasHydrated: (hydrated) => set({ _hasHydrated: hydrated }), + set: (taskId, messages) => + set((state) => { + // Empty queue → drop the key so the store only holds live queues. + if (messages.length === 0) { + if (!(taskId in state.byTaskId)) { + return state; + } + const { [taskId]: _removed, ...rest } = state.byTaskId; + return { byTaskId: rest }; + } + return { + byTaskId: capTasks({ + ...state.byTaskId, + [taskId]: capMessages(messages), + }), + }; + }), + clear: (taskId) => + set((state) => { + if (!(taskId in state.byTaskId)) { + return state; + } + const { [taskId]: _removed, ...rest } = state.byTaskId; + return { byTaskId: rest }; + }), + clearAll: () => set({ byTaskId: {} }), + }), + { + name: "queued-messages", + storage: electronStorage, + partialize: (state) => ({ byTaskId: state.byTaskId }), + onRehydrateStorage: () => (state, error) => { + if (error) { + useQueuedMessageStore.getState().setHasHydrated(true); + } else { + state?.setHasHydrated(true); + } + }, + }, + ), +); + +export const queuedMessageStoreApi = { + set: (taskId: string, messages: QueuedMessage[]) => + useQueuedMessageStore.getState().set(taskId, messages), + get: (taskId: string): QueuedMessage[] => + useQueuedMessageStore.getState().byTaskId[taskId] ?? [], + clear: (taskId: string) => useQueuedMessageStore.getState().clear(taskId), + clearAll: () => useQueuedMessageStore.getState().clearAll(), + whenHydrated: (): Promise => { + if (useQueuedMessageStore.getState()._hasHydrated) { + return Promise.resolve(); + } + return new Promise((resolve) => { + const unsubscribe = useQueuedMessageStore.subscribe((state) => { + if (state._hasHydrated) { + unsubscribe(); + resolve(); + } + }); + }); + }, +}; diff --git a/packages/ui/src/features/sessions/sessionServiceHost.test.ts b/packages/ui/src/features/sessions/sessionServiceHost.test.ts index a61454e9fb..f9e2e43413 100644 --- a/packages/ui/src/features/sessions/sessionServiceHost.test.ts +++ b/packages/ui/src/features/sessions/sessionServiceHost.test.ts @@ -3784,8 +3784,12 @@ describe("SessionService", () => { }); describe("sendPrompt", () => { - it("throws when offline", async () => { + it("throws when offline for an idle local session", async () => { mockGetIsOnline.mockReturnValue(false); + // A connected, non-busy local session can't run offline. + mockSessionStoreSetters.getSessionByTaskId.mockReturnValue( + createMockSession({ status: "connected", isPromptPending: false }), + ); const service = getSessionService(); await expect(service.sendPrompt("task-123", "Hello")).rejects.toThrow( @@ -3793,6 +3797,27 @@ describe("SessionService", () => { ); }); + it("queues (does not throw) an offline follow-up to a cloud run", async () => { + mockGetIsOnline.mockReturnValue(false); + mockSessionStoreSetters.getSessionByTaskId.mockReturnValue( + createMockSession({ + isCloud: true, + cloudStatus: "in_progress", + status: "connected", + }), + ); + const service = getSessionService(); + + const result = await service.sendPrompt("task-123", "Hello later"); + + expect(result).toEqual({ stopReason: "queued" }); + expect(mockSessionStoreSetters.enqueueMessage).toHaveBeenCalledWith( + "task-123", + "Hello later", + expect.anything(), + ); + }); + it("throws when no session exists", async () => { const service = getSessionService(); mockSessionStoreSetters.getSessionByTaskId.mockReturnValue(undefined); diff --git a/packages/ui/src/features/sessions/sessionServiceHost.ts b/packages/ui/src/features/sessions/sessionServiceHost.ts index 9a57f913e7..1596716d5f 100644 --- a/packages/ui/src/features/sessions/sessionServiceHost.ts +++ b/packages/ui/src/features/sessions/sessionServiceHost.ts @@ -23,6 +23,11 @@ import { fetchAuthState } from "@posthog/ui/features/auth/authQueries"; import { useUsageLimitStore } from "@posthog/ui/features/billing/usageLimitStore"; import { useAddDirectoryDialogStore } from "@posthog/ui/features/folder-picker/addDirectoryDialogStore"; import { NotificationBus } from "@posthog/ui/features/notifications/notifications"; +import { + resetQueuedMessagePersistenceTracking, + startQueuedMessagePersistence, +} from "@posthog/ui/features/sessions/queuedMessagePersistence"; +import { queuedMessageStoreApi } from "@posthog/ui/features/sessions/queuedMessageStore"; import { useSessionAdapterStore } from "@posthog/ui/features/sessions/sessionAdapterStore"; import { getPersistedConfigOptions, @@ -155,6 +160,9 @@ let serviceInstance: SessionService | null = null; export function getSessionService(): SessionService { if (!serviceInstance) { + // Mirror the in-memory queue to durable storage so follow-ups survive a + // restart. Idempotent; the subscription lives for the process lifetime. + startQueuedMessagePersistence(); serviceInstance = new SessionService(buildSessionServiceDeps()); } return serviceInstance; @@ -167,6 +175,12 @@ export function resetSessionService(): void { } sessionStoreSetters.clearAll(); + // clearAll wipes the sessions map wholesale (no per-session mutation the + // persistence subscription would mirror), so explicitly wipe the durable + // queue and its tracking to avoid leaking queued follow-ups across accounts + // / projects. + queuedMessageStoreApi.clearAll(); + resetQueuedMessagePersistenceTracking(); hostClient() .agent.resetAll.mutate()