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 const tasks: Array<() => void> = [];
68 const registry = AtomRegistry.make({
69 scheduleTask: (task) => {
70 tasks.push(task);
71 return () => {};
72 },
73 });
74 const flush = () => {
75 for (let task = tasks.shift(); task !== undefined; task = tasks.shift()) task();
76 };
77 const environmentIds = Atom.make<ReadonlyArray<EnvironmentId>>([LOCAL, REMOTE]).pipe(
78 Atom.keepAlive,
79 );
80 const threads = Atom.family((_environmentId: EnvironmentId) =>
81 Atom.make<ReadonlyArray<ReturnType<typeof shell>>>([]).pipe(Atom.keepAlive),
82 );
83 const keys = new Set<string>();
84 const states = Atom.family((_key: string) =>
85 Atom.make<AsyncResult.AsyncResult<EnvironmentThreadState>>(
86 AsyncResult.success(EMPTY_ENVIRONMENT_THREAD_STATE),
87 ),
88 );
89 const stateAtom = (environmentId: EnvironmentId, threadId: string) => {
90 const key = `${environmentId}:${threadId}`;
91 keys.add(key);
92 return states(key);
93 };
94 const keepAlive = createRunningThreadKeepAliveAtom({
95 environmentIdsAtom: environmentIds,
96 threadsAtom: threads,
97 stateAtom,
98 });
99 registry.mount(keepAlive);
100 return {
101 registry,
102 environmentIds,
103 threads,
104 stateAtom,
105 keepAlive,
106 openStreams: () => {
107 flush();
108 return [...keys].filter((key) => registry.getNodes().has(states(key))).toSorted();
109 },
110 };
111}
112
113describe("createRunningThreadKeepAliveAtom", () => {
114 it("keeps running threads open across shell updates and thread view visits", () => {
115 const h = makeHarness();
116 h.registry.set(h.threads(LOCAL), [shell("a", "running"), shell("b", "idle"), shell("c", null)]);
117 h.registry.set(h.threads(REMOTE), [shell("d", "starting")]);
118 expect(h.openStreams()).toEqual(["local:a", "remote:d"]);
119
120 const live = detail("a", "running");
121 h.registry.set(h.stateAtom(LOCAL, "a"), live);
122 h.registry.mount(h.stateAtom(LOCAL, "a"))();
123
124 const kept = h.registry.get(h.keepAlive);
125 h.registry.set(h.threads(LOCAL), [shell("a", "running"), shell("b", "idle")]);
126 expect(h.registry.get(h.keepAlive)).toBe(kept);
127 expect(h.openStreams()).toEqual(["local:a", "remote:d"]);
128 expect(h.registry.get(h.stateAtom(LOCAL, "a"))).toBe(live);
129 });
130
131 it("holds a stopped thread until its own stream is live and shows the stop", () => {
132 const h = makeHarness();
133 h.registry.set(h.threads(LOCAL), [
134 shell("a", "running"),
135 shell("b", "running"),
136 shell("c", "running"),
137 ]);
138 h.registry.set(h.stateAtom(LOCAL, "a"), detail("a", "running"));
139 h.registry.set(
140 h.stateAtom(LOCAL, "c"),
141 detail("c", "running", { status: "cached", error: Option.some("Could not sync.") }),
142 );
143
144 h.registry.set(h.threads(LOCAL), [shell("a", "idle"), shell("b", "idle"), shell("c", "idle")]);
145 expect(h.openStreams()).toEqual(["local:a", "local:b"]);
146
147 h.registry.set(h.stateAtom(LOCAL, "a"), detail("a", "idle"));
148 h.registry.set(h.stateAtom(LOCAL, "b"), detail("b", "idle", { status: "synchronizing" }));
149 expect(h.openStreams()).toEqual(["local:b"]);
150 h.registry.set(h.stateAtom(LOCAL, "b"), detail("b", "idle"));
151 expect(h.openStreams()).toEqual([]);
152 });
153
154 it("follows environments that connect and go away", () => {
155 const h = makeHarness();
156 h.registry.set(h.environmentIds, [LOCAL]);
157 h.registry.set(h.threads(REMOTE), [shell("d", "running")]);
158 expect(h.openStreams()).toEqual([]);
159
160 h.registry.set(h.environmentIds, [LOCAL, REMOTE]);
161 expect(h.openStreams()).toEqual(["remote:d"]);
162
163 h.registry.set(h.stateAtom(REMOTE, "d"), detail("d", "running"));
164 h.registry.set(h.threads(REMOTE), [shell("d", "idle")]);
165 expect(h.openStreams()).toEqual(["remote:d"]);
166 h.registry.set(h.environmentIds, [LOCAL]);
167 expect(h.openStreams()).toEqual([]);
168 });
169});