omnirush 0.9.0 → 0.9.1
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/assets/CHANGELOG.md +35 -0
- package/assets/extensions/omnirush/agents-lib.ts +370 -53
- package/assets/extensions/omnirush/agents.ts +92 -11
- package/assets/extensions/omnirush/bgshell-lib.ts +32 -0
- package/assets/extensions/omnirush/capture/workspace-collector.ts +86 -22
- package/assets/extensions/omnirush/collector.ts +203 -32
- package/assets/extensions/omnirush/index.ts +22 -8
- package/assets/extensions/omnirush/mcp.ts +103 -8
- package/assets/extensions/omnirush/memory-lib.ts +516 -0
- package/assets/extensions/omnirush/pi-engine.ts +50 -2
- package/assets/extensions/omnirush/plan-lib.ts +38 -9
- package/assets/extensions/omnirush/plan.ts +85 -30
- package/assets/extensions/omnirush/subagent-marker.ts +21 -0
- package/package.json +1 -1
- package/scripts/smoke-packages.py +1 -0
- package/src/bin.js +28 -1
- package/src/sessions.js +222 -0
|
@@ -32,10 +32,20 @@
|
|
|
32
32
|
// session reaches the parent's trace as session.child events, like the
|
|
33
33
|
// desktop's task sub-agents.
|
|
34
34
|
//
|
|
35
|
+
// A turn's prompt is only what the user submitted. A sub-agent process
|
|
36
|
+
// marks its own session (an `omnirush-subagent` entry naming its parent),
|
|
37
|
+
// so when that session is later opened as a main session (`--continue`,
|
|
38
|
+
// `--resume`, `--session`) its first user message, the task the parent
|
|
39
|
+
// agent wrote, is recorded as synthetic; for sub-agents that wrote no
|
|
40
|
+
// marker, the parent session file naming it in a spawn result confirms
|
|
41
|
+
// it. A prompt an extension sends (pi's input event, source "extension")
|
|
42
|
+
// is synthetic too.
|
|
43
|
+
//
|
|
35
44
|
// Silent no-op without device credentials: collection only runs for
|
|
36
45
|
// signed-in omnirush identities.
|
|
37
46
|
|
|
38
47
|
import { createHash } from "node:crypto";
|
|
48
|
+
import markSubagentSession from "./subagent-marker";
|
|
39
49
|
import { createReadStream, readdirSync, readFileSync, rmSync } from "node:fs";
|
|
40
50
|
import { readdir, stat } from "node:fs/promises";
|
|
41
51
|
import os from "node:os";
|
|
@@ -51,13 +61,16 @@ import {
|
|
|
51
61
|
engineMessagesFromEntries,
|
|
52
62
|
MAX_TURN_MESSAGE_BYTES,
|
|
53
63
|
messageParts,
|
|
64
|
+
subagentTask,
|
|
65
|
+
SUBAGENT_ENTRY_TYPE,
|
|
54
66
|
turnHistory,
|
|
55
67
|
turnModelFromMessages,
|
|
68
|
+
userText,
|
|
56
69
|
type EngineMessage,
|
|
57
70
|
type PiEntry,
|
|
58
71
|
} from "./pi-engine";
|
|
59
72
|
import { sharedRefresher } from "./refresh";
|
|
60
|
-
import { AGENTS_HOOK, type AgentsHook } from "./agents-lib";
|
|
73
|
+
import { AGENTS_HOOK, SUBAGENT_SESSIONS_DIR, type AgentsHook } from "./agents-lib";
|
|
61
74
|
import { BGSHELL_HOOK, type BgShellHook } from "./bgshell-lib";
|
|
62
75
|
|
|
63
76
|
/** info.agent of the root session's messages (a sub-agent's is its role). */
|
|
@@ -203,6 +216,71 @@ async function findSessionFile(sessionId: string, dirs: readonly string[]): Prom
|
|
|
203
216
|
return null;
|
|
204
217
|
}
|
|
205
218
|
|
|
219
|
+
/** Session files read at most when looking for a sub-agent's parent, and for how long. */
|
|
220
|
+
const MAX_PARENT_SCAN_FILES = 64;
|
|
221
|
+
const PARENT_SCAN_BUDGET_MS = 3_000;
|
|
222
|
+
|
|
223
|
+
/** Whether `file` contains `needle` (streamed in chunks, never held whole). */
|
|
224
|
+
async function fileContains(file: string, needle: string, deadline: number): Promise<boolean> {
|
|
225
|
+
const input = createReadStream(file, { encoding: "utf8", highWaterMark: 256 * 1024 });
|
|
226
|
+
let tail = "";
|
|
227
|
+
try {
|
|
228
|
+
for await (const chunk of input) {
|
|
229
|
+
const text = tail + (chunk as string);
|
|
230
|
+
if (text.includes(needle)) return true;
|
|
231
|
+
if (Date.now() > deadline) return false;
|
|
232
|
+
tail = text.slice(-needle.length);
|
|
233
|
+
}
|
|
234
|
+
} catch {
|
|
235
|
+
/* unreadable: not the parent */
|
|
236
|
+
} finally {
|
|
237
|
+
input.destroy();
|
|
238
|
+
}
|
|
239
|
+
return false;
|
|
240
|
+
}
|
|
241
|
+
|
|
242
|
+
/**
|
|
243
|
+
* The session that started sub-agent `childId`: the session file in `dir`
|
|
244
|
+
* whose spawn_agents (or agents_*) results name it, among the files
|
|
245
|
+
* written since the sub-agent started (`since`, epoch ms; the newest
|
|
246
|
+
* first, at most MAX_PARENT_SCAN_FILES, within PARENT_SCAN_BUDGET_MS).
|
|
247
|
+
*/
|
|
248
|
+
export async function findSpawningSession(childId: string, dir: string, since: number): Promise<string | null> {
|
|
249
|
+
const deadline = Date.now() + PARENT_SCAN_BUDGET_MS;
|
|
250
|
+
const found = await findSpawningSessionIn(childId, dir, since, deadline);
|
|
251
|
+
if (found) return found;
|
|
252
|
+
// A sub-agent session moved into <session dir>/subagents/ (src/sessions.js,
|
|
253
|
+
// or written there by a newer parent): its parent sits one level up.
|
|
254
|
+
if (path.basename(dir) === SUBAGENT_SESSIONS_DIR) {
|
|
255
|
+
return findSpawningSessionIn(childId, path.dirname(dir), since, deadline);
|
|
256
|
+
}
|
|
257
|
+
return null;
|
|
258
|
+
}
|
|
259
|
+
|
|
260
|
+
async function findSpawningSessionIn(childId: string, dir: string, since: number, deadline: number): Promise<string | null> {
|
|
261
|
+
let names: string[];
|
|
262
|
+
try {
|
|
263
|
+
names = await readdir(dir);
|
|
264
|
+
} catch {
|
|
265
|
+
return null;
|
|
266
|
+
}
|
|
267
|
+
const candidates: Array<{ file: string; id: string; mtime: number }> = [];
|
|
268
|
+
for (const name of names) {
|
|
269
|
+
const match = /_([A-Za-z0-9._:-]{8,128})\.jsonl$/.exec(name);
|
|
270
|
+
if (!match || match[1] === childId) continue;
|
|
271
|
+
const file = path.join(dir, name);
|
|
272
|
+
const stats = await stat(file).catch(() => null);
|
|
273
|
+
if (stats && stats.mtimeMs >= since) candidates.push({ file, id: match[1]!, mtime: stats.mtimeMs });
|
|
274
|
+
}
|
|
275
|
+
candidates.sort((a, b) => b.mtime - a.mtime);
|
|
276
|
+
const needle = `"session_id":${JSON.stringify(childId)}`;
|
|
277
|
+
for (const candidate of candidates.slice(0, MAX_PARENT_SCAN_FILES)) {
|
|
278
|
+
if (Date.now() > deadline) break;
|
|
279
|
+
if (await fileContains(candidate.file, needle, deadline)) return candidate.id;
|
|
280
|
+
}
|
|
281
|
+
return null;
|
|
282
|
+
}
|
|
283
|
+
|
|
206
284
|
/** Parse a session JSONL file without buffering the raw transcript or its line array. */
|
|
207
285
|
export async function parseSessionFileStream(file: string): Promise<{ header: PiEntry | null; entries: PiEntry[] }> {
|
|
208
286
|
let header: PiEntry | null = null;
|
|
@@ -253,7 +331,8 @@ export async function captureChildSessions(input: {
|
|
|
253
331
|
visited: Set<string>;
|
|
254
332
|
/** Unchanged session files are not read again in this process. */
|
|
255
333
|
seenFiles: Map<string, string>;
|
|
256
|
-
|
|
334
|
+
/** Before each event is recorded (may wait: the caller flushes and drains the upload queue). */
|
|
335
|
+
beforeRecord: (bytes: number) => void | Promise<void>;
|
|
257
336
|
}): Promise<void> {
|
|
258
337
|
if (input.depth > MAX_CHILD_DEPTH) return;
|
|
259
338
|
for (const child of spawnedChildren(input.parentEntries).slice(0, MAX_CHILDREN)) {
|
|
@@ -291,7 +370,7 @@ export async function captureChildSessions(input: {
|
|
|
291
370
|
});
|
|
292
371
|
}
|
|
293
372
|
if (history.delta.length > 0 || !input.known.has(child.id)) {
|
|
294
|
-
input.beforeRecord(history.sizes.reduce((total, bytes) => total + bytes, 0));
|
|
373
|
+
await input.beforeRecord(history.sizes.reduce((total, bytes) => total + bytes, 0));
|
|
295
374
|
input.collector.recordChildSession(input.rootSessionId, {
|
|
296
375
|
childSessionId: child.id,
|
|
297
376
|
parentSessionId: input.parentSessionId,
|
|
@@ -309,7 +388,13 @@ export async function captureChildSessions(input: {
|
|
|
309
388
|
|
|
310
389
|
export default function (pi: any) {
|
|
311
390
|
// A sub-agent: the parent session captures it (session.child events).
|
|
312
|
-
|
|
391
|
+
const parentSession = (process.env.OMNIRUSH_PARENT_SESSION || "").trim();
|
|
392
|
+
if (parentSession) {
|
|
393
|
+
// Its session says so (subagent-marker.ts; index.ts loads it for
|
|
394
|
+
// sub-agent processes, which do not load the collector).
|
|
395
|
+
markSubagentSession(pi, parentSession);
|
|
396
|
+
return;
|
|
397
|
+
}
|
|
313
398
|
const origin = resolveOrigin(process.env);
|
|
314
399
|
const gatewayUrl = process.env.OMNIRUSH_GATEWAY_URL || gatewayUrlForOrigin(origin);
|
|
315
400
|
// The shared auth file is the source of truth (another process may have
|
|
@@ -429,21 +514,34 @@ export default function (pi: any) {
|
|
|
429
514
|
pendingPrompts: Array<{ at: number; model: { providerID: string; modelID: string } | null; variant: string | null }>;
|
|
430
515
|
/** User entry id -> what its prompt was sent with (pi-engine promptSettings). */
|
|
431
516
|
promptSettings: Map<string, { model: { providerID: string; modelID: string } | null; variant: string | null }>;
|
|
517
|
+
/** Prompts an extension sent (pi's input event, source "extension"), until their user entry is found. */
|
|
518
|
+
extensionInputs: Array<{ at: number; text: string }>;
|
|
519
|
+
/** User entries the user did not submit -> their parts' metadata (pi-engine injectedPrompts). */
|
|
520
|
+
injectedPrompts: Map<string, Record<string, unknown>>;
|
|
521
|
+
/** Whether the session is a sub-agent's (its first prompt the parent agent's task), once looked up. */
|
|
522
|
+
origin: Promise<void> | null;
|
|
432
523
|
} | null = null;
|
|
433
524
|
let settling: Promise<void> = Promise.resolve();
|
|
434
525
|
|
|
526
|
+
// Where sub-agent session files are looked up: the session's own
|
|
527
|
+
// subagents/ folder first (where spawn_agents puts them, and where the
|
|
528
|
+
// launcher moves older releases' children), then the session dirs
|
|
529
|
+
// themselves (children of releases up to 0.9.0 not moved yet), then
|
|
530
|
+
// every project's.
|
|
435
531
|
const sessionDirs = (ctx: any): string[] => {
|
|
436
532
|
const dirs: string[] = [];
|
|
533
|
+
const add = (dir: unknown) => {
|
|
534
|
+
if (typeof dir !== "string" || !dir) return;
|
|
535
|
+
for (const candidate of [path.join(dir, SUBAGENT_SESSIONS_DIR), dir]) {
|
|
536
|
+
if (!dirs.includes(candidate)) dirs.push(candidate);
|
|
537
|
+
}
|
|
538
|
+
};
|
|
437
539
|
const file = ctx?.sessionManager?.getSessionFile?.();
|
|
438
|
-
if (typeof file === "string" && file)
|
|
439
|
-
|
|
440
|
-
if (typeof sessionDir === "string" && sessionDir && !dirs.includes(sessionDir)) dirs.push(sessionDir);
|
|
540
|
+
if (typeof file === "string" && file) add(path.dirname(file));
|
|
541
|
+
add(ctx?.sessionManager?.getSessionDir?.());
|
|
441
542
|
try {
|
|
442
543
|
const root = path.join(agentDir, "sessions");
|
|
443
|
-
for (const name of readdirSync(root))
|
|
444
|
-
const dir = path.join(root, name);
|
|
445
|
-
if (!dirs.includes(dir)) dirs.push(dir);
|
|
446
|
-
}
|
|
544
|
+
for (const name of readdirSync(root)) add(path.join(root, name));
|
|
447
545
|
} catch {
|
|
448
546
|
/* no sessions yet */
|
|
449
547
|
}
|
|
@@ -463,7 +561,23 @@ export default function (pi: any) {
|
|
|
463
561
|
// never wrote a user entry (handled by an extension) is let go.
|
|
464
562
|
else if (pending.at > Date.now() - 60_000) state.pendingPrompts.push(pending);
|
|
465
563
|
}
|
|
466
|
-
|
|
564
|
+
// An extension's prompt is the first user entry with its text written at or after it went out.
|
|
565
|
+
for (const input of state.extensionInputs.splice(0)) {
|
|
566
|
+
const entry = entries.find((candidate: any) =>
|
|
567
|
+
candidate?.type === "message" && candidate.message?.role === "user" && typeof candidate.id === "string"
|
|
568
|
+
&& !state.injectedPrompts.has(candidate.id) && userText(candidate.message) === input.text
|
|
569
|
+
&& Number(candidate.message.timestamp ?? Date.parse(candidate.timestamp)) >= input.at - 1_000);
|
|
570
|
+
if (entry) state.injectedPrompts.set(entry.id, { source: "extension" });
|
|
571
|
+
else if (input.at > Date.now() - 60_000) state.extensionInputs.push(input);
|
|
572
|
+
}
|
|
573
|
+
return engineMessagesFromEntries(entries, {
|
|
574
|
+
sessionId: state.id,
|
|
575
|
+
agent: ROOT_AGENT,
|
|
576
|
+
cwd: state.root,
|
|
577
|
+
toolStarts: state.toolStarts,
|
|
578
|
+
promptSettings: state.promptSettings,
|
|
579
|
+
injectedPrompts: state.injectedPrompts,
|
|
580
|
+
});
|
|
467
581
|
};
|
|
468
582
|
|
|
469
583
|
/**
|
|
@@ -476,7 +590,21 @@ export default function (pi: any) {
|
|
|
476
590
|
if (!id) return null;
|
|
477
591
|
if (session?.id !== id) {
|
|
478
592
|
const root = ctx?.cwd || process.cwd();
|
|
479
|
-
session = {
|
|
593
|
+
session = {
|
|
594
|
+
id,
|
|
595
|
+
root,
|
|
596
|
+
ctx,
|
|
597
|
+
started: false,
|
|
598
|
+
running: false,
|
|
599
|
+
lastMessageId: undefined,
|
|
600
|
+
toolStarts: new Map(),
|
|
601
|
+
seenFiles: new Map(),
|
|
602
|
+
pendingPrompts: [],
|
|
603
|
+
promptSettings: new Map(),
|
|
604
|
+
extensionInputs: [],
|
|
605
|
+
injectedPrompts: new Map(),
|
|
606
|
+
origin: null,
|
|
607
|
+
};
|
|
480
608
|
}
|
|
481
609
|
session.ctx = ctx;
|
|
482
610
|
// Idempotent: a session already started in this process keeps its segment.
|
|
@@ -492,7 +620,7 @@ export default function (pi: any) {
|
|
|
492
620
|
* sub-agents keep running past the turn that started them: each later
|
|
493
621
|
* settle (and the session's end) records what they did since.
|
|
494
622
|
*/
|
|
495
|
-
const captureChildren = async (state: NonNullable<typeof session>, room: (bytes: number) => void): Promise<void> => {
|
|
623
|
+
const captureChildren = async (state: NonNullable<typeof session>, room: (bytes: number) => void | Promise<void>): Promise<void> => {
|
|
496
624
|
const { id } = state;
|
|
497
625
|
try {
|
|
498
626
|
await captureChildSessions({
|
|
@@ -513,6 +641,26 @@ export default function (pi: any) {
|
|
|
513
641
|
}
|
|
514
642
|
};
|
|
515
643
|
|
|
644
|
+
/**
|
|
645
|
+
* Room for the next sub-agent event: once the recorded messages would pass
|
|
646
|
+
* one flush (MAX_TURN_MESSAGE_BYTES), they flush, and the next child is
|
|
647
|
+
* read only after that upload went out. A turn with many sub-agents (20
|
|
648
|
+
* children, MiBs each) used to read and queue all of them at once, holding
|
|
649
|
+
* every child's messages, their trace documents and envelopes in memory
|
|
650
|
+
* together.
|
|
651
|
+
*/
|
|
652
|
+
const childRoom = (id: string) => {
|
|
653
|
+
let pending = 0;
|
|
654
|
+
return async (bytes: number): Promise<void> => {
|
|
655
|
+
if (pending > 0 && pending + bytes > MAX_TURN_MESSAGE_BYTES) {
|
|
656
|
+
collector.flushTrace(id);
|
|
657
|
+
pending = 0;
|
|
658
|
+
await collector.drained(id);
|
|
659
|
+
}
|
|
660
|
+
pending += bytes;
|
|
661
|
+
};
|
|
662
|
+
};
|
|
663
|
+
|
|
516
664
|
/**
|
|
517
665
|
* Everything a settled run records, in the desktop observer's order:
|
|
518
666
|
* messages, sub-agents, the idle event, the snapshot, the trace and the
|
|
@@ -540,22 +688,17 @@ export default function (pi: any) {
|
|
|
540
688
|
for (const [index, part] of parts.slice(0, -1).entries()) {
|
|
541
689
|
collector.recordTrace(id, "turn.messages", { messages: part.messages, part: index + 1, parts: parts.length });
|
|
542
690
|
collector.flushTrace(id);
|
|
691
|
+
// One part waits in memory at a time (a resumed session's backlog is up to 32 MiB).
|
|
692
|
+
await collector.drained(id);
|
|
543
693
|
}
|
|
544
694
|
} catch (error) {
|
|
545
695
|
collector.recordTrace(id, "session.messages_failed", { error: error instanceof Error ? error.message : "unknown" });
|
|
546
696
|
}
|
|
547
697
|
const newest = parts.at(-1)!;
|
|
548
698
|
// Sub-agent messages flush on their own before they would crowd the turn's out.
|
|
549
|
-
|
|
550
|
-
const room = (bytes: number) => {
|
|
551
|
-
if (pending > 0 && pending + bytes > MAX_TURN_MESSAGE_BYTES) {
|
|
552
|
-
collector.flushTrace(id);
|
|
553
|
-
pending = 0;
|
|
554
|
-
}
|
|
555
|
-
pending += bytes;
|
|
556
|
-
};
|
|
699
|
+
const room = childRoom(id);
|
|
557
700
|
await captureChildren(state, room);
|
|
558
|
-
room(newest.bytes);
|
|
701
|
+
await room(newest.bytes);
|
|
559
702
|
collector.recordTrace(id, "session.idle", { status });
|
|
560
703
|
// The turn snapshot runs first so the artifacts it discovers are part of
|
|
561
704
|
// the trace flushed right behind it.
|
|
@@ -598,14 +741,7 @@ export default function (pi: any) {
|
|
|
598
741
|
if (typeof lastId !== "string" || lastId === state.lastMessageId) {
|
|
599
742
|
// No new message of its own: what its sub-agents did since the last settle.
|
|
600
743
|
settling = settling.then(async () => {
|
|
601
|
-
|
|
602
|
-
await captureChildren(state, (bytes) => {
|
|
603
|
-
if (pending > 0 && pending + bytes > MAX_TURN_MESSAGE_BYTES) {
|
|
604
|
-
collector.flushTrace(state.id);
|
|
605
|
-
pending = 0;
|
|
606
|
-
}
|
|
607
|
-
pending += bytes;
|
|
608
|
-
});
|
|
744
|
+
await captureChildren(state, childRoom(state.id));
|
|
609
745
|
collector.flushTrace(state.id);
|
|
610
746
|
}).catch(() => undefined);
|
|
611
747
|
await settling;
|
|
@@ -649,11 +785,45 @@ export default function (pi: any) {
|
|
|
649
785
|
clearTimeout(budget);
|
|
650
786
|
};
|
|
651
787
|
|
|
788
|
+
/**
|
|
789
|
+
* A session opened as a main session that is a sub-agent's own (its file
|
|
790
|
+
* sits with the main sessions, so `--continue` or the resume picker can
|
|
791
|
+
* reopen it): its first prompt is the task the parent agent wrote, not
|
|
792
|
+
* something the user typed. Looked up once, before the first prompt of
|
|
793
|
+
* this process is written (every user entry then on the branch came
|
|
794
|
+
* before this process).
|
|
795
|
+
*/
|
|
796
|
+
const lookUpOrigin = (state: NonNullable<typeof session>): Promise<void> =>
|
|
797
|
+
(state.origin ??= (async () => {
|
|
798
|
+
const entries: PiEntry[] = state.ctx?.sessionManager?.getBranch?.() ?? [];
|
|
799
|
+
const task = subagentTask(entries);
|
|
800
|
+
if (!task) return;
|
|
801
|
+
let parent = task.parentSessionId;
|
|
802
|
+
if (!task.marked) {
|
|
803
|
+
// Sub-agents that wrote no marker: the parent session names it.
|
|
804
|
+
const file = state.ctx?.sessionManager?.getSessionFile?.();
|
|
805
|
+
const started = Date.parse(String(state.ctx?.sessionManager?.getHeader?.()?.timestamp ?? entries[0]?.timestamp ?? ""));
|
|
806
|
+
parent = typeof file === "string" && file
|
|
807
|
+
? await findSpawningSession(state.id, path.dirname(file), Number.isFinite(started) ? started - 10 * 60_000 : 0).catch(() => null)
|
|
808
|
+
: null;
|
|
809
|
+
if (!parent) return;
|
|
810
|
+
}
|
|
811
|
+
state.injectedPrompts.set(task.entryId, { subagentTask: true, ...(parent ? { parentSessionID: parent } : {}) });
|
|
812
|
+
})().catch(() => undefined));
|
|
813
|
+
|
|
652
814
|
pi.on("session_start", (_event: any, ctx: any) => {
|
|
653
815
|
follow(ctx, false);
|
|
654
816
|
});
|
|
655
817
|
|
|
656
|
-
|
|
818
|
+
// What the user submits comes as input; an extension's prompt
|
|
819
|
+
// (pi.sendUserMessage) is marked so it never reads as the user's.
|
|
820
|
+
pi.on("input", (event: any, ctx: any) => {
|
|
821
|
+
if (event?.source !== "extension" || typeof event?.text !== "string") return;
|
|
822
|
+
const state = follow(ctx, false);
|
|
823
|
+
state?.extensionInputs.push({ at: Date.now(), text: event.text });
|
|
824
|
+
});
|
|
825
|
+
|
|
826
|
+
pi.on("before_agent_start", async (event: any, ctx: any) => {
|
|
657
827
|
const state = follow(ctx, true);
|
|
658
828
|
if (!state) return;
|
|
659
829
|
state.running = true;
|
|
@@ -710,6 +880,7 @@ export default function (pi: any) {
|
|
|
710
880
|
},
|
|
711
881
|
});
|
|
712
882
|
archive.turnFollowed(state.id);
|
|
883
|
+
await lookUpOrigin(state);
|
|
713
884
|
});
|
|
714
885
|
|
|
715
886
|
// A run that no prompt started (a background sub-agent's results
|
|
@@ -5,29 +5,43 @@
|
|
|
5
5
|
// extension (without it, none of these files were ever registered).
|
|
6
6
|
// EVERY module that registers commands/tools/handlers must be wired here:
|
|
7
7
|
// a module left out of this list is silently dead in sessions.
|
|
8
|
+
//
|
|
9
|
+
// A sub-agent (a headless spawn_agents child: OMNIRUSH_PARENT_SESSION is
|
|
10
|
+
// set) does not load what only the interactive top agent uses — the slash
|
|
11
|
+
// commands (commands), the session capture (collector: the parent records
|
|
12
|
+
// its children) and voice input (TUI only) — so each of the many child
|
|
13
|
+
// processes a fan-out starts stays smaller. They are imported on demand,
|
|
14
|
+
// in the same order as before, so the top agent registers exactly what it
|
|
15
|
+
// did.
|
|
8
16
|
|
|
9
17
|
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
|
|
10
|
-
import collector from "./collector";
|
|
11
|
-
import commands from "./commands";
|
|
12
|
-
import mcp from "./mcp";
|
|
13
18
|
import recorder from "./recorder";
|
|
14
19
|
import sota from "./sota";
|
|
20
|
+
import mcp from "./mcp";
|
|
15
21
|
import agents from "./agents";
|
|
16
22
|
import bgshell from "./bgshell";
|
|
17
23
|
import subagents from "./subagents";
|
|
18
24
|
import swarm from "./swarm";
|
|
19
25
|
import plan from "./plan";
|
|
20
|
-
import voice from "./voice";
|
|
21
26
|
import webfetch from "./webfetch";
|
|
22
27
|
import websearch from "./websearch";
|
|
23
28
|
import guard from "./guard";
|
|
24
29
|
import yolo from "./yolo";
|
|
25
30
|
|
|
26
|
-
|
|
27
|
-
|
|
31
|
+
/** A spawn_agents child process (the parent session captures it; no TUI). */
|
|
32
|
+
export function isSubagentProcess(env: NodeJS.ProcessEnv = process.env): boolean {
|
|
33
|
+
return Boolean(String(env.OMNIRUSH_PARENT_SESSION ?? "").trim());
|
|
34
|
+
}
|
|
35
|
+
|
|
36
|
+
export default async function (pi: ExtensionAPI) {
|
|
37
|
+
const subagent = isSubagentProcess();
|
|
38
|
+
if (!subagent) (await import("./commands")).default(pi as any);
|
|
28
39
|
recorder(pi as any);
|
|
29
40
|
sota(pi as any);
|
|
30
|
-
collector(pi as any);
|
|
41
|
+
if (!subagent) (await import("./collector")).default(pi as any);
|
|
42
|
+
// Children still mark their own session as a sub-agent's (the collector
|
|
43
|
+
// does that for the top agent's sessions; children do not load it).
|
|
44
|
+
else (await import("./subagent-marker")).default(pi as any);
|
|
31
45
|
mcp(pi as any);
|
|
32
46
|
webfetch(pi as any);
|
|
33
47
|
websearch(pi as any);
|
|
@@ -35,7 +49,7 @@ export default function (pi: ExtensionAPI) {
|
|
|
35
49
|
agents(pi as any);
|
|
36
50
|
subagents(pi as any);
|
|
37
51
|
bgshell(pi as any);
|
|
38
|
-
voice(pi as any);
|
|
52
|
+
if (!subagent) (await import("./voice")).default(pi as any);
|
|
39
53
|
yolo(pi as any);
|
|
40
54
|
guard(pi as any);
|
|
41
55
|
swarm(pi as any);
|
|
@@ -18,10 +18,21 @@
|
|
|
18
18
|
// The MCP SDK is imported from the committed sibling bundle
|
|
19
19
|
// (mcp-client.bundle.js): pi copies this directory into
|
|
20
20
|
// ~/.pi/agent/extensions/omnirush/ and jiti there resolves only sibling
|
|
21
|
-
// files, never npm packages.
|
|
21
|
+
// files, never npm packages. It is loaded on first connect, so a process
|
|
22
|
+
// without MCP servers never pays for it.
|
|
23
|
+
//
|
|
24
|
+
// Sub-agents (spawn_agents children) do not start every server up front:
|
|
25
|
+
// each child used to spawn its own copy of every configured stdio server
|
|
26
|
+
// at session start (20 children x 3 servers = 60 extra processes). A child
|
|
27
|
+
// registers the tools from the list its top agent saved
|
|
28
|
+
// (<omnirush dir>/mcp-tools.json, per server config) and connects a server
|
|
29
|
+
// only when it first calls one of its tools; without a saved list it
|
|
30
|
+
// connects at start as before.
|
|
22
31
|
|
|
32
|
+
import { createHash, randomUUID } from "node:crypto";
|
|
33
|
+
import { mkdirSync, readFileSync, renameSync, writeFileSync } from "node:fs";
|
|
23
34
|
import path from "node:path";
|
|
24
|
-
import { Client
|
|
35
|
+
import type { Client } from "./mcp-client.bundle.js";
|
|
25
36
|
import { omniDir } from "./auth";
|
|
26
37
|
import {
|
|
27
38
|
dedupeToolNames,
|
|
@@ -53,7 +64,62 @@ const SHUTDOWN_BUDGET_MS = 3_000;
|
|
|
53
64
|
// Stderr lines kept per stdio server for failure diagnostics.
|
|
54
65
|
const STDERR_TAIL_LINES = 5;
|
|
55
66
|
|
|
56
|
-
|
|
67
|
+
/** "deferred": tools registered from the saved list, connects on first use (sub-agents). */
|
|
68
|
+
type ServerStatus = "connecting" | "connected" | "failed" | "deferred";
|
|
69
|
+
|
|
70
|
+
type McpSdk = typeof import("./mcp-client.bundle.js");
|
|
71
|
+
let sdkPromise: Promise<McpSdk> | null = null;
|
|
72
|
+
let loadedSdk: McpSdk | null = null;
|
|
73
|
+
/** The MCP SDK bundle, imported on first use. */
|
|
74
|
+
function loadSdk(): Promise<McpSdk> {
|
|
75
|
+
sdkPromise ??= import("./mcp-client.bundle.js").then((sdk) => {
|
|
76
|
+
loadedSdk = sdk;
|
|
77
|
+
return sdk;
|
|
78
|
+
});
|
|
79
|
+
return sdkPromise;
|
|
80
|
+
}
|
|
81
|
+
|
|
82
|
+
/** Saved tool lists (by the process that connected), keyed by server name. */
|
|
83
|
+
export const MCP_TOOLS_CACHE_FILE = "mcp-tools.json";
|
|
84
|
+
/** A saved list older than this is not trusted (the server may have changed). */
|
|
85
|
+
const TOOLS_CACHE_MAX_AGE_MS = 7 * 24 * 3600_000;
|
|
86
|
+
|
|
87
|
+
/** Identity of a server config: a saved tool list only counts for the same config. */
|
|
88
|
+
export function serverConfigKey(config: McpServerConfig): string {
|
|
89
|
+
return createHash("sha256").update(JSON.stringify(config)).digest("hex").slice(0, 32);
|
|
90
|
+
}
|
|
91
|
+
|
|
92
|
+
function readToolsCache(dir: string): Record<string, { key: string; at: number; tools: McpToolInfo[] }> {
|
|
93
|
+
try {
|
|
94
|
+
const parsed = JSON.parse(readFileSync(path.join(dir, MCP_TOOLS_CACHE_FILE), "utf8"));
|
|
95
|
+
return parsed && typeof parsed === "object" && !Array.isArray(parsed) ? parsed : {};
|
|
96
|
+
} catch {
|
|
97
|
+
return {};
|
|
98
|
+
}
|
|
99
|
+
}
|
|
100
|
+
|
|
101
|
+
/** Save one server's tool list (atomic; other servers' entries kept). */
|
|
102
|
+
function saveToolsCache(dir: string, state: ServerState): void {
|
|
103
|
+
try {
|
|
104
|
+
const cache = readToolsCache(dir);
|
|
105
|
+
cache[state.name] = { key: serverConfigKey(state.config), at: Date.now(), tools: state.tools };
|
|
106
|
+
mkdirSync(dir, { recursive: true, mode: 0o700 });
|
|
107
|
+
const file = path.join(dir, MCP_TOOLS_CACHE_FILE);
|
|
108
|
+
const tmp = `${file}.${randomUUID()}.tmp`;
|
|
109
|
+
writeFileSync(tmp, JSON.stringify(cache), { mode: 0o600 });
|
|
110
|
+
renameSync(tmp, file);
|
|
111
|
+
} catch {
|
|
112
|
+
/* only an optimization for sub-agents */
|
|
113
|
+
}
|
|
114
|
+
}
|
|
115
|
+
|
|
116
|
+
/** A saved tool list for this server config, or null. */
|
|
117
|
+
function cachedTools(dir: string, state: ServerState, now = Date.now()): McpToolInfo[] | null {
|
|
118
|
+
const entry = readToolsCache(dir)[state.name];
|
|
119
|
+
if (!entry || entry.key !== serverConfigKey(state.config) || !Array.isArray(entry.tools)) return null;
|
|
120
|
+
if (!(now - Number(entry.at) < TOOLS_CACHE_MAX_AGE_MS)) return null;
|
|
121
|
+
return entry.tools.filter((tool) => tool && typeof tool.name === "string");
|
|
122
|
+
}
|
|
57
123
|
|
|
58
124
|
/** OMNIRUSH_DEBUG=1 diagnostics (same convention as the sota guard). */
|
|
59
125
|
function debug(line: string): void {
|
|
@@ -79,6 +145,8 @@ interface ServerState {
|
|
|
79
145
|
tools: McpToolInfo[];
|
|
80
146
|
client?: Client;
|
|
81
147
|
stderrTail: string[];
|
|
148
|
+
/** A deferred server's connect in flight (single-flight for parallel tool calls). */
|
|
149
|
+
connecting?: Promise<void>;
|
|
82
150
|
}
|
|
83
151
|
|
|
84
152
|
function withTimeout<T>(promise: Promise<T>, ms: number, label: string): Promise<T> {
|
|
@@ -138,6 +206,7 @@ function contentText(content: Array<Record<string, any>>): string {
|
|
|
138
206
|
|
|
139
207
|
/** Connect one server (bounded). Throws on failure — caller isolates. */
|
|
140
208
|
async function connectServer(state: ServerState): Promise<void> {
|
|
209
|
+
const { Client, StdioClientTransport, StreamableHTTPClientTransport } = await loadSdk();
|
|
141
210
|
const client = new Client({
|
|
142
211
|
name: "omnirush",
|
|
143
212
|
version: (process.env.OMNIRUSH_VERSION || "").trim() || "0.0.0",
|
|
@@ -194,6 +263,11 @@ async function connectServer(state: ServerState): Promise<void> {
|
|
|
194
263
|
}
|
|
195
264
|
}
|
|
196
265
|
|
|
266
|
+
/** A spawn_agents child process. */
|
|
267
|
+
function isSubagent(): boolean {
|
|
268
|
+
return Boolean(String(process.env.OMNIRUSH_PARENT_SESSION ?? "").trim());
|
|
269
|
+
}
|
|
270
|
+
|
|
197
271
|
function describeServer(state: ServerState): string {
|
|
198
272
|
if (state.config.url) return `${state.name} (${state.config.url})`;
|
|
199
273
|
const args = (state.config.args ?? []).join(" ");
|
|
@@ -231,6 +305,11 @@ export default function (pi: any) {
|
|
|
231
305
|
void toolCallId;
|
|
232
306
|
void onUpdate;
|
|
233
307
|
const current = servers.get(state.name);
|
|
308
|
+
if (current?.status === "deferred") {
|
|
309
|
+
// A sub-agent's first call to this server: connect now.
|
|
310
|
+
current.connecting ??= connectServer(current).finally(() => { current.connecting = undefined; });
|
|
311
|
+
await current.connecting.catch(() => undefined);
|
|
312
|
+
}
|
|
234
313
|
const client = current?.client;
|
|
235
314
|
if (current?.status !== "connected" || !client) {
|
|
236
315
|
throw new Error(
|
|
@@ -248,7 +327,7 @@ export default function (pi: any) {
|
|
|
248
327
|
{ timeout: CALL_TIMEOUT_MS, signal },
|
|
249
328
|
);
|
|
250
329
|
} catch (error: any) {
|
|
251
|
-
const detail = error instanceof McpError ? error.message : error?.message ?? String(error);
|
|
330
|
+
const detail = loadedSdk && error instanceof loadedSdk.McpError ? error.message : error?.message ?? String(error);
|
|
252
331
|
throw new Error(`MCP call failed on "${state.name}": ${detail}`);
|
|
253
332
|
}
|
|
254
333
|
const content = mcpContentToPiContent(result?.content);
|
|
@@ -284,6 +363,7 @@ export default function (pi: any) {
|
|
|
284
363
|
try {
|
|
285
364
|
await connectServer(state);
|
|
286
365
|
registerServerTools(state);
|
|
366
|
+
saveToolsCache(omniDir(), state);
|
|
287
367
|
} catch {
|
|
288
368
|
/* connectServer recorded the failure; isolation keeps going */
|
|
289
369
|
}
|
|
@@ -316,7 +396,9 @@ export default function (pi: any) {
|
|
|
316
396
|
const stateLabel =
|
|
317
397
|
state.status === "connected"
|
|
318
398
|
? `connected, ${state.tools.length} tool${state.tools.length === 1 ? "" : "s"}`
|
|
319
|
-
: state.status === "
|
|
399
|
+
: state.status === "deferred"
|
|
400
|
+
? `connects on first use, ${state.tools.length} tool${state.tools.length === 1 ? "" : "s"}`
|
|
401
|
+
: state.status === "connecting"
|
|
320
402
|
? "connecting…"
|
|
321
403
|
: `failed: ${state.error ?? "unknown error"}`;
|
|
322
404
|
lines.push(` ${describeServer(state)} — ${stateLabel}`);
|
|
@@ -367,10 +449,23 @@ export default function (pi: any) {
|
|
|
367
449
|
for (const error of errors) {
|
|
368
450
|
ctx?.ui?.notify?.(`mcp config: ${error}`, "warning");
|
|
369
451
|
}
|
|
370
|
-
if (servers.size
|
|
371
|
-
|
|
372
|
-
|
|
452
|
+
if (servers.size === 0) return;
|
|
453
|
+
if (isSubagent()) {
|
|
454
|
+
// Tools from the saved list; the server starts on first use.
|
|
455
|
+
const dir = omniDir();
|
|
456
|
+
for (const state of servers.values()) {
|
|
457
|
+
const tools = cachedTools(dir, state);
|
|
458
|
+
if (!tools) continue;
|
|
459
|
+
state.tools = tools;
|
|
460
|
+
state.status = "deferred";
|
|
461
|
+
registerServerTools(state);
|
|
462
|
+
}
|
|
463
|
+
const rest = [...servers.values()].filter((state) => state.status !== "deferred");
|
|
464
|
+
await Promise.allSettled(rest.map((state) => connectAll(state.name)));
|
|
465
|
+
return;
|
|
373
466
|
}
|
|
467
|
+
// Isolated + bounded per server; failures surface via /mcp only.
|
|
468
|
+
await connectAll();
|
|
374
469
|
});
|
|
375
470
|
|
|
376
471
|
pi.on("session_shutdown", async () => {
|