Improve connection selection and agent hydration reliability
This commit is contained in:
@@ -0,0 +1,34 @@
|
|||||||
|
import { test, expect } from './fixtures';
|
||||||
|
import { createAgent, ensureHostSelected, gotoHome, setWorkingDirectory } from './helpers/app';
|
||||||
|
import { createTempGitRepo } from './helpers/workspace';
|
||||||
|
|
||||||
|
test('agent timeline hydrates after reload via fetch_agent_timeline_request', async ({ page }) => {
|
||||||
|
const repo = await createTempGitRepo();
|
||||||
|
const marker = 'TIMELINE_HYDRATION_OK';
|
||||||
|
const prompt = `Respond with exactly: ${marker}`;
|
||||||
|
|
||||||
|
try {
|
||||||
|
await gotoHome(page);
|
||||||
|
await setWorkingDirectory(page, repo.path);
|
||||||
|
await ensureHostSelected(page);
|
||||||
|
await createAgent(page, prompt);
|
||||||
|
|
||||||
|
const assistantMessage = page
|
||||||
|
.getByTestId('assistant-message')
|
||||||
|
.filter({ hasText: marker })
|
||||||
|
.first();
|
||||||
|
await expect(assistantMessage).toBeVisible({ timeout: 120000 });
|
||||||
|
|
||||||
|
await page.reload({ waitUntil: 'commit' });
|
||||||
|
await expect(page).toHaveURL(/\/agent\//);
|
||||||
|
await expect(page.getByTestId('agent-loading')).toHaveCount(0, { timeout: 30000 });
|
||||||
|
await expect(page.getByText(prompt, { exact: true }).first()).toBeVisible({
|
||||||
|
timeout: 30000,
|
||||||
|
});
|
||||||
|
await expect(
|
||||||
|
page.getByTestId('assistant-message').filter({ hasText: marker }).first()
|
||||||
|
).toBeVisible({ timeout: 30000 });
|
||||||
|
} finally {
|
||||||
|
await repo.cleanup();
|
||||||
|
}
|
||||||
|
});
|
||||||
@@ -1,16 +1,25 @@
|
|||||||
import { SessionProvider } from "@/contexts/session-context";
|
import { SessionProvider } from "@/contexts/session-context";
|
||||||
import {
|
import {
|
||||||
useDaemonRegistry,
|
useDaemonRegistry,
|
||||||
|
type HostConnection,
|
||||||
type HostProfile,
|
type HostProfile,
|
||||||
} from "@/contexts/daemon-registry-context";
|
} from "@/contexts/daemon-registry-context";
|
||||||
import {
|
import {
|
||||||
buildDaemonWebSocketUrl,
|
buildDaemonWebSocketUrl,
|
||||||
buildRelayWebSocketUrl,
|
buildRelayWebSocketUrl,
|
||||||
} from "@/utils/daemon-endpoints";
|
} 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 { useMemo } from "react";
|
||||||
import type { ActiveConnection } from "@/contexts/daemon-connections-context";
|
import type { ActiveConnection } from "@/contexts/daemon-connections-context";
|
||||||
|
|
||||||
type Candidate = {
|
type Candidate = {
|
||||||
|
connection: HostConnection;
|
||||||
connectionId: string;
|
connectionId: string;
|
||||||
url: string;
|
url: string;
|
||||||
activeConnection: ActiveConnection;
|
activeConnection: ActiveConnection;
|
||||||
@@ -29,28 +38,27 @@ function sortConnectionsByPreference<T extends { id: string }>(
|
|||||||
|
|
||||||
function buildCandidates(host: HostProfile): Candidate[] {
|
function buildCandidates(host: HostProfile): Candidate[] {
|
||||||
const preferred = host.preferredConnectionId ?? null;
|
const preferred = host.preferredConnectionId ?? null;
|
||||||
|
const connections = sortConnectionsByPreference(host.connections, preferred);
|
||||||
const direct = sortConnectionsByPreference(
|
|
||||||
host.connections.filter((c) => c.type === "direct"),
|
|
||||||
preferred
|
|
||||||
);
|
|
||||||
const relay = sortConnectionsByPreference(
|
|
||||||
host.connections.filter((c) => c.type === "relay"),
|
|
||||||
preferred
|
|
||||||
);
|
|
||||||
|
|
||||||
const out: Candidate[] = [];
|
const out: Candidate[] = [];
|
||||||
|
|
||||||
for (const conn of direct) {
|
for (const conn of connections) {
|
||||||
out.push({
|
if (conn.type === "direct") {
|
||||||
connectionId: conn.id,
|
out.push({
|
||||||
url: buildDaemonWebSocketUrl(conn.endpoint),
|
connection: conn,
|
||||||
activeConnection: { type: "direct", endpoint: conn.endpoint, display: conn.endpoint },
|
connectionId: conn.id,
|
||||||
});
|
url: buildDaemonWebSocketUrl(conn.endpoint),
|
||||||
}
|
activeConnection: {
|
||||||
|
type: "direct",
|
||||||
|
endpoint: conn.endpoint,
|
||||||
|
display: conn.endpoint,
|
||||||
|
},
|
||||||
|
});
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
for (const conn of relay) {
|
|
||||||
out.push({
|
out.push({
|
||||||
|
connection: conn,
|
||||||
connectionId: conn.id,
|
connectionId: conn.id,
|
||||||
url: buildRelayWebSocketUrl({
|
url: buildRelayWebSocketUrl({
|
||||||
endpoint: conn.relayEndpoint,
|
endpoint: conn.relayEndpoint,
|
||||||
@@ -66,10 +74,69 @@ function buildCandidates(host: HostProfile): Candidate[] {
|
|||||||
|
|
||||||
function ManagedDaemonSession({ daemon }: { daemon: HostProfile }) {
|
function ManagedDaemonSession({ daemon }: { daemon: HostProfile }) {
|
||||||
const candidates = useMemo(() => buildCandidates(daemon), [daemon]);
|
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 =
|
const active =
|
||||||
candidates.find((candidate) => candidate.connectionId === daemon.preferredConnectionId) ??
|
candidates.find((candidate) => candidate.connectionId === activeConnectionId) ?? null;
|
||||||
candidates[0] ??
|
|
||||||
null;
|
|
||||||
const activeUrl = active?.url ?? null;
|
const activeUrl = active?.url ?? null;
|
||||||
|
|
||||||
if (!activeUrl) {
|
if (!activeUrl) {
|
||||||
|
|||||||
@@ -996,12 +996,6 @@ export function SessionProvider({
|
|||||||
});
|
});
|
||||||
return;
|
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(
|
const unsubPermissionRequest = client.on(
|
||||||
|
|||||||
@@ -63,6 +63,10 @@ import { useCheckoutStatusQuery } from "@/hooks/use-checkout-status-query";
|
|||||||
import { useAgentInitialization } from "@/hooks/use-agent-initialization";
|
import { useAgentInitialization } from "@/hooks/use-agent-initialization";
|
||||||
import { useToast } from "@/contexts/toast-context";
|
import { useToast } from "@/contexts/toast-context";
|
||||||
import { getInitDeferred, getInitKey } from "@/utils/agent-initialization";
|
import { getInitDeferred, getInitKey } from "@/utils/agent-initialization";
|
||||||
|
import {
|
||||||
|
derivePendingPermissionKey,
|
||||||
|
normalizeAgentSnapshot,
|
||||||
|
} from "@/utils/agent-snapshots";
|
||||||
import {
|
import {
|
||||||
DropdownMenu,
|
DropdownMenu,
|
||||||
DropdownMenuContent,
|
DropdownMenuContent,
|
||||||
@@ -369,6 +373,10 @@ function AgentScreenContent({
|
|||||||
const allPendingPermissions = useSessionStore(
|
const allPendingPermissions = useSessionStore(
|
||||||
(state) => state.sessions[serverId]?.pendingPermissions
|
(state) => state.sessions[serverId]?.pendingPermissions
|
||||||
);
|
);
|
||||||
|
const setAgents = useSessionStore((state) => state.setAgents);
|
||||||
|
const setPendingPermissions = useSessionStore(
|
||||||
|
(state) => state.setPendingPermissions
|
||||||
|
);
|
||||||
const pendingPermissions = useMemo(() => {
|
const pendingPermissions = useMemo(() => {
|
||||||
if (!allPendingPermissions || !resolvedAgentId) return new Map();
|
if (!allPendingPermissions || !resolvedAgentId) return new Map();
|
||||||
const filtered = new Map();
|
const filtered = new Map();
|
||||||
@@ -662,7 +670,46 @@ function AgentScreenContent({
|
|||||||
const attemptToken = ++initAttemptTokenRef.current;
|
const attemptToken = ++initAttemptTokenRef.current;
|
||||||
|
|
||||||
ensureAgentIsInitialized(resolvedAgentId)
|
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) {
|
if (attemptToken !== initAttemptTokenRef.current) {
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
@@ -681,10 +728,14 @@ function AgentScreenContent({
|
|||||||
});
|
});
|
||||||
}, [
|
}, [
|
||||||
agent,
|
agent,
|
||||||
|
client,
|
||||||
ensureAgentIsInitialized,
|
ensureAgentIsInitialized,
|
||||||
isConnected,
|
isConnected,
|
||||||
missingAgentState.kind,
|
missingAgentState.kind,
|
||||||
resolvedAgentId,
|
resolvedAgentId,
|
||||||
|
serverId,
|
||||||
|
setAgents,
|
||||||
|
setPendingPermissions,
|
||||||
shouldUseOptimisticStream,
|
shouldUseOptimisticStream,
|
||||||
]);
|
]);
|
||||||
|
|
||||||
@@ -754,6 +805,17 @@ function AgentScreenContent({
|
|||||||
</View>
|
</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 (
|
return (
|
||||||
<View style={styles.container} testID="agent-loading">
|
<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;
|
||||||
|
}
|
||||||
@@ -1,20 +1,5 @@
|
|||||||
import type { AgentTimelineItem, DaemonClient } from '@getpaseo/server'
|
import type { AgentTimelineItem, DaemonClient } from '@getpaseo/server'
|
||||||
|
|
||||||
type TimelineFetchResult = {
|
|
||||||
entries: Array<{ item: AgentTimelineItem }>
|
|
||||||
}
|
|
||||||
|
|
||||||
type TimelineFetchClient = {
|
|
||||||
fetchAgentTimeline: (
|
|
||||||
agentId: string,
|
|
||||||
options: {
|
|
||||||
direction: 'tail'
|
|
||||||
limit: 0
|
|
||||||
projection: 'projected'
|
|
||||||
}
|
|
||||||
) => Promise<TimelineFetchResult>
|
|
||||||
}
|
|
||||||
|
|
||||||
type FetchProjectedTimelineItemsInput = {
|
type FetchProjectedTimelineItemsInput = {
|
||||||
client: DaemonClient
|
client: DaemonClient
|
||||||
agentId: string
|
agentId: string
|
||||||
@@ -23,8 +8,7 @@ type FetchProjectedTimelineItemsInput = {
|
|||||||
export async function fetchProjectedTimelineItems(
|
export async function fetchProjectedTimelineItems(
|
||||||
input: FetchProjectedTimelineItemsInput
|
input: FetchProjectedTimelineItemsInput
|
||||||
): Promise<AgentTimelineItem[]> {
|
): Promise<AgentTimelineItem[]> {
|
||||||
const timelineClient = input.client as unknown as TimelineFetchClient
|
const timeline = await input.client.fetchAgentTimeline(input.agentId, {
|
||||||
const timeline = await timelineClient.fetchAgentTimeline(input.agentId, {
|
|
||||||
direction: 'tail',
|
direction: 'tail',
|
||||||
limit: 0,
|
limit: 0,
|
||||||
projection: 'projected',
|
projection: 'projected',
|
||||||
|
|||||||
@@ -104,9 +104,11 @@ export function createEncryptedTransport(
|
|||||||
onerror: (error) => emitError(error),
|
onerror: (error) => emitError(error),
|
||||||
});
|
});
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
logger.warn({ err: error }, "relay_e2ee_handshake_failed");
|
logger.warn({ err: normalizeTransportError(error) }, "relay_e2ee_handshake_failed");
|
||||||
emitError(error);
|
emitError(error);
|
||||||
base.close(1011, "E2EE handshake failed");
|
// Browser WebSocket.close only accepts 1000 or 3000-4999.
|
||||||
|
// Use an app-defined code so this path works in browser and Node runtimes.
|
||||||
|
base.close(4001, "E2EE handshake failed");
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
@@ -189,3 +191,14 @@ function invokeHandler<TArgs extends unknown[]>(
|
|||||||
// no-op
|
// no-op
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
function normalizeTransportError(error: unknown): Record<string, string> {
|
||||||
|
if (error instanceof Error) {
|
||||||
|
return {
|
||||||
|
name: error.name,
|
||||||
|
message: error.message,
|
||||||
|
...(typeof error.stack === "string" ? { stack: error.stack } : {}),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
return { message: String(error) };
|
||||||
|
}
|
||||||
|
|||||||
@@ -1,5 +1,6 @@
|
|||||||
import { describe, expect, test, vi } from "vitest";
|
import { describe, expect, test, vi } from "vitest";
|
||||||
import {
|
import {
|
||||||
|
createEncryptedTransport,
|
||||||
createWebSocketTransportFactory,
|
createWebSocketTransportFactory,
|
||||||
decodeMessageData,
|
decodeMessageData,
|
||||||
describeTransportClose,
|
describeTransportClose,
|
||||||
@@ -8,7 +9,48 @@ import {
|
|||||||
extractRelayMessageData,
|
extractRelayMessageData,
|
||||||
} from "./daemon-client-transport.js";
|
} from "./daemon-client-transport.js";
|
||||||
|
|
||||||
|
const createClientChannelMock = vi.hoisted(() => vi.fn());
|
||||||
|
|
||||||
|
vi.mock("@getpaseo/relay/e2ee", () => ({
|
||||||
|
createClientChannel: createClientChannelMock,
|
||||||
|
}));
|
||||||
|
|
||||||
describe("daemon-client transport helpers", () => {
|
describe("daemon-client transport helpers", () => {
|
||||||
|
test("createEncryptedTransport closes handshake failures with browser-safe code", async () => {
|
||||||
|
createClientChannelMock.mockReset();
|
||||||
|
createClientChannelMock.mockRejectedValueOnce(new Error("handshake failed"));
|
||||||
|
|
||||||
|
let openHandler: (() => void) | null = null;
|
||||||
|
const close = vi.fn();
|
||||||
|
|
||||||
|
createEncryptedTransport(
|
||||||
|
{
|
||||||
|
send: vi.fn(),
|
||||||
|
close,
|
||||||
|
onOpen: (handler) => {
|
||||||
|
openHandler = handler;
|
||||||
|
return () => {
|
||||||
|
if (openHandler === handler) {
|
||||||
|
openHandler = null;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
},
|
||||||
|
onClose: () => () => {},
|
||||||
|
onError: () => () => {},
|
||||||
|
onMessage: () => () => {},
|
||||||
|
},
|
||||||
|
"daemon-public-key",
|
||||||
|
{ warn: vi.fn() }
|
||||||
|
);
|
||||||
|
|
||||||
|
expect(openHandler).not.toBeNull();
|
||||||
|
openHandler?.();
|
||||||
|
|
||||||
|
await vi.waitFor(() => {
|
||||||
|
expect(close).toHaveBeenCalledWith(4001, "E2EE handshake failed");
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|
||||||
test("createWebSocketTransportFactory forwards sends when socket is open", () => {
|
test("createWebSocketTransportFactory forwards sends when socket is open", () => {
|
||||||
const addEventListener = vi.fn();
|
const addEventListener = vi.fn();
|
||||||
const removeEventListener = vi.fn();
|
const removeEventListener = vi.fn();
|
||||||
|
|||||||
@@ -185,10 +185,14 @@ export class VoiceAssistantWebSocketServer {
|
|||||||
server,
|
server,
|
||||||
path: "/ws",
|
path: "/ws",
|
||||||
verifyClient: ({ req }, callback) => {
|
verifyClient: ({ req }, callback) => {
|
||||||
const origin = req.headers.origin;
|
const requestMetadata = extractSocketRequestMetadata(req);
|
||||||
const requestHost = typeof req.headers.host === "string" ? req.headers.host : null;
|
const origin = requestMetadata.origin;
|
||||||
|
const requestHost = requestMetadata.host ?? null;
|
||||||
if (requestHost && !isHostAllowed(requestHost, allowedHosts)) {
|
if (requestHost && !isHostAllowed(requestHost, allowedHosts)) {
|
||||||
this.logger.warn({ host: requestHost }, "Rejected connection from disallowed host");
|
this.logger.warn(
|
||||||
|
{ ...requestMetadata, host: requestHost },
|
||||||
|
"Rejected connection from disallowed host"
|
||||||
|
);
|
||||||
callback(false, 403, "Host not allowed");
|
callback(false, 403, "Host not allowed");
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
@@ -200,7 +204,10 @@ export class VoiceAssistantWebSocketServer {
|
|||||||
if (!origin || allowedOrigins.has(origin) || sameOrigin) {
|
if (!origin || allowedOrigins.has(origin) || sameOrigin) {
|
||||||
callback(true);
|
callback(true);
|
||||||
} else {
|
} else {
|
||||||
this.logger.warn({ origin }, "Rejected connection from origin");
|
this.logger.warn(
|
||||||
|
{ ...requestMetadata, origin },
|
||||||
|
"Rejected connection from origin"
|
||||||
|
);
|
||||||
callback(false, 403, "Origin not allowed");
|
callback(false, 403, "Origin not allowed");
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
@@ -282,7 +289,7 @@ export class VoiceAssistantWebSocketServer {
|
|||||||
|
|
||||||
private async attachSocket(
|
private async attachSocket(
|
||||||
ws: WebSocketLike,
|
ws: WebSocketLike,
|
||||||
_request?: unknown,
|
request?: unknown,
|
||||||
metadata?: ExternalSocketMetadata
|
metadata?: ExternalSocketMetadata
|
||||||
): Promise<void> {
|
): Promise<void> {
|
||||||
const externalSessionKey =
|
const externalSessionKey =
|
||||||
@@ -320,7 +327,24 @@ export class VoiceAssistantWebSocketServer {
|
|||||||
}
|
}
|
||||||
|
|
||||||
const clientId = `client-${++this.clientIdCounter}`;
|
const clientId = `client-${++this.clientIdCounter}`;
|
||||||
const connectionLogger = this.logger.child({ clientId });
|
const requestMetadata = extractSocketRequestMetadata(request);
|
||||||
|
const connectionLoggerFields: Record<string, string> = {
|
||||||
|
clientId,
|
||||||
|
transport: externalSessionKey ? "relay" : "direct",
|
||||||
|
};
|
||||||
|
if (requestMetadata.host) {
|
||||||
|
connectionLoggerFields.host = requestMetadata.host;
|
||||||
|
}
|
||||||
|
if (requestMetadata.origin) {
|
||||||
|
connectionLoggerFields.origin = requestMetadata.origin;
|
||||||
|
}
|
||||||
|
if (requestMetadata.userAgent) {
|
||||||
|
connectionLoggerFields.userAgent = requestMetadata.userAgent;
|
||||||
|
}
|
||||||
|
if (requestMetadata.remoteAddress) {
|
||||||
|
connectionLoggerFields.remoteAddress = requestMetadata.remoteAddress;
|
||||||
|
}
|
||||||
|
const connectionLogger = this.logger.child(connectionLoggerFields);
|
||||||
const socketRef = { current: ws };
|
const socketRef = { current: ws };
|
||||||
|
|
||||||
const session = new Session({
|
const session = new Session({
|
||||||
@@ -510,17 +534,17 @@ export class VoiceAssistantWebSocketServer {
|
|||||||
data: Buffer | ArrayBuffer | Buffer[] | string
|
data: Buffer | ArrayBuffer | Buffer[] | string
|
||||||
): Promise<void> {
|
): Promise<void> {
|
||||||
try {
|
try {
|
||||||
|
const activeConnection = this.sessions.get(ws);
|
||||||
const buffer = bufferFromWsData(data);
|
const buffer = bufferFromWsData(data);
|
||||||
const asBytes = asUint8Array(buffer);
|
const asBytes = asUint8Array(buffer);
|
||||||
if (asBytes) {
|
if (asBytes) {
|
||||||
const frame = decodeBinaryMuxFrame(asBytes);
|
const frame = decodeBinaryMuxFrame(asBytes);
|
||||||
if (frame) {
|
if (frame) {
|
||||||
const connection = this.sessions.get(ws);
|
if (!activeConnection) {
|
||||||
if (!connection) {
|
|
||||||
this.logger.error("No session found for client");
|
this.logger.error("No session found for client");
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
connection.session.handleBinaryFrame(frame);
|
activeConnection.session.handleBinaryFrame(frame);
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -535,8 +559,10 @@ export class VoiceAssistantWebSocketServer {
|
|||||||
"type" in parsed &&
|
"type" in parsed &&
|
||||||
(parsed as { type?: unknown }).type === "session";
|
(parsed as { type?: unknown }).type === "session";
|
||||||
|
|
||||||
this.logger.warn(
|
const log = activeConnection?.connectionLogger ?? this.logger;
|
||||||
|
log.warn(
|
||||||
{
|
{
|
||||||
|
clientId: activeConnection?.clientId,
|
||||||
requestId: requestInfo?.requestId,
|
requestId: requestInfo?.requestId,
|
||||||
requestType: requestInfo?.requestType,
|
requestType: requestInfo?.requestType,
|
||||||
error: parsedMessage.error.message,
|
error: parsedMessage.error.message,
|
||||||
@@ -585,14 +611,13 @@ export class VoiceAssistantWebSocketServer {
|
|||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
|
||||||
const connection = this.sessions.get(ws);
|
if (!activeConnection) {
|
||||||
if (!connection) {
|
|
||||||
this.logger.error("No session found for client");
|
this.logger.error("No session found for client");
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
|
||||||
if (message.type === "session") {
|
if (message.type === "session") {
|
||||||
await connection.session.handleMessage(message.message);
|
await activeConnection.session.handleMessage(message.message);
|
||||||
}
|
}
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
const err = error instanceof Error ? error : new Error(String(error));
|
const err = error instanceof Error ? error : new Error(String(error));
|
||||||
@@ -751,6 +776,49 @@ export class VoiceAssistantWebSocketServer {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
type SocketRequestMetadata = {
|
||||||
|
host?: string;
|
||||||
|
origin?: string;
|
||||||
|
userAgent?: string;
|
||||||
|
remoteAddress?: string;
|
||||||
|
};
|
||||||
|
|
||||||
|
function extractSocketRequestMetadata(request: unknown): SocketRequestMetadata {
|
||||||
|
if (!request || typeof request !== "object") {
|
||||||
|
return {};
|
||||||
|
}
|
||||||
|
|
||||||
|
const record = request as {
|
||||||
|
headers?: {
|
||||||
|
host?: unknown;
|
||||||
|
origin?: unknown;
|
||||||
|
"user-agent"?: unknown;
|
||||||
|
};
|
||||||
|
socket?: {
|
||||||
|
remoteAddress?: unknown;
|
||||||
|
};
|
||||||
|
};
|
||||||
|
|
||||||
|
const host = typeof record.headers?.host === "string" ? record.headers.host : undefined;
|
||||||
|
const origin =
|
||||||
|
typeof record.headers?.origin === "string" ? record.headers.origin : undefined;
|
||||||
|
const userAgent =
|
||||||
|
typeof record.headers?.["user-agent"] === "string"
|
||||||
|
? record.headers["user-agent"]
|
||||||
|
: undefined;
|
||||||
|
const remoteAddress =
|
||||||
|
typeof record.socket?.remoteAddress === "string"
|
||||||
|
? record.socket.remoteAddress
|
||||||
|
: undefined;
|
||||||
|
|
||||||
|
return {
|
||||||
|
...(host ? { host } : {}),
|
||||||
|
...(origin ? { origin } : {}),
|
||||||
|
...(userAgent ? { userAgent } : {}),
|
||||||
|
...(remoteAddress ? { remoteAddress } : {}),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
function stringifyCloseReason(reason: unknown): string | null {
|
function stringifyCloseReason(reason: unknown): string | null {
|
||||||
if (typeof reason === "string") {
|
if (typeof reason === "string") {
|
||||||
return reason.length > 0 ? reason : null;
|
return reason.length > 0 ? reason : null;
|
||||||
|
|||||||
@@ -78,6 +78,28 @@ describe("shared messages stream parsing", () => {
|
|||||||
}
|
}
|
||||||
});
|
});
|
||||||
|
|
||||||
|
it("rejects removed initialize_agent_request inbound payload", () => {
|
||||||
|
const parsed = SessionInboundMessageSchema.safeParse({
|
||||||
|
type: "initialize_agent_request",
|
||||||
|
agentId: "agent-legacy",
|
||||||
|
requestId: "req-legacy-1",
|
||||||
|
});
|
||||||
|
expect(parsed.success).toBe(false);
|
||||||
|
});
|
||||||
|
|
||||||
|
it("rejects removed initialize_agent_request outbound payload", () => {
|
||||||
|
const parsed = SessionOutboundMessageSchema.safeParse({
|
||||||
|
type: "initialize_agent_request",
|
||||||
|
payload: {
|
||||||
|
requestId: "req-legacy-1",
|
||||||
|
agentId: "agent-legacy",
|
||||||
|
agentStatus: "running",
|
||||||
|
timelineSize: 12,
|
||||||
|
},
|
||||||
|
});
|
||||||
|
expect(parsed.success).toBe(false);
|
||||||
|
});
|
||||||
|
|
||||||
it("rejects websocket envelope for removed agent_stream_snapshot message type", () => {
|
it("rejects websocket envelope for removed agent_stream_snapshot message type", () => {
|
||||||
const fixture = {
|
const fixture = {
|
||||||
type: "agent_stream_snapshot",
|
type: "agent_stream_snapshot",
|
||||||
|
|||||||
Reference in New Issue
Block a user