refactor(cloud): share session orchestration

Generated-By: PostHog Code
Task-Id: 40c57a59-b4e1-4760-8e56-ecd03e9c2f0f

Author: richardsolomou

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 { useQueryClient } from "@tanstack/react-query";
 import * as Haptics from "expo-haptics";
 import { useLocalSearchParams, useRouter } from "expo-router";
@@ -39,17 +43,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 type { Task } from "@/features/tasks/types";
 import {
@@ -70,6 +75,7 @@ import { logger } from "@/lib/logger";
 import { useThemeColors } from "@/lib/theme";
 
 const log = logger.scope("task-detail");
+type QueuedMessage = CloudTaskQueuedMessage<PendingAttachment>;
 
 function getFirstParam(value?: string | string[]): string | undefined {
   return Array.isArray(value) ? value[0] : value;
@@ -108,12 +114,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;
@@ -128,7 +135,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
@@ -177,9 +184,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();
 
@@ -368,11 +373,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;
       }
@@ -394,7 +401,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)
@@ -447,7 +454,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,
@@ -458,16 +465,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],
   );
@@ -476,8 +483,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],
@@ -524,7 +531,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) {
@@ -542,7 +549,7 @@ export default function TaskDetailScreen() {
         })
         .catch(() => {});
     });
-  }, [taskId, stopRun, analytics, getSessionForTask]);
+  }, [taskId, stopRun, analytics]);
 
   const canStopRun =
     !!task &&

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";

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,19 +16,21 @@ 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<PendingAttachment>;
 
 interface QueuedMessagesDockProps {
   taskId: string;
   canSteer: boolean;
   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<string | null>(null);
 
-  if (!queued || queued.length === 0) return null;
+  if (queued.length === 0) return null;
   const active = queued.find((m) => m.id === activeId) ?? null;
 
   return (

apps/mobile/src/features/tasks/hooks/useMessagingMode.ts

@@ -1,18 +1,18 @@
 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 {
   return useMessagingModeStore((s) => s.getEffectiveMode(taskId));
 }
 
 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]);
 }

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<PendingAttachment>[];
+  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],
+  };
+}

apps/mobile/src/features/tasks/index.ts

@@ -18,10 +18,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

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",
+    );
+  });
+});