1import {
2 ORCHESTRATION_V2_WS_METHODS,
3 type EnvironmentId,
4 type OrchestrationV2ShellSnapshot,
5 type OrchestrationV2ThreadFamily,
6 type OrchestrationV2ThreadShell,
7 type ThreadId,
8} from "@t3tools/contracts";
9import * as Effect from "effect/Effect";
10import * as Option from "effect/Option";
11import * as SubscriptionRef from "effect/SubscriptionRef";
12import { AsyncResult, Atom, type AtomRegistry } from "effect/reactivity";
13
14import * as EnvironmentSupervisor from "../connection/supervisor.ts";
15import type * as EnvironmentRegistry from "../connection/registry.ts";
16import { EnvironmentRpcUnavailableError, request } from "../rpc/client.ts";
17import { arrayElementsEqual } from "./entities.ts";
18import { createEnvironmentQueryAtomFamily, selectEnvironmentQueryData } from "./runtime.ts";
19import type { EnvironmentShellState } from "./shell.ts";
20
21const EMPTY_FAMILY: OrchestrationV2ThreadFamily = { snapshotSequence: 0, threads: [] };
22const EMPTY_READS: ReadonlyArray<OrchestrationV2ThreadFamily> = [];
23
24interface MountedRows {
25 readonly rows: Atom.Atom<ReadonlyArray<OrchestrationV2ThreadFamily>>;
26 readonly count: number;
27}
28const EMPTY_MOUNTS: ReadonlyMap<string, MountedRows> = new Map();
29// Generation 0 is the cached shell from before this session's first authoritative snapshot. A read
30// for it would be thrown away once that snapshot lands, so it waits instead.
31const AWAITING_AUTHORITATIVE_SHELL = AsyncResult.initial<never, never>(true);
32
33/** Whether a family read can decide anything yet. Only a settled read can prove a thread is gone. */
34export type ThreadFamilyRead = "pending" | "failed" | "settled";
35
36export function threadFamilyRead(
37 result: AsyncResult.AsyncResult<OrchestrationV2ThreadFamily, unknown>,
38): ThreadFamilyRead {
39 // A failure whose retry is in flight is pending, so callers do not restart that retry.
40 if (result.waiting) return "pending";
41 if (AsyncResult.isFailure(result)) return "failed";
42 return AsyncResult.isSuccess(result) ? "settled" : "pending";
43}
44
45/**
46 * Family rows the shell does not have. A shell row is at least as new as any family copy, since
47 * shell deltas carry every change; a family copy read before the shell removed that thread is
48 * stale.
49 */
50export function threadFamilyRowsMissingFromShell(input: {
51 readonly shellThreads: ReadonlyArray<OrchestrationV2ThreadShell>;
52 readonly families: ReadonlyArray<OrchestrationV2ThreadFamily>;
53 readonly removedAt: ReadonlyMap<ThreadId, number>;
54}): ReadonlyArray<OrchestrationV2ThreadShell> {
55 const present = new Set<ThreadId>(input.shellThreads.map((thread) => thread.id));
56 const missing: OrchestrationV2ThreadShell[] = [];
57 for (const family of input.families) {
58 for (const thread of family.threads) {
59 if (present.has(thread.id)) continue;
60 if ((input.removedAt.get(thread.id) ?? -1) > family.snapshotSequence) continue;
61 present.add(thread.id);
62 missing.push(thread);
63 }
64 }
65 return missing;
66}
67
68/**
69 * Shells with `omitIdleSubagents` or `omitOlderSettled` leave rows to on-demand reads: a family is
70 * every active thread sharing a `lineage.rootThreadId`, read once per shell generation; shelf pages
71 * and search rows mount through `mountRows`. All merge behind shell rows, which win (deltas carry
72 * every change), and are never persisted.
73 */
74export function createThreadFamilyAtoms<R, ER, E>(
75 runtime: Atom.AtomRuntime<EnvironmentRegistry.EnvironmentRegistry | R, ER>,
76 shellStateAtom: (
77 environmentId: EnvironmentId,
78 ) => Atom.Atom<AsyncResult.AsyncResult<EnvironmentShellState, E>>,
79) {
80 const shellStateValue = (state: AsyncResult.AsyncResult<EnvironmentShellState, E>) =>
81 AsyncResult.getOrElse(state, () => null);
82 const generationAtom = Atom.family((environmentId: EnvironmentId) =>
83 Atom.make((get) => shellStateValue(get(shellStateAtom(environmentId)))?.generation ?? 0).pipe(
84 Atom.withLabel(`environment-shell-generation:${environmentId}`),
85 ),
86 );
87
88 // The generation is part of the key, so a read from before the latest authoritative snapshot is
89 // never used again, even while a new read is pending or failed.
90 const familyQueryAtom = createEnvironmentQueryAtomFamily(runtime, {
91 label: "environment-data:orchestration:thread-family",
92 staleTimeMs: Number.POSITIVE_INFINITY,
93 idleTtlMs: 60_000,
94 execute: Effect.fn("ThreadFamilies.read")(function* (input: {
95 readonly threadId: ThreadId;
96 readonly generation: number;
97 }) {
98 const supervisor = yield* EnvironmentSupervisor.EnvironmentSupervisor;
99 const session = yield* SubscriptionRef.get(supervisor.session);
100 if (Option.isNone(session)) {
101 return yield* new EnvironmentRpcUnavailableError({
102 environmentId: supervisor.target.environmentId,
103 message: "The environment has no session.",
104 });
105 }
106 // A config failure fails the read: an empty family would read as "the thread is gone".
107 const config = yield* session.value.initialConfig;
108 // Servers without families send every child in the shell, so absent means gone.
109 if (config.shellThreadFamilies !== true) return EMPTY_FAMILY;
110 return yield* request(ORCHESTRATION_V2_WS_METHODS.getThreadFamily, {
111 threadId: input.threadId,
112 });
113 }),
114 });
115 const queryFor = (environmentId: EnvironmentId, threadId: ThreadId, generation: number) =>
116 familyQueryAtom({ environmentId, input: { threadId, generation } });
117
118 const familyAtom = Atom.family((key: string) => {
119 const [environmentId, threadId] = JSON.parse(key) as [EnvironmentId, ThreadId];
120 return Atom.make((get) => {
121 const generation = get(generationAtom(environmentId));
122 return generation === 0
123 ? AWAITING_AUTHORITATIVE_SHELL
124 : get(queryFor(environmentId, threadId, generation));
125 }).pipe(Atom.withLabel(`environment-thread-family:${key}`));
126 });
127 const familyKey = (environmentId: EnvironmentId, threadId: ThreadId) =>
128 JSON.stringify([environmentId, threadId]);
129
130 /** Reads the current generation's family again, so a failed read does not last the session. */
131 const refresh = (
132 registry: AtomRegistry.AtomRegistry,
133 environmentId: EnvironmentId,
134 threadId: ThreadId,
135 ) => {
136 const generation = registry.get(generationAtom(environmentId));
137 if (generation !== 0) registry.refresh(queryFor(environmentId, threadId, generation));
138 };
139
140 const mountsAtom = Atom.family((environmentId: EnvironmentId) =>
141 Atom.make(EMPTY_MOUNTS).pipe(
142 Atom.keepAlive,
143 Atom.withLabel(`environment-thread-family-mounts:${environmentId}`),
144 ),
145 );
146
147 const familyRowsAtom = Atom.family((key: string) =>
148 Atom.make((get): ReadonlyArray<OrchestrationV2ThreadFamily> => {
149 const family = selectEnvironmentQueryData(get(familyAtom(key)));
150 return family === null ? EMPTY_READS : [family];
151 }).pipe(Atom.withLabel(`environment-thread-family-rows:${key}`)),
152 );
153
154 /**
155 * Merges `rows` behind the shell rows until the returned release runs. Reads read under `key`
156 * share one entry; each read's `snapshotSequence` decides which removals drop its rows.
157 */
158 const mountRows = (
159 registry: AtomRegistry.AtomRegistry,
160 environmentId: EnvironmentId,
161 key: string,
162 rows: Atom.Atom<ReadonlyArray<OrchestrationV2ThreadFamily>>,
163 ): (() => void) => {
164 const atom = mountsAtom(environmentId);
165 const mounts = registry.get(atom);
166 const entry = mounts.get(key);
167 registry.set(
168 atom,
169 new Map(mounts).set(key, { rows: entry?.rows ?? rows, count: (entry?.count ?? 0) + 1 }),
170 );
171 let released = false;
172 return () => {
173 if (released) return;
174 released = true;
175 const next = new Map(registry.get(atom));
176 const entry = next.get(key);
177 if (entry !== undefined && entry.count > 1)
178 next.set(key, { ...entry, count: entry.count - 1 });
179 else next.delete(key);
180 registry.set(atom, next);
181 };
182 };
183
184 /** Keeps a thread's family in the merged rows until the returned release runs. */
185 const mount = (
186 registry: AtomRegistry.AtomRegistry,
187 environmentId: EnvironmentId,
188 threadId: ThreadId,
189 ): (() => void) => {
190 const key = familyKey(environmentId, threadId);
191 const release = mountRows(registry, environmentId, `family:${key}`, familyRowsAtom(key));
192 // Mounting again after a failure retries instead of keeping the failure until reconnect.
193 if (threadFamilyRead(registry.get(familyAtom(key))) === "failed") {
194 refresh(registry, environmentId, threadId);
195 }
196 return release;
197 };
198
199 const snapshotAtom = Atom.family((environmentId: EnvironmentId) => {
200 let previous: {
201 readonly snapshot: OrchestrationV2ShellSnapshot;
202 readonly removedAt: ReadonlyMap<ThreadId, number>;
203 readonly families: ReadonlyArray<OrchestrationV2ThreadFamily>;
204 readonly merged: OrchestrationV2ShellSnapshot;
205 } | null = null;
206 return Atom.make((get): OrchestrationV2ShellSnapshot | null => {
207 const state = shellStateValue(get(shellStateAtom(environmentId)));
208 const snapshot = state === null ? null : Option.getOrNull(state.snapshot);
209 const mounts = get(mountsAtom(environmentId));
210 if (snapshot === null || state === null || mounts.size === 0) return snapshot;
211 // A read from before the latest authoritative snapshot can hold threads removed before it,
212 // and that snapshot cleared those removals; such reads (a cached search) wait for new ones.
213 const cut = snapshot.olderSettled?.asOfSequence ?? 0;
214 const families = [...mounts.values()].flatMap(({ rows }) =>
215 get(rows).filter((read) => read.snapshotSequence >= cut),
216 );
217 // Most shell updates move only the cursor. With the same rows, removals and families, keep
218 // the merged rows so thread lists skip their rebuild.
219 const unchanged =
220 previous !== null &&
221 previous.snapshot.threads === snapshot.threads &&
222 previous.removedAt === state.removedAt &&
223 arrayElementsEqual(previous.families, families);
224 if (unchanged && previous !== null && previous.snapshot === snapshot) return previous.merged;
225 let merged: OrchestrationV2ShellSnapshot;
226 if (unchanged && previous !== null) {
227 merged =
228 previous.merged === previous.snapshot
229 ? snapshot
230 : { ...snapshot, threads: previous.merged.threads };
231 } else {
232 const extras = threadFamilyRowsMissingFromShell({
233 shellThreads: snapshot.threads,
234 removedAt: state.removedAt,
235 families,
236 });
237 merged =
238 extras.length === 0
239 ? snapshot
240 : { ...snapshot, threads: [...snapshot.threads, ...extras] };
241 }
242 previous = { snapshot, removedAt: state.removedAt, families, merged };
243 return merged;
244 }).pipe(Atom.withLabel(`environment-snapshot-with-families:${environmentId}`));
245 });
246
247 return {
248 /** The shell snapshot plus the rows of every mounted family. */
249 snapshotAtom,
250 /** The current generation's family read, for decisions that need a settled answer. */
251 familyAtom: (environmentId: EnvironmentId, threadId: ThreadId) =>
252 familyAtom(familyKey(environmentId, threadId)),
253 mount,
254 mountRows,
255 refresh,
256 };
257}