prefect-client 3.1.12__tar.gz → 3.1.14__tar.gz

{prefect-client-3.1.12/src/prefect_client.egg-info → prefect-client-3.1.14}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: prefect-client
-Version: 3.1.12
+Version: 3.1.14
 Summary: Workflow orchestration and management.
 Home-page: https://www.prefect.io
 Author: Prefect Technologies, Inc.

{prefect-client-3.1.12 → prefect-client-3.1.14}/requirements-client.txt

@@ -20,7 +20,7 @@ packaging >= 21.3, < 24.3
 pathspec >= 0.8.0
 pendulum >= 3.0.0, <4
 prometheus-client >= 0.20.0
-pydantic >= 2.7, < 3.0.0, != 2.10.0
+pydantic >= 2.9, < 3.0.0, != 2.10.0
 pydantic_core >= 2.12.0, < 3.0.0
 pydantic_extra_types >= 2.8.2, < 3.0.0
 pydantic_settings > 2.2.1

{prefect-client-3.1.12 → prefect-client-3.1.14}/src/prefect/_experimental/lineage.py

@@ -178,3 +178,66 @@ async def emit_result_write_event(
             downstream_resources=downstream_resources,
             direction_of_run_from_event="upstream",
         )
+
+
+async def emit_external_resource_lineage(
+    upstream_resources: Optional[UpstreamResources] = None,
+    downstream_resources: Optional[DownstreamResources] = None,
+) -> None:
+    """Emit lineage events connecting external resources to Prefect context resources.
+
+    This function emits events that place the current Prefect context resources
+    (like flow runs, task runs) as:
+    1. Downstream of any provided upstream external resources
+    2. Upstream of any provided downstream external resources
+
+    Args:
+        upstream_resources: Optional sequence of resources that are upstream of the
+            current Prefect context
+        downstream_resources: Optional sequence of resources that are downstream of
+            the current Prefect context
+    """
+    from prefect.client.orchestration import get_client
+
+    if not get_current_settings().experiments.lineage_events_enabled:
+        return
+
+    upstream_resources = list(upstream_resources) if upstream_resources else []
+    downstream_resources = list(downstream_resources) if downstream_resources else []
+
+    # Get the current Prefect context resources (flow runs, task runs, etc.)
+    async with get_client() as client:
+        context_resources = await related_resources_from_run_context(client)
+
+    # Add lineage group label to all resources
+    for res in upstream_resources + downstream_resources:
+        if "prefect.resource.lineage-group" not in res:
+            res["prefect.resource.lineage-group"] = "global"
+
+    # For each context resource, emit an event showing it as downstream of upstream resources
+    if upstream_resources:
+        for context_resource in context_resources:
+            emit_kwargs: Dict[str, Any] = {
+                "event": "prefect.resource.consumed",
+                "resource": context_resource,
+                "related": upstream_resources,
+            }
+            emit_event(**emit_kwargs)
+
+    # For each downstream resource, emit an event showing it as downstream of context resources
+    for downstream_resource in downstream_resources:
+        emit_kwargs: Dict[str, Any] = {
+            "event": "prefect.resource.produced",
+            "resource": downstream_resource,
+            "related": context_resources,
+        }
+        emit_event(**emit_kwargs)
+
+        # For each downstream resource, emit an event showing it as downstream of upstream resources
+        if upstream_resources:
+            direct_emit_kwargs = {
+                "event": "prefect.resource.direct-lineage",
+                "resource": downstream_resource,
+                "related": upstream_resources,
+            }
+            emit_event(**direct_emit_kwargs)

prefect-client-3.1.14/src/prefect/_experimental/sla/client.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from prefect.client.orchestration.base import BaseAsyncClient, BaseClient
+
+if TYPE_CHECKING:
+    from uuid import UUID
+
+    from prefect._experimental.sla.objects import SlaMergeResponse, SlaTypes
+
+
+class SlaClient(BaseClient):
+    def apply_slas_for_deployment(
+        self, deployment_id: "UUID", slas: "list[SlaTypes]"
+    ) -> "SlaMergeResponse":
+        """
+        Applies service level agreements for a deployment. Performs matching by SLA name. If a SLA with the same name already exists, it will be updated. If a SLA with the same name does not exist, it will be created. Existing SLAs that are not in the list will be deleted.
+        Args:
+            deployment_id: The ID of the deployment to update SLAs for
+            slas: List of SLAs to associate with the deployment
+        Raises:
+            httpx.RequestError: if the SLAs were not updated for any reason
+        Returns:
+            SlaMergeResponse: The response from the backend, containing the names of the created, updated, and deleted SLAs
+        """
+        resource_id = f"prefect.deployment.{deployment_id}"
+
+        for sla in slas:
+            sla.set_deployment_id(deployment_id)
+
+        slas_spec_list = [
+            sla.model_dump(mode="json", exclude_unset=True) for sla in slas
+        ]
+
+        response = self.request(
+            "POST",
+            f"/slas/apply-resource-slas/{resource_id}",
+            json=slas_spec_list,
+        )
+        response.raise_for_status()
+
+        response_json = response.json()
+
+        from prefect._experimental.sla.objects import SlaMergeResponse
+
+        return SlaMergeResponse(
+            created=[sla.get("name") for sla in response_json.get("created")],
+            updated=[sla.get("name") for sla in response_json.get("updated")],
+            deleted=[sla.get("name") for sla in response_json.get("deleted")],
+        )
+
+
+class SlaAsyncClient(BaseAsyncClient):
+    async def apply_slas_for_deployment(
+        self, deployment_id: "UUID", slas: "list[SlaTypes]"
+    ) -> "UUID":
+        """
+        Applies service level agreements for a deployment. Performs matching by SLA name. If a SLA with the same name already exists, it will be updated. If a SLA with the same name does not exist, it will be created. Existing SLAs that are not in the list will be deleted.
+        Args:
+            deployment_id: The ID of the deployment to update SLAs for
+            slas: List of SLAs to associate with the deployment
+        Raises:
+            httpx.RequestError: if the SLAs were not updated for any reason
+        Returns:
+            SlaMergeResponse: The response from the backend, containing the names of the created, updated, and deleted SLAs
+        """
+        resource_id = f"prefect.deployment.{deployment_id}"
+
+        for sla in slas:
+            sla.set_deployment_id(deployment_id)
+
+        slas_spec_list = [
+            sla.model_dump(mode="json", exclude_unset=True) for sla in slas
+        ]
+
+        response = await self.request(
+            "POST",
+            f"/slas/apply-resource-slas/{resource_id}",
+            json=slas_spec_list,
+        )
+        response.raise_for_status()
+
+        response_json = response.json()
+
+        from prefect._experimental.sla.objects import SlaMergeResponse
+
+        return SlaMergeResponse(
+            created=[sla.get("name") for sla in response_json.get("created")],
+            updated=[sla.get("name") for sla in response_json.get("updated")],
+            deleted=[sla.get("name") for sla in response_json.get("deleted")],
+        )

{prefect-client-3.1.12 → prefect-client-3.1.14}/src/prefect/_experimental/sla/objects.py

@@ -5,7 +5,7 @@ from typing import Literal, Optional, Union
 from uuid import UUID
 
 from pydantic import Field, PrivateAttr, computed_field
-from typing_extensions import TypeAlias
+from typing_extensions import Self, TypeAlias
 
 from prefect._internal.schemas.bases import PrefectBaseModel
 
@@ -28,7 +28,7 @@ class ServiceLevelAgreement(PrefectBaseModel, abc.ABC):
         description="Whether the SLA is enabled.",
     )
 
-    def set_deployment_id(self, deployment_id: UUID):
+    def set_deployment_id(self, deployment_id: UUID) -> Self:
         self._deployment_id = deployment_id
         return self
 
@@ -49,5 +49,13 @@ class TimeToCompletionSla(ServiceLevelAgreement):
     )
 
 
+class SlaMergeResponse(PrefectBaseModel):
+    """A response object for the apply_slas_for_deployment method. Contains the names of the created, updated, and deleted SLAs."""
+
+    created: list[str]
+    updated: list[str]
+    deleted: list[str]
+
+
 # Concrete SLA types
 SlaTypes: TypeAlias = Union[TimeToCompletionSla]

{prefect-client-3.1.12 → prefect-client-3.1.14}/src/prefect/_internal/concurrency/services.py

@@ -32,7 +32,7 @@ class _QueueServiceBase(abc.ABC, Generic[T]):
         self._task: Optional[asyncio.Task[None]] = None
         self._stopped: bool = False
         self._started: bool = False
-        self._key = hash(args)
+        self._key = hash((self.__class__, *args))
         self._lock = threading.Lock()
         self._queue_get_thread = WorkerThread(
             # TODO: This thread should not need to be a daemon but when it is not, it
@@ -256,7 +256,7 @@ class _QueueServiceBase(abc.ABC, Generic[T]):
         If an instance already exists with the given arguments, it will be returned.
         """
         with cls._instance_lock:
-            key = hash(args)
+            key = hash((cls, *args))
             if key not in cls._instances:
                 cls._instances[key] = cls._new_instance(*args)
 

{prefect-client-3.1.12 → prefect-client-3.1.14}/src/prefect/_internal/concurrency/threads.py

@@ -2,6 +2,8 @@
 Utilities for managing worker threads.
 """
 
+from __future__ import annotations
+
 import asyncio
 import atexit
 import concurrent.futures
@@ -197,6 +199,10 @@ class EventLoopThread(Portal):
     def running(self) -> bool:
         return not self._shutdown_event.is_set()
 
+    @property
+    def loop(self) -> asyncio.AbstractEventLoop | None:
+        return self._loop
+
     def _entrypoint(self):
         """
         Entrypoint for the thread.

{prefect-client-3.1.12 → prefect-client-3.1.14}/src/prefect/_internal/retries.py

@@ -29,7 +29,7 @@ def retry_async_fn(
     retry_on_exceptions: tuple[type[Exception], ...] = (Exception,),
     operation_name: Optional[str] = None,
 ) -> Callable[
-    [Callable[P, Coroutine[Any, Any, R]]], Callable[P, Coroutine[Any, Any, Optional[R]]]
+    [Callable[P, Coroutine[Any, Any, R]]], Callable[P, Coroutine[Any, Any, R]]
 ]:
     """A decorator for retrying an async function.
 
@@ -48,9 +48,9 @@ def retry_async_fn(
 
     def decorator(
         func: Callable[P, Coroutine[Any, Any, R]],
-    ) -> Callable[P, Coroutine[Any, Any, Optional[R]]]:
+    ) -> Callable[P, Coroutine[Any, Any, R]]:
         @wraps(func)
-        async def wrapper(*args: P.args, **kwargs: P.kwargs) -> Optional[R]:
+        async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
             name = operation_name or func.__name__
             for attempt in range(max_attempts):
                 try:
@@ -67,6 +67,9 @@ def retry_async_fn(
                         f"Retrying in {delay:.2f} seconds..."
                     )
                     await asyncio.sleep(delay)
+            # Technically unreachable, but this raise helps pyright know that this function
+            # won't return None.
+            raise Exception(f"Function {name!r} failed after {max_attempts} attempts")
 
         return wrapper
 

{prefect-client-3.1.12 → prefect-client-3.1.14}/src/prefect/_internal/schemas/validators.py

@@ -6,6 +6,8 @@ format.
 This will be subject to consolidation and refactoring over the next few months.
 """
 
+from __future__ import annotations
+
 import os
 import re
 import urllib.parse
@@ -627,18 +629,18 @@ def validate_name_present_on_nonanonymous_blocks(values: M) -> M:
 
 
 @overload
-def validate_command(v: str) -> Path:
+def validate_working_dir(v: str) -> Path:
     ...
 
 
 @overload
-def validate_command(v: None) -> None:
+def validate_working_dir(v: None) -> None:
     ...
 
 
-def validate_command(v: Optional[str]) -> Optional[Path]:
+def validate_working_dir(v: Optional[Path | str]) -> Optional[Path]:
     """Make sure that the working directory is formatted for the current platform."""
-    if v is not None:
+    if isinstance(v, str):
         return relative_path_to_current_platform(v)
     return v
 

{prefect-client-3.1.12 → prefect-client-3.1.14}/src/prefect/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2025-01-09T10:09:15-0800",
+ "date": "2025-01-23T13:22:04-0800",
  "dirty": true,
  "error": null,
- "full-revisionid": "e299e5a781867735d62685e7c190b5d100a28b62",
- "version": "3.1.12"
+ "full-revisionid": "5f1ebb57222bee1537b4d2c64b6dc9e3791ebafe",
+ "version": "3.1.14"
 }
 '''  # END VERSION_JSON
 

{prefect-client-3.1.12 → prefect-client-3.1.14}/src/prefect/artifacts.py

@@ -20,7 +20,10 @@ from prefect.logging.loggers import get_logger
 from prefect.utilities.asyncutils import sync_compatible
 from prefect.utilities.context import get_task_and_flow_run_ids
 
-logger = get_logger("artifacts")
+if TYPE_CHECKING:
+    import logging
+
+logger: "logging.Logger" = get_logger("artifacts")
 
 if TYPE_CHECKING:
     from prefect.client.orchestration import PrefectClient

{prefect-client-3.1.12 → prefect-client-3.1.14}/src/prefect/automations.py

@@ -137,7 +137,7 @@ class Automation(AutomationCore):
             self.id = client.create_automation(automation=automation)
             return self
 
-    async def aupdate(self: Self):
+    async def aupdate(self: Self) -> None:
         """
         Updates an existing automation.
 

{prefect-client-3.1.12 → prefect-client-3.1.14}/src/prefect/blocks/abstract.py

@@ -15,7 +15,7 @@ from typing import (
     Union,
 )
 
-from typing_extensions import Self, TypeAlias
+from typing_extensions import TYPE_CHECKING, Self, TypeAlias
 
 from prefect.blocks.core import Block
 from prefect.exceptions import MissingContextError
@@ -26,7 +26,10 @@ T = TypeVar("T")
 if sys.version_info >= (3, 12):
     LoggingAdapter = logging.LoggerAdapter[logging.Logger]
 else:
-    LoggingAdapter = logging.LoggerAdapter
+    if TYPE_CHECKING:
+        LoggingAdapter = logging.LoggerAdapter[logging.Logger]
+    else:
+        LoggingAdapter = logging.LoggerAdapter
 
 LoggerOrAdapter: TypeAlias = Union[Logger, LoggingAdapter]
 

{prefect-client-3.1.12 → prefect-client-3.1.14}/src/prefect/blocks/notifications.py

@@ -880,6 +880,7 @@ class SendgridEmail(AbstractAppriseNotificationBlock):
         sendgrid_block = SendgridEmail.load("BLOCK_NAME")
 
         sendgrid_block.notify("Hello from Prefect!")
+        ```
     """
 
     _description = "Enables sending notifications via Sendgrid email service."

{prefect-client-3.1.12 → prefect-client-3.1.14}/src/prefect/cache_policies.py

@@ -2,7 +2,7 @@ import inspect
 from copy import deepcopy
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, Literal, Optional, Union
 
 from typing_extensions import Self
 
@@ -73,17 +73,19 @@ class CachePolicy:
     def compute_key(
         self,
         task_ctx: TaskRunContext,
-        inputs: Dict[str, Any],
-        flow_parameters: Dict[str, Any],
+        inputs: dict[str, Any],
+        flow_parameters: dict[str, Any],
         **kwargs: Any,
     ) -> Optional[str]:
         raise NotImplementedError
 
     def __sub__(self, other: str) -> "CachePolicy":
+        "No-op for all policies except Inputs and Compound"
+
+        # for interface compatibility
         if not isinstance(other, str):  # type: ignore[reportUnnecessaryIsInstance]
             raise TypeError("Can only subtract strings from key policies.")
-        new = Inputs(exclude=[other])
-        return CompoundCachePolicy(policies=[self, new])
+        return self
 
     def __add__(self, other: "CachePolicy") -> "CachePolicy":
         # adding _None is a no-op
@@ -132,14 +134,14 @@ class CacheKeyFnPolicy(CachePolicy):
 
     # making it optional for tests
     cache_key_fn: Optional[
-        Callable[["TaskRunContext", Dict[str, Any]], Optional[str]]
+        Callable[["TaskRunContext", dict[str, Any]], Optional[str]]
     ] = None
 
     def compute_key(
         self,
         task_ctx: TaskRunContext,
-        inputs: Dict[str, Any],
-        flow_parameters: Dict[str, Any],
+        inputs: dict[str, Any],
+        flow_parameters: dict[str, Any],
         **kwargs: Any,
     ) -> Optional[str]:
         if self.cache_key_fn:
@@ -155,13 +157,30 @@ class CompoundCachePolicy(CachePolicy):
     Any keys that return `None` will be ignored.
     """
 
-    policies: List[CachePolicy] = field(default_factory=list)
+    policies: list[CachePolicy] = field(default_factory=list)
+
+    def __post_init__(self) -> None:
+        # flatten any CompoundCachePolicies
+        self.policies = [
+            policy
+            for p in self.policies
+            for policy in (p.policies if isinstance(p, CompoundCachePolicy) else [p])
+        ]
+
+        # deduplicate any Inputs policies
+        inputs_policies = [p for p in self.policies if isinstance(p, Inputs)]
+        self.policies = [p for p in self.policies if not isinstance(p, Inputs)]
+        if inputs_policies:
+            all_excludes: set[str] = set()
+            for inputs_policy in inputs_policies:
+                all_excludes.update(inputs_policy.exclude)
+            self.policies.append(Inputs(exclude=sorted(all_excludes)))
 
     def compute_key(
         self,
         task_ctx: TaskRunContext,
-        inputs: Dict[str, Any],
-        flow_parameters: Dict[str, Any],
+        inputs: dict[str, Any],
+        flow_parameters: dict[str, Any],
         **kwargs: Any,
     ) -> Optional[str]:
         keys: list[str] = []
@@ -178,6 +197,35 @@ class CompoundCachePolicy(CachePolicy):
             return None
         return hash_objects(*keys, raise_on_failure=True)
 
+    def __add__(self, other: "CachePolicy") -> "CachePolicy":
+        # Call the superclass add method to handle validation
+        super().__add__(other)
+
+        if isinstance(other, CompoundCachePolicy):
+            policies = [*self.policies, *other.policies]
+        else:
+            policies = [*self.policies, other]
+
+        return CompoundCachePolicy(
+            policies=policies,
+            key_storage=self.key_storage or other.key_storage,
+            isolation_level=self.isolation_level or other.isolation_level,
+            lock_manager=self.lock_manager or other.lock_manager,
+        )
+
+    def __sub__(self, other: str) -> "CachePolicy":
+        if not isinstance(other, str):  # type: ignore[reportUnnecessaryIsInstance]
+            raise TypeError("Can only subtract strings from key policies.")
+
+        inputs_policies = [p for p in self.policies if isinstance(p, Inputs)]
+
+        if inputs_policies:
+            new = Inputs(exclude=[other])
+            return CompoundCachePolicy(policies=[*self.policies, new])
+        else:
+            # no dependency on inputs already
+            return self
+
 
 @dataclass
 class _None(CachePolicy):
@@ -189,8 +237,8 @@ class _None(CachePolicy):
     def compute_key(
         self,
         task_ctx: TaskRunContext,
-        inputs: Dict[str, Any],
-        flow_parameters: Dict[str, Any],
+        inputs: dict[str, Any],
+        flow_parameters: dict[str, Any],
         **kwargs: Any,
     ) -> Optional[str]:
         return None
@@ -209,8 +257,8 @@ class TaskSource(CachePolicy):
     def compute_key(
         self,
         task_ctx: TaskRunContext,
-        inputs: Optional[Dict[str, Any]],
-        flow_parameters: Optional[Dict[str, Any]],
+        inputs: Optional[dict[str, Any]],
+        flow_parameters: Optional[dict[str, Any]],
         **kwargs: Any,
     ) -> Optional[str]:
         if not task_ctx:
@@ -236,8 +284,8 @@ class FlowParameters(CachePolicy):
     def compute_key(
         self,
         task_ctx: TaskRunContext,
-        inputs: Dict[str, Any],
-        flow_parameters: Dict[str, Any],
+        inputs: dict[str, Any],
+        flow_parameters: dict[str, Any],
         **kwargs: Any,
     ) -> Optional[str]:
         if not flow_parameters:
@@ -255,8 +303,8 @@ class RunId(CachePolicy):
     def compute_key(
         self,
         task_ctx: TaskRunContext,
-        inputs: Dict[str, Any],
-        flow_parameters: Dict[str, Any],
+        inputs: dict[str, Any],
+        flow_parameters: dict[str, Any],
         **kwargs: Any,
     ) -> Optional[str]:
         if not task_ctx:
@@ -273,13 +321,13 @@ class Inputs(CachePolicy):
     Policy that computes a cache key based on a hash of the runtime inputs provided to the task..
     """
 
-    exclude: List[str] = field(default_factory=list)
+    exclude: list[str] = field(default_factory=list)
 
     def compute_key(
         self,
         task_ctx: TaskRunContext,
-        inputs: Dict[str, Any],
-        flow_parameters: Dict[str, Any],
+        inputs: dict[str, Any],
+        flow_parameters: dict[str, Any],
         **kwargs: Any,
     ) -> Optional[str]:
         hashed_inputs = {}

{prefect-client-3.1.12 → prefect-client-3.1.14}/src/prefect/client/orchestration/_automations/client.py

@@ -45,6 +45,8 @@ class AutomationClient(BaseClient):
         return Automation.model_validate_list(response.json())
 
     def find_automation(self, id_or_name: "str | UUID") -> "Automation | None":
+        from uuid import UUID
+
         if isinstance(id_or_name, str):
             name = id_or_name
             try:
@@ -202,6 +204,8 @@ class AutomationAsyncClient(BaseAsyncClient):
         return Automation.model_validate_list(response.json())
 
     async def find_automation(self, id_or_name: "str | UUID") -> "Automation | None":
+        from uuid import UUID
+
         if isinstance(id_or_name, str):
             name = id_or_name
             try:

{prefect-client-3.1.12 → prefect-client-3.1.14}/src/prefect/client/orchestration/_deployments/client.py

@@ -1031,8 +1031,8 @@ class DeploymentAsyncClient(BaseAsyncClient):
         deployment_ids: list["UUID"],
         scheduled_before: "datetime.datetime | None" = None,
         limit: int | None = None,
-    ) -> list["FlowRunResponse"]:
-        from prefect.client.schemas.responses import FlowRunResponse
+    ) -> list["FlowRun"]:
+        from prefect.client.schemas.objects import FlowRun
 
         body: dict[str, Any] = dict(deployment_ids=[str(id) for id in deployment_ids])
         if scheduled_before:
@@ -1046,7 +1046,7 @@ class DeploymentAsyncClient(BaseAsyncClient):
             json=body,
         )
 
-        return FlowRunResponse.model_validate_list(response.json())
+        return FlowRun.model_validate_list(response.json())
 
     async def create_flow_run_from_deployment(
         self,

{prefect-client-3.1.12 → prefect-client-3.1.14}/src/prefect/client/utilities.py

@@ -5,7 +5,7 @@ Utilities for working with clients.
 # This module must not import from `prefect.client` when it is imported to avoid
 # circular imports for decorators such as `inject_client` which are widely used.
 
-from collections.abc import Awaitable, Coroutine
+from collections.abc import Coroutine
 from functools import wraps
 from typing import TYPE_CHECKING, Any, Callable, Optional, Union
 
@@ -61,8 +61,8 @@ def get_or_create_client(
 
 
 def client_injector(
-    func: Callable[Concatenate["PrefectClient", P], Awaitable[R]],
-) -> Callable[P, Awaitable[R]]:
+    func: Callable[Concatenate["PrefectClient", P], Coroutine[Any, Any, R]],
+) -> Callable[P, Coroutine[Any, Any, R]]:
     @wraps(func)
     async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
         client, _ = get_or_create_client()

{prefect-client-3.1.12 → prefect-client-3.1.14}/src/prefect/context.py

@@ -337,10 +337,15 @@ class EngineContext(RunContext):
         flow: The flow instance associated with the run
         flow_run: The API metadata for the flow run
         task_runner: The task runner instance being used for the flow run
-        task_run_futures: A list of futures for task runs submitted within this flow run
-        task_run_states: A list of states for task runs created within this flow run
         task_run_results: A mapping of result ids to task run states for this flow run
-        flow_run_states: A list of states for flow runs created within this flow run
+        log_prints: Whether to log print statements from the flow run
+        parameters: The parameters passed to the flow run
+        detached: Flag indicating if context has been serialized and sent to remote infrastructure
+        result_store: The result store used to persist results
+        persist_result: Whether to persist the flow run result
+        task_run_dynamic_keys: Counter for task calls allowing unique keys
+        observed_flow_pauses: Counter for flow pauses
+        events: Events worker to emit events
     """
 
     flow: Optional["Flow[Any, Any]"] = None
@@ -373,7 +378,7 @@ class EngineContext(RunContext):
     __var__: ClassVar[ContextVar[Self]] = ContextVar("flow_run")
 
     def serialize(self: Self, include_secrets: bool = True) -> dict[str, Any]:
-        return self.model_dump(
+        serialized = self.model_dump(
             include={
                 "flow_run",
                 "flow",
@@ -381,13 +386,18 @@ class EngineContext(RunContext):
                 "log_prints",
                 "start_time",
                 "input_keyset",
-                "result_store",
                 "persist_result",
             },
             exclude_unset=True,
-            serialize_as_any=True,
             context={"include_secrets": include_secrets},
         )
+        if self.result_store:
+            serialized["result_store"] = self.result_store.model_dump(
+                serialize_as_any=True,
+                exclude_unset=True,
+                context={"include_secrets": include_secrets},
+            )
+        return serialized
 
 
 FlowRunContext = EngineContext  # for backwards compatibility