@staix/agent-hub 0.12.2 → 0.12.3

package/src/local/tools.ts

@@ -14,6 +14,8 @@ export interface ToolContext {
   sandboxProfile: string;
   /** Extra environment for sandboxed commands: the egress proxy's variables when network goes through it (#65). */
   sandboxEnv?: Record<string, string>;
+  /** Turn cancellation signal, including the execution-budget wall cap. */
+  signal?: AbortSignal;
   /** Ask the console. Resolves false on deny or timeout. */
   permit: (title: string) => Promise<boolean>;
   /** Publish a message to other peers mid-turn. Returns a one-line receipt. */
@@ -106,6 +108,7 @@ export async function runTool(name: string, rawArgs: string, ctx: ToolContext):
     return "error: tool arguments were not valid JSON";
   }
   try {
+    if (ctx.signal?.aborted) return "error: turn cancelled before tool execution";
     switch (name) {
       case "read": {
         const file = guardPath(ctx, String(a.path), "read");
@@ -119,6 +122,7 @@ export async function runTool(name: string, rawArgs: string, ctx: ToolContext):
         let file = guardPath(ctx, String(a.path), "write");
         const content = String(a.content ?? "");
         if (!(await ctx.permit(`write ${a.path} (${content.length} chars):\n${preview(content)}`))) return "error: the user did not approve this write";
+        if (ctx.signal?.aborted) return "error: turn cancelled before write";
         file = guardPath(ctx, String(a.path), "write");
         mkdirSync(dirname(file), { recursive: true });
         writeFileSync(file, content);
@@ -130,6 +134,7 @@ export async function runTool(name: string, rawArgs: string, ctx: ToolContext):
         let count = text.split(String(a.old)).length - 1;
         if (!a.old || count !== 1) return `error: \`old\` must match exactly once, it matched ${count} times`;
         if (!(await ctx.permit(`edit ${a.path}:\n- ${preview(String(a.old), 600)}\n+ ${preview(String(a.new ?? ""), 600)}`))) return "error: the user did not approve this edit";
+        if (ctx.signal?.aborted) return "error: turn cancelled before edit";
         // Approval can outlive another peer's edit or a path change. Apply only the approved fragment to current bytes.
         file = guardPath(ctx, String(a.path), "write");
         text = readFileSync(file, "utf8");
@@ -144,7 +149,7 @@ export async function runTool(name: string, rawArgs: string, ctx: ToolContext):
         // The approver sees the whole command, not a prefix: what is hidden cannot be approved.
         if (command.length > 4000) return "error: command longer than 4000 characters; put it in a script file with write, then run that";
         if (!(await ctx.permit(`bash: ${command}`))) return "error: the user did not approve this command";
-        const res = await sandboxedExec(["/bin/bash", "-c", command], { cwd: ctx.cwd, profile: ctx.sandboxProfile, timeoutMs: (Number(a.timeout_s) || 120) * 1000, ...(ctx.sandboxEnv ? { env: ctx.sandboxEnv } : {}) });
+        const res = await sandboxedExec(["/bin/bash", "-c", command], { cwd: ctx.cwd, profile: ctx.sandboxProfile, timeoutMs: (Number(a.timeout_s) || 120) * 1000, ...(ctx.sandboxEnv ? { env: ctx.sandboxEnv } : {}), ...(ctx.signal ? { signal: ctx.signal } : {}) });
         return `${res.output}\n(exit ${res.code})`;
       }
       case "git": {
@@ -155,7 +160,7 @@ export async function runTool(name: string, rawArgs: string, ctx: ToolContext):
         if (problem) return `error: git ${sub}: ${problem}`;
         if (GIT_WRITE.has(sub) && !(await ctx.permit(`git ${args.join(" ")}`))) return "error: the user did not approve this git command";
         // Sandboxed like bash: flags such as --output or an editor cannot write outside the project or reach the network.
-        const res = await sandboxedExec(["git", "--no-pager", ...args], { cwd: ctx.cwd, profile: ctx.sandboxProfile, ...(ctx.sandboxEnv ? { env: ctx.sandboxEnv } : {}) });
+        const res = await sandboxedExec(["git", "--no-pager", ...args], { cwd: ctx.cwd, profile: ctx.sandboxProfile, ...(ctx.sandboxEnv ? { env: ctx.sandboxEnv } : {}), ...(ctx.signal ? { signal: ctx.signal } : {}) });
         return `${res.output}\n(exit ${res.code})`;
       }
       case "hub_send":

package/src/models/relay.ts

@@ -32,6 +32,8 @@ export interface ModelRelayStatus {
 
 export interface ModelRelayOptions {
   omni: OmniRoute;
+  /** Authoritative admission immediately before each upstream request, including fallbacks. */
+  admitRequest?: () => Promise<{ allowed: boolean; reason?: string; remainingMs?: number }>;
   host?: string;
   port?: number;
   token?: string;
@@ -60,6 +62,8 @@ interface ActiveRequest {
   cleanup: () => void;
 }
 
+class ExecutionAdmissionError extends Error {}
+
 const safeHeader = (value: string | null): string | undefined => value && value.length < 256 ? value : undefined;
 
 function assertLoopback(host: string): void {
@@ -111,8 +115,9 @@ function sseResponse(response: Response, release: () => void, onModel?: (model:
     for (const line of lines) {
       if (!line.startsWith("data:") || line.slice(5).trim() === "[DONE]") continue;
       try {
-        const value = JSON.parse(line.slice(5).trim()) as { model?: unknown };
-        if (typeof value.model === "string" && value.model.length < 256) {
+        const value = JSON.parse(line.slice(5).trim()) as { model?: unknown; choices?: unknown[] };
+        // Gateway heartbeat events can name a synthetic "keepalive" model with no choices.
+        if (Array.isArray(value.choices) && value.choices.length && typeof value.model === "string" && value.model.length < 256) {
           inspectedModel = true;
           onModel(value.model);
           return;
@@ -220,7 +225,10 @@ export async function startModelRelay(options: ModelRelayOptions): Promise<Model
       const boundedRequest = isOllama
         ? { ...request, max_tokens: request.max_tokens ?? options.mlx!.maxTokens ?? 2048, reasoning_effort: request.reasoning_effort ?? "none" }
         : request;
-      response = await fetch(`${base.replace(/\/$/, "")}/chat/completions`, { method: "POST", ...(isOllama ? { redirect: "error" as const } : {}), headers, body: JSON.stringify(bodyForUpstream(boundedRequest, model)), signal: AbortSignal.any([signal, AbortSignal.timeout(180_000)]) });
+      const decision = await options.admitRequest?.();
+      if (decision && !decision.allowed) throw new ExecutionAdmissionError(decision.reason ?? "execution budget exhausted");
+      const deadline = decision?.remainingMs === undefined ? 180_000 : Math.max(1, Math.min(180_000, decision.remainingMs));
+      response = await fetch(`${base.replace(/\/$/, "")}/chat/completions`, { method: "POST", ...(isOllama ? { redirect: "error" as const } : {}), headers, body: JSON.stringify(bodyForUpstream(boundedRequest, model)), signal: AbortSignal.any([signal, AbortSignal.timeout(deadline)]) });
     } catch (error) {
       releaseOnce();
       setState(backend, { state: "error", lastError: error instanceof Error ? error.message.slice(0, 160) : "upstream request failed" });
@@ -307,7 +315,7 @@ export async function startModelRelay(options: ModelRelayOptions): Promise<Model
         record.release = release;
         return sseResponse(result.response, release, result.onModel, (cancel) => { record.cancel = cancel; });
       } catch (error) {
-        if (!fallback || controller.signal.aborted) {
+        if (!fallback || controller.signal.aborted || error instanceof ExecutionAdmissionError) {
           record.cleanup();
           activeRequests.delete(record);
           return Response.json({ error: error instanceof Error ? error.message : "backend unavailable" }, { status: 502 });

package/src/omniroute/client.ts

@@ -1,4 +1,5 @@
 import { readFileSync } from "node:fs";
+import { normalizeUsage, safeModelLabel, type NormalizedUsage } from "./usage.ts";
 
 export interface OmniRouteConfig {
   /** Candidate base URLs ending in /v1, in order of preference. `AGENTHUB_OMNIROUTE_URL` replaces the list. */
@@ -26,6 +27,10 @@ export interface ChatMessage {
 }
 export interface ChatResult {
   message: ChatMessage;
+  /** Usage counters returned by the provider, when valid. No missing counter is inferred as zero. */
+  usage?: NormalizedUsage;
+  /** Model named by the completed response, when the gateway reports it. */
+  servedModel?: string;
   /** `x-omniroute-provider` (direct path) */
   provider?: string;
   /** `x-model-router-selected-model` (through Switchyard) */
@@ -171,16 +176,22 @@ export class OmniRoute {
       throw e;
     }
     if (!res.ok) throw new Error(`model call failed: HTTP ${res.status} ${(await res.text()).slice(0, 300)}`);
-    const json = (await res.json()) as { choices?: { message?: ChatMessage }[] };
+    const json = (await res.json()) as { choices?: { message?: ChatMessage }[]; usage?: unknown; model?: unknown };
     const message = json.choices?.[0]?.message;
     if (!message) throw new Error("model call returned no message");
     const provider = res.headers.get("x-omniroute-provider") ?? undefined;
     const selectedModel = res.headers.get("x-model-router-selected-model") ?? undefined;
+    const servedModel = safeModelLabel(json.model) ?? safeModelLabel(selectedModel);
+    const safeProvider = safeModelLabel(provider);
+    const safeSelectedModel = safeModelLabel(selectedModel);
+    const usage = normalizeUsage(json.usage);
     // reasoning_content and other extras are dropped here: only role, content and tool_calls travel on.
     return {
       message: { role: "assistant", content: message.content ?? null, ...(message.tool_calls?.length ? { tool_calls: message.tool_calls } : {}) },
-      ...(provider ? { provider } : {}),
-      ...(selectedModel ? { selectedModel } : {}),
+      ...(usage ? { usage } : {}),
+      ...(safeProvider ? { provider: safeProvider } : {}),
+      ...(safeSelectedModel ? { selectedModel: safeSelectedModel } : {}),
+      ...(servedModel ? { servedModel } : {}),
     };
   }
 }

package/src/omniroute/usage.ts

@@ -0,0 +1,38 @@
+/** Provider usage is optional. Missing counters stay missing; they are never guessed as zero. */
+export interface NormalizedUsage {
+  inputTokens?: number;
+  outputTokens?: number;
+  cacheReadTokens?: number;
+  cacheWriteTokens?: number;
+  totalTokens?: number;
+}
+
+const count = (value: unknown): number | undefined =>
+  typeof value === "number" && Number.isSafeInteger(value) && value >= 0 ? value : undefined;
+
+/** Accept known OpenAI-compatible and Anthropic-compatible counter spellings, dropping malformed fields. */
+export function normalizeUsage(value: unknown): NormalizedUsage | undefined {
+  if (!value || typeof value !== "object" || Array.isArray(value)) return undefined;
+  const raw = value as Record<string, unknown>;
+  const inputTokens = count(raw.input_tokens) ?? count(raw.prompt_tokens) ?? count(raw.inputTokens);
+  const outputTokens = count(raw.output_tokens) ?? count(raw.completion_tokens) ?? count(raw.outputTokens);
+  const cacheReadTokens = count(raw.cache_read_input_tokens) ?? count(raw.cached_read_tokens) ?? count(raw.cacheReadTokens)
+    ?? count((raw.prompt_tokens_details as Record<string, unknown> | undefined)?.cached_tokens);
+  const cacheWriteTokens = count(raw.cache_creation_input_tokens) ?? count(raw.cache_write_tokens) ?? count(raw.cacheWriteTokens);
+  const totalTokens = count(raw.total_tokens) ?? count(raw.totalTokens);
+  const result = {
+    ...(inputTokens !== undefined ? { inputTokens } : {}),
+    ...(outputTokens !== undefined ? { outputTokens } : {}),
+    ...(cacheReadTokens !== undefined ? { cacheReadTokens } : {}),
+    ...(cacheWriteTokens !== undefined ? { cacheWriteTokens } : {}),
+    ...(totalTokens !== undefined ? { totalTokens } : {}),
+  };
+  return Object.keys(result).length ? result : undefined;
+}
+
+/** Model/provider labels are metadata, but still untrusted: keep a short printable identifier only. */
+export function safeModelLabel(value: unknown): string | undefined {
+  if (typeof value !== "string") return undefined;
+  const label = value.trim();
+  return label.length > 0 && label.length <= 160 && /^[A-Za-z0-9][A-Za-z0-9._:/+@-]*$/.test(label) ? label : undefined;
+}

package/src/pi/extension.ts

@@ -2,6 +2,7 @@ import { assistantTokens } from "./usage.ts";
 type ExtensionAPI = any;
 import { execFileSync } from "node:child_process";
 import { createHash } from "node:crypto";
+type BudgetUnit = "model_calls" | "tool_calls";
 
 type BridgeEvent = { type: string; text?: string; sessionId?: string; sessionFile?: string; error?: string };
 
@@ -14,10 +15,12 @@ const models = (() => {
   try { return JSON.parse(process.env.AGENTHUB_PI_MODELS ?? "[]") as unknown[]; } catch { return []; }
 })();
 let pollStarted = false;
+let pollStopped = false;
 let toolSteps = 0;
 let lastActivity = 0;
 let usageSeq = 0;
 let forcedFailure = "";
+let turnGeneration = 0;
 const maxSteps = Number(process.env.AGENTHUB_PI_MAX_STEPS ?? 30);
 let shutdown: (() => void) | undefined;
 let runtimeCtx: any;
@@ -29,10 +32,27 @@ async function post(path: string, body: unknown): Promise<any> {
   if (!response.ok) throw new Error(`Pi bridge HTTP ${response.status}`);
   return response.json();
 }
+async function admitBudget(unit: BudgetUnit, idleUserBash = false): Promise<{ allowed: boolean; reservation?: string }> {
+  try {
+    const result = await post("/budget", { unit, generation: turnGeneration, ...(idleUserBash ? { idleUserBash: true } : {}) });
+    const decisions = Array.isArray(result?.decisions) ? result.decisions : [];
+    const denied = decisions.find((item: any) => item?.allowed === false);
+    if (!denied) return { allowed: true, ...(idleUserBash && typeof result?.reservation === "string" ? { reservation: result.reservation } : {}) };
+    const reason = `execution budget ${denied.reason ?? "exhausted"}: ${denied.scope} ${denied.unit ?? unit} used ${denied.used}${denied.limit === null ? "" : ` of ${denied.limit}`}`;
+    if (idleUserBash) return { allowed: false };
+    forcedFailure = reason;
+  } catch (error) {
+    if (idleUserBash) return { allowed: false };
+    forcedFailure = `execution budget admission unavailable: ${(error as Error).message}`;
+  }
+  try { await post("/event", { type: "agent_end", generation: turnGeneration, failed: true, error: forcedFailure }); } catch { /* abort remains authoritative when the bridge is unavailable */ }
+  runtimeCtx?.abort?.();
+  return { allowed: false };
+}
 function processSignature(): string | undefined { try { const text = execFileSync("ps", ["-p", String(process.pid), "-o", "lstart=,comm="], { encoding: "utf8" }).trim(); return text ? createHash("sha256").update(text).digest("hex") : undefined; } catch { return undefined; } }
 
 async function poll(pi: ExtensionAPI): Promise<void> {
-  while (true) {
+  while (!pollStopped) {
     try {
       const response = await fetch(`${bridgeUrl}/commands`, { headers: { authorization: `Bearer ${bridgeToken}` } });
       const payload = await response.json() as { command?: any };
@@ -54,7 +74,15 @@ async function poll(pi: ExtensionAPI): Promise<void> {
           const model = modelRegistry?.find(String(command.provider), String(command.modelId));
           if (!model || !(await (pi as any).setModel?.(model))) throw new Error("Pi model is not available");
         } else if (command.type === "shutdown") {
+          pollStopped = true;
           shutdown?.();
+        } else if (command.type === "abort_budget") {
+          if (Number.isSafeInteger(command.generation) && command.generation === turnGeneration) {
+            const reason = "execution budget exhausted: elapsed_ms wall cap reached";
+            forcedFailure = reason;
+            await post("/event", { type: "agent_end", generation: turnGeneration, failed: true, error: reason });
+            runtimeCtx?.abort?.();
+          }
         }
         await post("/ack", { id: command.id, ok: true });
       } catch (error) { await post("/ack", { id: command.id, ok: false, error: (error as Error).message }); }
@@ -89,7 +117,7 @@ export default function(pi: ExtensionAPI): void {
     const text = message?.content?.filter((c: any) => c?.type === "text").map((c: any) => c.text).join("")?.trim();
     const failed = !!forcedFailure || message?.stopReason === "error";
     const cancelled = !forcedFailure && message?.stopReason === "aborted";
-    await post("/event", { type: "agent_end", text: text ?? "", failed, cancelled, ...(failed ? { error: forcedFailure || message?.errorMessage || message?.stopReason } : {}) });
+    await post("/event", { type: "agent_end", generation: turnGeneration, text: text ?? "", failed, cancelled, ...(failed ? { error: forcedFailure || message?.errorMessage || message?.stopReason } : {}) });
   });
   pi.on("message_end", async (event: any) => {
     const tokens = assistantTokens(event.message);
@@ -104,21 +132,28 @@ export default function(pi: ExtensionAPI): void {
     lastActivity = now;
     try { await post("/event", { type: "activity" }); } catch { /* shutdown owns bridge cleanup */ }
   });
-  pi.on("agent_start", async () => { toolSteps = 0; forcedFailure = ""; await post("/event", { type: "agent_start" }); });
+  pi.on("agent_start", async () => { toolSteps = 0; forcedFailure = ""; turnGeneration++; await post("/event", { type: "agent_start", generation: turnGeneration }); });
   pi.on("model_select", async (_event: any, ctx: any) => { if (ctx.model?.provider && ctx.model.provider !== "agent-hub-local") ctx.shutdown?.(); });
   pi.on("before_agent_start", async (event: any, ctx: any) => {
     if (ctx.model?.provider && ctx.model.provider !== "agent-hub-local") { ctx.abort?.(); return { systemPrompt: event.systemPrompt }; }
     return undefined;
   });
-  pi.on("before_provider_request", async (_event: any, ctx: any) => { if (ctx.model?.provider && ctx.model.provider !== "agent-hub-local") ctx.abort?.(); return undefined; });
+  pi.on("before_provider_request", async (_event: any, ctx: any) => {
+    if (ctx.model?.provider && ctx.model.provider !== "agent-hub-local") { ctx.abort?.(); return undefined; }
+    try { if (!(await admitBudget("model_calls")).allowed) ctx.abort?.(); }
+    catch (error) { forcedFailure = `execution budget admission failed: ${(error as Error).message}`; ctx.abort?.(); }
+    return undefined;
+  });
   pi.on("agent_settled", async (_event: any, ctx: any) => {
     const entries = ctx.sessionManager?.getEntries?.() ?? [];
     const message = entries.slice().reverse().find((entry: any) => entry.type === "message" && entry.message?.role === "assistant")?.message;
     const text = message?.content?.filter((c: any) => c?.type === "text").map((c: any) => c.text).join("")?.trim();
-    await post("/event", { type: "agent_settled", ...(text ? { text } : {}) });
+    await post("/event", { type: "agent_settled", generation: turnGeneration, ...(text ? { text } : {}) });
   });
   pi.on("user_bash", async (event: any) => {
-    const result = await post("/tool", { name: "bash", args: { command: event.command, cwd: event.cwd }, toolCallId: `pi-shell-${Date.now()}` });
+    const admission = await admitBudget("tool_calls", true);
+    if (!admission.allowed || !admission.reservation) return { cancel: true };
+    const result = await post("/tool", { name: "bash", args: { command: event.command, cwd: event.cwd }, toolCallId: `pi-shell-${Date.now()}`, purpose: "idle_user_bash", generation: turnGeneration, reservation: admission.reservation });
     return { result: { output: String(result.text ?? result), exitCode: 0, cancelled: false, truncated: false } };
   });
   // Managed sessions may only be handed over by PiPeer after it has fenced the
@@ -126,11 +161,12 @@ export default function(pi: ExtensionAPI): void {
   // otherwise silently detach the hub from its recorded session.
   pi.on("session_before_switch", async () => ({ cancel: true }));
   pi.on("session_before_fork", async () => ({ cancel: true }));
-  pi.on("session_shutdown", async () => { await post("/event", { type: "session_shutdown" }); });
+  pi.on("session_shutdown", async () => { pollStopped = true; await post("/event", { type: "session_shutdown" }); });
   for (const raw of (() => { try { return JSON.parse(process.env.AGENTHUB_PI_TOOLS ?? "[]") as any[]; } catch { return []; } })()) {
     if (!raw || typeof raw.name !== "string" || !raw.parameters) continue;
     pi.registerTool({ name: raw.name, label: raw.name, description: raw.description ?? raw.name, parameters: raw.parameters, async execute(toolCallId: string, params: unknown) {
-      if (toolSteps++ >= maxSteps) { const reason = `Pi tool step limit ${maxSteps} reached`; forcedFailure = reason; await post("/event", { type: "agent_end", failed: true, error: reason }); runtimeCtx?.abort?.(); return { content: [{ type: "text", text: `error: ${reason}` }], details: {}, isError: true }; }
+      if (!(await admitBudget("tool_calls")).allowed) return { content: [{ type: "text", text: `error: ${forcedFailure}` }], details: {}, isError: true };
+      if (toolSteps++ >= maxSteps) { const reason = `Pi tool step limit ${maxSteps} reached`; forcedFailure = reason; await post("/event", { type: "agent_end", generation: turnGeneration, failed: true, error: reason }); runtimeCtx?.abort?.(); return { content: [{ type: "text", text: `error: ${reason}` }], details: {}, isError: true }; }
       const result = await post("/tool", { name: raw.name, args: params, toolCallId });
       return { content: [{ type: "text", text: String(result.text ?? result) }], details: {} };
     } });