From 04c66355251a0520c5d09075356aee3ca433daa6 Mon Sep 17 00:00:00 2001
From: Jeff Huang <huangjeff@google.com>
Date: Thu, 20 Aug 2026 13:34:28 -0500
Subject: [PATCH 1/5] feat(py): accept Operation or dump in check and cancel
 operation

---
 py/packages/genkit/src/genkit/_ai/_aio.py     |  15 +-
 .../genkit/src/genkit/_core/_background.py    | 147 +++++++++++++++---
 .../genkit/ai/background_generate_test.py     |  19 +++
 .../genkit/tests/genkit/ai/genkit_api_test.py | 146 ++++++++++++++++-
 4 files changed, 292 insertions(+), 35 deletions(-)

diff --git a/py/packages/genkit/src/genkit/_ai/_aio.py b/py/packages/genkit/src/genkit/_ai/_aio.py
index 84ddab6f5b..73eba3a7df 100644
--- a/py/packages/genkit/src/genkit/_ai/_aio.py
+++ b/py/packages/genkit/src/genkit/_ai/_aio.py
@@ -92,9 +92,9 @@
     CancelModelOpFn,
     CheckModelOpFn,
     StartModelOpFn,
+    cancel_operation,
     check_operation,
     define_background_model,
-    lookup_background_action,
 )
 from genkit._core._channel import Channel, run_loop
 from genkit._core._dap import (
@@ -1627,20 +1627,13 @@ async def run(
                 # the exception details.
                 raise
 
-    async def check_operation(self, operation: Operation) -> Operation:
+    async def check_operation(self, operation: Operation | Mapping[str, Any]) -> Operation:
         """Check the status of a long-running background operation."""
         return await check_operation(self.registry, operation)
 
-    async def cancel_operation(self, operation: Operation) -> Operation:
+    async def cancel_operation(self, operation: Operation | Mapping[str, Any]) -> Operation:
         """Cancel a long-running background operation."""
-        if not operation.action:
-            raise ValueError('Provided operation is missing original request information')
-
-        background_action = await lookup_background_action(self.registry, operation.action)
-        if background_action is None:
-            raise ValueError(f'Failed to resolve background action from original request: {operation.action}')
-
-        return await background_action.cancel(operation)
+        return await cancel_operation(self.registry, operation)
 
     @overload
     async def generate_operation(
diff --git a/py/packages/genkit/src/genkit/_core/_background.py b/py/packages/genkit/src/genkit/_core/_background.py
index 8df1f411d3..536f59e89b 100644
--- a/py/packages/genkit/src/genkit/_core/_background.py
+++ b/py/packages/genkit/src/genkit/_core/_background.py
@@ -19,10 +19,11 @@
 from __future__ import annotations
 
 import time
-from collections.abc import Awaitable, Callable
-from typing import Any, Generic, TypeVar
+from collections.abc import Awaitable, Callable, Mapping
+from typing import Any, ClassVar, Generic, TypeVar, cast
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict, ValidationError
+from pydantic.alias_generators import to_camel
 
 from genkit._core._action import Action, ActionKind, ActionRunContext
 from genkit._core._error import GenkitError
@@ -30,6 +31,7 @@
 from genkit._core._registry import Registry
 from genkit._core._schema import to_json_schema
 from genkit._core._typing import (
+    Error,
     ModelInfo,
     Operation,
 )
@@ -137,17 +139,20 @@ async def check(self, operation: Operation) -> Operation:
     async def cancel(self, operation: Operation) -> Operation:
         """Cancel a background operation.
 
-        If cancellation is not supported, returns the operation unchanged.
-
         Args:
             operation: The operation to cancel.
 
         Returns:
             Updated Operation reflecting cancellation attempt.
+
+        Raises:
+            GenkitError: If this action does not implement cancel.
         """
         if self.cancel_action is None:
-            # Return operation unchanged if cancel not supported
-            return operation
+            raise GenkitError(
+                status='UNIMPLEMENTED',
+                message=f'Background action {operation.action} does not support cancellation.',
+            )
         result = await self.cancel_action.run(operation)
         return _ensure_operation(response=result.response, name=self.cancel_action.name)
 
@@ -180,6 +185,20 @@ class DefineBackgroundModelOptions(BaseModel):
     config_schema: type | dict[str, Any] | None = None
 
 
+class OperationInput(Operation):
+    """Poll handle. Leftover dump keys like latencyMs are ignored."""
+
+    model_config: ClassVar[ConfigDict] = ConfigDict(
+        alias_generator=to_camel,
+        extra='ignore',
+        populate_by_name=True,
+        json_schema_extra={'title': 'Operation'},
+    )
+
+
+OperationInput.model_rebuild(_types_namespace={'Error': Error})
+
+
 def define_background_model(
     registry: Registry,
     name: str,
@@ -271,8 +290,8 @@ async def wrapped_start(request: ModelRequest, ctx: ActionRunContext) -> Operati
         return op
 
     # Wrap the check function (no ctx parameter)
-    async def wrapped_check(op: Operation, ctx: ActionRunContext) -> Operation:
-        updated = await check(op)
+    async def wrapped_check(op: OperationInput, ctx: ActionRunContext) -> Operation:
+        updated = await check(operation_from_handle(op))
         # Preserve action key
         updated.action = action_key
         return updated
@@ -301,8 +320,8 @@ async def wrapped_check(op: Operation, ctx: ActionRunContext) -> Operation:
         # Capture cancel in local scope for the nested function
         cancel_fn = cancel
 
-        async def wrapped_cancel(op: Operation, ctx: ActionRunContext) -> Operation:
-            cancelled = await cancel_fn(op)
+        async def wrapped_cancel(op: OperationInput, ctx: ActionRunContext) -> Operation:
+            cancelled = await cancel_fn(operation_from_handle(op))
             cancelled.action = action_key
             return cancelled
 
@@ -367,29 +386,113 @@ async def lookup_background_action(
     )
 
 
+def operation_from_handle(value: object) -> Operation:
+    """Read a persisted poll handle.
+
+    Callers save the Operation (or a dump of it) and pass it back. A
+    generate() ModelResponse is not a handle — pass response.operation.
+    Extra keys from a dump are ignored so a persist/reload does not 500.
+    """
+    if isinstance(value, ModelResponse):
+        raise GenkitError(
+            status='INVALID_ARGUMENT',
+            message='got ModelResponse; pass response.operation',
+        )
+    if isinstance(value, Operation):
+        return Operation.model_validate(value.model_dump())
+    if isinstance(value, Mapping):
+        mapping = cast('Mapping[str, object]', value)
+        nested = mapping.get('operation')
+        if isinstance(nested, Operation | Mapping):
+            raise GenkitError(
+                status='INVALID_ARGUMENT',
+                message="got a generate() envelope; pass the 'operation' field",
+            )
+        try:
+            known = {key: item for key, item in mapping.items() if is_operation_field(key)}
+            return Operation.model_validate(known)
+        except ValidationError as exc:
+            raise GenkitError(
+                status='INVALID_ARGUMENT',
+                message='Provided operation is not a valid Operation.',
+                cause=exc,
+            ) from exc
+    raise GenkitError(
+        status='INVALID_ARGUMENT',
+        message=f'got {type(value).__name__}, expected Operation | Mapping',
+    )
+
+
+def is_operation_field(key: str) -> bool:
+    fields = Operation.model_fields
+    if key in fields:
+        return True
+    return any(field.alias == key for field in fields.values())
+
+
+async def resolve_operation_action(
+    registry: Registry,
+    operation: Operation | Mapping[str, Any],
+) -> tuple[Operation, BackgroundAction]:
+    """Turn a poll handle into the background action that owns it."""
+    resolved = operation_from_handle(operation)
+    if not resolved.action:
+        raise GenkitError(
+            status='INVALID_ARGUMENT',
+            message='Provided operation is missing original request information',
+        )
+
+    background_action = await lookup_background_action(registry, resolved.action)
+    if background_action is None:
+        raise GenkitError(
+            status='INVALID_ARGUMENT',
+            message=f'Failed to resolve background action from original request: {resolved.action}',
+        )
+    return resolved, background_action
+
+
 async def check_operation(
     registry: Registry,
-    operation: Operation,
+    operation: Operation | Mapping[str, Any],
 ) -> Operation:
     """Check the status of a background operation.
 
-    Matches JS checkOperation from js/ai/src/check-operation.ts.
-
     Args:
         registry: The registry to look up actions from.
-        operation: The operation to check.
+        operation: A live Operation, or a dump of one.
 
     Returns:
         Updated Operation with current status.
 
     Raises:
-        ValueError: If operation is missing action or action not found.
+        GenkitError: If the handle cannot be read, is missing action, or
+            the action is not found.
     """
-    if not operation.action:
-        raise ValueError('Provided operation is missing original request information')
+    resolved, background_action = await resolve_operation_action(registry, operation)
+    return await background_action.check(resolved)
 
-    background_action = await lookup_background_action(registry, operation.action)
-    if background_action is None:
-        raise ValueError(f'Failed to resolve background action from original request: {operation.action}')
 
-    return await background_action.check(operation)
+async def cancel_operation(
+    registry: Registry,
+    operation: Operation | Mapping[str, Any],
+) -> Operation:
+    """Cancel a background operation.
+
+    Args:
+        registry: The registry to look up actions from.
+        operation: A live Operation, or a dump of one.
+
+    Returns:
+        Updated Operation reflecting the cancel attempt.
+
+    Raises:
+        GenkitError: If the handle cannot be read, is missing action, or
+            the action is not found.
+    """
+    resolved, background_action = await resolve_operation_action(registry, operation)
+    if not background_action.supports_cancel:
+        raise GenkitError(
+            status='UNIMPLEMENTED',
+            message=f'Background action {resolved.action} does not support cancellation.',
+        )
+    return await background_action.cancel(resolved)
diff --git a/py/packages/genkit/tests/genkit/ai/background_generate_test.py b/py/packages/genkit/tests/genkit/ai/background_generate_test.py
index 89492189c2..8993d85f92 100644
--- a/py/packages/genkit/tests/genkit/ai/background_generate_test.py
+++ b/py/packages/genkit/tests/genkit/ai/background_generate_test.py
@@ -222,6 +222,7 @@ async def check(op: Operation) -> Operation:
     assert response.operation is None
 
 
+<<<<<<< HEAD
 @pytest.mark.asyncio
 async def test_generate_persists_clean_history_without_injected_docs(ai: Genkit) -> None:
     """Injected RAG text stays off response.request.messages."""
@@ -373,3 +374,21 @@ def test_model_response_eq_uses_operation_id() -> None:
     c = ModelResponse(operation=Operation(id='unique-b'))
     assert a == b
     assert a != c
+
+
+@pytest.mark.asyncio
+async def test_check_action_accepts_dumped_operation_with_extra_keys(ai: Genkit) -> None:
+    """The Dev UI check action ignores leftover dump keys like latencyMs."""
+    action = await register_bg_model(ai)
+    dumped = {
+        'id': 'bg-op-123',
+        'done': False,
+        'action': '/background-model/bg-model',
+        'latencyMs': 42,
+    }
+
+    result = await action.check_action.run(dumped)
+
+    assert result.response.id == 'bg-op-123'
+    assert result.response.action == '/background-model/bg-model'
+
diff --git a/py/packages/genkit/tests/genkit/ai/genkit_api_test.py b/py/packages/genkit/tests/genkit/ai/genkit_api_test.py
index ee3ae2d601..005ccbabeb 100644
--- a/py/packages/genkit/tests/genkit/ai/genkit_api_test.py
+++ b/py/packages/genkit/tests/genkit/ai/genkit_api_test.py
@@ -12,6 +12,8 @@
 
 from genkit import Genkit
 from genkit._core._action import _action_context
+from genkit._core._error import GenkitError
+from genkit._core._model import ModelResponse
 from genkit._core._typing import Operation
 
 
@@ -67,8 +69,9 @@ async def test_genkit_check_operation_no_action() -> None:
     ai = Genkit()
     op = Operation(id='123', done=False)  # action is None
 
-    with pytest.raises(ValueError, match='Provided operation is missing original request information'):
+    with pytest.raises(GenkitError, match='Provided operation is missing original request information') as exc_info:
         await ai.check_operation(op)
+    assert exc_info.value.status == 'INVALID_ARGUMENT'
 
 
 @pytest.mark.asyncio
@@ -78,8 +81,147 @@ async def test_genkit_check_operation_not_found() -> None:
     op = Operation(id='123', done=False, action='missing')
     ai.registry.resolve_action_by_key = AsyncMock(return_value=None)  # type: ignore[assignment]
 
-    with pytest.raises(ValueError, match='Failed to resolve background action from original request: missing'):
+    with pytest.raises(
+        GenkitError, match='Failed to resolve background action from original request: missing'
+    ) as exc_info:
         await ai.check_operation(op)
+    assert exc_info.value.status == 'INVALID_ARGUMENT'
+
+
+@pytest.mark.asyncio
+async def test_check_operation_rejects_boxed_response() -> None:
+    """Poll the handle, not the generate() envelope. Pass response.operation."""
+    ai = Genkit()
+    op = Operation(id='123', done=False, action='/background-model/test_action')
+
+    with pytest.raises(GenkitError, match='got ModelResponse; pass response.operation') as exc_info:
+        await ai.check_operation(ModelResponse(operation=op))  # type: ignore[arg-type]
+    assert exc_info.value.status == 'INVALID_ARGUMENT'
+
+
+@pytest.mark.asyncio
+async def test_check_operation_rejects_empty_response() -> None:
+    ai = Genkit()
+    with pytest.raises(GenkitError, match='got ModelResponse; pass response.operation') as exc_info:
+        await ai.check_operation(ModelResponse())  # type: ignore[arg-type]
+    assert exc_info.value.status == 'INVALID_ARGUMENT'
+
+
+@pytest.mark.asyncio
+async def test_check_operation_accepts_dumped_operation() -> None:
+    """A persisted dump still polls, even with leftover keys like latencyMs."""
+    ai = Genkit()
+    dumped = {
+        'id': '123',
+        'done': False,
+        'action': '/background-model/test_action',
+        'latencyMs': 42,
+    }
+    mock_background_action = MagicMock()
+    mock_background_action.check = AsyncMock(return_value=Operation(id='123', done=True))
+
+    with mock.patch(
+        'genkit._core._background.lookup_background_action',
+        new=AsyncMock(return_value=mock_background_action),
+    ):
+        updated = await ai.check_operation(dumped)
+
+    assert updated.done is True
+
+
+@pytest.mark.asyncio
+async def test_check_operation_rejects_dumped_response() -> None:
+    """A dumped generate() envelope is not a handle."""
+    ai = Genkit()
+    dumped = {
+        'operation': {
+            'id': '123',
+            'done': False,
+            'action': '/background-model/test_action',
+        },
+        'finishReason': 'stop',
+    }
+
+    with pytest.raises(GenkitError, match="got a generate\\(\\) envelope; pass the 'operation' field") as exc_info:
+        await ai.check_operation(dumped)
+    assert exc_info.value.status == 'INVALID_ARGUMENT'
+
+
+@pytest.mark.asyncio
+async def test_check_operation_rejects_envelope_with_wrapper_id() -> None:
+    """A persist wrapper with its own id is still an envelope, not the handle."""
+    ai = Genkit()
+    dumped = {
+        'id': 'job-99',
+        'action': '/background-model/test_action',
+        'operation': {
+            'id': '123',
+            'done': False,
+            'action': '/background-model/test_action',
+        },
+    }
+
+    with pytest.raises(GenkitError, match="got a generate\\(\\) envelope; pass the 'operation' field") as exc_info:
+        await ai.check_operation(dumped)
+    assert exc_info.value.status == 'INVALID_ARGUMENT'
+
+
+@pytest.mark.asyncio
+async def test_check_operation_rejects_unreadable_handle() -> None:
+    ai = Genkit()
+    with pytest.raises(GenkitError, match='got str, expected Operation \\| Mapping') as exc_info:
+        await ai.check_operation('not-an-op')  # type: ignore[arg-type]
+    assert exc_info.value.status == 'INVALID_ARGUMENT'
+
+
+@pytest.mark.asyncio
+async def test_cancel_operation_accepts_dumped_operation() -> None:
+    """Cancel accepts the same persisted Operation dump as check."""
+    ai = Genkit()
+    dumped = {
+        'id': '123',
+        'done': False,
+        'action': '/background-model/test_action',
+        'latencyMs': 42,
+    }
+    mock_background_action = MagicMock()
+    mock_background_action.cancel = AsyncMock(return_value=Operation(id='123', done=True))
+
+    with mock.patch(
+        'genkit._core._background.lookup_background_action',
+        new=AsyncMock(return_value=mock_background_action),
+    ):
+        updated = await ai.cancel_operation(dumped)
+
+    assert updated.done is True
+
+
+@pytest.mark.asyncio
+async def test_cancel_operation_rejects_boxed_response() -> None:
+    ai = Genkit()
+    op = Operation(id='123', done=False, action='/background-model/test_action')
+
+    with pytest.raises(GenkitError, match='got ModelResponse; pass response.operation') as exc_info:
+        await ai.cancel_operation(ModelResponse(operation=op))  # type: ignore[arg-type]
+    assert exc_info.value.status == 'INVALID_ARGUMENT'
+
+
+@pytest.mark.asyncio
+async def test_cancel_operation_without_cancel_is_unimplemented() -> None:
+    ai = Genkit()
+    op = Operation(id='123', done=False, action='/background-model/test_action')
+    mock_background_action = MagicMock()
+    mock_background_action.supports_cancel = False
+    mock_background_action.cancel = AsyncMock()
+
+    with mock.patch(
+        'genkit._core._background.lookup_background_action',
+        new=AsyncMock(return_value=mock_background_action),
+    ):
+        with pytest.raises(GenkitError, match='does not support cancellation') as exc_info:
+            await ai.cancel_operation(op)
+    assert exc_info.value.status == 'UNIMPLEMENTED'
+    mock_background_action.cancel.assert_not_awaited()
 
 
 @pytest.mark.asyncio

From 136760c41c128610a639119df213164cbb084e0d Mon Sep 17 00:00:00 2001
From: Jeff Huang <huangjeff@google.com>
Date: Fri, 21 Aug 2026 10:02:30 -0500
Subject: [PATCH 2/5] simplify operation_from_handle

---
 .../genkit/src/genkit/_core/_background.py    | 75 +++++--------------
 .../genkit/src/genkit/_core/_typing.py        |  2 +-
 .../genkit/ai/background_generate_test.py     |  1 +
 .../genkit/tests/genkit/ai/genkit_api_test.py | 33 ++------
 py/scripts/schema_to_typing.py                | 14 +++-
 5 files changed, 39 insertions(+), 86 deletions(-)

diff --git a/py/packages/genkit/src/genkit/_core/_background.py b/py/packages/genkit/src/genkit/_core/_background.py
index 536f59e89b..4d4e36f288 100644
--- a/py/packages/genkit/src/genkit/_core/_background.py
+++ b/py/packages/genkit/src/genkit/_core/_background.py
@@ -20,10 +20,9 @@
 
 import time
 from collections.abc import Awaitable, Callable, Mapping
-from typing import Any, ClassVar, Generic, TypeVar, cast
+from typing import Any, Generic, TypeVar
 
-from pydantic import BaseModel, ConfigDict, ValidationError
-from pydantic.alias_generators import to_camel
+from pydantic import BaseModel, ValidationError
 
 from genkit._core._action import Action, ActionKind, ActionRunContext
 from genkit._core._error import GenkitError
@@ -31,7 +30,6 @@
 from genkit._core._registry import Registry
 from genkit._core._schema import to_json_schema
 from genkit._core._typing import (
-    Error,
     ModelInfo,
     Operation,
 )
@@ -185,20 +183,6 @@ class DefineBackgroundModelOptions(BaseModel):
     config_schema: type | dict[str, Any] | None = None
 
 
-class OperationInput(Operation):
-    """Poll handle. Leftover dump keys like latencyMs are ignored."""
-
-    model_config: ClassVar[ConfigDict] = ConfigDict(
-        alias_generator=to_camel,
-        extra='ignore',
-        populate_by_name=True,
-        json_schema_extra={'title': 'Operation'},
-    )
-
-
-OperationInput.model_rebuild(_types_namespace={'Error': Error})
-
-
 def define_background_model(
     registry: Registry,
     name: str,
@@ -290,8 +274,8 @@ async def wrapped_start(request: ModelRequest, ctx: ActionRunContext) -> Operati
         return op
 
     # Wrap the check function (no ctx parameter)
-    async def wrapped_check(op: OperationInput, ctx: ActionRunContext) -> Operation:
-        updated = await check(operation_from_handle(op))
+    async def wrapped_check(op: Operation, ctx: ActionRunContext) -> Operation:
+        updated = await check(op)
         # Preserve action key
         updated.action = action_key
         return updated
@@ -320,8 +304,8 @@ async def wrapped_check(op: OperationInput, ctx: ActionRunContext) -> Operation:
         # Capture cancel in local scope for the nested function
         cancel_fn = cancel
 
-        async def wrapped_cancel(op: OperationInput, ctx: ActionRunContext) -> Operation:
-            cancelled = await cancel_fn(operation_from_handle(op))
+        async def wrapped_cancel(op: Operation, ctx: ActionRunContext) -> Operation:
+            cancelled = await cancel_fn(op)
             cancelled.action = action_key
             return cancelled
 
@@ -389,45 +373,20 @@ async def lookup_background_action(
 def operation_from_handle(value: object) -> Operation:
     """Read a persisted poll handle.
 
-    Callers save the Operation (or a dump of it) and pass it back. A
-    generate() ModelResponse is not a handle — pass response.operation.
-    Extra keys from a dump are ignored so a persist/reload does not 500.
+    The Dev UI posts start output back here with leftover keys like
+    latencyMs that aren't on the Operation wire format. We ignore extras
+    instead of failing loudly: persist the handle as a dict, then
+    revalidate it on ai.check_operation() — leftover keys shouldn't 500
+    that path.
     """
-    if isinstance(value, ModelResponse):
+    try:
+        return Operation.model_validate(value)
+    except ValidationError as exc:
         raise GenkitError(
             status='INVALID_ARGUMENT',
-            message='got ModelResponse; pass response.operation',
-        )
-    if isinstance(value, Operation):
-        return Operation.model_validate(value.model_dump())
-    if isinstance(value, Mapping):
-        mapping = cast('Mapping[str, object]', value)
-        nested = mapping.get('operation')
-        if isinstance(nested, Operation | Mapping):
-            raise GenkitError(
-                status='INVALID_ARGUMENT',
-                message="got a generate() envelope; pass the 'operation' field",
-            )
-        try:
-            known = {key: item for key, item in mapping.items() if is_operation_field(key)}
-            return Operation.model_validate(known)
-        except ValidationError as exc:
-            raise GenkitError(
-                status='INVALID_ARGUMENT',
-                message='Provided operation is not a valid Operation.',
-                cause=exc,
-            ) from exc
-    raise GenkitError(
-        status='INVALID_ARGUMENT',
-        message=f'got {type(value).__name__}, expected Operation | Mapping',
-    )
-
-
-def is_operation_field(key: str) -> bool:
-    fields = Operation.model_fields
-    if key in fields:
-        return True
-    return any(field.alias == key for field in fields.values())
+            message='Provided operation is not a valid Operation.',
+            cause=exc,
+        ) from exc
 
 
 async def resolve_operation_action(
diff --git a/py/packages/genkit/src/genkit/_core/_typing.py b/py/packages/genkit/src/genkit/_core/_typing.py
index b283404ea5..b02d49ebd6 100644
--- a/py/packages/genkit/src/genkit/_core/_typing.py
+++ b/py/packages/genkit/src/genkit/_core/_typing.py
@@ -567,7 +567,7 @@ class MultipartToolResponse(GenkitModel):
 class Operation(GenkitModel):
     """Model for operation data."""
 
-    model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True)
+    model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='ignore', populate_by_name=True)
     action: str | None = None
     id: str = Field(...)
     done: bool | None = None
diff --git a/py/packages/genkit/tests/genkit/ai/background_generate_test.py b/py/packages/genkit/tests/genkit/ai/background_generate_test.py
index 8993d85f92..d4771bb22b 100644
--- a/py/packages/genkit/tests/genkit/ai/background_generate_test.py
+++ b/py/packages/genkit/tests/genkit/ai/background_generate_test.py
@@ -391,4 +391,5 @@ async def test_check_action_accepts_dumped_operation_with_extra_keys(ai: Genkit)
 
     assert result.response.id == 'bg-op-123'
     assert result.response.action == '/background-model/bg-model'
+    assert 'latencyMs' not in result.response.model_dump()
 
diff --git a/py/packages/genkit/tests/genkit/ai/genkit_api_test.py b/py/packages/genkit/tests/genkit/ai/genkit_api_test.py
index 005ccbabeb..dda7e1e8ef 100644
--- a/py/packages/genkit/tests/genkit/ai/genkit_api_test.py
+++ b/py/packages/genkit/tests/genkit/ai/genkit_api_test.py
@@ -90,11 +90,11 @@ async def test_genkit_check_operation_not_found() -> None:
 
 @pytest.mark.asyncio
 async def test_check_operation_rejects_boxed_response() -> None:
-    """Poll the handle, not the generate() envelope. Pass response.operation."""
+    """A generate() ModelResponse is not an Operation dump."""
     ai = Genkit()
     op = Operation(id='123', done=False, action='/background-model/test_action')
 
-    with pytest.raises(GenkitError, match='got ModelResponse; pass response.operation') as exc_info:
+    with pytest.raises(GenkitError, match='not a valid Operation') as exc_info:
         await ai.check_operation(ModelResponse(operation=op))  # type: ignore[arg-type]
     assert exc_info.value.status == 'INVALID_ARGUMENT'
 
@@ -102,7 +102,7 @@ async def test_check_operation_rejects_boxed_response() -> None:
 @pytest.mark.asyncio
 async def test_check_operation_rejects_empty_response() -> None:
     ai = Genkit()
-    with pytest.raises(GenkitError, match='got ModelResponse; pass response.operation') as exc_info:
+    with pytest.raises(GenkitError, match='not a valid Operation') as exc_info:
         await ai.check_operation(ModelResponse())  # type: ignore[arg-type]
     assert exc_info.value.status == 'INVALID_ARGUMENT'
 
@@ -131,7 +131,7 @@ async def test_check_operation_accepts_dumped_operation() -> None:
 
 @pytest.mark.asyncio
 async def test_check_operation_rejects_dumped_response() -> None:
-    """A dumped generate() envelope is not a handle."""
+    """A dumped generate() envelope is not an Operation (no id)."""
     ai = Genkit()
     dumped = {
         'operation': {
@@ -142,26 +142,7 @@ async def test_check_operation_rejects_dumped_response() -> None:
         'finishReason': 'stop',
     }
 
-    with pytest.raises(GenkitError, match="got a generate\\(\\) envelope; pass the 'operation' field") as exc_info:
-        await ai.check_operation(dumped)
-    assert exc_info.value.status == 'INVALID_ARGUMENT'
-
-
-@pytest.mark.asyncio
-async def test_check_operation_rejects_envelope_with_wrapper_id() -> None:
-    """A persist wrapper with its own id is still an envelope, not the handle."""
-    ai = Genkit()
-    dumped = {
-        'id': 'job-99',
-        'action': '/background-model/test_action',
-        'operation': {
-            'id': '123',
-            'done': False,
-            'action': '/background-model/test_action',
-        },
-    }
-
-    with pytest.raises(GenkitError, match="got a generate\\(\\) envelope; pass the 'operation' field") as exc_info:
+    with pytest.raises(GenkitError, match='not a valid Operation') as exc_info:
         await ai.check_operation(dumped)
     assert exc_info.value.status == 'INVALID_ARGUMENT'
 
@@ -169,7 +150,7 @@ async def test_check_operation_rejects_envelope_with_wrapper_id() -> None:
 @pytest.mark.asyncio
 async def test_check_operation_rejects_unreadable_handle() -> None:
     ai = Genkit()
-    with pytest.raises(GenkitError, match='got str, expected Operation \\| Mapping') as exc_info:
+    with pytest.raises(GenkitError, match='not a valid Operation') as exc_info:
         await ai.check_operation('not-an-op')  # type: ignore[arg-type]
     assert exc_info.value.status == 'INVALID_ARGUMENT'
 
@@ -201,7 +182,7 @@ async def test_cancel_operation_rejects_boxed_response() -> None:
     ai = Genkit()
     op = Operation(id='123', done=False, action='/background-model/test_action')
 
-    with pytest.raises(GenkitError, match='got ModelResponse; pass response.operation') as exc_info:
+    with pytest.raises(GenkitError, match='not a valid Operation') as exc_info:
         await ai.cancel_operation(ModelResponse(operation=op))  # type: ignore[arg-type]
     assert exc_info.value.status == 'INVALID_ARGUMENT'
 
diff --git a/py/scripts/schema_to_typing.py b/py/scripts/schema_to_typing.py
index 7bffe66dfb..19ede3c0fa 100644
--- a/py/scripts/schema_to_typing.py
+++ b/py/scripts/schema_to_typing.py
@@ -133,6 +133,17 @@ def _models_allowing_extra(schema: dict) -> set[str]:
     return result
 
 
+def _extra_policy(name: str, allow: set[str]) -> str:
+    if name in allow:
+        return 'allow'
+    # Dev UI posts start output (latencyMs) back into check. That key
+    # isn't on the Operation wire format; ignore extras so a persisted
+    # dict can revalidate on check_operation.
+    if name == 'Operation':
+        return 'ignore'
+    return 'forbid'
+
+
 def _typed_map_aliases(defs: dict) -> dict[str, str]:
     """Inline object schemas with typed scalar ``additionalProperties`` -> Python dict alias.
 
@@ -274,7 +285,8 @@ def _emit_model(
         req = req - omit - {_camel_to_snake(k) for k in omit}
     ext = ', protected_namespaces=()' if any(_camel_to_snake(k) in ('schema', 'schema_') for k in props) else ''
     frz = ', frozen=True' if name == 'PathMetadata' else ''
-    cfg = f"ConfigDict(alias_generator=to_camel, extra='{'allow' if name in allow else 'forbid'}', populate_by_name=True{ext}{frz})"
+    extra = _extra_policy(name, allow)
+    cfg = f"ConfigDict(alias_generator=to_camel, extra='{extra}', populate_by_name=True{ext}{frz})"
     lines = [
         f'class {name}(GenkitModel):',
         f'    """Model for {name.lower().replace("_", " ")} data."""',

From 462d4c1f117f4ac866d6e35ed19a0edba4b87022 Mon Sep 17 00:00:00 2001
From: Jeff Huang <huangjeff@google.com>
Date: Fri, 21 Aug 2026 10:37:38 -0500
Subject: [PATCH 3/5] remove the union types

---
 py/packages/genkit/src/genkit/_ai/_aio.py     | 21 ++----
 .../genkit/src/genkit/_core/_background.py    | 67 ++++---------------
 .../genkit/ai/background_generate_test.py     |  2 +-
 .../genkit/tests/genkit/ai/genkit_api_test.py | 60 +----------------
 py/scripts/schema_to_typing.py                |  5 +-
 5 files changed, 25 insertions(+), 130 deletions(-)

diff --git a/py/packages/genkit/src/genkit/_ai/_aio.py b/py/packages/genkit/src/genkit/_ai/_aio.py
index 73eba3a7df..fc74f8453b 100644
--- a/py/packages/genkit/src/genkit/_ai/_aio.py
+++ b/py/packages/genkit/src/genkit/_ai/_aio.py
@@ -150,18 +150,10 @@
 def _model_supports_long_running(model_action: Action) -> bool:
     """Check if a model action supports long-running operations."""
     model_info = model_action.metadata.get('model') if model_action.metadata else None
-    if not model_info:
+    if not isinstance(model_info, dict):
         return False
-    # Handle ModelInfo object
-    if hasattr(model_info, 'supports'):
-        supports = getattr(model_info, 'supports', None)
-        return bool(getattr(supports, 'long_running', False)) if supports else False
-    # Handle dict (cast needed because isinstance narrows too much for type checkers)
-    if isinstance(model_info, dict):
-        model_dict = cast(dict[str, Any], model_info)
-        supports = model_dict.get('supports')
-        return bool(supports.get('longRunning', False)) if isinstance(supports, dict) else False
-    return False
+    supports = cast(dict[str, Any], model_info).get('supports')
+    return bool(supports.get('longRunning', False)) if isinstance(supports, dict) else False
 
 
 class Genkit:
@@ -1627,11 +1619,11 @@ async def run(
                 # the exception details.
                 raise
 
-    async def check_operation(self, operation: Operation | Mapping[str, Any]) -> Operation:
+    async def check_operation(self, operation: Operation) -> Operation:
         """Check the status of a long-running background operation."""
         return await check_operation(self.registry, operation)
 
-    async def cancel_operation(self, operation: Operation | Mapping[str, Any]) -> Operation:
+    async def cancel_operation(self, operation: Operation) -> Operation:
         """Cancel a long-running background operation."""
         return await cancel_operation(self.registry, operation)
 
@@ -1745,8 +1737,7 @@ async def generate_operation(
             docs=docs,
         )
 
-        # Extract operation from response
-        if not hasattr(response, 'operation') or not response.operation:
+        if not response.operation:
             raise GenkitError(
                 status='FAILED_PRECONDITION',
                 message=f"Model '{model_action.name}' did not return an operation.",
diff --git a/py/packages/genkit/src/genkit/_core/_background.py b/py/packages/genkit/src/genkit/_core/_background.py
index 4d4e36f288..b43a297caa 100644
--- a/py/packages/genkit/src/genkit/_core/_background.py
+++ b/py/packages/genkit/src/genkit/_core/_background.py
@@ -19,11 +19,9 @@
 from __future__ import annotations
 
 import time
-from collections.abc import Awaitable, Callable, Mapping
+from collections.abc import Awaitable, Callable
 from typing import Any, Generic, TypeVar
 
-from pydantic import BaseModel, ValidationError
-
 from genkit._core._action import Action, ActionKind, ActionRunContext
 from genkit._core._error import GenkitError
 from genkit._core._model import ModelRequest, ModelResponse
@@ -165,23 +163,6 @@ def _ensure_operation(*, response: object, name: str) -> Operation:
     )
 
 
-class DefineBackgroundModelOptions(BaseModel):
-    """Options for defining a background model.
-
-    Attributes:
-        name: Unique name for this background model.
-        label: Human-readable label (defaults to name).
-        versions: Known version names for this model.
-        supports: Model capability information.
-        config_schema: Custom options schema for this model.
-    """
-
-    name: str
-    label: str | None = None
-    versions: list[str] | None = None
-    supports: dict[str, Any] | None = None
-    config_schema: type | dict[str, Any] | None = None
-
 
 def define_background_model(
     registry: Registry,
@@ -370,62 +351,42 @@ async def lookup_background_action(
     )
 
 
-def operation_from_handle(value: object) -> Operation:
-    """Read a persisted poll handle.
-
-    The Dev UI posts start output back here with leftover keys like
-    latencyMs that aren't on the Operation wire format. We ignore extras
-    instead of failing loudly: persist the handle as a dict, then
-    revalidate it on ai.check_operation() — leftover keys shouldn't 500
-    that path.
-    """
-    try:
-        return Operation.model_validate(value)
-    except ValidationError as exc:
-        raise GenkitError(
-            status='INVALID_ARGUMENT',
-            message='Provided operation is not a valid Operation.',
-            cause=exc,
-        ) from exc
-
-
 async def resolve_operation_action(
     registry: Registry,
-    operation: Operation | Mapping[str, Any],
+    operation: Operation,
 ) -> tuple[Operation, BackgroundAction]:
     """Turn a poll handle into the background action that owns it."""
-    resolved = operation_from_handle(operation)
-    if not resolved.action:
+    if not operation.action:
         raise GenkitError(
             status='INVALID_ARGUMENT',
             message='Provided operation is missing original request information',
         )
 
-    background_action = await lookup_background_action(registry, resolved.action)
+    background_action = await lookup_background_action(registry, operation.action)
     if background_action is None:
         raise GenkitError(
             status='INVALID_ARGUMENT',
-            message=f'Failed to resolve background action from original request: {resolved.action}',
+            message=f'Failed to resolve background action from original request: {operation.action}',
         )
-    return resolved, background_action
+    return operation, background_action
 
 
 async def check_operation(
     registry: Registry,
-    operation: Operation | Mapping[str, Any],
+    operation: Operation,
 ) -> Operation:
     """Check the status of a background operation.
 
     Args:
         registry: The registry to look up actions from.
-        operation: A live Operation, or a dump of one.
+        operation: The poll handle.
 
     Returns:
         Updated Operation with current status.
 
     Raises:
-        GenkitError: If the handle cannot be read, is missing action, or
-            the action is not found.
+        GenkitError: If the handle is missing action, or the action is
+            not found.
     """
     resolved, background_action = await resolve_operation_action(registry, operation)
     return await background_action.check(resolved)
@@ -433,20 +394,20 @@ async def check_operation(
 
 async def cancel_operation(
     registry: Registry,
-    operation: Operation | Mapping[str, Any],
+    operation: Operation,
 ) -> Operation:
     """Cancel a background operation.
 
     Args:
         registry: The registry to look up actions from.
-        operation: A live Operation, or a dump of one.
+        operation: The poll handle.
 
     Returns:
         Updated Operation reflecting the cancel attempt.
 
     Raises:
-        GenkitError: If the handle cannot be read, is missing action, or
-            the action is not found.
+        GenkitError: If the handle is missing action, the action is not
+            found, or cancel is not implemented.
     """
     resolved, background_action = await resolve_operation_action(registry, operation)
     if not background_action.supports_cancel:
diff --git a/py/packages/genkit/tests/genkit/ai/background_generate_test.py b/py/packages/genkit/tests/genkit/ai/background_generate_test.py
index d4771bb22b..0e8ecc8dcb 100644
--- a/py/packages/genkit/tests/genkit/ai/background_generate_test.py
+++ b/py/packages/genkit/tests/genkit/ai/background_generate_test.py
@@ -378,7 +378,7 @@ def test_model_response_eq_uses_operation_id() -> None:
 
 @pytest.mark.asyncio
 async def test_check_action_accepts_dumped_operation_with_extra_keys(ai: Genkit) -> None:
-    """The Dev UI check action ignores leftover dump keys like latencyMs."""
+    """A persisted dump still checks, even with leftover keys like latencyMs."""
     action = await register_bg_model(ai)
     dumped = {
         'id': 'bg-op-123',
diff --git a/py/packages/genkit/tests/genkit/ai/genkit_api_test.py b/py/packages/genkit/tests/genkit/ai/genkit_api_test.py
index dda7e1e8ef..f0c3a6cdda 100644
--- a/py/packages/genkit/tests/genkit/ai/genkit_api_test.py
+++ b/py/packages/genkit/tests/genkit/ai/genkit_api_test.py
@@ -13,7 +13,6 @@
 from genkit import Genkit
 from genkit._core._action import _action_context
 from genkit._core._error import GenkitError
-from genkit._core._model import ModelResponse
 from genkit._core._typing import Operation
 
 
@@ -88,25 +87,6 @@ async def test_genkit_check_operation_not_found() -> None:
     assert exc_info.value.status == 'INVALID_ARGUMENT'
 
 
-@pytest.mark.asyncio
-async def test_check_operation_rejects_boxed_response() -> None:
-    """A generate() ModelResponse is not an Operation dump."""
-    ai = Genkit()
-    op = Operation(id='123', done=False, action='/background-model/test_action')
-
-    with pytest.raises(GenkitError, match='not a valid Operation') as exc_info:
-        await ai.check_operation(ModelResponse(operation=op))  # type: ignore[arg-type]
-    assert exc_info.value.status == 'INVALID_ARGUMENT'
-
-
-@pytest.mark.asyncio
-async def test_check_operation_rejects_empty_response() -> None:
-    ai = Genkit()
-    with pytest.raises(GenkitError, match='not a valid Operation') as exc_info:
-        await ai.check_operation(ModelResponse())  # type: ignore[arg-type]
-    assert exc_info.value.status == 'INVALID_ARGUMENT'
-
-
 @pytest.mark.asyncio
 async def test_check_operation_accepts_dumped_operation() -> None:
     """A persisted dump still polls, even with leftover keys like latencyMs."""
@@ -124,37 +104,11 @@ async def test_check_operation_accepts_dumped_operation() -> None:
         'genkit._core._background.lookup_background_action',
         new=AsyncMock(return_value=mock_background_action),
     ):
-        updated = await ai.check_operation(dumped)
+        updated = await ai.check_operation(Operation.model_validate(dumped))
 
     assert updated.done is True
 
 
-@pytest.mark.asyncio
-async def test_check_operation_rejects_dumped_response() -> None:
-    """A dumped generate() envelope is not an Operation (no id)."""
-    ai = Genkit()
-    dumped = {
-        'operation': {
-            'id': '123',
-            'done': False,
-            'action': '/background-model/test_action',
-        },
-        'finishReason': 'stop',
-    }
-
-    with pytest.raises(GenkitError, match='not a valid Operation') as exc_info:
-        await ai.check_operation(dumped)
-    assert exc_info.value.status == 'INVALID_ARGUMENT'
-
-
-@pytest.mark.asyncio
-async def test_check_operation_rejects_unreadable_handle() -> None:
-    ai = Genkit()
-    with pytest.raises(GenkitError, match='not a valid Operation') as exc_info:
-        await ai.check_operation('not-an-op')  # type: ignore[arg-type]
-    assert exc_info.value.status == 'INVALID_ARGUMENT'
-
-
 @pytest.mark.asyncio
 async def test_cancel_operation_accepts_dumped_operation() -> None:
     """Cancel accepts the same persisted Operation dump as check."""
@@ -172,21 +126,11 @@ async def test_cancel_operation_accepts_dumped_operation() -> None:
         'genkit._core._background.lookup_background_action',
         new=AsyncMock(return_value=mock_background_action),
     ):
-        updated = await ai.cancel_operation(dumped)
+        updated = await ai.cancel_operation(Operation.model_validate(dumped))
 
     assert updated.done is True
 
 
-@pytest.mark.asyncio
-async def test_cancel_operation_rejects_boxed_response() -> None:
-    ai = Genkit()
-    op = Operation(id='123', done=False, action='/background-model/test_action')
-
-    with pytest.raises(GenkitError, match='not a valid Operation') as exc_info:
-        await ai.cancel_operation(ModelResponse(operation=op))  # type: ignore[arg-type]
-    assert exc_info.value.status == 'INVALID_ARGUMENT'
-
-
 @pytest.mark.asyncio
 async def test_cancel_operation_without_cancel_is_unimplemented() -> None:
     ai = Genkit()
diff --git a/py/scripts/schema_to_typing.py b/py/scripts/schema_to_typing.py
index 19ede3c0fa..ae2bea59e1 100644
--- a/py/scripts/schema_to_typing.py
+++ b/py/scripts/schema_to_typing.py
@@ -136,9 +136,8 @@ def _models_allowing_extra(schema: dict) -> set[str]:
 def _extra_policy(name: str, allow: set[str]) -> str:
     if name in allow:
         return 'allow'
-    # Dev UI posts start output (latencyMs) back into check. That key
-    # isn't on the Operation wire format; ignore extras so a persisted
-    # dict can revalidate on check_operation.
+    # A saved handle dump can carry leftover keys that aren't Operation
+    # fields. Drop them so Operation.model_validate(saved) does not 500.
     if name == 'Operation':
         return 'ignore'
     return 'forbid'

From 5a765a033d69a70655860753c3adc8f14b27c71b Mon Sep 17 00:00:00 2001
From: Jeff Huang <huangjeff@google.com>
Date: Fri, 21 Aug 2026 13:07:45 -0500
Subject: [PATCH 4/5] operation check returns genkit error

---
 .../genkit/src/genkit/_core/_background.py    | 23 +++++-
 .../genkit/tests/genkit/ai/genkit_api_test.py | 73 ++++++++++++++++++-
 2 files changed, 94 insertions(+), 2 deletions(-)

diff --git a/py/packages/genkit/src/genkit/_core/_background.py b/py/packages/genkit/src/genkit/_core/_background.py
index b43a297caa..512a0d2434 100644
--- a/py/packages/genkit/src/genkit/_core/_background.py
+++ b/py/packages/genkit/src/genkit/_core/_background.py
@@ -19,7 +19,7 @@
 from __future__ import annotations
 
 import time
-from collections.abc import Awaitable, Callable
+from collections.abc import Awaitable, Callable, Mapping
 from typing import Any, Generic, TypeVar
 
 from genkit._core._action import Action, ActionKind, ActionRunContext
@@ -351,11 +351,32 @@ async def lookup_background_action(
     )
 
 
+def require_operation(*, value: object) -> Operation:
+    """A poll handle is an Operation. A dump or generate() box is not."""
+    if isinstance(value, Operation):
+        return value
+    if isinstance(value, ModelResponse):
+        raise GenkitError(
+            status='INVALID_ARGUMENT',
+            message='got ModelResponse; pass response.operation',
+        )
+    if isinstance(value, Mapping):
+        raise GenkitError(
+            status='INVALID_ARGUMENT',
+            message='got a dump; pass Operation.model_validate(...)',
+        )
+    raise GenkitError(
+        status='INVALID_ARGUMENT',
+        message=f'got {type(value).__name__}, expected Operation',
+    )
+
+
 async def resolve_operation_action(
     registry: Registry,
     operation: Operation,
 ) -> tuple[Operation, BackgroundAction]:
     """Turn a poll handle into the background action that owns it."""
+    operation = require_operation(value=operation)
     if not operation.action:
         raise GenkitError(
             status='INVALID_ARGUMENT',
diff --git a/py/packages/genkit/tests/genkit/ai/genkit_api_test.py b/py/packages/genkit/tests/genkit/ai/genkit_api_test.py
index f0c3a6cdda..f0845cb465 100644
--- a/py/packages/genkit/tests/genkit/ai/genkit_api_test.py
+++ b/py/packages/genkit/tests/genkit/ai/genkit_api_test.py
@@ -11,8 +11,9 @@
 import pytest
 
 from genkit import Genkit
-from genkit._core._action import _action_context
+from genkit._core._action import ActionRunContext, _action_context
 from genkit._core._error import GenkitError
+from genkit._core._model import ModelRequest, ModelResponse
 from genkit._core._typing import Operation
 
 
@@ -109,6 +110,42 @@ async def test_check_operation_accepts_dumped_operation() -> None:
     assert updated.done is True
 
 
+@pytest.mark.asyncio
+async def test_check_operation_dump_is_invalid_argument() -> None:
+    """A saved dict is not an Operation until model_validate."""
+    ai = Genkit()
+    dumped = {
+        'id': '123',
+        'done': False,
+        'action': '/background-model/test_action',
+        'latencyMs': 42,
+    }
+
+    with pytest.raises(GenkitError, match='got a dump; pass Operation.model_validate') as exc_info:
+        await ai.check_operation(dumped)  # type: ignore[arg-type]
+    assert exc_info.value.status == 'INVALID_ARGUMENT'
+
+
+@pytest.mark.asyncio
+async def test_check_operation_boxed_response_is_invalid_argument() -> None:
+    """generate() returns a ModelResponse; the handle is response.operation."""
+    ai = Genkit()
+    boxed = ModelResponse(operation=Operation(id='123', action='/background-model/test_action'))
+
+    with pytest.raises(GenkitError, match='got ModelResponse; pass response.operation') as exc_info:
+        await ai.check_operation(boxed)  # type: ignore[arg-type]
+    assert exc_info.value.status == 'INVALID_ARGUMENT'
+
+
+@pytest.mark.asyncio
+async def test_check_operation_str_is_invalid_argument() -> None:
+    ai = Genkit()
+
+    with pytest.raises(GenkitError, match='got str, expected Operation') as exc_info:
+        await ai.check_operation('not-an-op')  # type: ignore[arg-type]
+    assert exc_info.value.status == 'INVALID_ARGUMENT'
+
+
 @pytest.mark.asyncio
 async def test_cancel_operation_accepts_dumped_operation() -> None:
     """Cancel accepts the same persisted Operation dump as check."""
@@ -131,6 +168,21 @@ async def test_cancel_operation_accepts_dumped_operation() -> None:
     assert updated.done is True
 
 
+@pytest.mark.asyncio
+async def test_cancel_operation_dump_is_invalid_argument() -> None:
+    ai = Genkit()
+    dumped = {
+        'id': '123',
+        'done': False,
+        'action': '/background-model/test_action',
+        'latencyMs': 42,
+    }
+
+    with pytest.raises(GenkitError, match='got a dump; pass Operation.model_validate') as exc_info:
+        await ai.cancel_operation(dumped)  # type: ignore[arg-type]
+    assert exc_info.value.status == 'INVALID_ARGUMENT'
+
+
 @pytest.mark.asyncio
 async def test_cancel_operation_without_cancel_is_unimplemented() -> None:
     ai = Genkit()
@@ -149,6 +201,25 @@ async def test_cancel_operation_without_cancel_is_unimplemented() -> None:
     mock_background_action.cancel.assert_not_awaited()
 
 
+@pytest.mark.asyncio
+async def test_background_action_cancel_without_fn_is_unimplemented() -> None:
+    """A real no-cancel BackgroundAction raises UNIMPLEMENTED from .cancel."""
+
+    async def start(_request: ModelRequest, _ctx: ActionRunContext) -> Operation:
+        return Operation(id='1', done=False)
+
+    async def check(op: Operation) -> Operation:
+        return op
+
+    ai = Genkit()
+    action = ai.define_background_model(name='no-cancel', start=start, check=check)
+    op = Operation(id='1', action='/background-model/no-cancel')
+
+    with pytest.raises(GenkitError, match='does not support cancellation') as exc_info:
+        await action.cancel(op)
+    assert exc_info.value.status == 'UNIMPLEMENTED'
+
+
 @pytest.mark.asyncio
 async def test_current_context() -> None:
     """Test Genkit.current_context method."""

From 6f8098260c14b3915f69b1ec41110886355c9b01 Mon Sep 17 00:00:00 2001
From: Jeff Huang <huangjeff@google.com>
Date: Fri, 21 Aug 2026 16:02:14 -0500
Subject: [PATCH 5/5] require operation

---
 .../genkit/src/genkit/_core/_background.py    |  3 +-
 .../genkit/ai/background_generate_test.py     | 18 ++++++--
 .../genkit/tests/genkit/ai/genkit_api_test.py | 46 +++++++++++++++++++
 3 files changed, 62 insertions(+), 5 deletions(-)

diff --git a/py/packages/genkit/src/genkit/_core/_background.py b/py/packages/genkit/src/genkit/_core/_background.py
index 512a0d2434..a3ebe7df64 100644
--- a/py/packages/genkit/src/genkit/_core/_background.py
+++ b/py/packages/genkit/src/genkit/_core/_background.py
@@ -129,6 +129,7 @@ async def check(self, operation: Operation) -> Operation:
         Returns:
             Updated Operation with current status.
         """
+        operation = require_operation(value=operation)
         result = await self.check_action.run(operation)
         return _ensure_operation(response=result.response, name=self.check_action.name)
 
@@ -144,6 +145,7 @@ async def cancel(self, operation: Operation) -> Operation:
         Raises:
             GenkitError: If this action does not implement cancel.
         """
+        operation = require_operation(value=operation)
         if self.cancel_action is None:
             raise GenkitError(
                 status='UNIMPLEMENTED',
@@ -163,7 +165,6 @@ def _ensure_operation(*, response: object, name: str) -> Operation:
     )
 
 
-
 def define_background_model(
     registry: Registry,
     name: str,
diff --git a/py/packages/genkit/tests/genkit/ai/background_generate_test.py b/py/packages/genkit/tests/genkit/ai/background_generate_test.py
index 0e8ecc8dcb..8b66264bc5 100644
--- a/py/packages/genkit/tests/genkit/ai/background_generate_test.py
+++ b/py/packages/genkit/tests/genkit/ai/background_generate_test.py
@@ -17,6 +17,7 @@
 """generate() / generate_operation() against a define_background_model fake."""
 
 from collections.abc import Awaitable, Callable
+from typing import Any, cast
 
 import pytest
 
@@ -222,7 +223,6 @@ async def check(op: Operation) -> Operation:
     assert response.operation is None
 
 
-<<<<<<< HEAD
 @pytest.mark.asyncio
 async def test_generate_persists_clean_history_without_injected_docs(ai: Genkit) -> None:
     """Injected RAG text stays off response.request.messages."""
@@ -330,7 +330,7 @@ async def test_define_model_returning_operation_raises(ai: Genkit) -> None:
     async def model_fn(_request: ModelRequest, _ctx: ActionRunContext) -> Operation:
         return Operation(id='sneaky', done=False)
 
-    ai.define_model(name='plain', fn=model_fn)
+    ai.define_model(name='plain', fn=cast(Any, model_fn))
 
     with pytest.raises(GenkitError, match='define_background_model') as exc_info:
         await ai.generate(model='plain', prompt='hi')
@@ -379,7 +379,18 @@ def test_model_response_eq_uses_operation_id() -> None:
 @pytest.mark.asyncio
 async def test_check_action_accepts_dumped_operation_with_extra_keys(ai: Genkit) -> None:
     """A persisted dump still checks, even with leftover keys like latencyMs."""
-    action = await register_bg_model(ai)
+
+    async def start(_request: ModelRequest, _ctx: ActionRunContext) -> Operation:
+        return Operation(id='bg-op-123', done=False)
+
+    async def check(op: Operation) -> Operation:
+        return op
+
+    action = ai.define_background_model(
+        name='bg-model',
+        start=start,
+        check=check,
+    )
     dumped = {
         'id': 'bg-op-123',
         'done': False,
@@ -392,4 +403,3 @@ async def test_check_action_accepts_dumped_operation_with_extra_keys(ai: Genkit)
     assert result.response.id == 'bg-op-123'
     assert result.response.action == '/background-model/bg-model'
     assert 'latencyMs' not in result.response.model_dump()
-
diff --git a/py/packages/genkit/tests/genkit/ai/genkit_api_test.py b/py/packages/genkit/tests/genkit/ai/genkit_api_test.py
index f0845cb465..897564171b 100644
--- a/py/packages/genkit/tests/genkit/ai/genkit_api_test.py
+++ b/py/packages/genkit/tests/genkit/ai/genkit_api_test.py
@@ -220,6 +220,52 @@ async def check(op: Operation) -> Operation:
     assert exc_info.value.status == 'UNIMPLEMENTED'
 
 
+@pytest.mark.asyncio
+async def test_background_action_check_rejects_non_operation() -> None:
+    """BackgroundAction.check uses the same require_operation gate as the veneer."""
+
+    async def start(_request: ModelRequest, _ctx: ActionRunContext) -> Operation:
+        return Operation(id='1', done=False)
+
+    async def check(op: Operation) -> Operation:
+        return op
+
+    ai = Genkit()
+    action = ai.define_background_model(name='bg-check', start=start, check=check)
+    dumped = {'id': '1', 'action': '/background-model/bg-check'}
+    boxed = ModelResponse(operation=Operation(id='1', action='/background-model/bg-check'))
+
+    with pytest.raises(GenkitError, match='got a dump; pass Operation.model_validate') as dump_exc:
+        await action.check(dumped)  # type: ignore[arg-type]
+    assert dump_exc.value.status == 'INVALID_ARGUMENT'
+
+    with pytest.raises(GenkitError, match='got ModelResponse; pass response.operation') as box_exc:
+        await action.check(boxed)  # type: ignore[arg-type]
+    assert box_exc.value.status == 'INVALID_ARGUMENT'
+
+    with pytest.raises(GenkitError, match='got str, expected Operation') as str_exc:
+        await action.check('not-an-op')  # type: ignore[arg-type]
+    assert str_exc.value.status == 'INVALID_ARGUMENT'
+
+
+@pytest.mark.asyncio
+async def test_background_action_cancel_rejects_non_operation() -> None:
+    """A dump must not AttributeError on .action before UNIMPLEMENTED."""
+
+    async def start(_request: ModelRequest, _ctx: ActionRunContext) -> Operation:
+        return Operation(id='1', done=False)
+
+    async def check(op: Operation) -> Operation:
+        return op
+
+    ai = Genkit()
+    action = ai.define_background_model(name='no-cancel', start=start, check=check)
+
+    with pytest.raises(GenkitError, match='got a dump; pass Operation.model_validate') as exc_info:
+        await action.cancel({'id': '1', 'action': '/background-model/no-cancel'})  # type: ignore[arg-type]
+    assert exc_info.value.status == 'INVALID_ARGUMENT'
+
+
 @pytest.mark.asyncio
 async def test_current_context() -> None:
     """Test Genkit.current_context method."""