lsd-pi 1.2.4 → 1.3.6

package/dist/resources/extensions/subagent/in-process-runner.js

@@ -0,0 +1,387 @@
+import { createAgentSession, getAgentDir, SessionManager } from "@gsd/pi-coding-agent";
+import { resolveConfiguredSubagentModel } from "./configured-model.js";
+import { resolveSubagentModel } from "./model-resolution.js";
+import { readBudgetSubagentModelFromSettings } from "./legacy-runner.js";
+import { loadEffectivePreferences } from "../shared/preferences.js";
+const MAX_AGENT_DURATION_MS = 30 * 60 * 1000;
+export const MAX_IN_PROCESS_SUBAGENT_DEPTH = 3;
+export const MAX_ACTIVE_DESCENDANTS = 8;
+const NESTED_EXTENSION_TOOL_ALLOWLIST = new Set([
+    "subagent",
+    "await_subagent",
+    "bg_shell",
+    "fetch_page",
+    "resolve_library",
+    "get_library_docs",
+    "tool_search",
+    "tool_enable",
+    "Skill",
+]);
+const parentSessionIdByChildSessionId = new Map();
+const childSessionIdsByParentSessionId = new Map();
+const handleBySessionId = new Map();
+let activeDescendantCount = 0;
+function getFinalOutput(messages) {
+    for (let i = messages.length - 1; i >= 0; i--) {
+        const msg = messages[i];
+        if (msg.role === "assistant") {
+            for (const part of msg.content) {
+                if (part.type === "text")
+                    return part.text;
+            }
+        }
+    }
+    return "";
+}
+function toErrorString(error) {
+    return error instanceof Error ? error.message : String(error);
+}
+function createUnknownAgentResult(agents, agentName, task, step, parentSessionFile) {
+    const available = agents.map((a) => `"${a.name}"`).join(", ") || "none";
+    return {
+        agent: agentName,
+        agentSource: "unknown",
+        task,
+        exitCode: 1,
+        messages: [],
+        stderr: `Unknown agent: "${agentName}". Available agents: ${available}.`,
+        usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0, contextTokens: 0, turns: 0 },
+        step,
+        parentSessionFile,
+    };
+}
+function createResolvedHandle(result) {
+    const resultPromise = Promise.resolve(result);
+    const handle = {
+        result: resultPromise,
+        isBusy: () => false,
+        prompt: async () => undefined,
+        steer: async () => undefined,
+        followUp: async () => undefined,
+        abort: () => undefined,
+        dispose: () => undefined,
+    };
+    return { handle, currentResult: result, resultPromise };
+}
+function reserveDescendantSlot(depth) {
+    if (depth <= 1)
+        return () => undefined;
+    if (activeDescendantCount >= MAX_ACTIVE_DESCENDANTS)
+        return null;
+    activeDescendantCount += 1;
+    let released = false;
+    return () => {
+        if (released)
+            return;
+        released = true;
+        activeDescendantCount = Math.max(0, activeDescendantCount - 1);
+    };
+}
+function linkParentChildSession(parentSessionId, childSessionId) {
+    parentSessionIdByChildSessionId.set(childSessionId, parentSessionId);
+    const children = childSessionIdsByParentSessionId.get(parentSessionId) ?? new Set();
+    children.add(childSessionId);
+    childSessionIdsByParentSessionId.set(parentSessionId, children);
+}
+function unlinkSessionFromParent(sessionId) {
+    const parentSessionId = parentSessionIdByChildSessionId.get(sessionId);
+    if (!parentSessionId)
+        return;
+    parentSessionIdByChildSessionId.delete(sessionId);
+    const siblings = childSessionIdsByParentSessionId.get(parentSessionId);
+    if (!siblings)
+        return;
+    siblings.delete(sessionId);
+    if (siblings.size === 0) {
+        childSessionIdsByParentSessionId.delete(parentSessionId);
+    }
+}
+function abortDescendantTree(sessionId, visited = new Set()) {
+    if (visited.has(sessionId))
+        return;
+    visited.add(sessionId);
+    const descendants = childSessionIdsByParentSessionId.get(sessionId);
+    if (!descendants || descendants.size === 0)
+        return;
+    for (const childSessionId of descendants) {
+        handleBySessionId.get(childSessionId)?.abort();
+        abortDescendantTree(childSessionId, visited);
+    }
+}
+function resolveRequestedToolNames(agent, defaultActiveToolNames) {
+    const requested = agent.tools && agent.tools.length > 0 ? agent.tools : defaultActiveToolNames;
+    return [...new Set(requested)];
+}
+function applyNestedExtensionToolPolicy(toolNames, extensionToolNames, nestingDepth) {
+    if (nestingDepth <= 1) {
+        return { activeToolNames: toolNames, droppedTools: [] };
+    }
+    const activeToolNames = [];
+    const droppedTools = [];
+    for (const toolName of toolNames) {
+        if (!extensionToolNames.has(toolName)) {
+            activeToolNames.push(toolName);
+            continue;
+        }
+        if (NESTED_EXTENSION_TOOL_ALLOWLIST.has(toolName)) {
+            activeToolNames.push(toolName);
+        }
+        else {
+            droppedTools.push(toolName);
+        }
+    }
+    return { activeToolNames, droppedTools };
+}
+export async function startInProcessSingleAgent(defaultCwd, agents, agentName, task, cwd, step, modelOverride, parentModel, signal, onUpdate, makeDetails, parentSessionFile, depth, parentSessionId, ancestry, onSubagentEvent) {
+    const agent = agents.find((candidate) => candidate.name === agentName);
+    if (!agent) {
+        return createResolvedHandle(createUnknownAgentResult(agents, agentName, task, step, parentSessionFile));
+    }
+    const nestingDepth = depth ?? 1;
+    if (nestingDepth > MAX_IN_PROCESS_SUBAGENT_DEPTH) {
+        return createResolvedHandle({
+            agent: agentName,
+            agentSource: agent.source,
+            task,
+            exitCode: 1,
+            messages: [],
+            stderr: `Max subagent depth (${MAX_IN_PROCESS_SUBAGENT_DEPTH}) exceeded.`,
+            usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0, contextTokens: 0, turns: 0 },
+            step,
+            parentSessionFile,
+        });
+    }
+    const parentAgentName = ancestry && ancestry.length > 0 ? ancestry[ancestry.length - 1] : undefined;
+    if (parentAgentName && parentAgentName === agentName) {
+        return createResolvedHandle({
+            agent: agentName,
+            agentSource: agent.source,
+            task,
+            exitCode: 1,
+            messages: [],
+            stderr: `Subagent "${agentName}" cannot spawn another subagent with the same name as itself.`,
+            usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0, contextTokens: 0, turns: 0 },
+            step,
+            parentSessionFile,
+        });
+    }
+    const releaseDescendantSlot = reserveDescendantSlot(nestingDepth);
+    if (nestingDepth > 1 && !releaseDescendantSlot) {
+        return createResolvedHandle({
+            agent: agentName,
+            agentSource: agent.source,
+            task,
+            exitCode: 1,
+            messages: [],
+            stderr: `Maximum active descendant subagents (${MAX_ACTIVE_DESCENDANTS}) reached. Wait for running descendants to finish or cancel one before spawning more.`,
+            usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0, contextTokens: 0, turns: 0 },
+            step,
+            parentSessionFile,
+        });
+    }
+    const releaseDescendantSlotOnce = (() => {
+        let released = false;
+        return () => {
+            if (released)
+                return;
+            released = true;
+            releaseDescendantSlot?.();
+        };
+    })();
+    const preferences = loadEffectivePreferences()?.preferences;
+    const settingsBudgetModel = readBudgetSubagentModelFromSettings();
+    const resolvedModel = resolveConfiguredSubagentModel(agent, preferences, settingsBudgetModel);
+    const inferredModel = resolveSubagentModel({ name: agent.name, model: resolvedModel }, { overrideModel: modelOverride, parentModel });
+    const currentResult = {
+        agent: agentName,
+        agentSource: agent.source,
+        task,
+        exitCode: 0,
+        messages: [],
+        stderr: "",
+        usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0, contextTokens: 0, turns: 0 },
+        model: inferredModel,
+        step,
+        parentSessionFile,
+    };
+    const emitUpdate = () => {
+        if (!onUpdate)
+            return;
+        onUpdate({
+            content: [{ type: "text", text: getFinalOutput(currentResult.messages) || "(running...)" }],
+            details: makeDetails([currentResult]),
+        });
+    };
+    const effectiveCwd = cwd ?? defaultCwd;
+    const sessionManager = SessionManager.inMemory(effectiveCwd);
+    let session;
+    try {
+        const created = await createAgentSession({
+            cwd: effectiveCwd,
+            agentDir: getAgentDir(),
+            sessionManager,
+        });
+        session = created.session;
+        await session.bindExtensions({});
+    }
+    catch (error) {
+        releaseDescendantSlotOnce();
+        currentResult.exitCode = 1;
+        currentResult.stderr = `Failed to initialize in-process subagent session: ${toErrorString(error)}`;
+        return createResolvedHandle(currentResult);
+    }
+    if (inferredModel) {
+        const [provider, modelId] = inferredModel.split("/");
+        if (provider && modelId) {
+            const model = session.modelRegistry.find(provider, modelId);
+            if (!model) {
+                session.dispose();
+                releaseDescendantSlotOnce();
+                currentResult.exitCode = 1;
+                currentResult.stderr = `Unable to resolve model ${inferredModel} for in-process subagent.`;
+                return createResolvedHandle(currentResult);
+            }
+            await session.setModel(model);
+            currentResult.model = `${model.provider}/${model.id}`;
+        }
+    }
+    const defaultActiveToolNames = session.getActiveToolNames();
+    const requestedToolNames = resolveRequestedToolNames(agent, defaultActiveToolNames);
+    const extensionToolNames = new Set(session.extensionRunner?.getAllRegisteredTools().map((registered) => registered.definition.name) ?? []);
+    const { activeToolNames: resolvedActiveToolNames, droppedTools } = applyNestedExtensionToolPolicy(requestedToolNames, extensionToolNames, nestingDepth);
+    session.setActiveToolsByName(resolvedActiveToolNames);
+    const ancestryChain = ancestry && ancestry.length > 0 ? ancestry.join(" -> ") : "root";
+    const nestedPolicyLine = droppedTools.length > 0
+        ? `Nested safety policy disabled extension tools: ${droppedTools.join(", ")}.`
+        : "";
+    const antiRecursionPrompt = [
+        `You are already the ${agentName} subagent for this session.`,
+        "Do not spawn or delegate to another subagent with the same name as yourself.",
+        `If the user asks you to continue ${agentName} work, do that work directly in this session.`,
+        `Original delegated task: ${task}`,
+        `Current subagent depth: ${nestingDepth}/${MAX_IN_PROCESS_SUBAGENT_DEPTH}.`,
+        `Subagent ancestry: ${ancestryChain}.`,
+        nestedPolicyLine,
+    ].filter(Boolean).join("\n");
+    const appendedParts = [antiRecursionPrompt, agent.systemPrompt.trim()].filter((part) => part.length > 0);
+    if (appendedParts.length > 0) {
+        const appendedPrompt = `${session.systemPrompt}\n\n${appendedParts.join("\n\n")}`;
+        session.agent.setSystemPrompt(appendedPrompt);
+    }
+    let completed = false;
+    let disposed = false;
+    const unsubscribe = session.subscribe((event) => {
+        onSubagentEvent?.(event, currentResult);
+        if (event?.type !== "message_end" || !event.message)
+            return;
+        const message = event.message;
+        currentResult.messages.push(message);
+        if (message.role === "assistant") {
+            currentResult.usage.turns++;
+            const usage = message.usage;
+            if (usage) {
+                currentResult.usage.input += usage.input || 0;
+                currentResult.usage.output += usage.output || 0;
+                currentResult.usage.cacheRead += usage.cacheRead || 0;
+                currentResult.usage.cacheWrite += usage.cacheWrite || 0;
+                currentResult.usage.cost += usage.cost?.total || 0;
+                currentResult.usage.contextTokens = usage.totalTokens || 0;
+            }
+            if (message.model)
+                currentResult.model = message.model;
+            if (message.stopReason)
+                currentResult.stopReason = message.stopReason;
+            if (message.errorMessage)
+                currentResult.errorMessage = message.errorMessage;
+        }
+        emitUpdate();
+    });
+    const sessionId = session.sessionId;
+    const dispose = () => {
+        if (disposed)
+            return;
+        disposed = true;
+        if (sessionId) {
+            handleBySessionId.delete(sessionId);
+            unlinkSessionFromParent(sessionId);
+        }
+        releaseDescendantSlotOnce();
+        unsubscribe();
+        session.dispose();
+    };
+    const handle = {
+        result: Promise.resolve(currentResult),
+        sessionId,
+        isBusy: () => session.isStreaming,
+        prompt: (message, images) => session.prompt(message, { images }),
+        steer: (message, images) => session.steer(message, images),
+        followUp: (message, images) => session.followUp(message, images),
+        abort: () => {
+            if (sessionId) {
+                abortDescendantTree(sessionId);
+            }
+            void session.abort();
+        },
+        dispose,
+    };
+    if (sessionId) {
+        handleBySessionId.set(sessionId, handle);
+        if (parentSessionId) {
+            linkParentChildSession(parentSessionId, sessionId);
+        }
+    }
+    let timeoutTimer;
+    const timeoutPromise = new Promise((_, reject) => {
+        timeoutTimer = setTimeout(() => {
+            handle.abort();
+            reject(new Error(`In-process subagent timed out after ${MAX_AGENT_DURATION_MS}ms.`));
+        }, MAX_AGENT_DURATION_MS);
+        if (typeof timeoutTimer === "object" && "unref" in timeoutTimer)
+            timeoutTimer.unref();
+    });
+    const runPromise = (async () => {
+        if (signal?.aborted) {
+            handle.abort();
+            throw new Error("Subagent was aborted");
+        }
+        const onAbort = () => handle.abort();
+        signal?.addEventListener("abort", onAbort, { once: true });
+        try {
+            await session.prompt(task);
+            completed = true;
+        }
+        finally {
+            signal?.removeEventListener("abort", onAbort);
+        }
+    })();
+    const resultPromise = (async () => {
+        try {
+            await Promise.race([runPromise, timeoutPromise]);
+        }
+        catch (error) {
+            currentResult.exitCode = 1;
+            currentResult.stderr = toErrorString(error);
+            if (signal?.aborted)
+                currentResult.stopReason = "aborted";
+        }
+        finally {
+            if (timeoutTimer)
+                clearTimeout(timeoutTimer);
+            if (!completed && signal?.aborted) {
+                currentResult.stopReason = "aborted";
+            }
+            dispose();
+        }
+        return currentResult;
+    })();
+    handle.result = resultPromise;
+    return {
+        handle,
+        currentResult,
+        resultPromise,
+    };
+}
+export async function runInProcessSingleAgent(defaultCwd, agents, agentName, task, cwd, step, modelOverride, parentModel, signal, onUpdate, makeDetails, parentSessionFile, depth, parentSessionId, ancestry, onSubagentEvent) {
+    const started = await startInProcessSingleAgent(defaultCwd, agents, agentName, task, cwd, step, modelOverride, parentModel, signal, onUpdate, makeDetails, parentSessionFile, depth, parentSessionId, ancestry, onSubagentEvent);
+    return started.resultPromise;
+}