fix: apply text transforms to tool calls

ZengWen-DT · ZengWen-DT · commit 051e827c6977 · 2026-06-29T12:55:03.000Z
diff --git a/packages/agent-core/src/agent-loop.ts b/packages/agent-core/src/agent-loop.ts
@@ -88,6 +88,17 @@ function removeNonExecutableToolCalls(message: AssistantMessage): AssistantMessa
   return content.length === message.content.length ? message : { ...message, content };
 }
 
+async function finalizeAssistantMessageForLoop(
+  message: AssistantMessage,
+  config: AgentLoopConfig,
+  signal: AbortSignal | undefined,
+): Promise<AssistantMessage> {
+  const executableMessage = removeNonExecutableToolCalls(message);
+  return config.transformAssistantMessage
+    ? await config.transformAssistantMessage(executableMessage, signal)
+    : executableMessage;
+}
+
 /**
  * Start an agent loop with a new prompt message.
  * The prompt is added to the context and events are emitted for it.
@@ -516,7 +527,11 @@ async function streamAssistantResponse(
 
       case "done":
       case "error": {
-        const finalMessage = removeNonExecutableToolCalls(await response.result());
+        const finalMessage = await finalizeAssistantMessageForLoop(
+          await response.result(),
+          config,
+          signal,
+        );
         if (addedPartial) {
           context.messages[context.messages.length - 1] = finalMessage;
         } else {
@@ -531,7 +546,11 @@ async function streamAssistantResponse(
     }
   }
 
-  const finalMessage = removeNonExecutableToolCalls(await response.result());
+  const finalMessage = await finalizeAssistantMessageForLoop(
+    await response.result(),
+    config,
+    signal,
+  );
   if (addedPartial) {
     context.messages[context.messages.length - 1] = finalMessage;
   } else {
diff --git a/packages/agent-core/src/agent.ts b/packages/agent-core/src/agent.ts
@@ -110,6 +110,8 @@ export interface AgentOptions {
   convertToLlm?: (messages: AgentMessage[]) => Message[] | Promise<Message[]>;
   /** Optionally rewrite context before each provider request. */
   transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => Promise<AgentMessage[]>;
+  /** Optionally rewrite finalized assistant messages before emitting/executing. */
+  transformAssistantMessage?: AgentLoopConfig["transformAssistantMessage"];
   /** Injected stream runtime used when streamFn is not supplied. */
   runtime?: AgentCoreStreamRuntimeDeps;
   /** Explicit stream implementation, preferred over runtime.streamSimple. */
@@ -214,6 +216,7 @@ export class Agent {
     messages: AgentMessage[],
     signal?: AbortSignal,
   ) => Promise<AgentMessage[]>;
+  public transformAssistantMessage?: AgentLoopConfig["transformAssistantMessage"];
   public runtime?: AgentCoreStreamRuntimeDeps;
   public streamFn: StreamFn;
   public getApiKey?: (provider: string) => Promise<string | undefined> | string | undefined;
@@ -247,6 +250,7 @@ export class Agent {
     this.mutableState = createMutableAgentState(options.initialState);
     this.convertToLlm = options.convertToLlm ?? defaultConvertToLlm;
     this.transformContext = options.transformContext;
+    this.transformAssistantMessage = options.transformAssistantMessage;
     this.runtime = options.runtime;
     this.streamFn = resolveAgentCoreStreamFn(options.runtime, options.streamFn);
     this.getApiKey = options.getApiKey;
@@ -495,6 +499,7 @@ export class Agent {
         : undefined,
       convertToLlm: this.convertToLlm,
       transformContext: this.transformContext,
+      transformAssistantMessage: this.transformAssistantMessage,
       getApiKey: this.getApiKey,
       getSteeringMessages: async () => {
         if (skipInitialSteeringPoll) {
diff --git a/packages/agent-core/src/types.ts b/packages/agent-core/src/types.ts
@@ -195,6 +195,18 @@ export interface AgentLoopConfig extends SimpleStreamOptions {
    */
   transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => Promise<AgentMessage[]>;
 
+  /**
+   * Optional transform applied to each finalized assistant message before it is
+   * emitted, stored in loop state, or used for tool execution.
+   *
+   * Contract: must not throw or reject. Return the original message or another
+   * safe fallback value instead.
+   */
+  transformAssistantMessage?: (
+    message: AssistantMessage,
+    signal?: AbortSignal,
+  ) => Promise<AssistantMessage> | AssistantMessage;
+
   /**
    * Resolves an API key dynamically for each LLM call.
    *
diff --git a/src/agents/embedded-agent-runner/run/attempt.ts b/src/agents/embedded-agent-runner/run/attempt.ts
@@ -47,6 +47,7 @@ import { resolveHeartbeatSummaryForAgent } from "../../../infra/heartbeat-summar
 import { getMachineDisplayName } from "../../../infra/machine-name.js";
 import { createCodexNativeWebSearchWrapper } from "../../../llm/providers/stream-wrappers/openai.js";
 import type { AssistantMessage } from "../../../llm/types.js";
+import type { PluginTextReplacement } from "../../../plugins/cli-backend.types.js";
 import { listRegisteredPluginAgentPromptGuidance } from "../../../plugins/command-registry-state.js";
 import { getCurrentPluginMetadataSnapshot } from "../../../plugins/current-plugin-metadata-snapshot.js";
 import {
@@ -178,7 +179,10 @@ import {
 import { resolveModelAuthMode } from "../../model-auth.js";
 import { resolveDefaultModelForAgent } from "../../model-selection.js";
 import { supportsModelTools } from "../../model-tool-support.js";
-import { wrapStreamFnTextTransforms } from "../../plugin-text-transforms.js";
+import {
+  transformPluginMessageText,
+  wrapStreamFnTextTransforms,
+} from "../../plugin-text-transforms.js";
 import { resolveAgentPromptSurfaceForSessionKey } from "../../prompt-surface.js";
 import { describeProviderRequestRoutingSummary } from "../../provider-attribution.js";
 import { registerProviderStreamForModel } from "../../provider-stream.js";
@@ -534,6 +538,37 @@ export {
 const MAX_BTW_SNAPSHOT_MESSAGES = 100;
 const PROMPT_TOOL_RESULT_AGGREGATE_CAP_MULTIPLIER = 4;
 
+type AssistantMessageTransform = (
+  message: AssistantMessage,
+  signal?: AbortSignal,
+) => AssistantMessage | Promise<AssistantMessage>;
+type AssistantMessageTransformAgent = {
+  transformAssistantMessage?: AssistantMessageTransform;
+};
+
+const originalAssistantMessageTransformByAgent = new WeakMap<
+  AssistantMessageTransformAgent,
+  AssistantMessageTransform | undefined
+>();
+
+function installAttemptAssistantMessageTextTransform(
+  agent: AssistantMessageTransformAgent,
+  output?: PluginTextReplacement[],
+): void {
+  if (!originalAssistantMessageTransformByAgent.has(agent)) {
+    originalAssistantMessageTransformByAgent.set(agent, agent.transformAssistantMessage);
+  }
+  const originalTransform = originalAssistantMessageTransformByAgent.get(agent);
+  if (!output || output.length === 0) {
+    agent.transformAssistantMessage = originalTransform;
+    return;
+  }
+  agent.transformAssistantMessage = async (message, signal) => {
+    const baseMessage = originalTransform ? await originalTransform(message, signal) : message;
+    return transformPluginMessageText(baseMessage, output);
+  };
+}
+
 function pluginMetadataSnapshotCoversProvider(
   snapshot: PluginMetadataSnapshot | undefined,
   provider: string,
@@ -2876,8 +2911,13 @@ export async function runEmbeddedAttempt(
           input: providerTextTransforms.input,
           output: providerTextTransforms.output,
           transformSystemPrompt: false,
+          transformFinalResult: false,
         });
       }
+      installAttemptAssistantMessageTextTransform(
+        activeSession.agent,
+        providerTextTransforms?.output,
+      );
       const nativeWebSearchPolicyContext = {
         sessionKey: sandboxSessionKey,
         sandboxToolPolicy: sandbox?.tools,
diff --git a/src/agents/plugin-text-transforms.test.ts b/src/agents/plugin-text-transforms.test.ts
@@ -1,15 +1,18 @@
 // Verifies plugin text transforms rewrite prompts and streamed assistant output.
-import type { StreamFn } from "openclaw/plugin-sdk/agent-core";
+import { runAgentLoop, type AgentEvent, type StreamFn } from "openclaw/plugin-sdk/agent-core";
 import {
   createAssistantMessageEventStream,
   type AssistantMessage,
   type Context,
   type Model,
+  type ToolCall,
 } from "openclaw/plugin-sdk/llm";
+import { Type } from "typebox";
 import { describe, expect, it } from "vitest";
 import {
   applyPluginTextReplacements,
   mergePluginTextTransforms,
+  transformPluginMessageText,
   wrapStreamFnTextTransforms,
 } from "./plugin-text-transforms.js";
 
@@ -40,6 +43,14 @@ function makeAssistantMessage(text: string): AssistantMessage {
   };
 }
 
+function makeAssistantToolMessage(toolCall: ToolCall): AssistantMessage {
+  return {
+    ...makeAssistantMessage("unused"),
+    content: [toolCall],
+    stopReason: "toolUse",
+  };
+}
+
 describe("plugin text transforms", () => {
   it("merges registered transform groups in order", () => {
     const merged = mergePluginTextTransforms(
@@ -180,4 +191,180 @@ describe("plugin text transforms", () => {
     expect(firstEvent?.delta).toBe("red basket on the left shelf");
     expect(result.content).toEqual([{ type: "text", text: "final red basket on the left shelf" }]);
   });
+
+  it("wraps streamed tool call deltas and argument strings with outbound replacements", async () => {
+    const streamedToolCall: ToolCall = {
+      type: "toolCall",
+      id: "call_[MASKED]",
+      name: "send_[MASKED]",
+      arguments: {
+        text: "Message for [MASKED]",
+        nested: { title: "[MASKED] follow-up", count: 2 },
+        recipients: ["[MASKED]", "ops", true],
+      },
+      partialArgs: '{"text":"Message for [MASKED]"}',
+      partialJson: '{"text":"Message for [MASKED]"}',
+    } as ToolCall;
+    const finalToolCall: ToolCall = {
+      ...streamedToolCall,
+      arguments: {
+        text: "Final message for [MASKED]",
+        nested: { title: "[MASKED] final", enabled: false },
+      },
+      partialArgs: '{"text":"Final message for [MASKED]"}',
+      partialJson: '{"text":"Final message for [MASKED]"}',
+    } as ToolCall;
+    const partial = makeAssistantToolMessage(streamedToolCall);
+    const finalMessage = makeAssistantToolMessage(finalToolCall);
+    const baseStreamFn: StreamFn = () => {
+      const stream = createAssistantMessageEventStream();
+      queueMicrotask(() => {
+        stream.push({
+          type: "toolcall_delta",
+          contentIndex: 0,
+          delta: '{"text":"Message for [MASKED]"}',
+          partial,
+        });
+        stream.push({
+          type: "toolcall_end",
+          contentIndex: 0,
+          toolCall: streamedToolCall,
+          partial,
+        });
+        stream.push({
+          type: "done",
+          reason: "toolUse",
+          message: finalMessage,
+        });
+        stream.end();
+      });
+      return stream;
+    };
+
+    const wrapped = wrapStreamFnTextTransforms({
+      streamFn: baseStreamFn,
+      output: [{ from: /\[MASKED\]/g, to: "John Smith" }],
+    });
+    const stream = await Promise.resolve(wrapped(model, { messages: [] } as Context, undefined));
+    const events = [];
+    for await (const event of stream) {
+      events.push(event);
+    }
+    const result = await stream.result();
+
+    const delta = events[0] as { type?: string; delta?: string; partial?: AssistantMessage };
+    expect(delta.type).toBe("toolcall_delta");
+    expect(delta.delta).toBe('{"text":"Message for John Smith"}');
+    expect(delta.partial?.content[0]).toMatchObject({
+      id: "call_[MASKED]",
+      name: "send_[MASKED]",
+      partialArgs: '{"text":"Message for John Smith"}',
+      partialJson: '{"text":"Message for John Smith"}',
+      arguments: {
+        text: "Message for John Smith",
+        nested: { title: "John Smith follow-up", count: 2 },
+        recipients: ["John Smith", "ops", true],
+      },
+    });
+
+    const end = events[1] as { type?: string; toolCall?: ToolCall; partial?: AssistantMessage };
+    expect(end.type).toBe("toolcall_end");
+    expect(end.toolCall).toMatchObject({
+      id: "call_[MASKED]",
+      name: "send_[MASKED]",
+      partialArgs: '{"text":"Message for John Smith"}',
+      partialJson: '{"text":"Message for John Smith"}',
+      arguments: {
+        text: "Message for John Smith",
+        nested: { title: "John Smith follow-up", count: 2 },
+        recipients: ["John Smith", "ops", true],
+      },
+    });
+    expect(result.content[0]).toMatchObject({
+      id: "call_[MASKED]",
+      name: "send_[MASKED]",
+      partialArgs: '{"text":"Final message for John Smith"}',
+      partialJson: '{"text":"Final message for John Smith"}',
+      arguments: {
+        text: "Final message for John Smith",
+        nested: { title: "John Smith final", enabled: false },
+      },
+    });
+  });
+
+  it("rewrites finalized tool call arguments before agent tool execution", async () => {
+    const output = [{ from: /cat/g, to: "black cat" }];
+    const capturedEvents: AgentEvent[] = [];
+    const executedArgs: unknown[] = [];
+    let turn = 0;
+    const baseStreamFn: StreamFn = () => {
+      turn += 1;
+      const stream = createAssistantMessageEventStream();
+      queueMicrotask(() => {
+        const message =
+          turn === 1
+            ? makeAssistantToolMessage({
+                type: "toolCall",
+                id: "call_read",
+                name: "read",
+                arguments: { path: "cat.txt" },
+                partialArgs: '{"path":"cat.txt"}',
+              } as ToolCall)
+            : makeAssistantMessage("done");
+        stream.push({
+          type: "done",
+          reason: message.stopReason === "toolUse" ? "toolUse" : "stop",
+          message,
+        });
+        stream.end();
+      });
+      return stream;
+    };
+    const wrapped = wrapStreamFnTextTransforms({
+      streamFn: baseStreamFn,
+      output,
+      transformFinalResult: false,
+    });
+
+    await runAgentLoop(
+      [{ role: "user", content: "read it", timestamp: 1 }],
+      {
+        systemPrompt: "",
+        messages: [],
+        tools: [
+          {
+            name: "read",
+            label: "read",
+            description: "read",
+            parameters: Type.Object({ path: Type.String() }, { additionalProperties: false }),
+            execute: async (_toolCallId, args) => {
+              executedArgs.push(args);
+              return {
+                content: [{ type: "text", text: "ok" }],
+                details: args,
+                terminate: true,
+              };
+            },
+          },
+        ],
+      },
+      {
+        model,
+        convertToLlm: (messages) => messages as never,
+        transformAssistantMessage: (message) => transformPluginMessageText(message, output),
+      },
+      (event) => {
+        capturedEvents.push(event);
+      },
+      undefined,
+      wrapped,
+    );
+
+    const toolStart = capturedEvents.find(
+      (event): event is Extract<AgentEvent, { type: "tool_execution_start" }> =>
+        event.type === "tool_execution_start",
+    );
+    expect(toolStart?.args).toEqual({ path: "black cat.txt" });
+    expect(executedArgs).toEqual([{ path: "black cat.txt" }]);
+  });
 });
diff --git a/src/agents/plugin-text-transforms.ts b/src/agents/plugin-text-transforms.ts
