Skip to content
Open
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 26 additions & 0 deletions packages/core/src/agent/agent.spec-d.ts
Original file line number Diff line number Diff line change
Expand Up @@ -628,6 +628,32 @@ describe("Agent Type System", () => {
expectTypeOf(hooks).toMatchTypeOf<AgentHooks>();
});

it("should validate agent toolGuard option", () => {
const agentOptions: AgentOptions = {
name: "GuardedAgent",
instructions: "Test",
model: "openai/gpt-4o-mini",
toolGuard: async ({ agent, tool, args, context }) => {
expectTypeOf(agent).toMatchTypeOf<Agent>();
expectTypeOf(tool.name).toEqualTypeOf<string>();
expectTypeOf(args).toBeAny();
expectTypeOf(context).toMatchTypeOf<OperationContext>();
return { denied: true, reason: "read-only" };
},
};

expectTypeOf(agentOptions).toMatchTypeOf<AgentOptions>();

const booleanGuardOptions: AgentOptions = {
name: "BooleanGuardAgent",
instructions: "Test",
model: "openai/gpt-4o-mini",
toolGuard: () => true,
};

expectTypeOf(booleanGuardOptions).toMatchTypeOf<AgentOptions>();
});

it("should allow sync and async hooks", () => {
const syncHooks: AgentHooks = {
onStart: () => {
Expand Down
71 changes: 71 additions & 0 deletions packages/core/src/agent/agent.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1489,6 +1489,77 @@ Use pandas and summarize findings.`.split("\n"),
operationContext.traceContext.end("completed");
});

it("blocks tool execution when toolGuard denies the tool", async () => {
const execute = vi.fn().mockResolvedValue("should-not-run");
const toolGuard = vi.fn().mockResolvedValue({
denied: true,
reason: "read-only agent",
});
const onToolError = vi.fn();
const onToolEnd = vi.fn();
const agent = new Agent({
name: "TestAgent",
instructions: "Test",
model: mockModel as any,
toolGuard,
hooks: createHooks({ onToolError, onToolEnd }),
});

const protectedTool = new Tool({
name: "delete-note",
description: "Deletes a note",
parameters: z.object({ id: z.string() }),
execute,
});

const operationContext = (agent as any).createOperationContext("input");
const executeFactory = (agent as any).createToolExecutionFactory(
operationContext,
agent.hooks,
);

const result = await executeFactory(protectedTool)({ id: "note-1" });

expect(execute).not.toHaveBeenCalled();
expect(toolGuard).toHaveBeenCalledWith(
expect.objectContaining({
agent,
tool: protectedTool,
args: { id: "note-1" },
context: operationContext,
}),
);
expect(result).toMatchObject({
error: true,
toolName: "delete-note",
code: "TOOL_FORBIDDEN",
});
expect(result.message).toContain("read-only agent");
expect(onToolError).toHaveBeenCalledTimes(1);
const toolErrorArgs = onToolError.mock.calls[0][0];
expect(toolErrorArgs.tool).toBe(protectedTool);
expect(toolErrorArgs.args).toEqual({ id: "note-1" });
expect(toolErrorArgs.originalError).toMatchObject({
code: "TOOL_FORBIDDEN",
message: "read-only agent",
});
expect(toolErrorArgs.error).toMatchObject({
message: "read-only agent",
stage: "tool_execution",
});

expect(onToolEnd).toHaveBeenCalledTimes(1);
const toolEndArgs = onToolEnd.mock.calls[0][0];
expect(toolEndArgs.tool).toBe(protectedTool);
expect(toolEndArgs.output).toBeUndefined();
expect(toolEndArgs.error).toMatchObject({
message: "read-only agent",
stage: "tool_execution",
});

operationContext.traceContext.end("completed");
});

it("calls onToolError when a tool throws", async () => {
const onToolError = vi.fn();
const onToolEnd = vi.fn();
Expand Down
54 changes: 52 additions & 2 deletions packages/core/src/agent/agent.ts
Original file line number Diff line number Diff line change
Expand Up @@ -106,7 +106,12 @@ import {
type EnqueueEvalScoringArgs,
enqueueEvalScoring as enqueueEvalScoringHelper,
} from "./eval";
import type { AgentHooks, OnToolEndHookResult, OnToolErrorHookResult } from "./hooks";
import type {
AgentHooks,
AgentToolGuard,
OnToolEndHookResult,
OnToolErrorHookResult,
} from "./hooks";
import { stripDanglingOpenAIReasoningFromModelMessages } from "./model-message-normalizer";
import { AgentTraceContext, addModelAttributesToSpan } from "./open-telemetry/trace-context";
import {
Expand Down Expand Up @@ -1044,6 +1049,7 @@ export class Agent {
private readonly workspaceToolkitOptions: AgentOptions["workspaceToolkits"];
private readonly workspaceSkillsPromptOption: AgentOptions["workspaceSkillsPrompt"];
private readonly configuredHooks?: AgentHooks;
private readonly toolGuard?: AgentToolGuard;
private readonly maxStepsConfigured: boolean;
private defaultObservability?: VoltAgentObservability;
private readonly toolManager: ToolManager;
Expand Down Expand Up @@ -1078,6 +1084,7 @@ export class Agent {
this.workspaceToolkitOptions = options.workspaceToolkits;
this.workspaceSkillsPromptOption = options.workspaceSkillsPrompt;
this.configuredHooks = options.hooks;
this.toolGuard = options.toolGuard;
this.maxStepsConfigured = options.maxSteps !== undefined;
const globalWorkspace = AgentRegistry.getInstance().getGlobalWorkspace();
const workspaceOption = options.workspace === undefined ? globalWorkspace : options.workspace;
Expand Down Expand Up @@ -6421,6 +6428,46 @@ export class Agent {
return parseResult.data;
}

private async assertToolGuardAllows(
tool: BaseTool | ProviderTool,
args: any,
oc: OperationContext,
options?: ToolExecuteOptions,
): Promise<void> {
if (!this.toolGuard) {
return;
}

const result = await this.toolGuard({
agent: this,
tool: tool as any,
context: oc,
args,
options,
});

const denied =
result === false ||
(typeof result === "object" &&
result !== null &&
(result.denied === true || result.allowed === false));
if (!denied) {
return;
}

const reason =
typeof result === "object" && result !== null && typeof result.reason === "string"
? result.reason
: "Tool execution denied by toolGuard.";

throw new ToolDeniedError({
toolName: tool.name,
message: reason,
code: "TOOL_FORBIDDEN",
httpStatus: 403,
});
}

private createToolExecutionFactory(
oc: OperationContext,
hooks: AgentHooks,
Expand Down Expand Up @@ -6623,6 +6670,7 @@ export class Agent {
try {
await this.waitForSpeculativeInputGuardrail(oc);
await oc.traceContext.withSpan(toolSpan, async () => {
await this.assertToolGuardAllows(tool, args, oc, executionOptions);
Comment thread
cubic-dev-ai[bot] marked this conversation as resolved.
await runToolStartHooks();
Comment on lines +6665 to 6666

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win

Call the tool end hook once for denied local calls.

When toolGuard denies a local tool, execution enters handleToolError. That handler invokes tool.hooks.onEnd twice at Lines 6617-6631. A denied call can therefore create duplicate tool-level audit records or duplicate cleanup side effects.

Remove the duplicate invocation. Add a test that asserts tool.hooks.onEnd runs exactly once for a denied call.

Also applies to: 6734-6736

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@packages/core/src/agent/agent.ts` around lines 6673 - 6674, Remove the
duplicate tool.hooks.onEnd invocation in handleToolError for denied local calls,
including the corresponding path near the alternate call site, so each denied
call triggers the end hook exactly once. Add or update a test covering a denied
local tool call and assert tool.hooks.onEnd is invoked once.

});

Expand Down Expand Up @@ -6683,7 +6731,8 @@ export class Agent {
return oc.traceContext.withSpan(toolSpan, async () => {
try {
await this.waitForSpeculativeInputGuardrail(oc);
// Call tool start hook - can throw ToolDeniedError
// Call tool guard and start hook - both can throw ToolDeniedError
await this.assertToolGuardAllows(tool, args, oc, executionOptions);
await runToolStartHooks();

// Execute tool with merged options
Expand Down Expand Up @@ -7242,6 +7291,7 @@ export class Agent {
`Provider tool "${tool.name}" received arguments that do not match callTool input.`,
);
}
await this.assertToolGuardAllows(tool, callInput, oc, executionOptions);
Comment thread
cubic-dev-ai[bot] marked this conversation as resolved.
Outdated

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1: For provider tools this guard runs after the tool has already been executed, so it does not actually authorize the call.

In executeProviderToolViaCallTool, runInternalGenerateText (~line 7284) is invoked before this guard. It calls generateText({ tools: { [tool.name]: tool }, toolChoice: { type: "tool", toolName: tool.name } }), which makes the AI SDK invoke the provider tool's own execute and record its result. By the time assertToolGuardAllows is reached here, the provider tool has already run with all its side effects. A denial only prevents the result from being surfaced to the caller — it cannot stop the tool from executing. This is inconsistent with the local-tool paths where the guard runs before tool.execute, and it gives a false sense of authorization for provider tools routed via callTool.

Additionally, provider tools exposed directly to the model are passed through untouched in ToolManager.prepareToolsForExecution (tools[tool.name] = tool;), so they never go through createToolExecutionFactory and never hit assertToolGuardAllows at all on the direct-execution path. Providers routed only through callTool are the sole case that touches this guard, and that happens after execution. Consider moving the guard to before the provider tool's callTool execution begins (e.g., before runInternalGenerateText) and documenting/covering the direct pass-through case.

Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At packages/core/src/agent/agent.ts, line 7294:

<comment>For provider tools this guard runs after the tool has already been executed, so it does not actually authorize the call.

In `executeProviderToolViaCallTool`, `runInternalGenerateText` (~line 7284) is invoked before this guard. It calls `generateText({ tools: { [tool.name]: tool }, toolChoice: { type: "tool", toolName: tool.name } })`, which makes the AI SDK invoke the provider tool's own `execute` and record its result. By the time `assertToolGuardAllows` is reached here, the provider tool has already run with all its side effects. A denial only prevents the result from being surfaced to the caller — it cannot stop the tool from executing. This is inconsistent with the local-tool paths where the guard runs before `tool.execute`, and it gives a false sense of authorization for provider tools routed via `callTool`.

Additionally, provider tools exposed directly to the model are passed through untouched in `ToolManager.prepareToolsForExecution` (`tools[tool.name] = tool;`), so they never go through `createToolExecutionFactory` and never hit `assertToolGuardAllows` at all on the direct-execution path. Providers routed only through `callTool` are the sole case that touches this guard, and that happens after execution. Consider moving the guard to before the provider tool's `callTool` execution begins (e.g., before `runInternalGenerateText`) and documenting/covering the direct pass-through case.</comment>

<file context>
@@ -7242,6 +7291,7 @@ export class Agent {
           `Provider tool "${tool.name}" received arguments that do not match callTool input.`,
         );
       }
+      await this.assertToolGuardAllows(tool, callInput, oc, executionOptions);
       await hooks.onToolStart?.({
         agent: this,
</file context>

await hooks.onToolStart?.({
agent: this,
tool: tool as any,
Expand Down
13 changes: 13 additions & 0 deletions packages/core/src/agent/hooks/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,19 @@ export interface OnToolStartHookArgs {
options?: ToolExecuteOptions;
}

export interface ToolGuardArgs extends OnToolStartHookArgs {}

export type ToolGuardResult =
| boolean
| {
allowed?: boolean;
Comment thread
cubic-dev-ai[bot] marked this conversation as resolved.
Outdated
denied?: boolean;
reason?: string;
}
| undefined;

export type AgentToolGuard = (args: ToolGuardArgs) => Promise<ToolGuardResult> | ToolGuardResult;
Comment thread
coderabbitai[bot] marked this conversation as resolved.

export interface OnToolEndHookArgs {
agent: Agent;
tool: AgentTool;
Expand Down
7 changes: 6 additions & 1 deletion packages/core/src/agent/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,7 @@ import type {
WorkspaceSkillsToolkitOptions,
} from "../workspace";
import type { ContextInput } from "./agent";
import type { AgentHooks } from "./hooks";
import type { AgentHooks, AgentToolGuard } from "./hooks";
import type { AgentTraceContext } from "./open-telemetry/trace-context";

// Re-export for backward compatibility
Expand Down Expand Up @@ -717,6 +717,11 @@ export type AgentOptions = {

// Hooks
hooks?: AgentHooks;
/**
* Optional per-tool authorization guard.
* Return `false`, `{ allowed: false }`, or `{ denied: true }` to block execution.
*/
toolGuard?: AgentToolGuard;

// Guardrails
inputGuardrails?: InputGuardrail[];
Expand Down