fix(coding-agent): scope SDK extension packages

This commit is contained in:
Sunil Srivatsa
2026-08-01 13:50:53 -04:00
parent df38e5b48e
commit f22c0edbfa
4 changed files with 157 additions and 8 deletions
@@ -15,6 +15,7 @@
* @see ./omp-plugins.ts
* @see ./builtin.ts `loadExtensionModules`
*/
import { AsyncLocalStorage } from "node:async_hooks";
import * as fs from "node:fs/promises";
import * as path from "node:path";
import { getAgentDir, isEnoent, logger, tryParseJson } from "@oh-my-pi/pi-utils";
@@ -41,8 +42,18 @@ interface InjectedRoot {
level: "user" | "project";
}
export type OmpExtensionRootMode = "merge" | "explicit-only";
interface InvocationRootScope {
/** Raw SDK spellings, resolved against the LoadContext that performs discovery. */
paths: readonly string[];
mode: OmpExtensionRootMode;
}
const invocationRootScope = new AsyncLocalStorage<InvocationRootScope>();
let injectedCliRoots: InjectedRoot[] = [];
let injectedCliRootMode: "merge" | "explicit-only" = "merge";
let injectedCliRootMode: OmpExtensionRootMode = "merge";
export interface InjectOmpExtensionCliRootOptions {
/**
@@ -50,11 +61,25 @@ export interface InjectOmpExtensionCliRootOptions {
* with `--no-extensions` so configured and installed packages cannot
* contribute sibling capabilities through the `omp-plugins` provider.
*/
mode?: "merge" | "explicit-only";
mode?: OmpExtensionRootMode;
/** Replace roots from an earlier invocation instead of extending them. */
replace?: boolean;
}
/**
* Run one SDK invocation with its own extension-package roots. Async resources
* started inside `callback` retain this scope, including discovery deliberately
* deferred until the end of session startup. Raw relative paths are resolved by
* {@link listOmpExtensionRoots} against that invocation's active cwd.
*/
export function withOmpExtensionRootScope<T>(
paths: readonly string[],
mode: OmpExtensionRootMode,
callback: () => T,
): T {
return invocationRootScope.run({ paths: [...paths], mode }, callback);
}
/**
* Register CLI-provided extension package paths (e.g. from `--extension`/`-e`)
* so the sub-discovery providers can find their sibling `skills/`, `hooks/`,
@@ -146,7 +171,8 @@ async function isDirectory(p: string): Promise<boolean> {
* Sources, in order of precedence (later entries with the same absolute path
* are dropped):
*
* 1. CLI roots injected via {@link injectOmpExtensionCliRoots}
* 1. Invocation-scoped SDK roots, when present; otherwise CLI roots injected
* via {@link injectOmpExtensionCliRoots}
* 2. Project `<cwd>/.omp/settings.json#extensions`
* 3. User `~/.omp/agent/settings.json#extensions`
* 4. Enabled npm/link plugins installed under `<plugins>/node_modules/` (for
@@ -159,10 +185,14 @@ async function isDirectory(p: string): Promise<boolean> {
* other sources still surface.
*/
export async function listOmpExtensionRoots(ctx: LoadContext): Promise<OmpExtensionRoot[]> {
let candidates: InjectedRoot[] = injectedCliRoots.map(root =>
root.relativePath ? { ...root, path: path.resolve(ctx.cwd, root.relativePath) } : root,
);
if (injectedCliRootMode === "merge") {
const scopedRoots = invocationRootScope.getStore();
const rootMode = scopedRoots?.mode ?? injectedCliRootMode;
let candidates: InjectedRoot[] = scopedRoots
? scopedRoots.paths.map(raw => ({ path: resolveAgainst(raw, ctx), level: "user" }))
: injectedCliRoots.map(root =>
root.relativePath ? { ...root, path: path.resolve(ctx.cwd, root.relativePath) } : root,
);
if (rootMode === "merge") {
const { project, user } = scopeDirs(ctx);
const [projectExtensions, userExtensions, installedPlugins] = await Promise.all([
readSettingsExtensions(path.join(project, "settings.json")),
@@ -177,7 +207,7 @@ export async function listOmpExtensionRoots(ctx: LoadContext): Promise<OmpExtens
];
}
// First-seen-wins dedup preserves CLI > project-settings > user-settings > installed precedence.
// First-seen-wins dedup preserves invocation/CLI > project-settings > user-settings > installed precedence.
const seen = new Set<string>();
const unique: InjectedRoot[] = [];
for (const candidate of candidates) {
+8
View File
@@ -66,6 +66,7 @@ import { CursorExecHandlers, type CursorMcpResourceAdapter } from "./cursor";
import { createBridgeEditTool, createBridgeGrepFactory } from "./cursor-bridge-tools";
import "./discovery";
import { initializeWithSettings } from "./discovery";
import { withOmpExtensionRootScope } from "./discovery/omp-extension-roots";
import { disposeAllJuliaKernelSessions, disposeJuliaKernelSessionsByOwner } from "./eval/jl/executor";
import { disposeVmContextsByOwner } from "./eval/js/context-manager";
import { disposeAllKernelSessions, disposeKernelSessionsByOwner } from "./eval/py/executor";
@@ -1217,6 +1218,13 @@ export function createAutoLearnCaptureRunner(
* ```
*/
export async function createAgentSession(options: CreateAgentSessionOptions = {}): Promise<CreateAgentSessionResult> {
const rootMode = options.disableExtensionDiscovery ? "explicit-only" : "merge";
return await withOmpExtensionRootScope(options.additionalExtensionPaths ?? [], rootMode, () =>
createAgentSessionScoped(options),
);
}
async function createAgentSessionScoped(options: CreateAgentSessionOptions): Promise<CreateAgentSessionResult> {
const cwd = options.cwd ?? getProjectDir();
const agentDir = options.agentDir ?? getAgentDir();
const eventBus = options.eventBus ?? new EventBus();
@@ -32,6 +32,7 @@ import {
clearOmpExtensionCliRoots,
injectOmpExtensionCliRoots,
listOmpExtensionRoots,
withOmpExtensionRootScope,
} from "@oh-my-pi/pi-coding-agent/discovery/omp-extension-roots";
import { discoverExtensionPaths } from "@oh-my-pi/pi-coding-agent/extensibility/extensions/loader";
import { getConfigRootDir, removeSyncWithRetries, setAgentDir } from "@oh-my-pi/pi-utils";
@@ -213,6 +214,45 @@ test("explicit-only CLI roots replace stale state and exclude every ambient pack
).toBe(false);
});
test("invocation scopes isolate concurrent SDK roots and merge ambient roots only when requested", async () => {
const otherExplicit = path.join(tempDir, "other-explicit-extension");
const projectExt = path.join(tempDir, "project-extension");
const installed = path.join(home, ".omp", "plugins", "node_modules", "installed-extension");
const staleCli = path.join(tempDir, "stale-cli-extension");
buildExtensionPackage(otherExplicit, "other-explicit-skill");
buildExtensionPackage(projectExt, "project-skill");
buildExtensionPackage(installed, "installed-skill");
buildExtensionPackage(staleCli, "stale-cli-skill");
writeFile(path.join(project, ".omp", "settings.json"), JSON.stringify({ extensions: [projectExt] }));
writeFile(
path.join(home, ".omp", "plugins", "package.json"),
JSON.stringify({ name: "omp-plugins", dependencies: { "installed-extension": "1.0.0" } }),
);
injectOmpExtensionCliRoots([staleCli], home, project);
const firstEntered = Promise.withResolvers<void>();
const secondEntered = Promise.withResolvers<void>();
const [firstRoots, secondRoots] = await Promise.all([
withOmpExtensionRootScope([ext], "explicit-only", async () => {
firstEntered.resolve();
await secondEntered.promise;
return listOmpExtensionRoots(ctx());
}),
withOmpExtensionRootScope([otherExplicit], "explicit-only", async () => {
secondEntered.resolve();
await firstEntered.promise;
return listOmpExtensionRoots(ctx());
}),
]);
expect(firstRoots.map(root => root.path)).toEqual([ext]);
expect(secondRoots.map(root => root.path)).toEqual([otherExplicit]);
const mergedRoots = await withOmpExtensionRootScope([ext], "merge", () => listOmpExtensionRoots(ctx()));
expect(mergedRoots.map(root => root.path)).toEqual(expect.arrayContaining([ext, projectExt, installed]));
expect(mergedRoots.map(root => root.path)).not.toContain(staleCli);
});
test("file-extension entrypoints contribute zero sub-surface (the file has no siblings to scan)", async () => {
const standaloneFile = path.join(tempDir, "standalone.ts");
fs.writeFileSync(standaloneFile, "export default function (_pi) {}\n");
@@ -7,6 +7,7 @@ import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings";
import { getActiveSkills } from "@oh-my-pi/pi-coding-agent/extensibility/skills";
import type { Skill } from "@oh-my-pi/pi-coding-agent/sdk";
import { createAgentSession } from "@oh-my-pi/pi-coding-agent/sdk";
import type { AgentSession } from "@oh-my-pi/pi-coding-agent/session/agent-session";
import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage";
import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager";
import { removeSyncWithRetries } from "@oh-my-pi/pi-utils";
@@ -24,6 +25,19 @@ function createIsolatedSkillsSettings(): Settings {
});
}
function createExtensionSkill(packageDir: string, skillName: string): void {
fs.mkdirSync(path.join(packageDir, "skills", skillName), { recursive: true });
fs.writeFileSync(
path.join(packageDir, "package.json"),
JSON.stringify({ name: path.basename(packageDir), omp: { extensions: ["./extension.ts"] } }),
);
fs.writeFileSync(path.join(packageDir, "extension.ts"), "export default function extension() {}\n");
fs.writeFileSync(
path.join(packageDir, "skills", skillName, "SKILL.md"),
`---\nname: ${skillName}\ndescription: SDK extension package skill\n---\nbody\n`,
);
}
describe("createAgentSession skills option", () => {
let tempDir: string;
let skillsDir: string;
@@ -107,6 +121,63 @@ Loaded via symbolic link.
expect(session.skills.some((s: Skill) => s.name === "test-skill")).toBe(true);
});
it("SDK invocation root scope isolates disabled discovery and merges normal discovery", async () => {
const explicitPackage = path.join(tempDir, "sdk-explicit-extension");
const settingsPackage = path.join(tempDir, "sdk-settings-extension");
const installedPackage = path.join(tempHomeDir, ".omp", "plugins", "node_modules", "sdk-installed-extension");
createExtensionSkill(explicitPackage, "sdk-explicit-skill");
createExtensionSkill(settingsPackage, "sdk-settings-skill");
createExtensionSkill(installedPackage, "sdk-installed-skill");
fs.writeFileSync(path.join(tempDir, ".omp", "settings.json"), JSON.stringify({ extensions: [settingsPackage] }));
fs.mkdirSync(path.join(tempHomeDir, ".omp", "plugins"), { recursive: true });
fs.writeFileSync(
path.join(tempHomeDir, ".omp", "plugins", "package.json"),
JSON.stringify({ name: "omp-plugins", dependencies: { "sdk-installed-extension": "1.0.0" } }),
);
const previousAgentDir = getAgentDir();
setAgentDir(path.join(tempHomeDir, ".omp", "agent"));
const baseSessionOptions = {
cwd: tempDir,
agentDir: path.join(tempHomeDir, ".omp", "agent"),
modelRegistry: sharedModelRegistry,
additionalExtensionPaths: [explicitPackage],
enableMCP: false,
enableLsp: false,
contextFiles: [],
promptTemplates: [],
slashCommands: [],
rules: [],
};
let session: AgentSession | undefined;
try {
({ session } = await createAgentSession({
...baseSessionOptions,
sessionManager: SessionManager.inMemory(),
settings: createIsolatedSkillsSettings(),
disableExtensionDiscovery: true,
}));
const isolatedSkillNames = session.skills.map(skill => skill.name);
expect(isolatedSkillNames).toContain("sdk-explicit-skill");
expect(isolatedSkillNames).not.toEqual(expect.arrayContaining(["sdk-settings-skill", "sdk-installed-skill"]));
await session.dispose();
session = undefined;
({ session } = await createAgentSession({
...baseSessionOptions,
sessionManager: SessionManager.inMemory(),
settings: createIsolatedSkillsSettings(),
}));
const mergedSkillNames = session.skills.map(skill => skill.name);
expect(mergedSkillNames).toEqual(expect.arrayContaining(["sdk-explicit-skill", "sdk-settings-skill"]));
} finally {
await session?.dispose();
setAgentDir(previousAgentDir);
}
});
it("should discover skills when skill directory is a symlink", async () => {
const { session } = await createAgentSession({
cwd: tempDir,