Add RPC helper for agent initialization

packages/app/src/app/agent/[serverId]/[agentId].tsx

@@ -259,6 +259,12 @@ function AgentScreenContent({ session, agentId, routeServerId, onBack }: AgentSc
       return;
     }
 
+    console.log("[AgentScreen] initializeAgent effect", {
+      agentId: resolvedAgentId,
+      hasStreamState,
+      initializationState,
+    });
+
     if (initializationState !== undefined) {
       return;
     }

packages/app/src/contexts/session-context.tsx

@@ -2,6 +2,7 @@ import { createContext, useState, useRef, ReactNode, useCallback, useEffect, use
 import AsyncStorage from "@react-native-async-storage/async-storage";
 import { useWebSocket, type UseWebSocketReturn } from "@/hooks/use-websocket";
 import { useDaemonRequest } from "@/hooks/use-daemon-request";
+import { useSessionRpc } from "@/hooks/use-session-rpc";
 import { useAudioPlayer } from "@/hooks/use-audio-player";
 import { reduceStreamUpdate, generateMessageId, hydrateStreamState, type StreamItem } from "@/types/stream";
 import type { PendingPermission } from "@/types/shared";
@@ -58,6 +59,8 @@ type FileExplorerResponseMessage = Extract<
   { type: "file_explorer_response" }
 >;
 
+type StatusMessage = Extract<SessionOutboundMessage, { type: "status" }>;
+
 export type MessageEntry =
   | {
       type: "user";
@@ -417,6 +420,13 @@ export function SessionProvider({ children, serverUrl, serverId }: SessionProvid
     updateConnectionStatus(serverId, { status: "offline" });
   }, [serverId, updateConnectionStatus, ws.isConnected, ws.isConnecting, ws.lastError]);
 
+  // If the socket drops mid-initialization, clear pending flags so screens don't spin forever
+  useEffect(() => {
+    if (!ws.isConnected) {
+      setInitializingAgents(new Map());
+    }
+  }, [ws.isConnected]);
+
   useEffect(() => {
     return () => {
       updateConnectionStatus(serverId, { status: "offline", lastError: null });
@@ -462,7 +472,6 @@ export function SessionProvider({ children, serverUrl, serverId }: SessionProvid
   const activeAudioGroupsRef = useRef<Set<string>>(new Set());
   const previousAgentStatusRef = useRef<Map<string, AgentLifecycleStatus>>(new Map());
   const providerModelRequestIdsRef = useRef<Map<AgentProvider, string>>(new Map());
-  const pendingAgentLifecycleRequestsRef = useRef<PendingAgentLifecycleRequest[]>([]);
   const hasHydratedSnapshotRef = useRef(false);
   const hasRequestedInitialSnapshotRef = useRef(false);
 
@@ -601,6 +610,38 @@ export function SessionProvider({ children, serverUrl, serverId }: SessionProvid
     keepPreviousData: false,
   });
 
+  const refreshAgentRequest = useDaemonRequest<
+    { agentId: string },
+    { agentId: string; lifecycle: AgentLifecycleStatus | undefined },
+    StatusMessage
+  >({
+    ws,
+    responseType: "status",
+    buildRequest: ({ params, requestId }) => ({
+      type: "session",
+      message: {
+        type: "refresh_agent_request",
+        agentId: params?.agentId ?? "",
+        requestId,
+      },
+    }),
+    matchResponse: (message, context) =>
+      message.payload.status === "agent_initialized" &&
+      (message.payload as { agentId?: string }).agentId === context.params?.agentId &&
+      (message.payload as { requestId?: string }).requestId === context.requestId,
+    getRequestKey: (params) => params?.agentId ?? "default",
+    selectData: (message) => ({
+      agentId: (message.payload as { agentId?: string }).agentId ?? "",
+      lifecycle: (message.payload as { agentStatus?: AgentLifecycleStatus }).agentStatus,
+    }),
+    extractError: (message) =>
+      message.payload.status === "error"
+        ? new Error((message.payload as { message?: string }).message ?? "Refresh failed")
+        : null,
+    timeoutMs: 15000,
+    keepPreviousData: false,
+  });
+
   const directoryListingRequest = useDaemonRequest<
     { agentId: string; path: string },
     FileExplorerResponseMessage["payload"],
@@ -659,40 +700,10 @@ export function SessionProvider({ children, serverUrl, serverId }: SessionProvid
     keepPreviousData: false,
   });
 
-  const sendAgentLifecycleRequest = useCallback(
+  const initializeAgentRpc = useSessionRpc({
-    (request: PendingAgentLifecycleRequest) => {
+    ws,
-      const { kind, params } = request;
+    type: "initialize_agent_request",
-      const messageType = kind === "initialize" ? "initialize_agent_request" : "refresh_agent_request";
+  });
-
-      const msg: WSInboundMessage = {
-        type: "session",
-        message: {
-          type: messageType,
-          agentId: params.agentId,
-          requestId: params.requestId,
-        },
-      };
-      ws.send(msg);
-    },
-    [ws]
-  );
-
-  useEffect(() => {
-    if (!wsIsConnected) {
-      return;
-    }
-
-    if (pendingAgentLifecycleRequestsRef.current.length === 0) {
-      return;
-    }
-
-    const pending = [...pendingAgentLifecycleRequestsRef.current];
-    pendingAgentLifecycleRequestsRef.current = [];
-
-    for (const request of pending) {
-      sendAgentLifecycleRequest(request);
-    }
-  }, [wsIsConnected, sendAgentLifecycleRequest]);
 
   useEffect(() => {
     if (!wsIsConnected) {
@@ -720,6 +731,7 @@ export function SessionProvider({ children, serverUrl, serverId }: SessionProvid
       const { agents: agentsList, commands: commandsList } = message.payload;
 
       console.log("[Session] Session state:", agentsList.length, "agents,", commandsList.length, "commands");
+      setInitializingAgents(new Map());
 
       const { agents: hydratedAgents, pendingPermissions: hydratedPermissions } = buildSessionStateFromSnapshots(
         serverId,
@@ -804,6 +816,8 @@ export function SessionProvider({ children, serverUrl, serverId }: SessionProvid
       const { agentId, event, timestamp } = message.payload;
       const parsedTimestamp = new Date(timestamp);
 
+      console.log("[Session] agent_stream", { agentId, eventType: event.type });
+
       setAgentStreamState((prev) => {
         const currentStream = prev.get(agentId) || [];
         const newStream = reduceStreamUpdate(
@@ -845,6 +859,11 @@ export function SessionProvider({ children, serverUrl, serverId }: SessionProvid
       if (message.type !== "agent_stream_snapshot") return;
       const { agentId, events } = message.payload;
 
+      console.log("[Session] agent_stream_snapshot", {
+        agentId,
+        eventCount: events.length,
+      });
+
       const hydrated = hydrateStreamState(
         events.map(({ event, timestamp }) => ({
           event: event as AgentStreamEventPayload,
@@ -871,6 +890,12 @@ export function SessionProvider({ children, serverUrl, serverId }: SessionProvid
     const unsubStatus = ws.on("status", (message) => {
       if (message.type !== "status") return;
       const status = message.payload.status;
+      if (status === "agent_initialized" && "agentId" in message.payload) {
+        console.log("[Session] status agent_initialized", {
+          agentId: (message.payload as any).agentId,
+          requestId: (message.payload as any).requestId,
+        });
+      }
       if (status === "agent_initialized" && "agentId" in message.payload) {
         const agentId = (message.payload as any).agentId as string | undefined;
         if (agentId) {
@@ -1290,6 +1315,7 @@ export function SessionProvider({ children, serverUrl, serverId }: SessionProvid
   }, [ws, audioPlayer, setIsPlayingAudio, updateExplorerState]);
 
   const initializeAgent = useCallback(({ agentId, requestId }: { agentId: string; requestId?: string }) => {
+    console.log("[Session] initializeAgent called", { agentId, requestId });
     setInitializingAgents((prev) => {
       const next = new Map(prev);
       next.set(agentId, true);
@@ -1302,19 +1328,17 @@ export function SessionProvider({ children, serverUrl, serverId }: SessionProvid
       return next;
     });
 
-    if (!wsIsConnected) {
+    initializeAgentRpc
-      pendingAgentLifecycleRequestsRef.current.push({
+      .send({ agentId })
-        kind: "initialize",
+      .catch((error) => {
-        params: { agentId, requestId },
+        console.warn("[Session] initializeAgent failed", { agentId, error });
+        setInitializingAgents((prev) => {
+          const next = new Map(prev);
+          next.set(agentId, false);
+          return next;
+        });
       });
-      return;
+  }, [initializeAgentRpc, setAgentStreamState, setInitializingAgents]);
-    }
-
-    sendAgentLifecycleRequest({
-      kind: "initialize",
-      params: { agentId, requestId },
-    });
-  }, [wsIsConnected, sendAgentLifecycleRequest]);
 
   const refreshAgent = useCallback(({ agentId, requestId }: { agentId: string; requestId?: string }) => {
     setInitializingAgents((prev) => {
@@ -1329,19 +1353,17 @@ export function SessionProvider({ children, serverUrl, serverId }: SessionProvid
       return next;
     });
 
-    if (!wsIsConnected) {
+    refreshAgentRequest
-      pendingAgentLifecycleRequestsRef.current.push({
+      .execute({ agentId }, { requestKeyOverride: agentId, dedupe: false })
-        kind: "refresh",
+      .catch((error) => {
-        params: { agentId, requestId },
+        console.warn("[Session] refreshAgent failed", { agentId, error });
+        setInitializingAgents((prev) => {
+          const next = new Map(prev);
+          next.set(agentId, false);
+          return next;
+        });
       });
-      return;
+  }, [refreshAgentRequest, setAgentStreamState, setInitializingAgents]);
-    }
-
-    sendAgentLifecycleRequest({
-      kind: "refresh",
-      params: { agentId, requestId },
-    });
-  }, [wsIsConnected, sendAgentLifecycleRequest]);
 
   const requestProviderModels = useCallback((provider: AgentProvider, options?: { cwd?: string }) => {
     const requestId = generateMessageId();

packages/app/src/hooks/use-session-rpc.ts

@@ -0,0 +1,88 @@
+import { useCallback, useEffect, useMemo, useRef, useState } from "react";
+import type { SessionInboundMessage, SessionOutboundMessage } from "@server/server/messages";
+import type { UseWebSocketReturn } from "./use-websocket";
+import { generateMessageId } from "@/types/stream";
+
+type SharedType = SessionInboundMessage["type"] & SessionOutboundMessage["type"];
+
+type RequestOf<TType extends SharedType> = Extract<SessionInboundMessage, { type: TType }>;
+type ResponseOf<TType extends SharedType> = Extract<SessionOutboundMessage, { type: TType }>;
+
+type RpcState<T> =
+  | { status: "idle"; requestId: null }
+  | { status: "loading"; requestId: string }
+  | { status: "success"; requestId: string; data: T }
+  | { status: "error"; requestId: string | null; error: Error };
+
+type ResponseWithRequestId<TType extends SharedType> = ResponseOf<TType> & {
+  payload: { requestId?: string; error?: string };
+};
+
+type SelectResponse<TType extends SharedType, TData> = (message: ResponseWithRequestId<TType>) => TData;
+
+export function useSessionRpc<TType extends SharedType, TData = ResponseWithRequestId<TType>["payload"]>(options: {
+  ws: UseWebSocketReturn;
+  type: TType;
+  select?: SelectResponse<TType, TData>;
+}) {
+  const { ws, type, select } = options;
+  const [state, setState] = useState<RpcState<TData>>({ status: "idle", requestId: null });
+  const activeRequestIdRef = useRef<string | null>(null);
+  const resolveRef = useRef<((value: TData) => void) | null>(null);
+  const rejectRef = useRef<((reason?: any) => void) | null>(null);
+
+  useEffect(() => {
+    const unsubscribe = ws.on(type, (message) => {
+      const payload = (message as ResponseWithRequestId<TType>).payload;
+      if (!payload || payload.requestId !== activeRequestIdRef.current) {
+        return;
+      }
+
+      if (payload.error) {
+        const error = new Error(payload.error);
+        setState({ status: "error", requestId: payload.requestId ?? null, error });
+        rejectRef.current?.(error);
+        activeRequestIdRef.current = null;
+        return;
+      }
+
+      const data = select ? select(message as ResponseWithRequestId<TType>) : ((payload as unknown) as TData);
+      setState({ status: "success", requestId: payload.requestId ?? null, data });
+      resolveRef.current?.(data);
+      activeRequestIdRef.current = null;
+    });
+
+    return () => {
+      unsubscribe();
+    };
+  }, [select, type, ws]);
+
+  const send = useCallback(
+    (params: Omit<RequestOf<TType>, "type" | "requestId">) => {
+      const requestId = generateMessageId();
+      activeRequestIdRef.current = requestId;
+      setState({ status: "loading", requestId });
+
+      const request = {
+        type,
+        ...params,
+        requestId,
+      } as RequestOf<TType>;
+
+      ws.send({ type: "session", message: request });
+
+      return new Promise<TData>((resolve, reject) => {
+        resolveRef.current = resolve;
+        rejectRef.current = reject;
+      });
+    },
+    [type, ws]
+  );
+
+  const reset = useCallback(() => {
+    activeRequestIdRef.current = null;
+    setState({ status: "idle", requestId: null });
+  }, []);
+
+  return useMemo(() => ({ state, send, reset }), [reset, send, state]);
+}

packages/server/src/server/messages.ts

@@ -379,6 +379,18 @@ export const InitializeAgentRequestMessageSchema = z.object({
   requestId: z.string().optional(),
 });
 
+export const InitializeAgentResponseMessageSchema = z.object({
+  // Response shares the same type literal as the request for simple requestId matching
+  type: z.literal("initialize_agent_request"),
+  payload: z.object({
+    agentId: z.string(),
+    agentStatus: z.string().optional(),
+    timelineSize: z.number().optional(),
+    requestId: z.string().optional(),
+    error: z.string().optional(),
+  }),
+});
+
 export const SetAgentModeMessageSchema = z.object({
   type: z.literal("set_agent_mode"),
   agentId: z.string(),
@@ -717,6 +729,7 @@ export const SessionOutboundMessageSchema = z.discriminatedUnion("type", [
   AudioOutputMessageSchema,
   TranscriptionResultMessageSchema,
   StatusMessageSchema,
+  InitializeAgentResponseMessageSchema,
   ArtifactMessageSchema,
   ConversationLoadedMessageSchema,
   AgentStateMessageSchema,
@@ -764,6 +777,7 @@ export type ListPersistedAgentsResponseMessage = z.infer<typeof ListPersistedAge
 export type ListProviderModelsResponseMessage = z.infer<
   typeof ListProviderModelsResponseMessageSchema
 >;
+export type InitializeAgentResponseMessage = z.infer<typeof InitializeAgentResponseMessageSchema>;
 
 // Type exports for payload types
 export type ActivityLogPayload = z.infer<typeof ActivityLogPayloadSchema>;

packages/server/src/server/session.ts

@@ -1183,18 +1183,15 @@ export class Session {
       // Send timeline snapshot after hydration (if any)
       const timelineSize = this.emitAgentTimelineSnapshot(snapshot);
 
-      if (requestId) {
+      this.emit({
-        this.emit({
+        type: "initialize_agent_request",
-          type: "status",
+        payload: {
-          payload: {
+          agentId,
-            status: "agent_initialized",
+          agentStatus: snapshot.lifecycle,
-            agentId,
+          timelineSize,
-            agentStatus: snapshot.lifecycle,
+          requestId,
-            requestId,
+        },
-            timelineSize,
+      });
-          },
-        });
-      }
 
       console.log(
         `[Session ${this.clientId}] Agent ${agentId} initialized with ${timelineSize} timeline item(s); status=${snapshot.lifecycle}`
@@ -1205,12 +1202,11 @@ export class Session {
         error
       );
       this.emit({
-        type: "activity_log",
+        type: "initialize_agent_request",
         payload: {
-          id: uuidv4(),
+          agentId,
-          timestamp: new Date(),
+          requestId,
-          type: "error",
+          error: error?.message ?? "Failed to initialize agent",
-          content: `Failed to initialize agent: ${error.message}`,
         },
       });
     }