Skip to content

feat: add streaming chunk callback support to generation middleware - #171

Merged
xavidop merged 1 commit into
xavier/middleware-v2from
pj/middleware-stream-intercept
May 8, 2026
Merged

feat: add streaming chunk callback support to generation middleware#171
xavidop merged 1 commit into
xavier/middleware-v2from
pj/middleware-stream-intercept

Conversation

@pavelgj

@pavelgj pavelgj commented May 8, 2026

Copy link
Copy Markdown
Member

Add onChunk callback to GenerateParams and withStreamCallback to ModelParams, enabling middleware to observe, transform, or replace streaming chunks. This mirrors the JS SDK's onChunk middleware hook.

  • Thread onChunk through GenerateParams and into ModelParams
  • Add withStreamCallback to ModelParams for middleware wrapping
  • Wire streaming callback through GenerateAction tool loop
  • Add StreamingMiddlewareTest covering intercept, transform, and multi-middleware streaming scenarios

Add `onChunk` callback to `GenerateParams` and `withStreamCallback` to
`ModelParams`, enabling middleware to observe, transform, or replace
streaming chunks. This mirrors the JS SDK's `onChunk` middleware hook.

- Thread `onChunk` through `GenerateParams` and into `ModelParams`
- Add `withStreamCallback` to `ModelParams` for middleware wrapping
- Wire streaming callback through `GenerateAction` tool loop
- Add `StreamingMiddlewareTest` covering intercept, transform, and
  multi-middleware streaming scenarios

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

Code Review

This pull request refactors the streaming implementation to propagate streaming callbacks through the entire middleware chain, enabling middleware to observe or transform response chunks. Key changes include adding onChunk to GenerateParams, streamCallback to ModelParams, and refactoring Genkit.generateStream to delegate to a unified generateInternal method. Comprehensive tests demonstrate middleware capabilities such as re-chunking and content filtering. Feedback suggests moving the supportsStreaming() validation deeper into the call chain to correctly handle middleware-modified model names and eliminate redundant checks.

}

final String spanPath = "/generate/" + modelName;
final java.util.function.Consumer<ModelResponseChunk> sc = mparams.getStreamCallback();

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

medium

Adding the supportsStreaming() check here ensures that it is performed on the final resolved model, after any middleware has had a chance to modify the model name. This is more robust than checking it at the entry point of generateStream (lines 1587-1590), which is now redundant and can be removed.

          final java.util.function.Consumer<ModelResponseChunk> sc = mparams.getStreamCallback();
          if (sc != null && !model.supportsStreaming()) {
            throw new GenkitException("Model " + modelName + " does not support streaming");
          }

// The callback flows through the full middleware chain:
// generateInternal → GenerateParams.onChunk → wrapGenerate hooks
// → rawGenerate → ModelParams.streamCallback → wrapModel hooks → model.run(ctx, req, cb)
return generateInternal(options, streamCallback);

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

medium

The supportsStreaming() check at the beginning of this method (lines 1587-1590) is now redundant if the validation is moved into the model call chain (see suggestion in buildWrappedModelCall). Removing it also avoids a redundant model lookup and ensures that middleware-modified model names are correctly validated.

@xavidop
xavidop merged commit e1316d0 into xavier/middleware-v2 May 8, 2026
4 checks passed
@xavidop
xavidop deleted the pj/middleware-stream-intercept branch May 18, 2026 16:45
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.

2 participants