RestatePlugin: after_model_callback never invoked - before_model_callback returns LlmResponse causing ADK to skip post-model hooks

Open
#191 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Assessment

Difficulty
4/5
Estimated time
3-5 days
Newbie friendliness
48/100
Issue type
Bug
Clarity
Mostly clear
Activity status
Quiet
Tech stack
python
Domain
backend

Research direction

Start by tracing RestatePlugin.before_model_callback and the Google ADK handling of before_model_callback return values. Compare that flow with after_model_callback and the reported subclass workaround. Done means the post-model hook can receive the model response and usage metadata when additional plugins are registered, without breaking Restate replay semantics.

Written by the indexing model from the issue text.

Description

Bug Description

When registering additional plugins, alongside Restate (in our case it was for ADK-level observability), after_model_callback is never invoked. This makes it impossible to access the model response and usage metadata in post-call hook, which is the natural place to capture $ai_generation events.

Environment
  • restate-sdk (version 0.16.0)
  • google-adk (1.28.1)
  • RestatePlugin
Root Cause

The ADK framework uses the return value of before_model_callback as a signal:

  • None -> proceed with the real model call, then invoke after_model_callback
  • LlmResponse -> short-circuit: skip the model call and skip after_model_callback

RestatePlugin.before_model_callback wraps the LLM call inside a Restate durable step, which means it executes the model call itself and returns a LlmResponse. This is correct for Restate's replay semantics, but it has a side effect: after_model_callback is never reached, regardless of what subclasses implement there.

Workaround

Our current workaround is to subclass (decorate) the RestatePlugin itself, so we can simulate the after_model_callback trigger inside the wrapped before_model_callback call, but we believe this is worth tracking and resolving.

Dominant language
Python
Stars
81
Forks
22
Avg merge
2d 2h
Merged PRs (30d)
2

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from restatedev/sdk-python

All issues in restatedev/sdk-python

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.