feat(py): box background operation before wrap_model

py/packages/genkit-google-genai/tests/google_genai_plugin_test.py

@@ -50,7 +50,7 @@
 from genkit_google_genai.models.imagen import ImagenConfigSchema
 from genkit_google_genai.models.veo import VeoConfigSchema, VeoModel
 
-from genkit import ActionKind, GenkitError, Message, ModelRequest, Part, Role, TextPart
+from genkit import ActionKind, Genkit, GenkitError, Message, ModelRequest, Part, Role, TextPart
 from genkit.model import Operation
 from genkit.plugin_api import Action, to_json_schema
 
@@ -563,6 +563,21 @@ async def test_vertexai_resolve_veo_as_model_returns_none(mock_list_models: Magi
     assert action is None
 
 
+@patch('genkit_google_genai.google.genai.client.Client')
+@patch('genkit_google_genai.google._list_genai_models')
+@pytest.mark.asyncio
+async def test_resolve_model_finds_veo_as_background(mock_list_models: MagicMock, mock_client: MagicMock) -> None:
+    """resolve(MODEL, veo) is None so resolve_model can see the background start action."""
+    mock_list_models.return_value = GenaiModels()
+
+    ai = Genkit(plugins=[GoogleAI(api_key='test-key')])
+    action = await ai.registry.resolve_model('googleai/veo-3.0-generate-001')
+
+    assert action is not None
+    assert action.kind == ActionKind.BACKGROUND_MODEL
+    assert action.name == 'googleai/veo-3.0-generate-001'
+
+
 @patch('genkit_google_genai.google.genai.client.Client')
 @patch('genkit_google_genai.google._list_genai_models')
 @pytest.mark.asyncio

py/packages/genkit/src/genkit/_ai/_generate.py

@@ -72,6 +72,7 @@
     FinishReason,
     MiddlewareRef,
     MultipartToolResponse,
+    Operation,
     Part,
     Role,
     TextPart,
@@ -766,14 +767,19 @@ async def run_one_iteration(
             request = _augment_with_context(request)
 
         async def next_fn(params: ModelHookParams, c: GenerateMiddlewareContext) -> ModelResponse:
-            return (
+            raw = (
                 await model.run(
                     input=params.request,
                     context=c.custom_context,
                     on_chunk=c.on_chunk,
                     abort_signal=c.abort_signal,
                 )
             ).response
+            # start() returns a poll handle. wrap_model is documented as
+            # seeing a ModelResponse, so wrap before the hook runs.
+            if isinstance(raw, Operation):
+                return ModelResponse(operation=raw, request=params.request)
+            return raw
 
         with chunks.intercept_model_stream(ctx, role=Role.MODEL):
             model_response = await dispatch_model(
@@ -782,6 +788,21 @@ async def next_fn(params: ModelHookParams, c: GenerateMiddlewareContext) -> Mode
                 next_fn,
             )
 
+        # A background start is a poll handle, not a conversation turn.
+        # define_model LRO replies that also carry a message still go
+        # through persist and the tool loop.
+        if model.kind == ActionKind.BACKGROUND_MODEL:
+            if model_response.operation is None:
+                raise GenkitError(
+                    status='FAILED_PRECONDITION',
+                    message=(
+                        'wrap_model returned no operation for a background model; pass through response.operation'
+                    ),
+                )
+            if model_response.request is None:
+                model_response.request = request
+            return model_response
+
         def message_parser(msg: Message) -> Any:  # noqa: ANN401
             if formatter is None:
                 return None
@@ -891,7 +912,13 @@ def message_parser(msg: Message) -> Any:  # noqa: ANN401
         iteration=current_turn,
         message_index=chunks.message_index,
     )
-    return await dispatch_generate(generate_params, run_ctx, run_one_iteration)
+    response = await dispatch_generate(generate_params, run_ctx, run_one_iteration)
+    if model.kind == ActionKind.BACKGROUND_MODEL and response.operation is None:
+        raise GenkitError(
+            status='FAILED_PRECONDITION',
+            message=('wrap_generate returned no operation for a background model; pass through response.operation'),
+        )
+    return response
 
 
 def apply_format(

py/packages/genkit/src/genkit/_core/_model.py

@@ -517,9 +517,16 @@ def assert_valid_schema(self) -> None:
         pass
 
     def __eq__(self, other: object) -> bool:
-        """Compare responses by message and finish_reason."""
+        """Compare responses by message, finish_reason, and operation.
+
+        Two start handles with different job ids are not the same response.
+        """
         if isinstance(other, ModelResponse):
-            return self.message == other.message and self.finish_reason == other.finish_reason
+            return (
+                self.message == other.message
+                and self.finish_reason == other.finish_reason
+                and self.operation == other.operation
+            )
         return super().__eq__(other)
 
     def __hash__(self) -> int:

py/packages/genkit/tests/genkit/ai/background_generate_test.py

@@ -0,0 +1,241 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""generate() / generate_operation() against a define_background_model fake."""
+
+from collections.abc import Awaitable, Callable
+
+import pytest
+
+from genkit import Genkit, Message
+from genkit._core._action import ActionRunContext
+from genkit._core._background import BackgroundAction
+from genkit._core._error import GenkitError
+from genkit._core._middleware import BaseMiddleware, GenerateHookParams, GenerateMiddlewareContext, ModelHookParams
+from genkit._core._model import ModelRequest, ModelResponse
+from genkit._core._typing import (
+    FinishReason,
+    ModelInfo,
+    Operation,
+    Part,
+    Role,
+    Supports,
+    TextPart,
+    ToolRequest,
+    ToolRequestPart,
+)
+
+
+@pytest.fixture
+def ai() -> Genkit:
+    return Genkit()
+
+
+async def register_bg_model(ai: Genkit, *, op_id: str = 'bg-op-123') -> BackgroundAction:
+    async def start(_request: ModelRequest, _ctx: ActionRunContext) -> Operation:
+        return Operation(id=op_id, done=False)
+
+    async def check(op: Operation) -> Operation:
+        return op
+
+    return ai.define_background_model(
+        name='bg-model',
+        start=start,
+        check=check,
+    )
+
+
+@pytest.mark.asyncio
+async def test_generate_returns_operation_for_background_model(ai: Genkit) -> None:
+    """generate() wraps the start handle. message stays empty."""
+    await register_bg_model(ai)
+
+    response = await ai.generate(model='bg-model', prompt='a cat surfing')
+
+    assert response.operation is not None
+    assert response.operation.id == 'bg-op-123'
+    assert response.operation.done is False
+    assert response.operation.action == '/background-model/bg-model'
+    assert response.message is None
+
+
+@pytest.mark.asyncio
+async def test_generate_operation_with_background_model(ai: Genkit) -> None:
+    """generate_operation() returns that same handle."""
+    await register_bg_model(ai, op_id='bg-op-456')
+
+    operation = await ai.generate_operation(model='bg-model', prompt='a cat surfing')
+
+    assert isinstance(operation, Operation)
+    assert operation.id == 'bg-op-456'
+    assert operation.action == '/background-model/bg-model'
+
+
+@pytest.mark.asyncio
+async def test_generate_returns_the_job_without_polling(ai: Genkit) -> None:
+    """generate() hands back the job now. It does not wait until the job is done.
+
+    A background model (video, and anything registered with
+    ``define_background_model``) starts a job and returns a handle. You
+    poll later with ``check_operation``. ``generate()`` and
+    ``generate_operation()`` only start; they must not call ``check``
+    on the way out, or a long render would block the first call.
+    """
+    checks = 0
+
+    async def start(_request: ModelRequest, _ctx: ActionRunContext) -> Operation:
+        return Operation(id='bg-op-123', done=False)
+
+    async def check(op: Operation) -> Operation:
+        nonlocal checks
+        checks += 1
+        return Operation(id=op.id, done=True)
+
+    ai.define_background_model(name='bg-model', start=start, check=check)
+
+    response = await ai.generate(model='bg-model', prompt='a cat surfing')
+    operation = await ai.generate_operation(model='bg-model', prompt='a cat surfing')
+
+    assert response.operation is not None
+    assert response.operation.done is False
+    assert operation.done is False
+    assert checks == 0
+
+
+class ReadsMessage(BaseMiddleware):
+    async def wrap_model(
+        self,
+        params: ModelHookParams,
+        ctx: GenerateMiddlewareContext,
+        next_fn: Callable[[ModelHookParams, GenerateMiddlewareContext], Awaitable[ModelResponse]],
+    ) -> ModelResponse:
+        response = await next_fn(params, ctx)
+        _ = response.message
+        return response
+
+
+@pytest.mark.asyncio
+async def test_generate_boxes_before_wrap_model(ai: Genkit) -> None:
+    """wrap_model sees a ModelResponse, so reading .message does not crash."""
+    await register_bg_model(ai)
+
+    response = await ai.generate(model='bg-model', prompt='a cat surfing', use=[ReadsMessage()])
+
+    assert response.operation is not None
+    assert response.operation.id == 'bg-op-123'
+    assert response.message is None
+
+
+class DropsOperation(BaseMiddleware):
+    async def wrap_model(
+        self,
+        params: ModelHookParams,
+        ctx: GenerateMiddlewareContext,
+        next_fn: Callable[[ModelHookParams, GenerateMiddlewareContext], Awaitable[ModelResponse]],
+    ) -> ModelResponse:
+        resp = await next_fn(params, ctx)
+        return ModelResponse(message=resp.message, finish_reason=resp.finish_reason)
+
+
+@pytest.mark.asyncio
+async def test_generate_fails_when_wrap_model_drops_operation(ai: Genkit) -> None:
+    """A hook that rebuilds ModelResponse without operation must not blame the model."""
+    await register_bg_model(ai)
+
+    with pytest.raises(GenkitError, match='wrap_model returned no operation') as exc_info:
+        await ai.generate(model='bg-model', prompt='a cat surfing', use=[DropsOperation()])
+
+    assert exc_info.value.status == 'FAILED_PRECONDITION'
+
+    with pytest.raises(GenkitError, match='wrap_model returned no operation') as op_exc:
+        await ai.generate_operation(model='bg-model', prompt='a cat surfing', use=[DropsOperation()])
+
+    assert op_exc.value.status == 'FAILED_PRECONDITION'
+    assert 'did not return an operation' not in str(op_exc.value)
+
+
+class DropsGenerate(BaseMiddleware):
+    async def wrap_generate(
+        self,
+        params: GenerateHookParams,
+        ctx: GenerateMiddlewareContext,
+        next_fn: Callable[[GenerateHookParams, GenerateMiddlewareContext], Awaitable[ModelResponse]],
+    ) -> ModelResponse:
+        resp = await next_fn(params, ctx)
+        return ModelResponse(message=resp.message, finish_reason=resp.finish_reason)
+
+
+@pytest.mark.asyncio
+async def test_generate_fails_when_wrap_generate_drops_operation(ai: Genkit) -> None:
+    """wrap_generate can still rebuild the final response; dropping the handle must not blame the model."""
+    await register_bg_model(ai)
+
+    with pytest.raises(GenkitError, match='wrap_generate returned no operation') as exc_info:
+        await ai.generate(model='bg-model', prompt='a cat surfing', use=[DropsGenerate()])
+
+    assert exc_info.value.status == 'FAILED_PRECONDITION'
+    assert 'did not return an operation' not in str(exc_info.value)
+
+
+@pytest.mark.asyncio
+async def test_generate_on_lro_define_model_still_runs_tools(ai: Genkit) -> None:
+    """A define_model LRO that also returns a tool request still runs the tool loop."""
+    tool_ran = 0
+
+    @ai.tool(name='ping')
+    async def ping() -> str:
+        nonlocal tool_ran
+        tool_ran += 1
+        return 'pong'
+
+    turns = 0
+
+    async def model_fn(request: ModelRequest, ctx: ActionRunContext) -> ModelResponse:
+        nonlocal turns
+        turns += 1
+        if turns == 1:
+            return ModelResponse(
+                message=Message(
+                    role=Role.MODEL,
+                    content=[Part(root=ToolRequestPart(tool_request=ToolRequest(name='ping', input={}, ref='1')))],
+                ),
+                operation=Operation(id='lro-1', done=False),
+                finish_reason=FinishReason.STOP,
+            )
+        return ModelResponse(
+            message=Message(role=Role.MODEL, content=[Part(root=TextPart(text='done'))]),
+            finish_reason=FinishReason.STOP,
+        )
+
+    ai.define_model(
+        name='lro-model',
+        fn=model_fn,
+        info=ModelInfo(supports=Supports(long_running=True, tools=True)),
+    )
+
+    response = await ai.generate(model='lro-model', prompt='x', tools=['ping'])
+
+    assert tool_ran == 1
+    assert response.text == 'done'
+    assert response.operation is None
+
+
+def test_model_response_eq_includes_operation() -> None:
+    """Two start handles with different job ids are not the same response."""
+    a = ModelResponse(operation=Operation(id='unique-a'))
+    b = ModelResponse(operation=Operation(id='unique-b'))
+    assert a != b
+    assert a == ModelResponse(operation=Operation(id='unique-a'))