@chainingintention/pi-web-cn 1.202606.3 → 1.202606.5

package/dist/server/sessions/piSessionService.js

@@ -1,5 +1,9 @@
-import { readFile, writeFile } from "node:fs/promises";
-import { AuthStorage, createAgentSessionFromServices, createAgentSessionRuntime, createAgentSessionServices, createEditToolDefinition, defineTool, getAgentDir, ModelRegistry, SessionManager, } from "@earendil-works/pi-coding-agent";
+import { spawn } from "node:child_process";
+import { constants } from "node:fs";
+import { access as fsAccess, mkdir as fsMkdir, readFile, readdir as fsReaddir, realpath as fsRealpath, stat as fsStat, writeFile } from "node:fs/promises";
+import { basename, dirname, isAbsolute, relative, resolve, sep } from "node:path";
+import { AuthStorage, createAgentSessionFromServices, createAgentSessionRuntime, createAgentSessionServices, createEditToolDefinition, createFindToolDefinition, createGrepToolDefinition, createLsToolDefinition, createReadToolDefinition, createWriteToolDefinition, defineTool, getAgentDir, ModelRegistry, SessionManager, } from "@earendil-works/pi-coding-agent";
+import { Type } from "typebox";
 import { pageMessagesAtSafeBoundary } from "./messagePaging.js";
 import { BUILTIN_COMMANDS } from "./builtinCommands.js";
 import { SessionCommandService } from "./sessionCommandService.js";
@@ -7,6 +11,9 @@ import { SessionArchiveStore } from "./sessionArchiveStore.js";
 import { findArchiveCandidateByIdOrPrefix, planSessionArchiveTree } from "./sessionArchiveTree.js";
 import { fallbackSessionName, generateShortSessionName } from "./sessionNameGenerator.js";
 import { computeEditPreview } from "./editPreview.js";
+import { managementToolAllowed } from "../managementEmbed.js";
+import { bubblewrapUnavailableReason, createBubblewrapPythonInvocation, createManagedPythonFallbackPrelude, createManagedSandboxEnvironment, DEFAULT_BUBBLEWRAP_PATHS } from "./managementSandbox.js";
+import { PI_PERMISSION_SYSTEM_POLICY_AGENT_DIR, managementAgentToolNames, withRuntimeCreationEnvironment, writeManagementPermissionSystemPolicy } from "./managementPermissionSystem.js";
 function noop() {
     // Intentionally empty default unsubscribe callback.
 }
@@ -20,17 +27,47 @@ function defaultCreateAgentRuntime(createRuntime, options) {
 }
 function createDefaultRuntimeFactory(authStorage, modelRegistry) {
     return async ({ cwd, agentDir, sessionManager, sessionStartEvent }) => {
-        const services = await createAgentSessionServices({ cwd, agentDir, authStorage, modelRegistry });
-        const customTools = [createPiWebEditToolDefinition(cwd)];
-        const options = sessionStartEvent === undefined
-            ? { services, sessionManager, customTools }
-            : { services, sessionManager, sessionStartEvent, customTools };
-        const result = await createAgentSessionFromServices(options);
-        return { ...result, services, diagnostics: services.diagnostics };
+        return withRuntimeCreationEnvironment({}, async () => {
+            const services = await createAgentSessionServices({ cwd, agentDir, authStorage, modelRegistry });
+            const customTools = [createPiWebEditToolDefinition(cwd)];
+            const options = sessionStartEvent === undefined
+                ? { services, sessionManager, customTools }
+                : { services, sessionManager, sessionStartEvent, customTools };
+            const result = await createAgentSessionFromServices(options);
+            return { ...result, services, diagnostics: services.diagnostics };
+        });
+    };
+}
+function createManagementRuntimeFactory(authStorage, modelRegistry, managementContext) {
+    return async ({ cwd, agentDir, sessionManager, sessionStartEvent }) => {
+        const policyAgentDir = await writeManagementPermissionSystemPolicy(agentDir, cwd, managementContext);
+        return withRuntimeCreationEnvironment({ [PI_PERMISSION_SYSTEM_POLICY_AGENT_DIR]: policyAgentDir }, async () => {
+            const services = await createAgentSessionServices({ cwd, agentDir, authStorage, modelRegistry });
+            const customTools = createManagementSandboxToolDefinitions(cwd, managementContext);
+            const options = sessionStartEvent === undefined
+                ? { services, sessionManager, customTools, tools: managementAgentToolNames(managementContext) }
+                : { services, sessionManager, sessionStartEvent, customTools, tools: managementAgentToolNames(managementContext) };
+            // @ts-expect-error SDK customTools accepts concrete ToolDefinition instances at runtime, but the published type is invariant in render callbacks.
+            const result = await createAgentSessionFromServices(options);
+            return { ...result, services, diagnostics: services.diagnostics };
+        });
     };
 }
-function createPiWebEditToolDefinition(cwd) {
-    const editTool = createEditToolDefinition(cwd);
+export { managementAgentToolNames };
+export function createManagementSandboxToolDefinitions(cwd, context) {
+    const operations = createManagedFileOperations(cwd);
+    return [
+        createReadToolDefinition(cwd, { operations: operations.read }),
+        createWriteToolDefinition(cwd, { operations: operations.write }),
+        createPiWebEditToolDefinition(cwd, operations.edit),
+        createLsToolDefinition(cwd, { operations: operations.ls }),
+        createGrepToolDefinition(cwd, { operations: operations.grep }),
+        createFindToolDefinition(cwd, { operations: operations.find }),
+        createManagedPythonToolDefinition(cwd, context),
+    ];
+}
+function createPiWebEditToolDefinition(cwd, operations) {
+    const editTool = createEditToolDefinition(cwd, operations === undefined ? undefined : { operations });
     return defineTool({
         name: editTool.name,
         label: editTool.label,
@@ -50,6 +87,211 @@ function createPiWebEditToolDefinition(cwd) {
         },
     });
 }
+const pythonSchema = Type.Object({
+    code: Type.String({ description: "Python code to run in the managed project workspace" }),
+    timeoutMs: Type.Optional(Type.Number({ description: "Execution timeout in milliseconds" })),
+});
+function createManagedPythonToolDefinition(cwd, context) {
+    return defineTool({
+        name: "python",
+        label: "python",
+        description: "Run Python code inside the managed project workspace. Shell commands and paths outside the project are blocked.",
+        promptSnippet: "Run Python code in the current project",
+        promptGuidelines: ["Use python for scripts and calculations. Do not use it to run shell commands."],
+        parameters: pythonSchema,
+        async execute(_toolCallId, params, signal) {
+            const configuredPython = context.sandbox?.pythonExecutable?.trim();
+            const pythonExecutable = configuredPython !== undefined && configuredPython !== "" ? configuredPython : "python3";
+            const timeoutMs = Math.max(1_000, Math.min(params.timeoutMs ?? 30_000, 120_000));
+            const configuredBubblewrap = process.env["PI_WEB_BWRAP_EXECUTABLE"]?.trim();
+            const bubblewrapExecutable = configuredBubblewrap !== undefined && configuredBubblewrap !== "" ? configuredBubblewrap : "bwrap";
+            const env = createManagedSandboxEnvironment({ hostEnv: process.env, context });
+            return runManagedPython({ pythonExecutable, bubblewrapExecutable, cwd, code: params.code, timeoutMs, env, signal });
+        },
+    });
+}
+async function runManagedPython(options) {
+    const root = await fsRealpath(options.cwd);
+    const invocation = createBubblewrapPythonInvocation({
+        bubblewrapExecutable: options.bubblewrapExecutable,
+        pythonExecutable: options.pythonExecutable,
+        workspaceRoot: root,
+        env: options.env,
+        readOnlyPaths: await readableBubblewrapPaths(),
+    });
+    const result = await runPythonProcess({ command: invocation.command, args: invocation.args, cwd: root, env: options.env, code: options.code, timeoutMs: options.timeoutMs, signal: options.signal });
+    const unavailable = bubblewrapUnavailableReason(result.output);
+    if (unavailable !== undefined) {
+        return runManagedPythonFallback({ pythonExecutable: options.pythonExecutable, cwd: root, code: options.code, timeoutMs: options.timeoutMs, env: options.env, signal: options.signal });
+    }
+    return pythonToolResult(result.code, result.output);
+}
+async function runManagedPythonFallback(options) {
+    const code = `${createManagedPythonFallbackPrelude(options.cwd)}\n${options.code}`;
+    const result = await runPythonProcess({ command: options.pythonExecutable, args: ["-I", "-"], cwd: options.cwd, env: options.env, code, timeoutMs: options.timeoutMs, signal: options.signal });
+    return pythonToolResult(result.code, result.output);
+}
+async function runPythonProcess(options) {
+    return new Promise((resolvePromise, reject) => {
+        if (options.signal?.aborted === true) {
+            reject(new Error("Operation aborted"));
+            return;
+        }
+        const child = spawn(options.command, options.args, { cwd: options.cwd, env: options.env, stdio: ["pipe", "pipe", "pipe"], shell: false });
+        let stdout = "";
+        let stderr = "";
+        const timer = setTimeout(() => {
+            child.kill();
+            reject(new Error(`Python execution timed out after ${String(options.timeoutMs)}ms`));
+        }, options.timeoutMs);
+        const onAbort = () => {
+            child.kill();
+            reject(new Error("Operation aborted"));
+        };
+        options.signal?.addEventListener("abort", onAbort, { once: true });
+        child.stdout.on("data", (chunk) => { stdout += chunk.toString("utf8"); });
+        child.stderr.on("data", (chunk) => { stderr += chunk.toString("utf8"); });
+        child.on("error", (error) => {
+            clearTimeout(timer);
+            options.signal?.removeEventListener("abort", onAbort);
+            if (isNodeErrorWithCode(error, "ENOENT"))
+                reject(new Error("Python sandbox is unavailable"));
+            else
+                reject(error);
+        });
+        child.on("close", (codeValue) => {
+            clearTimeout(timer);
+            options.signal?.removeEventListener("abort", onAbort);
+            const output = truncateToolOutput([stdout.trimEnd(), stderr.trimEnd()].filter((part) => part !== "").join("\n"));
+            resolvePromise({ code: codeValue, output });
+        });
+        child.stdin.end(options.code);
+    });
+}
+function pythonToolResult(codeValue, output) {
+    const prefix = codeValue === 0 ? "" : `Python exited with code ${String(codeValue)}\n`;
+    return { content: [{ type: "text", text: `${prefix}${output}`.trimEnd() }], details: undefined };
+}
+async function readableBubblewrapPaths() {
+    const paths = await Promise.all(DEFAULT_BUBBLEWRAP_PATHS.map(async (path) => {
+        try {
+            await fsAccess(path, constants.R_OK);
+            return path;
+        }
+        catch {
+            return undefined;
+        }
+    }));
+    return paths.filter(isDefined);
+}
+function createManagedFileOperations(cwd) {
+    const read = {
+        readFile: async (absolutePath) => readFile(await assertExistingInside(cwd, absolutePath)),
+        access: async (absolutePath) => { await fsAccess(await assertExistingInside(cwd, absolutePath), constants.R_OK); },
+    };
+    const write = {
+        writeFile: async (absolutePath, content) => { await writeFile(await assertWritableInside(cwd, absolutePath), content, "utf8"); },
+        mkdir: async (dir) => { await fsMkdir(await assertPotentialInside(cwd, dir), { recursive: true }); },
+    };
+    const edit = {
+        readFile: read.readFile,
+        writeFile: write.writeFile,
+        access: async (absolutePath) => { await fsAccess(await assertExistingInside(cwd, absolutePath), constants.R_OK | constants.W_OK); },
+    };
+    const ls = {
+        exists: async (absolutePath) => pathExistsInside(cwd, absolutePath),
+        stat: async (absolutePath) => fsStat(await assertExistingInside(cwd, absolutePath)),
+        readdir: async (absolutePath) => fsReaddir(await assertExistingInside(cwd, absolutePath)),
+    };
+    const grep = {
+        isDirectory: async (absolutePath) => (await fsStat(await assertExistingInside(cwd, absolutePath))).isDirectory(),
+        readFile: async (absolutePath) => (await readFile(await assertExistingInside(cwd, absolutePath), "utf8")),
+    };
+    const find = {
+        exists: async (absolutePath) => pathExistsInside(cwd, absolutePath),
+        glob: async (pattern, searchPath, options) => managedGlob(cwd, pattern, searchPath, options.limit),
+    };
+    return { read, write, edit, ls, grep, find };
+}
+async function pathExistsInside(rootPath, targetPath) {
+    try {
+        await assertExistingInside(rootPath, targetPath);
+        return true;
+    }
+    catch {
+        return false;
+    }
+}
+async function managedGlob(rootPath, pattern, searchPath, limit) {
+    const root = await fsRealpath(rootPath);
+    const start = await assertExistingInside(root, searchPath);
+    const regex = globPatternToRegExp(pattern);
+    const results = [];
+    async function walk(dir) {
+        if (results.length >= limit)
+            return;
+        const entries = await fsReaddir(dir, { withFileTypes: true });
+        for (const entry of entries) {
+            if (results.length >= limit)
+                return;
+            if (entry.name === ".git" || entry.name === "node_modules")
+                continue;
+            const fullPath = resolve(dir, entry.name);
+            const safePath = await assertExistingInside(root, fullPath);
+            const rel = relative(start, safePath).split(sep).join("/");
+            if (entry.isDirectory()) {
+                await walk(safePath);
+            }
+            else if (regex.test(rel)) {
+                results.push(safePath);
+            }
+        }
+    }
+    await walk(start);
+    return results;
+}
+function globPatternToRegExp(pattern) {
+    const escaped = pattern
+        .replace(/[.+^${}()|[\]\\]/g, "\\$&")
+        .replace(/\*\*/g, "\0")
+        .replace(/\*/g, "[^/]*")
+        .replace(/\?/g, "[^/]")
+        .replace(/\0/g, ".*");
+    return new RegExp(`^${escaped}$`);
+}
+async function assertExistingInside(rootPath, targetPath) {
+    const [root, target] = await Promise.all([fsRealpath(rootPath), fsRealpath(targetPath)]);
+    ensurePathInside(root, target);
+    return target;
+}
+async function assertWritableInside(rootPath, targetPath) {
+    try {
+        return await assertExistingInside(rootPath, targetPath);
+    }
+    catch {
+        const parent = await assertExistingInside(rootPath, dirname(targetPath));
+        const target = resolve(parent, basename(targetPath));
+        ensurePathInside(await fsRealpath(rootPath), target);
+        return target;
+    }
+}
+async function assertPotentialInside(rootPath, targetPath) {
+    const root = await fsRealpath(rootPath);
+    const target = resolve(targetPath);
+    ensurePathInside(root, target);
+    return target;
+}
+function ensurePathInside(root, target) {
+    const rel = relative(root, target);
+    if (rel === "" || (!rel.startsWith("..") && !isAbsolute(rel) && (sep === "/" || !rel.split(sep).includes(".."))))
+        return;
+    throw new Error("Path is outside the managed project sandbox");
+}
+function truncateToolOutput(value, limit = 64_000) {
+    if (value.length <= limit)
+        return value;
+    return `${value.slice(0, limit)}\n[output truncated]`;
+}
 export class PiSessionService {
     constructor(events, deps = {}) {
         this.events = events;
@@ -58,6 +300,7 @@ export class PiSessionService {
         this.compactionPromptQueues = new Map();
         this.compactionDrainTimers = new Map();
         this.authLossWarnings = new Set();
+        this.managementContexts = new Map();
         this.archiveStore = deps.archiveStore ?? new SessionArchiveStore();
         this.agentDir = deps.agentDir ?? getAgentDir();
         this.sessionManager = deps.sessionManager ?? SessionManager;
@@ -88,6 +331,7 @@ export class PiSessionService {
         this.activities.clear();
         this.compactionPromptQueues.clear();
         this.authLossWarnings.clear();
+        this.managementContexts.clear();
         await Promise.all(activeSessions.map(async (active) => {
             active.unsubscribe();
             this.workspaceActivity?.removeSession(active.runtime.session.sessionId, active.runtime.session.sessionManager.getCwd());
@@ -110,9 +354,11 @@ export class PiSessionService {
             .filter(isDefined);
         return [...unarchivedSessions, ...archivedSessions];
     }
-    async start(cwd) {
-        const active = await this.create(this.sessionManager.create(cwd), cwd);
+    async start(cwd, managementContext) {
+        const active = await this.create(this.sessionManager.create(cwd), cwd, managementContext);
         const { session } = active.runtime;
+        if (managementContext !== undefined)
+            this.managementContexts.set(session.sessionId, managementContext);
         return {
             id: session.sessionId,
             path: session.sessionFile ?? "",
@@ -200,9 +446,9 @@ export class PiSessionService {
         }
         return commands.sort((a, b) => a.name.localeCompare(b.name));
     }
-    async prompt(sessionId, text, streamingBehavior) {
+    async prompt(sessionId, text, streamingBehavior, managementContext) {
         await this.assertWritable(sessionId);
-        const session = await this.getOrOpen(sessionId);
+        const session = await this.getOrOpen(sessionId, managementContext);
         this.maybeGenerateSessionName(session, text);
         const isQueued = session.isStreaming || session.isCompacting;
         const behavior = isQueued ? streamingBehavior ?? "followUp" : undefined;
@@ -236,8 +482,12 @@ export class PiSessionService {
         this.publishActivity(session, "message queued during compaction", "active");
         this.publishStatus(session);
     }
-    async shell(sessionId, text) {
+    async shell(sessionId, text, managementContext) {
         await this.assertWritable(sessionId);
+        this.assertManagedSessionAccess(sessionId, managementContext);
+        const effectiveContext = managementContext ?? this.managementContexts.get(sessionId);
+        if (effectiveContext !== undefined && !managementToolAllowed(effectiveContext, "shell"))
+            throw new Error("Shell commands are disabled in management embed mode");
         const active = await this.getActive(sessionId);
         const { session } = active.runtime;
         const isExcluded = text.startsWith("!!");
@@ -271,8 +521,9 @@ export class PiSessionService {
             this.publishStatus(session);
         });
     }
-    async runCommand(sessionId, text) {
+    async runCommand(sessionId, text, managementContext) {
         await this.assertWritable(sessionId);
+        this.assertManagedSessionAccess(sessionId, managementContext);
         return this.commandService.run(sessionId, text);
     }
     async respondToCommand(sessionId, requestId, value) {
@@ -285,8 +536,23 @@ export class PiSessionService {
             throw new Error("Stop current session activity before archiving");
         const archiveInput = await this.archiveInputForSession(session);
         await this.closeActive(session.sessionId);
+        this.managementContexts.delete(session.sessionId);
         await this.archiveStore.archive(archiveInput);
     }
+    assertManagedSessionAccess(sessionId, context) {
+        const existing = this.managementContexts.get(sessionId);
+        if (context === undefined || existing === undefined)
+            return;
+        if (existing.user.rootUserId !== context.user.rootUserId)
+            throw new Error("Session is outside the managed embed authorization scope");
+    }
+    rememberManagedSessionAccess(sessionId, context) {
+        if (context === undefined)
+            return;
+        this.assertManagedSessionAccess(sessionId, context);
+        if (!this.managementContexts.has(sessionId))
+            this.managementContexts.set(sessionId, context);
+    }
     async archiveTree(sessionId) {
         const session = await this.getOrOpen(sessionId);
         const catalog = await this.workspaceArchiveCandidates(session.sessionManager.getCwd());
@@ -417,6 +683,7 @@ export class PiSessionService {
         if (!active)
             return;
         this.active.delete(sessionId);
+        this.managementContexts.delete(sessionId);
         this.activities.delete(sessionId);
         this.workspaceActivity?.removeSession(sessionId, active.runtime.session.sessionManager.getCwd());
         this.clearAuthLossWarningsForSession(sessionId);
@@ -434,23 +701,41 @@ export class PiSessionService {
         if (await this.archiveStore.isArchived(sessionId))
             throw new Error("Archived sessions are read-only. Restore the session to continue.");
     }
-    async getOrOpen(sessionId) {
-        return (await this.getActive(sessionId)).runtime.session;
+    async getOrOpen(sessionId, managementContext) {
+        return (await this.getActive(sessionId, managementContext)).runtime.session;
     }
-    async getActive(sessionId) {
+    async getActive(sessionId, managementContext) {
         const active = this.active.get(sessionId);
-        if (active)
+        if (active) {
+            const activeSessionId = active.runtime.session.sessionId;
+            const existingContext = this.managementContexts.get(activeSessionId);
+            if (managementContext !== undefined && existingContext === undefined) {
+                const sessionFile = active.runtime.session.sessionFile;
+                if (sessionFile === undefined || sessionFile === "")
+                    throw new Error("Managed embed session must be persisted before it can be resumed safely");
+                const activeCwd = active.runtime.session.sessionManager.getCwd();
+                await this.closeActive(activeSessionId);
+                return this.create(this.sessionManager.open(sessionFile), activeCwd, managementContext);
+            }
+            this.rememberManagedSessionAccess(active.runtime.session.sessionId, managementContext);
             return active;
+        }
         const archived = await this.archiveStore.get(sessionId);
         if (archived?.archivePath !== undefined)
-            return this.create(this.sessionManager.open(archived.archivePath), archived.cwd);
+            return this.create(this.sessionManager.open(archived.archivePath), archived.cwd, managementContext);
         const match = (await this.sessionManager.listAll()).find((s) => s.id === sessionId || s.id.startsWith(sessionId));
         if (!match)
             throw new Error("Session not found");
-        return this.create(this.sessionManager.open(match.path), match.cwd);
-    }
-    async create(sessionManager, cwd) {
-        const runtime = await this.createAgentRuntime(this.createRuntime, { cwd, agentDir: this.agentDir, sessionManager });
+        return this.create(this.sessionManager.open(match.path), match.cwd, managementContext);
+    }
+    async create(sessionManager, cwd, managementContext) {
+        const createRuntime = managementContext === undefined
+            ? this.createRuntime
+            : createManagementRuntimeFactory(this.modelRegistry.authStorage, this.modelRegistry, managementContext);
+        const runtimeOptions = managementContext === undefined
+            ? { cwd, agentDir: this.agentDir, sessionManager }
+            : { cwd, agentDir: this.agentDir, sessionManager, managementContext };
+        const runtime = await this.createAgentRuntime(createRuntime, runtimeOptions);
         const active = { runtime, unsubscribe: noop };
         this.bindRuntime(active);
         runtime.setRebindSession(() => {
@@ -458,6 +743,7 @@ export class PiSessionService {
             return Promise.resolve();
         });
         this.active.set(runtime.session.sessionId, active);
+        this.rememberManagedSessionAccess(runtime.session.sessionId, managementContext);
         this.publishStatus(runtime.session);
         return active;
     }
@@ -467,6 +753,11 @@ export class PiSessionService {
         for (const [sessionId, candidate] of this.active.entries()) {
             if (candidate === active) {
                 this.active.delete(sessionId);
+                const context = this.managementContexts.get(sessionId);
+                if (context !== undefined && sessionId !== session.sessionId) {
+                    this.managementContexts.delete(sessionId);
+                    this.managementContexts.set(session.sessionId, context);
+                }
                 if (sessionId !== session.sessionId)
                     this.clearCompactionPromptQueue(sessionId);
             }
@@ -890,6 +1181,9 @@ function archivedTimestamp(record) {
 function isDefined(value) {
     return value !== undefined;
 }
+function isNodeErrorWithCode(error, code) {
+    return typeof error === "object" && error !== null && "code" in error && error.code === code;
+}
 async function clearParentSession(sessionFile) {
     const content = await readFile(sessionFile, "utf8");
     const newlineIndex = content.indexOf("\n");