Skip to content

Commit 7b74407

Browse files
jacobjmcjacobjmc
committed
fix(workspaces): retry restore and prevent duplicate reconnects
1 parent f710e5f commit 7b74407

File tree

2 files changed

+250
-6
lines changed

2 files changed

+250
-6
lines changed

‎src/features/workspaces/hooks/useWorkspaceRestore.test.tsx‎

Lines changed: 147 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,147 @@
1+
// @vitest-environment jsdom
2+
import { act, renderHook } from "@testing-library/react";
3+
import { afterEach, describe, expect, it, vi } from "vitest";
4+
import type { WorkspaceInfo } from "../../../types";
5+
import { useWorkspaceRestore } from "./useWorkspaceRestore";
6+
7+
const RETRY_DELAY_MS = 1500;
8+
9+
function flushMicrotasks() {
10+
return Promise.resolve().then(() => Promise.resolve());
11+
}
12+
13+
function createWorkspace(overrides?: Partial<WorkspaceInfo>): WorkspaceInfo {
14+
return {
15+
id: "ws-1",
16+
name: "Workspace",
17+
path: "/tmp/workspace",
18+
connected: false,
19+
kind: "main",
20+
parentId: null,
21+
worktree: null,
22+
settings: { sidebarCollapsed: false },
23+
...overrides,
24+
};
25+
}
26+
27+
describe("useWorkspaceRestore", () => {
28+
afterEach(() => {
29+
vi.useRealTimers();
30+
vi.clearAllMocks();
31+
});
32+
33+
it("retries connect/list after a startup failure", async () => {
34+
vi.useFakeTimers();
35+
const workspace = createWorkspace();
36+
const connectWorkspace = vi
37+
.fn<WorkspaceRestoreOptions["connectWorkspace"]>()
38+
.mockRejectedValueOnce(new Error("temporary failure"))
39+
.mockResolvedValue(undefined);
40+
const listThreadsForWorkspace = vi
41+
.fn<WorkspaceRestoreOptions["listThreadsForWorkspace"]>()
42+
.mockResolvedValue(undefined);
43+
44+
renderHook(() =>
45+
useWorkspaceRestore({
46+
workspaces: [workspace],
47+
hasLoaded: true,
48+
connectWorkspace,
49+
listThreadsForWorkspace,
50+
}),
51+
);
52+
53+
await act(async () => {
54+
await flushMicrotasks();
55+
});
56+
57+
expect(connectWorkspace).toHaveBeenCalledTimes(1);
58+
expect(listThreadsForWorkspace).toHaveBeenCalledTimes(0);
59+
60+
await act(async () => {
61+
vi.advanceTimersByTime(RETRY_DELAY_MS);
62+
await flushMicrotasks();
63+
});
64+
65+
expect(connectWorkspace).toHaveBeenCalledTimes(2);
66+
expect(listThreadsForWorkspace).toHaveBeenCalledTimes(1);
67+
});
68+
69+
it("does not start duplicate reconnects while one is in flight", async () => {
70+
vi.useFakeTimers();
71+
const workspace = createWorkspace();
72+
let resolveConnect: (() => void) | null = null;
73+
const connectWorkspace = vi.fn<WorkspaceRestoreOptions["connectWorkspace"]>(() =>
74+
new Promise<void>((resolve) => {
75+
resolveConnect = resolve;
76+
}),
77+
);
78+
const listThreadsForWorkspace = vi
79+
.fn<WorkspaceRestoreOptions["listThreadsForWorkspace"]>()
80+
.mockResolvedValue(undefined);
81+
82+
const { rerender } = renderHook(
83+
({ workspaces }) =>
84+
useWorkspaceRestore({
85+
workspaces,
86+
hasLoaded: true,
87+
connectWorkspace,
88+
listThreadsForWorkspace,
89+
}),
90+
{
91+
initialProps: { workspaces: [workspace] as WorkspaceInfo[] },
92+
},
93+
);
94+
95+
await act(async () => {
96+
await flushMicrotasks();
97+
});
98+
99+
rerender({ workspaces: [{ ...workspace }] });
100+
101+
await act(async () => {
102+
await flushMicrotasks();
103+
});
104+
105+
expect(connectWorkspace).toHaveBeenCalledTimes(1);
106+
107+
await act(async () => {
108+
resolveConnect?.();
109+
await flushMicrotasks();
110+
});
111+
112+
expect(listThreadsForWorkspace).toHaveBeenCalledTimes(1);
113+
});
114+
115+
it("loads thread list on open for already connected workspaces", async () => {
116+
const workspace = createWorkspace({ connected: true });
117+
const connectWorkspace = vi.fn<WorkspaceRestoreOptions["connectWorkspace"]>();
118+
const listThreadsForWorkspace = vi
119+
.fn<WorkspaceRestoreOptions["listThreadsForWorkspace"]>()
120+
.mockResolvedValue(undefined);
121+
122+
renderHook(() =>
123+
useWorkspaceRestore({
124+
workspaces: [workspace],
125+
hasLoaded: true,
126+
connectWorkspace,
127+
listThreadsForWorkspace,
128+
}),
129+
);
130+
131+
await act(async () => {
132+
await flushMicrotasks();
133+
});
134+
135+
expect(connectWorkspace).not.toHaveBeenCalled();
136+
expect(listThreadsForWorkspace).toHaveBeenCalledTimes(1);
137+
expect(listThreadsForWorkspace).toHaveBeenCalledWith(workspace);
138+
});
139+
});
140+
141+
type WorkspaceRestoreOptions = {
142+
connectWorkspace: (workspace: WorkspaceInfo) => Promise<void>;
143+
listThreadsForWorkspace: (
144+
workspace: WorkspaceInfo,
145+
options?: { preserveState?: boolean },
146+
) => Promise<void>;
147+
};

‎src/features/workspaces/hooks/useWorkspaceRestore.ts‎

Lines changed: 103 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
import { useEffect, useRef } from "react";
22
import type { WorkspaceInfo } from "../../../types";
33

4+
const RESTORE_RETRY_DELAY_MS = 1500;
5+
const MAX_RESTORE_RETRIES = 5;
6+
47
type WorkspaceRestoreOptions = {
58
workspaces: WorkspaceInfo[];
69
hasLoaded: boolean;
@@ -18,26 +21,120 @@ export function useWorkspaceRestore({
1821
listThreadsForWorkspace,
1922
}: WorkspaceRestoreOptions) {
2023
const restoredWorkspaces = useRef(new Set<string>());
24+
const restoringWorkspaces = useRef(new Set<string>());
25+
const retryCountByWorkspace = useRef(new Map<string, number>());
26+
const retryTimers = useRef(new Map<string, ReturnType<typeof setTimeout>>());
27+
const optionsRef = useRef({
28+
workspaces,
29+
hasLoaded,
30+
connectWorkspace,
31+
listThreadsForWorkspace,
32+
});
33+
34+
useEffect(() => {
35+
optionsRef.current = {
36+
workspaces,
37+
hasLoaded,
38+
connectWorkspace,
39+
listThreadsForWorkspace,
40+
};
41+
});
42+
43+
useEffect(
44+
() => () => {
45+
retryTimers.current.forEach((timer) => {
46+
clearTimeout(timer);
47+
});
48+
retryTimers.current.clear();
49+
},
50+
[],
51+
);
2152

2253
useEffect(() => {
2354
if (!hasLoaded) {
2455
return;
2556
}
26-
workspaces.forEach((workspace) => {
27-
if (restoredWorkspaces.current.has(workspace.id)) {
57+
58+
const activeWorkspaceIds = new Set(workspaces.map((workspace) => workspace.id));
59+
restoredWorkspaces.current.forEach((workspaceId) => {
60+
if (!activeWorkspaceIds.has(workspaceId)) {
61+
restoredWorkspaces.current.delete(workspaceId);
62+
}
63+
});
64+
restoringWorkspaces.current.forEach((workspaceId) => {
65+
if (!activeWorkspaceIds.has(workspaceId)) {
66+
restoringWorkspaces.current.delete(workspaceId);
67+
}
68+
});
69+
retryCountByWorkspace.current.forEach((_count, workspaceId) => {
70+
if (!activeWorkspaceIds.has(workspaceId)) {
71+
retryCountByWorkspace.current.delete(workspaceId);
72+
}
73+
});
74+
retryTimers.current.forEach((timer, workspaceId) => {
75+
if (activeWorkspaceIds.has(workspaceId)) {
76+
return;
77+
}
78+
clearTimeout(timer);
79+
retryTimers.current.delete(workspaceId);
80+
});
81+
82+
const restoreWorkspace = (workspaceId: string) => {
83+
const {
84+
workspaces: latestWorkspaces,
85+
hasLoaded: latestHasLoaded,
86+
connectWorkspace: latestConnectWorkspace,
87+
listThreadsForWorkspace: latestListThreadsForWorkspace,
88+
} = optionsRef.current;
89+
if (!latestHasLoaded) {
90+
return;
91+
}
92+
const workspace = latestWorkspaces.find((entry) => entry.id === workspaceId);
93+
if (!workspace) {
94+
return;
95+
}
96+
if (restoredWorkspaces.current.has(workspaceId)) {
97+
return;
98+
}
99+
if (restoringWorkspaces.current.has(workspaceId)) {
28100
return;
29101
}
30-
restoredWorkspaces.current.add(workspace.id);
102+
103+
restoringWorkspaces.current.add(workspaceId);
31104
void (async () => {
32105
try {
33106
if (!workspace.connected) {
34-
await connectWorkspace(workspace);
107+
await latestConnectWorkspace(workspace);
108+
}
109+
await latestListThreadsForWorkspace(workspace);
110+
restoredWorkspaces.current.add(workspaceId);
111+
retryCountByWorkspace.current.delete(workspaceId);
112+
const retryTimer = retryTimers.current.get(workspaceId);
113+
if (retryTimer) {
114+
clearTimeout(retryTimer);
115+
retryTimers.current.delete(workspaceId);
35116
}
36-
await listThreadsForWorkspace(workspace);
37117
} catch {
38-
// Silent: connection errors show in debug panel.
118+
const attempts = retryCountByWorkspace.current.get(workspaceId) ?? 0;
119+
if (attempts >= MAX_RESTORE_RETRIES) {
120+
return;
121+
}
122+
retryCountByWorkspace.current.set(workspaceId, attempts + 1);
123+
if (!retryTimers.current.has(workspaceId)) {
124+
const timer = setTimeout(() => {
125+
retryTimers.current.delete(workspaceId);
126+
restoreWorkspace(workspaceId);
127+
}, RESTORE_RETRY_DELAY_MS);
128+
retryTimers.current.set(workspaceId, timer);
129+
}
130+
} finally {
131+
restoringWorkspaces.current.delete(workspaceId);
39132
}
40133
})();
134+
};
135+
136+
workspaces.forEach((workspace) => {
137+
restoreWorkspace(workspace.id);
41138
});
42139
}, [connectWorkspace, hasLoaded, listThreadsForWorkspace, workspaces]);
43140
}

0 commit comments

Comments
 (0)