@oh-my-pi/pi-agent-core 16.0.5 → 16.0.6

package/CHANGELOG.md

@@ -2,6 +2,12 @@
 
 ## [Unreleased]
 
+## [16.0.6] - 2026-06-18
+
+### Added
+
+- Added `transformAssistantMessage` hook to `AgentOptions` and `Agent` to allow mutating the finalized assistant message before UI emission, context appending, or tool dispatch
+
 ## [16.0.5] - 2026-06-17
 
 ### Breaking Changes

package/dist/types/agent.d.ts

@@ -158,6 +158,12 @@ export interface AgentOptions {
      * message are emitted. See {@link AgentLoopConfig.afterToolCall} for full semantics.
      */
     afterToolCall?: AgentLoopConfig["afterToolCall"];
+    /**
+     * Called once an assistant message is finalized, before it reaches the
+     * context, the UI, or tool dispatch. May mutate the message in place (text +
+     * tool-call arguments). See {@link AgentLoopConfig.transformAssistantMessage}.
+     */
+    transformAssistantMessage?: AgentLoopConfig["transformAssistantMessage"];
     /**
      * Opt-in OpenTelemetry instrumentation. Passing `{}` enables the loop's
      * GenAI-semantic-convention spans using the global tracer provider. See
@@ -187,6 +193,11 @@ export declare class Agent {
      * message emission. Reassign at any time to swap the implementation.
      */
     afterToolCall?: AgentLoopConfig["afterToolCall"];
+    /**
+     * Hook invoked once an assistant message is finalized, before context append,
+     * UI emission, and tool dispatch. Reassign at any time to swap the implementation.
+     */
+    transformAssistantMessage?: AgentLoopConfig["transformAssistantMessage"];
     constructor(opts?: AgentOptions);
     /**
      * Get the current session ID used for provider caching.

package/dist/types/index.d.ts

@@ -6,5 +6,6 @@ export * from "./proxy";
 export * from "./run-collector";
 export * from "./telemetry";
 export * from "./thinking";
+export * from "./tokenizer";
 export * from "./types";
 export * from "./utils/yield";

package/dist/types/tokenizer.d.ts

@@ -0,0 +1 @@
+export declare function countTokens(text: string | string[]): number;

package/dist/types/types.d.ts

@@ -245,6 +245,19 @@ export interface AgentLoopConfig extends SimpleStreamOptions {
      * backlog so advice produced during the wait is injected as an aside).
      */
     onTurnEnd?: (messages: AgentMessage[], signal?: AbortSignal) => Promise<void> | void;
+    /**
+     * Called once an assistant message is finalized from the model stream, before
+     * it is appended to the context, emitted as `message_end`, or its tool calls
+     * are validated and dispatched. The hook may mutate the message in place —
+     * both its text content and its tool-call arguments — and those edits are seen
+     * by the transcript, the UI, and tool execution alike (single source of truth).
+     *
+     * Used for inline macro expansion: rewriting `@[[runtime.name(args)]]` tokens
+     * to host-computed values before anything downstream consumes the message.
+     * Runs at most once per assistant message; must not throw (a throw would abort
+     * the turn).
+     */
+    transformAssistantMessage?: (message: AssistantMessage, signal?: AbortSignal) => Promise<void> | void;
     /**
      * Called after a tool finishes executing, before `tool_execution_end` and the
      * tool-result message are emitted.

package/package.json

@@ -1,7 +1,7 @@
 {
 	"type": "module",
 	"name": "@oh-my-pi/pi-agent-core",
-	"version": "16.0.5",
+	"version": "16.0.6",
 	"description": "General-purpose agent with transport abstraction, state management, and attachment support",
 	"homepage": "https://omp.sh",
 	"author": "Can Boluk",
@@ -35,11 +35,11 @@
 		"fmt": "biome format --write ."
 	},
 	"dependencies": {
-		"@oh-my-pi/pi-ai": "16.0.5",
-		"@oh-my-pi/pi-catalog": "16.0.5",
-		"@oh-my-pi/pi-natives": "16.0.5",
-		"@oh-my-pi/pi-utils": "16.0.5",
-		"@oh-my-pi/snapcompact": "16.0.5",
+		"@oh-my-pi/pi-ai": "16.0.6",
+		"@oh-my-pi/pi-catalog": "16.0.6",
+		"@oh-my-pi/pi-natives": "16.0.6",
+		"@oh-my-pi/pi-utils": "16.0.6",
+		"@oh-my-pi/snapcompact": "16.0.6",
 		"@opentelemetry/api": "^1.9.1"
 	},
 	"devDependencies": {

package/src/agent-loop.ts

@@ -5,9 +5,11 @@
 import {
 	type AssistantMessage,
 	type AssistantMessageEvent,
+	arkToWireSchema,
 	type Context,
 	EventStream,
 	isApiKeyResolver,
+	isArkSchema,
 	isZodSchema,
 	resolveApiKeyOnce,
 	seedApiKeyResolver,
@@ -566,6 +568,9 @@ export function normalizeTools(
 			if (isZodSchema(parameters)) {
 				const wired = zodToWireSchema(parameters);
 				parameters = injectIntentIntoSchema(wired, intentMode) as TSchema;
+			} else if (isArkSchema(parameters)) {
+				const wired = arkToWireSchema(parameters);
+				parameters = injectIntentIntoSchema(wired, intentMode) as TSchema;
 			} else {
 				parameters = injectIntentIntoSchema(parameters, intentMode) as TSchema;
 			}
@@ -1226,6 +1231,12 @@ async function streamAssistantResponse(
 							}
 						}
 						finalMessage = snapshotAssistantMessage(finalMessage);
+						// Expand inline macros (and any other registered rewrite) on the
+						// finalized message before it reaches the context, the UI, or tool
+						// dispatch — so a single mutation is the source of truth for all three.
+						if (config.transformAssistantMessage) {
+							await config.transformAssistantMessage(finalMessage, requestSignal);
+						}
 						if (addedPartial) {
 							context.messages[context.messages.length - 1] = finalMessage;
 						} else {
@@ -1249,9 +1260,19 @@ async function streamAssistantResponse(
 					switch (event.type) {
 						case "start":
 							partialMessage = event.partial;
-							context.messages.push(partialMessage);
-							addedPartial = true;
-							stream.push({ type: "message_start", message: snapshotAssistantMessage(partialMessage) });
+							if (addedPartial) {
+								context.messages[context.messages.length - 1] = partialMessage;
+								completedToolCallIds.clear();
+								stream.push({
+									type: "message_update",
+									assistantMessageEvent: snapshotAssistantMessageEvent(event),
+									message: snapshotAssistantMessage(partialMessage),
+								});
+							} else {
+								context.messages.push(partialMessage);
+								addedPartial = true;
+								stream.push({ type: "message_start", message: snapshotAssistantMessage(partialMessage) });
+							}
 							break;
 
 						case "text_start":

package/src/agent.ts

@@ -257,6 +257,13 @@ export interface AgentOptions {
 	 */
 	afterToolCall?: AgentLoopConfig["afterToolCall"];
 
+	/**
+	 * Called once an assistant message is finalized, before it reaches the
+	 * context, the UI, or tool dispatch. May mutate the message in place (text +
+	 * tool-call arguments). See {@link AgentLoopConfig.transformAssistantMessage}.
+	 */
+	transformAssistantMessage?: AgentLoopConfig["transformAssistantMessage"];
+
 	/**
 	 * Opt-in OpenTelemetry instrumentation. Passing `{}` enables the loop's
 	 * GenAI-semantic-convention spans using the global tracer provider. See
@@ -358,6 +365,11 @@ export class Agent {
 	 * message emission. Reassign at any time to swap the implementation.
 	 */
 	afterToolCall?: AgentLoopConfig["afterToolCall"];
+	/**
+	 * Hook invoked once an assistant message is finalized, before context append,
+	 * UI emission, and tool dispatch. Reassign at any time to swap the implementation.
+	 */
+	transformAssistantMessage?: AgentLoopConfig["transformAssistantMessage"];
 
 	constructor(opts: AgentOptions = {}) {
 		this.#state = { ...this.#state, ...opts.initialState };
@@ -402,6 +414,7 @@ export class Agent {
 		this.#onHarmonyLeak = opts.onHarmonyLeak;
 		this.beforeToolCall = opts.beforeToolCall;
 		this.afterToolCall = opts.afterToolCall;
+		this.transformAssistantMessage = opts.transformAssistantMessage;
 		this.#telemetry = opts.telemetry;
 		this.#appendOnlyContext = opts.appendOnlyContext;
 		this.#transformProviderContext = opts.transformProviderContext;
@@ -1051,6 +1064,9 @@ export class Agent {
 			appendOnlyContext: this.#appendOnlyContext,
 			beforeToolCall: this.beforeToolCall ? (ctx, signal) => this.beforeToolCall?.(ctx, signal) : undefined,
 			afterToolCall: this.afterToolCall ? (ctx, signal) => this.afterToolCall?.(ctx, signal) : undefined,
+			transformAssistantMessage: this.transformAssistantMessage
+				? (message, signal) => this.transformAssistantMessage?.(message, signal)
+				: undefined,
 			onAssistantMessageEvent: this.#onAssistantMessageEvent,
 			onHarmonyLeak: this.#onHarmonyLeak,
 			onTurnEnd: (messages, signal) => this.#onTurnEnd?.(messages, signal),

package/src/compaction/compaction.ts

@@ -20,11 +20,11 @@ import {
 } from "@oh-my-pi/pi-ai";
 import { preferredDialect } from "@oh-my-pi/pi-catalog/identity";
 import { clampThinkingLevelForModel } from "@oh-my-pi/pi-catalog/model-thinking";
-import { countTokens } from "@oh-my-pi/pi-natives";
 import { logger, prompt } from "@oh-my-pi/pi-utils";
 import * as snapcompact from "@oh-my-pi/snapcompact";
 import { type AgentTelemetry, instrumentedCompleteSimple } from "../telemetry";
 import { ThinkingLevel } from "../thinking";
+import { countTokens } from "../tokenizer";
 import type { AgentMessage } from "../types";
 import type { CompactionEntry, SessionEntry } from "./entries";
 import { type ConvertToLlm, createBranchSummaryMessage, createCustomMessage, defaultConvertToLlm } from "./messages";

package/src/compaction/openai.ts

@@ -13,7 +13,7 @@
  */
 
 import { ProviderHttpError } from "@oh-my-pi/pi-ai/errors";
-import { parseTextSignature } from "@oh-my-pi/pi-ai/providers/openai-responses-shared";
+import { parseTextSignature } from "@oh-my-pi/pi-ai/providers/openai-shared";
 import { transformMessages } from "@oh-my-pi/pi-ai/providers/transform-messages";
 import type { AssistantMessage, FetchImpl, Message, Model } from "@oh-my-pi/pi-ai/types";
 import {

package/src/compaction/shake.ts

@@ -11,7 +11,7 @@
  */
 
 import type { TextContent, ToolResultMessage } from "@oh-my-pi/pi-ai";
-import { countTokens } from "@oh-my-pi/pi-natives";
+import { countTokens } from "../tokenizer";
 import type { AgentMessage } from "../types";
 import { estimateTokens } from "./compaction";
 import type { CustomMessageEntry, SessionEntry, SessionMessageEntry } from "./entries";

package/src/index.ts

@@ -14,6 +14,8 @@ export * from "./run-collector";
 export * from "./telemetry";
 // Thinking selectors
 export * from "./thinking";
+// Tokenizer choice
+export * from "./tokenizer";
 // Types
 export * from "./types";
 // Yield utilities for Bun event-loop busy-wait prevention

package/src/proxy.ts

@@ -210,6 +210,18 @@ function processProxyEvent(
 ): AssistantMessageEvent | undefined {
 	switch (proxyEvent.type) {
 		case "start":
+			partial.content.length = 0;
+			partial.usage = {
+				input: 0,
+				output: 0,
+				cacheRead: 0,
+				cacheWrite: 0,
+				totalTokens: 0,
+				cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+			};
+			delete (partial as { stopReason?: string }).stopReason;
+			delete (partial as { errorMessage?: string }).errorMessage;
+			delete (partial as { duration?: number }).duration;
 			return { type: "start", partial };
 
 		case "text_start":

package/src/tokenizer.ts

@@ -0,0 +1,17 @@
+import { countTokens as countTokensNat } from "@oh-my-pi/pi-natives";
+
+const accurate = process.env.PI_TOKENIZER_ACCURATE === "1" && Bun.env.NODE_ENV !== "test";
+
+function estimateTokens(text: string) {
+	return (Buffer.byteLength(text, "utf-8") + 3) >> 2;
+}
+
+export function countTokens(text: string | string[]): number {
+	if (accurate) {
+		return countTokensNat(text);
+	} else if (Array.isArray(text)) {
+		return text.reduce((sum, t) => sum + estimateTokens(t), 0);
+	} else {
+		return estimateTokens(text);
+	}
+}

package/src/types.ts

@@ -293,6 +293,20 @@ export interface AgentLoopConfig extends SimpleStreamOptions {
 	 */
 	onTurnEnd?: (messages: AgentMessage[], signal?: AbortSignal) => Promise<void> | void;
 
+	/**
+	 * Called once an assistant message is finalized from the model stream, before
+	 * it is appended to the context, emitted as `message_end`, or its tool calls
+	 * are validated and dispatched. The hook may mutate the message in place —
+	 * both its text content and its tool-call arguments — and those edits are seen
+	 * by the transcript, the UI, and tool execution alike (single source of truth).
+	 *
+	 * Used for inline macro expansion: rewriting `@[[runtime.name(args)]]` tokens
+	 * to host-computed values before anything downstream consumes the message.
+	 * Runs at most once per assistant message; must not throw (a throw would abort
+	 * the turn).
+	 */
+	transformAssistantMessage?: (message: AssistantMessage, signal?: AbortSignal) => Promise<void> | void;
+
 	/**
 	 * Called after a tool finishes executing, before `tool_execution_end` and the
 	 * tool-result message are emitted.