pi-famulus 0.1.3 → 0.1.4

package/src/index.ts

@@ -5,7 +5,7 @@
  * M2: NotifyCenter + monitor tool.
  * M3: subagent tool (InProcessRunner + tasks/chain + budget-to-async) + fleet widget.
  */
-import { applyBehaviorGuidelines } from "./behavior-guidelines";
+import { registerBehaviorGuidelines } from "./behavior-guidelines";
 import { realClock } from "./clock";
 import { ExitNotifyGate } from "./exit-notify-gate";
 import { ExitWatchdog } from "./exit-watchdog";
@@ -27,6 +27,7 @@ import { describeManagerSearch, getFamulusHome, loadConfig, resolveManagerPath,
 import type { TaskExitInfo } from "./format";
 import { ManagerClient, type ManagerEvent, type TaskRecord } from "./manager-client";
 import { createMonitorTool, exitEventFromRecord, MonitorRegistry } from "./monitor";
+import { admitAgentChild, reregisterAgentLeases } from "./subagent/admission";
 import { NotifyCenter } from "./notify";
 import { onWakeMessageEnd, registryStatusLookup } from "./wake-delivery";
 import { createChildBashTool } from "./subagent/child-bash";
@@ -35,13 +36,14 @@ import {
   agentEndReason,
   headOf,
   tailOf,
+  updateAgentChildMetrics,
   writeAgentChildRecord,
   type AgentChildRecord,
 } from "./subagent/agent-records";
 import { TranscriptWriter } from "./subagent/transcript";
 import { FleetWidget } from "./subagent/fleet-widget";
 import { WorkIndex, type WorkItem } from "./work-index";
-import { createPiSessionFn, modelCandidates } from "./subagent/pi-runtime";
+import { createPiSessionFn, isFamulusChildSession, modelCandidates } from "./subagent/pi-runtime";
 import { SubagentRegistry } from "./subagent/registry";
 import { InProcessRunner } from "./subagent/runner";
 import { createSubagentTool } from "./subagent/tool";
@@ -68,6 +70,8 @@ function toExitStatus(event: ManagerEvent): TaskExitInfo["status"] {
 }
 
 export default function (pi: ExtensionAPI): void {
+  if (isFamulusChildSession(pi)) return;
+  registerBehaviorGuidelines(pi);
   const home = getFamulusHome();
   const config = loadConfig(home);
   const managerPath = resolveManagerPath(config, home);
@@ -78,6 +82,29 @@ export default function (pi: ExtensionAPI): void {
 
   let ctx: ExtensionContext | null = null;
   let client: ManagerClient | null = null;
+  const globalAgentLeases = new Map<string, string>();
+  const pendingAgentLeases = new Set<string>();
+  const pendingAgentReregistrations = new Map<string, AbortController>();
+  let leaseReconciliation: Promise<void> | null = null;
+  const reconcileAgentLeases = (manager: ManagerClient): Promise<void> => {
+    if (leaseReconciliation) return leaseReconciliation;
+    leaseReconciliation = reregisterAgentLeases(
+      manager,
+      globalAgentLeases,
+      undefined,
+      pendingAgentLeases,
+      pendingAgentReregistrations,
+    ).finally(() => {
+      leaseReconciliation = null;
+    });
+    return leaseReconciliation;
+  };
+  const capacityNotices = new Set<string>();
+  const capacityNotice = (reason: "manager unavailable" | "daemon too old"): void => {
+    if (capacityNotices.has(reason)) return;
+    capacityNotices.add(reason);
+    console.warn(`pi-famulus: machine-wide agent capacity inactive (${reason}); using per-session limit only`);
+  };
   let notifyCenter: NotifyCenter | null = null;
   let monitorRegistry: MonitorRegistry | null = null;
   let subagentRegistry: SubagentRegistry | null = null;
@@ -407,10 +434,6 @@ export default function (pi: ExtensionAPI): void {
 
   pi.on("session_start", async (_event, startCtx) => {
     ctx = startCtx;
-    const base = (
-      startCtx as { getSystemPromptOptions?: () => { sections?: Record<string, string> } }
-    ).getSystemPromptOptions?.();
-    if (base) applyBehaviorGuidelines(base);
     notifyCenter?.dispose();
     notifyCenter = new NotifyCenter({
       sendMessage: (msg, opts) => pi.sendMessage(msg, opts),
@@ -426,8 +449,13 @@ export default function (pi: ExtensionAPI): void {
     });
     fleetWidget?.dispose();
     fleetWidget = null;
-    subagentRegistry?.disposeAll();
+    comms.dispose(); // Release decision waiters before awaiting child aborts.
+    await subagentRegistry?.disposeAll().catch(() => {});
     subagentRegistry = null;
+    for (const controller of pendingAgentReregistrations.values()) controller.abort();
+    pendingAgentReregistrations.clear();
+    globalAgentLeases.clear();
+    pendingAgentLeases.clear();
 
     client = new ManagerClient({
       home,
@@ -483,7 +511,21 @@ export default function (pi: ExtensionAPI): void {
       }
     });
     client.onReconnect(() => {
+      // A daemon restart may drop permits; a same-session rebind may retain
+      // them. Keep each kind for idempotent re-registration and mark all local
+      // leases pending until the daemon confirms them.
+      for (const childId of globalAgentLeases.keys()) pendingAgentLeases.add(childId);
       void monitorRegistry?.rewatchAll().then(() => syncWithManager());
+      void (async () => {
+        const manager = client;
+        if (!manager || manager.protocolLevel() < 4) {
+          capacityNotice("daemon too old");
+          return;
+        }
+        // The registry and its child ids survive socket reconnects in this
+        // extension process; held ids include resumable interrupted children.
+        await reconcileAgentLeases(manager);
+      })().catch(() => {});
     });
 
     // M3: subagent registry + in-process runner + fleet widget. The runner's
@@ -519,6 +561,7 @@ export default function (pi: ExtensionAPI): void {
       getParentThinkingLevel: () => ctx?.thinkingLevel,
       getScopedModels: () => ctx?.scopedModels ?? [],
       getCwd: () => ctx?.cwd ?? process.cwd(),
+      getProjectTrusted: () => ctx?.isProjectTrusted() ?? false,
       customTools: (req) => {
         const tools: Array<ToolDefinition<any, any, any>> = [
           // M4: every child can reach the supervisor and its siblings.
@@ -527,21 +570,19 @@ export default function (pi: ExtensionAPI): void {
         ];
         // The no-background bash variant replaces the built-in bash inside
         // child sessions (custom tools override builtins by name).
-        if (req.agent.tools.includes("bash")) {
-          tools.push(
-            createChildBashTool({
-              getClient: () => client,
-              home,
-              sessionId: () => ctx?.sessionManager.getSessionId() ?? "",
-              sessionEnv: () => (ctx ? sessionEnv(ctx) : {}),
-              trackTask,
-              childId: req.childId,
-              runId: req.runId,
-              clock,
-              shells: childShells,
-            }),
-          );
-        }
+        tools.push(
+          createChildBashTool({
+            getClient: () => client,
+            home,
+            sessionId: () => ctx?.sessionManager.getSessionId() ?? "",
+            sessionEnv: () => (ctx ? sessionEnv(ctx) : {}),
+            trackTask,
+            childId: req.childId,
+            runId: req.runId,
+            clock,
+            shells: childShells,
+          }),
+        );
         return tools;
       },
     });
@@ -559,9 +600,37 @@ export default function (pi: ExtensionAPI): void {
       stallRetryDelayMs: subagentConfig.stallRetryDelayMs,
       overrunRepeatMs: subagentConfig.overrunRepeatMs,
       hardTimeoutMs: subagentConfig.hardTimeoutMs,
-      acquire: (req, ticket) => registry.admitChild(req.childId, ticket),
+      acquire: (req, ticket) => admitAgentChild({
+        childId: req.childId,
+        workKind: req.workKind ?? "other",
+        reserveLocal: () => registry.reserveChildSlot(req.childId, ticket),
+        manager: client,
+        leases: globalAgentLeases,
+        pendingLeases: pendingAgentLeases,
+        pendingReregistrations: pendingAgentReregistrations,
+        ticket,
+        notice: capacityNotice,
+      }),
       onActivity: (childId) => {
         syncTranscript(childId);
+        const handle = registry.handle(childId);
+        const usage = handle?.tokenUsage();
+        const wall = handle?.wallUsage();
+        if (usage && wall) {
+          const metrics = {
+            tokens_input: usage.input,
+            tokens_output: usage.output,
+            tokens_cache_read: usage.cacheRead,
+            tokens_cache_write: usage.cacheWrite,
+            llm_ms: wall.llmMs,
+            tool_ms: wall.toolMs,
+            queue_ms: wall.queueMs,
+            wall_other_ms: wall.otherMs,
+            wall_approximate: wall.approximate,
+          };
+          updateAgentChildMetrics(home, sessionIdForAgents(), childId, metrics);
+          logEvent("agent.usage", { child_id: childId, ...metrics });
+        }
       },
       onStall: (childId, attempt) => {
         // One event per stall detection: an auto-resume follows unless the
@@ -600,6 +669,8 @@ export default function (pi: ExtensionAPI): void {
             run_id: run.runId,
             name: c.name,
             agent: c.agent,
+            work_kind: c.workKind ?? "other",
+            queue_ms: c.queueMs ?? 0,
             ...(c.model ? { model: c.model } : {}),
           });
         }
@@ -611,8 +682,11 @@ export default function (pi: ExtensionAPI): void {
             child_id: c.childId,
             status: c.status,
             ...(error ? { error } : {}),
+            work_kind: c.workKind ?? "other",
             ...(c.result?.stalls ? { stalls: c.result.stalls } : {}),
+            ...(c.queueMs !== undefined ? { queue_ms: c.queueMs } : {}),
             duration_ms: c.result?.durationMs ?? Math.max(0, clock.now() - c.startedAt),
+            ...(c.result?.queueMs !== undefined ? { queue_ms: c.result.queueMs } : {}),
           });
         }
         previousAgentStatus.set(c.childId, c.status);
@@ -624,9 +698,20 @@ export default function (pi: ExtensionAPI): void {
           session_id: sid,
           name: c.name,
           agent: c.agent,
+          work_kind: c.workKind ?? "other",
+          ...(c.queueMs !== undefined ? { queue_ms: c.queueMs } : {}),
           ...(c.model !== undefined ? { model: c.model } : {}),
           status: c.status,
           started_at: c.startedAt,
+          tokens_input: c.tokenUsage?.input ?? 0,
+          tokens_output: c.tokenUsage?.output ?? 0,
+          tokens_cache_read: c.tokenUsage?.cacheRead ?? 0,
+          tokens_cache_write: c.tokenUsage?.cacheWrite ?? 0,
+          llm_ms: c.wallUsage?.llmMs ?? 0,
+          tool_ms: c.wallUsage?.toolMs ?? 0,
+          queue_ms: c.wallUsage?.queueMs ?? 0,
+          wall_other_ms: c.wallUsage?.otherMs ?? 0,
+          wall_approximate: c.wallUsage?.approximate ?? false,
           ...(c.endedAt !== undefined ? { ended_at: c.endedAt } : {}),
           ...(c.result?.error ? { error: c.result.error } : {}),
           ...(c.result?.attempts !== undefined && c.result.attempts > 1
@@ -655,6 +740,8 @@ export default function (pi: ExtensionAPI): void {
           runId: run.runId,
           name: c.name,
           agent: c.agent,
+          workKind: c.workKind ?? "other",
+          ...(c.queueMs !== undefined ? { queueMs: c.queueMs } : {}),
           ...(c.model !== undefined ? { model: c.model } : {}),
           cwd: startCtx.cwd,
           ...(c.prompt !== undefined ? { prompt: c.prompt } : {}),
@@ -679,7 +766,11 @@ export default function (pi: ExtensionAPI): void {
     void c
       .connect()
       .then((ok) => {
-        if (!ok && startCtx.hasUI) {
+        if (ok) {
+          if (c.protocolLevel() >= 4) void reconcileAgentLeases(c).catch(() => {});
+          return;
+        }
+        if (startCtx.hasUI) {
           const detail = c.lastError();
           const reason = detail ? ` (${detail})` : "";
           const searched = managerPath ? `using ${managerPath}` : `looked in: ${describeManagerSearch(config, home)}`;
@@ -698,19 +789,22 @@ export default function (pi: ExtensionAPI): void {
     exitWatchdog.dispose();
     fleetWidget?.dispose();
     fleetWidget = null;
-    subagentRegistry?.disposeAll();
+    comms.dispose(); // Release decision waiters before awaiting child aborts.
+    await subagentRegistry?.disposeAll().catch(() => {});
     subagentRegistry = null;
     agentLoader = null;
-    comms.dispose(); // resolve orphaned need_decision waiters
     monitorRegistry?.disposeAll();
     notifyCenter?.dispose();
     const current = client;
     client = null;
     if (current) {
-      if (current.isAvailable()) {
-        await current.shutdownSession().catch(() => {});
+      try {
+        if (current.isAvailable()) await current.shutdownSession();
+      } catch {
+        // Child teardown is best-effort; still release the manager connection.
+      } finally {
+        await current.close().catch(() => {});
       }
-      await current.close();
     }
     ctx = null;
   });
@@ -728,8 +822,4 @@ export default function (pi: ExtensionAPI): void {
     });
     return message ? { message } : undefined;
   });
-
-  pi.on("before_agent_start", async (event) => {
-    applyBehaviorGuidelines(event.systemPromptOptions as { sections?: Record<string, string> });
-  });
 }