Skip to content

feat(hooks): add BeforeToolCallback to intercept tool calls before execution#533

Open
ousamabenyounes wants to merge 2 commits intomistralai:mainfrom
ousamabenyounes:feat/before-tool-callback
Open

feat(hooks): add BeforeToolCallback to intercept tool calls before execution#533
ousamabenyounes wants to merge 2 commits intomistralai:mainfrom
ousamabenyounes:feat/before-tool-callback

Conversation

@ousamabenyounes
Copy link
Copy Markdown

Summary

Adds a before_tool_callback hook to AgentLoop that lets callers intercept every tool invocation before execution and optionally rewrite the arguments.

Motivation: External tools like RTK (Rust Token Killer) need to transparently rewrite shell commands before they run (e.g. git statusrtk git status). Claude Code and Gemini CLI both provide equivalent hooks (PreToolUse and BeforeTool). This PR brings Vibe to parity. Tracked in rtk-ai/rtk#800.

API

async def before_bash_rewrite(
    tool_name: str, args: dict[str, Any]
) -> dict[str, Any] | None:
    if tool_name == "bash":
        cmd = args.get("command", "")
        return {**args, "command": rewrite_to_rtk(cmd)}
    return None  # passthrough for other tools

agent_loop.set_before_tool_callback(before_bash_rewrite)

The callback receives (tool_name, args_dict) and may:

  • Return None → args unchanged (passthrough)
  • Return a modified dict → the tool receives the rewritten args

Changes

File What changed
vibe/core/types.py New BeforeToolCallback type alias
vibe/core/tools/base.py before_tool_callback field in InvokeContext
vibe/core/agent_loop.py Callback invocation before invoke(), set_before_tool_callback() public method
tests/test_before_tool_callback.py 5 new tests

Design notes

  • Follows the existing set_approval_callback / set_user_input_callback pattern exactly
  • Zero impact on callers that don't set the callback (None check guards every path)
  • Applied after permission checks (tool is already approved), before args are passed to run()
  • Async to allow I/O in the callback (e.g. subprocess rewrite logic)

Test plan

  • test_callback_receives_tool_name_and_args — correct inputs to callback
  • test_callback_none_leaves_args_unchanged — passthrough behaviour
  • test_callback_can_rewrite_args — rewrite reaches the tool
  • test_no_callback_tool_executes_normally — no regression without callback
  • test_rewritten_args_produce_successful_tool_result — no errors with rewritten args
  • Full existing test suite: 345 passed, 0 regressions

🤖 Generated with Claude Code

Introduces a `before_tool_callback` mechanism on `AgentLoop` that lets
callers intercept every tool invocation before execution and optionally
rewrite the arguments.

The callback signature is:
  async def cb(tool_name: str, args: dict) -> dict | None

Returning `None` leaves the args unchanged; returning a modified dict
replaces the args passed to the tool. This enables transparent command
rewriting (e.g. `git status` → `rtk git status`) without patching any
tool internals.

Changes:
- `types.py`: add `BeforeToolCallback` type alias
- `tools/base.py`: add `before_tool_callback` field to `InvokeContext`
- `agent_loop.py`: apply callback before `tool_instance.invoke()`,
  expose `set_before_tool_callback()` public method (mirrors
  `set_approval_callback` pattern)

Tests: 5 new tests covering passthrough, rewriting, and error-free
execution with rewritten args.

Closes mistralai#531
- args_dict is now shallow-copied before being passed to the callback so
  the callback cannot mutate the frozen ResolvedToolCall internals
- callback return is merged with original args ({**original, **modified})
  instead of replacing them wholesale; partial returns (only the changed
  key) no longer silently drop the other fields
- exceptions raised by the callback are caught and logged as warnings;
  the tool call continues with the original args (silent fallback)
- two new tests: partial-return merge and callback-exception fallback
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant