fix(sessions): persist queued follow-ups across restart and offline

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

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}

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<string>();
+const lastQueueByTask = new Map<string, QueuedMessage[]>();
+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<string>();
+    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();
+}

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({});
+  });
+});