1// @effect-diagnostics nodeBuiltinImport:off -- unit names and RAM size come from the host.
2import * as NodeCrypto from "node:crypto";
3import * as NodeOS from "node:os";
4
5import {
6 AgentScope,
7 agentScopeCommand,
8 type AgentScopeShape,
9 type AgentSystemdScope,
10} from "@t3tools/shared/AgentScope";
11import * as HostProcess from "@t3tools/shared/HostProcess";
12import { resolveCommandPath } from "@t3tools/shared/shell";
13import * as Effect from "effect/Effect";
14import * as FileSystem from "effect/FileSystem";
15import * as Layer from "effect/Layer";
16import * as Path from "effect/Path";
17
18import * as ServerConfig from "../config.ts";
19import * as ProcessRunner from "../processRunner.ts";
20
21/**
22 * Agents and terminals share this slice. It sits in app.slice next to the
23 * server's own unit, so the server's CPUWeight and IOWeight compete with the
24 * agents directly.
25 */
26const AGENT_SLICE = "app-t3code-agents.slice";
27
28/**
29 * The child slice for one server's agents, named after its state directory.
30 * Only one server uses a state directory at a time, so an agent in this slice
31 * when the server starts was left by an earlier server on the same state,
32 * for example one the launcher killed during an update. Other servers on the
33 * machine, such as dev servers, have their own slice.
34 */
35export function ownedAgentSlice(stateDir: string): string {
36 const hash = NodeCrypto.createHash("sha256").update(stateDir).digest("hex").slice(0, 12);
37 return `app-t3code-agents-${hash}.slice`;
38}
39
40const GiB = 1024 ** 3;
41
42/**
43 * Memory limits for the agents slice. Agents together leave about 6 GB free
44 * before they are throttled and 4 GB before the kernel kills one, so the
45 * server, sshd, and tailscaled keep working. Small machines still give agents
46 * at least half of RAM before throttling and three quarters before a kill.
47 */
48export function agentSliceMemoryLimits(totalBytes: number) {
49 return {
50 high: Math.max(totalBytes - 6 * GiB, Math.floor(totalBytes / 2)),
51 max: Math.max(totalBytes - 4 * GiB, Math.floor(totalBytes * 0.75)),
52 };
53}
54
55export interface ScopeState {
56 readonly loadState: string;
57 readonly activeState: string;
58 readonly result: string;
59 /** `oom_kill` from the scope's memory.events, while its cgroup exists. */
60 readonly oomKills: number;
61 /** `populated` from the scope's cgroup.events, while its cgroup exists. */
62 readonly populated: boolean | undefined;
63}
64
65/**
66 * Reads one scope as the server sees it right after the agent process exits.
67 * Our pipe can close before systemd handles the kill, so a scope that is
68 * still stopping, is empty, or already counted an OOM kill is "stopping":
69 * systemd records `Result=oom-kill` for it in a moment.
70 */
71export function classifyScope(state: ScopeState): "oom-killed" | "gone" | "running" | "stopping" {
72 if (state.result === "oom-kill") return "oom-killed";
73 if (
74 state.loadState !== "loaded" ||
75 state.activeState === "inactive" ||
76 state.activeState === "failed"
77 ) {
78 return "gone";
79 }
80 if (state.activeState === "deactivating" || state.oomKills > 0 || state.populated === false) {
81 return "stopping";
82 }
83 return "running";
84}
85
86function parseKeyValues(text: string): Map<string, string> {
87 const values = new Map<string, string>();
88 for (const line of text.split("\n")) {
89 const index = line.indexOf("=");
90 if (index > 0) values.set(line.slice(0, index), line.slice(index + 1).trim());
91 }
92 return values;
93}
94
95function parseCgroupCounter(text: string, key: string): number | undefined {
96 const match = new RegExp(`^${key} (\\d+)$`, "m").exec(text);
97 return match === null ? undefined : Number(match[1]);
98}
99
100// A stop sends SIGTERM, then SIGKILL after this long. systemd's 90 s default
101// would outlast the wait for a previous server's agents below.
102const STOP_TIMEOUT_OPTION = "--property=TimeoutStopSec=10s";
103const STOPPING_POLL_INTERVAL = "100 millis";
104const STOPPING_POLL_ATTEMPTS = 20;
105// A stopping scope can wait out its TimeoutStopSec for a stuck process.
106const CLEAR_FAILED_INTERVAL = "1 second";
107const CLEAR_FAILED_ATTEMPTS = 120;
108const MAX_TRACKED_THREADS = 500;
109
110export const make = Effect.gen(function* () {
111 const platform = yield* HostProcess.Platform;
112 if (platform !== "linux") return AgentScope.defaultValue();
113
114 const runner = yield* ProcessRunner.ProcessRunner;
115 const fileSystem = yield* FileSystem.FileSystem;
116 const path = yield* Path.Path;
117 const hostEnvironment = yield* HostProcess.Environment;
118 const { stateDir } = yield* ServerConfig.ServerConfig;
119 const ownedSlice = ownedAgentSlice(stateDir);
120 const layerScope = yield* Effect.scope;
121
122 const resolve = (command: string, env: NodeJS.ProcessEnv) =>
123 resolveCommandPath(command, { env }).pipe(
124 Effect.provideService(FileSystem.FileSystem, fileSystem),
125 Effect.provideService(Path.Path, path),
126 Effect.provideService(HostProcess.Platform, platform),
127 );
128
129 const run = (command: string, args: ReadonlyArray<string>) =>
130 runner.run({ command, args, timeout: "5 seconds" }).pipe(Effect.option);
131
132 const readCgroupFile = (controlGroup: string, file: string) =>
133 fileSystem
134 .readFileString(path.join("/sys/fs/cgroup", controlGroup, file))
135 .pipe(Effect.orElseSucceed(() => ""));
136
137 // Old agents keep their scopes when the server is killed, and the restarted
138 // server resumes their sessions. Stop them first, so two copies of one
139 // session never work in the same worktree.
140 const stopLeftoverAgents = (systemctl: string) =>
141 Effect.gen(function* () {
142 const shown = yield* run(systemctl, [
143 "--user",
144 "show",
145 ownedSlice,
146 "--property=ActiveState,ControlGroup",
147 ]);
148 if (shown._tag === "None") return;
149 const values = parseKeyValues(shown.value.stdout);
150 if (values.get("ActiveState") !== "active") return;
151 // An active slice can be empty once its last agent exited.
152 const controlGroup = values.get("ControlGroup") ?? "";
153 const events = controlGroup ? yield* readCgroupFile(controlGroup, "cgroup.events") : "";
154 if (parseCgroupCounter(events, "populated") === 0) return;
155 yield* Effect.logWarning("Stopping agents left by a previous server", { slice: ownedSlice });
156 // Agents get SIGTERM, then SIGKILL after their stop timeout.
157 const stopped = yield* runner
158 .run({ command: systemctl, args: ["--user", "stop", ownedSlice], timeout: "30 seconds" })
159 .pipe(Effect.option);
160 if (stopped._tag === "None" || stopped.value.code !== 0) {
161 yield* Effect.logWarning("Could not stop agents left by a previous server", {
162 slice: ownedSlice,
163 stderr: stopped._tag === "Some" ? stopped.value.stderr.trim() : undefined,
164 });
165 }
166 });
167
168 // Runs once, on the first agent or terminal launch. It proves this process
169 // can reach a systemd user manager, stops agents a previous server left,
170 // creates the agents slice, and sets its limits.
171 const systemd = yield* Effect.cached(
172 Effect.gen(function* () {
173 const runtimeDir = hostEnvironment.XDG_RUNTIME_DIR;
174 if (runtimeDir === undefined || runtimeDir.length === 0) return undefined;
175 const systemdRun = yield* resolve("systemd-run", hostEnvironment);
176 const systemctl = yield* resolve("systemctl", hostEnvironment);
177 if (systemdRun === null || systemctl === null) return undefined;
178 yield* stopLeftoverAgents(systemctl);
179 const probe = (options: ReadonlyArray<string>) =>
180 run(systemdRun, [
181 "--user",
182 "--scope",
183 "--quiet",
184 "--collect",
185 `--slice=${AGENT_SLICE}`,
186 ...options,
187 "--",
188 "true",
189 ]).pipe(Effect.map((result) => result._tag === "Some" && result.value.code === 0));
190 // Stop the whole scope when the kernel kills one of its processes, and
191 // record the result as oom-kill. Scopes take OOMPolicy from systemd 253;
192 // older user managers still get scopes, without the OOM label.
193 const stopOptions = [STOP_TIMEOUT_OPTION];
194 const oomOptions = [...stopOptions, "--property=OOMPolicy=stop"];
195 const scopeOptions = (yield* probe(oomOptions))
196 ? oomOptions
197 : (yield* probe(stopOptions))
198 ? stopOptions
199 : undefined;
200 if (scopeOptions === undefined) {
201 yield* Effect.logInfo("Agent scopes are off: no systemd user manager");
202 return undefined;
203 }
204 const limits = agentSliceMemoryLimits(NodeOS.totalmem());
205 const limited = yield* run(systemctl, [
206 "--user",
207 "set-property",
208 "--runtime",
209 AGENT_SLICE,
210 `MemoryHigh=${limits.high}`,
211 `MemoryMax=${limits.max}`,
212 ]);
213 if (limited._tag === "None" || limited.value.code !== 0) {
214 yield* Effect.logWarning("Could not set memory limits on the agents slice", {
215 stderr: limited._tag === "Some" ? limited.value.stderr.trim() : undefined,
216 });
217 }
218 yield* Effect.logInfo("Agent scopes are on", {
219 slice: AGENT_SLICE,
220 ...limits,
221 oomPolicy: scopeOptions.length > 0,
222 });
223 return {
224 systemdRun,
225 systemctl,
226 runtimeDir,
227 scopeOptions,
228 };
229 }),
230 );
231
232 // The newest scope per thread, oldest thread first. A thread runs one agent
233 // at a time, so only its latest scope can explain a failure.
234 const scopeByThread = new Map<string, { unit: string; oomKilled: boolean }>();
235
236 const readScope = (systemctl: string, unit: string) =>
237 Effect.gen(function* () {
238 const shown = yield* run(systemctl, [
239 "--user",
240 "show",
241 unit,
242 "--property=LoadState,ActiveState,Result,ControlGroup",
243 ]);
244 // A failed query says nothing about the scope; the caller keeps the unit.
245 if (shown._tag === "None" || shown.value.code !== 0) return undefined;
246 const values = parseKeyValues(shown.value.stdout);
247 const controlGroup = values.get("ControlGroup") ?? "";
248 const memoryEvents = controlGroup ? yield* readCgroupFile(controlGroup, "memory.events") : "";
249 const cgroupEvents = controlGroup ? yield* readCgroupFile(controlGroup, "cgroup.events") : "";
250 const populated = parseCgroupCounter(cgroupEvents, "populated");
251 return {
252 loadState: values.get("LoadState") ?? "",
253 activeState: values.get("ActiveState") ?? "",
254 result: values.get("Result") ?? "",
255 oomKills: parseCgroupCounter(memoryEvents, "oom_kill") ?? 0,
256 populated: populated === undefined ? undefined : populated === 1,
257 } satisfies ScopeState;
258 });
259
260 const clearFailed = (systemctl: string, unit: string) =>
261 Effect.gen(function* () {
262 for (let attempt = 0; attempt < CLEAR_FAILED_ATTEMPTS; attempt++) {
263 const state = yield* readScope(systemctl, unit);
264 if (state === undefined) return;
265 if (state.activeState === "failed") {
266 yield* run(systemctl, ["--user", "reset-failed", unit]);
267 return;
268 }
269 if (state.activeState !== "deactivating") return;
270 yield* Effect.sleep(CLEAR_FAILED_INTERVAL);
271 }
272 });
273
274 // Failed scopes stay loaded so we can read their result. Clear them once
275 // systemd finishes stopping them.
276 const forget = (systemctl: string, unit: string) =>
277 clearFailed(systemctl, unit).pipe(Effect.forkIn(layerScope), Effect.asVoid);
278
279 // Drops the oldest thread's scope once it has ended. A scope that still
280 // runs moves to the back, so its thread keeps its OOM answer and cleanup.
281 const evictOldest = (systemctl: string) =>
282 Effect.gen(function* () {
283 const oldest = scopeByThread.entries().next();
284 if (oldest.done === true) return;
285 const [threadId, tracked] = oldest.value;
286 const state = yield* readScope(systemctl, tracked.unit);
287 // A newer launch replaced it while we read, and forgets it itself.
288 if (scopeByThread.get(threadId) !== tracked) return;
289 scopeByThread.delete(threadId);
290 const status = state === undefined ? "running" : classifyScope(state);
291 if (status === "running" || status === "stopping") {
292 scopeByThread.set(threadId, tracked);
293 return;
294 }
295 yield* forget(systemctl, tracked.unit);
296 });
297
298 // Records a thread's newest scope, or forgets the thread when `unit` is
299 // undefined. The replaced scope has its failed state cleared.
300 const track = (systemctl: string, threadId: string, unit: string | undefined) =>
301 Effect.gen(function* () {
302 const previous = scopeByThread.get(threadId);
303 scopeByThread.delete(threadId);
304 if (unit !== undefined) scopeByThread.set(threadId, { unit, oomKilled: false });
305 if (previous !== undefined) yield* forget(systemctl, previous.unit);
306 if (scopeByThread.size > MAX_TRACKED_THREADS) yield* evictOldest(systemctl);
307 });
308
309 const oomKilled: AgentScopeShape["oomKilled"] = (threadId) =>
310 Effect.gen(function* () {
311 const tracked = scopeByThread.get(threadId);
312 if (tracked === undefined) return false;
313 if (tracked.oomKilled) return true;
314 const scope = yield* systemd;
315 if (scope === undefined) return false;
316 for (let attempt = 0; attempt < STOPPING_POLL_ATTEMPTS; attempt++) {
317 const state = yield* readScope(scope.systemctl, tracked.unit);
318 if (state === undefined) return false;
319 const status = classifyScope(state);
320 if (status === "running") return false;
321 if (status === "stopping") {
322 yield* Effect.sleep(STOPPING_POLL_INTERVAL);
323 continue;
324 }
325 yield* forget(scope.systemctl, tracked.unit);
326 if (status === "gone") return false;
327 yield* Effect.logWarning("Agent scope was killed: out of memory", {
328 threadId,
329 unit: tracked.unit,
330 });
331 tracked.oomKilled = true;
332 return true;
333 }
334 return false;
335 });
336
337 const clear: AgentScopeShape["clear"] = (threadId) =>
338 Effect.gen(function* () {
339 if (!scopeByThread.has(threadId)) return;
340 const scope = yield* systemd;
341 if (scope !== undefined) yield* track(scope.systemctl, threadId, undefined);
342 });
343
344 return {
345 wrap: Effect.fn("AgentScope.wrap")(function* ({
346 command,
347 args,
348 name,
349 threadId,
350 env,
351 survivesRestart,
352 }) {
353 const resolved = yield* resolve(command, env ?? hostEnvironment);
354 // Leave unresolved commands alone so the spawn reports the missing binary.
355 if (resolved === null) return { command, args };
356 const scope = yield* systemd;
357 if (scope === undefined) return agentScopeCommand({ command: resolved, args });
358 const unit = `t3code-${name}-${NodeCrypto.randomUUID().slice(0, 8)}.scope`;
359 if (threadId !== undefined) yield* track(scope.systemctl, threadId, unit);
360 const systemdScope: AgentSystemdScope = {
361 systemdRun: scope.systemdRun,
362 runtimeDir: scope.runtimeDir,
363 slice: survivesRestart === true ? AGENT_SLICE : ownedSlice,
364 unit,
365 // Without a thread nothing reads the result, so let systemd unload
366 // the unit even when it fails.
367 options: threadId === undefined ? [...scope.scopeOptions, "--collect"] : scope.scopeOptions,
368 };
369 return agentScopeCommand({ command: resolved, args, scope: systemdScope });
370 }),
371 oomKilled,
372 clear,
373 } satisfies AgentScopeShape;
374});
375
376export const layer = Layer.effect(AgentScope, make).pipe(Layer.provide(ProcessRunner.layer));