pi-famulus 0.1.2 → 0.1.4

package/src/subagent/pi-runtime.ts

@@ -11,12 +11,12 @@
  * allowlist plus injected custom tools (child bash, M4 comms), and wraps it
  * into a ChildSessionAdapter.
  *
- * Child sessions receive an isolated DefaultResourceLoader with extensions
- * disabled; createAgentSession binds whatever the loader returns. Thus the
- * subagent tool is never present in a child session — the depth-1 cap holds
- * by construction.
+ * Child sessions discover user resources and built-in extensions normally.
+ * A session-local event bus identifies child loaders before the famulus
+ * extension initializes, preventing recursive parent-tool registration.
  */
 import type {
+  ExtensionAPI,
   ExtensionContext,
   ModelRegistry,
   ToolDefinition,
@@ -28,7 +28,8 @@ import {
   type ModelCandidate,
 } from "./model-spec";
 import { turnsFromMessages } from "./conversation";
-import type { ChildRunRequest, ChildSessionAdapter, CreateSessionFn } from "./types";
+import { createTokenUsageAccumulator, normalizedTokenUsage } from "./usage";
+import type { ChildRunRequest, ChildSessionAdapter, ChildSessionEvent, CreateSessionFn } from "./types";
 
 /** Structural subset of the pi module namespace we rely on. */
 type PiModule = typeof import("@earendil-works/pi-coding-agent");
@@ -37,6 +38,34 @@ type PiAgentSession = Awaited<ReturnType<PiModule["createAgentSession"]>>["sessi
 
 type Model = NonNullable<ExtensionContext["model"]>;
 
+export function createPiUsageAdapter() {
+  const totals = createTokenUsageAccumulator();
+  return {
+    tokenUsage: () => totals.snapshot(),
+    adapt(event: {
+      type: string;
+      toolCallId?: string;
+      message?: {
+        role?: string;
+        usage?: { input?: unknown; output?: unknown; cacheRead?: unknown; cacheWrite?: unknown };
+      };
+    }): ChildSessionEvent {
+      const message = event.message;
+      let usage: ReturnType<typeof normalizedTokenUsage> | undefined;
+      if (event.type === "message_end" && message?.usage) {
+        usage = normalizedTokenUsage(message.usage);
+        totals.add(usage);
+      }
+      return {
+        type: event.type,
+        ...(typeof event.toolCallId === "string" ? { toolCallId: event.toolCallId } : {}),
+        ...(typeof message?.role === "string" ? { role: message.role } : {}),
+        ...(usage ? { usage } : {}),
+      };
+    },
+  };
+}
+
 export interface PiRuntimeDeps {
   /** Parent session model registry (ctx.modelRegistry). */
   getModelRegistry: () => ModelRegistry | null;
@@ -58,6 +87,10 @@ export interface PiRuntimeDeps {
   getScopedModels?: () => readonly { model: Model; thinkingLevel?: string }[];
   /** Working directory for child sessions. */
   getCwd: () => string;
+  /** Global resource directory; defaults to pi.getAgentDir(). */
+  getAgentDir?: () => string;
+  /** Preserve the parent's project trust decision. */
+  getProjectTrusted?: () => boolean;
   /**
    * Custom tools injected into every child session (e.g. the no-background
    * bash variant, M4 contact_supervisor). Called per child request.
@@ -137,13 +170,13 @@ function resolveModel(deps: PiRuntimeDeps, req: ChildRunRequest): ResolvedModel
     if (!res.ok) throw new Error(modelResolutionError(req.model, res));
     const model = findModel(res.provider, res.id);
     if (!model) throw new Error(`model ${res.provider}/${res.id} is not usable in this session`);
-    return { model, thinkingOverride: res.thinking };
+    return { model, thinkingOverride: res.thinking, warning: res.warning };
   }
   if (req.agent.model !== undefined) {
     const res = resolveModelSpec(req.agent.model, candidates);
     if (res.ok) {
       const model = findModel(res.provider, res.id);
-      if (model) return { model, thinkingOverride: res.thinking };
+      if (model) return { model, thinkingOverride: res.thinking, warning: res.warning };
     }
     return {
       model: deps.getParentModel(),
@@ -156,11 +189,17 @@ function resolveModel(deps: PiRuntimeDeps, req: ChildRunRequest): ResolvedModel
 /** Wrap a real AgentSession into the pi-free ChildSessionAdapter. */
 function wrapSession(
   session: PiAgentSession,
-  extras: { warning?: string; resolvedModel?: string } = {},
+  extras: { warning?: string; resolvedModel?: string; effectiveThinkingLevel?: string } = {},
 ): ChildSessionAdapter {
+  const usageAdapter = createPiUsageAdapter();
+  let disposal: Promise<void> | undefined;
+  // Children receive all queued steering messages in one drain/turn rather
+  // than the interactive one-at-a-time default, which can strand wake bursts.
+  session.setSteeringMode("all");
   return {
     ...(extras.warning !== undefined ? { warning: extras.warning } : {}),
     ...(extras.resolvedModel !== undefined ? { resolvedModel: extras.resolvedModel } : {}),
+    ...(extras.effectiveThinkingLevel !== undefined ? { effectiveThinkingLevel: extras.effectiveThinkingLevel } : {}),
     // pi ≥0.99 returns QueuedInputDisposition ("handled" | "queued") from
     // input methods; the adapter contract stays Promise<void>.
     prompt: async (text) => { await session.prompt(text); },
@@ -180,33 +219,64 @@ function wrapSession(
       };
     },
     getConversation: () => turnsFromMessages(session.messages),
+    tokenUsage: () => usageAdapter.tokenUsage(),
     getActiveToolNames: () => session.getActiveToolNames(),
     getSystemPrompt: () => session.systemPrompt,
     isStreaming: () => session.isStreaming,
-    subscribe: (listener) => session.subscribe((event) => listener({ type: event.type })),
-    dispose: () => session.dispose(),
+    getSteeringMode: () => session.steeringMode,
+    subscribe: (listener) => session.subscribe((event) => {
+      listener(usageAdapter.adapt(event as Parameters<typeof usageAdapter.adapt>[0]));
+    }),
+    dispose: () => disposal ??= (async () => {
+      try {
+        await session.abort();
+        await session.extensionRunner.emit({ type: "session_shutdown", reason: "quit" });
+      } finally {
+        session.dispose();
+      }
+    })(),
   };
 }
 
-/** Options passed to createAgentSession. Tested without loading pi. */
-export function childResourceLoaderOptions(input: { cwd: string; agentDir: string }): {
+const CHILD_ROLE_CHANNEL = "pi-famulus:child-role";
+
+/** The loader's bus answers synchronously, before any extension initialization. */
+export function isFamulusChildSession(pi: Pick<ExtensionAPI, "events">): boolean {
+  if (!pi.events) {
+    throw new Error("pi-famulus requires the extension event bus to identify child sessions safely");
+  }
+  const probe = { child: false };
+  pi.events.emit(CHILD_ROLE_CHANNEL, probe);
+  return probe.child;
+}
+
+/** Keep configuration scopes and paths intact; never persist child overrides. */
+export async function createChildResources(input: {
   cwd: string;
   agentDir: string;
-  noExtensions: true;
-  noSkills: true;
-  noPromptTemplates: true;
-  noThemes: true;
-  noContextFiles: true;
-} {
-  return {
+  projectTrusted: boolean;
+}) {
+  const pi = await importPi();
+  const settingsManager = pi.SettingsManager.create(input.cwd, input.agentDir, {
+    projectTrusted: input.projectTrusted,
+  });
+  const eventBus = pi.createEventBus();
+  eventBus.on(CHILD_ROLE_CHANNEL, (probe) => {
+    if (probe && typeof probe === "object" && "child" in probe) probe.child = true;
+  });
+  const resourceLoader = new pi.DefaultResourceLoader({
     cwd: input.cwd,
     agentDir: input.agentDir,
-    noExtensions: true,
-    noSkills: true,
-    noPromptTemplates: true,
-    noThemes: true,
-    noContextFiles: true,
-  };
+    settingsManager,
+    eventBus,
+    extensionFactories: [
+      { name: "codemode", factory: pi.createCodemodeExtension(), builtin: true, replaceable: true },
+      { name: "tool-search", factory: pi.createToolSearchExtension(), builtin: true, replaceable: true },
+      { name: "mcp", factory: pi.createMcpExtension(), builtin: true, replaceable: true },
+    ],
+  });
+  await resourceLoader.reload();
+  return { resourceLoader, settingsManager };
 }
 
 export function childSessionCreateOptions(input: {
@@ -222,13 +292,18 @@ export function childSessionCreateOptions(input: {
   model: unknown;
   thinkingLevel: unknown;
   tools?: string[];
+  excludeTools?: string[];
   customTools?: unknown[];
   modelRuntime?: unknown;
   resourceLoader?: unknown;
 } {
-  const tools = input.tools && input.tools.length > 0 ? [...input.tools] : undefined;
-  const customTools = input.customTools?.map((customTool) => {
+  const tools = input.tools !== undefined ? [...input.tools] : undefined;
+  const customTools = input.customTools?.filter((customTool) => {
+    const name = (customTool as { name?: unknown }).name;
+    return name !== "bash" || tools === undefined || tools.includes("bash");
+  }).map((customTool) => {
     const tool = customTool as { name?: unknown; promptGuidelines?: string[] };
+    if (tool.name === "bash") return { ...tool, defaultActive: false };
     if (tool.name === "contact_supervisor") {
       return {
         ...tool,
@@ -247,7 +322,11 @@ export function childSessionCreateOptions(input: {
     cwd: input.cwd,
     model: input.model,
     thinkingLevel: input.thinkingLevel,
-    ...(tools ? { tools } : {}),
+    ...(tools !== undefined ? { tools } : {}),
+    // pi 1.0.4+ otherwise retains MCP tools behind non-MCP allowlists.
+    ...(tools !== undefined && !tools.some((name) => name.startsWith("mcp__"))
+      ? { excludeTools: ["mcp__*"] }
+      : {}),
     ...(customTools && customTools.length > 0 ? { customTools } : {}),
     ...(input.modelRuntime ? { modelRuntime: input.modelRuntime } : {}),
     ...(input.resourceLoader ? { resourceLoader: input.resourceLoader } : {}),
@@ -267,31 +346,37 @@ export function createPiSessionFn(deps: PiRuntimeDeps): CreateSessionFn {
       req.agent.thinking ??
       deps.getParentThinkingLevel();
     const cwd = deps.getCwd();
-    const resourceLoader = new pi.DefaultResourceLoader(
-      childResourceLoaderOptions({ cwd, agentDir: pi.getAgentDir() }),
-    );
-    await resourceLoader.reload();
+    const agentDir = deps.getAgentDir?.() ?? pi.getAgentDir();
+    const { resourceLoader, settingsManager } = await createChildResources({
+      cwd,
+      agentDir,
+      projectTrusted: deps.getProjectTrusted?.() ?? false,
+    });
     const customTools = deps.customTools?.(req) ?? [];
     const modelRuntime = deps.getModelRuntime?.();
     const options = childSessionCreateOptions({
       cwd,
       model: resolved.model,
       thinkingLevel,
-      tools: req.agent.tools.length > 0 ? [...req.agent.tools] : undefined,
+      tools: req.agent.tools,
       customTools: customTools.length > 0 ? customTools : undefined,
       modelRuntime,
       resourceLoader,
     });
     const { session } = await pi.createAgentSession({
       ...options,
+      agentDir,
+      settingsManager,
       sessionManager: pi.SessionManager.inMemory(cwd),
     } as never);
+    await session.bindExtensions({});
     const resolvedModel = resolved.model
       ? `${resolved.model.provider}/${resolved.model.id}`
       : undefined;
     return wrapSession(session, {
       ...(resolved.warning !== undefined ? { warning: resolved.warning } : {}),
       ...(resolvedModel !== undefined ? { resolvedModel } : {}),
+      ...(typeof thinkingLevel === "string" ? { effectiveThinkingLevel: thinkingLevel } : {}),
     });
   };
 }