feat(py): require a live Operation on check and cancel

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 (
@@ -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:
@@ -1633,14 +1625,7 @@ async def check_operation(self, operation: Operation) -> Operation:
 
     async def cancel_operation(self, operation: Operation) -> 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(
@@ -1752,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.",

py/packages/genkit/src/genkit/_core/_background.py

@@ -19,12 +19,11 @@
 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 pydantic import BaseModel
-
 from genkit._core._action import Action, ActionKind, ActionRunContext
+from genkit._core._error import GenkitError
 from genkit._core._model import ModelRequest, ModelResponse
 from genkit._core._registry import Registry
 from genkit._core._schema import to_json_schema
@@ -119,7 +118,7 @@ async def start(
             An Operation with an ID to track the job.
         """
         result = await self.start_action.run(input)
-        return _ensure_operation(result.response)
+        return result.response
 
     async def check(self, operation: Operation) -> Operation:
         """Check the status of a background operation.
@@ -130,52 +129,30 @@ 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(result.response)
+        return result.response
 
     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.
         """
+        operation = require_operation(value=operation)
         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(result.response)
-
-
-def _ensure_operation(response: Any) -> Operation:  # noqa: ANN401
-    """Convert response to Operation type."""
-    if isinstance(response, Operation):
-        return response
-    if isinstance(response, dict):
-        return Operation.model_validate(response)
-    raise TypeError(f'Expected Operation, got {type(response)}')
-
-
-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
+        return result.response
 
 
 def define_background_model(
@@ -365,29 +342,89 @@ 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',
+            message='Provided operation is missing original request information',
+        )
+
+    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: {operation.action}',
+        )
+    return operation, background_action
+
+
 async def check_operation(
     registry: Registry,
     operation: Operation,
 ) -> 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: The poll handle.
 
     Returns:
         Updated Operation with current status.
 
     Raises:
-        ValueError: If operation is missing action or action not found.
+        GenkitError: If the handle 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,
+) -> Operation:
+    """Cancel a background operation.
+
+    Args:
+        registry: The registry to look up actions from.
+        operation: The poll handle.
+
+    Returns:
+        Updated Operation reflecting the cancel attempt.
+
+    Raises:
+        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:
+        raise GenkitError(
+            status='UNIMPLEMENTED',
+            message=f'Background action {resolved.action} does not support cancellation.',
+        )
+    return await background_action.cancel(resolved)

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

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

@@ -239,3 +239,21 @@ def test_model_response_eq_includes_operation() -> None:
     b = ModelResponse(operation=Operation(id='unique-b'))
     assert a != b
     assert a == ModelResponse(operation=Operation(id='unique-a'))
+
+
+@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)
+    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'
+    assert 'latencyMs' not in result.response.model_dump()