1import {
2 EMPTY_ENVIRONMENT_THREAD_STATE,
3 type EnvironmentThreadState,
4} from "@t3tools/client-runtime/state/threads";
5import {
6 EnvironmentId,
7 MessageId,
8 RunId,
9 ThreadId,
10 type OrchestrationV2ThreadShell,
11} from "@t3tools/contracts";
12import { makeThreadProjectionFixture } from "../test-fixtures";
13import * as Option from "effect/Option";
14import { AsyncResult, Atom, AtomRegistry } from "effect/reactivity";
15import { describe, expect, it } from "vite-plus/test";
16
17import { createRunningThreadKeepAliveAtom } from "./threads";
18
19const LOCAL = EnvironmentId.make("local");
20const REMOTE = EnvironmentId.make("remote");
21
22type Status = "running" | "starting" | "preparing" | "queued" | "waiting" | "idle";
23function shell(id: string, status: Status | null) {
24 return { id: ThreadId.make(id), status: status ?? "idle" } satisfies Pick<
25 OrchestrationV2ThreadShell,
26 "id" | "status"
27 >;
28}
29function detail(id: string, status: Status, overrides: Partial<EnvironmentThreadState> = {}) {
30 const projection = makeThreadProjectionFixture();
31 const threadId = ThreadId.make(id);
32 const thread = {
33 ...projection,
34 thread: { ...projection.thread, id: threadId },
35 runs:
36 status === "idle"
37 ? []
38 : [
39 {
40 id: RunId.make(`run-${id}`),
41 threadId,
42 ordinal: 1,
43 providerInstanceId: projection.thread.providerInstanceId,
44 modelSelection: projection.thread.modelSelection,
45 providerThreadId: null,
46 userMessageId: MessageId.make(`message-${id}`),
47 rootNodeId: null,
48 activeAttemptId: null,
49 status,
50 requestedAt: projection.updatedAt,
51 startedAt: null,
52 completedAt: null,
53 checkpointId: null,
54 contextHandoffId: null,
55 },
56 ],
57 };
58 return AsyncResult.success<EnvironmentThreadState>({
59 ...EMPTY_ENVIRONMENT_THREAD_STATE,
60 status: "live",
61 data: Option.some(thread),
62 ...overrides,
63 });
64}
65
66function makeHarness() {
67 // Registry cleanup runs only on `flush`, like the real deferred task.
68 const tasks: Array<() => void> = [];
69 const registry = AtomRegistry.make({
70 scheduleTask: (task) => {
71 tasks.push(task);
72 return () => {};
73 },
74 });
75 const flush = () => {
76 for (let task = tasks.shift(); task !== undefined; task = tasks.shift()) task();
77 };
78 const environmentIds = Atom.make<ReadonlyArray<EnvironmentId>>([LOCAL, REMOTE]).pipe(
79 Atom.keepAlive,
80 );
81 const threads = Atom.family((_environmentId: EnvironmentId) =>
82 Atom.make<ReadonlyArray<ReturnType<typeof shell>>>([]).pipe(Atom.keepAlive),
83 );
84 // Stand-ins for the thread state atoms. Each one lives only while mounted,
85 // as the real stream does.
86 const keys = new Set<string>();
87 const states = Atom.family((_key: string) =>
88 Atom.make<AsyncResult.AsyncResult<EnvironmentThreadState>>(
89 AsyncResult.success(EMPTY_ENVIRONMENT_THREAD_STATE),
90 ),
91 );
92 const stateAtom = (environmentId: EnvironmentId, threadId: string) => {
93 const key = `${environmentId}:${threadId}`;
94 keys.add(key);
95 return states(key);
96 };
97 const keepAlive = createRunningThreadKeepAliveAtom({
98 environmentIdsAtom: environmentIds,
99 threadsAtom: threads,
100 stateAtom,
101 });
102 registry.mount(keepAlive);
103 return {
104 registry,
105 environmentIds,
106 threads,
107 stateAtom,
108 keepAlive,
109 openStreams: () => {
110 flush();
111 return [...keys].filter((key) => registry.getNodes().has(states(key))).toSorted();
112 },
113 };
114}
115
116describe("createRunningThreadKeepAliveAtom", () => {
117 it("keeps running threads open across shell updates and thread view visits", () => {
118 const h = makeHarness();
119 h.registry.set(h.threads(LOCAL), [shell("a", "running"), shell("b", "idle"), shell("c", null)]);
120 h.registry.set(h.threads(REMOTE), [shell("d", "starting")]);
121 expect(h.openStreams()).toEqual(["local:a", "remote:d"]);
122
123 // A thread view that comes and goes shares the kept stream.
124 const live = detail("a", "running");
125 h.registry.set(h.stateAtom(LOCAL, "a"), live);
126 h.registry.mount(h.stateAtom(LOCAL, "a"))();
127
128 // A shell update that starts or stops nothing does not rebuild the set.
129 const kept = h.registry.get(h.keepAlive);
130 h.registry.set(h.threads(LOCAL), [shell("a", "running"), shell("b", "idle")]);
131 expect(h.registry.get(h.keepAlive)).toBe(kept);
132 expect(h.openStreams()).toEqual(["local:a", "remote:d"]);
133 expect(h.registry.get(h.stateAtom(LOCAL, "a"))).toBe(live);
134 });
135
136 it("holds a stopped thread until its own stream is live and shows the stop", () => {
137 const h = makeHarness();
138 h.registry.set(h.threads(LOCAL), [
139 shell("a", "running"),
140 shell("b", "running"),
141 shell("c", "running"),
142 ]);
143 // "b" has not loaded yet. "c" hit a stream error.
144 h.registry.set(h.stateAtom(LOCAL, "a"), detail("a", "running"));
145 h.registry.set(
146 h.stateAtom(LOCAL, "c"),
147 detail("c", "running", { status: "cached", error: Option.some("Could not sync.") }),
148 );
149
150 // The shell reports the stops first. A failed stream cannot deliver its
151 // stop, so only it is released now.
152 h.registry.set(h.threads(LOCAL), [shell("a", "idle"), shell("b", "idle"), shell("c", "idle")]);
153 expect(h.openStreams()).toEqual(["local:a", "local:b"]);
154
155 h.registry.set(h.stateAtom(LOCAL, "a"), detail("a", "idle"));
156 h.registry.set(h.stateAtom(LOCAL, "b"), detail("b", "idle", { status: "synchronizing" }));
157 expect(h.openStreams()).toEqual(["local:b"]);
158 h.registry.set(h.stateAtom(LOCAL, "b"), detail("b", "idle"));
159 expect(h.openStreams()).toEqual([]);
160 });
161
162 it("keeps queued threads open", () => {
163 const h = makeHarness();
164 h.registry.set(h.threads(LOCAL), [shell("a", "queued")]);
165 expect(h.openStreams()).toEqual(["local:a"]);
166 });
167
168 it("follows environments that connect and go away", () => {
169 const h = makeHarness();
170 h.registry.set(h.environmentIds, [LOCAL]);
171 h.registry.set(h.threads(REMOTE), [shell("d", "running")]);
172 expect(h.openStreams()).toEqual([]);
173
174 h.registry.set(h.environmentIds, [LOCAL, REMOTE]);
175 expect(h.openStreams()).toEqual(["remote:d"]);
176
177 // Removal drops every mount, including one still waiting for its stop.
178 h.registry.set(h.stateAtom(REMOTE, "d"), detail("d", "running"));
179 h.registry.set(h.threads(REMOTE), [shell("d", "idle")]);
180 expect(h.openStreams()).toEqual(["remote:d"]);
181 h.registry.set(h.environmentIds, [LOCAL]);
182 expect(h.openStreams()).toEqual([]);
183 });
184});