Improve connection selection and agent hydration reliability

2026-02-14 11:31:57 +07:00
parent 21f191583f
commit c4bebb97ab
11 changed files with 523 additions and 59 deletions
@@ -1,16 +1,25 @@
 import { SessionProvider } from "@/contexts/session-context";
 import {
  useDaemonRegistry,
+  type HostConnection,
  type HostProfile,
 } from "@/contexts/daemon-registry-context";
 import {
  buildDaemonWebSocketUrl,
  buildRelayWebSocketUrl,
 } from "@/utils/daemon-endpoints";
+import { measureConnectionLatency } from "@/utils/test-daemon-connection";
+import {
+  selectBestConnection,
+  type ConnectionCandidate,
+  type ConnectionProbeState,
+} from "@/utils/connection-selection";
+import { useQueries } from "@tanstack/react-query";
 import { useMemo } from "react";
 import type { ActiveConnection } from "@/contexts/daemon-connections-context";

 type Candidate = {
+  connection: HostConnection;
  connectionId: string;
  url: string;
  activeConnection: ActiveConnection;
@@ -29,28 +38,27 @@ function sortConnectionsByPreference<T extends { id: string }>(

 function buildCandidates(host: HostProfile): Candidate[] {
  const preferred = host.preferredConnectionId ?? null;
-
-  const direct = sortConnectionsByPreference(
-    host.connections.filter((c) => c.type === "direct"),
-    preferred
-  );
-  const relay = sortConnectionsByPreference(
-    host.connections.filter((c) => c.type === "relay"),
-    preferred
-  );
+  const connections = sortConnectionsByPreference(host.connections, preferred);

  const out: Candidate[] = [];

-  for (const conn of direct) {
-    out.push({
-      connectionId: conn.id,
-      url: buildDaemonWebSocketUrl(conn.endpoint),
-      activeConnection: { type: "direct", endpoint: conn.endpoint, display: conn.endpoint },
-    });
-  }
+  for (const conn of connections) {
+    if (conn.type === "direct") {
+      out.push({
+        connection: conn,
+        connectionId: conn.id,
+        url: buildDaemonWebSocketUrl(conn.endpoint),
+        activeConnection: {
+          type: "direct",
+          endpoint: conn.endpoint,
+          display: conn.endpoint,
+        },
+      });
+      continue;
+    }

-  for (const conn of relay) {
    out.push({
+      connection: conn,
      connectionId: conn.id,
      url: buildRelayWebSocketUrl({
        endpoint: conn.relayEndpoint,
@@ -66,10 +74,69 @@ function buildCandidates(host: HostProfile): Candidate[] {

 function ManagedDaemonSession({ daemon }: { daemon: HostProfile }) {
  const candidates = useMemo(() => buildCandidates(daemon), [daemon]);
+  const latencyQueries = useQueries({
+    queries: candidates.map((candidate) => ({
+      queryKey: ["connection-selection-latency", daemon.serverId, candidate.connectionId],
+      queryFn: () =>
+        measureConnectionLatency(candidate.connection, {
+          serverId: daemon.serverId,
+        }),
+      refetchInterval: 10_000,
+      staleTime: 9_000,
+      gcTime: 60_000,
+      retry: 1,
+    })),
+  });
+
+  const probeByConnectionId = useMemo(() => {
+    const next = new Map<string, ConnectionProbeState>();
+
+    candidates.forEach((candidate, index) => {
+      const query = latencyQueries[index];
+      if (!query) {
+        next.set(candidate.connectionId, { status: "pending", latencyMs: null });
+        return;
+      }
+
+      if (query.isSuccess && typeof query.data === "number") {
+        next.set(candidate.connectionId, {
+          status: "available",
+          latencyMs: query.data,
+        });
+        return;
+      }
+
+      if (query.isError) {
+        next.set(candidate.connectionId, { status: "unavailable", latencyMs: null });
+        return;
+      }
+
+      next.set(candidate.connectionId, { status: "pending", latencyMs: null });
+    });
+
+    return next;
+  }, [candidates, latencyQueries]);
+
+  const candidateInputs = useMemo<ConnectionCandidate[]>(
+    () =>
+      candidates.map((candidate) => ({
+        connectionId: candidate.connectionId,
+        connection: candidate.connection,
+      })),
+    [candidates]
+  );
+
+  const activeConnectionId = useMemo(
+    () =>
+      selectBestConnection({
+        candidates: candidateInputs,
+        preferredConnectionId: daemon.preferredConnectionId,
+        probeByConnectionId,
+      }),
+    [candidateInputs, daemon.preferredConnectionId, probeByConnectionId]
+  );
  const active =
-    candidates.find((candidate) => candidate.connectionId === daemon.preferredConnectionId) ??
-    candidates[0] ??
-    null;
+    candidates.find((candidate) => candidate.connectionId === activeConnectionId) ?? null;
  const activeUrl = active?.url ?? null;

  if (!activeUrl) {
@@ -996,12 +996,6 @@ export function SessionProvider({
        });
        return;
      }
-      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,
-        });
-      }
    });

    const unsubPermissionRequest = client.on(
@@ -63,6 +63,10 @@ import { useCheckoutStatusQuery } from "@/hooks/use-checkout-status-query";
 import { useAgentInitialization } from "@/hooks/use-agent-initialization";
 import { useToast } from "@/contexts/toast-context";
 import { getInitDeferred, getInitKey } from "@/utils/agent-initialization";
+import {
+  derivePendingPermissionKey,
+  normalizeAgentSnapshot,
+} from "@/utils/agent-snapshots";
 import {
  DropdownMenu,
  DropdownMenuContent,
@@ -369,6 +373,10 @@ function AgentScreenContent({
  const allPendingPermissions = useSessionStore(
    (state) => state.sessions[serverId]?.pendingPermissions
  );
+  const setAgents = useSessionStore((state) => state.setAgents);
+  const setPendingPermissions = useSessionStore(
+    (state) => state.setPendingPermissions
+  );
  const pendingPermissions = useMemo(() => {
    if (!allPendingPermissions || !resolvedAgentId) return new Map();
    const filtered = new Map();
@@ -662,7 +670,46 @@ function AgentScreenContent({
    const attemptToken = ++initAttemptTokenRef.current;

    ensureAgentIsInitialized(resolvedAgentId)
-      .then(() => {
+      .then(async () => {
+        if (attemptToken !== initAttemptTokenRef.current) {
+          return;
+        }
+        const currentAgent = useSessionStore
+          .getState()
+          .sessions[serverId]
+          ?.agents.get(resolvedAgentId);
+        if (!currentAgent && client) {
+          const snapshot = await client.fetchAgent(resolvedAgentId);
+          if (attemptToken !== initAttemptTokenRef.current) {
+            return;
+          }
+          if (!snapshot) {
+            setMissingAgentState({
+              kind: "not_found",
+              message: `Agent not found: ${resolvedAgentId}`,
+            });
+            return;
+          }
+          const normalized = normalizeAgentSnapshot(snapshot, serverId);
+          setAgents(serverId, (prev) => {
+            const next = new Map(prev);
+            next.set(normalized.id, normalized);
+            return next;
+          });
+          setPendingPermissions(serverId, (prev) => {
+            const next = new Map(prev);
+            for (const [key, pending] of next.entries()) {
+              if (pending.agentId === normalized.id) {
+                next.delete(key);
+              }
+            }
+            for (const request of normalized.pendingPermissions) {
+              const key = derivePendingPermissionKey(normalized.id, request);
+              next.set(key, { key, agentId: normalized.id, request });
+            }
+            return next;
+          });
+        }
        if (attemptToken !== initAttemptTokenRef.current) {
          return;
        }
@@ -681,10 +728,14 @@ function AgentScreenContent({
      });
  }, [
    agent,
+    client,
    ensureAgentIsInitialized,
    isConnected,
    missingAgentState.kind,
    resolvedAgentId,
+    serverId,
+    setAgents,
+    setPendingPermissions,
    shouldUseOptimisticStream,
  ]);

@@ -754,6 +805,17 @@ function AgentScreenContent({
        </View>
      );
    }
+    if (missingAgentState.kind === "error") {
+      return (
+        <View style={styles.container} testID="agent-load-error">
+          <MenuHeader title="Agent" />
+          <View style={styles.errorContainer}>
+            <Text style={styles.errorText}>Failed to load agent</Text>
+            <Text style={styles.statusText}>{missingAgentState.message}</Text>
+          </View>
+        </View>
+      );
+    }

    return (
      <View style={styles.container} testID="agent-loading">
@@ -0,0 +1,113 @@
+import { describe, expect, it } from "vitest";
+import { type HostConnection } from "@/contexts/daemon-registry-context";
+import {
+  selectBestConnection,
+  type ConnectionCandidate,
+  type ConnectionProbeState,
+} from "./connection-selection";
+
+function makeDirect(id: string, endpoint: string): HostConnection {
+  return { id, type: "direct", endpoint };
+}
+
+function makeRelay(
+  id: string,
+  relayEndpoint: string,
+  daemonPublicKeyB64 = "abc"
+): HostConnection {
+  return { id, type: "relay", relayEndpoint, daemonPublicKeyB64 };
+}
+
+function probes(
+  input: Record<string, ConnectionProbeState>
+): Map<string, ConnectionProbeState> {
+  return new Map(Object.entries(input));
+}
+
+describe("selectBestConnection", () => {
+  it("picks the available connection with lowest latency regardless of transport type", () => {
+    const candidates: ConnectionCandidate[] = [
+      { connectionId: "direct:a", connection: makeDirect("direct:a", "a:6767") },
+      {
+        connectionId: "relay:b",
+        connection: makeRelay("relay:b", "relay.example:443"),
+      },
+    ];
+
+    const selected = selectBestConnection({
+      candidates,
+      preferredConnectionId: "direct:a",
+      probeByConnectionId: probes({
+        "direct:a": { status: "available", latencyMs: 84 },
+        "relay:b": { status: "available", latencyMs: 34 },
+      }),
+    });
+
+    expect(selected).toBe("relay:b");
+  });
+
+  it("ignores unavailable and pending probes", () => {
+    const candidates: ConnectionCandidate[] = [
+      { connectionId: "direct:a", connection: makeDirect("direct:a", "a:6767") },
+      {
+        connectionId: "relay:b",
+        connection: makeRelay("relay:b", "relay.example:443"),
+      },
+      { connectionId: "direct:c", connection: makeDirect("direct:c", "c:6767") },
+    ];
+
+    const selected = selectBestConnection({
+      candidates,
+      preferredConnectionId: "direct:a",
+      probeByConnectionId: probes({
+        "direct:a": { status: "pending", latencyMs: null },
+        "relay:b": { status: "unavailable", latencyMs: null },
+        "direct:c": { status: "available", latencyMs: 41 },
+      }),
+    });
+
+    expect(selected).toBe("direct:c");
+  });
+
+  it("falls back to preferred connection when no candidates are available", () => {
+    const candidates: ConnectionCandidate[] = [
+      { connectionId: "direct:a", connection: makeDirect("direct:a", "a:6767") },
+      {
+        connectionId: "relay:b",
+        connection: makeRelay("relay:b", "relay.example:443"),
+      },
+    ];
+
+    const selected = selectBestConnection({
+      candidates,
+      preferredConnectionId: "relay:b",
+      probeByConnectionId: probes({
+        "direct:a": { status: "pending", latencyMs: null },
+        "relay:b": { status: "unavailable", latencyMs: null },
+      }),
+    });
+
+    expect(selected).toBe("relay:b");
+  });
+
+  it("falls back to first candidate when preferred is missing and none are available", () => {
+    const candidates: ConnectionCandidate[] = [
+      { connectionId: "direct:a", connection: makeDirect("direct:a", "a:6767") },
+      {
+        connectionId: "relay:b",
+        connection: makeRelay("relay:b", "relay.example:443"),
+      },
+    ];
+
+    const selected = selectBestConnection({
+      candidates,
+      preferredConnectionId: "missing",
+      probeByConnectionId: probes({
+        "direct:a": { status: "unavailable", latencyMs: null },
+        "relay:b": { status: "pending", latencyMs: null },
+      }),
+    });
+
+    expect(selected).toBe("direct:a");
+  });
+});
@@ -0,0 +1,65 @@
+import type { HostConnection } from "@/contexts/daemon-registry-context";
+
+export type ConnectionCandidate = {
+  connectionId: string;
+  connection: HostConnection;
+};
+
+export type ConnectionProbeState =
+  | { status: "pending"; latencyMs: null }
+  | { status: "unavailable"; latencyMs: null }
+  | { status: "available"; latencyMs: number };
+
+export type SelectBestConnectionInput = {
+  candidates: ConnectionCandidate[];
+  preferredConnectionId: string | null;
+  probeByConnectionId: Map<string, ConnectionProbeState>;
+};
+
+export function selectBestConnection(
+  input: SelectBestConnectionInput
+): string | null {
+  const { candidates, preferredConnectionId, probeByConnectionId } = input;
+  if (candidates.length === 0) {
+    return null;
+  }
+
+  const available = candidates
+    .map((candidate, index) => ({ candidate, index }))
+    .filter(({ candidate }) => {
+      const probe = probeByConnectionId.get(candidate.connectionId);
+      return probe?.status === "available";
+    })
+    .sort((left, right) => {
+      const leftProbe = probeByConnectionId.get(left.candidate.connectionId);
+      const rightProbe = probeByConnectionId.get(right.candidate.connectionId);
+
+      const leftLatency =
+        leftProbe && leftProbe.status === "available"
+          ? leftProbe.latencyMs
+          : Number.POSITIVE_INFINITY;
+      const rightLatency =
+        rightProbe && rightProbe.status === "available"
+          ? rightProbe.latencyMs
+          : Number.POSITIVE_INFINITY;
+
+      if (leftLatency === rightLatency) {
+        return left.index - right.index;
+      }
+
+      return leftLatency - rightLatency;
+    });
+
+  if (available.length > 0) {
+    return available[0]!.candidate.connectionId;
+  }
+
+  const preferred = candidates.find(
+    (candidate) => candidate.connectionId === preferredConnectionId
+  );
+  if (preferred) {
+    return preferred.connectionId;
+  }
+
+  return candidates[0]!.connectionId;
+}