Skip to content

Commit 2c83b4e

Browse files
steipetesteipete
committed
fix(agents): preserve structured replay provenance
1 parent baeb8d3 commit 2c83b4e

15 files changed

Lines changed: 306 additions & 70 deletions

‎src/agents/agent-tool-definition-adapter.after-tool-call.fires-once.test.ts‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ const beforeToolCallMocks = vi.hoisted(() => ({
3131
},
3232
consumeAdjustedParamsForToolCall: vi.fn((_: string): unknown => undefined),
3333
recordAdjustedParamsForToolCall: vi.fn(),
34+
recordStructuredReplayTrustForToolCall: vi.fn(),
3435
isToolWrappedWithBeforeToolCallHook: vi.fn(() => false),
3536
runBeforeToolCallHook: vi.fn(async ({ params }: { params: unknown }) => ({
3637
blocked: false,
@@ -97,14 +98,22 @@ async function loadFreshAfterToolCallModulesForTest() {
9798
emitAgentEvent: vi.fn(),
9899
emitAgentItemEvent: vi.fn(),
99100
}));
101+
vi.doMock("./agent-tools.before-tool-call.state.js", () => ({
102+
consumeAdjustedParamsForToolCall: beforeToolCallMocks.consumeAdjustedParamsForToolCall,
103+
consumePreExecutionBlockedToolCall: vi.fn(() => false),
104+
consumeStructuredReplaySafeToolCall: vi.fn(() => false),
105+
}));
100106
vi.doMock("./agent-tools.before-tool-call.js", () => ({
101107
BeforeToolCallBlockedError: beforeToolCallMocks.BeforeToolCallBlockedError,
102108
buildBlockedToolResult: ({ reason }: { reason: string }) => ({
103109
content: [{ type: "text", text: reason }],
104110
details: { status: "blocked", deniedReason: "plugin-before-tool-call", reason },
105111
}),
106112
consumeAdjustedParamsForToolCall: beforeToolCallMocks.consumeAdjustedParamsForToolCall,
113+
consumePreExecutionBlockedToolCall: vi.fn(() => false),
107114
recordAdjustedParamsForToolCall: beforeToolCallMocks.recordAdjustedParamsForToolCall,
115+
recordStructuredReplayTrustForToolCall:
116+
beforeToolCallMocks.recordStructuredReplayTrustForToolCall,
108117
isBeforeToolCallBlockedError: (error: unknown) =>
109118
error instanceof beforeToolCallMocks.BeforeToolCallBlockedError,
110119
isToolWrappedWithBeforeToolCallHook: beforeToolCallMocks.isToolWrappedWithBeforeToolCallHook,

‎src/agents/agent-tool-definition-adapter.after-tool-call.test.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ const hookMocks = vi.hoisted(() => ({
2525
isToolWrappedWithBeforeToolCallHook: vi.fn(() => false),
2626
consumeAdjustedParamsForToolCall: vi.fn((_: string) => undefined as unknown),
2727
recordAdjustedParamsForToolCall: vi.fn(),
28+
recordStructuredReplayTrustForToolCall: vi.fn(),
2829
runBeforeToolCallHook: vi.fn(async ({ params }: { params: unknown }) => ({
2930
blocked: false,
3031
params,
@@ -43,6 +44,7 @@ vi.mock("./agent-tools.before-tool-call.js", () => ({
4344
}),
4445
consumeAdjustedParamsForToolCall: hookMocks.consumeAdjustedParamsForToolCall,
4546
recordAdjustedParamsForToolCall: hookMocks.recordAdjustedParamsForToolCall,
47+
recordStructuredReplayTrustForToolCall: hookMocks.recordStructuredReplayTrustForToolCall,
4648
isBeforeToolCallBlockedError: (error: unknown) =>
4749
error instanceof hookMocks.BeforeToolCallBlockedError,
4850
isToolWrappedWithBeforeToolCallHook: hookMocks.isToolWrappedWithBeforeToolCallHook,
@@ -72,6 +74,7 @@ describe("agent tool definition adapter after_tool_call", () => {
7274
hookMocks.consumeAdjustedParamsForToolCall.mockClear();
7375
hookMocks.consumeAdjustedParamsForToolCall.mockReturnValue(undefined);
7476
hookMocks.recordAdjustedParamsForToolCall.mockClear();
77+
hookMocks.recordStructuredReplayTrustForToolCall.mockClear();
7578
hookMocks.runBeforeToolCallHook.mockClear();
7679
hookMocks.runBeforeToolCallHook.mockImplementation(async ({ params }) => ({
7780
blocked: false,

‎src/agents/agent-tool-definition-adapter.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import {
1313
isToolWrappedWithBeforeToolCallHook,
1414
isBeforeToolCallBlockedError,
1515
recordAdjustedParamsForToolCall,
16+
recordStructuredReplayTrustForToolCall,
1617
runBeforeToolCallHook,
1718
} from "./agent-tools.before-tool-call.js";
1819
import {
@@ -372,6 +373,7 @@ export function toToolDefinitions(
372373
executionMode: tool.executionMode,
373374
execute: async (...args: ToolExecuteArgs): Promise<AgentToolResult<unknown>> => {
374375
const { toolCallId, params, onUpdate, signal } = splitToolExecuteArgs(args);
376+
recordStructuredReplayTrustForToolCall(toolCallId, tool, hookContext?.runId);
375377
let executeParams = params;
376378
try {
377379
if (!beforeHookWrapped) {

‎src/agents/agent-tools.before-tool-call.integration.e2e.test.ts‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import { addTestHook, createMockPluginRegistry } from "../plugins/hooks.test-hel
1717
import { patchPluginSessionExtension } from "../plugins/host-hook-state.js";
1818
import { createEmptyPluginRegistry } from "../plugins/registry.js";
1919
import { setActivePluginRegistry } from "../plugins/runtime.js";
20+
import { setPluginToolMeta } from "../plugins/tools.js";
2021
import type { PluginHookRegistration } from "../plugins/types.js";
2122
import { toClientToolDefinitions, toToolDefinitions } from "./agent-tool-definition-adapter.js";
2223
import { wrapToolWithAbortSignal } from "./agent-tools.abort.js";
@@ -31,6 +32,7 @@ import { normalizeToolParameters } from "./agent-tools.schema.js";
3132
import { markCodeModeControlTool } from "./code-mode-control-tools.js";
3233
import { CODE_MODE_EXEC_TOOL_NAME, createCodeModeTools } from "./code-mode.js";
3334
import { splitSdkTools } from "./embedded-agent-runner.js";
35+
import type { ExtensionContext } from "./sessions/index.js";
3436
import { setToolTerminalPresentation } from "./tool-terminal-presentation.js";
3537

3638
type BeforeToolCallHandlerMock = ReturnType<typeof vi.fn>;
@@ -92,6 +94,7 @@ describe("before_tool_call hook integration", () => {
9294
resetGlobalHookRunner();
9395
resetDiagnosticSessionStateForTest();
9496
beforeToolCallTesting.adjustedParamsByToolCallId.clear();
97+
beforeToolCallTesting.structuredReplaySafeToolCallIds.clear();
9598
beforeToolCallHook = installBeforeToolCallHook();
9699
});
97100

@@ -115,6 +118,54 @@ describe("before_tool_call hook integration", () => {
115118
);
116119
});
117120

121+
it("records structured replay trust only for concrete core-owned tools", async () => {
122+
beforeToolCallHook = installBeforeToolCallHook({ enabled: false });
123+
const execute = vi.fn().mockResolvedValue({ content: [], details: { ok: true } });
124+
const coreTool = wrapToolWithBeforeToolCallHook({ name: "search", execute } as any, {
125+
runId: "run-core",
126+
});
127+
const pluginSource = { name: "search", execute } as any;
128+
setPluginToolMeta(pluginSource, { pluginId: "example", optional: false });
129+
const pluginTool = wrapToolWithBeforeToolCallHook(pluginSource, {
130+
runId: "run-plugin",
131+
});
132+
133+
const [coreDefinition] = toToolDefinitions([coreTool], { runId: "run-core" });
134+
const [pluginDefinition] = toToolDefinitions([pluginTool], { runId: "run-plugin" });
135+
const extensionContext = {} as ExtensionContext;
136+
await coreDefinition?.execute(
137+
"call-core",
138+
{ query: "core" },
139+
undefined,
140+
undefined,
141+
extensionContext,
142+
);
143+
await pluginDefinition?.execute(
144+
"call-plugin",
145+
{ query: "plugin" },
146+
undefined,
147+
undefined,
148+
extensionContext,
149+
);
150+
151+
expect(
152+
beforeToolCallTesting.structuredReplaySafeToolCallIds.has(
153+
beforeToolCallTesting.buildAdjustedParamsKey({
154+
runId: "run-core",
155+
toolCallId: "call-core",
156+
}),
157+
),
158+
).toBe(true);
159+
expect(
160+
beforeToolCallTesting.structuredReplaySafeToolCallIds.has(
161+
beforeToolCallTesting.buildAdjustedParamsKey({
162+
runId: "run-plugin",
163+
toolCallId: "call-plugin",
164+
}),
165+
),
166+
).toBe(false);
167+
});
168+
118169
it("allows hook to modify parameters", async () => {
119170
beforeToolCallHook = installBeforeToolCallHook({
120171
runBeforeToolCallImpl: async () => ({ params: { mode: "safe" } }),

‎src/agents/agent-tools.before-tool-call.state.ts‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,52 @@
55
*/
66
export const adjustedParamsByToolCallId = new Map<string, unknown>();
77
export const preExecutionBlockedToolCallIds = new Set<string>();
8+
export const structuredReplaySafeToolCallIds = new Set<string>();
9+
10+
export function buildAdjustedParamsKey(params: { runId?: string; toolCallId: string }): string {
11+
if (params.runId && params.runId.trim()) {
12+
return `${params.runId}:${params.toolCallId}`;
13+
}
14+
return params.toolCallId;
15+
}
16+
17+
/** Consume and remove hook-adjusted params for a completed tool call. */
18+
export function consumeAdjustedParamsForToolCall(toolCallId: string, runId?: string): unknown {
19+
const key = buildAdjustedParamsKey({ runId, toolCallId });
20+
const params = adjustedParamsByToolCallId.get(key);
21+
adjustedParamsByToolCallId.delete(key);
22+
return params;
23+
}
24+
25+
/** Snapshot hook-adjusted params without consuming later outcome bookkeeping. */
26+
export function peekAdjustedParamsForToolCall(toolCallId: string, runId?: string): unknown {
27+
const key = buildAdjustedParamsKey({ runId, toolCallId });
28+
const params = adjustedParamsByToolCallId.get(key);
29+
return params === undefined ? undefined : structuredClone(params);
30+
}
31+
32+
/** Consume whether policy prevented the target tool from starting. */
33+
export function consumePreExecutionBlockedToolCall(toolCallId: string, runId?: string): boolean {
34+
const key = buildAdjustedParamsKey({ runId, toolCallId });
35+
const blocked = preExecutionBlockedToolCallIds.has(key);
36+
preExecutionBlockedToolCallIds.delete(key);
37+
return blocked;
38+
}
39+
40+
export function recordStructuredReplaySafeToolCall(toolCallId: string, runId?: string): void {
41+
structuredReplaySafeToolCallIds.add(buildAdjustedParamsKey({ runId, toolCallId }));
42+
}
43+
44+
export function consumeStructuredReplaySafeToolCall(toolCallId: string, runId?: string): boolean {
45+
const key = buildAdjustedParamsKey({ runId, toolCallId });
46+
const replaySafe = structuredReplaySafeToolCallIds.has(key);
47+
structuredReplaySafeToolCallIds.delete(key);
48+
return replaySafe;
49+
}
850

951
/** Clear adjusted tool parameters between isolated tests. */
1052
export function resetAdjustedParamsByToolCallIdForTests(): void {
1153
adjustedParamsByToolCallId.clear();
1254
preExecutionBlockedToolCallIds.clear();
55+
structuredReplaySafeToolCallIds.clear();
1356
}

‎src/agents/agent-tools.before-tool-call.ts‎

Lines changed: 28 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,15 @@ import { resolveSkillWorkshopToolApproval } from "../skills/workshop/policy.js";
6161
import { isPlainObject, truncateUtf16Safe } from "../utils.js";
6262
import {
6363
adjustedParamsByToolCallId,
64+
buildAdjustedParamsKey,
6465
preExecutionBlockedToolCallIds,
66+
recordStructuredReplaySafeToolCall,
67+
structuredReplaySafeToolCallIds,
68+
} from "./agent-tools.before-tool-call.state.js";
69+
export {
70+
consumeAdjustedParamsForToolCall,
71+
consumePreExecutionBlockedToolCall,
72+
peekAdjustedParamsForToolCall,
6573
} from "./agent-tools.before-tool-call.state.js";
6674
import { copyChannelAgentToolMeta, getChannelAgentToolMeta } from "./channel-tools.js";
6775
import {
@@ -352,6 +360,25 @@ export function recordAdjustedParamsForToolCall(
352360
}
353361
}
354362

363+
/** Record that one concrete core-owned tool call may use structured replay classification. */
364+
export function recordStructuredReplayTrustForToolCall(
365+
toolCallId: string | undefined,
366+
tool: AnyAgentTool,
367+
runId?: string,
368+
): void {
369+
if (!toolCallId || getPluginToolMeta(tool) || getChannelAgentToolMeta(tool as never)) {
370+
return;
371+
}
372+
recordStructuredReplaySafeToolCall(toolCallId, runId);
373+
while (structuredReplaySafeToolCallIds.size > MAX_TRACKED_ADJUSTED_PARAMS) {
374+
const oldest = structuredReplaySafeToolCallIds.values().next().value;
375+
if (!oldest) {
376+
break;
377+
}
378+
structuredReplaySafeToolCallIds.delete(oldest);
379+
}
380+
}
381+
355382
/**
356383
* Returns true when an error represents an intentional before_tool_call veto.
357384
*/
@@ -364,13 +391,6 @@ const loadBeforeToolCallRuntime = createLazyRuntimeSurface(
364391
({ beforeToolCallRuntime }) => beforeToolCallRuntime,
365392
);
366393

367-
function buildAdjustedParamsKey(params: { runId?: string; toolCallId: string }): string {
368-
if (params.runId && params.runId.trim()) {
369-
return `${params.runId}:${params.toolCallId}`;
370-
}
371-
return params.toolCallId;
372-
}
373-
374394
function mergeParamsWithApprovalOverrides(
375395
originalParams: unknown,
376396
approvalParams?: unknown,
@@ -1538,29 +1558,6 @@ export function copyBeforeToolCallHookMarker(source: AnyAgentTool, target: AnyAg
15381558
});
15391559
}
15401560

1541-
/** Consume and remove hook-adjusted params for a completed tool call. */
1542-
export function consumeAdjustedParamsForToolCall(toolCallId: string, runId?: string): unknown {
1543-
const adjustedParamsKey = buildAdjustedParamsKey({ runId, toolCallId });
1544-
const params = adjustedParamsByToolCallId.get(adjustedParamsKey);
1545-
adjustedParamsByToolCallId.delete(adjustedParamsKey);
1546-
return params;
1547-
}
1548-
1549-
/** Snapshot hook-adjusted params without consuming later outcome bookkeeping. */
1550-
export function peekAdjustedParamsForToolCall(toolCallId: string, runId?: string): unknown {
1551-
const adjustedParamsKey = buildAdjustedParamsKey({ runId, toolCallId });
1552-
const params = adjustedParamsByToolCallId.get(adjustedParamsKey);
1553-
return params === undefined ? undefined : structuredClone(params);
1554-
}
1555-
1556-
/** Consume whether policy prevented the target tool from starting. */
1557-
export function consumePreExecutionBlockedToolCall(toolCallId: string, runId?: string): boolean {
1558-
const key = buildAdjustedParamsKey({ runId, toolCallId });
1559-
const blocked = preExecutionBlockedToolCallIds.has(key);
1560-
preExecutionBlockedToolCallIds.delete(key);
1561-
return blocked;
1562-
}
1563-
15641561
function recordPreExecutionBlockedToolCall(toolCallId?: string, runId?: string): void {
15651562
if (!toolCallId) {
15661563
return;
@@ -1584,6 +1581,7 @@ export const testing = {
15841581
buildAdjustedParamsKey,
15851582
adjustedParamsByToolCallId,
15861583
preExecutionBlockedToolCallIds,
1584+
structuredReplaySafeToolCallIds,
15871585
runBeforeToolCallHook,
15881586
mergeParamsWithApprovalOverrides,
15891587
isPlainObject,

‎src/agents/agent-tools.deferred-followup-guidance.test.ts‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,10 @@
33
* Protects the model-facing text selected after tool filtering.
44
*/
55
import { describe, expect, it } from "vitest";
6+
import { getPluginToolMeta, setPluginToolMeta } from "../plugins/tools.js";
67
import { applyDeferredFollowupToolDescriptions } from "./agent-tools.deferred-followup.js";
78
import type { AnyAgentTool } from "./agent-tools.types.js";
9+
import { getChannelAgentToolMeta, setChannelAgentToolMeta } from "./channel-tool-metadata.js";
810

911
function findToolDescription(toolName: string, includeCron: boolean) {
1012
const tools = applyDeferredFollowupToolDescriptions([
@@ -45,4 +47,21 @@ describe("createOpenClawCodingTools deferred follow-up guidance", () => {
4547
"Manage running exec sessions for commands already started: list, poll, log, write, send-keys, submit, paste, kill. Use poll/log when you need status, logs, quiet-success confirmation, or completion confirmation when automatic completion wake is unavailable. Use poll/log also for input-wait hints. Use write/send-keys/submit/paste/kill for input or intervention.",
4648
);
4749
});
50+
51+
it("preserves ownership metadata when replacing process descriptions", () => {
52+
const processTool = {
53+
name: "process",
54+
description: "plugin process",
55+
} as AnyAgentTool;
56+
setPluginToolMeta(processTool, { pluginId: "example", optional: false });
57+
setChannelAgentToolMeta(processTool as never, { channelId: "example-channel" });
58+
59+
const [updated] = applyDeferredFollowupToolDescriptions([processTool]);
60+
61+
expect(updated).not.toBe(processTool);
62+
expect(getPluginToolMeta(updated)).toEqual({ pluginId: "example", optional: false });
63+
expect(getChannelAgentToolMeta(updated as never)).toEqual({
64+
channelId: "example-channel",
65+
});
66+
});
4867
});

‎src/agents/agent-tools.deferred-followup.ts‎

Lines changed: 15 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,23 @@
1+
import { copyPluginToolMeta } from "../plugins/tools.js";
2+
import { copyBeforeToolCallHookMarker } from "./agent-tools.before-tool-call.js";
13
/**
24
* Adjusts exec/process tool descriptions for long-running follow-up behavior.
35
* Cron-aware runs can point models at scheduled follow-ups; cronless runs keep
46
* guidance constrained to process polling and wake handling.
57
*/
68
import type { AnyAgentTool } from "./agent-tools.types.js";
79
import { describeExecTool, describeProcessTool } from "./bash-tools.descriptions.js";
10+
import { copyChannelAgentToolMeta } from "./channel-tools.js";
11+
import { copyToolTerminalPresentation } from "./tool-terminal-presentation.js";
12+
13+
function replaceDescription(tool: AnyAgentTool, description: string): AnyAgentTool {
14+
const updated = { ...tool, description };
15+
copyPluginToolMeta(tool, updated);
16+
copyChannelAgentToolMeta(tool as never, updated as never);
17+
copyBeforeToolCallHookMarker(tool, updated);
18+
copyToolTerminalPresentation(tool, updated);
19+
return updated;
20+
}
821

922
/** Return tools with exec/process descriptions adjusted for cron availability. */
1023
export function applyDeferredFollowupToolDescriptions(
@@ -14,16 +27,10 @@ export function applyDeferredFollowupToolDescriptions(
1427
const hasCronTool = tools.some((tool) => tool.name === "cron");
1528
return tools.map((tool) => {
1629
if (tool.name === "exec") {
17-
return {
18-
...tool,
19-
description: describeExecTool({ agentId: params?.agentId, hasCronTool }),
20-
};
30+
return replaceDescription(tool, describeExecTool({ agentId: params?.agentId, hasCronTool }));
2131
}
2232
if (tool.name === "process") {
23-
return {
24-
...tool,
25-
description: describeProcessTool({ hasCronTool }),
26-
};
33+
return replaceDescription(tool, describeProcessTool({ hasCronTool }));
2734
}
2835
return tool;
2936
});

0 commit comments

Comments
 (0)