google / google/adk-python

Add around-call plugin hooks for durable model/tool execution

オープン
#6,224 コメント 1 件 リアクション 0 件 担当者 1 名 @sanketpatil06 が担当を希望しています GitHub で見る
core needs review
主要言語
Python
スター
21.5k
フォーク
4k
平均マージ
1日 14時間
マージ済み PR(30日)
37

説明

## Use case

We are [adding Google ADK support](https://github.com/zenml-io/kitaru/pull/480) to [Kitaru](https://zenml.io/kitaru), an OSS durable execution/checkpointing runtime for AI workflows.

For ADK, we can already support two useful paths:

1. Whole-runner checkpointing: wrap an ADK runner invocation as one durable operation.
2. Explicit model/tool checkpointing: users pass wrapped ADK objects such as `KitaruADKModel` and `KitaruADKTool`, and ADK calls those wrapped objects directly.

The second path works because the actual model/tool operation happens inside the wrapper method:

```text
ADK runner
-> ADK LlmAgent
-> KitaruADKModel.generate_content_async(...)
-> Kitaru checkpoint
-> real model.generate_content_async(...)

ADK runner
-> ADK tool dispatch
-> KitaruADKTool.run_async(...)
-> Kitaru checkpoint
-> real tool.run_async(...)
```

That gives replay-safe behavior: if a process crashes after a model/tool call completes, Kitaru can replay from the stored checkpoint instead of repeating the provider call or tool side effect.

## What is missing

For arbitrary ADK runners, we do not see a public plugin hook that lets a plugin wrap the actual model/tool call body.

The current plugin/callback shape appears to be before/after/error:

```text
before_model_callback(callback_context, llm_request)
ADK calls the model
after_model_callback(callback_context, llm_response)
```

and similarly for tools:

```text
before_tool_callback(tool, tool_args, tool_context)
ADK runs the tool
after_tool_callback(tool, tool_args, tool_context, result)
```

Those hooks are useful for logging, guardrails, short-circuiting, and response modification. But they do not let a durability runtime run the actual provider/tool call inside its own checkpoint.

The failure case is:

```text
1. before_model_callback runs.
2. It returns None, so ADK calls the provider.
3. The provider returns successfully.
4. The process crashes before after_model_callback persists the response.
5. Replay calls the provider again.
```

For a durability system, that means duplicate model cost or duplicated tool side effects.

Returning a cached response from `before_model_callback` helps when the cache already exists, but it does not solve the first successful uncached call unless the call itself can happen inside the durable checkpoint.

## Requested API shape

Would ADK consider adding middleware-style around-call plugin hooks that receive both the call input and a continuation/proceed callable?

For example, model calls:

```python
async def wrap_model_call(
self,
*,
callback_context: CallbackContext,
llm_request: LlmRequest,
proceed: Callable[[LlmRequest], Awaitable[LlmResponse]],
) -> LlmResponse:
...
```

Tool calls:

```python
async def wrap_tool_call(
self,
*,
tool: BaseTool,
tool_args: dict[str, Any],
tool_context: ToolContext,
proceed: Callable[[], Awaitable[dict]],
) -> dict:
...
```

Then a durability plugin can do this:

```text
open checkpoint
-> call proceed(...)
-> persist returned response/result as checkpoint output
return response/result to ADK
```

This differs from before/after callbacks because the plugin controls when the side effect happens and can place that side effect inside its own durable boundary.

## Why this matters

This would enable replay-safe integrations for:

- durable model-call checkpointing
- durable tool-call checkpointing
- exact retry/replay behavior after process crashes
- provider-cost deduplication after successful calls
- avoiding duplicated tool side effects when tools are idempotent only at the runtime/checkpoint layer

It would also let ADK support durable execution systems without requiring users to manually wrap every model/tool object before constructing their agents.

## Current workaround

The current safe workaround is explicit wrapping:

```text
user passes wrapped model/tool objects into ADK
ADK calls the wrappers
wrappers checkpoint their own method bodies
```

That works, but it is opt-in per model/tool and does not give a runner-level integration for arbitrary ADK runners.

We intentionally do not treat before/after metadata callbacks as replay-safe checkpointing, because that would make users think provider/tool side effects are protected when they may still be repeated on replay.

## Questions

1. Is there already a public around-call hook or continuation-style plugin API that we missed?
2. If not, would the ADK team be open to a `wrap_model_call` / `wrap_tool_call` style plugin API?
3. Are there constraints in ADK's runner/model/tool execution model that would make this unsafe or hard to support?

コントリビューションガイド

コントリビューションガイドを開く

評価

この issue はまだ評価されていません。

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。