Merge remote-tracking branch 'upstream/main' into feat/guardrail_inte…

…rface # Conflicts: # src/agents/pi-embedded-runner/run.ts # src/agents/pi-embedded-runner/run/attempt.ts # src/agents/pi-tool-definition-adapter.ts # src/plugins/hooks.ts # src/plugins/types.ts
Reapor-Yurnero · Feb 1, 2026 · Feb 1, 2026 · Feb 1, 2026 · Feb 1, 2026 · Feb 1, 2026
commit 1141b52d6ce99caa2f8f4c59abf53f69143bc7b7
diff --git a/src/agents/pi-embedded-runner/run.ts b/src/agents/pi-embedded-runner/run.ts
@@ -536,10 +536,12 @@ export async function runEmbeddedPiAgent(
               }
             : undefined;
 
-          mergeUsageIntoAccumulator(
-            usageAccumulator,
-            attempt.attemptUsage ?? normalizeUsage(lastAssistant?.usage as UsageLike),
-          );
+          const lastAssistantUsage = normalizeUsage(lastAssistant?.usage as UsageLike);
+          const attemptUsage = attempt.attemptUsage ?? lastAssistantUsage;
+          mergeUsageIntoAccumulator(usageAccumulator, attemptUsage);
+          // Keep prompt size from the latest model call so session totalTokens
+          // reflects current context usage, not accumulated tool-loop usage.
+          lastRunPromptUsage = lastAssistantUsage ?? attemptUsage;
           autoCompactionCount += Math.max(0, attempt.compactionCount ?? 0);
 
           const formattedAssistantErrorText = lastAssistant

diff --git a/src/agents/pi-embedded-runner/run/attempt.ts b/src/agents/pi-embedded-runner/run/attempt.ts
@@ -150,69 +150,6 @@ export function injectHistoryImagesIntoMessages(
   return didMutate;
 }
 
-function summarizeMessagePayload(msg: AgentMessage): { textChars: number; imageBlocks: number } {
-  const content = (msg as { content?: unknown }).content;
-  if (typeof content === "string") {
-    return { textChars: content.length, imageBlocks: 0 };
-  }
-  if (!Array.isArray(content)) {
-    return { textChars: 0, imageBlocks: 0 };
-  }
-
-  let textChars = 0;
-  let imageBlocks = 0;
-  for (const block of content) {
-    if (!block || typeof block !== "object") {
-      continue;
-    }
-    const typedBlock = block as { type?: unknown; text?: unknown };
-    if (typedBlock.type === "image") {
-      imageBlocks++;
-      continue;
-    }
-    if (typeof typedBlock.text === "string") {
-      textChars += typedBlock.text.length;
-    }
-  }
-
-  return { textChars, imageBlocks };
-}
-
-function summarizeSessionContext(messages: AgentMessage[]): {
-  roleCounts: string;
-  totalTextChars: number;
-  totalImageBlocks: number;
-  maxMessageTextChars: number;
-} {
-  const roleCounts = new Map<string, number>();
-  let totalTextChars = 0;
-  let totalImageBlocks = 0;
-  let maxMessageTextChars = 0;
-
-  for (const msg of messages) {
-    const role = typeof msg.role === "string" ? msg.role : "unknown";
-    roleCounts.set(role, (roleCounts.get(role) ?? 0) + 1);
-
-    const payload = summarizeMessagePayload(msg);
-    totalTextChars += payload.textChars;
-    totalImageBlocks += payload.imageBlocks;
-    if (payload.textChars > maxMessageTextChars) {
-      maxMessageTextChars = payload.textChars;
-    }
-  }
-
-  return {
-    roleCounts:
-      [...roleCounts.entries()]
-        .toSorted((a, b) => a[0].localeCompare(b[0]))
-        .map(([role, count]) => `${role}:${count}`)
-        .join(",") || "none",
-    totalTextChars,
-    totalImageBlocks,
-    maxMessageTextChars,
-  };
-}
-
 export async function runEmbeddedAttempt(
   params: EmbeddedRunAttemptParams,
 ): Promise<EmbeddedRunAttemptResult> {
@@ -542,11 +479,18 @@ export async function runEmbeddedAttempt(
         model: params.model,
       });
 
-      const toolHookOptions = {
-        context: toolHookContext,
-        getMessages: () => sessionManager?.buildSessionContext().messages ?? [],
-        systemPrompt: appendPrompt,
-      };
+      const skipGuardrailHooks =
+        isGuardrailRunId(params.sessionId) || isGuardrailRunId(params.runId);
+      // Get hook runner early so it's available when creating tools.
+      // When invoked by a guardrail itself, skip hooks to prevent recursion.
+      const hookRunner = skipGuardrailHooks ? null : getGlobalHookRunner();
+      const toolHookOptions = skipGuardrailHooks
+        ? undefined
+        : {
+            context: toolHookContext,
+            getMessages: () => sessionManager?.buildSessionContext().messages ?? [],
+            systemPrompt: appendPrompt,
+          };
       const { builtInTools, customTools } = splitSdkTools({
         tools,
         sandboxEnabled: !!sandbox?.enabled,
@@ -742,7 +686,7 @@ export async function runEmbeddedAttempt(
       const subscription = subscribeEmbeddedPiSession({
         session: activeSession,
         runId: params.runId,
-        hookRunner: getGlobalHookRunner() ?? undefined,
+        hookRunner: hookRunner ?? undefined,
         verboseLevel: params.verboseLevel,
         reasoningMode: params.reasoningLevel ?? "off",
         toolResultFormat: params.toolResultFormat,
@@ -831,10 +775,7 @@ export async function runEmbeddedAttempt(
         }
       }
 
-      // Get hook runner once for both before_agent_start and agent_end hooks
-      const skipGuardrailHooks =
-        isGuardrailRunId(params.sessionId) || isGuardrailRunId(params.runId);
-      const hookRunner = skipGuardrailHooks ? null : getGlobalHookRunner();
+      // Hook runner was already obtained earlier before tool creation
       const hookAgentId =
         typeof params.agentId === "string" && params.agentId.trim()
           ? normalizeAgentId(params.agentId)
@@ -904,8 +845,9 @@ export async function runEmbeddedAttempt(
             },
             {
               agentId: sessionAgentId,
+              sessionId: params.sessionId,
               sessionKey: params.sessionKey,
-              workspaceDir: params.workspaceDir,
+              workspaceDir: effectiveWorkspace,
               messageProvider: params.messageProvider ?? undefined,
             },
           );
@@ -978,7 +920,10 @@ export async function runEmbeddedAttempt(
               historyMessages: activeSession.messages,
               maxBytes: MAX_IMAGE_BYTES,
               // Enforce sandbox path restrictions when sandbox is enabled
-              sandboxRoot: sandbox?.enabled ? sandbox.workspaceDir : undefined,
+              sandbox:
+                sandbox?.enabled && sandbox?.fsBridge
+                  ? { root: sandbox.workspaceDir, bridge: sandbox.fsBridge }
+                  : undefined,
             });
 
             // Inject history images into their original message positions.

diff --git a/src/agents/pi-embedded-subscribe.handlers.tools.ts b/src/agents/pi-embedded-subscribe.handlers.tools.ts
@@ -1,4 +1,4 @@
-import type { AgentEvent } from "@mariozechner/pi-agent-core";
+import type { AgentEvent, AgentToolResult } from "@mariozechner/pi-agent-core";
 import type { PluginHookAfterToolCallEvent } from "../plugins/types.js";
 import type { EmbeddedPiSubscribeContext } from "./pi-embedded-subscribe.handlers.types.js";
 import { emitAgentEvent } from "../infra/agent-events.js";
@@ -13,6 +13,7 @@ import {
   sanitizeToolResult,
 } from "./pi-embedded-subscribe.tools.js";
 import { inferToolMetaFromArgs } from "./pi-embedded-utils.js";
+import { consumeAfterToolCallHookHandled } from "./pi-tools.before-tool-call.js";
 import { normalizeToolName } from "./tool-policy.js";
 
 /** Track tool execution start times and args for after_tool_call hook */
@@ -237,22 +238,29 @@ export async function handleToolExecutionEnd(
   // Run after_tool_call plugin hook (fire-and-forget)
   const hookRunnerAfter = ctx.hookRunner ?? getGlobalHookRunner();
   if (hookRunnerAfter?.hasHooks("after_tool_call")) {
+    if (consumeAfterToolCallHookHandled(toolCallId)) {
+      toolStartData.delete(toolCallId);
+      return;
+    }
+
     const startData = toolStartData.get(toolCallId);
     toolStartData.delete(toolCallId);
     const durationMs = startData?.startTime != null ? Date.now() - startData.startTime : undefined;
     const toolArgs = startData?.args;
     const hookEvent: PluginHookAfterToolCallEvent = {
       toolName,
+      toolCallId,
       params: (toolArgs && typeof toolArgs === "object" ? toolArgs : {}) as Record<string, unknown>,
-      result: sanitizedResult,
+      result: sanitizedResult as AgentToolResult<unknown>,
       error: isToolError ? extractToolErrorMessage(sanitizedResult) : undefined,
       durationMs,
+      messages: ctx.params.session?.messages ?? [],
     };
     void hookRunnerAfter
       .runAfterToolCall(hookEvent, {
         toolName,
         agentId: undefined,
-        sessionKey: undefined,
+        sessionKey: ctx.params.sessionKey,
       })
       .catch((err) => {
         ctx.log.warn(`after_tool_call hook failed: tool=${toolName} error=${String(err)}`);