From b36f870ee93f54d2649073bd59dab38b8f7d17de Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Tue, 2 Dec 2025 03:57:36 +0000 Subject: [PATCH] Add RPC helper for agent initialization --- .../src/app/agent/[serverId]/[agentId].tsx | 6 + packages/app/src/contexts/session-context.tsx | 140 ++++++++++-------- packages/app/src/hooks/use-session-rpc.ts | 88 +++++++++++ packages/server/src/server/messages.ts | 14 ++ packages/server/src/server/session.ts | 30 ++-- 5 files changed, 202 insertions(+), 76 deletions(-) create mode 100644 packages/app/src/hooks/use-session-rpc.ts diff --git a/packages/app/src/app/agent/[serverId]/[agentId].tsx b/packages/app/src/app/agent/[serverId]/[agentId].tsx index edd0ba7ec..774a044d4 100644 --- a/packages/app/src/app/agent/[serverId]/[agentId].tsx +++ b/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; } diff --git a/packages/app/src/contexts/session-context.tsx b/packages/app/src/contexts/session-context.tsx index 7ee426955..77747b5ac 100644 --- a/packages/app/src/contexts/session-context.tsx +++ b/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; + 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>(new Set()); const previousAgentStatusRef = useRef>(new Map()); const providerModelRequestIdsRef = useRef>(new Map()); - const pendingAgentLifecycleRequestsRef = useRef([]); 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( - (request: PendingAgentLifecycleRequest) => { - const { kind, params } = 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]); + const initializeAgentRpc = useSessionRpc({ + ws, + type: "initialize_agent_request", + }); 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) { - pendingAgentLifecycleRequestsRef.current.push({ - kind: "initialize", - params: { agentId, requestId }, + initializeAgentRpc + .send({ agentId }) + .catch((error) => { + console.warn("[Session] initializeAgent failed", { agentId, error }); + setInitializingAgents((prev) => { + const next = new Map(prev); + next.set(agentId, false); + return next; + }); }); - return; - } - - sendAgentLifecycleRequest({ - kind: "initialize", - params: { agentId, requestId }, - }); - }, [wsIsConnected, sendAgentLifecycleRequest]); + }, [initializeAgentRpc, setAgentStreamState, setInitializingAgents]); 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) { - pendingAgentLifecycleRequestsRef.current.push({ - kind: "refresh", - params: { agentId, requestId }, + refreshAgentRequest + .execute({ agentId }, { requestKeyOverride: agentId, dedupe: false }) + .catch((error) => { + console.warn("[Session] refreshAgent failed", { agentId, error }); + setInitializingAgents((prev) => { + const next = new Map(prev); + next.set(agentId, false); + return next; + }); }); - return; - } - - sendAgentLifecycleRequest({ - kind: "refresh", - params: { agentId, requestId }, - }); - }, [wsIsConnected, sendAgentLifecycleRequest]); + }, [refreshAgentRequest, setAgentStreamState, setInitializingAgents]); const requestProviderModels = useCallback((provider: AgentProvider, options?: { cwd?: string }) => { const requestId = generateMessageId(); diff --git a/packages/app/src/hooks/use-session-rpc.ts b/packages/app/src/hooks/use-session-rpc.ts new file mode 100644 index 000000000..e90c2963f --- /dev/null +++ b/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 = Extract; +type ResponseOf = Extract; + +type RpcState = + | { status: "idle"; requestId: null } + | { status: "loading"; requestId: string } + | { status: "success"; requestId: string; data: T } + | { status: "error"; requestId: string | null; error: Error }; + +type ResponseWithRequestId = ResponseOf & { + payload: { requestId?: string; error?: string }; +}; + +type SelectResponse = (message: ResponseWithRequestId) => TData; + +export function useSessionRpc["payload"]>(options: { + ws: UseWebSocketReturn; + type: TType; + select?: SelectResponse; +}) { + const { ws, type, select } = options; + const [state, setState] = useState>({ status: "idle", requestId: null }); + const activeRequestIdRef = useRef(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).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) : ((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, "type" | "requestId">) => { + const requestId = generateMessageId(); + activeRequestIdRef.current = requestId; + setState({ status: "loading", requestId }); + + const request = { + type, + ...params, + requestId, + } as RequestOf; + + ws.send({ type: "session", message: request }); + + return new Promise((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]); +} diff --git a/packages/server/src/server/messages.ts b/packages/server/src/server/messages.ts index 07da81618..e0f576c5e 100644 --- a/packages/server/src/server/messages.ts +++ b/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; +export type InitializeAgentResponseMessage = z.infer; // Type exports for payload types export type ActivityLogPayload = z.infer; diff --git a/packages/server/src/server/session.ts b/packages/server/src/server/session.ts index 97c763bfb..f2cd79ec2 100644 --- a/packages/server/src/server/session.ts +++ b/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({ - type: "status", - payload: { - status: "agent_initialized", - agentId, - agentStatus: snapshot.lifecycle, - requestId, - timelineSize, - }, - }); - } + this.emit({ + type: "initialize_agent_request", + payload: { + agentId, + agentStatus: snapshot.lifecycle, + timelineSize, + requestId, + }, + }); 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(), - timestamp: new Date(), - type: "error", - content: `Failed to initialize agent: ${error.message}`, + agentId, + requestId, + error: error?.message ?? "Failed to initialize agent", }, }); }