1import type { ReactNode } from "react";
2import { act, useSyncExternalStore } from "react";
3import { createRoot } from "react-dom/client";
4import { AsyncResult } from "effect/reactivity";
5import { expect, it, vi } from "vite-plus/test";
6
7interface Snapshot {
8 readonly projects: ReadonlyArray<{
9 readonly environmentId: string;
10 readonly id: string;
11 readonly workspaceRoot: string;
12 }>;
13 readonly catalog: {
14 readonly isLoadingConnections: boolean;
15 readonly hasLoadedShellSnapshot: boolean;
16 readonly hasConnectingEnvironment: boolean;
17 readonly connectionError: string | null;
18 };
19}
20
21const harness = vi.hoisted(() => {
22 const listeners = new Set<() => void>();
23 const state: { snapshot: Snapshot | null } = { snapshot: null };
24 return {
25 state,
26 set(next: Snapshot) {
27 state.snapshot = next;
28 for (const listener of listeners) listener();
29 },
30 subscribe(listener: () => void) {
31 listeners.add(listener);
32 return () => listeners.delete(listener);
33 },
34 route: { key: "thread-route", name: "Thread", params: {} as object },
35 reset: () => undefined,
36 checkout: vi.fn(),
37 draftProjectRefs: [] as unknown[],
38 // Hook results stay stable across renders, as the real hooks' are. Only
39 // main's screen reads these (see the main-comparison mocks below).
40 layout: { deferThreadSelectionUntil: () => undefined },
41 navigation: { goBack: () => undefined },
42 preventContext: { preventedRoutes: {} },
43 };
44});
45
46function useSnapshot(): Snapshot {
47 return useSyncExternalStore(harness.subscribe, () => harness.state.snapshot!);
48}
49
50vi.mock("react-native", () => ({
51 View: ({ children }: { children: ReactNode }) => <div>{children}</div>,
52}));
53vi.mock("../../components/ui/AppText", () => ({
54 AppText: ({ children }: { children: ReactNode }) => <span>{children}</span>,
55}));
56vi.mock("../../components/ui/Dialog", () => ({ alert: vi.fn() }));
57vi.mock("../../components/ui/AppBar", () => ({ AppBar: () => null }));
58vi.mock("../layout/workspace/useAndroidWorkspaceNavigation", () => ({
59 useWorkspaceBackNavigation: () => () => undefined,
60}));
61vi.mock("./ThreadScreenHost", () => ({
62 ThreadScreenFallback: ({ children }: { children: ReactNode }) => <>{children}</>,
63}));
64vi.mock("./NewTaskDraftScreen", () => ({
65 NewTaskDraftScreen: (props: { initialProjectRef: { branch?: string | null } }) => {
66 harness.draftProjectRefs.push(props.initialProjectRef);
67 return <p>draft on {props.initialProjectRef.branch}</p>;
68 },
69}));
70vi.mock("./new-task-context-presentation", () => ({
71 resolveNewTaskBranchWorktreePath: () => null,
72}));
73vi.mock("./new-task-flow-provider", () => ({
74 useNewTaskFlow: () => harness,
75}));
76vi.mock("./checkout-new-task-branch", () => ({ checkoutNewTaskBranch: harness.checkout }));
77// Main-comparison mocks: this screen no longer uses them, but main's screen
78// did (deferThreadSelectionUntil, useAtomCommand, useNavigation). They stay so
79// this test can run against main's screen and fail there for the right reason
80// (a second checkout) instead of on an unmockable import.
81vi.mock("../layout/AdaptiveWorkspaceLayout", () => ({
82 useAdaptiveWorkspaceLayout: () => harness.layout,
83}));
84vi.mock("../../state/use-atom-command", () => ({ useAtomCommand: () => harness.checkout }));
85vi.mock("../../state/entities", () => ({ useProjects: () => useSnapshot().projects }));
86vi.mock("../../state/workspace", () => ({
87 useWorkspaceState: () => ({ state: useSnapshot().catalog }),
88}));
89vi.mock("../../state/vcs", () => ({ vcsEnvironment: { readRefs: {}, switchRef: {} } }));
90vi.mock("@react-navigation/native", () => ({
91 CommonActions: { goBack: () => ({ type: "GO_BACK" }) },
92 createNavigationContainerRef: () => ({ current: null }),
93 // Main-comparison only, like the mocks above.
94 useNavigation: () => harness.navigation,
95 useRoute: () => harness.route,
96 usePreventRemove: () => undefined,
97 usePreventRemoveContext: () => harness.preventContext,
98}));
99
100import { NewTaskDraftRouteScreen } from "./NewTaskDraftRouteScreen";
101
102const project = { environmentId: "env-1", id: "project-1", workspaceRoot: "/repo" };
103const connected: Snapshot["catalog"] = {
104 isLoadingConnections: false,
105 hasLoadedShellSnapshot: true,
106 hasConnectingEnvironment: false,
107 connectionError: null,
108};
109
110// The project can vanish while the catalog still loads, or while it reports ready.
111it.each([
112 ["while the catalog reloads", { ...connected, isLoadingConnections: true }],
113 ["while the catalog stays ready", connected],
114])(
115 "checks the branch out once when a reconnect drops and restores the project %s",
116 async (gap, gapCatalog) => {
117 // Checkouts live in a module table keyed by route, so each case uses its own route.
118 const routeKey = `thread-route ${gap}`;
119 (globalThis as { IS_REACT_ACT_ENVIRONMENT?: boolean }).IS_REACT_ACT_ENVIRONMENT = true;
120 const firstCheckout = Promise.withResolvers<unknown>();
121 harness.checkout.mockReturnValueOnce(firstCheckout.promise);
122 harness.checkout.mockResolvedValue(
123 AsyncResult.success({ name: "feature", current: true, isDefault: false, worktreePath: null }),
124 );
125 harness.set({ projects: [project], catalog: connected });
126 const draft = { environmentId: "env-1", projectId: "project-1", branch: "feature" };
127 harness.route = { key: routeKey, name: "Thread", params: { draft } };
128
129 const container = document.createElement("div");
130 const root = createRoot(container);
131 using _root = { [Symbol.dispose]: () => act(() => root.unmount()) };
132 act(() =>
133 root.render(
134 <NewTaskDraftRouteScreen
135 route={{ key: routeKey, name: "Thread", params: draft } as never}
136 />,
137 ),
138 );
139 await vi.waitFor(() => expect(harness.checkout).toHaveBeenCalledTimes(1));
140
141 // The connection drops while the checkout runs, then the project comes back.
142 act(() => harness.set({ projects: [], catalog: gapCatalog }));
143 act(() => harness.set({ projects: [project], catalog: connected }));
144 await act(async () => {
145 firstCheckout.resolve(
146 AsyncResult.success({
147 name: "feature",
148 current: true,
149 isDefault: false,
150 worktreePath: null,
151 }),
152 );
153 });
154
155 await vi.waitFor(() =>
156 expect(
157 harness.checkout.mock.calls.length > 1 || container.textContent === "draft on feature",
158 ).toBe(true),
159 );
160 expect(harness.checkout).toHaveBeenCalledTimes(1);
161 expect(container.textContent).toBe("draft on feature");
162 },
163);
164
165it("hands the draft screen the same project ref across re-renders", async () => {
166 (globalThis as { IS_REACT_ACT_ENVIRONMENT?: boolean }).IS_REACT_ACT_ENVIRONMENT = true;
167 harness.checkout.mockReset();
168 harness.draftProjectRefs.length = 0;
169 harness.checkout.mockResolvedValue(
170 AsyncResult.success({ name: "other", current: true, isDefault: false, worktreePath: null }),
171 );
172 harness.set({ projects: [project], catalog: connected });
173 const draft = { environmentId: "env-1", projectId: "project-1", branch: "other" };
174 harness.route = { key: "stable-route", name: "Thread", params: { draft } };
175
176 const container = document.createElement("div");
177 const root = createRoot(container);
178 using _root = { [Symbol.dispose]: () => act(() => root.unmount()) };
179 act(() =>
180 root.render(
181 <NewTaskDraftRouteScreen
182 route={{ key: "stable-route", name: "Thread", params: draft } as never}
183 />,
184 ),
185 );
186 await vi.waitFor(() => expect(container.textContent).toBe("draft on other"));
187
188 // The draft screen re-applies its route project when this identity changes,
189 // and that write re-renders the route; an unstable ref loops forever.
190 act(() => harness.set({ projects: [{ ...project }], catalog: { ...connected } }));
191 // A stable ref lets the compiled screen skip re-rendering the draft screen at all; an unstable
192 // one re-renders it with a second ref.
193 expect(new Set(harness.draftProjectRefs).size).toBe(1);
194});