1import {
2 ORCHESTRATION_V2_WS_METHODS,
3 type EnvironmentId,
4 type ThreadId,
5 type OrchestrationV2ThreadLinkMetadata,
6 type OrchestrationGetTurnDiffResult,
7 type OrchestrationV2GetTurnItemInput,
8} from "@t3tools/contracts";
9import * as Effect from "effect/Effect";
10import * as Option from "effect/Option";
11import * as Stream from "effect/Stream";
12import * as SubscriptionRef from "effect/SubscriptionRef";
13import type { HttpClient } from "effect/http";
14import { Atom } from "effect/reactivity";
15
16import {
17 applyCommandOutputFrame,
18 EMPTY_COMMAND_OUTPUT,
19 type CommandOutputView,
20} from "./commandOutput.ts";
21
22import {
23 createEnvironmentQueryAtomFamily,
24 createEnvironmentRpcCommand,
25 createEnvironmentRpcQueryAtomFamily,
26 createEnvironmentRpcSubscriptionAtomFamily,
27 createEnvironmentCommand,
28 type QueryAnswerHandler,
29} from "./runtime.ts";
30import type { EnvironmentRegistry } from "../connection/registry.ts";
31import * as EnvironmentSupervisor from "../connection/supervisor.ts";
32import * as RemoteEnvironmentAuthorization from "../authorization/service.ts";
33import * as ManagedRelayDpopSigner from "../relay/managedRelay.ts";
34import { EnvironmentRpcUnavailableError, request } from "../rpc/client.ts";
35import type { EnvironmentShellState } from "./shell.ts";
36import { fetchEnvironmentThreadTranscript } from "./threadTranscriptHttp.ts";
37
38export function createOrchestrationEnvironmentAtoms<R, E>(
39 runtime: Atom.AtomRuntime<EnvironmentRegistry | HttpClient.HttpClient | R, E>,
40 options: {
41 /** Re-checks `threadTitles` when one of their threads leaves the shell. */
42 readonly shellState?: (
43 environmentId: EnvironmentId,
44 ) => Atom.Atom<Pick<EnvironmentShellState, "removedAt">>;
45 /** Runs on each turn-range and full-thread diff the server answers. */
46 readonly onCheckpointDiffAnswer?: QueryAnswerHandler<
47 { readonly ignoreWhitespace?: boolean },
48 OrchestrationGetTurnDiffResult
49 >;
50 } = {},
51) {
52 const onCheckpointDiffAnswer =
53 options.onCheckpointDiffAnswer === undefined
54 ? {}
55 : { onAnswer: options.onCheckpointDiffAnswer };
56 const shellState = options.shellState;
57 const liveTurnItem = createEnvironmentRpcQueryAtomFamily(runtime, {
58 label: "environment-data:orchestration:turn-item",
59 tag: ORCHESTRATION_V2_WS_METHODS.getTurnItem,
60 staleTimeMs: 60_000,
61 idleTtlMs: 60_000,
62 });
63 // A settled item's revision is its updatedAt, so the answer for it never goes stale.
64 const settledTurnItem = createEnvironmentRpcQueryAtomFamily(runtime, {
65 label: "environment-data:orchestration:turn-item",
66 tag: ORCHESTRATION_V2_WS_METHODS.getTurnItem,
67 staleTimeMs: Number.POSITIVE_INFINITY,
68 idleTtlMs: 5 * 60_000,
69 });
70 const refreshThreadLookup = ({
71 environmentId,
72 input,
73 }: {
74 readonly environmentId: EnvironmentId;
75 readonly input: { readonly threadIds: ReadonlyArray<ThreadId> };
76 }) => {
77 if (shellState === undefined) return undefined;
78 const state = shellState(environmentId);
79 return Atom.make((get) => {
80 const { removedAt } = get(state);
81 return Math.max(-1, ...input.threadIds.map((threadId) => removedAt.get(threadId) ?? -1));
82 });
83 };
84 return {
85 threadTitles: createEnvironmentRpcQueryAtomFamily(runtime, {
86 label: "environment-data:orchestration:thread-titles",
87 tag: ORCHESTRATION_V2_WS_METHODS.getThreadTitles,
88 staleTimeMs: 0,
89 idleTtlMs: 0,
90 // Deleting a thread the shell already left out keeps these ids, so the query would not rerun.
91 // A removal can also be an archive, so it only prompts a re-check.
92 refreshTrigger: refreshThreadLookup,
93 }),
94 threadLinkMetadata: createEnvironmentQueryAtomFamily(runtime, {
95 label: "environment-data:orchestration:thread-link-metadata",
96 staleTimeMs: 0,
97 idleTtlMs: 0,
98 refreshTrigger: refreshThreadLookup,
99 execute: Effect.fn("orchestration.threadLinkMetadata")(function* (input: {
100 readonly threadIds: ReadonlyArray<ThreadId>;
101 }) {
102 const supervisor = yield* EnvironmentSupervisor.EnvironmentSupervisor;
103 const session = yield* SubscriptionRef.get(supervisor.session);
104 if (Option.isNone(session)) {
105 return yield* new EnvironmentRpcUnavailableError({
106 environmentId: supervisor.target.environmentId,
107 message: "The environment has no session.",
108 });
109 }
110 const config = yield* session.value.initialConfig;
111 if (config.threadLinkMetadata === true) {
112 return yield* request(ORCHESTRATION_V2_WS_METHODS.getThreadLinkMetadata, input);
113 }
114 const titles = yield* request(ORCHESTRATION_V2_WS_METHODS.getThreadTitles, input);
115 return Object.fromEntries(
116 Object.entries(titles).map(([id, title]) => [
117 id,
118 { title, projectId: null } satisfies OrchestrationV2ThreadLinkMetadata,
119 ]),
120 );
121 }),
122 }),
123 threadTranscript: createEnvironmentCommand(runtime, {
124 label: "environment-data:orchestration:thread-transcript",
125 execute: Effect.fn("orchestration.execute")(function* (input: {
126 readonly threadId: ThreadId;
127 }) {
128 const supervisor = yield* EnvironmentSupervisor.EnvironmentSupervisor;
129 const prepared = yield* SubscriptionRef.get(supervisor.prepared);
130 if (Option.isNone(prepared)) {
131 return yield* new EnvironmentRpcUnavailableError({
132 environmentId: supervisor.target.environmentId,
133 message: "The source environment is not connected.",
134 });
135 }
136 return yield* fetchEnvironmentThreadTranscript({
137 prepared: prepared.value,
138 threadId: input.threadId,
139 signer: yield* Effect.serviceOption(ManagedRelayDpopSigner.ManagedRelayDpopSigner),
140 remoteAuthorization: yield* Effect.serviceOption(
141 RemoteEnvironmentAuthorization.RemoteEnvironmentAuthorization,
142 ),
143 });
144 }),
145 }),
146 v2: {
147 dispatchCommand: createEnvironmentRpcCommand(runtime, {
148 label: "environment-data:orchestration-v2:dispatch-command",
149 tag: ORCHESTRATION_V2_WS_METHODS.dispatchCommand,
150 }),
151 shell: createEnvironmentRpcSubscriptionAtomFamily(runtime, {
152 label: "environment-data:orchestration-v2:shell",
153 tag: ORCHESTRATION_V2_WS_METHODS.subscribeShell,
154 }),
155 thread: createEnvironmentRpcSubscriptionAtomFamily(runtime, {
156 label: "environment-data:orchestration-v2:thread",
157 tag: ORCHESTRATION_V2_WS_METHODS.subscribeThread,
158 idleTtlMs: 0,
159 }),
160 // One expanded command row's output. A running command's subscription ends
161 // as soon as its row closes, so hidden rows receive nothing; a finished
162 // command's final output stays cached so reopening it does not refetch.
163 commandOutput: createEnvironmentRpcSubscriptionAtomFamily(runtime, {
164 label: "environment-data:orchestration-v2:command-output",
165 tag: ORCHESTRATION_V2_WS_METHODS.subscribeCommandOutput,
166 idleTtlMs: 0,
167 completeWhen: (view: CommandOutputView) => !view.running,
168 completedTtlMs: 5 * 60_000,
169 // `scan` emits its seed first; skipping it keeps the row "loading" until the first frame.
170 transform: (frames) =>
171 frames.pipe(
172 Stream.scan(() => EMPTY_COMMAND_OUTPUT, applyCommandOutputFrame),
173 Stream.filter((view) => view !== EMPTY_COMMAND_OUTPUT),
174 ),
175 }),
176 // The open thread's predicted next prompt, while its composer is mounted.
177 promptSuggestion: createEnvironmentRpcSubscriptionAtomFamily(runtime, {
178 label: "environment-data:orchestration-v2:prompt-suggestion",
179 tag: ORCHESTRATION_V2_WS_METHODS.subscribePromptSuggestion,
180 idleTtlMs: 0,
181 }),
182 },
183 turnDiff: createEnvironmentRpcQueryAtomFamily(runtime, {
184 label: "environment-data:orchestration:turn-diff",
185 tag: ORCHESTRATION_V2_WS_METHODS.getTurnDiff,
186 ...onCheckpointDiffAnswer,
187 }),
188 workflowScript: createEnvironmentRpcQueryAtomFamily(runtime, {
189 label: "environment-data:orchestration:workflow-script",
190 tag: ORCHESTRATION_V2_WS_METHODS.getWorkflowScript,
191 staleTimeMs: 300_000,
192 idleTtlMs: 300_000,
193 }),
194 // Mounted only while a subagent inspector is open; running children change between opens.
195 nativeChildWork: createEnvironmentRpcQueryAtomFamily(runtime, {
196 label: "environment-data:orchestration:native-child-work",
197 tag: ORCHESTRATION_V2_WS_METHODS.getNativeChildWork,
198 staleTimeMs: 0,
199 idleTtlMs: 0,
200 }),
201 stopWorkflow: createEnvironmentRpcCommand(runtime, {
202 label: "environment-data:orchestration:stop-workflow",
203 tag: ORCHESTRATION_V2_WS_METHODS.stopWorkflow,
204 }),
205 forceStopThread: createEnvironmentRpcCommand(runtime, {
206 label: "environment-data:orchestration:force-stop-thread",
207 tag: ORCHESTRATION_V2_WS_METHODS.forceStopThread,
208 }),
209 turnItem: (target: {
210 readonly environmentId: EnvironmentId;
211 readonly input: OrchestrationV2GetTurnItemInput;
212 }) =>
213 target.input.revision === undefined || target.input.revision === "live"
214 ? liveTurnItem(target)
215 : settledTurnItem(target),
216 fullThreadDiff: createEnvironmentRpcQueryAtomFamily(runtime, {
217 label: "environment-data:orchestration:full-thread-diff",
218 tag: ORCHESTRATION_V2_WS_METHODS.getFullThreadDiff,
219 ...onCheckpointDiffAnswer,
220 }),
221 threadFind: createEnvironmentRpcQueryAtomFamily(runtime, {
222 label: "environment-data:orchestration:thread-find",
223 tag: ORCHESTRATION_V2_WS_METHODS.searchThread,
224 staleTimeMs: 0,
225 idleTtlMs: 0,
226 }),
227 threadFindProgressive: createEnvironmentRpcSubscriptionAtomFamily(runtime, {
228 label: "environment-data:orchestration:thread-find-progressive",
229 tag: ORCHESTRATION_V2_WS_METHODS.searchThreadStream,
230 completeWhen: (result) => result.complete !== false,
231 idleTtlMs: 0,
232 }),
233 threadSearch: createEnvironmentRpcQueryAtomFamily(runtime, {
234 label: "environment-data:orchestration:thread-search",
235 tag: ORCHESTRATION_V2_WS_METHODS.searchThreads,
236 staleTimeMs: 30_000,
237 idleTtlMs: 60_000,
238 }),
239 archivedShellSnapshot: createEnvironmentRpcQueryAtomFamily(runtime, {
240 label: "environment-data:orchestration:archived-shell-snapshot",
241 tag: ORCHESTRATION_V2_WS_METHODS.getArchivedShellSnapshot,
242 }),
243 threadFamily: createEnvironmentRpcQueryAtomFamily(runtime, {
244 label: "environment-data:orchestration:thread-family-for-delete",
245 tag: ORCHESTRATION_V2_WS_METHODS.getThreadFamily,
246 staleTimeMs: 0,
247 idleTtlMs: 0,
248 }),
249 // Read right before deleting a worktree, so never served from cache.
250 worktreeThreadIds: createEnvironmentRpcQueryAtomFamily(runtime, {
251 label: "environment-data:orchestration:worktree-thread-ids",
252 tag: ORCHESTRATION_V2_WS_METHODS.getWorktreeThreadIds,
253 staleTimeMs: 0,
254 idleTtlMs: 0,
255 }),
256 };
257}